ServiceNow가 15B 추론 모델을 Mamba 하이브리드로 변환해 2.1배 처리량 증가 달성

Context

기존 15B 추론 모델을 효율적으로 만들어야 했으나, 20T 토큰 사전학습이나 아키텍처 공동설계를 위한 무한 컴퓨팅 자원이 없었다. 기존 모델을 기반으로 증류(Distillation)를 통해 효율성을 확보할 수 있는지가 핵심 과제였다.

Technical Solution

• 주의: 일반 사전학습 데이터 대신 고품질 추론 데이터(수학 증명, 코딩 작업, 과학 분석)로만 증류 수행
• Leave-One-Out(LOO) 분석을 통해 50개 레이어 중 중요도 낮은 25개 레이어를 Mamba 레이어로 초기화
• Stage 1: 25개 Mamba 레이어까지만 LOO 기반 정적 분석으로 교체 (H-25 체크포인트)
• Stage 2: 25개 이상에서는 MIL-Mamba-Replacement(MMR) 동적 휴리스틱 도입하여 단계별 진행(25→27→30→34→37→40)
• Stage 3: 최종 Mamba 레이어 개수 도달 후 SFT 데이터로 end-to-end 학습
• Reverse KL 발산(온도 1) 사용: Forward KL 대신 Teacher 확신도 높은 예제에서 Student의 높은 확신도 유도
• Transformers와 vLLM 통합: 교체 가능한 Attention/Mamba 레이어 클래스 제공 및 연속 배칭, 프리픽스 캐싱 지원
• 총 학습: 55.9B 토큰 증류 + 20.9B 토큰 SFT

Impact

• 처리량 2.1배 증가(문맥 길이에 따라 1.89~2.09배)
• MATH500: 0.90→0.92 (+0.02점)
• MTBench: 8.30→8.58 (+0.28점)
• GSM8k: 0.97→0.95 (-0.02점)
• GPQA: 0.59→0.55 (-0.04점)
• AIME24: 0.70→0.65 (-0.05점)

Key Takeaway

기존 강력한 모델을 효율화할 때는 일반 사전학습 데이터가 아닌 보존하려는 능력(추론)과 정확히 일치하는 고품질 태스크 데이터로 증류해야 한다. 추론 모델의 복잡한 Attention 메커니즘(장거리 의존성, 논리 체인 추적)을 Mamba 선형 재귀로 대체할 때, 명시적으로 추론 구조가 드러나는 학습 예제가 새로운 경로 발견을 가능하게 한다.