효율적인 선형 시퀀스 모델링을 위한 게이트형 슬롯 어텐션

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

저자: Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, Guohong Fu

cs.AI

초록

선형 주의 트랜스포머와 그들의 게이트 변형은 병렬 훈련과 효율적인 순환 추론을 가능하게 함으로써 칭송받았지만, 전통적인 트랜스포머에 비해 회상 집중적 작업에서는 아직 부족하며, 처음부터 훈련에 상당한 자원을 요구합니다. 본 논문은 게이트형 슬롯 주의 (GSA)를 소개하는데, 이는 게이트 선형 주의 (GLA)에서 영감을 받은 게이팅 메커니즘을 통해 주의를 Bounded-memory-Control (ABC)로 향상시킴으로써 주의를 강화합니다. 기본적으로 GSA는 소프트맥스를 통해 연결된 두 층 GLA로 구성되며, 컨텍스트 인식 메모리 읽기와 적응적인 잊기를 활용하여 기억 용량을 향상시키면서도 간결한 순환 상태 크기를 유지합니다. 이 설계는 GLA의 하드웨어 효율적인 훈련 알고리즘과 상태 크기 축소를 통해 훈련과 추론 효율성을 크게 향상시킵니다. 또한 소프트맥스 연산을 유지하는 것은 "미세 조정된 사전 훈련된 트랜스포머를 RNN에 미세 조정" (T2R) 설정에서 특히 유익하며, 처음부터의 광범위한 훈련 필요성을 줄입니다. 광범위한 실험은 GSA가 문맥 기억과 T2R 설정에서 우수한 성능을 보여주는 것을 확인합니다.

English

Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared to traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with Bounded-memory-Control (ABC) by incorporating a gating mechanism inspired by Gated Linear Attention (GLA). Essentially, GSA comprises a two-layer GLA linked via softmax, utilizing context-aware memory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and inference efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the softmax operation is particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R) settings, reducing the need for extensive training from scratch. Extensive experiments confirm GSA's superior performance in scenarios requiring in-context recall and in T2R settings.