기술 요약: Nemotron-Labs-Diffusion

문제 정의

자기회귀(Autoregressive, AR) 언어 모델(LM)은 엄격하게 순차적인 토큰 단위 디코딩 프로세스로 인해 근본적인 한계를 지니며, 이는 추론의 병렬성을 제한하고 특히 낮은 배치 크기 시나리오에서 자원 활용도를 저하시킵니다. 확산(Diffusion) LM은 한 번의 순방향 패스(forward pass) 동안 다중 토큰을 디코딩함으로써 병렬 생성을 가능하게 하여 유망한 대안으로 부상했으나, 현재 정확도와 학습 효율성 측면에서 AR 모델에 뒤처져 있습니다. 기존의 확산 LM은 대등한 성능에 도달하기 위해 실질적으로 더 많은 데이터를 필요로 하며, 실제적인 효율성-정확도 트레이드오프 측면에서 다중 토큰 예측(Multi-Token Prediction, MTP) 방식보다 뛰어난 성능을 내는 데 어려움을 겪습니다. 또한, AR과 확산 패러다임 사이의 관계는 여전히 불분clear합니다. 즉, 이들이 경쟁 관계인지, 아니면 양측의 강점을 활용하기 위해 조화될 수 있는 것인지가 불확실합니다.

방법론

저자들은 AR, 확산, 그리고 자기 추측(self-speculation) 디코딩을 단일 아키텍처 내에 통합한 삼중 모드(tri-mode) 언어 모델 제품군인 Nemotron-Labs-Diffusion을 소개합니다. 핵심 방법론은 공동 학습 프레임워크와 유연한 추론 전략을 포함합니다.

1. 공동 AR-확산 학습 (Joint AR-Diffusion Training)

모델은 AR 차기 토큰 손실(next-token loss)과 블록 단위 확산 디노이징 손실의 가중 조합을 사용하여 학습됩니다:
$$\mathcal{L}(\theta) = \mathcal{L}_{AR}(\theta) + \alpha \mathcal{L}_{diff}(\theta)$$

• 2단계 학습: 강력한 좌우(left-to-right) 언어적 사전 지식(priors)을 구축하기 위해, 저자들은 2단계 전략을 채택합니다. 1단계에서는 AR 목적 함수만으로 학습합니다 ($\alpha=0$). 2단계에서는 확산 감독(diffusion supervision)을 도입하여 ($\alpha=0.3$), AR 사전 지식을 덮어쓰지 않으면서 병렬 디코딩 능력을 통합합니다.
• 글로벌 손실 평균화 (Global Loss Averaging): 확산 학습 시 가변적인 마스킹된 토큰 수로 인한 최적화 불안정성을 해결하기 위해, 저자들은 샘플의 노이즈가 적은 토큰 수로 인한 그래디언트 분산을 방지하고자 시퀀스별 평균화 대신 배치 전체에 걸친 글로벌 손실 평균화 전략을 사용합니다.
• 어텐션 패턴 (Attention Pattern): 모델은 구조화된 어텐션 마스크를 가진 이중 스트림 입력(노이즈가 섞인 뷰와 깨끗한 뷰)을 활용합니다. 결정적으로, 깨끗한 스트림은 엄격한 인과적(causal) 마스크를 강제하여 레이블 누출 없이 깨끗한 토큰에 대해 AR 목적 함수를 계산할 수 있게 하며, 노이즈가 섞인 스트림은 병렬 디노이징을 위해 블록 내에서 양방향 어텐션을 사용합니다.

2. 삼중 모드 추론 (Tri-Mode Inference)

학습된 모델은 추론 시 선택 가능한 세 가지 디코딩 모드를 지원합니다:

• AR 모드: 인과적 어텐션을 사용하는 표준 좌우 생성 방식입니다. 이 모드는 고동시성(high-concurrency) 클라우드 서빙에 최적화되어 있습니다.
• 확산 모드: 병렬 블록 단위 디노이징을 수행합니다. 모델은 신뢰도 임계값을 기반으로 반복적으로 여러 토큰을 디노이징합니다. 상위 1위(top-1) 예측이 올바른지 예측하기 위해 경량화된 학습된 샘플러를 사용하여, 토큰당 순방향 패스(TPF)와 에러율 사이의 트레이드오프를 최적화합니다.
• 자기 추측 모드 (Self-Speculation Mode): 확산 경로가 여러 후보 토큰을 병렬로 초안(draft)을 작성하고, AR 경로가 이를 검증하는 하이브리드 접근 방식입니다.
  • 선형 자기 추측 (Linear Self-Speculation): 확산 모델이 $k$개의 토큰을 초안으로 작성하면, AR 모델이 두 번째 순방향 패스에서 이를 검증합니다. 저자들은 확산 초안 작성 모델을 AR 검증기에 더 잘 정렬시키기 위해 LoRA 어댑터를 튜닝하여 이를 개선했습니다.
  • 이차 자기 추측 (Quadratic Self-Speculation): 단일 순방향 패스 내에서 가능한 모든 수용 위치를 준비하기 위해 이차 디코딩을 사용하는 변형이지만, 현재 커널 최적화 문제로 인해 효율성 과제에 직면해 있습니다.

3. 빛의 속도(Speed-of-Light, SOL) 분석

저자들은 확산 디코딩의 이론적 상한선을 정량화하기 위해 "빛의 속도" 분석을 제안합니다. 재귀적 동적 압축(recursive dynamic compaction) 방법을 사용하여 직렬 디노이징 대상과 일치하는 가장 큰 안전한 부분 집합을 식별함으로써, AR 검증 없이 달성 가능한 최대 TPF를 추정합니다. 이 분석은 현재의 신뢰도 기반 샘플링이 상당한 병렬성을 남겨두고 있음을 보여줍니다.

주요 기여

1. 통합 삼중 모드 아키텍처: AR, 확산, 자기 추측 디코딩을 원활하게 통합하여 배포 요구 사항(예: 고동시성 vs 저동시성)에 따라 동적 전환이 가능한 최초의 모델입니다.
2. 목적 함수의 상보성: AR과 확산 목적 함수가 경쟁 관계가 아니라 상보적임을 입증합니다. AR은 강력한 언어적 사전 지식과 미래 계획을 제공하며, 확산은 병렬 생성을 가능하게 합니다. 결합된 목적 함수는 AR 정확도를 유지하거나 약간 향상시키면서 고속 확산 모드를 활성화합니다.
3. MTP 대비 우월성: 자기 추측(확산 초안 작성 + AR 검증)이 특히 낮은 동시성 환경에서 기존 MTP 방식(Eagle3 등)보다 수용률(acceptance rates)과 실제 장치 효율성 측면에서 더 뛰어남을 보여줍니다.
4. 확장성 및 다재다유: 모델 제품군은 3B에서 14B 파라미터까지 확장되며 시각-언어 모델(VLM)로도 확장되어, 최첨단 오픈 소스 AR 및 확산 LM을 일관되게 능가합니다.

결과

• 정확도 및 효율성: Nemotron-Labs-Diffusion-8B 모델은 일반, 코딩, 수학 벤치마크 전반에서 Qwen3-8B와 대등하거나 더 나은 정확도를 달면서 확산 모드에서 순방향 패스당 6배 더 많은 토큰을 디코딩합니다.
• 처리량 이득: SPEED-Bench에서 이 모델은 NVIDIA GB200 GPU 상의 SGLang을 사용하여 Qwen3-8B보다 4배 높은 처리량을 달성합니다.
• 자기 추측 성능: LoRA 튜닝을 적용한 선형 자기 추측은 8B 모델에서 5.99×, 14B 모델에서 **6.38×**의 평균 TPF를 달성하여 베이스라인을 크게 앞섭/니다.
• SOL 잠재력: SOL 분석에 따르면, 최적의 샘플러를 사용할 경우 확산 디코딩은 현재의 자기 추측 방식보다 76.5% 더 많은 토큰을 순방향 패스당 예측할 수 있는 잠재력이 있으며, 이는 상당한 미래 발전 여지를 시사합니다.
• VLM 성능: VLM 변형 모델은 AR 베이스라인과 경쟁력 있는 정확도를 유지하면서, 확산 모드에서 2.46×–3.15× TPF를, 긴 응답에 대해 선형 자기 추측 모드에서 최대 7.45× TPF를 제공합니다.

의의 및 주장

본 논문은 Nemotron-Labs-Diffusion이 AR과 확산이 서로 배타적인 것이 아니라, 더 유연하고 효율적인 추론 엔진을 만들기 위해 통합될 수 있다는 패러다임의 전환을 나타낸다고 주장합니다. 저자들은 다음과 같이 단언합니다:

• 조화 가능성: AR과 확산 목적 함수를 통합하는 것은 순수 확산 모델과 관련된 정확도 저하 문제를 피할 수 있는 유망한 경로입니다.
• 배포 유연성: 삼중 모드 설계는 단일 모델이 고동시성 환경에서는 AR 모델의 드롭인 교체재 역할을 하는 동시에, 저동시성 개인용 AI 시나리오에서는 자기 추측을 통해 상당한 속도 향상을 제공할 수 있게 합니다.
• 미래 잠재력: "빛의 속도" 분석은 확산 디코딩의 전체 잠재력이 아직 완전히 실현되지 않았음을 시사합니다. 현재 성능과 SOL 천장 사이의 간극은 샘플러 최적화에 대한 향후 연구가 더 큰 병렬성을 확보하여, 프리픽스(prefix) 기반의 추측적 디코딩의 한계를 넘어설 수 있음을 나타냅니다.

결론적으로, 현재의 인프라는 선형 자기 추측에 유리하지만, 확산 디코딩의 이론적 상한선은 매우 병렬적인 비-프리픽스(non-prefix) 생성으로 향하는 장기적 궤적을 제시하며, 이는 LLM 추론 효율성을 근본적으로 변화시킬 수 있습니다.