라플라시안 메커니즘으로 토큰 기하학을 재구성한 트랜스포머 성능 향상

초록

본 논문은 기존 어텐션에 라플라시안 메커니즘을 도입해 토큰의 분산을 직접 조절함으로써, 비전·언어 벤치마크에서 일관된 정확도 향상을 달성하고, 토큰 임베딩이 클래스별로 붕괴(NC)하는 새로운 “Neural Token Collapse” 형태로 기하학적 구조를 재편한다는 것을 실증한다.

상세 분석

이 연구는 트랜스포머의 핵심 연산인 어텐션이 토큰들의 평균을 계산하고, 잔차 연결과 레이어 정규화를 통해 토큰 분산을 간접적으로 제어한다는 점에 착안한다. 저자들은 어텐션 출력 (PV) 와 원본 값 (V) 의 차이를 직접 사용해 토큰 분산을 조절하는 라플라시안 메커니즘 (L(X)=V-PV) 을 제안한다. 이 변형은 기존 어텐션 헤드 중 (k) 개를 라플라시안 헤드로 교체하는 방식으로 구현되며, 추가 파라미터 없이 구현 난이도가 낮다.

실험에서는 DeiT‑3 기반 ViT‑B 모델에 0~12개의 라플라시안 헤드를 삽입해 CIFAR‑10/100 및 ImageNet‑1k에서 Top‑1 정확도가 평균 0.5 ~ 1 %p 상승함을 보였다. 특히 CIFAR‑10/100에서는 라플라시안 헤드 수가 늘어날수록 정확도가 단조 증가했으며, ImageNet에서는 적절한 혼합(예: 9L)으로도 의미 있는 개선을 얻었다. 언어 영역에서는 GPT‑2‑style 디코더에 라플라시안 헤드를 삽입해 ARC, MMLU, GSM8K, HumanEval 등 6개 벤치마크에서 평균 ≈ 3 %p 상승을 기록했다.

토큰 기하학 분석에서는 네 가지 도구를 사용했다. (1) PCA 시각화는 라플라시안 헤드가 많을수록 클래스별 클러스터가 명확히 구분되는 것을 보여준다. (2) 토큰 분산(ANOVA) 분석에서는 전체 분산이 Between‑Class 비중으로 이동하고, Within‑Seq 분산이 급감함을 확인했다. (3) 층별 코사인 유사도 측정은 라플라시안 메커니즘이 깊은 층으로 갈수록 토큰 간 정렬을 가속화한다는 점을 시사한다. (4) Neural Collapse(NC) 메트릭을 확장한 “Neural Token Collapse”(NTC) 정의에 따르면, 클래스 평균이 단순한 ETF(Equiangular Tight Frame) 구조를 이루고, 클래스별 토큰이 해당 평균으로 붕괴한다. 즉, 라플라시안 메커니즘은 트랜스포머가 이상적인 NC 구조에 가까워지도록 토큰 공간을 재구성한다.

이러한 결과는 토큰 분산을 직접 제어함으로써 어텐션이 학습 과정에서 자연스럽게 “클래스 간 최대 분리”를 달성하게 만든다. 라플라시안 헤드가 잔차 연결과 결합될 때, 평균‑분산 비율을 조정하는 기존 방식보다 더 효율적인 변형이 가능해진다. 또한, 라플라시안 메커니즘은 기존 어텐션의 장점(다양한 패턴 학습)과 보완적으로 작동해, 특히 복잡한 데이터셋에서 헤드 수를 적절히 조절하면 성능과 토큰 기하학 모두에서 최적점을 찾을 수 있음을 시사한다.