abstract: Pretrained transformer models have demonstrated remarkable performance across various natural language processing tasks. These models leverage the attention mechanism to capture long- and short-range dependencies in the sequence. However, the (full) attention mechanism incurs high computational cost – quadratic in the sequence length, which is not affordable in tasks with long sequences, e.g., inputs with 8k tokens. Although sparse attention can be used to improve computational efficiency, as suggested in existing work, it has limited modeling capacity and often fails to capture complicated dependencies in long sequences. To tackle this challenge, we propose MASFormer, an easy-to-implement transformer variant with Mixed Attention Spans. Specifically, MASFormer is equipped with full attention to capture long-range dependencies, but only at a small number of layers. For the remaining layers, MASformer only employs sparse attention to capture short-range dependencies. Our experiments on natural language modeling and generation tasks show that a decoder-only MASFormer model of 1.3B parameters can achieve competitive performance to vanilla transformers with full attention while significantly reducing computational cost (up to 75%). Additionally, we investigate the effectiveness of continual training with long sequence data and how sequence length impacts downstream generation performance, which may be of independent interests.
TL;DR
(문제 제기: 일반적인 바닐라 어텐션의 Time Complexity 문제) Pretrained transformer model은 attention mechanism을 활용하여 시퀀스의 장/단거리 종속성을 포착하지만, (전체) attention mechanism은 시퀀스 길이에 따라 선형적으로 높은 계산 비용이 발생하므로(바닐라 어텐션의 일반적인 문제, Zephyr 및 Mistral, Qwen 참고), 8천 개의 토큰이 입력된 긴 시퀀스가 있는 작업에서는 적합하지 않음.
(선행 연구 중 sparse attention의 문제) 기존 연구에서 제안된 바와 같이 sparse attention은 계산 효율성을 개선하는 데 사용할 수 있지만, 모델링 용량이 제한적이며 긴 시퀀스에서 복잡한 종속성을 포착하지 못하는 경우가 많음.
(아키텍처 제안: Mixed Attention Spans, MASFormer) 이런 문제를 해결하기 위해 Mixed Attention Spans를 사용하여 구현하기 쉬운 트랜스포머 변형인 MASFormer를 제안.
(MASFormer 아키텍처 요약)MASFormer는 멀리 떨어진 거리의 종속성을 캡처할 수 있는 full attention을 갖추고 있지만, 소수의 레이어에서만 사용 가능하고, 나머지 레이어에 대해서는 단거리 종속성을 포착하기 위해 sparse attention만 사용.
Full attention: to capture long-range dependencies (only at a small number of layers)
Sparse attention: to capture short-range dependencies
(성능 검증) 1.3B Decoder 모델에서 Full attention을 사용하는 바닐라 트랜스포머와 동일한 성능.
자연어 모델링 및 생성 작업에 대한 실험 결과, 13억 개의 파라미터로 구성된 디코더 전용 MASFormer 모델(1.3B Decoder only model)이 full attention을 사용하는 바닐라 트랜스포머와 동등한 성능을 달성하면서도 계산 비용을 크게 절감(최대 75%)
긴 시퀀스 데이터에 대한 지속적인 훈련의 효과와 시퀀스 길이가 downstream 생성 성능에 미치는 영향 추가 조사할 예정이라고 함.