[논문 리뷰] Block-Recurrent Transformers

DeLesley Hutchins, Imanol Schlag외 3명|arXiv (Cornell University)|2022. 03. 11.

Topic Modeling인용 수 27

한 줄 요약

블록-재귀 트랜스포머는 블록 단위로 트랜스포머 레이어를 재귀적으로 적용하여 긴 시퀀스를 처리하며, 시퀀스 길이에 대해 선형 복잡도를 달성하고 유사하거나 낮은 계산으로 긴 문서에서 더 큰 perplexity 개선을 달성한다.

ABSTRACT

We introduce the Block-Recurrent Transformer, which applies a transformer layer in a recurrent fashion along a sequence, and has linear complexity with respect to sequence length. Our recurrent cell operates on blocks of tokens rather than single tokens during training, and leverages parallel computation within a block in order to make efficient use of accelerator hardware. The cell itself is strikingly simple. It is merely a transformer layer: it uses self-attention and cross-attention to efficiently compute a recurrent function over a large set of state vectors and tokens. Our design was inspired in part by LSTM cells, and it uses LSTM-style gates, but it scales the typical LSTM cell up by several orders of magnitude. Our implementation of recurrence has the same cost in both computation time and parameter count as a conventional transformer layer, but offers dramatically improved perplexity in language modeling tasks over very long sequences. Our model out-performs a long-range Transformer XL baseline by a wide margin, while running twice as fast. We demonstrate its effectiveness on PG19 (books), arXiv papers, and GitHub source code. Our code has been released as open source.

연구 동기 및 목표

• 긴 거리의 기억을 유지하면서 어텐션 비용이 제곱적으로 증가하지 않는 트랜스포머 아키텍처를 동기부여하고 개발한다.
• 토큰과 상태 벡터의 블록에서 작동하는 순환 단위를 도입하여 긴 컨텍스트 모델링을 가능하게 한다.
• 표준 트랜스포머와 계산적 동등성을 유지하면서 긴 시퀀스 언어 모델링 작업에서 유의미한 perplexity 개선을 달성한다.

제안 방법

• 길이 N의 구간을 W 크기의 블록으로 처리하는 블록 단위 재발과 함께 슬라이딩 윈도우(인과) 어텐션을 사용한다.
• 수직 방향(자체 어텐션 + 재귀 상태에 대한 교차 어텐션)과 수평 방향(상태에서 상태로의 자체 어텐션과 입력에 대한 교차 어텐션)을 결합하고 정보를 잊거나 기억하기 위한 게이트를 가진 순환 단위를 구현한다.
• 수직과 수평 방향 간에 키와 값을 공유하되 쿼리는 분리된 상태로 유지한다.
• 긴 시퀀스에서 학습의 안정성을 확보하기 위해 학습된 상태 ID와 상대 위치 바이어스를 도입한다.
• 게이트 구성(고정 게이트, LSTM 스타일 게이트)을 실험하고 안정성, 스케일링 및 애블레이션을 평가한다.

실험 결과

연구 질문

• RQ1 트랜스포머의 블록 단위 재발이 매우 긴 맥락(예: 60k 토큰)의 기억을 제공하고 Transformer-XL 기준선과 비교하여 perplexity가 경쟁력 있거나 더 나은가?
• RQ2 재발 게이트 유형과 상태 크기가 perplexity, 학습 안정성 및 계산 비용에 미치는 영향은 무엇인가?
• RQ3 하나의 순환 계층과 다수의 비순환 계층 추가 간의 성능과 비용은 긴 문서 데이터셋에서 어떻게 비교되는가?

주요 결과

모델	세그먼트 길이	윈도우 길이	스텝 시간	PG19 (바이트/토큰)	arXiv (토큰)	GitHub (토큰)
XL:512	512	512	0.88	1.01	3.62±0.01	1.45	1.21
XL:1024	1024	1024	1.20	0.997	3.59±0.01	1.37	1.08
XL:2048	2048	2048	2.11	0.990	3.58±0.01	1.31	1.01
Slide:12L	4096	512	0.93	0.989	3.60	1.43	1.19
Slide:13L	-	-	1.00	0.989	3.58±0.01	1.42	1.17
Rec:lstm:dual	4096	512	1.06	0.985	3.54±0.01	1.26	1.01
Rec:lstm:single	-	-	1.05	0.962	3.54±0.01	1.29	1.03
Rec:lstm:skip	-	-	1.00	0.969	3.56±0.01	1.31	1.10
Rec:fixed:dual	-	-	1.01	0.957	3.52±0.01	1.27	0.991
Rec:fixed:single	-	-	1.02	0.966	3.58±0.01	1.25	1.00
Rec:fixed:skip	-	-	0.99	0.952	3.53±0.01	1.24	0.976
Feedback:lstm:single	4096	512	1.40	0.977	3.50	1.22	-
Feedback:fixed:skip	-	-	1.35	0.935	3.49	1.24	-
Memorizing Trans. 64k	512	512	1.94	0.950	3.53	1.22	-

• 단일 계층의 재발은 비순환 계층을 추가하는 것보다 더 큰 perplexity 이득을 얻는다.
• PG19, arXiv, GitHub에서 블록-재귀 변형은 유사하거나 더 낮은 스텝 시간으로 Transformer-XL 기준선을 능가하며, 더 강한 넓은 컨텍스트 기준선보다 더 빠르게 실행될 수 있다.
• 최적 구성(Rec:fixed:skip)은 종종 다른 설정을 능가하고 XL:2048를 상회하면서 두 배 이상 빠르게 실행되기도 한다; 또한 1.3B 매개변수 모델로 PG19에서 단어 수준 perplexity 26.50로 최첨단 성능에 도달한다.
• 스케일링 결과는 재발이 모델 크기에 걸쳐 이점을 제공하며 매개변수의 두 배 이상 커짐에 따라 이득이 커진다.
• 하나의 계층으로 재발하는 것이 일반적으로 대부분의 이점을 얻기에 충분하며, 하나의 재발 계층 이상 추가하면 이익은 제한적이거나 데이터셋에 따라 다르다.
• 블록-재귀 모델은 특수 CUDA 커널 없이도 베이스라인의 성능에 맞추거나 능가할 수 있으며 오픈소스 릴리스로 제공된다.