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

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

왜 중요한가

핵심 기여

Dual-Architecture의 내재적 병렬화

AR backbone을 고정하고 diffusion head를 학습시켜 두 뷰가 단일 공유 KV Cache에서 작동하도록 구성한다. 학습 시 diffusion head의 가중치만 업데이트되며 전체 모델의 약 16% 파라미터를 차지한다.

Intra-Model Consensus로 Lossless Inference 보장

diffusion head의 병렬 예측(pd i ff )와 AR head의 정확한 목표분포(pAR)를 좌측에서 우측으로 검증하는 합의 메커니즘을 도입해 분포의 차이를 제거하고, 필요 시 reject 샘플링으로 완전한 일치를 보장한다.

메모리 및 파라미터 효율 극대화

공유 KV 캐시를 통해 추가적인 KV 캐시 저장이 필요 없고, peak 메모리 오버헤드는 약 ∼100 MiB에 불과하며 KV 캐시 오버헤드는 상수(O(1)) 규모인 약 4.5 MiB로 고정된다.

학습 기반의 경량화된 파이프라인

2에폭 학습, 앵커 블록 256개, 블록 크기 K=32를 사용한 Dual-Pass Block Masking으로 AR 교사 분포를 diffusion 뷰에 소프트 디스틸레이션한다.

Plug-and-Play 확장성

Qwen3-8B 계열과 같은 높은 품질의 AR 모델에 diffusion 헤드를 주입해 병렬 토큰 생성을 가능하게 하며, 24시간 이내에 8xH200 노드에서 배치 가능하도록 설계되었다.

핵심 아이디어 이해하기

시작점: 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 교사의 정확한 예측 분포를 모사하도록만 학습되며, 원래의 바이어스나 한계를 그대로 상속한다. 이는 추론 가속의 한계로 보이며, 학습으로부터의 개선은 제한적이다.

실무 활용

Orthrus는 frozen AR 백본에 경량 diffusion head를 주입하는plug-and-play 프레임워크로, 단일 KV Cache를 공유하면서 고속 병렬 토큰 생성을 가능하게 한다. 이로써 프로덕션 수준의 고속 inference에 적합한 무손실(losless) 가속이 가능하다.

대화형 AI 시스템의 실시간 응답 가속
대규모 코드/문서 생성 도구의 파이프라인 throughput 향상
Long-context 조건의 자동 코드 작성 및 문서 요약
실시간 AI 코딩 보조 도구의 응답 속도 향상
대규모 다중 도메인 인퍼런스 워크플로우의 효율화

코드 공개 여부: 공개

코드 저장소 보기

키워드

autoregressive Large Language Modelsdiffusion modelsparallel token generationTransformerKey-Value cacheconsensus mechanismlossless inferencedual-architecture

추가 이미지 분석

Diagram
병렬 블록 크기(K)가 Throughput에 미치는 영향을 정량화하여, K=32에서 최적임을 뒷받침한다.
Ablation과 Throughput-Latency 관계를 보여주는 도표—K 값에 따른 Throughput 증가와 Latency의 변화를 동시에 제시

용어 해설

KV 캐시(KV cache): — Transformer의 키-값 캐시를 재활용해 컨텍스트를 빠르게 재구성하는 메모리 구조. AR와 diffusion 뷰가 동일 KV 캐시를 공유함으로써 중복 저장을 피하고 대용량 컨텍스트를 효과적으로 처리할 수 있다.
Diffusion 헤드(Diffusion Head): — AR 백본의 고정된 표현 공간 위에 학습 가능한 diffusion 모듈을 추가해 병렬 토큰 생성을 수행하는 구성요소. 학습 시에는 AR의 예측분포를 확률적 분포로 소프트 디스틸레이션한다.
intra-모델 합의(Intra-Model Consensus): — Diffusion 뷰의 예측이 AR 뷰의 정확한 목표분포와 일치하도록 내부적으로 검증하고 채택/거절을 통해 최종 출력을 동기화하는 메커니즘.
앵커 블록 마스킹(Anchor-Block Masking): — 훈련 시 각 블록에서 첫 토큰을 anchor로 두고 남은 토큰을 <mask>로 대체해 diffusion 뷰가 병렬 예측하도록 하는 마스킹 전략.

시작점: 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를 달성한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

추가 이미지 분석

용어 해설

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형

Grouped Query Attention의 작동 원리와 KV cache 최적화

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형

Grouped Query Attention의 작동 원리와 KV cache 최적화