AWS, 사전 훈련 비용 99% 절감하는 하이브리드 SSM 변환 기법 'priming' 공개

AWS 연구진이 사전 훈련된 Transformer 가중치를 하이브리드 상태 공간 모델(SSM)로 변환하는 priming 기법을 공개했다. 이 방법은 사전 훈련 비용을 99% 이상 절감하며, 개발자가 주의(attention) 레이어와 SSM 레이어를 혼합한 하이브리드 아키텍처의 설계 공간을 저비용으로 탐색할 수 있게 한다.

Priming 작동 방식

Priming은 Transformer와 SSM의 구조적 유사성을 활용하여 Transformer를 상태 공간 실현의 특수한 경우로 간주한다. 과정은 세 단계로 구성된다:

• 레이어 선택: 슬라이딩 윈도우 주의 프록시를 사용하여 고정 크기 상태로 압축 가능한 주의 레이어를 식별한다.
• 가중치 이전: Transformer 가중치에서 SSM 레이어를 초기화하고 증류(distillation)를 통해 정렬한다.
• 엔드투엔드 적응: 다음 토큰 예측을 위해 하이브리드 모델을 미세 조정한다.

이 파이프라인은 원본 모델 사전 훈련 토큰 예산의 0.5% 미만만 소비하며, Qwen, Llama, Mistral 등 다양한 모델 가족에서 작동한다.

성능 및 발견 사항

연구팀은 세 가지 SSM 유형인 Gated KalmaNet (GKA), Gated DeltaNet (GDN), Mamba-2를 비교했다. 주요 발견 사항은 다음과 같다:

• GKA가 가장 우수한 성능을 보였고, GDN과 Mamba-2가 뒤를 이었으며, 긴 컨텍스트 추론 작업에서 격차는 더 벌어졌다.
• 320억 매개변수의 GKA 기반 하이브리드는 추론 벤치마크에서 원본 Transformer와 동일한 성능을 기록하면서 128,000 토큰 컨텍스트에서 디코딩 속도를 최대 2배까지 높였다.
• SSM 기반 하이브리드는 메모리 예산이 두 배인 슬라이딩 윈도우 주의(SWA) 하이브리드보다 일관되게 더 높은 성능을 보였다.

가변 연산 및 컨텍스트 확장

GKA는 가변 길이 추론 연산을 제공하여 재훈련 없이 서빙 시점에 속도-정확도 트레이드오프를 조정할 수 있다. 예를 들어, 반복 횟수를 30에서 10으로 줄이면 AIME 2025에서 정확도는 약 1% 감소하지만 디코딩 속도는 10% 빨라진다. 또한, 팀은 청크 간 상태를 결합하여 추가 훈련 없이 컨텍스트 길이를 128K에서 256K 토큰으로 확장하는 방법을 공개했다.

오픈 소스 릴리스

AWS는 GitHub와 Hugging Face에 Apache 2.0 라이선스 하에 priming된 하이브리드 모델(8B 및 32B 매개변수), 훈련 코드, 최적화된 커널, vLLM 서빙 플러그인을 공개했다.