Attention to Mamba: A Recipe for Cross-Architecture Distillation

Apple 연구진은 Pythia-1B Transformer로부터 순수 Mamba 계열 구조로 지식을 전이하기 위해 HedgeMamba라는 2단계 교차 구조 증류 방법을 개발했다.

Apple 연구진은 Pythia-1B Transformer로부터 순수 Mamba 계열 구조로 지식을 전이하기 위해 HedgeMamba라는 2단계 교차 구조 증류 방법을 개발했다. 이 접근법은 교사 모델 원래 사전학습 토큰 예산의 단 2.7%만 사용하면서, Mamba 학생 모델이 교사의 13.86에 근접한 14.11의 퍼플렉서티를 달성하도록 했다.

핵심 요약

• Transformer 모델은 강력하지만 시퀀스 길이가 늘어남에 따라 2차 계산 비용이 발생하여, 긴 시퀀스와 추론에 대한 효율성이 제한된다.
• Mamba 같은 상태공간 모델(SSM)을 비롯한 대안적 선형 구조는 효율성을 제공하지만, 대규모에서 다운스트림 과제에 대해 Transformer보다 성능이 떨어지는 경우가 많다.
• Transformer에서 Mamba 구조로의 교차 구조 증류에 대한 이전 시도들은 성능 보존에 어려움을 겪었고, 흔히 하이브리드 Attention-SSM 해법을 필요로 했다.
• 원칙에 입각한 2단계 증류 레시피가 개발되었다. 먼저 Softmax Attention에서 학습된 Linear Attention(Hedgehog)으로, 그다음 Linear Attention에서 적응된 Mamba(HedgeMamba)로 전이한다.
• Mamba SSM 믹서 파라미터는 Linear Attention 단계에서 학습된 특징 맵을 사용해 초기화되었고, 추가적인 Mamba 구성요소는 항등 연산자로 작동하도록 설정되었다.
• 이후 전체 HedgeMamba 구조는 정답에 대한 표준 교차 엔트로피 손실을 사용해 미세조정되어, 모든 Mamba 구성요소를 활성화했다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.