NVIDIA's Nemotron-TwoTower Hits 2.42x Faster Generation Without Retraining

NVIDIA splits a 30B model into two specialized towers to generate text 2.42x faster while keeping 98.7% of the original model's quality

·
·
NVIDIA's Nemotron-TwoTower Hits 2.42x Faster Generation Without Retraining
  • NVIDIA released Nemotron-Labs-TwoTower, a diffusion LLM that generates tokens in parallel blocks instead of one at a time.
  • It achieves 2.42x faster generation while retaining 98.7% of the original model's benchmark quality at the default operating point.
  • The architecture splits a single 30B pretrained model into two towers: a frozen context tower and a trainable denoiser tower, trained on only ~2.1T tokens.
  • Code and math see the largest quality drops; commonsense and multilingual tasks are nearly unchanged or improved.
  • Requires 2x H100/A100 80GB GPUs (~59GB per GPU) for full diffusion mode; single GPU supported for AR-only mode.
  • Available now on Hugging Face under NVIDIA's open model license for commercial use, with the paper also released.

Autoregressive language models have a fundamental speed problem: they generate one token at a time, sequentially, no matter how much compute you throw at them. NVIDIA Research just released Nemotron-Labs-TwoTower, a new architecture that sidesteps this bottleneck by splitting a single pretrained model into two specialized copies that work in tandem, achieving 2.42x faster generation while retaining 98.7% of the original model's benchmark quality.

The core problem with diffusion language models

Diffusion models for text are a promising alternative to autoregressive generation. Instead of writing tokens left-to-right one at a time, they start with a fully masked sequence and iteratively "denoise" it, predicting multiple tokens in parallel. The catch is that existing approaches use a single network for both context representation and iterative denoising, forcing one model to serve both roles and limiting its capacity for either.

Think of it like asking one person to both hold the conversation history in their head and simultaneously write the next paragraph. Both tasks compete for the same mental bandwidth. NVIDIA's insight is to just use two people.

Two towers, one pretrained checkpoint

TwoTower decouples these roles into two towers: a frozen AR context tower that causally processes clean tokens, and a trainable diffusion denoiser tower with bidirectional block attention that refines noisy blocks via cross-attention to the context. Both towers are initialized from the same pretrained Nemotron-3-Nano-30B-A3B checkpoint, a 30B hybrid model that interleaves Mamba-2 state-space layers, standard attention, and mixture-of-experts (MoE) layers.

The key engineering decision: only the denoiser tower is trained. The context tower stays completely frozen, preserving all the knowledge baked in during the original 25T-token pretraining. The denoiser is then adapted via a masked diffusion objective on roughly 2.1T tokens, a fraction of the original training budget.

Here is how generation actually works at inference time:

  1. The context tower encodes the prompt causally, producing per-layer KV caches and Mamba state vectors.
  2. The denoiser receives a block of 16 masked tokens ([MASK], [MASK], ...).
  3. Over multiple denoising steps, it predicts all masked positions in parallel, using bidirectional attention within the block and cross-attending to the context tower layer-by-layer.
  4. Tokens predicted with confidence above a threshold (default γ=0.8) are "committed" and locked in. The rest stay masked for the next step.
  5. Once the block is complete, the context tower updates its cache with the new tokens, and the process repeats for the next block.

This confidence-based unmasking is what drives the speedup. Early in each block's denoising, the model is highly confident about most tokens and commits them all at once. Only the ambiguous positions require additional refinement steps. The result is that the model commits far more than one token per compute step, which is the core advantage over autoregressive decoding.

Bar chart comparing accuracy and throughput between Nemotron-3-Nano-30B-A3B baseline and Nemotron-TwoTower across benchmark categories

What the numbers actually look like

Built on Nemotron-3-Nano-30B-A3B and trained on approximately 2.1T tokens, Nemotron-TwoTower retains 98.7% of the autoregressive baseline's quality while offering 2.42x higher wall-clock generation throughput. The benchmark breakdown tells a more nuanced story:

  • General knowledge (MMLU): 78.56 → 78.24, essentially flat
  • Commonsense (ARC-Challenge): 91.72 → 92.66, actually improved
  • Code (HumanEval): 79.27 → 75.58, modest drop
  • Math (MATH-500): 84.40 → 80.60, the largest degradation
  • Multilingual (MMLU Global Lite): 73.97 → 73.94, nearly identical

Code and math show modest degradation, while commonsense and multilingual scores are recovered or slightly improved. The throughput dial is also adjustable: lowering the confidence threshold commits more tokens per step and pushes throughput beyond 3x, at the cost of more quality loss. Raising it recovers quality closer to the AR baseline but reduces speed.

The architecture innovations that make it work

Several specific design choices were validated through ablations in the accompanying paper:

  • Layer-aligned cross-attention: The denoiser at layer i attends to the context tower's layer i KV cache. Because both towers start from the same checkpoint, same-index layers operate at comparable representation levels, making this a natural pairing. Prior approaches only broadcast the final hidden state.
  • Bidirectional in-block attention: Within the current noisy block, tokens attend to each other in both directions. Across past committed blocks, attention remains causal. This adds zero parameters.
  • Time conditioning via adaLN-single: A small MLP (~1.5M parameters) maps the current diffusion timestep to per-layer scale, shift, and gate values, telling the denoiser how noisy the current block is. This improved generation, code, and math scores significantly in ablations.
  • Causal Mamba (not bidirectional): The team tested running Mamba-2 layers in both directions and averaging the outputs. It barely helped quality and doubled the SSM compute cost, so they kept Mamba causal and relied on bidirectional attention for denoising context.

The real breakthrough: converting any AR model into a diffusion model

The deeper implication here is methodological. Diffusion language models have emerged as a promising paradigm that enables parallel generation, but their learning efficiency lags behind autoregressive models when trained from scratch. TwoTower studies AR-to-diffusion conversion to transform pretrained AR models into efficient diffusion models that excel in speed while preserving task accuracy.

The TwoTower approach sidesteps the expensive from-scratch training problem entirely. You take an existing, well-trained AR model, freeze it, clone it, and train only the clone as a denoiser. The frozen copy acts as a rich, stable representation backbone that the denoiser can query at every layer. The paper explicitly frames this as a general technique applicable to any pretrained autoregressive model, not just Nemotron.

The ablations also reveal something important about architecture choices: tying the two towers together under a joint AR+diffusion loss was substantially worse than keeping them decoupled. Sharing weights forced a single set of parameters to serve both the causal context role and the bidirectional denoising role simultaneously, which is exactly the problem TwoTower was designed to avoid.

Hardware requirements and how to run it

Full two-tower diffusion uses 2 GPUs, about 59GB per GPU in BF16. AR-only mode runs on a single 80GB GPU. The released checkpoint ships both towers together (~60B total parameters, ~3B active per token per tower due to the MoE routing). The model is available now on Hugging Face under the NVIDIA Nemotron Open Model License and is ready for commercial use.

You can load and run diffusion generation with standard Transformers (requires trust_remote_code=True):

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_name = "nvidia/Nemotron-Labs-TwoTower-30B-A3B-Base-BF16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.bfloat16, trust_remote_code=True
)
# Place each tower on its own GPU
model.place_towers_on_devices("cuda:0", "cuda:1")
model.eval()
prompt = "Explain the theory of relativity:"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
# Block-wise diffusion generation
outputs = model.generate_mask_diffusion(
    inputs["input_ids"],
    max_new_tokens=256,
    block_size=16,             # tokens generated per block
    steps_per_block=16,        # denoising iterations per block
    confidence_threshold=0.8,  # commit tokens above this confidence
    temperature=0.1,
    eos_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

The model also supports generate_mock_ar() for two-tower autoregressive decoding (one token per step, useful for debugging) and generate_ar() for single-GPU standard AR generation using only the context tower. The checkpoint supports all three modes without any changes to the weights.

Where this fits in the diffusion LLM landscape

The diffusion LLM space has been heating up rapidly. LLaDA scaled masked diffusion to 8B parameters, demonstrating competitive downstream performance with LLaMA3. Gemini Diffusion achieved commercial-grade generation with 5x speed improvements. Mercury Coder from Inception Labs showed commercial viability for code generation. What sets TwoTower apart is the explicit separation of concerns between context and denoising, and the demonstration that this works at 30B scale on a hybrid Mamba-MoE architecture, not just on standard dense transformers.

The practical implication for teams running inference at scale: if you have a large AR model and need higher throughput without retraining from scratch, the TwoTower recipe offers a credible path. The cost is doubling your GPU memory footprint for the model weights, but the sequence-length-dependent KV cache scales the same as the original AR model. For long-context generation workloads where throughput is the bottleneck, the 2.42x speedup at near-identical quality is a compelling trade-off.

Comments

avatar