Orthrus: AR Transformer를 위한 확산 기반 추론 가속 기법

섹션별 상세

Orthrus는 고정된 AR Transformer의 각 층에 학습 가능한 확산 어텐션 모듈을 주입하여 작동한다. 확산 헤드가 32개의 토큰을 병렬로 투영하고, AR 헤드가 이를 검증하여 가장 긴 일치 접두사를 선택한다. 이 방식은 출력 분포를 원본 모델과 동일하게 유지하면서 추론 속도를 높인다.

기존 확산 기반 언어 모델들은 기본 가중치를 수정하여 정확도 손실이 발생했으나, Orthrus는 백본을 고정하여 Qwen3-8B와 동일한 정확도를 유지한다. 8xH200 환경에서 24시간 학습으로 16%의 파라미터를 훈련하여 MATH-500 벤치마크에서 7.8배의 TPF 향상을 달성했다.

기존 Speculative Decoding 기법인 EAGLE-3나 DFlash와 달리 외부 드래프터 모델이나 별도의 캐시가 필요하지 않다. 이로 인해 TTFT 페널티가 없으며, KV 캐시 오버헤드는 약 4.5 MiB로 매우 낮다. MATH-500에서 11.7의 수용 길이를 기록하여 DFlash(7.9)와 EAGLE-3(3.5)를 상회했다.

용어 해설

추론 가속 기법(Speculative Decoding): — 작은 모델이 여러 토큰을 미리 생성하고 큰 모델이 이를 한 번에 검증하여 추론 속도를 높이는 기법. Orthrus는 외부 모델 없이 이를 구현함.
확산 모델(Diffusion Model): — 데이터의 노이즈를 점진적으로 제거하여 샘플을 생성하는 확률적 모델. Orthrus는 이를 언어 모델의 추론 가속에 활용함.
KV 캐시(KV Cache): — Transformer 모델 추론 시 이전 토큰의 Key와 Value 값을 저장하여 중복 계산을 방지하는 메모리 구조.
자기회귀 트랜스포머(AR Transformer): — 이전 토큰들을 기반으로 다음 토큰을 순차적으로 예측하는 구조의 모델.

언급된 리소스

논문Orthrus Paper

GitHubOrthrus GitHub