[논문리뷰] Nemotron-Labs-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

arXiv 2026. [Paper] [Hugging Face]
Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro
NVIDIA
25 Jun 2026

Introduction

본 논문에서는 TwoTower를 소개하고, Mamba-2, attention, Mixture-of-Experts (MoE) layer를 인터리빙한 하이브리드 모델인 Nemotron-Labs-TwoTower를 Nemotron-3-Nano-30BA3B에 구현했다. 저자들은 사전 학습된 네트워크의 두 복사본을 가져와 상호 보완적인 역할을 부여했다. AR 컨텍스트 타워는 깨끗한 토큰에 대한 인과적인 모델로 고정되어 backbone의 AR (autoregressive) 기능을 유지한다. Diffusion denoiser 타워는 각 noisy block 내에서 양방향 self-attention을, 그리고 과거의 깨끗한 컨텍스트에 대한 causal cross-attention을 사용하는 masked diffusion objective를 통해 학습된다.

타워들은 layer별로 연결된다. 각 denoiser layer는 컨텍스트 타워의 해당 layer와 cross-attention되어 denoiser가 backbone의 표현에 멀티스케일로 접근할 수 있도록 한다. 이는 마지막 hidden state만 브로드캐스팅하는 기존 접근 방식과 대조적이다. 또한 denoiser는 adaLN을 통해 diffusion timestep에 따라 변조되는데, 이는 이미지 DiT에서는 일반적이지만 masked diffusion 언어 모델에서는 덜 일반적이다.

약 2.1조 개의 토큰으로 학습된 Nemotron-Labs-TwoTower는 AR 기반 알고리즘의 품질을 대부분 유지하면서 30B 하이브리드 MoE 아키텍처에서 훨씬 높은 처리량을 제공한다. Nemotron-Labs-TwoTower는 디코딩 초기 단계에서 정제 단계마다 여러 토큰을 커밋하여, 한 번에 하나의 토큰만 처리하는 AR 디코딩 방식보다 실제 처리 속도가 향상된다.

Method

1. Two-Tower Architecture

TwoTower는 사전 학습된 AR 언어 모델에 적용할 수 있는 일반적인 접근 방식이다. 본 논문에서는 52개의 layer (Mamba-2 layer 23개, self-attention layer 6개, MoE layer 23개)로 구성된 Nemotron-3-Nano-30B-A3B 모델에 TwoTower를 적용했다.

이 네트워크의 복사본 두 개를 만들고 각각에 상호 보완적인 역할을 부여한다. 프롬프트가 주어지면 AR 컨텍스트 타워는 프롬프트와 이전 토큰에 대해 AR 모델처럼 작동하여 layer별 KV 쌍과 최종 Mamba state를 생성한다. 생성은 block 단위로 진행된다. 각 새 block은 $S$개의 [MASK] 토큰으로 초기화되고 diffusion denoiser 타워에서 $T$번의 denoising step을 거쳐 반복적으로 정제된다.

Denoiser의 각 layer에서 attention layer는 해당 컨텍스트 타워 KV 캐시에 cross-attention하고, Mamba-2 layer는 해당 컨텍스트 타워 Mamba state에서 초기 state를 시드한다. Block의 모든 토큰이 정제되면 block이 커밋된다. 컨텍스트 타워는 커밋된 토큰을 인과적으로 처리하여 캐시를 업데이트하고, 다음 block으로 생성이 계속된다.

Inference 시, TwoTower는 컨텍스트 타워 가중치에 고정된 메모리 공간을 사용하면서 하나의 prefix cache만 유지한다. 컨텍스트 타워만 block 간에 KV 및 Mamba state를 유지하고 block이 커밋될 때마다 업데이트하므로 시퀀스 길이에 따라 달라지는 캐시 메모리는 AR 모델처럼 확장된다.

컨텍스트 타워 본체는 변경하지 않고 그대로 유지한다. 최종 vocabulary projection과 언어 모델(LM) head는 선택 사항이다. 컨텍스트 타워가 state만 생성하면 되는 기본 diffusion 생성 경로에서는 생략할 수 있으며, 컨텍스트 타워가 speculative decoding, 검증, likelihood 평가, AR scoring에 사용될 때는 유지한다.

이와 대조적으로, denoiser 타워는 diffusion 학습 및 inference에 적합하도록 몇 가지 아키텍처를 수정했다.

양방향 Attention

정제 대상인 block 내에서 denoiser는 causal mask를 완화한다. 구체적으로, noisy block의 토큰은 이전의 깨끗한 block에 대해서는 인과 관계를 유지하면서 다른 noisy block 내의 토큰에 대해 양방향으로 attention한다. Denoiser layer $i$에서 noisy block의 query는 컨텍스트 타워의 layer $i$에 있는 이전 block $0, \ldots, b-1$에 대한 KV와 denoiser 자체 layer $i$에 있는 KV의 concat된 key/value 시퀀스에 대해 attention한다.

\[\begin{equation} \textrm{Attn} (\textbf{Q}_b^{(i)}, [\textbf{K}_{<b}^{\textrm{ctx}, (i)}; \textbf{K}_b^{\textrm{den}, (i)}], [\textbf{V}_{<b}^{\textrm{ctx}, (i)}; \textbf{V}_b^{\textrm{den}, (i)}]) \end{equation}\]

Cross-attention은 layer 정렬 방식으로 작동한다. 즉, denoiser layer $i$는 context layer $i$에 attention을 적용한다. 두 타워 모두 동일한 사전 학습된 체크포인트에서 초기화되므로, 동일 인덱스의 layer는 유사한 표현 수준에서 작동하여 layer 정렬 방식의 cross-attention이 자연스러운 조합을 이룬다.

양방향 Mamba

저자들은 사전 학습된 Mamba-2 가중치를 zero state에서 왼쪽에서 오른쪽으로, 오른쪽에서 왼쪽으로 각각 실행한 후 출력을 평균화하는 방식으로 파라미터가 없는 양방향 Mamba-2 변형을 테스트했다. 이 방식은 layer당 SSM FLOPs를 대략 두 배로 증가시켰다. 품질 향상이 미미했기 때문에 최종 설계에서는 Mamba의 인과적 특성을 유지했다.

Timestep 컨디셔닝

DiT의 adaLN-single을 사용하여 timestep $t$에 따라 denoiser를 컨디셔닝한다. 글로벌 MLP는 $t$를 공유 scale, shift, gate 파라미터에 매핑하고, 각 layer별로 학습된 임베딩을 사용하여 layer별 변조를 수행한다. 30B backbone에서 이는 1.5M 파라미터만 추가한다. 이러한 작은 모듈들을 샤딩하는 대신 각 tensor-parallel rank에 복제하여 추가 통신 비용을 방지하였다.

Expert Routing

저자들은 시퀀스 레벨 로드 밸런싱을 포함하여 backbone의 기존 MoE 라우팅 메커니즘을 사용하였다. 토큰은 adaLN 변조를 통해 noise-aware하다. 명시적인 라우팅 선택은 하지 않으며, expert가 denoising objective를 통해 전문화를 학습하도록 하였다.

2. Block Diffusion Language Modeling

저자들은 각 AR block에 적용되는 masked diffusion 프레임워크 하에서 denoiser를 학습시켰다. $S$개의 토큰으로 이루어진 $b$번째 block을 \(\textbf{x}_b = (x_b^1, \ldots, x_b^S)\)로 나타내고, 이전 깨끗한 block에 대한 컨텍스트 타워 attention KV 캐시와 block $b-1$ 이후의 layer별 Mamba state를 \(\textbf{c}_{< b} = (\textrm{KV}_{< b}, \textrm{State}_{b-1})\)로 나타내자. 시퀀스 분포를 block-autoregressive로 모델링한다.

\[\begin{equation} \log p_\theta (\textbf{x}) = \prod_{b=1}^B \log p_\theta (\textbf{x}_b \mid \textbf{x}_{<b}) \end{equation}\]

Diffusion time $t \in (0, 1]$ 동안, forward process는 각 토큰을 독립적으로 대체하여 깨끗한 block \(\textbf{x}_b\)를 noise가 있는 block \(\textbf{z}_t^b\)로 변형시킨다. Schedule \(\alpha_t\)는 noise level $t$에서 깨끗한 토큰이 마스킹되지 않은 상태로 유지될 확률이므로, \(1 − \alpha_t\)는 마스킹 확률이다. 저자들은 \(\alpha_t = 1 - t\)를 사용하였다. 토큰 $y$에서의 점 질량을 \(\delta_y\)로, [MASK] 점 질량을 $\textbf{m}$이라 할 때, 토큰별 marginal은 다음과 같다.

\[\begin{equation} q (\textbf{z}_t^{b, \ell} \mid x_b^{\ell}) = \textrm{Concat} \left( \textbf{z}_t^{b, \ell}; \alpha_t \delta_{x_b^{\ell}} + (1 - \alpha_t) \textbf{m} \right) \end{equation}\]

$t$가 작으면 block 손상이 적은 반면, $t$가 크면 완전히 마스킹된 블록에 가까워진다.

Denoiser는 noise가 포함된 block, diffusion time, 컨텍스트 캐시를 조건으로 마스킹된 위치에서 깨끗한 토큰을 예측한다. Masked diffusion의 ELBO를 유도해 보면, 결국 마스킹된 토큰들에 대한 negative log-likelihood (NLL)를 시간 가중하는 objective가 자연스럽게 나온다. \(\alpha_t = 1 - t\)의 경우 이론적 가중치는 $1/t$이다. 학습 과정에서 안정성을 위해 이 중요도 가중치를 생략하고 마스킹된 위치에 대한 평균 NLL을 최적화한다.

\[\begin{equation} \mathcal{L}_\textrm{MD} = \mathbb{E}_{t, \textbf{z}_t} \left[ \frac{1}{\vert \mathcal{M}_t \vert} \sum_{(b, \ell) \in \mathcal{M}_t} -\log p_\theta (x_b^\ell \mid \textbf{z}_t^b, t, \textbf{c}_{<b}) \right] \\ \textrm{where} \quad \mathcal{M}_t = \{ (b, \ell) : \textbf{z}_t^{b, \ell} = [\textrm{MASK}] \} \end{equation}\]

샘플링

샘플러는 컨텍스트 타워에서 prefix cache \((\textrm{KV}_{< b}, \textrm{State}_{b-1})\)를 받아 block $b$를 정제하는 동안 고정된 상태로 유지한다. Prefix cache는 block이 커밋된 후에만 이동한다.

또한 저자들은 block diffusion model에서 사용되는 predict-and-noise 샘플러를 고려하였다. 이 샘플러는 깨끗한 block을 예측한 다음 noise schedule에 따라 신뢰도가 낮은 위치를 다시 마스킹한다. 구체적으로, 각 step에서 denoiser는 현재 마스킹된 모든 토큰을 병렬로 예측하고, threshold $\gamma$보다 높은 예측값은 커밋하며, 나머지 불확실한 위치는 나중에 정제하기 위해 마스킹된 상태로 유지한다. 이를 통해 커밋되는 토큰의 수는 모델의 신뢰도에 따라 조정되면서 block이 $T$ step 내에 완료되도록 한다.

Experiments

• 학습 디테일
  • 2단계 학습
    • 1단계: Nemotron-3-Nano 1단계 혼합 데이터의 일부를 사용하여 denoiser를 학습
    • 2단계: 1단계에서 학습된 denoiser를 계속 사용하면서, Nemotron-3-Nano 2단계 혼합 데이터로 전환
  • 총 약 14억 개의 토큰으로 학습
  • BF16 precision
  • optimizer: AdamW
  • learning rate: Warmup-Stable-Decay (peak $1 \times 10^{-4}$, final $1 \times 10^{-6}$)
  • block 크기 $S = 16$

1. Quality–Throughput Trade-off

다음은 카테고리별로 baseline model인 Nemotron-3-Nano-30B-A3B와 비교한 결과이다.

2. Design Choices

다음은 주요 디자인 선택에 대한 ablation 결과이다.

다음은 사전 학습된 AR 표현이 고정된 상태로 유지되어야 하는지 여부를 테스트한 결과이다.

3. Training Block Size

다음은 학습 block 크기에 따른 성능을 비교한 결과이다.

4. Sampling Block Size

다음은 $S=16$으로 학습된 2단계 체크포인트에 대하여, 샘플링 block 크기에 따른 성능을 비교한 결과이다.

5. Sampling Dynamics

다음은 (a) 각 diffusion step에서 완료되는 답변 생성 block의 비율과 (b) 위치 및 diffusion step별 토큰 확정 분포이다. 일반적으로 초기 위치의 토큰이 더 빨리 확정되어 autoregressive한 패턴을 나타낸다.

다음은 diffusion step마다 평균 커밋된 토큰 수를 나타낸 그래프이다.