Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models

AI summary

Overview

Research area: Efficient small language model (SLM) architecture design and training, with a focus on real-device inference latency rather than parameter count alone.

Technical level: Advanced. The paper assumes familiarity with transformer architecture, linear attention and state-space operators (Mamba, Mamba2, DeltaNet, Gated DeltaNet, GLA), scaling laws, and neural architecture search.

Scope: The paper studies which architectural factors (depth–width ratio, operator selection) and training techniques determine real-device latency in SLMs, then combines them into a new model family called Nemotron-Flash.

What This Paper Is About

Most small language models are designed to minimize parameter count on the assumption that fewer parameters means proportionally faster inference. That assumption does not hold on hardware accelerators such as GPUs and TPUs, where deep-thin designs that are parameter-efficient can still be slow. This paper sets out to identify what actually governs real-device latency in SLMs and to derive general design principles and automated methods for building models that sit on the accuracy–latency frontier, rather than the accuracy–parameter frontier.

Key Contributions

1. A systematic study of depth–width ratios under latency constraints. The authors train Llama models at depths of 6, 12, 18, 24, and 30 blocks on 100B tokens from the Smollm-corpus, varying width, and show that the depth–width configuration that maximizes accuracy per parameter is not the one that maximizes accuracy per unit of latency.

2. An augmented scaling law that decouples depth and width. Existing scaling laws parameterize loss by model size and data size; this work reformulates the loss as L(D, W, N) = L0 + aD^−α + bW^−β + cN^−γ, allowing the sweet-spot depth–width ratio to be selected for a target latency budget.

3. An evolutionary search framework for hybrid operator combinations. Using DeltaNet, attention, and Mamba2 as candidate operators, the framework searches block types, operator ratios, FFN counts, and repetition counts across early, middle, and late architecture stages. Short-training perplexity serves as a search proxy, validated by an 88.8% Spearman correlation with full-training perplexity.

4. Two general training techniques. Weight normalization, which projects weight matrices onto a unit norm sphere after each training iteration, and 256 learnable meta tokens used for cache initialization. Both are reported as reusable add-ons across model families.

Main Findings

• Deep-thin models are not latency-optimal. Deeper models generally achieve a better accuracy–parameter trade-off over a wide depth range, although the benefit gradually saturates. For the accuracy–latency trade-off, however, there exists an optimal depth for a given latency budget: when the latency budget is 3 seconds, a depth of 12 achieved the best accuracy among the evaluated settings.

• The optimal depth–width ratio scales with the latency budget. The paper reports that the optimal ratio generally increases as the allowed latency grows, meaning depth and width should be selected deliberately per deployment constraint rather than defaulting to deep-thin designs.

• The augmented scaling law extrapolates. Fitting the reformulated law on a subset of depth/width settings and validating on larger unseen settings stayed within 5.3% of the ground-truth perplexity.

• DeltaNet and Gated DeltaNet lie on the PPL–latency Pareto frontier. Among the 500M models trained on 100B tokens with Mamba, Mamba2, GLA, DeltaNet, Gated DeltaNet, RWKV7, and sliding window attention with a 512 window, DeltaNet (Wiki PPL 23.87, CR accuracy 47.83%) and Gated DeltaNet (23.80, 47.96%) were the strongest pure linear-attention candidates, versus Mamba2 (24.36, 47.72%), attention (24.44, 48.02%), and GLA (25.16, 46.82%).

• Pairing with Mamba2 is more reliable than pairing with attention. DeltaNet combined with Mamba2 reached 23.37 PPL and 48.03% CR accuracy, and Gated DeltaNet with Mamba2 reached the same 23.37 PPL and 48.03% accuracy, both exceeding the corresponding pure models. Combinations with attention were less stable.

• Operator gaps narrow inside hybrid models. Although Gated DeltaNet outperformed DeltaNet in pure language modeling, their task performance became comparable when integrated with Mamba2, making the more efficient DeltaNet preferable in hybrid designs.

• Searched architectures beat latency-matched baselines. The latency-searched architecture (837M parameters, 17.71 s decoding latency, 20.70 Wiki PPL, 51.04% CR accuracy) outperformed SWA (616M, 23.33 PPL, 48.72%), GLA (862M, 22.67, 48.43%), DeltaNet (852M, 20.90, 50.38%), Gated DeltaNet (672M, 21.98, 49.99%), Mamba2 (601M, 23.14, 48.61%), and Mamba2 + FFN (889M, 21.43, 50.04%) at comparable latency.

• The search metric shapes the architecture found. Searching for decoding latency produced a model built from DeltaNet-FFN-Mamba2-FFN and Attention-FFN-Mamba2-FFN blocks stacked in an interleaved manner. Searching with a 500M parameter budget instead produced a deeper, more attention-heavy design that achieved over 1.21% higher CR accuracy and a reduction of more than 0.74 in PPL compared to all 500M baselines.

• Weight normalization improves final convergence. Averaged across Llama 1B, DeltaNet 1B, and Mamba2 1B trained on 100B tokens, it improved CR accuracy by +1.20% and reduced perplexity by 0.66. It reduces the L2 norm of weights while slightly increasing gradient norm, producing larger relative weight updates that help in late training stages.

• Activation normalization is the expensive part of nGPT. Applying weight normalization alone achieved final task performance comparable to the full nGPT solution, while nGPT's additional activation normalization layers increase SLM training time by more than 20%.

• Meta tokens help linear attention too. Prepending 256 meta tokens improved language modeling and reasoning accuracy by +0.45% on average with negligible overhead, serving as learned cache initialization for recurrent-form linear attention.

• Nemotron-Flash advances the frontier. With all models accelerated using TensorRT-LLM's AutoDeploy kernels and CUDA Graph, Nemotron-Flash-3B achieved +2.0%/+5.5% higher average accuracy, 1.7×/1.3× lower latency, and 6.4×/18.7× higher throughput than Qwen2.5-3B/Qwen3-1.7B. Nemotron-Flash-1B achieved +5.5% higher average accuracy, 1.9× lower latency, and 45.6× higher throughput than Qwen3-0.6B.

• Most competitive in-domain accuracy at scale. Nemotron-Flash-3B attained the highest accuracy in commonsense reasoning, math, coding, and recall tasks among models larger than 1.5B parameters. Despite containing only 2 and 3 full-attention layers respectively, Nemotron-Flash-1B and 3B maintained the most competitive recall accuracy, indicating full KV cache across all layers is not necessary.

• Long-context retrieval still needs full attention. In the attention configuration ablation (continuous pretraining at 29k context for 25B tokens), the 1FA+2SWA setting achieved 1.6× higher throughput than 3FA (4657 versus 2939 tok/s) while general benchmark accuracy stayed largely unaffected, but NIAH performance dropped significantly at longer context lengths when FA layers were reduced to one. The authors recommend keeping at least two full attention layers even in SLMs.

Methodology in Plain English

The authors attack the problem in two stages: architecture and training.

For architecture, they first run controlled experiments. They train Llama models at five different depths across a range of widths, all on the same 100B-token corpus, and plot accuracy against both parameter count and measured decoding latency. This isolates whether depth or width is buying the speed-up. They then extend a standard scaling law so that loss is written as a function of depth and width separately rather than a single lumped "model size" term, which lets them fit a curve on some configurations and predict the best depth–width point for an unseen latency budget.

For operators, they train 500M-parameter models using roughly the same recipe but swap in each candidate layer type — Mamba, Mamba2, GLA, DeltaNet, Gated DeltaNet, RWKV7, and sliding window attention — using the original implementations and official kernels. They measure both perplexity and actual decoding latency, then test promising operators in pairs within hybrid models.

For the search, they build an evolutionary engine using aging evolution. An initial population of architectures is short-trained; in each cycle, tournament selection picks parents that score well on the short-training proxy and fit a latency budget; mutations then change operator ratios, FFN ratios, or block type counts; offspring are short-trained and their latency is looked up from a precomputed table. The proxy is justified by the observation that architecture rankings stabilize early in training and by the 88.8% rank correlation with full-training perplexity.

For training, they project weight matrices onto a unit norm sphere after each optimizer step, row-wise for weights applied to hidden features and column-wise for weights whose outputs are added back to hidden features. This removes the radial component of the update and forces learning to happen through angle changes. They also prepend 256 learnable meta tokens.

The final Nemotron-Flash-1B (0.96B parameters, hidden size 2048, 12 blocks, 24 operators) and Nemotron-Flash-3B (2.7B parameters, hidden size 3072, 36 operators) were trained for 4.5T tokens on 256 NVIDIA H100 GPUs with the Adam optimizer, no weight decay, a cosine schedule with initial learning rate 1e-3, batch size of 2M tokens, and a 4096 context length for all but the final 25B tokens, which used a 29000 context length.

Why This Matters

Impact on research. The paper reframes SLM design around deployment hardware rather than parameter count, and it contributes reusable machinery: a depth-and-width scaling law, evidence that early-training rankings are a reliable search signal, a weight normalization scheme that captures most of nGPT's benefit at lower cost, and a demonstration that meta tokens help linear attention as well as softmax attention. The negative result on deep-thin models challenges a design default used by prior SLMs.

Real-world applications:

• On-device assistants and chatbots running on consumer GPUs or laptops, where interactive response time at batch size 1 is the binding constraint.
• Retrieval and recall workloads such as document question answering and extraction, where the paper shows competitive recall accuracy is achievable with only 2–3 full attention layers.
• Latency-sensitive industrial and embedded deployments that the paper explicitly cites, where hardware is resource-constrained and high latency blocks practical use.
• High-throughput serving, where the throughput-optimized Nemotron-Flash-3B-TP reached 4657 tok/s and 10.1×/29.7× higher throughput than Qwen2.5-3B and Qwen3-1.7B.

Industry relevance. The released family (Nemotron-Flash-1B, 3B, and 3B-Instruct on Hugging Face) is benchmarked using production deployment tooling — TensorRT-LLM AutoDeploy kernels and CUDA Graph — so reported latency and throughput reflect realistic serving conditions rather than theoretical FLOP counts. The finding that at least two full attention layers should be retained gives practitioners a concrete configuration rule.

Future Directions

• Whether the depth–width scaling law generalizes beyond the tested model family, devices, and generation lengths. The paper notes that detailed trade-off curves shift across devices and generation lengths, which complicates selecting depth and width.

• Broadening the hybrid search space. The search was restricted to three candidate operators (DeltaNet, attention, Mamba2) and a maximum of three block types placed in early, middle, or late stages; whether other efficient operators or finer-grained placements yield further gains is untested.

• Closing the long-context gap. Reducing full attention layers to one preserved general benchmark accuracy but caused NIAH performance to drop significantly at longer context lengths — an open problem for latency-optimal long-context SLMs.

• Extending the general-purpose training techniques. Weight normalization and meta tokens were validated on Llama, DeltaNet, and Mamba2 1B models and on the searched architecture; the paper presents them as generalizable components for future SLMs but does not establish their behavior at larger scale or on other architectures.

Target Audience

Researchers and engineers working on efficient language model architecture, on-device or edge inference, and neural architecture search. It is also relevant to practitioners who must choose a small model under a concrete latency or throughput budget, and to those evaluating hybrid attention–state-space designs, since the paper provides operator-level perplexity and latency comparisons plus a documented search procedure. Readers without background in linear attention variants or scaling laws will find the architectural sections demanding.

Authors’ abstract

Efficient deployment of small language models (SLMs) is essential for numerous real-world applications with stringent latency constraints. While previous work on SLM design has primarily focused on reducing the number of parameters to achieve parameter-optimal SLMs, parameter efficiency does not necessarily translate into proportional real-device speed-ups. This work aims to identify the key determinants of SLMs' real-device latency and offer generalizable principles and methodologies for SLM design and training when real-device latency is the primary consideration. Specifically, we identify two central architectural factors: depth-width ratios and operator choices. The former is crucial for small-batch-size latency, while the latter affects both latency and large-batch-size throughput. In light of this, we first study latency-optimal depth-width ratios, with the key finding that although deep-thin models generally achieve better accuracy under the same parameter budget, they may not lie on the accuracy-latency trade-off frontier. Next, we explore emerging efficient attention alternatives to evaluate their potential as candidate building operators. Using the identified promising operators, we construct an evolutionary search framework to automatically discover latency-optimal combinations of these operators within hybrid SLMs, thereby advancing the accuracy-latency frontier. In addition to architectural improvements, we further enhance SLM training using a weight normalization technique that enables more effective weight updates and improves final convergence. Combining these methods, we introduce a new family of hybrid SLMs, called Nemotron-Flash, which significantly advances the accuracy-efficiency frontier of state-of-the-art SLMs, e.g., achieving over +5.5% average accuracy, 1.3x/1.9x lower latency, and 18.7x/45.6x higher throughput compared to Qwen3-1.7B/0.6B, respectively.