Google의 Infini-Attention을 Llama 3 8B에 구현한 결과, 메모리 압축 횟수 증가에 따라 성능이 저하되어 Ring Attention과 YaRN이 더 나은 선택지임을 실증

Context

언어모델의 컨텍스트 길이 확장 시 표준 어텐션의 메모리 복잡도가 O(seq_len²)로 증가하여, 시퀀스 길이 10배 증가 시 메모리 요구량이 100배 증가한다. Llama 3 8B를 1백만 토큰 컨텍스트로 확장하려면 512개 GPU가 필요한 리소스 제약이 존재한다.

Technical Solution

• Infini-Attention 구현: 시퀀스를 세그먼트로 분할하고 이전 세그먼트를 고정 크기 버퍼로 압축하여 총 메모리 사용량 상한선 설정
• 세그먼트 내 어텐션 제한: 현재 세그먼트 내 단어 간 어텐션과 압축된 메모리에서의 정보 검색을 동일 쿼리로 수행
• 게이팅 메커니즘 조정: 수렴 문제 해결을 위해 게이팅 훈련 신호 개선
• 기준 모델 구축: Ring Attention, YaRN, RoPE Scaling과의 성능 비교를 위한 기준선 설정
• 재현 원칙 적용: 최소 모델 크기부터 시작 → 신호 확인 후 확장 → 단일 변수만 수정하여 영향 측정

Impact

Infini-Attention의 장기 컨텍스트 성능은 메모리 압축 횟수가 증가할수록 감소한다. 게이팅 수렴 개선 후에도 성능이 충분하지 않았다.

Key Takeaway

신경망 훈련에서 손실함수 감소만으로는 충분하지 않으며, 실제 평가 실행이 필수이다. 차원 불일치로 인한 버그는 손실 감소 추세를 보이면서도 생성 성능을 완전히 훼손할 수 있으므로, 구현 검증 시 수학적 정확성과 실제 동작을 모두 확인해야 한다.