Orthrus: Dual-View Diffusion으로 메모리 효율적인 병렬 토큰 생성을 위한 프레임워크

AR 기반 대형 언어 모델은 토큰 생성이 순차적으로 이뤄져 인퍼런스가 병목 현상을 겪는 한계가 있다. diffusion 기반의 병렬 생성은 속도를 올릴 수 있지만 품질 저하와 학습 비용 문제가 있다. Orthrus는 frozen AR 백본과 경량 diffusion head를 결합해 두 뷰가 동일한 고정 KV Cache를 공유하고, intra-model consensus로 lossless inference를 보장하며 병렬 생성을 가능하게 한다.

1. 시작점: Autoregressive(AR) 모델은 컨텍스트를 잘 구성하는 반면, 토큰 생성은 순차적이다. 이로 인해 메모리 대역폭과 지연이 발생한다. Diffusion Language Models(DLMs)은 토큰 블록을 병렬로 생성하지만, 조건독립성으로 인해 AR 수준의 예측 분포를 벗어나고 품질 저하가 나타난다. 2) 해결 원리: Orthrus는 AR 백본을 고정하고, trainable diffusion head를 추가해 두 뷰가 같은 컨텍스트에서 작동하도록 한다. diffusion head는 AR의 KV 캐시를 조건으로 하여 K개의 토큰을 한 번에 병렬로 예측하고, AR head의 target 분포와의 차이를 soft-distillation과 intra-model consensus로 맞춘다. 3) 달라지는 점: 이를 통해 메모리 오버헤드 없이 병렬 토큰 생성을 달성하고, exact causal distribution을 유지하므로 lossless한 속도 증가를 실현한다. 4) 학습/추론의 구체적 메커니즘은 Dual-Pass Block Masking과 Left-to-Right 합의 검증으로 구현된다. 5) 결과적으로, AR의 품질을 유지하면서 Diffusion 기반 병렬 생성을 실용적 속도로 가속하고, 8B 스케일에서 평균적으로 5.39의 TPF를 기록하며 최대 7.8×의 속도up를 달성한다.

단락 1: 전체 접근 방식과 핵심 아이디어—AR backbone을 고정하고 diffusion head를 추가하여 dual-view Attention을 구현하며, 두 뷰가 동일한 Shared KV Cache를 활용한다. 단락 2: 핵심 메커니즘—Diffusion Head의 Q/K/V를 AR KV Cache(KAR, VAR)에 대해 병렬로 주입하고, Block Masking으로 데이터 흐름을 엄격히 차단한다. 단락 3: 학습 전략—Dual-Pass Block Masking을 도입하고 forward KL(Diffusion) 대 AR의 완전한 분포를 Distill한다. 단락 4: 추론—Step 1에서 diffusion head가 K개의 토큰을 병렬 예측하고, Step 2에서 AR head가 이 예측을 검증하여 최종 토큰을 Accept/Reject 하며, Reject 시 prefix를 AR의 yJ로 수정하고 KV Cache를 축소한다. 수식: LOrthrus = Ex,{ab} Σ^B_b=1 Σ^K_k=1 DKL(pAR(· | x≤ab+k−1) ∥ pdiff(· | x<ab, y˜b))를 사용해 diffusion 뷰를 AR 분포로 소프트 디스틸레이션한다. 패러다임: 입력 x, y˜b는 anchor 포함 로 구성되며, pdiff는 x<t, y˜의 병렬 분포를 따른다. 그라디언트는 diffusion 모듈로만 흐르게 하고 AR 백본은 고정된다.

주요 벤치마크에서 AR와 동일한 정확도와 함께 병렬화로 인한 속도 향상을 달성한다. GSM8K, MATH-500, AIME-24, AIME-25, HumanEval, MBPP 등에서 평균적으로 수공정량의 속도up를 보이며, 8B 스케일에서 평균 TPF가 5.39에 도달한다. 또 다른 벤치에서 단일 스텝 예측( single-step) 전략이 다단 denoising 대비 더 높은 Throughput을 제공한다(예: 6.35 TPF). Inference 시 40K 토큰 컨텍스트에서도 DFlash 등 기존 speculative decoding 대비 안정적인 Throughput를 유지한다. AR 기반 baseline과의 exact Distribution Parity를 통해 lossless한 속도 증가를 달성하며, AR 8B 기반 벤치에서 동일한 zero-shot MATH-500 정확도를 보인다.

전체 아키텍처는 두 개의 Parallel Attention 경로(AR-Head, Diffusion-Head)와 Shared KV Cache로 구성된다. Diffusion-Head의 Q/K/V는 AR Head의 파라미터에서 초기화되며, Training은 AR backbone을 고정하고 Diffusion Head의 파라미터만 학습한다. LOrthrus = Ex,{ab} Σ^B_b=1 Σ^K_k=1 DKL(pAR(· | x≤ab+k−1) ∥ pdiff(· | x<ab, y˜b))로 soft distillation을 수행한다. Inference 시에는 Step1에서 pdiff를 통해 K개의 후보 토큰을 병렬로 예측하고, Step2에서 AR Head의 pAR 분포와의 일치성을 검사해 Accept 조건을 충족하는 토큰만 KV 캐시에 반영한다. 이로써 lossless한 분포 일치를 유지하며 속도up를 달성한다.

AR 백본을 Strictly Frozen하기 때문에 생성 능력은 기반 모델의 한계를 초과하지 못한다. diffusion head는 AR 교사의 정확한 예측 분포를 모사하도록만 학습되며, 원래의 바이어스나 한계를 그대로 상속한다. 이는 추론 가속의 한계로 보이며, 학습으로부터의 개선은 제한적이다.