Nemotron-Labs-Diffusion: AR·확산·Self-Speculation을 단일 아키텍처로 통합한 트라이모드 언어모델

왜 중요한가

AR 방식은 좌→우 언어적 우선성을 학습해 정확도가 높지만 병렬성이 낮아 저배치 환경에서 처리량이 제한된다. 확산 기반 디코딩은 블록 내 병렬 토큰 예측으로 처리량을 크게 높일 잠재력이 있으나 좌→우 정보 활용 부족으로 학습 효율과 정확도에서 AR에 뒤처지는 경향이 있다. 본 논문은 두 손실을 공동 최적화해 양쪽의 장점을 살리고, 자기-추측 모드를 통해 실장 환경에서 검증 비용을 통제하면서 병렬 처리를 실용화 가능한 수준으로 끌어올렸다는 점에서 의의가 있다.

핵심 기여

AR-확산 공동 목적함수와 두 단계 학습 파이프라인

모델을 먼저 순수 AR로 사전 학습한 뒤 알파(α)=0.3으로 확산 손실을 함께 최적화하는 두 단계 전략을 채택해 좌→우 언어적 편향을 유지하면서 확산 기반 병렬 예측 능력을 확보했다. 이 절차는 global loss averaging과 DP-rank 마스킹 비율 등 분산성 제어 기법과 결합되어 안정적인 학습을 이끌어냈다.

블록-단위 확산과 인과적 클린 스트림을 결합한 이중 스트림 주의 패턴

입력에 노이즈 스트림과 클린 스트림을 결합한 dual-stream을 적용해 블록 내에서는 양방향(attend)으로 병렬 예측을 지원하고 블록 간에는 인과적(attention causal) 관계를 유지해 KV 캐시를 재사용할 수 있게 했다. 클린 스트림에 엄격한 causal mask를 적용하여 동일 순전파에서 AR 손실과 확산 손실을 계산하면서 레이블 누출(label leakage)을 방지했다.

자기-추측(self-speculation)과 LoRA 기반 drafter 정렬

확산으로 k개 토큰을 병렬 초안 생성(draft)하고 동일 모델의 AR 모드로 검증(verify)해 접수하는 선형 및 이차적 자기-추측 알고리즘을 제안했다. 추가로 drafter의 o_proj에 LoRA(rank=128, α=512) 어댑터를 학습해 드래프트와 검증 로짓 분포 간 정합을 높였고, LK-hybrid + CE 손실로 수렴 안정성과 수용(prefix) 길이를 늘렸다.

Speed-of-Light(SOL) 분석과 실용 샘플러 설계

확산 모드 자체의 직렬 수렴 출력을 기준 목표로 삼아 병렬 정책이 달성할 수 있는 이론적 TPF 상한을 계산하는 SOL 방식을 도입했다. SOL은 블록 길이 B=32에서 평균 7.60×의 내재적 병렬성을 보였고 논문에서는 신뢰도 기반 샘플러와 학습된 샘플러를 비교해 실제 샘플러가 SOL과 큰 격차를 남긴다는 정량적 근거를 제공했다.

핵심 아이디어 이해하기

언어 생성에서 AR 모델은 좌→우 조건부 확률을 최대화하여 다음 토큰 예측에 탁월한 언어적 사전(left-to-right prior)을 학습한다. 이 사전은 문맥을 이용한 계획(planning) 능력을 내재화시켜 단일 토큰 예측의 정확도를 보장하지만, 토큰 단위 순차 디코딩은 병렬 처리 능력을 제한해 낮은 동시성 환경에서 처리량 저하를 초래한다. 반대로 확산(discrete diffusion) 방식은 블록 내 다수 토큰을 동시에 노이즈 제거함으로써 병렬 디코딩을 가능하게 하나, 모든 토큰 순열을 동등 처리하는 학습 특성으로 인해 좌→우 구조를 충분히 활용하지 못해 학습 효율과 정확도에서 손해를 보는 경향이 있다.

방법론

전체 접근은 AR 손실과 블록 단위 확산 손실을 가중 합한 공동 목적함수 L(θ)=L_AR(θ)+α L_diff(θ)를 사용하고, α를 0.3으로 설정해 두 목적의 크기를 정렬했다. 학습은 두 단계로 구성되어 먼저 순수 AR 학습으로 강한 좌→우 프라이어를 구축하고 다음에 joint objective로 전환해 확산 기반 병렬 능력을 통합했다. 블록-단위 확산은 시퀀스를 B개의 연속 블록으로 나누고, 노이즈 스트림에서는 각 블록 내 양방향 attention과 블록 간 인과적 attention을 허용하며 클린 스트림에는 엄격한 causal mask를 적용해 동일 전향-역전파에서 AR과 diffusion 손실을 동시에 계산할 수 있게 했다. 확산 디코딩에서는 confidence-based sampling과 학습된 샘플러를 도입하여 각 마스크 위치의 확정 여부를 예측하고, 자기-추측 모드에서는 diffusion으로 초안 생성 후 AR로 검증하는 두 패스 전략을 사용해 accepted prefix를 커밋한다. 자기-추측의 drafter 정렬을 위해 LoRA 어댑터를 o_proj에 적용하고 LK-hybrid 분포 정합 손실과 토큰 수준 CE를 결합해 드래프트의 수용 확률과 분포 간 차이를 줄였다.

주요 결과

학습 및 평가 결과에서 α=0.3 설정이 두 모드(AR·diffusion) 모두에 최적의 균형을 제공하여 양쪽 성능을 함께 유지하거나 소폭 향상시켰다. 논문은 25B 토큰 연속 프리트레이닝 및 45B 토큰 SFT 환경에서 Ablation을 수행했고, global loss averaging과 two-stage training, AR 손실 추가가 누적되어 평균 정확도를 16.05% 포인트 향상시킨 점을 보고했다. Nemotron-Labs-Diffusion-8B(instruct)는 여러 벤치마크에서 동급 AR 모델 대비 동등하거나 더 높은 정확도를 유지하면서 diffusion 모드에서 약 2.57× TPF, LoRA 보강 자기-추측에서는 최대 5.99× TPF, 이차 자기-추측에서는 6.38× TPF를 달성해 실효 처리량을 크게 개선했다. SOL 분석에서는 최적 샘플러를 가정할 때 diffusion이 자기-추측 대비 최대 76.5% 더 많은 토큰을 단일 전방 패스로 예측할 수 있는 잠재력이 있음을 확인했다.

기술 상세

전체 아키텍처는 표준 Transformer 기반의 LM 백본을 공유하되 입력을 노이즈 스트림과 클린 스트림으로 중첩하는 dual-stream 레이아웃을 사용한다. 블록-단위 확산에서는 시퀀스를 B 블록으로 나누고 노이즈 스트림의 각 블록 내에서는 bidirectional attention, 블록 간에는 causal attention을 적용해 블록별 병렬 예측과 KV 캐시 재사용을 동시에 달성했다. 학습 목적함수는 AR 손실 L_AR(θ)=E_x[-∑{i=1}^{|x|} log p_θ(x_i|x{<i})]와 확산 손실 L_diff(θ)=E_{t∼U[0,1], ilde{x}t^b∼q(⋅|x^b)}[-(1/t) ∑{b=1}^B log p_θ(x^b | ilde{x}_t^b, x^{<b})]를 합한 형태이며 α로 확산 손실의 기여도를 조절한다. 수식 해석 예: AR 손실은 각 위치 i에서 조건부 확률 p_θ를 통해 -log 확률을 누적한 평균이므로 모델이 예측을 잘 할수록 값이 작아진다; 확산 손실은 노이즈 레벨 t에 대해 1/t로 재가중하면 작은 t에서 부과된 토큰이 더 큰 영향력을 가지게 된다. 이 때문에 전역 평균(global loss averaging)을 채택해 마스크된 토큰 수가 다른 샘플들이 배치 손실에 비대칭적으로 기여하는 문제를 완화했다. 자기-추측의 LoRA 설정은 o_proj에 rank=128, α=512로 약 36M 학습 파라미터(약 0.4% 백본)를 추가하여 drafter의 분포를 AR 검증자에 맞춰 정렬하도록 학습하며, LK-hybrid와 CE 손실을 결합해 truncation된 상위-K 지원에서 분포 정합과 modal 토큰 추종을 동시에 확보했다.

한계점

논문이 명시한 한계로는 현재 실용 샘플러들이 SOL 상한과 큰 격차를 보이며, 확산 모드에서 안전하게 비접두(prefix가 아닌) 위치를 대량 커밋하는 방법이 아직 미해결 문제로 남아 있다. 또한 블록 길이가 커질수록 벤치마크 정확도가 감소하는 경향이 관찰되어 대형 블록에서의 확산 성능 유지가 추가 연구 과제로 제시되었다. 자기-추측 모드는 AR 검증의 추가 전방 패스 비용 때문에 실효 TPF가 SOL보다 낮게 제한되는 구조적 한계가 있으며, 이 제약을 해소하려면 검증 단계의 구조적 변경 또는 대체 검증기 설계가 필요하다.

실무 활용

트라이모드 설계는 기존 AR 파이프라인을 변경하지 않고도 동일 모델이 저동시성 환경에서는 자기-추측, 고동시성 환경에서는 AR, 중간 지점에서는 확산 기반 디코딩으로 전환해 운영 조건에 맞는 처리량/정확도 균형을 제공한다. LoRA 어댑터와 학습된 샘플러는 실무 배포에서 경량 파라미터 조정으로 병렬성 이득을 끌어올리는 현실적인 경로를 제공한다.

저동시성 대화 서비스에서 자기-추측 모드를 사용해 응답 지연을 낮추고 사용자당 처리량을 늘리는 배포
코딩·수학처럼 구조적 제약이 강한 태스크에 대해 블록-단위 확산 모드를 적용해 배치 크기 1 환경에서 토큰 병렬성을 활용하는 추론 파이프라인
비전-언어 응용에서 고해상도 이미지 토큰을 클린 스트림에 유지하고 텍스트 응답 부분만 확산 적용해 연산 오버헤드를 줄이는 VLM 배포

코드 공개 여부: 미확인

키워드

LLM(대형 언어 모델)Diffusion(확산 모델)Autoregressive(자기회귀)Self-Speculation(자기-추측 디코딩)LoRA(저순위 파인튜닝)SPEED-Bench(추론 벤치마크)

용어 해설

블록 단위 확산(Block-wise Diffusion): — 시퀀스를 여러 개의 연속 블록으로 분할한 뒤 각 블록을 독립적으로 노이즈-되돌리기(denoising)하도록 학습하는 방식이다. 블록 내에서는 양방향(attention)을 사용해 병렬로 토큰을 예측하고 블록 간에는 인과적 연결을 유지해 KV 캐시를 재사용할 수 있게 한다. 이 구조는 확산 기반 병렬 디코딩에서 토큰 병렬성을 확보하면서도 AR 모델의 캐시 이점을 살리는 데 중요하다.
자기-추측 디코딩(Self-Speculation): — 확산 모드로 다수의 토큰 초안을 병렬로 생성한 뒤 동일 모델의 AR 모드로 그 초안들을 검증하여 접수(prefix)하는 방식이다. 초안(draft)은 병렬로 생성되어 GPU 활용도를 높이고 AR 검증은 정확도를 확보하므로, 드래프트와 검증 간 alignment가 높을수록 단위 전방 패스당 커밋되는 토큰 수가 증가한다. MTP 계열과 달리 별도 보조 예측 헤드를 요구하지 않으며 단일 모델에서 드래프트·검증을 모두 수행한다.
속도 한계 분석(Speed-of-Light Analysis): — 확산 디코딩이 가질 수 있는 이론적 상한을 계산하는 기법으로, 직렬(한 위치씩) 디노이징으로 얻어지는 수열을 기준 목표로 삼아 병렬 정책이 동일한 목표를 재현할 때의 토큰-퍼-포워드(TPF)를 산정한다. 이 분석은 현재 샘플러가 남긴 병렬성 여지를 정량화하여 최적 샘플러의 잠재력을 평가한다. 논문에서는 이 분석을 통해 실용적 샘플러와의 간격(예: 76.5% 향상 여지)을 도출했다.
신뢰도 기반 샘플링(Confidence-Based Sampling): — 각 마스크 위치에서 모델의 확률 분포가 가진 최고값(peak probability) 또는 학습된 분류기의 출력값을 이용해 해당 위치를 즉시 확정(커밋)할지 결정하는 정책이다. 임계값을 통해 TPF와 오류율 사이를 조절하며, 학습된 샘플러는 이 정책을 데이터에 맞춰 보정해 병렬성 확보를 시도한다. 논문에서는 기본 신뢰도 방식이 SOL 상한 대비 큰 격차를 남긴다고 관찰되었다.
전체 토큰 평균 손실(Global Loss Averaging): — 배치 내 모든 토큰 손실을 균등하게 평균하여 샘플마다 마스크된 토큰 수가 달라지는 경우에 특정 샘플(예: 작은 t 값으로 가중치가 큰 항)을 과도하게 영향력 있게 만드는 현상을 완화하는 집계 방식이다. 시퀀스 단위 평균과 대비되어 배치 간 변동성을 줄이고 학습 안정성을 높이는 역할을 한다. 논문에서 이 기법은 diffusion 학습에서 성능 향상에 기여했다.

Nemotron-Labs-Diffusion: AR·확산·Self-Speculation을 단일 아키텍처로 통합한 트라이모드 언어모델

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드