Apple, Attention 모델을 Mamba로 증류해 성능 대부분 살렸다
Attention to Mamba: A Recipe for Cross-Architecture Distillation
Apple 연구진이 Transformer를 Mamba 계열로 증류하는 2단계 레시피를 제안했다. 먼저 kernel trick을 응용해 Transformer 지식을 선형화된 Attention으로 옮기고, 이를 다시 Attention 블록이 없는 Mamba로 증류한다. 그 결과 Pythia-1B의 성능을 대부분 유지하며 perplexity 14.01로 교사 모델의 13.86에 근접했다. 다만 in-context recall에 의존하는 과제에서는 격차가 나타난다.
우리 연구의 핵심 주장은 Mamba에 원리 기반의 초기화를 부여하면, 교차 아키텍처 증류를 위한 전반적으로 더 나은 레시피를 회복할 수 있다는 것이다.
- buildbot
내 생각엔, 모든 state space attention 유형에 대한 비슷한 아이디어를 더 좋고 더 포괄적으로 다룬 논문은 AWS의 이 논문이다: https://arxiv.org/abs/2605.08301
- kooi
멋지네. 나는 state space 모델이 선형 동역학 시스템에 대한 제어와 모델링의 힘에 제대로 경의를 표하는 "느낌"이라서 좋아한다. (Mamba는 비선형 활성화로 연결된 시변 선형 시스템들의 묶음이다).
지금 OG 논문 중 하나를 읽고 있는데, 몇 가지 흥미로운 구절이 있다. 데이터 센터에서의 추론이 전통적인 transformer로 얼마나, SSM으로 얼마나 이루어지는지 궁금하다. 효율성 이득이 꽤 극적이니까.
> 우리는 선택적 state space 모델의 새로운 부류를 제안하며, 이는 여러 축에서 선행 연구를 개선해 시퀀스 길이에 선형으로 스케일링하면서도 Transformer의 모델링 능력을 달성한다.
> 우리의 Mamba 언어 모델은 비슷한 크기의 Transformer 대비 5배의 생성 처리량을 가지며, Mamba-3B의 품질은 두 배 크기의 Transformer와 맞먹는다.
그리고 흥미롭게도 학습에 관해서는:
> 일반적으로, 모델[SSM]은 효율적인 병렬화 가능 학습을 위해 합성곱 모드 (3)를 사용하고(전체 입력 시퀀스를 미리 본다), 효율적인 자기회귀 추론을 위해 순환 모드 (2)로 전환한다(입력을 한 타임스텝씩 본다).
- vlovich123
1B 모델만이다. 이것이 더 큰 모델로도 계속 스케일링되는지 / Mamba 쪽에서 그것이 어떤 모습인지에 대한 언급은 없다.
> 이는 대체로 다운스트림 작업으로 이어지지만, in-context recall에 더 의존하는 작업에서는 약간의 격차가 나타난다
그리고 그렇다, SSM의 유명한 아킬레스건 - 형편없는 recall.