Nemotron-TwoTower: Autoregressive Diffusion Model

• The paper introduces Nemotron-TwoTower, a method that decouples autoregressive context representation from diffusion denoising by using a frozen context tower and a trainable denoiser tower.
• The model retains 98.7% of its autoregressive baseline’s quality while achieving 2.42× faster generation through block-wise iterative refinement.
• The architecture leverages a layer-aligned cross-attention mechanism and confidence unmasking to progressively commit tokens within blocks.

Nemotron-TwoTower is a block-wise autoregressive diffusion LLM that separates causal context representation from iterative denoising by using two distinct towers: a frozen autoregressive context tower and a trainable diffusion denoiser tower (Reda et al., 25 Jun 2026). It is built on Nemotron-3-Nano-30B-A3B, an open-weight 30B hybrid Mamba-Transformer MoE model, and is presented as a method for preserving the language understanding and generation quality of a large pretrained autoregressive model while replacing one-token-at-a-time decoding with block-wise diffusion-style generation. The reported result is a system that retains 98.7% of the autoregressive baseline’s aggregate benchmark quality while delivering 2.42X higher wall-clock generation throughput (Reda et al., 25 Jun 2026).

1. Conceptual basis and problem setting

Nemotron-TwoTower is motivated by a specific limitation in standard autoregressive language modeling: inference throughput. In a conventional autoregressive model, every new token depends on the previous one, so decoding is inherently sequential. Diffusion LLMs offer an alternative because they can refine many tokens in parallel, but the paper argues that prior approaches commonly use a single network both to represent clean context and to denoise corrupted tokens. That arrangement is described as an entanglement of functions, in which one parameter set is required to support both causal context modeling and bidirectional denoising (Reda et al., 25 Jun 2026).

The central architectural claim is that these two roles should be decoupled. Nemotron-TwoTower therefore assigns clean left-to-right prefix processing to a frozen autoregressive tower and delegates noisy-block refinement to a separate trainable denoiser. This preserves the pretrained autoregressive representation rather than forcing it to adapt jointly to the diffusion objective. A common misconception is that the model ceases to be autoregressive; in fact, it remains autoregressive at the block level. The change is not from autoregression to fully non-causal generation, but from token-level sequential decoding to iterative refinement within causally ordered blocks.

This framing places Nemotron-TwoTower between standard autoregressive LLMs and earlier masked diffusion LLMs. Relative to autoregressive systems, it changes the generation algorithm. Relative to earlier diffusion systems, its key innovation is not merely masked denoising, but the architectural separation between clean-context encoding and noisy-block denoising.

2. Two-tower architecture and pretrained backbone

The architecture is instantiated by duplicating a pretrained Nemotron-3-Nano-30B-A3B checkpoint, which is described as a 30B hybrid model with 52 layers composed of 23 Mamba-2 layers, 6 attention layers, and 23 MoE layers (Reda et al., 25 Jun 2026). The two copies are then assigned asymmetric functions.

The first copy becomes the frozen autoregressive context tower. It processes the prompt and any previously committed tokens causally and serves only as a representation model for clean left-to-right context. In practice, it produces per-layer key/value caches for attention and recurrent Mamba states summarizing the committed prefix. The second copy becomes the trainable diffusion denoiser tower. It is initialized from the same checkpoint but adapted for diffusion-style refinement of noisy blocks. The towers are therefore not symmetric replicas used in the same manner; one is fixed to preserve the original autoregressive inductive structure, while the other is optimized for the new block-refinement procedure.

A distinctive feature is the layer-aligned cross-attention interface. At denoiser layer $i$, queries from the noisy block attend jointly to past clean context from the context tower and to the denoiser’s own representation of the current block:

$$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$

Because both towers start from the same pretrained checkpoint, layer $i$ in each tower operates at a comparable representational level, which the paper identifies as the basis for natural alignment (Reda et al., 25 Jun 2026). Within the current noisy block, attention is bidirectional, allowing masked tokens to condition on one another while still remaining causal with respect to prior committed blocks.

The handling of Mamba-2 is also specific. The final design keeps Mamba causal in the denoiser rather than making it bidirectional, because bidirectional Mamba is reported to roughly double SSM compute while providing only marginal gains. Instead, bidirectional interaction is supplied by in-block attention, while the causal Mamba pathway is seeded from the context tower’s Mamba state. Time conditioning is injected using adaLN-single: a small global MLP maps the timestep $t$ to scale, shift, and gate parameters, with per-layer learned embeddings for modulation. The paper states that this adds only about 1.5M parameters on the 30B model and improves quality. MoE routing is retained from the backbone, and tokens are not manually rerouted; expert specialization is left to emerge under the denoising objective.

3. Block-autoregressive diffusion formulation

Nemotron-TwoTower partitions a sequence into blocks of size $S$ and generates one block at a time (Reda et al., 25 Jun 2026). The overall likelihood is factorized block-autoregressively as

$$\log p_\theta(\mathbf{x}) = \sum_{b=1}^{B} \log p_\theta(\mathbf{x}_b \mid \mathbf{x}_{<b}),$$

where $\mathbf{x}_b$ denotes the $b$-th block and $\mathbf{x}_{<b}$ the clean prefix of earlier blocks. This preserves left-to-right causal structure at the level of blocks rather than individual tokens.

Inside each block, the model uses a masked discrete diffusion process. For a clean block $\mathbf{x}_b=(x_b^1,\dots,x_b^S)$, the forward corruption process replaces tokens with $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$0 according to a noise level $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$1. The per-token marginal is

$$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$2

Here $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$3 is the noisy token at position $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$4 in block $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$5, $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$6 is a point mass on the clean token, and $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$7 is a point mass on $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$8. The schedule is linear, $$\text{Attn}\big(\mathbf{Q}_b^{(i)},\; [\mathbf{K}_{<b}^{\text{ctx},(i)};\, \mathbf{K}_b^{\text{den},(i)}],\; [\mathbf{V}_{<b}^{\text{ctx},(i)};\, \mathbf{V}_b^{\text{den},(i)}]\big).$$9, so small $i$0 corresponds to light corruption and large $i$1 to strong masking.

The training objective is a simplified masked-diffusion loss:

$i$2

This objective samples a diffusion time $i$3, corrupts the block, and trains the denoiser to recover the original token at every masked position, conditioned on the noisy block and on the context-tower state $i$4 for prior committed blocks. The paper notes that the theoretically correct diffusion weighting simplifies to $i$5 under the linear schedule, but that importance weight is omitted for stability; instead, the negative log-likelihood is averaged over masked positions (Reda et al., 25 Jun 2026).

The resulting formulation has an important interpretive consequence. The model does not attempt to replace causal language modeling with a purely denoising-based objective over full sequences. Rather, it embeds denoising inside a left-to-right block factorization, preserving a causal outer structure while introducing parallel iterative refinement inside each block.

4. Training procedure and optimization regime

Training is applied only to the denoiser tower, while the context tower remains frozen as a pure autoregressive representation model (Reda et al., 25 Jun 2026). This design operationalizes the paper’s central hypothesis: context representation and diffusion denoising should not be forced into a shared parameterization once a strong autoregressive backbone already exists.

The base checkpoint is Nemotron-3-Nano-30B-A3B, and the training corpus totals roughly 2.1T tokens, which the paper explicitly contrasts with the 25T tokens used to pretrain the original backbone. Training follows a two-stage curriculum aligned with the backbone’s data phases. Phase 1 uses a broad, diverse blend; phase 2 shifts toward higher-quality and STEM-heavy data. The denoiser is trained in BF16 with AdamW and a Warmup-Stable-Decay learning rate schedule, peaking at $i$6 and ending at $i$7.

The default block size for the final model is $i$8. Some training stages use $i$9 for efficiency and then adapt to $t$0 for final sampling. This detail is operationally significant because the paper later reports that sampling with a block size much larger than the training block size can harm performance, especially on generation-heavy tasks such as HumanEval and math benchmarks. A plausible implication is that the denoiser learns refinement dynamics tied to a particular locality and masking geometry, so aggressive extrapolation in block size creates a train–test mismatch.

The training setup also clarifies what is and is not updated. The model is not jointly optimized as a dual-tower system with both towers adapting to the diffusion loss. Nor is it trained by continuing standard autoregressive optimization on the backbone. The paper identifies the frozen-context / trainable-denoiser split as the superior regime in ablation.

5. Inference algorithm and sampling behavior

Inference uses a two-level loop (Reda et al., 25 Jun 2026). The outer loop is block autoregression: generate one block, commit it, update the frozen context tower, then move to the next block. The inner loop performs diffusion refinement within the current block. A new block is initialized as all $t$1 tokens and refined over $t$2 denoising steps using the denoiser together with the frozen context cache.

The default sampler is a confidence-unmasking variant. At each refinement step, the denoiser predicts all currently masked tokens in parallel, commits the high-confidence ones above threshold $t$3, and leaves uncertain positions masked for later steps. The default quality-throughput operating point uses $t$4 and block size $t$5. This differs from predict-and-noise samplers discussed in earlier block-diffusion work because the number of tokens committed per step is adaptive rather than fixed.

The reported sampling dynamics are front-loaded. Many blocks finish in the first few diffusion steps, and the model often commits many tokens immediately before spending later steps resolving harder residual positions. Within a block, earlier positions are typically committed earlier than later positions, producing an upper-left triangular pattern that resembles partial autoregressive behavior. The paper suggests that this may reflect the strong left-to-right inductive bias inherited from the causal backbone. That interpretation is explicitly presented as a suggestion rather than a formal proof of mechanism.

An important clarification is that the context tower’s LM head is optional. It is not required for the default diffusion path, but it can be used for speculative decoding, likelihood evaluation, or autoregressive scoring. This means the architecture preserves a path back to conventional autoregressive utilities even though its default generation procedure is diffusion-based.

6. Empirical results, ablations, and practical constraints

The principal empirical claim is that Nemotron-TwoTower retains 98.7% of the autoregressive baseline’s aggregate benchmark quality while achieving 2.42× wall-clock generation throughput relative to one-token-at-a-time decoding (Reda et al., 25 Jun 2026). The evaluation spans general knowledge, code, math, commonsense, and multilingual tasks. The paper attributes the throughput gain to the ability to commit multiple tokens per refinement step, especially early in decoding.

The ablations support several design choices. Bidirectional Mamba yields almost no quality benefit while increasing compute, so the final model keeps Mamba causal. Time conditioning improves generation, code, and math. Continuing training on the phase-2 blend improves results further. Most importantly, the tower-decoupling ablation indicates that keeping the context tower frozen and training a separate denoiser is better than continuing autoregressive training on the backbone or tying the towers under a joint loss. This is the direct empirical support for the paper’s main architectural thesis.

The reported trade-offs are not trivial. Inference has a fixed memory cost for keeping both towers’ weights resident, although sequence-length-dependent cache memory remains similar to an autoregressive model because only one prefix cache is maintained. Quality and throughput depend jointly on block size and confidence threshold. Larger blocks and lower thresholds increase throughput but can degrade quality; smaller blocks improve quality but reduce speed gains because the context tower must be updated more often. The paper also shows that sampling with a block size far larger than the training block size can degrade results severely, particularly on generation-heavy tasks.

These constraints address another common misunderstanding: the 2.42× throughput improvement is not an unconditional property of diffusion-style decoding. It is tied to a particular operating point, namely confidence unmasking with $t$6 and $t$7, and depends on the interaction among architecture, sampler, and block-size choice.

7. Position within diffusion language modeling

Nemotron-TwoTower is situated as a diffusion-language-model adaptation of a large pretrained autoregressive backbone rather than as a de novo diffusion model (Reda et al., 25 Jun 2026). Compared with earlier masked diffusion LLMs, the paper emphasizes that prior systems often used one network both to encode context and to denoise masked tokens. TwoTower instead freezes a dedicated autoregressive context pathway and allows the denoiser to cross-attend to layer-aligned caches rather than only a final hidden state. Compared with block diffusion and encoder-decoder diffusion approaches, this is presented as a stronger form of architectural decoupling.

Relative to standard autoregressive LLMs, the model preserves much of the pretrained capability but changes the decoding algorithm from strict token-wise generation to iterative block refinement. The paper’s main conclusion is therefore architectural rather than purely algorithmic: a large pretrained autoregressive model can be repurposed into a higher-throughput diffusion generator without substantial quality loss if causal context representation is preserved in a frozen tower and denoising is delegated to a separate trainable tower.

In summary, Nemotron-TwoTower consists of one frozen tower that maintains clean causal context, one trainable tower that denoises masked blocks using bidirectional in-block attention and layer-aligned cross-attention to the clean prefix state, and a confidence-based sampler that commits multiple tokens at once (Reda et al., 25 Jun 2026). This combination is the basis for the reported result that generation can be substantially accelerated while maintaining almost all of the autoregressive baseline’s benchmark quality.