무한 입력이 가능한 Infini-attention

어텐션 메커니즘의 기본 개념은 입력 시퀀스에서 출력을 예측할 때 전체 입력이 아닌 관련 부분에만 집중하는 것입니다. 예를 들어 기계번역을 할 때, 소스 문장의 특정 단어들에 주의(어텐션)를 기울여 타겟 문장을 생성합니다.

어텐션 메커니즘은 다음 세 가지 입력을 받습니다:

1. Query: 현재 상황에 대한 정보
2. Key: 모든 입력 토큰들의 정보
3. Value: 각 입력 토큰들의 실제 값

그리고 Query와 Key 벡터 사이의 유사도를 측정하여 Value에 가중치를 부여합니다. 이렇게 구한 가중치 합이 출력 벡터가 됩니다.

트랜스포머에서는 Self-Attention이라는 특수한 어텐션 메커니즘을 사용합니다. 여기서 Query, Key, Value는 모두 동일한 시퀀스에서 온 임베딩 벡터입니다. 각 단어가 다른 모든 단어들에 주의를 기울이면서 문맥을 파악하는 방식입니다.

하지만 이러한 모델들은 scaled dot-product attention 계층에서 키-값 상태를 저장하기 때문에 입력 길이가 길어질수록 메모리 사용량과 계산 복잡도가 급격히 증가하는 문제가 있었습니다.

Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

이 리서치는 기존 트랜스포머 기반 대규모 언어모델이 가진 제한된 컨텍스트 길이 문제를 해결하기 위해 새로운 어텐션 메커니즘인 Infini-attention을 소개합니다.

Infini-attention은 압축 메모리 모듈을 기존 scaled dot-product attention 메커니즘에 통합하였습니다. 구체적으로 현재 입력에 대해서는 일반 로컬 어텐션을 계산하고, 이전 입력의 정보는 압축 메모리에 누적 저장한 후 필요시 복원하여 두 정보를 결합하는 방식입니다. 압축 메모리의 업데이트와 복원 과정에는 선형 어텐션 메커니즘의 계산식을 차용하여 효율성을 높였습니다.

Scaled Dot-Product Attention:
Q, K, V = XWQ, XWK, XWV
Adot = softmax(QKT/√dmodel)V

압축 메모리 복원:
Amem = σ(Q)M(s-1) / (σ(Q)z(s-1))

압축 메모리 업데이트:
M(s) = M(s-1) + σ(K)T(V - σ(K)M(s-1)/(σ(K)z(s-1)))
z(s) = z(s-1) + Σσ(K)

최종 출력:
A = sigmoid(β) ⊙ Amem + (1 - sigmoid(β)) ⊙ Adot
(β: 학습되는 gating scalar 가중치)

σ: 비선형 활성화 함수(ELU+1)
M: 압축 메모리
z: 정규화 텀