Nemotron Elastic: Towards Efficient Many-in-One Reasoning LLMs | Summary

Contents

Summary
Introduction
Methodology
Experiments and Results
Related Work
Conclusions
Appendix
Brief Thoughts

This article explains the key points of Nemotron Elastic: Towards Efficient Many-in-One Reasoning LLMs.

2025-11-20 (arXiv)
Taghibakhshi, Ali, Sreenivas, Sharath Turuvekere, Muralidharan, Saurav, Cai, Ruisi, Chochowski, Marcin, Mahabaleshwarkar, Ameya Sunil, Suhara, Yoshi, Olabiyi, Oluwatobi, Korzekwa, Daniel, Patwary, Mostofa, et al.
NVIDIA
Paper
Project Page

Summary

Producing several sizes of reasoning large language models, which generate intermediate steps to solve complex problems, usually requires separate training or compression runs. Starting from a pretrained 12B-parameter parent, Nemotron Elastic jointly trains nested 6B, 9B, and 12B variants—smaller models that reuse subsets of the parent’s weights—with 110B training tokens; under reasoning-enabled evaluation on six benchmarks, its 12B variant averages 77.41 versus 77.38 for NanoV2-12B.
The method combines weight sharing, a learned router that selects model dimensions for a target parameter budget, and knowledge distillation, which trains students to match a teacher’s predictions. Training progresses from sequences of 8192 tokens to 49152 tokens, with the second stage allocating more samples to the largest variant.
The reported 110B-token budget compares with 750B for sequential compression and 40T for training the smaller models from scratch, although several baseline budgets are estimates. Shared BF16 weights, a 16-bit parameter representation, occupy 24 GB for all three variants versus 42 GB for separate NanoV2 9B and 12B weights; the comparison does not quantify runtime caches, activations, or inference speed.

Introduction

Independently training each model size repeats expensive optimization, while conventional compression requires separate pruning and recovery runs. Pruning removes model components; elastic or Matryoshka-style networks instead embed smaller models inside a larger model, sharing parameters and allowing extraction without additional training.

Hybrid Mamba-Attention architectures combine attention, which computes interactions between sequence positions, with state-space models (SSMs), sequence-processing components such as Mamba. Mamba supports linear-time sequence processing, but its group and channel constraints complicate compression.
Reasoning models need long sequences for intermediate steps. The paper targets both architectural flexibility across resource budgets and preservation of extended-context reasoning.
Figure 1 averages MATH-500, mathematics problems; AIME-2024 and AIME-2025, American Invitational Mathematics Examination problems; GPQA, graduate-level science questions; LiveCodeBench v5, code generation and problem solving; and MMLU-Pro, college-level multiple-choice reasoning. Its second panel illustrates training-token and stored-weight scaling with family size.

Average reasoning-benchmark accuracy by model size and illustrated training-token and deployment-memory scaling with model-family size.

Methodology

The pipeline first ranks components by importance, then trains a shared student model and budget-conditioned routers to select width, the sizes of component dimensions, and depth, the number of retained layers. A separate frozen original model supplies distillation targets, while differentiable masks select components and carry training gradients to the routers.

Figure 2 connects training-time masking to zero-shot slicing, extraction without further training: one trained checkpoint contains the supported model sizes.

Budget-conditioned masking and knowledge distillation during training, followed by zero-shot extraction of nested 12B, 9B, and 6B models.

Importance Estimation and Model Preparation

Width importance is estimated from forward-pass activation magnitudes for embedding channels, the dimensions of token representations; feed-forward network (FFN) intermediate neurons, the units between its linear projections; attention heads, parallel attention components; and Mamba heads and their channels. Mamba channel scores aggregate across heads, head scores use the top-ranked channels, and heads are ranked independently within their groups to preserve SSM structure.

Embedding scores aggregate absolute layer-normalized input activations across batch and sequence dimensions; layer normalization rescales a layer’s input features. FFN scores use the first linear projection’s intermediate activations, and attention-head scores aggregate query-projection activations.
Depth importance uses iterative layer ablation, removal to measure a layer’s effect, and normalized mean squared error (MSE). It compares full-model and ablated-model logits, the scores before output probabilities are computed, and divides their squared difference by the full-model squared magnitude.
For layer j, the score is s_j = sum_{B,L}(M_full - M_-j)^2 / sum_{B,L}M_full^2. Here B and L index batch and sequence dimensions, and M denotes model logits; layers with the highest scores are preferentially retained.

Elastic Formulation

Elastic width selects embedding dimensions, FFN intermediate dimensions, attention-head counts, Mamba-head counts, and Mamba head-channel counts from predefined choices. Smaller variants retain contiguous subsets of importance-ranked components used by larger variants, preserving nested weight sharing.

Elastic depth uses binary layer-retention coefficients. Inactive layers are bypassed through residual connections, paths that carry a layer’s input forward.
Mamba selection preserves group structure and channel consistency, while attention selection operates on heads. Homogeneous configurations use the same compression ratio across instances of a component type; heterogeneous configurations allow ratios to vary by layer.

Elastic Training

Router Architecture and Design assigns a separate two-layer network with LeakyReLU, an activation function between its linear layers, to each dynamic dimension. Its input is a one-hot representation identifying the target compression level; outputs specify shared component choices in homogeneous mode or separate choices for each layer in heterogeneous mode, while embedding selection remains global.

Loss Formulation uses Gumbel-Softmax, a differentiable relaxation of discrete configuration selection. Annealing its temperature shifts choices from soft exploration toward sharper decisions.
The router objective penalizes deviation between selected resource cost and the target constraint. The framework lists parameter count, memory, latency, and throughput as possible costs; the reported experiment targets parameter budgets.

Versatile Training Options includes cross-entropy training, which fits the probabilities of training targets; distillation from a trainable teacher; distillation from a frozen teacher; and combinations of these modes. Distillation minimizes forward Kullback–Leibler divergence, a discrepancy between the teacher’s and student’s output probability distributions; the reported nested-model experiment uses a frozen original 12B teacher.

A trainable teacher also receives cross-entropy supervision to prevent collapse during self-distillation, training with supervision from the model itself. Frozen teachers provide fixed targets without gradient updates.
Final Optimization Target combines task loss with a weighted router loss: L_total = L_task(theta) + lambda * L_router(psi). Here theta denotes model parameters, psi denotes router parameters, and lambda weights the resource penalty.

Two-Stage Training with Curriculum-Based Sampling assigns each sample to one target budget. Stage 1 samples budgets uniformly to stabilize architecture learning; Stage 2 uses longer sequences and favors the full-budget student to address the performance imbalance observed under uniform sampling.

The experimental Stage 2 probabilities are p(12B) = 0.5, p(9B) = 0.3, and p(6B) = 0.2.
The authors interpret the imbalance as gradient competition, conflicting training updates between shared variants. The sampling ablation supports a tradeoff in budget allocation but does not directly measure that mechanism; the frozen teacher remains unchanged.

Implementation

Dynamic Model Formulation applies dimension-specific masks to activations and weight matrices within a shared architecture, with binary coefficients controlling layer retention. Dynamic Mamba preserves group-aware head selection and consistent head channels through projections, causal convolution, which processes only current and preceding positions, state-space updates, and masked outputs.

Dynamic Attention masks embedding dimensions and the query, key, and value projections used to compute attention, followed by output projection and masking.
Dynamic FFN masks embedding and intermediate dimensions across its two linear transformations and output.
Mask Generation converts router choices into masks retaining the highest-ranked components. Heterogeneous mode constructs per-layer masks; depth selection activates the highest-ranked layers.
The implementation supports hard selection scaled by the selected logit or probability-weighted soft masks. Deployment uses argmax, selection of the highest-probability option, with the selected logit set to 1.

Elastic Model Deployment

For a budget included during training, the router identifies parameters to retain, and the remaining components are permanently sliced out of the elastic checkpoint. The resulting model requires no additional retraining, fine-tuning, or distillation.

The extraction procedure applies to the trained budget set. Equivalent quality for arbitrary unseen budgets is not evaluated.
A shared parent checkpoint avoids separately storing every variant. Selecting among trained budgets is distinct from task-specific dynamic inference-time routing, which the conclusions list as future work.

Experiments and Results

Experiments compress NVIDIA Nemotron Nano V2 12B into nested 9B and 6B models, corresponding to 25% and 50% parameter reductions. The paper states that both base and reasoning variants are considered; its main quantitative results report six reasoning and knowledge benchmarks for the nested family.

Experimental Setup

Training uses the same compression data blend as NanoV2-9B for calibration, estimating component importance, and distillation. Importance estimation processes 1024 samples of length 8192; Phase 1 uses batch size 1536 and sequence length 8192 for approximately 65B tokens, while Phase 2 uses batch size 512 and sequence length 49152 for approximately 45B tokens.

Evaluation enables reasoning and reports pass@1, success for a single generated answer. The paper states that results are averaged over 4 to 16 shots as appropriate for each of MATH-500, AIME-2024, AIME-2025, GPQA, LiveCodeBench v5, and MMLU-Pro.
Model and router learning rates are 9e-5 and 1e-2, respectively, with a 60-step linear warmup. Gumbel-Softmax temperature decreases from 1.0 to 0.05, router loss weight is 1.0, router-logit scaling increases from 1.0 to 10.0, and router hidden dimension is 256.
Stage 1 gives each of the 6B, 9B, and 12B budgets probability 1/3. Stage 2 uses probabilities 0.2, 0.3, and 0.5, respectively.

Six-benchmark accuracy comparison for nested Nemotron-Elastic variants and NanoV2 and QWen3 baselines.

Exploratory, final, and total training-token budgets for deriving 6B and 9B models.

BF16 weight-memory comparison between shared Nemotron Elastic variants and separate NanoV2 checkpoints.

Results

Table 1 reports average scores of 70.61, 75.95, and 77.41 for Nemotron-Elastic-6B, 9B, and 12B. The 9B variant closely matches NanoV2-9B’s 75.99 average, and the 12B variant closely matches NanoV2-12B’s 77.38, with benchmark-specific gains and losses.

Elastic-9B scores 75.42 on AIME-2025 versus 71.43 for NanoV2-9B, but scores slightly lower on the other five benchmarks. Elastic-12B improves mathematics and coding scores while falling below NanoV2-12B on GPQA and MMLU-Pro.
Table 2 compares 110B elastic-training tokens with 480B exploratory plus 270B final compression tokens, totaling 750B, and 40T pretraining tokens. The caption identifies estimated budgets for Minitron-SSM, the group-aware hybrid-model pruning and distillation baseline, at 6B and for NanoV2 6B/9B pretraining.
Table 3 compares 24 GB of shared BF16 weights for 6B + 9B + 12B with 42 GB for separate 9B + 12B NanoV2 weights, a reported 43% reduction. Router overhead is stated to be below 2% of largest-model size; the table does not quantify runtime caches and activations.

Stage 1 and Stage 2 accuracy, absolute gains, and relative improvements across three model sizes and three benchmarks.

Uniform versus adjusted budget-sampling accuracy, with NanoV2 reference scores.

Effects of Two-Stage Training

Table 4 shows that the extended-context stage raises 6B AIME-2025 accuracy from 56.88 to 68.13, an 11.25-percentage-point gain and 19.8% relative improvement. Corresponding 9B and 12B gains are 6.67 and 2.91 points, with smaller gains on MATH-500 and positive gains on GPQA.

These results support the usefulness of the second training stage in this setting. Because it also adds approximately 45B training tokens, the comparison does not isolate context length from additional optimization.
Table 5 shows that adjusted sampling raises 12B AIME-2025 from 72.29 to 75.83 and GPQA from 61.11 to 63.25, gains of 3.54 and 2.14 percentage points.
The adjustment has costs for smaller variants: 6B GPQA falls from 55.30 to 53.78, while 9B MATH-500 and GPQA fall slightly. The results support prioritizing full-model quality, with losses on some smaller-budget tasks.

The paper compares Nemotron Elastic with pruning-and-distillation methods that recover each target size separately, hybrid SSM-Transformer models that combine state-space processing with attention, and nested architectures such as MatFormer, Flextron, and MatMamba. Its proposed distinction is the combination of hybrid structural constraints, reasoning-oriented extended-context training, and end-to-end budget-conditioned architecture selection.

MatFormer uses nested Transformer models and MixnMatch, heuristics for selecting sub-networks; Flextron introduces input-adaptive routing; MatMamba extends nested sub-blocks to Mamba. The paper also cites Gemma 3n as a deployment example of nested models with conditional parameter loading.
The reasoning-training discussion contrasts architectural compression with prompting and reinforcement learning from reasoning traces. The main results do not provide controlled comparisons against all cited elastic architectures or establish that extended-context training is uniquely necessary for reasoning models.

Conclusions

The demonstrated outcome is a 6B/9B/12B reasoning-model family extracted from one elastic checkpoint after 110B tokens of joint training. The conclusions emphasize extended-context adaptation and shared-weight deployment, and identify larger families, task-specific architecture selection, dynamic inference-time routing, and quantization, reduced-precision parameter storage, as future directions.

Appendix

The supplied 17-page PDF contains no appendix; it ends with Acknowledgments and References. The model link provided on the first page is https://huggingface.co/nvidia/Nemotron-Elastic-12B, and the paper identifier is arXiv:2511.16664v1.

Brief Thoughts

The strongest evidence is the near-preservation of 9B and 12B average accuracy while obtaining all three sizes through one training run, together with the shared-weight storage comparison. The sampling ablation exposes a practical tradeoff: protecting the largest variant can reduce smaller-model scores on some tasks. The reported token savings depend partly on estimated baseline budgets and exclude the already pretrained parent’s cost. They describe model-family derivation budgets, without measuring end-to-end compute or monetary savings. Evaluation covers one 12B parent family and three budgets, with no reported uncertainty intervals or serving-latency measurements. Constant-token scaling to larger families remains an extrapolation beyond this configuration. The paper describes normalized-MSE depth ranking and heterogeneous width selection but provides no isolated quantitative ablation of their contributions. Their individual effects cannot be separated from the full training recipe.