Nemotron-Labs-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

TL;DR: NVIDIA takes an off-the-shelf 30B autoregressive model (a hybrid Mamba/attention/MoE), freezes a copy of it as a pure "context reader," and trains a second copy as a block-wise diffusion denoiser that cross-attends into the frozen copy layer by layer. The result keeps 98.7% of the AR baseline's benchmark quality while generating text 2.42× faster in wall-clock terms — and the ablations make an unusually clean case that the decoupling itself is what makes it work: sharing weights between the two roles craters quality by ~27 points. Weights and code are released.

The problem: diffusion LMs make one network do two jobs

Autoregressive decoding is a serial bottleneck: one forward pass, one token. Diffusion language models promise a way out — start a whole block of tokens as [MASK], predict all of them in parallel, keep the confident ones, and iterate. If a 16-token block resolves in 3–4 denoising steps instead of 16 sequential passes, you win.

But there's a structural tension buried in how existing diffusion LMs are built. At every denoising step, the network is doing two very different things with the same weights:

1. Representing clean context — the prompt and already-committed tokens, which is exactly what an autoregressive LM was pretrained to do.
2. Denoising corrupted tokens — reasoning bidirectionally about a partially-masked block, which no AR pretraining ever taught it.

Training one network to do both pulls the weights in opposing directions. Worse, if you adapt a pretrained AR model into a diffusion model by fine-tuning it on masked inputs, you're actively degrading the very context representations you wanted to inherit. This is the paper's diagnosis, and the ablations later show it's not just a story — joint training with tied weights loses ~26–28% relative accuracy under both decoding modes.

Arriola et al. (2025) noticed the two roles could be split into an encoder-decoder at 1.7B scale, but with tied weights. The open question was whether full decoupling — a genuinely frozen context model plus a genuinely separate denoiser — scales, and whether it works on the messy modern hybrids (Mamba-2 + attention + MoE) that are increasingly the backbones people actually deploy.

The idea: freeze the reader, train the writer

Take two copies of the pretrained Nemotron-3-Nano-30B-A3B (52 layers: 23 Mamba-2, 6 attention, 23 MoE).

Tower 1 — the context tower — is frozen. It runs exactly as it was pretrained: causally, over clean tokens only (prompt plus committed blocks). It never sees a [MASK] token in its life. Its job is to produce per-layer KV caches (for attention layers) and boundary states (for Mamba layers). Because it's frozen and only ever processes clean text, the pretrained representations are preserved perfectly, by construction.

Tower 2 — the denoiser — is trained. It takes a block of S=16 noisy tokens and refines them over up to T steps. Within the block, attention is bidirectional (masked tokens can look at each other). For context, it doesn't cross-attend to a single summary vector — each denoiser layer i attends to the context tower's layer-i KV cache:

where \mathbf{Q}_b^{(i)} are queries from the noisy block b at layer i, and the keys/values concatenate the context tower's cache for past blocks with the denoiser's own KV for the current block. Since both towers start from the same checkpoint, layer i in one tower speaks the same representational "dialect" as layer i in the other — layer-aligned cross-attention is essentially free residual wiring between two copies of the same brain. Mamba layers get the analogous treatment: the denoiser's Mamba-2 layers seed their initial recurrent state from the context tower's state at the block boundary.

Once every token in a block is confidently resolved, the block is committed, the frozen context tower ingests it causally (one clean forward chunk, updating its caches), and generation moves to the next block.

Two small but useful additions: the denoiser is conditioned on the noise level via adaLN-single (the DiT-style timestep modulation from image diffusion — a global MLP producing scale/shift/gate, plus per-layer embeddings, only 1.5M parameters on a 30B model), and the MoE routing is left entirely alone; experts learn noise-specialization on their own through the adaLN-modulated inputs.

One design choice worth noting because the ablation goes against intuition: they tried making Mamba bidirectional (run the pretrained scan left-to-right and right-to-left, average outputs). It doubled SSM compute and hurt code and math slightly (68.64→68.05 and 80.57→79.78 on their internal eval). Bidirectional attention within the block was enough; Mamba stays causal.

Training and sampling

Training is standard masked diffusion applied per block. The forward process masks each token in block \mathbf{x}_b independently with probability 1-\alpha_t under a linear schedule \alpha_t = 1-t. The loss is mean negative log-likelihood over masked positions, conditioned on the noisy block, the timestep, and the frozen context caches \mathbf{c}_{<b} — they drop the theoretical 1/t ELBO weight for stability. Only the denoiser gets gradients; the context tower runs once per sequence under no-grad. Training reuses the backbone's own two-phase data curriculum (broad coverage, then STEM-heavy) for ~2.1T tokens total — a fraction of the backbone's 25T-token pretraining, but not a small run either.

Sampling uses confidence unmasking: at each step, predict all masked positions in parallel, commit every prediction above a confidence threshold \gamma = 0.8, leave the rest masked for the next step. The number of tokens committed per step is adaptive — easy stretches of text resolve in one shot, hard tokens get more refinement.

The evidence

Headline: 98.7% of the AR baseline's aggregate benchmark quality (base models, before any instruction tuning), 2.42× wall-clock throughput on 2×H100 in BF16. General knowledge stays within ~1 point, code and math degrade modestly, commonsense and multilingual are recovered or slightly improved. Tuning \gamma traces a Pareto curve: throughput beyond 3× is reachable at larger quality cost.

The ablation that carries the paper is tower decoupling (Table 2, after only ~167B adaptation tokens):

This is the "aha" in numbers. Forcing one set of weights to serve both roles under a joint AR+diffusion loss loses ~21–28% relative accuracy in both decoding modes — the diffusion objective doesn't just fail to help the AR mode, it destroys it, and vice versa. Even plain continued AR training on the adaptation blend degrades things (probably data-distribution drift), which makes freezing the context tower look even better: it's not just cheaper, it's the highest-quality option they tested.

Block size controls the parallelism/quality trade-off in the way you'd expect, with a clean monotone trend:

More interesting is the asymmetry in sampling block size (Table 4). Fix the S{=}16-trained checkpoint and only change the sampling block: shrinking to S{=}8 is harmless or slightly helpful, but growing to S{=}64 is catastrophic on generative tasks — HumanEval falls 76.40 → 19.85, GSM8K and MATH-500 collapse to 2.20. The denoiser learned to resolve dependencies within a 16-token horizon; ask it to fantasize 64 tokens from a fixed prefix and it falls apart, while MMLU (short answers) barely notices (78.32 → 78.10). If you deploy one of these models, this is the sharp edge to know about.

Where does the speedup actually come from? Section 3.6 is the most illuminating part of the paper and I'd point you there first. Tracing generation on 100 problems per task: the first diffusion step commits the most tokens by far, then commitments drop off quickly as the sampler chews on the residual low-confidence positions. Most blocks finish within the first few steps; harder tasks (MMLU-Pro, code) have longer tails. So the model behaves like an adaptive-width decoder — sprinting through predictable text, slowing down on the tokens that matter.

And there's a lovely emergent finding: within each block, commitments follow a strongly left-to-right order — the position×step commitment heatmap is upper-left triangular. Nobody imposed this. The authors' hypothesis: the backbone has 23 causal Mamba layers and only 6 attention layers, and the Mamba layers stay causal even in the denoiser, so the model inherits a dominant left-to-right inductive bias. The "diffusion" model has, in effect, rediscovered something close to autoregression — but with adaptive parallelism where the text is easy. Whether that bias is a feature (it clearly doesn't hurt quality) or a ceiling on what block diffusion buys over smarter speculative decoding is an open question the paper doesn't engage with.

What this changes, and what to be skeptical of

The constructive takeaway is a recipe: you can convert any strong pretrained AR model into a fast block-diffusion generator without touching the original weights. The context tower is the unmodified checkpoint — it can still do AR scoring, likelihood evaluation, or serve as a verifier for speculative decoding (the paper explicitly notes the LM head is kept for those uses). Layer-aligned cross-attention between two copies of the same checkpoint is a cheap, natural interface that others will likely reuse; it's essentially the same trick that makes ControlNet work for image diffusion, applied to language. And this is the largest and most architecturally realistic diffusion-LM adaptation released with open weights — 30B hybrid Mamba/MoE, not a toy dense transformer.

Now the caveats, and there are real ones:

• The throughput number needs scrutiny. 2.42× is against one-token-at-a-time AR decoding on 2×H100 in BF16, measured as wall-clock to final answer. There is no comparison against speculative decoding, EAGLE-style drafting, or multi-token prediction — techniques that routinely deliver 2–3× on the same baseline with zero quality loss. The paper never claims to beat them, but the headline invites the comparison and doesn't make it. Also unspecified: batch size and serving regime. Diffusion decoding's parallelism advantage shrinks under high-batch serving where AR decoding is already compute-bound.
• Doubled weight memory. Two full copies of a 30B model resident at inference. The prefix cache stays single (only the context tower keeps sequence-length-dependent state, so cache memory scales like AR), but the fixed footprint doubles. For an MoE with 3B active parameters this may be tolerable; it's still a real deployment cost.
• 98.7% is an aggregate; code and math — the tasks people most want fast generation for — degrade modestly, and the exact per-benchmark numbers for the released checkpoint live in a figure rather than a table in the extracted text.
• Base models only. Everything is evaluated pre-instruction-tuning and pre-RL. Whether the denoiser survives post-training, and whether the quality gap widens on long chain-of-thought reasoning (where the sampling-block-size fragility hints at trouble), is deferred to future releases.
• 2.1T adaptation tokens is "a fraction of 25T," but it's still a serious training run — this is not a weekend fine-tune, and there's no data on how quality scales with adaptation budget beyond the 167B-token ablation.

If the result holds up under honest serving benchmarks and post-training, the framing shift is the durable contribution: diffusion stops being a rival pretraining paradigm competing with AR from scratch, and becomes a decoding adaptation you bolt onto the AR models you already have — inheriting their quality by freezing it in place, and spending new capacity only on the one thing AR pretraining never taught: filling in the blanks in parallel.

Where to spend your time in the paper: Section 3.6 and Appendix A (the sampling dynamics — genuinely informative about what block diffusion is actually doing), then Table 2 (the decoupling ablation that justifies the whole design), then Table 4 (the sampling-block-size cliff, which is the main practical gotcha).