쿼리 비선형화로 보는 어텐션 혁신

초록

최근 연구에서 Query 가중치 행렬 W_Q 가 아이덴티티로 대체돼도 성능 손실이 없다는 사실을 바탕으로, 저자는 Query에 비선형 잔차 MLP를 삽입한 새로운 설계를 제안한다. GPT‑3‑small 규모 실험에서 비선형 Query가 기존 모델 및 파라미터를 12.5 % 늘린 MLP‑확장 모델을 모두 능가하며, 학습 안정성도 향상됨을 보여준다.

상세 분석

이 논문은 트랜스포머 어텐션에서 Q, K, V 세 개의 선형 투사 행렬이 서로 독립적이지 않으며, 특히 W_Q 는 거의 완전한 자유도가 없다는 수학적 사실을 이용한다. Karbevski와 Mijoski(2025)의 정리(Θ = W_Q)를 적용하면 W_Q 를 아이덴티티 I 로 바꿔도 전체 네트워크 함수가 변하지 않는다. 이는 W_Q 가 파라미터 차원에서 중복된다는 의미이며, 대신 Q 경로에 비선형성을 부여해 모델 표현력을 높일 여지를 만든다.

저자는 Q(X) = (X + f_θ(X))/2 라는 잔차 형태를 채택한다. 여기서 f_θ 는 차원 d → r (d/2) → d 구조의 bottleneck MLP이며, GELU와 RMSNorm, LayerNorm을 조합해 안정성을 확보한다. 파라미터 수는 d² + O(d) 로 기존 W_Q 와 동일하고, FLOPs도 크게 증가하지 않는다.

실험은 NanoGPT 기반 GPT‑3‑small(12 layer, d=768) 모델을 60 k 스텝(≈29 B 토큰) 동안 학습시켰다. 비교 대상은 (1) 기본 모델, (2) MLP‑폭확장 모델(숨겨진 차원 4.75 d, 파라미터 12.5 % 증가). 비선형 Query 모델은 학습률을 3× 높이고 가중치 감쇠를 0.03으로 낮추어도 안정적으로 수렴했으며, 최종 검증 손실 2.915(1.40 % 개선)를 기록했다. 특히 초기 워밍업 단계와 후반 학습 단계에서 손실 감소가 두드러졌다.

제한점으로는 단일 규모·단일 시드 실험, 추론 속도 미측정, K와 V에 대한 비선형 확장은 아직 검증되지 않았음이 언급된다. 또한, 비선형 Query가 모델 깊이·폭이 커짐에 따라 더 큰 이점을 제공할 가능성이 있지만, 현재는 보수적인 설계(스케일링 ½, 1/2 스케일링)로 안정성을 우선시했다.

향후 연구 방향은 (1) 다양한 정규화·활성화 함수 탐색, (2) K·V에도 비선형 잔차 적용, (3) 대규모·다중 시드 실험을 통한 일반화 검증, (4) 커스텀 커널 구현을 통한 추론 효율성 개선, (5) 사전학습 모델에 대한 LoRA‑유사 비선형 적응 기법 적용 등이 제시된다.

전반적으로 이 논문은 트랜스포머 구조에서 “선형 W_Q 는 불필요”라는 이론적 통찰을 실용적인 비선형 설계로 전환함으로써, 파라미터 효율성을 유지하면서도 성능과 학습 안정성을 동시에 향상시킬 수 있음을 실증한다.