Softmax 가중치를 선형 ViT로 옮기려면 MLP는 복사하고 어텐션은 증류한다

무엇인가

이 논문이 푸는 문제는 명확하다. Vision Transformer의 Softmax 어텐션은 토큰 수 N에 대해 O(N²) 비용을 쓰기 때문에 고해상도·장문맥 작업에서 토큰 라우팅이 병목이 된다. 선형 ViT는 이를 선형 복잡도 어텐션 연산자로 바꿔 O(N)으로 줄이지만, 처음부터 사전학습해야 하고 대개 원래 Softmax 버전보다 성능이 낮다. 저자들은 여기서 "대부분의 파운데이션 ViT가 Softmax 어텐션 위에 세워져 있는데, 선형 ViT가 그 사전학습 가중치를 초기화에 활용할 수 있는가"라고 직접 묻는다. 기존 Attention Transfer 연구는 Softmax ViT 사이에서는 어텐션만 옮겨도 성능이 회복된다고 보고했지만, 저자들은 Softmax→선형 경계에서는 정반대 결과를 얻는다.

어떻게 동작하나

첫 번째 발견은 어텐션 가중치가 연산자 특이적(operator-specific)이라는 것이다. 사전학습된 Softmax의 Q·K·V 투영 가중치를 선형 ViT의 각 블록에 그대로 복사하고 나머지는 랜덤 초기화하는 Copy 설정은 거의 도움이 되지 않고, 때로는 순수 랜덤 초기화보다 나쁘다. 사전학습 소스를 전혀 쓰지 않고 구조적 사전지식만으로 만든 Mimetic, Impulse 같은 해석적 초기화가 Copy를 체계적으로 앞선다. 이 실패는 ELU·ReLU·InLine·MHLA·TTT 다섯 가지 선형 ViT 변형, Tiny/Small/Base 크기, CIFAR-10·CIFAR-100·STL-10·Food·Flowers·Pets 및 ImageNet-1K에서 일관되게 나타난다. DeiT 대신 DINO·MoCov3·iBOT·MAE 가중치를 소스로 써도, ImageNet-1K 분포 이동 벤치마크에서도 양상은 같다. 즉 사전학습 어텐션 가중치는 복사만으로는 연산자 변경을 넘어가지 못한다.

무엇과 다른가

두 번째 발견은 어텐션의 토큰 라우팅 자체는 증류로 회복된다는 것이다. 저자들은 동결된 사전학습 Softmax 교사와 학생의 어텐션 "출력"을 맞추는 증류 손실을 하위 작업 학습 중에 추가한다. Softmax ViT 간 전이에서 쓰던 어텐션 맵 매칭이 아니라 출력 매칭을 쓴 것이 핵심으로, 선형 연산자마다 맵의 성질이 다르고 어떤 연산자는 맵을 아예 만들지 않기 때문이다. 손실 가중치는 λ=50으로 고정해 모든 연산자·데이터셋에 공통 적용했으며, 이 증류는 Copy 대비 Softmax 목표와의 격차를 크게 줄이고 MHLA와 TTT에서는 일부 데이터셋에서 목표에 도달하기도 한다. 예외는 세밀한 분류인 Flowers로, 학습 샘플이 적어 어텐션 감독만으로는 표현을 다시 학습할 신호가 약하고 MHLA는 Copy보다도 낮아진다.

어떻게 쓰나

세 번째 발견은 MLP 가중치가 연산자 무관(operator-agnostic)이라는 것이다. 어텐션 관련 구성요소는 랜덤으로 두고 MLP 블록만 복사하는 MLPCopy는 Random과 Copy보다 모든 데이터셋에서 크게 앞서고, 어텐션까지 복사하는 FullCopy는 MLPCopy 위에 일관된 추가 이득을 준다. 저자들은 이를 MLP를 학습된 지식을 담는 key-value 메모리로 보는 관점과 연결한다. 즉 Softmax→선형 경계를 복사로 건너가는 것은 MLP 블록에 담긴 표현이며, 이것이 사전학습 가중치 이득의 대부분을 나른다.

전제와 한계

이 두 가지를 합친 것이 논문의 레시피다. MLP는 복사하고 어텐션은 증류하면, 증류만으로는 도달하지 못하던 Softmax 목표를 대부분 닫고 경우에 따라 능가한다. 흥미로운 점은 MLPCopy w/ Distill이 FullCopy w/ Distill에 근접한다는 것, 즉 라우팅 행동을 증류하고 나면 복사한 어텐션 가중치는 거의 중복이 된다는 것이다. 다만 Copy w/ Distill은 Random w/ Distill보다 꾸준히 높아, 복사된 어텐션 가중치가 라우팅 회복을 위한 쓸 만한 워밍 스타트를 제공한다는 해석이 붙는다. 손실 설계 절제에서는 ReLU·ELU처럼 어텐션 맵이 정상 분포인 경우 맵 매칭과 출력 매칭의 최고 회복률이 비슷하지만, 감산 정규화로 맵이 부호 있는 준분포가 되는 InLine에서는 맵 매칭이 일관되게 뒤진다. 따라서 연산자 무관한 어텐션 출력에서 맞추는 설계가 더 일반적이며, λ는 최고점 주변의 넓은 평탄 구간을 가져 50이 기본값으로 채택됐다.

개발자 관점에서 이 논문의 실용적 메시지는 두 갈래다. 선형 ViT를 사전학습된 Softmax 체크포인트에서 출발시키려면 어텐션 투영 가중치 복사를 기대하지 말고 MLP 블록을 그대로 가져오고, 어텐션은 교사의 어텐션 출력에 대한 증류 손실로 라우팅을 되살려야 한다. 다만 이 증류는 하위 작업 학습 시점에 동결된 Softmax 교사가 함께 있어야 하므로, 교사 없이 가중치만 배포하는 시나리오에는 그대로 적용되지 않는다. 또한 저자들이 제시한 정확도 수치는 본문에 인용된 그림·표(Figure 2~5, Table 2)에 담겨 있고 제공된 원문 텍스트에는 구체적 수치가 제시되지 않았으므로, 개선폭을 인용하려면 원문 그림과 표를 직접 확인해야 한다.

저자들이 밝힌 전제와 한계도 분명하다. MHLA와 TTT처럼 Softmax에 없는 추가 구성요소를 가진 변형에서는 대응되지 않는 파라미터를 기본 설정으로 초기화하며, Mimetic과 Impulse는 표준 QKᵀ 파라미터화를 가정하므로 TTT에는 적용할 수 없다. 어텐션 증류만으로는 Flowers 같은 소량·세밀 데이터셋에서 격차가 남고, 실험은 3개 랜덤 시드 평균으로 보고되며 데이터셋 규모에 따라 300 또는 600 에폭 학습한다. 윤리 성명에서는 이 결과가 "어텐션 전이는 연산자 간에 통하지 않는다"로 오독되어서는 안 되고, 교사 가중치가 물려주는 데이터 편향이나 안전성 문제를 제거하지는 못한다는 점을 명시한다.