텐서곱 어텐션: 필요한 모든 것
2025-01-22

더 긴 입력 시퀀스를 처리하는 언어 모델의 확장에는 일반적으로 대규모 키-값(KV) 캐시가 필요하며, 이로 인해 추론 중 메모리 오버헤드가 크게 증가합니다. 본 논문에서는 텐서 분해를 사용하여 쿼리, 키, 값을 간결하게 표현함으로써 추론 중 KV 캐시 크기를 크게 줄이는 새로운 어텐션 메커니즘인 텐서곱 어텐션(TPA)을 제안합니다. 이러한 표현을 컨텍스트 종속 저랭크 구성 요소(컨텍스트 분해)로 분해하고 RoPE와 원활하게 통합함으로써 TPA는 메모리 효율성을 유지하면서 모델의 품질을 향상시킵니다. TPA를 기반으로 시퀀스 모델링을 위한 새로운 모델 아키텍처인 텐서곱 어텐션 트랜스포머(T6)를 소개합니다. 언어 모델링 작업에 대한 광범위한 실험적 평가를 통해 T6이 퍼플렉서티와 다양한 유명한 평가 벤치마크를 포함한 다양한 지표에서 MHA, MQA, GQA, MLA와 같은 표준 트랜스포머 기준 모델을 능가함을 보여줍니다. 특히 TPA의 메모리 효율성을 통해 고정된 리소스 제약 조건 하에서 훨씬 더 긴 시퀀스를 처리할 수 있으므로 최신 언어 모델의 중요한 확장성 문제를 해결합니다. 코드는 공개되어 있습니다.