어텐션의 계산 효율을 높이는 MALA 메커니즘 제안
원제 MassAlloc Attention: Let Attention Allocate Its Own Compute
추천 63댓글 1
Key Point
Full Attention의 비효율적 계산 패턴을 식별하고 실제 기여도 기반 할당으로 2~3배 속도 향상을 달성하면서도 모델 성능을 유지하는 방법으로, 대규모 언어 모델의 추론·학습 효율을 개선하는 데 직결된다.
핵심 요약
- 기존 Full Attention은 QK 타일 계산 후 전체 후속 경로를 실행하지만 정규화 가중치에서 대부분의 값이 무시할 수준으로 작아 비효율적이다.
- MALA(MassAlloc Attention)는 정규화된 기여도에 따라 후속 계산을 선택적으로 할당하는 퓨즈드 어텐션 프리미티브로, 순전파는 온라인 소프트맥스 정규화기를, 역전파는 최종 정규화기를 재사용한다.
- 하나의 공통 임계값으로 학습과 추론에서 적응적 유지(retention)가 가능하며, 8K 토큰 비교에서 생략된 질량은 0.0188%로 기준 오라클의 0.0182%에 근접한다.
- 1K에서 32K 토큰 범위에서 같은 임계값으로 출력과 그래디언트 오류를 낮게 유지하며, 8K에서 연관성 회상 정확도는 89.67%로 Full Attention의 89.97%와 유사하다.
- 128K 토큰에서 텐서 병렬화 적용 시 학습 중 순전파는 2.2배, 역전파는 3.0배 빠르고 추론 디코딩은 1.6배 빠르다.
- 0.6B에서 14B 파라미터 모델까지 확장 법칙 학습에서 MALA는 Full Attention과 유사한 퍼플렉시티를 유지하면서 총 학습 FLOPs를 감소시킨다.
- 별도 계속 학습으로 생성된 14B, 32B 모델이 Full Attention과 비슷한 지식·추론·장문맥 검색 점수를 달성하며, 정규화된 어텐션 기여도에 따른 계산 할당으로 Full Attention의 성능을 유지하면서 어텐션 계산량을 줄일 수 있음을 보여준다.