Hugging Face Infini-Attention 재현 분석

Infini-Attention 메커니즘

Infini-Attention은 표준 자기 어텐션의 2차 메모리 증가를 고정 크기의 압축 메모리 버퍼로 대체함으로써 이론적으로 무한한 컨텍스트 길이를 달성하려고 합니다. 과정은 다음 기술 단계들을 따릅니다:

Segmentation: 입력 시퀀스를 고정 크기의 세그먼트로 나눕니다.
Local Attention: 현재 세그먼트 내에서 표준 인과 점곱 어텐션을 계산합니다.
Memory Retrieval: 모델은 현재 세그먼트의 쿼리 벡터($Q$)를 사용하여 압축 메모리 행렬($M_{s-1}$)에서 장기 컨텍스트를 검색하고, 비선형 활성화 함수(ELU + 1)를 통해 처리합니다.
Integration: 학습 가능한 스칼라 파라미터($\beta$)가 게이팅 메커니즘으로 작동하여 시그모이드 함수를 사용해 검색된 장기 메모리($A_{\text{mem}}$)와 로컬 점곱 어텐션($A_{\text{dot}}$)의 기여도를 균형 맞춥니다.
Memory Update: 압축 메모리는 현재 세그먼트의 키-값 상태를 기존 버퍼에 추가함으로써 업데이트됩니다.
State Transfer: 이전 세그먼트의 어텐션 상태는 버려지고, 업데이트된 압축 메모리만 다음 세그먼트에 전달됩니다.

재현 도전 과제 및 수렴 문제

게이팅 수렴

초기 실험에서는 균형 가중치의 약 95%가 0.5 주변에 집중되어 있어 게이팅 메커니즘이 수렴하지 않음을 나타냈습니다. 분석 결과, 표준 Llama 3 8B 하이퍼파라미터(학습률 $3.0 \times 10^{-4}$)가 균형 팩터가 초기값에서 크게 이동하도록 하기에는 충분하지 않았습니다.

이를 해결하기 위해 팀은 게이팅 함수에 별도의 높은 학습률(0.01)을 적용하고 전역 학습률은 $3.0 \times 10^{-4}$로 유지했습니다. 이로 인해 균형 팩터가 보다 이상적인 범위에 도달했지만, 200M 모델에서는 20B 토큰 이후 처음에 NaN 손실이 발생했습니다.

가중치 감소 및 롤아웃

추가 검토 결과, 가중치 감소가 균형 팩터의 작은 L2 노름을 유도하여 시그모이드 값이 0.5 주변에 집중되게 함을 발견했습니다. 이를 해결하기 위해 Hugging Face는 다음과 같은 변경을 적용했습니다:

Removed weight decay: 균형 팩터에서 가중치 감소를 제거했습니다.
Increased rollouts: 롤아웃을 16으로 늘렸으며(세그먼트 길이 64), 모델이 압축 메모리를 활용하도록 강력한 인센티브를 제공했습니다.

이러한 조정으로 전역 가중치가 0에서 1 전체 범위에 걸쳐 분포했으며, 헤드의 10%가 0.9에서 1.0 사이의 가중치를 달성했습니다.

실험 결과 및 평가

평가는 "패스키 검색 작업"을 사용하여 수행되었으며, 여기서 특정 "needle"(예: 패스키)이 컨텍스트의 다양한 위치에 있는 무관한 텍스트에 숨겨져 있습니다.

Small-scale (200M model): 초기 신호는 모델이 이전 세그먼트와 관련된 내용을 생성할 수 있음을 보여주었지만 정확한 needle을 검색하는 데 어려움을 겪었습니다.
Llama 3 8B: 모델은 초기에는 needle이 이전 세그먼트에 배치될 때 needle 평가에 실패했습니다.
Optimized Setup: 학습률을 조정하고 가중치 감소를 제거한 후, 모델은 이전 세그먼트에서 정확한 내용을 이어가는 능력이 향상되었으며 특정 프롬프트에서 일부 needle 패스키 테스트를 통과했지만 여전히 일관성이 부족했습니다.

결론 및 기술적 시사점

게이팅 수렴이 개선되었음에도 불구하고, Hugging Face는 Infini-Attention이 아직 프로덕션에 사용하기에 충분히 신뢰할 수 없다고 결론지었습니다. 팀은 Ring Attention, YaRN, 그리고 RoPE 스케일링이 사전 학습된 모델을 더 긴 컨텍스트로 확장하는 데 여전히 최고의 선택이라고 주장합니다.

핵심 기술 교훈:

Gating Sensitivity: 게이팅 함수의 적절한 수렴을 위해서는 특정 학습률 튜닝과 가중치 감소 제거가 필요하며, 이는 가중치가 0.5에 집중되는 것을 방지합니다.
Loss vs. Utility: 그래디언트 하강은 모델이 조건이 좋지 않거나 어텐션 출력에 차원 버그가 있더라도 훈련 손실을 감소시킬 수 있으므로, 손실만으로는 일관성을 평가하기에 충분하지 않으며 지속적인 평가가 필수적입니다.
Compression Trade-off: 메모리 압축 횟수가 증가함에 따라 성능이 본질적으로 저하됩니다.