Apriel-H1: Mamba 하이브리드를 통한 효율적인 추론 모델 증류
ServiceNow AI는 Apriel-H1을 개발했습니다. 이는 완전 어텐션 교사 모델을 Mamba 하이브리드 아키텍처로 변환하는 15B 추론 모델 제품군입니다. 이 접근 방식은 추론 품질의 최소한의 손실로 최대 2.1배 처리량 증가를 달성하며, 대상 증류를 통해 기존 강력한 모델에 효율성을 retrofit할 수 있음을 보여줍니다.
핵심 통찰: 능력 매칭 증류
추론 모델을 증류할 때는 보존하려는 특정 능력과 일치하는 데이터를 사용해야 하며, 범용 사전 학습 데이터를 사용해서는 안 됩니다. 사전 학습 데이터나 사전 학습과 SFT 데이터를 혼합하여 증류를 시도했을 때는 추론 품질이 크게 저하되었습니다.
성공적인 접근 방식은 교사의 SFT 데이터셋에서 고품질 추론 트레이스를 활용했는데, 다음을 포함합니다:
- 다단계 수학 증명
- 명확한 논리적 의존성을 가진 코딩 작업
- 상세한 설명 체인을 가진 과학적 분석
이는 검색 및 유도 헤드와 같은 어텐션 메커니즘을 Mamba의 선형 재귀로 대체할 때 추론에 사용되는 계산 경로가 방해되기 때문입니다. 고품질 추론 트레이스는 하이브리드 모델이 동일한 결과로 이어지는 새로운 경로를 발견할 수 있도록 필요한 명시적인 구조를 제공합니다. 또한 팀은 증류에 역방향 KL 발산 (온도 1)을 사용했는데, 이는 모드 탐색 동작으로 인해 학생 모델이 교사의 높은 신뢰도와 구조화된 예측에 전념하도록 유도하기 때문입니다.
Apriel-H1 성능 및 벤치마크
플래그십 모델인 Apriel-H1-15b-Thinker-SFT는 완전 어텐션 교사 모델과 비교하여 2.1배 처리량을 달성합니다. 추론 벤치마크에 미치는 영향은 다음과 같습니다:
| 벤치마크 | 교사 점수 | Apriel-H1-15b-Thinker-SFT 점수 |
|---|---|---|
| MATH500 | 0.90 $ | |
| ightarrow$ 0.92 | 향상 | |
| MTBench | 8.30 $ | |
| ightarrow$ 8.58 | 향상 | |
| GSM8k | 0.97 $ | |
| ightarrow$ 0.95 | 약간의 회귀 | |
| GPQA | 0.59 $ | |
| ightarrow$ 0.55 | 약간의 회귀 | |
| AIME24 | 0.70 $ | |
| ightarrow$ 0.65 | 약간의 회귀 |
Apriel-H1 제품군 전체에서 체크포인트는 25에서 40개의 Mamba 레이어(총 50개 중) 범위를 갖습니다. H-30 변형이 가장 좋은 균형을 제공하는 반면, H-40 변형은 지연 시간이 중요한 워크로드에 대해 처리량을 최대화합니다(최대 3.4배).
단계별 증류 과정
완전 어텐션 모델을 하이브리드로 변환하는 것은 단일 단계로 수행할 수 없습니다. ServiceNow AI는 세 단계 절차를 사용했습니다:
단계 1: 초기 레이어 교체
MMLU에서 Leave-One-Out (LOO) 분석을 사용하여 팀은 가장 중요도가 낮은 25개 레이어를 식별했습니다. 이러한 레이어는 Mamba-in-Llama (MIL)로 초기화된 믹서로 교체되고 엔드투엔드로 증류되었습니다.
단계 2: 점진적 변환
LOO 분석은 25개 레이어를 넘어가면 신뢰할 수 없게 되므로 팀은 MIL-Mamba-교체 (MMR) 를 구현했습니다. 이 동적 휴리스틱은 Mamba 믹서를 초기화하고 100개의 훈련 단계를 실행하여 증류 손실을 기록합니다. 손실이 낮은 쪽으로 수렴하는 레이어는 교체 우선 순위로 지정됩니다. 이 과정은 증분식으로 수행되었습니다(25 $ ightarrow$ 27 $ ightarrow$ 30 $ ightarrow$ 34 $ ightarrow$ 37 $ ightarrow$ 40 레이어).
단계 3: 최종 SFT
목표 Mamba 레이어 수에 도달한 후, 추론 성능이 안정될 때까지 최종 SFT 패스를 수행했습니다. 최종 Apriel-H1-15b-Thinker-SFT 모델은 55.9B 증류 토큰과 20.9B SFT 토큰을 처리한 후에 생성되었습니다.
구현 및 재현 가능성
Apriel-H1은 Fast-LLM을 사용하여 구축되었습니다. Fast-LLM은 어텐션과 Mamba를 모듈식 믹싱 인터페이스로 취급하는 오픈소스(Apache 2.0) 훈련 프레임워크입니다.これにより 구성의 pattern 필드를 통해 레이어 순서를 지정할 수 있어 블록 유형을 자유롭게 교체할 수 있습니다.
프로덕션 배포
Apriel-H1은 Hugging Face Transformers 및 vLLM과 통합되었습니다. vLLM 통합은 연속 배치, 프리픽스 캐싱, 청크된 프리필을 위해 Mamba 캐시 작업을 활용합니다. 저자들은 처리량 증가가 betyd하지만, 도구가 성숙 중이기 때문에 현재 하이브리드 모델 배포는 맞춤 코드와 신중한 수치 검증이 필요하다고 언급합니다.
기술 사양
- Mamba 하이퍼파라미터: 상태 크기 16, DT 순위 16, 내부 차원 4096.
- 아키텍처: Mamba-1과 어텐션 레이어의 하이브리드.
- 훈련 데이터: 플래그십 모델을 위한 총 76.8B 토큰.