원문 캡처 · amazon.science
원문 캡처 · amazon.science

아마존웹서비스(AWS) 연구팀은 사전 학습한 트랜스포머의 가중치를 재사용해 하이브리드 모델을 만드는 ‘프라이밍(priming)’ 기법을 공개했다. 연구팀은 이 방식으로 처음부터 학습할 때보다 사전 학습 비용을 99% 넘게 줄였으며, 320억 매개변수 모델에서 원본 트랜스포머와 비슷한 추론 정확도를 유지하면서 긴 문맥 처리량을 최대 두 배로 높였다고 밝혔다.

프라이밍은 트랜스포머의 일부 어텐션 층을 고정 크기 상태로 과거 정보를 압축하는 상태공간 모델(SSM) 층으로 바꾼다. 어텐션은 이전 토큰을 상세히 저장하는 대신, 필요한 정보만 유지하는 층을 섞어 메모리 사용량을 줄이는 방식이다. AWS 연구팀은 이 방법이 새 아키텍처를 매번 처음부터 학습해야 했던 비용을 낮춰, 어떤 층을 얼마나 교체할지 비교하기 쉽게 한다고 설명했다.

프라이밍 방식과 비교 결과

프라이밍은 먼저 슬라이딩 윈도우 어텐션(SWA)으로 각 층을 바꿔도 성능 저하가 작은지 확인한다. 과거 토큰 일부만 유지해도 되는 층을 상태공간 모델로 바꿀 후보로 고른 뒤, 대응하는 트랜스포머 가중치로 초기화하고 짧은 증류 과정을 거쳐 동작을 맞춘다. 마지막으로 전체 하이브리드 모델을 다음 토큰 예측으로 미세 조정한다. 연구팀은 전체 과정이 원본 모델 사전 학습 토큰 예산의 0.5% 미만을 사용한다고 밝혔다.

AWS 연구팀은 Mamba-2, Gated DeltaNet(GDN), Gated KalmaNet(GKA)을 비교했다. 통제된 조건에서 GKA가 가장 높은 성능을 보였고 GDN, Mamba-2가 뒤를 이었다. 320억 매개변수 규모에서는 GKA 하이브리드가 같은 데이터로 학습한 원본 트랜스포머와 정확도 차이 1% 이내를 보이면서 추론 시 키·값 캐시(KV cache)를 절반가량 사용했다고 연구팀은 보고했다.

연구팀은 상태공간 모델 기반 하이브리드가 긴 문맥 벤치마크에서 SWA 기반 하이브리드보다 나은 결과를 냈다고 밝혔다. SWA에 상태공간 모델보다 메모리를 두 배에서 네 배 더 할당한 경우에도 이 결과가 유지됐다.

GKA는 추론할 때 회귀 계산의 반복 횟수를 조정할 수 있다. 별도 재학습 없이 반복 횟수를 보통의 30회에서 10회로 낮추면 AIME 2025에서 정확도가 약 1% 떨어지는 대신 디코딩 속도가 10% 빨라졌다고 연구팀은 설명했다. 따라서 하나의 모델로 속도와 정확도 사이의 설정을 조절할 수 있다.

긴 문맥 처리와 활용 대상

AIME 2025에서 GKA 기반 320억 매개변수 모델은 원본 Qwen3-32B와 비슷한 길이의 추론 과정을 생성하면서 더 높은 처리량을 보였다. 실제 동시 처리 상황에서 목표 정확도에 도달하는 시간은 전체 벤치마크에서 1.2~1.4배, 가장 어려운 문제에서 최대 1.6배 빨랐다고 AWS 연구팀은 밝혔다. 문맥 길이 128,000토큰에서 처리량은 원본 대비 최대 2.06배였고, 80억 매개변수 모델은 16,000토큰에서 1.94배, 128,000토큰에서 2.29배였다. 강화학습 후처리에서도 같은 GPU 예산으로 학습 데이터를 생성하는 속도가 20~100% 늘었다고 덧붙였다.

연구팀은 학습 없이 하이브리드 모델의 문맥 길이를 128,000토큰에서 256,000토큰으로 확장하는 방법도 제시했다. 긴 입력을 모델의 기본 문맥 길이보다 짧은 여러 조각으로 나눠 각각 처리한 뒤, 조각별 KV 캐시를 연결하고 상태공간 모델 상태는 가중 평균으로 합친다. 연구팀은 학습하지 않은 길이에서도 RULER NIAH 검색 과제와 BABILong 추론 과제에서 상당한 정확도를 유지했다고 보고했다.

코드, 모델, 최적화 커널과 vLLM 서빙 플러그인은 Apache 2.0 라이선스로 공개됐다. 공개 자료에는 80억·320억 매개변수 모델, 프라이밍 학습 코드, 상태공간 모델용 시퀀스 병렬화 도구와 GKA 커널이 포함된다. 연구팀은 이 자료가 긴 문맥 서비스 운영, 강화학습 후처리, 새 아키텍처 연구에 활용될 수 있다고 설명했다. 모델은 Qwen, Llama, Mistral 계열과 밀집형·전문가 혼합(MoE) 모델에 적용할 수 있으며, 연구팀은 Llama·Mistral·MoE로 확장할 때 알고리즘 변경은 필요하지 않다고 밝혔다.

자세한 실험 결과는 논문에서 확인할 수 있다. 학습 코드와 도구는 GitHub 저장소에, 모델은 Hugging Face 컬렉션에 공개돼 있다. 하이브리드 구조의 배경은 관련 연구에서 볼 수 있다.