기술적 요약: 라플라시안 헤드는 토큰 표현을 평활화하여 트랜스포머를 개선합니다

문제 제기

표준 트랜스포머 아키텍처는 멀티헤드 어텐션과 잔차 연결을 통해 토큰 표현을 업데이트합니다. 이 공식에서 각 어텐션 헤드 $i$는 소프트맥스 어텐션 행렬 $P^{(i)}$를 계산하고, 값 projections 의 가중 합을 통해 입력 $X$를 업데이트합니다: $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

이 메커니즘은 효과적이지만, 토큰 표현의 **시퀀스 내 분산 (within-sequence variance)**에 대한 명시적 제어는 부족합니다. 그래프 신경망 (GNN) 과 심층 트랜스포머에 관한 문헌에서 지배적인 견해는 토큰 표현이 지나치게 유사해져 구별력을 잃는 "과평활화 (oversmoothing)"를 해로운 실패 모드로 간주합니다. 따라서 대부분의 기존 연구는 정규화, 엣지 드롭, 또는 다양성 손실 등을 통해 평활화를 방지하는 데 초점을 맞추고 있습니다. 본 논문은 토큰 과평활화가 본질적으로 해롭다는 가정에 도전하며, 대신 제어된 평활화가 표현 학습에 유익할 수 있음을 제안합니다.

방법론

저자들은 파라미터가 없는 아키텍처 수정을 제안합니다: 표준 어텐션 헤드의 일부 subset 을 라플라시안 헤드로 대체하는 것입니다.

라플라시안 연산자

표준 업데이트에서 헤드는 $P^{(i)}X$를 계산하는데, 이는 토큰들의 어텐션 가중 평균입니다. 제안된 라플라시안 헤드는 어텐션 행렬 $P^{(i)}$를 그래프 라플라시안 연산자 $I - P^{(i)}$로 대체합니다. 여기서 $I$는 단위 행렬입니다. 업데이트 방정식은 다음과 같습니다:
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
여기서 $A$는 표준 어텐션 헤드의 집합이고 $L$은 라플라시안 헤드의 집합입니다.

동기

이 논문은 이 수정을 위한 두 가지 이론적 동기를 제시합니다:

1. 평균과 분산의 직접적 제어: 표준 어텐션 헤드는 토큰 표현의 평균을 업데이트합니다. 반면, 항 $(I - P^{(i)})X$는 각 토큰과 그 어텐션 가중 평균 사이의 차이를 계산합니다. 이 차이를 추가함으로써 모델은 표준 트랜스포머에는 없는 시퀀스 내 토큰 표현의 분산을 직접 제어할 수 있게 됩니다.
2. 그래프 확산: 토큰을 $P^{(i)}$로 정의된 엣지 가중치를 가진 그래프의 노드로 간주하면, $I - P^{(i)}$는 랜덤 워크 정규화 그래프 라플라시안이 됩니다. 이 연산자를 적용하는 것은 토큰 그래프에서 **열 확산 (heat diffusion)**의 단일 단계를 수행하는 것이며, 시퀀스 전반에 걸쳐 신호 (토큰 표현) 를 점진적으로 평활화합니다.

주요 기여

1. 아키텍처 제안: 라플라시안 헤드를 표준 어텐션 헤드와 함께 통합하는 간단하고 파라미터가 없는 트랜스포머 아키텍처 수정안.
2. 실증적 검증: 이 수정이 지도 분류, 언어 모델링, 자기지도 학습이라는 세 가지 다른 학습 패러다임 전반에 걸쳐 성능을 향상시킨다는 입증.
3. 기하학적 특성화: 라플라시안 헤드가 학습된 토큰 표현의 기하학을 어떻게 변화시키는지 대한 상세 분석으로, 해로운 과평활화가 아닌 유익한 형태의 "붕괴 (collapse)"와 평활화를 유도함을 보여줌.

실험 결과

1. 이미지 분류 (지도 학습)

CIFAR-10, CIFAR-100, ImageNet-1k 에서 DeiT-3 계열 (ViT-B) 을 사용:

• 성능: 라플라시안 헤드를 추가하면 Top-1 정확도가 일관되게 향상되었습니다. 향상 정도는 라플라시안 헤드의 수에 따라 거의 단조 증가했습니다.
  • ImageNet 에서 12 개의 라플라시안 헤드를 가진 모델 (ViT-B-12L) 은 기준 모델의 81.45% 대비 **82.32%**의 정확도를 달성했습니다.
  • CIFAR-100 에서 향상 폭은 거의 2% 였습니다 (77.04% 대 75.12%).
• 기하학:
  • ANOVA: 라플라시안 헤드는 시퀀스 내 분산을 유의미하게 줄이는 반면 클래스 간 분산은 증가시켰습니다. 이는 동일한 시퀀스 내의 토큰들이 공통 평균을 향해 붕괴되는 반면, 서로 다른 시퀀스의 평균들은 더 효과적으로 분리됨을 나타냅니다.
  • 신경 붕괴 (Neural Collapse): 모델들은 신경 붕괴 기하학으로의 더 강력한 수렴을 보였습니다. 토큰 표현은 각 클래스별로 잘 분리된 심플렉스 형태의 군집 (Simplex ETF) 을 형성했으며, 시퀀스 평균들은 전역 클래스 기하학과 정렬되었습니다.

2. 언어 모델링 (자기회귀)

FineWeb-Edu 에서 사전 학습되고 추론/코드 벤치마크에서 파인튜닝된 GPT-2 스타일 모델 (561M 파라미터) 을 사용:

• 성능: 라플라시안 헤드를 가진 모델들은 ARC-Easy, ARC-Challenge, MMLU, GSM8K, HumanEval 에서 평균적으로 기준 모델보다 우수한 성능을 보였습니다.
  • 5 개의 라플라시안 헤드를 가진 모델이 가장 높은 평균 성능 (30.06%) 을 달성했으며, ARC-Easy (+16%) 와 ARC-Challenge (+7%) 에서 유의미한 향상을 보였습니다.
• 기하학:
  • 스펙트럼 감쇠: 라플라시안 헤드는 더 빠르게 감쇠하는 스펙트럼을 가진 토큰 표현을 생성했습니다. 모든 헤드가 라플라시안일 때 분산의 99% 를 설명하는 데 필요한 주성분의 수가 21% 감소했습니다 (547.7 에서 430.4 로).
  • 언어적 붕괴 (Linguistic Collapse): 모델들은 동일한 다음 토큰 예측을 공유하는 토큰 표현들 사이의 분리도가 증가하는 것을 보였으며, 이는 유사한 의미적 역할을 가진 토큰들이 군집하는 더 강력한 형태의 "언어적 붕괴"를 시사합니다.

3. 자기지도 학습 (DINO)

ImageNet-1k 에서 DINO 사전 학습을 위해 ViT-S 모델을 사용:

• 성능: 라플라시안 헤드는 ImageNet 에서 선형 프로빙 정확도와 ADE20K 에서의 분할 성능을 향상시켰습니다.
  • ViT-S-5L 은 기준 모델 (29.81%) 대비 가장 좋은 ADE20K mIoU (32.27%) 를 달성했습니다.
• 기하학:
  • 분할: PCA 시각화는 라플라시안 헤드가 기준 모델에 비해 더 선명한 객체 경계를 생성하고 더 미세한 세부 사항 (예: 헬멧 모양, 저지 번호) 을 보존함을 보여주었습니다.
  • 스펙트럼 감쇠: 언어 모델링과 유사하게, 이러한 모델들은 더 빠른 스펙트럼 감쇠를 보이며 이미지 내 토큰 표현들을 저차원 부분 공간으로 집중시켰습니다.

중요성과 주장

이 논문은 토큰 과평활화가 보편적으로 해롭다는 지배적인 견해가 불완전하다고 주장합니다. 대신 저자들은 평활화가 표현 학습을 위한 메커니즘이며 제어될 때 유익할 수 있다고 주장합니다.

• 작업별 이점: 라플라시안 헤드에 의해 유도된 "평활화"는 실패 모드가 아니라 작업 목표와 일치하는 기능입니다:
  • 지도 분류에서는 시퀀스 내 토큰들을 붕괴시키고 시퀀스 평균을 클래스 기하학과 정렬시킴으로써 신경 붕괴를 촉진합니다.
  • 언어 모델링에서는 동일한 다음 토큰 타겟을 가진 토큰들의 분리도를 향상시켜 언어적 붕괴를 강화합니다.
  • 자기지도 학습에서는 표현들을 저차원 부분 공간에 집중시켜 분할과 같은 하류 작업을 지원합니다.
• 메커니즘: 라플라시안을 통해 그래프 확산을 명시적으로 아키텍처에 통합함으로써, 모델은 시퀀스 내 분산을 직접 제어할 수 있는 능력을 얻게 되며, 이는 더 견고하고 기하학적으로 구조화된 표현으로 이어집니다.

저자들은 고정된 전략 구현이 아마도 최적은 아닐 것이라고 결론지었지만, 그 결과들은 평활화를 완화해야 할 문제로만 보지 않고 표현 학습의 근본적인 구성 요소로서 더 깊이 연구되어야 함을 시사합니다.