💬 NLP

Nemotron-TwoTower: Diffusion Language Modeling with Pretrained Autoregressive Context

Nemotron-TwoTower introduces a novel block-wise autoregressive diffusion architecture that decouples context processing and denoising into two specialized towers, achieving 2.42x faster generation throughput while retaining 98.7% of the quality of its autoregressive baseline.

Original authors: Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

Published 2026-06-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Fitsum Reda, John Kamalu, Roger Waleffe, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to write a long story, but you have a strict rule: you can only write one word at a time, and you must wait for the previous word to be fully finished before you can start the next one. This is how most current AI language models work (called "Autoregressive" models). It's like a very careful, slow typist who never makes a mistake but takes forever to finish a sentence.

Now, imagine a new method where you can type a whole sentence at once, look at it, and then fix the mistakes in all the words simultaneously. This is faster, but usually, the quality of the story suffers because the AI gets confused trying to fix everything at once.

Nemotron-TwoTower is a new invention that tries to get the best of both worlds: the speed of typing everything at once, with the high quality of the careful, one-word-at-a-time typist.

Here is how it works, using a simple analogy:

The Two Towers: The Librarian and The Editor

The researchers built a system with two distinct "towers" (parts of the AI) that work together, like a team of two people:

The Frozen Context Tower (The Librarian):
- Role: This is the "expert" who knows the story so far. It is a pre-trained AI that is frozen (locked in place). It acts like a strict librarian who reads the story word-by-word from the beginning to the end.
- Job: It doesn't change. It just keeps a perfect, organized record of everything written so far. It tells the other tower, "Here is the context of the story up to this point."
- Why it matters: Because it's frozen, it keeps all the original intelligence and knowledge of the massive 30-billion-parameter model it was built from. It never forgets or gets confused.
The Trainable Diffusion Denoiser Tower (The Editor):
- Role: This is the "fast worker." It is a trainable AI that is allowed to change and learn.
- Job: Instead of writing one word at a time, it takes a block of 16 words (or tokens) that are currently blank (masked). It looks at the "Librarian's" notes and tries to fill in all 16 blanks at once.
- The Process: It doesn't get it perfect immediately. It makes a guess, then looks at its own guess and the Librarian's notes, and refines the words. It does this repeatedly (iteratively) until the block of words is clean and makes sense.
- The Magic: Because it can look at all 16 words at the same time and fix them together, it is much faster than writing one by one.

How They Work Together (The Assembly Line)

Imagine a factory assembly line:

The Librarian reads the prompt and the words already written, keeping a perfect mental map of the story.
The Editor grabs a blank chunk of paper (16 words long).
The Editor asks the Librarian, "What comes next?" and the Librarian provides the context.
The Editor fills in the 16 blanks, checks its work, and fixes any errors.
Once the Editor is confident the 16 words are good, they are "committed" (finalized).
The Librarian updates its mental map to include these new 16 words.
The Editor grabs the next blank chunk of 16 words, and the cycle repeats.

The Results: Speed vs. Quality

The paper claims this system achieves a "sweet spot":

Speed: It is 2.42 times faster than the traditional one-word-at-a-time method. It's like switching from a snail to a sprinter.
Quality: It retains 98.7% of the original model's quality. It didn't lose much intelligence to gain that speed.
Efficiency: They trained this new "Editor" tower on about 2.1 trillion words, which is a tiny fraction of the 25 trillion words used to train the original "Librarian." This means they didn't have to re-teach the whole AI from scratch; they just taught the Editor how to use the Librarian's notes.

Why This is Different from Previous Attempts

Previous attempts tried to use one single brain to do both jobs: remembering the story and fixing the words at the same time. The paper argues this is like asking a person to be a perfect memory keeper and a fast editor simultaneously—it's too much work, and they end up doing neither job perfectly.

By splitting the roles into two towers (one frozen, one trainable), they let each part do what it's best at. The Librarian stays perfect, and the Editor gets fast.

The "Confidence" Trick

The system also has a smart feature called "confidence unmasking."

If the Editor is very confident about a word, it locks it in immediately.
If it's unsure, it leaves that word blank (masked) and tries to fix it in the next round of editing.
This means it doesn't waste time re-checking words it already knows are right, but it spends extra time polishing the tricky parts.

Summary

Nemotron-TwoTower is a way to make AI text generation much faster without making it "dumber." It does this by separating the job of "remembering the story" (which stays slow and perfect) from the job of "writing the next chunk" (which becomes fast and iterative). The result is a system that writes stories nearly as well as the old slow method, but more than twice as fast.

Technical Summary: Nemotron-TwoTower

Problem Statement

Autoregressive (AR) language models dominate text generation but suffer from a throughput bottleneck due to their sequential, one-token-at-a-time decoding process. Discrete diffusion language models offer a promising alternative by generating tokens in parallel and refining them iteratively, potentially yielding higher throughput and finer-grained controllability. However, existing diffusion approaches typically rely on a single network to perform two distinct roles simultaneously: representing clean context tokens and iteratively denoising corrupted tokens. This entanglement forces a single set of weights to serve conflicting objectives, limiting the model's capacity to excel at either role. Furthermore, it remains unclear whether fully decoupling these roles is feasible at large scales, particularly on modern hybrid architectures combining Mamba, attention, and Mixture-of-Experts (MoE) layers.

Methodology

The authors propose TwoTower, a block-wise autoregressive diffusion model that decouples the context representation and denoising roles into two separate towers, built upon the Nemotron-3-Nano-30B-A3B hybrid architecture (which interleaves Mamba-2, self-attention, and MoE layers).

Architecture

The model utilizes two copies of the pretrained network with complementary roles:

Frozen AR Context Tower: This tower remains frozen and acts as a causal autoregressive model over the prompt and previously committed tokens. It processes clean tokens to produce layer-aligned Key-Value (KV) caches and Mamba states. It preserves the backbone's original autoregressive capabilities and maintains reusable states across blocks.
Trainable Diffusion Denoiser Tower: This tower is trained to refine noisy blocks. It employs bidirectional block attention within the current noisy block (allowing tokens to attend to each other) while maintaining causal cross-attention to the corresponding layers of the frozen context tower. This layer-aligned cross-attention provides the denoiser with multi-scale access to the backbone's representations, contrasting with prior methods that only broadcast the final hidden state.

Generation Process

Generation proceeds block-by-block (Algorithm 1):

The context tower processes the prompt and committed tokens causally.
A new block is initialized with $S$ masked tokens.
The denoiser tower iteratively refines this block over $T$ steps. At each layer, the denoiser attends to the context tower's KV cache for past blocks and its own bidirectional KV cache for the current block.
Once a block is sufficiently refined (based on confidence thresholds), it is committed. The context tower then updates its states with the new clean tokens, and the process repeats for the next block.

Training and Conditioning

Masked Diffusion: The denoiser is trained using a masked diffusion objective. Clean blocks are corrupted by replacing tokens with [MASK] according to a linear schedule. The model minimizes the negative log-likelihood of the original tokens at masked positions.
Time Conditioning: The denoiser is conditioned on the diffusion timestep $t$ using adaptive layer normalization (adaLN), a technique common in image diffusion but less utilized in masked language models. This adds minimal parameters (~1.5M) but improves sampling quality.
Curriculum: Training follows a two-stage curriculum. Phase 1 uses a diverse data blend for broad coverage, while Phase 2 upweights high-quality and STEM-focused sources. The training block size is gradually reduced (from $S=32$ to $S=16$ ) to balance parallelism with conditioning quality.
Sampling: The authors employ a confidence-unmasking strategy (Algorithm 2). At each step, the denoiser predicts all masked tokens in parallel; tokens with confidence above a threshold $\gamma$ are committed, while uncertain tokens remain masked for further refinement. This makes the number of committed tokens adaptive.

Key Results

Evaluated on the Nemotron-3-Nano-30B-A3B backbone (30B parameters, hybrid MoE), Nemotron-TwoTower demonstrates the following:

Quality-Throughput Trade-off: The model retains 98.7% of the autoregressive baseline's aggregate benchmark quality while achieving a 2.42× improvement in wall-clock generation throughput.
Benchmark Performance:
- General Knowledge: Performance remains within ~1 point of the AR baseline.
- Code and Math: Show modest degradation compared to the baseline.
- Commonsense and Multilingual: Performance is recovered or slightly improved.
Design Ablations:
- Bidirectional Mamba: Testing a bidirectional Mamba variant showed marginal quality gains but doubled SSM compute; the final design keeps Mamba causal.
- Time Conditioning: Explicit time conditioning significantly improved generation, code, and math scores.
- Tower Decoupling: Keeping the context tower frozen and training a separate denoiser yielded the best results. Tying weights between towers or continuing AR training on the context tower led to significant accuracy drops.
Block Size Sensitivity: Smaller training block sizes ( $S=16$ ) improved quality over larger sizes ( $S=128$ ) but reduced throughput. The authors selected $S=16$ as the default to balance quality retention with speed. Sampling with blocks larger than the training block size resulted in severe performance degradation.

Significance and Claims

The paper claims that Nemotron-TwoTower successfully demonstrates that pretrained autoregressive representations can support a trainable diffusion denoiser without re-pretraining the entire model. By decoupling the context and denoising roles, the approach preserves the broad capabilities of the original 30B hybrid MoE model while unlocking the parallel generation benefits of diffusion.

The authors emphasize that this quality recovery is achieved through adaptation on a fraction of the data used for the original pretraining (~2.1T tokens vs. 25T tokens). The work establishes a new Pareto frontier for language models, offering a viable path to higher throughput generation without sacrificing the foundational quality of large-scale autoregressive models. The code and model weights are released to facilitate further research in this direction.