근사 어텐션 (Approximate Attention)을 위한 I/O 최적성 접근

우리는 대규모 언어 모델 (Large Language Models, LLMs)에서 어텐션 (Attention)의 I/O 복잡성 (I/O complexity)을 재검토합니다. 쿼리-키-값 (Query-Key-Value) 행렬 $Q,K,V o ext{R}^{n imes d}$와 빠른 메모리 크기가 $M$인 머신이 주어졌을 때, 목표는 빠른 메모리 (Fast memory)와 느린 메모리 (Slow memory) 사이의 데이터 전송 횟수를 최소화하면서 "어텐션 행렬 (Attention matrix)" $A= ext{softmax}(Q K ^{ op}/\sqrt{d}) V$를 계산하는 것입니다. 기존 문헌의 방법들, 특히 가장 유명한 FlashAttention 및 그 변형들은 $n$에 대해 이차적 (Quadratically)으로 의존하는 I/O 비용을 발생시키는 반면, 단순한 하한선 (Lower bound)은 입력을 읽고 출력을 쓰는 데 $Ω(nd)$의 I/O만을 요구합니다. 본 연구에서는 대부분의 파라미터 영역 (Parameter regimes)에서 I/O 비용이 $n$에 대해 거의 선형적 (Almost-linearly)으로만 의존하는 어텐션 계산 기술을 제시합니다. 이는 최근 Alman과 Song의 근사 어텐션 (Approximate attention) 프레임워크에서 영감을 얻은 I/O 효율적인 알고리즘을 개발함으로써 달성되었습니다. 또한 우리는 각 파라미터 영역에서 상응하는 하한선을 증명하여 우리의 알고리즘이 실제로 I/O 최적 (I/O-optimal)에 근접함을 보여줍니다.