검색 인식 기반 변환기‑SSM 하이브리드 효율화

초록

본 논문은 사전 학습된 대형 변환기 모델에서 인‑컨텍스트 검색에 핵심적인 소수의 어텐션 헤드(‘Gather‑and‑Aggregate’ 헤드)를 식별하고, 이를 보존한 채 나머지 헤드를 상태공간 모델(SSM) 기반 재귀 헤드로 대체하는 ‘검색‑인식(distillation)’ 방법을 제안한다. 1 B 규모 모델에서 전체 512개 헤드 중 2 %에 해당하는 10개만 유지해도 검색‑중심 벤치마크에서 교사 모델 성능의 95 % 이상을 회복한다. 또한 검색 전용 헤드를 확보하면 SSM의 상태 차원을 8배 축소해도 성능 저하가 미미해, 메모리 사용량을 5‑6배 절감할 수 있음을 보인다.

상세 분석

이 연구는 변환기와 상태공간 모델(SSM) 사이의 성능 격차가 전체 아키텍처의 한계가 아니라, ‘Gather‑and‑Aggregate(G&A)’라 명명된 소수의 어텐션 헤드가 담당하는 검색 기능의 부재에 기인한다는 최근 발견을 기반으로 한다. 저자들은 먼저 합성 KV‑retrieval 작업을 이용해 각 헤드의 중요도를 정량화한다. 이 작업은 키‑값 쌍을 저장하고 질의에 대해 정확히 매칭하는 능력을 측정하므로, G&A 헤드가 크게 손실될 경우 정확도가 급격히 떨어진다. 이러한 ablation 결과를 스코어링해 상위 2 % 헤드만을 선택하고, 나머지는 SSM 기반 재귀 헤드로 교체한다.

핵심 기술은 세 단계로 구성된다. ① 헤드 선정: KV‑retrieval 프로브에서 각 헤드의 성능 저하량을 측정해 ‘검색‑중요도 점수’를 부여한다. ② 구조 재구성: 선정된 헤드만 남기고, 남은 어텐션 헤드를 동일 레이어 내에서 SSM 믹서(Discrete‑Mamba‑2)로 대체한다. 이때 남은 헤드들의 출력과 SSM 출력의 평균·분산을 맞추는 파라미터‑프리 정규화(LayerNorm) 과정을 거쳐 혼합 안정성을 확보한다. ③ 지식 증류: MOHA‑WK 프레임워크를 활용해 교사의 토큰‑믹싱 행렬과 학생 SSM 믹서를 매트릭스‑오리엔테이션 단계에서 Frobenius 거리 최소화로 정렬하고, 이어서 히든‑스테이트 L2 정렬, 최종적으로 로그잇(KL) 손실을 이용한 엔드‑투‑엔드 미세조정을 수행한다.

실험에서는 Llama‑3.2‑1B와 Qwen2.5‑1.5B 두 모델에 적용했으며, 10개의 헤드(전체 512개 중 2 %)만 보존해도 검색‑중점 과제(SWDE, KV‑Retrieval, Lambada 등)에서 교사 대비 95 % 이상 커버리지를 달성했다. 반면 기존 하이브리드는 최소 25 % 이상의 헤드를 유지해야 비슷한 성능을 얻었다. 또한, 검색 전용 헤드가 충분히 확보되면 SSM의 상태 차원을 64→8로 8배 축소해도 정확도 저하가 미미했으며, 이로 인해 어텐션 캐시와 SSM 상태 메모리 모두 크게 감소한다. 결과적으로 짧은 시퀀스(128 토큰)에서는 5배, 긴 시퀀스(4K 토큰)에서는 6배의 메모리 절감과 함께 추론 속도도 개선되었다.

이 논문의 주요 기여는 (1) 교사의 내부 구조를 활용해 ‘어떤 어텐션이 검색에 필수적인가’를 정량적으로 식별하는 방법을 제시한 점, (2) 식별된 헤드만을 보존하고 나머지를 효율적인 재귀 연산으로 대체함으로써 하이브리드 모델의 메모리·연산 효율성을 크게 높인 점, (3) 검색 기능이 확보되면 SSM 백본 자체를 경량화할 수 있음을 실증한 점이다. 이러한 접근은 대규모 변환기 모델을 제한된 하드웨어 환경에 배포하거나, 장기 의존성을 요구하는 실시간 애플리케이션에 적용하기 위한 실용적인 로드맵을 제공한다.