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

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

언어 생성에서 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보다 낮게 제한되는 구조적 한계가 있으며, 이 제약을 해소하려면 검증 단계의 구조적 변경 또는 대체 검증기 설계가 필요하다.