Transformer-SSM 하이브리드 모델을 위한 검색 인식 지식 증류

Retrieval-Aware Distillation for Transformer-SSM Hybrids

상태 공간 모델(SSM)은 효율적인 시퀀스 모델링을 제공하지만, 문맥 내 검색이 필요한 벤치마크에서는 Transformer 모델에 비해 성능이 뒤쳐집니다. 이전 연구에서는 이러한 성능 차이를 '수집 및 집계(G&A)'라고 불리는 제한된 수의 어텐션 헤드에서 비롯된다고 밝혔으며, SSM은 이러한 헤드를 구현하는 데 어려움을 겪습니다. 본 연구에서는 '검색 인식 지식 증류' 방법을 제안합니다. 이 방법은 사전 학습된 Transformer 모델을, 검색에 중요한 헤드만 유지하고 나머지 헤드를 순환형 헤드로 변환하여 하이브리드 모델(학생 모델)을 생성합니다. 합성 검색 작업을 통해 필수적인 헤드를 식별하고, 희소하고 불균일한 어텐션 배치 구조를 가진 하이브리드 모델을 생성합니다. 실험 결과, **어텐션 헤드의 2%만 유지해도 95% 이상의 교사 모델 성능을 복구할 수 있습니다 (10억 파라미터 모델에서 10개의 헤드 사용)**. 이는 최소 25% 이상의 헤드를 유지하는 기존 하이브리드 모델보다 훨씬 적은 수의 헤드를 사용합니다. 또한, 큰 크기의 순환 상태가 누락된 검색 기능을 보완하는 경우가 많으며, 이러한 헤드를 통해 검색 기능을 처리하면 SSM의 핵심 부분을 단순화할 수 있습니다. 순환 상태의 차원을 8배 줄여도 성능 저하가 미미합니다. 결과적으로, 제안하는 하이브리드 모델은 어텐션 캐시와 SSM 상태를 줄여 동등한 하이브리드 모델보다 5~6배 더 메모리 효율적이며, Transformer와 SSM 간의 성능 격차를 메모리 비용의 일부로 좁힐 수 있습니다.

State-space models (SSMs) offer efficient sequence modeling but lag behind Transformers on benchmarks that require in-context retrieval. Prior work links this gap to a small set of attention heads, termed Gather-and-Aggregate (G&A), which SSMs struggle to reproduce. We propose *retrieval-aware distillation*, which converts a pretrained Transformer into a hybrid student by preserving only these retrieval-critical heads and distilling the rest into recurrent heads. We identify the essential heads via ablation on a synthetic retrieval task, producing a hybrid with sparse, non-uniform attention placement. We show that preserving **just 2% of attention heads recovers over 95% of teacher performance on retrieval-heavy tasks** (10 heads in a 1B model), requiring far fewer heads than hybrids that retain at least 25%. We further find that large recurrent states often compensate for missing retrieval: once retrieval is handled by these heads, the SSM backbone can be simplified with limited loss, even with an $8\times$ reduction in state dimension. By reducing both the attention cache and the SSM state, the resulting hybrid is $5$--$6\times$ more memory-efficient than comparable hybrids, closing the Transformer--SSM gap at a fraction of the memory cost.