MALA: 어텐션 기여도에 따라 계산량을 배분하다

긴 컨텍스트 어텐션의 비용은 QK 점수 행렬을 만드는 데서 끝나지 않는다. 점수 계산 이후에도 밀집 커널은 모든 합법적인 인과 상호작용에 대해 동일한 후속 연산을 수행한다. 하지만 최종적으로 정규화된 어텐션 질량이 거의 없는 위치도 많다. MassAlloc Attention, 즉 MALA는 어텐션 분포 자체를 활용해 어디에 후속 계산을 계속 투입할지 결정한다.

작동 방식

MALA는 점수 계산 전에 인과적 연결을 고정적으로 제거하지 않는다. 먼저 합법적인 인과 범위 전체에 대한 QK 점수 접근을 유지한다. 그다음 softmax 정규화 과정에서 드러나는 기여도를 사용해, 영향이 작은 영역의 post-score 계산을 줄인다.

• 순방향 계산: 처리 과정에서 계속 갱신되는 온라인 softmax 정규화값으로 각 영역의 기여도를 추정한다.
• 역전파: 최종 정규화값을 재사용해 표준 어텐션 상태만으로 중첩된 유지 영역을 계산한다.
• 공통 허용오차: 학습과 추론에 같은 허용 기준을 적용해 현재 어텐션 분포에 따라 작업량을 조절한다.
• 융합 연산 프리미티브: 동적 선택을 어텐션 커널에 결합해, 줄어든 후처리 작업이 실제 하드웨어 성능 향상으로 이어지도록 설계한다.

실험 결과

8K 컨텍스트에서 전체 post-score 작업량을 동일하게 맞춘 실험에서 MALA가 누락한 평균 어텐션 질량은 0.0188%였다. 인스턴스별 최종 질량을 알고 있다고 가정하는 참조 오라클의 0.0182%와 근접한 수치다. 총 작업량을 통제했기 때문에, 이 비교는 단순히 더 많은 계산을 수행한 효과가 아니라 분포에 맞춰 계산을 배분하는 능력을 보여준다.

1K부터 32K까지 컨텍스트 길이를 확장했을 때도 같은 허용오차로 FullAttn 대비 출력 오차와 그래디언트 오차를 낮게 유지했다. 더 넓은 연관 회수 평가에서 8K 컨텍스트 기준 MALA의 정확도는 89.67%, FullAttn은 89.97%였다. 보고된 실험 범위에서는 후속 계산을 줄여도 능력 저하가 크지 않았다.

시스템 성능도 측정됐다. 8개의 H100과 TP=8 환경에서 128K 토큰 어텐션 연산자를 비교한 결과, MALA는 FullAttn보다 학습 순방향에서 2.2배, 역전파에서 3.0배, 추론 디코딩에서 1.6배 빠른 결과를 보였다. 0.6B부터 14B까지의 모델을 대상으로 한 스케일링 실험에서는 14B 모델의 32K 컨텍스트 학습 시 전체 학습 FLOPs를 23.1% 줄이면서 평가된 능력을 FullAttn과 비슷하게 유지했다.

의미와 한계

MALA의 핵심은 고정된 희소 패턴을 가정하지 않는다는 점이다. 각 입력과 쿼리에서 실제로 형성된 어텐션 분포를 보고 중요도가 낮은 부분에 투입할 계산을 줄인다. 따라서 긴 과거 문맥 중 대부분이 현재 출력에 거의 영향을 주지 않는 경우, 메모리 접근뿐 아니라 점수 이후의 행렬 연산과 학습 역전파 비용도 절감할 수 있다.

다만 동적 유지 영역을 판단하는 제어 비용이 크면 절감 효과가 줄어들 수 있다. 실제 이득은 모델 구조, 작업, 컨텍스트 길이, 어텐션 분포의 집중도에 따라 달라질 가능성이 있다. 따라서 MALA는 새로운 어텐션 수식이라기보다, 기존 어텐션을 분포 적응형으로 실행하기 위한 커널 및 시스템 설계로 보는 편이 정확하다.

출처: Hugging Face Daily Papers