자유 에너지 믹서 채널별 선택을 위한 무손실 어텐션

초록

본 논문은 기존 어텐션이 키‑값 캐시를 손실 없이 저장하지만, 모든 채널에 동일한 가중치를 적용해 값을 읽는다는 근본적인 한계를 지적한다. 이를 해결하기 위해 자유 에너지(FEM)라는 새로운 읽기 연산을 제안한다. FEM은 쿼리·키로부터 얻은 빠른 사전 분포에 값‑주도 로그선형 기울기를 적용해 각 채널마다 별개의 후방 분포를 만들며, 온도 파라미터를 학습함으로써 평균 혼합에서 채널별 선택으로 부드럽게 전환한다. 복잡도는 기존 소프트맥스 어텐션(O(T²))이나 선형 어텐션(O(T))와 동일하게 유지된다.

상세 분석

논문은 먼저 “손실 없는 저장 vs 손실 있는 처리”라는 개념을 도입한다. KV‑cache는 시퀀스 전체를 그대로 보관하지만, 전통적인 어텐션은 각 헤드마다 하나의 확률 벡터 αₜ를 계산해 모든 채널에 동일하게 적용한다. 이때 출력 oₜ는 값들의 볼록합(convex hull) 안에 머물게 되며, 채널별로 서로 다른 인덱스를 선택하는 것이 불가능해진다. 저자들은 이를 “채널‑와이즈 셀렉터”라는 정의로 정형화하고, Lemma 2.2와 Corollary 2.3을 통해 기존 어텐션이 일반적인 채널‑와이즈 선택을 구현하지 못함을 수학적으로 증명한다.

기존의 해결책—헤드 수 증가, 깊이 증가, 차원별 쿼리·키, 복잡한 인‑헤드 믹서, 선형 RNN/SSM—을 체계적으로 분석한 결과, 모두 시간 복잡도 상승이나 저장 효율 저하, 혹은 여전히 토큰‑분리적인 읽기 구조에 머무르는 한계가 있음을 보여준다. 특히 헤드·깊이 확장은 채널 수 D에 비해 H·L이 충분히 크지 않으면 채널‑와이즈 선택 용량이 급격히 부족하다는 Proposition 2.6을 제시한다.

이러한 배경에서 FEM은 “사전‑후방 자유 에너지 최적화”라는 새로운 프레임을 제시한다. 사전 pₜ는 기존 어텐션의 소프트맥스 혹은 선형 커널 정규화에 의해 빠르게 계산된다. 각 채널 j마다 값 v_{i,j}를 이용해 로그선형 기울기 β_{t,j}를 곱한 뒤, pₜ와 결합해 로그합(exp) 형태의 자유 에너지

F_{t,j}(β) = (1/β) log ∑{i∈M_t} pₜ(i) exp(β v{i,j})

를 계산한다. 이 식은 KL 제약을 가진 변분 최적화의 해와 동일함을 Theorem 2.8이 증명한다. β가 커질수록 후방 분포 q_{t,β}는 argmax에 집중하고, F_{t,j}(β)는 채널별 최대값에 수렴한다. 따라서 β가 학습 가능한 역온도라면, 모델은 자동으로 평균 혼합에서 고정밀 채널‑와이즈 선택으로 전이한다.

복잡도 측면에서 FEM은 각 채널당 하나의 마스크드 로그합을 추가로 수행하지만, 이는 O(T) 혹은 O(T²)와 같은 기존 어텐션의 시간 복잡도에 상수배만큼 영향을 미친다. 또한 두 단계 게이팅(내부 λ, 외부 g)을 도입해 평균값 µₜ와 최대값 Fₜ를 적절히 혼합함으로써 학습 안정성을 확보한다. 구현은 저차원 로컬 컨볼루션(C)과 저랭크 선형 변환을 이용해 효율성을 높였으며, 기존 트랜스포머 블록에 플러그‑인 형태로 삽입 가능하도록 설계되었다.

실험에서는 NLP(언어 모델링, 텍스트 분류), 비전(이미지 분류, 비디오 인식), 시계열(멀티베리어 데이터) 등 다양한 도메인에서 동일 파라미터 예산 하에 기존 소프트맥스 어텐션, 라스팅 어텐션, 선형 어텐션, Mamba와 같은 SSM 기반 모델들을 능가하는 성능 향상을 보였다. 특히 채널‑와이즈 구조가 중요한 멀티베리어 시계열 예측에서 FEM은 기존 모델 대비 2~4% 이상의 절대 정확도 향상을 기록했다. 코드와 사전 훈련된 체크포인트는 공개 리포지터리에서 제공된다.

요약하면, FEM은 “선택 사전 + 값‑주도 자유 에너지 후방”이라는 새로운 읽기 메커니즘을 통해 어텐션의 근본적인 손실을 제거하면서도 복잡도와 메모리 효율성을 유지한다는 점에서 의미가 크다. 이는 차원별 선택이 요구되는 다양한 실제 응용에 바로 적용 가능하며, 앞으로 어텐션 기반 모델 설계에 새로운 패러다임을 제시한다.