From Sparsity to Simplicity: Enabling Simpler Sequential Replacements via Sparse Attention Distillation

기술 요약: 희소성에서 단순성으로 (S2S)

문제 제기
자기-어텐션 메커니즘은 대규모 트랜스포머 사전 학습의 기초로서, 유연한 토큰 혼합과 장거리 의존성 모델링을 가능하게 합니다. 그러나 시퀀스 길이에 대한 2 차 복잡도로 인해 높은 추론 지연 시간과 메모리 오버헤드가 발생합니다. 어텐션을 RNN, 상태 공간 모델과 같은 더 간단한 순차 모듈로 대체하는 것은 효율성을 위한 매력적인 방안이지만, 단순한 대체는 특히 대규모 규모에서 상당한 정확도 저하를 초래하는 경향이 있습니다. 기존 효율성 접근법들은 종종 어텐션 계산 자체를 수정하거나 토큰을 가지치기하지만, 어텐션을 핵심 연산자로 유지합니다. 본 논문은 성능에 상당한 손실 없이 어텐션 레이어를 완전히 순차 모듈로 대체하는 과제를 다룹니다.

방법론
본 논문은 사전 학습된 트랜스포머가 복잡한 토큰 의존성을 다양한 유효 복잡도를 가진 블록 단위 매핑으로 분해한다는 가설을 기반으로 하는 Sparsity to Simplicity (S2S) 라는 프레임워크를 제안합니다. 저자들은 희소한 어텐션 패턴을 보이는 레이어가 본질적으로 더 간단한 순차 모듈로 근사하기 쉽다고 가정합니다.

방법론은 세 가지 핵심 요소로 구성됩니다:

1. 통합 레이어별 어텐션 대체:

   • 저자들은 사전 학습된 비전 트랜스포머 (ViT) 의 특정 어텐션 레이어를 원래 레이어 인터페이스 (레이어 정규화, 잔차 연결, MLP) 를 유지하면서 순차 토큰 믹서 (구체적으로 BiMamba 및 BiLSTM) 로 교체하는 플러그 앤 플레이 프레임워크를 도입합니다.
   • 순차 모듈은 ViT 어텐션의 인과적이지 않고 전체 컨텍스트를 다루는 특성과 일치하도록 멀티헤드 양방향 구조를 채택합니다.
   • 레이어 정렬 감독은 증류 과정에서 사용됩니다: 교체된 순차 믹서의 출력은 동일한 깊이의 원래 어텐션 레이어의 출력과 직접 비교되어 평균 제곱 오차 (MSE) 손실 함수를 사용합니다.

2. 희소성 유도 증류:

   • 희소성의 영향을 통제하고 연구하기 위해 저자들은 이진 토큰 유지 마스크를 생성하기 위해 할팅 점수를 사용하는 A-ViT(Adaptive ViT) 를 통합합니다.
   • 1 단계 (증류): 학생 모델 (순차 대체 포함) 은 교사 (A-ViT) 가 생성한 동일한 레이어별 토큰 마스크 하에서 학습됩니다. 이는 토큰 라우팅의 차이에서 토큰 혼합 함수를 근사하는 어려움만 분리하여 학생과 교사가 동일한 활성 토큰에서 작동하도록 보장합니다. 손실 함수에는 유사성 손실, 할팅 정렬 항, 그리고 분류 손실이 포함됩니다.
   • 2 단계 (파인튜닝): 학생은 자체 학습된 할팅 모듈을 사용하여 마스크를 생성하도록 파인튜닝되어, 작업 성능을 유지하면서 자율적으로 작동할 수 있게 됩니다.

3. 레이어 인식 희소 대체 전략:

   • 균일한 대체 대신, 이 전략은 자연스러운 희소성 프로파일에 기반하여 특정 레이어 그룹을 대상으로 합니다. 저자들은 가장 깊은 레이어들 (가장 높은 자연 희소성을 보이는 곳) 만을 대체하고 이러한 특정 블록에 희소성 유도 증류 프레임워크를 적용합니다.

주요 기여

• S2S 가설: 본 논문은 어텐션 희소성이 매핑 복잡도의 징표임을 입증합니다. 사전 학습된 트랜스포머의 희소한 블록은 최소한의 손실로 더 간단한 순차 모듈로 대체될 수 있는 반면, 밀집된 블록은 근사하기 어렵습니다.
• 통제된 대체 프레임워크: 인터페이스 호환성을 유지하고 레이어 선택적 분석을 가능하게 하면서 어텐션을 순차 모듈 (Mamba/LSTM) 로 교체하는 통합 방법.
• 명시적 희소성 통제: 교사 유도 희소성 (A-ViT 마스크를 통해) 이 학생과 명시적으로 정렬되는 증류 파이프라인을 도입하여, 교사의 희소성을 증가시키는 것이 일관되게 학생 - 교사 간 성능 격차를 줄임을 입증했습니다.
• 최적화된 트레이드오프: 자연스러운 레이어별 희소성과 명시적 희소화를 결합하여 희소 어텐션 베이스라인보다 더 나은 정확도 - 효율성 트레이드오프를 달성하는 전략.

실험 결과
실험은 DeiT 백본 (Tiny, Small, Base) 을 사용하여 BiMamba 와 BiLSTM 을 대체제로 ImageNet 에서 수행되었습니다.

• 자연 희소성과 대체 가능성: 토큰 상호작용 맵의 시각화는 얕은 레이어는 밀집된 전역 혼합을 보이는 반면, 깊은 레이어는 현저히 희소하고 선택적이 됨을 보여줍니다. 서로 다른 레이어 그룹을 대체하는 통제 실험은 더 깊은 (희소한) 레이어를 대체하는 것이 더 얕은 (밀집된) 레이어를 대체하는 것보다 정확도 저하가 현저히 작음을 드러냈습니다.
• 명시적 희소성의 효과: A-ViT 를 통해 토큰 유지 비율이 감소될 때, 어텐션 기반 교사와 순차 학생 간의 정확도 격차는 일관되게 축소됩니다. 이는 명시적 희소화가 목표 매핑을 단순화하여 순차 모듈이 학습하기 쉽게 만든다는 것을 확인시켜 줍니다.
• 정확도 - 효율성 트레이드오프: 제안된 레이어 인식 희소 대체 (매칭된 희소성 예산 하에서 가장 깊은 3 개의 레이어만 BiMamba 로 대체) 는 다음과 같은 결과를 달성합니다:
  • 유사한 토큰 유지 비율에서 희소 어텐션 베이스라인 (A-ViT) 과 동등하거나 약간 더 높은 정확도.
  • 상당한 효율성 향상: 384² 해상도에서 밀집 베이스라인 대비 최대 1.71 배의 추정 속도 향상을 달성하며, 더 큰 토큰 수에서 희소 어텐션 모델보다 더 높은 토큰 혼합 처리량을 보입니다.

의의 및 주장
본 논문은 어텐션 대체가 균일한 아키텍처 교체가 아닌 구조화되고 레이어 의존적인 문제로 취급되어야 한다고 주장합니다. 이 연구의 의의는 다음을 입증하는 데 있습니다:

1. 희소성은 대체 가능성의 신호입니다: 사전 학습된 모델은 복잡성을 자연스럽게 분해하여 희소한 레이어가 단순화에 가장 적합하도록 만듭니다.
2. 명시적 희소화는 증류를 돕습니다: 교사에게 희소성을 강제하면 순차 학생의 학습 과제가 쉬워져 성능 격차가 줄어듭니다.
3. 전략적 대체가 균일한 접근법보다 우수합니다: 자연스럽게 희소한 레이어에 대체를 집중하고 통제된 희소성으로 과정을 안내함으로써, 단순히 어텐션을 가지치기하거나 전체 모델을 대체하는 것보다 정확도와 지연 시간의 더 나은 균형을 달성할 수 있습니다.

저자들은 이러한 관점이 사전 학습된 트랜스포머 지식을 활용하면서도 모델의 내부 구조가 허용하는 곳에서 더 간단하고 효율적인 시퀀스 모듈을 통합하는 하이브리드 모델에 대한 향후 연구를 장려한다고 결론지었습니다.