Parallelizing Linear Transformers with the Delta Rule over Sequence Length

Transformers with linear attention (i.e., linear transformers) and state-space models have recently been suggested as a viable linear-time alternative to transformers with softmax attention. However, these models still underperform transformers especially on tasks that require in-context retrieval. While more expressive variants of linear transformers which replace the additive update in linear transformers with the delta rule (DeltaNet) have been found to be more effective at associative recall, existing algorithms for training such models do not parallelize over sequence length and are thus inefficient to train on modern hardware. This work describes a hardware-efficient algorithm for training linear transformers with the delta rule, which exploits a memory-efficient representation for computing products of Householder matrices. This algorithm allows us to scale up DeltaNet to standard language modeling settings. We train a 1.3B model for 100B tokens and find that it outperforms recent linear-time baselines such as Mamba and GLA in terms of perplexity and zero-shot performance on downstream tasks. We also experiment with two hybrid models which combine DeltaNet layers with (1) sliding-window attention layers every other layer or (2) two global attention layers, and find that these hybrids outperform strong transformer baselines.

한 줄 요약

DeltaNet의 학습을 시퀀스 길이에 대해 병렬화한 하드웨어 효율 알고리즘을 제안하고, 1.3B 파라미터 모델로 실험하여 Mamba와 GLA를 상회하는 성능을 보인다.

핵심 기여도

DeltaNet의 학습 알고리즘을 시퀀스 길이에 대해 병렬화하여 하드웨어 효율성을 향상.
Householder 행렬의 곱을 메모리 효율적으로 표현하여 I/O 비용 감소.
1.3B 파라미터 DeltaNet 모델을 100B 토큰으로 학습, Mamba와 GLA보다 낮은 퍼플렉시티와 높은 제로샷 성능 달성.
DeltaNet과 슬라이딩-윈도우 어텐션, 글로벌 어텐션의 하이브리드 모델이 강력한 트랜스포머 베이스라인을 상회.

핵심 아이디어

기존 DeltaNet은 델타 규칙(delta rule)을 사용하여 연관적 추출(associative recall)을 개선했으나, 학습 알고리즘이 시퀀스 길이에 대해 병렬화되지 않아 하드웨어 효율성이 낮았다. 본 연구는 DeltaNet을 Householder 변환 기반의 행렬 RNN으로 재파라미터화함으로써, 시퀀스 길이에 대한 병렬 처리가 가능하도록 했다. 특히, Householder 행렬의 곱을 WY 표현식(WY representation)을 사용해 메모리 효율적으로 표현함으로써, 기존의 메모리 비효율적인 히든 상태 저장 문제를 해결했다. 이는 DeltaNet의 대규모 학습을 가능하게 했으며, 100B 토큰 데이터셋에서 1.3B 파라미터 모델을 학습하는 데 성공했다는 점에서 의미가 있다.

기술적 접근법

DeltaNet을 Householder 변환 기반의 행렬 RNN으로 재구성.
Householder 행렬의 곱을 WY 표현식을 사용해 메모리 효율적으로 계산.
시퀀스 길이에 대해 병렬화된 forward/backward pass 알고리즘 구현.
FlashLinearAttention 라이브러리에 병렬 DeltaNet 레이어 포함 (https://github.com/fla-org/flash-linear-attention).
1.3B 파라미터 모델을 100B 토큰으로 학습.

주요 결과

1.3B DeltaNet 모델은 Mamba와 GLA 대비 더 낮은 퍼플렉시티를 기록.
제로샷 성능에서 Mamba와 GLA를 상회.
DeltaNet + 슬라이딩-윈도우 어텐션 하이브리드 모델이 강력한 트랜스포머 베이스라인을 상회.
DeltaNet + 글로벌 어텐션 하이브리드 모델도 동일하게 성능 향상.

의의 및 한계

본 연구는 DeltaNet의 학습 알고리즘을 하드웨어 효율적으로 병렬화함으로써, 대규모 언어 모델링에 적용 가능하도록 했다는 점에서 학술적·실용적 의의가 있다. 특히, Householder 변환과 WY 표현식을 활용한 메모리 최적화는 기존의 선형 어텐션 모델 학습에 새로운 접근법을 제시한다. 다만, DeltaNet의 성능은 여전히 최신 트랜스포머 베이스라인과 비교하면 다소 낮으며, 더 큰 규모의 모델에서의 확장성은 추가 연구가 필요하다는 한계가 있다.

실용적 활용

DeltaNet은 긴 컨텍스트에서의 연관 추출이 필요한 언어 모델, 정보 검색, 리트리벌-오그먼트드 제너레이션(RAG) 등에 적용 가능하다. 특히, 메모리 효율적인 학습과 추론이 요구되는 실시간 시스템이나 임베디드 환경에서 유용할 것으로 기대된다.