PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ali Taghibakhshi, Sharath Turuvekere Sreenivas, Saurav Muralidharan, Ruisi Cai, Marcin Chochowski, Ameya Sunil Mahabaleshwarkar, Yoshi Suhara, Oluwatobi Olabiyi, Daniel Korzekwa, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro, Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov

arXiv:2511.16664v1cs.CL

TL;DR

Training separate model sizes is prohibitively expensive, and existing compression remains costly for reasoning models with long-context needs. Nemotron Elastic embeds multiple weight-sharing submodels in one parent using a trained router and reasoning-specific curriculum. It derives 6B and 9B variants from a 12B model with 110B tokens while preserving accuracy and constant deployment memory across the family.

  • Problem

    Training separate model sizes and compressing each variant require prohibitively large token budgets, while reasoning models also need elastic deployment and long-context capability.

    Source: Introduction · pp. 1–2

  • Method

    Nemotron Elastic uses an end-to-end router, nested weight sharing, long-context two-stage training, teacher distillation, and structural elastification for hybrid Mamba-Attention models.

    Source: Introduction · pp. 2–3

  • Results

    110B training tokens produce 6B and 9B variants from a 12B parent, yielding a 360× reduction versus training from scratch and a 7× reduction versus sequential compression.

    Source: Conclusions · p. 15

  • Takeaways & Limitations

    A single elastic checkpoint provides multiple reasoning models with competitive accuracy and constant deployment memory as the model family grows.

    Source: Introduction · p. 2, Conclusions · p. 15

Abstract

from arXiv · show

Training a family of large language models targeting multiple scales and deployment objectives is prohibitively expensive, requiring separate training runs for each different size. Recent work on model compression through pruning and knowledge distillation has reduced this cost; however, this process still incurs hundreds of billions of tokens worth of training cost per compressed model. In this paper, we present Nemotron Elastic, a framework for building reasoning-oriented LLMs, including hybrid Mamba-Attention architectures, that embed multiple nested submodels within a single parent model, each optimized for different deployment configurations and budgets. Each of these submodels shares weights with the parent model and can be extracted zero-shot during deployment without additional training or fine-tuning. We enable this functionality through an end-to-end trained router, tightly coupled to a two-stage training curriculum designed specifically for reasoning models. We additionally introduce group-aware SSM elastification that preserves Mamba's structural constraints, heterogeneous MLP elastification, normalized MSE-based layer importance for improved depth selection, and knowledge distillation enabling simultaneous multi-budget optimization. We apply Nemotron Elastic to the Nemotron Nano V2 12B model, simultaneously producing a 9B and a 6B model using only 110B training tokens; this results in over 360x cost reduction compared to training model families from scratch, and around 7x compared to SoTA compression techniques. Each of the nested models performs on par or better than the SoTA in accuracy. Moreover, unlike other compression methods, the nested capability of our approach allows having a many-in-one reasoning model that has constant deployment memory against the number of models in the family.

Introduction

Nemotron Elastic addresses the cost and flexibility limits of training separate model sizes by embedding nested, weight-sharing submodels in one reasoning-oriented hybrid LLM. Its router and training curriculum support long-context reasoning while reducing training cost and deployment-memory growth.

Introduction

  • Existing compression methods still require hundreds of billions of tokens per compressed model, while nested networks share parent weights and extract submodels without additional training.Source: Introduction · p. 1
  • Reasoning models require flexible computation and long-context processing, but existing compression methods address neither elastic deployment nor extended reasoning sequences.Source: Introduction · p. 2
  • Nemotron Elastic uses an end-to-end trained router and two-stage curriculum to produce multiple deployment budgets from one elastic training run, with extended-context training at 49K tokens.Source: Introduction · p. 2
  • The resulting nested models achieve competitive or superior accuracy against independently trained baselines while supporting faster inference.Source: Introduction · p. 2
  • The framework introduces normalized-MSE depth ranking, teacher-guided knowledge distillation, group-aware SSM elastification, and heterogeneous layer-wise configurations.Source: Introduction · pp. 2–3
  • 110B training tokens derive 6B and 9B variants from a 12B parent, reducing cost 7× versus NanoV2 Compression and 360× versus pretraining from scratch.Source: Introduction · p. 3

Methodology

Nemotron Elastic constructs nested hybrid-model variants by ranking components, selecting width and depth under structural constraints, and extracting weight-sharing sub-networks without fine-tuning. Its formulation covers Mamba, attention, FFN, embedding, and layer-depth choices.

Methodology

  • The methodology combines importance estimation, elastic width and depth selection, two-stage router training, and dynamic masking for multi-budget hybrid models.Source: Methodology · p. 3

Width

Depth

  • Normalized MSE between full-model and layer-ablated predictions ranks layers by contribution, with higher-ranked layers preferentially retained at smaller depths.Source: Depth · pp. 3–4

Elastic Formulation

  • Width elasticity varies embedding size, attention heads, FFN dimensions, Mamba heads, and Mamba channels, retaining contiguous ranked subsets for smaller models.Source: Elastic Formulation · p. 4
  • Depth elasticity uses binary layer-selection coefficients, bypassing inactive layers through residual skips while retaining important layers at lower budgets.Source: Elastic Formulation · pp. 4–5
  • Hybrid selection jointly balances Mamba’s efficient sequence processing with attention’s contextual reasoning while respecting group and head-wise constraints.Source: Elastic Formulation · p. 5

Elastic Training

Dedicated routers learn architecture choices for each dynamic dimension from target compression levels. Their outputs support homogeneous or layer-wise heterogeneous configurations across the joint model design space.

Router Architecture and Design

  • Each dynamic dimension has a dedicated two-layer router with leaky-ReLU activation that searches configurations for embeddings, Mamba, attention heads, FFNs, and depth.Source: Router Architecture and Design · p. 5
  • Routers receive one-hot vectors encoding the target compression level, allowing architecture decisions to depend on the requested model configuration.Source: Router Architecture and Design · p. 5
  • Homogeneous outputs assign one compression choice across a component type, whereas heterogeneous outputs allow each Mamba, attention, or FFN layer to choose independently.Source: Router Architecture and Design · p. 5

Loss Formulation

Nemotron Elastic combines router-based architecture selection with flexible task and distillation losses. The router uses relaxed discrete choices and resource-aware objectives to balance accuracy against deployment constraints.

Loss Formulation

  • Gumbel-Softmax converts router logits into relaxed configuration selections, allowing gradients to update router parameters during training.Temperature annealing moves from soft exploration toward sharper architectural decisions.Source: Loss Formulation · p. 5
  • The router objective maps selected configurations to resource costs such as parameters, memory, latency, or throughput.A target constraint guides search toward configurations balancing resource use and model performance.Source: Loss Formulation · pp. 5–6
  • End-to-end router optimization discovers Pareto-efficient configurations while accounting for hybrid Mamba-Attention trade-offs.The method combines importance-based ordering with learned policies rather than relying only on post-hoc scores.Source: Loss Formulation · p. 6

Versatile Training Options

  • Cross-entropy trains the model directly on corpus data and can operate alone or alongside other objectives.The loss uses the training dataset, model parameters, and predicted token distributions.Source: Versatile Training Options · p. 6
  • Knowledge distillation transfers teacher information through forward KL divergence between temperature-scaled student and teacher distributions.This supplies an alternative supervision signal for elastic submodels.Source: Versatile Training Options · p. 6
  • Trainable-teacher mode jointly updates the full-budget teacher and student, while cross-entropy prevents self-distillation collapse.The teacher adapts to the training distribution but provides moving supervision targets.Source: Versatile Training Options · p. 6
  • Frozen-teacher mode keeps teacher parameters fixed, reducing computational overhead and providing stable supervision.The teacher may be the original full model or an alternative architecture.Source: Versatile Training Options · p. 6
  • Mixed training modes can combine trainable and frozen teachers across phases, submodels, and loss combinations.This supports multi-objective training scenarios with different supervision strategies.Source: Versatile Training Options · p. 6

Final Optimization Target

Nemotron Elastic jointly optimizes task learning and router decisions while training multiple budgets through a two-stage curriculum. Structured masking then exposes deployable width- and depth-variable submodels without recompilation.

Final Optimization Target

  • The combined objective adds a router penalty to the task loss, balancing accuracy against resource constraints through λ.The task term may use cross-entropy, distillation, or both.Source: Final Optimization Target · p. 6
  • End-to-end coupling lets architecture discovery respond to the actual training signal instead of relying on frozen-checkpoint proxy scores.This distinguishes the method from decoupled importance-based searches such as Minitron and Minitron-SSM.Source: Final Optimization Target · pp. 6–7

Two-Stage Training with Curriculum-Based Sampling

  • Uniform short-context sampling stabilizes the router by giving each target budget balanced training exposure.This establishes diverse architectural patterns before reasoning becomes the dominant bottleneck.Source: Two-Stage Training with… · p. 7
  • The two-stage curriculum exposes elastic variants to long reasoning chains and guides architecture discovery toward extended-context capability.Stage 1 explores configurations, whereas Stage 2 enforces reasoning-specific constraints.Source: Two-Stage Training with… · p. 7
  • Extended-context training shifts to non-uniform sampling that prioritizes full-budget models while continuing to train smaller variants.Uniform extended-context sampling degraded the full model while improving smaller budgets, indicating gradient competition.Source: Two-Stage Training with… · p. 7

Implementation

  • Structured binary masks select active components during hybrid Mamba-Attention-MLP training, enabling simultaneous multi-budget optimization.The masking formulation preserves transparent architectures and straightforward subnetwork deployment.Source: Implementation · pp. 7–8

Dynamic Model Formulation

  • The framework supports comprehensive width and depth flexibility by dynamically masking both the layer stack and dimensions within each layer type.Depth uses layer-retention coefficients, while width uses dimension-specific masks.Source: Dynamic Model Formulation · p. 8
  • Embedding, attention, Mamba, and MLP dimensions can be selectively retained through binary masks, producing different parameter budgets from one parent model.The dynamic operator applies masks to layer outputs and component dimensions.Source: Dynamic Model Formulation · p. 8

Dynamic Mamba

Dynamic Mamba elastification applies masks to Mamba dimensions while preserving the grouped structure required by state-space computations. The forward pass masks normalized inputs, projections, intermediate states, and outputs.

Dynamic Mamba

  • The Mamba operator masks embedding dimensions and Mamba-specific dimensions before processing the layer output.The output mask is constructed from dynamic embedding and structural Mamba constraints.Source: Dynamic Mamba · p. 8
  • Embedding masks retain a prefix of hidden dimensions and broadcast across activation or weight-matrix dimensions.This operator is applied to layer-normalization outputs and matrices interfacing with the embedding dimension.Source: Dynamic Mamba · p. 8
  • Group-aware Mamba masks preserve permutation structure by retaining consistent head and channel selections within each Mamba group.The construction maintains head-channel consistency across heads.Source: Dynamic Mamba · p. 8
  • The Mamba forward pass applies masked layer normalization, computes z, x, B, C, and dt projections, then performs causal convolution and state-space updates.Gated RMSNorm and output projection complete the computation before final masking.Source: Dynamic Mamba · pp. 8–9
  • Dynamic masks are applied again at the layer output, yielding the elastified Mamba result while preserving the masked computation path.The complete output is represented as 𝒟(Mambaℓ(y)) = yout.Source: Dynamic Mamba · p. 9

Dynamic Attention

Dynamic attention elastification controls capacity through embedding and attention-head masks. The masked forward pass normalizes inputs, computes query-key-value projections, applies attention, and projects the output.

Dynamic Attention

  • Attention capacity is adjusted by masking attention heads and embedding dimensions in the hybrid architecture.The output mask combines dynamic embedding and attention-head constraints.Source: Dynamic Attention · p. 9
  • The attention-head mask selects active head and head-dimension entries through a structured binary index.Active dimensions are defined over the attention-head and per-head-dimension axes.Source: Dynamic Attention · p. 9
  • The dynamic attention forward pass applies masked layer normalization before computing query, key, and value projections.Projection matrices use the number of attention heads and head dimension to determine capacity.Source: Dynamic Attention · p. 9
  • Attention computation is followed by output projection and final application of the embedding and head masks.This produces an output whose capacity matches the selected elastic configuration.Source: Dynamic Attention · p. 9

Dynamic FFN

Dynamic FFN elastification masks embedding and intermediate dimensions so feed-forward layers can operate at selected widths. Depth adaptation retains important layers and bypasses skipped layers through residual connections, preserving signal propagation while reducing computation.

Dynamic FFN

  • Dynamic FFN masking applies to both embedding and intermediate dimensions before producing the layer output.The operator supports matrices oriented with either dimension derived from the FFN intermediate size.Source: Dynamic FFN · pp. 9–10
  • The forward pass combines masked layer normalization, a dynamically masked first linear transformation, activation, and a dynamically masked second transformation.The first matrix maps into the intermediate dimension, while the second maps back to the embedding dimension.Source: Dynamic FFN · p. 10
  • Depth adaptation selects a target number of active layers using importance-based retention and bypasses skipped layers with residual connections.For hybrid models, this allows different scales to retain complementary Mamba and attention components.Source: Dynamic FFN · p. 10

Mask Generation

The router converts budget decisions into differentiable or discrete masks for heterogeneous layer configurations, then uses those decisions to slice deployable nested models. Extracted variants require no retraining or fine-tuning and can be selected from one checkpoint.

Mask Generation

  • Gumbel-Softmax router outputs generate configuration selections for homogeneous budgets or per-layer heterogeneous architectures.Heterogeneous routing partitions selections across layers, allowing distinct compression ratios and depth choices.Source: Mask Generation · p. 10
  • Generated masks control embeddings, Mamba, attention heads, FFNs, and depth retention, enabling dynamic capacity adjustment across hybrid model components.Depth selection activates the highest-ranked target number of layers.Source: Mask Generation · p. 10
  • The framework supports hard argmax masks at inference and soft probabilistic masks during training so gradients reach all configuration options.Hard selection applies the chosen configuration, while soft masking combines candidates proportionally to their probabilities.Source: Mask Generation · p. 10

Elastic Model Deployment

  • A trained router slices the full checkpoint into target-budget sub-networks by permanently removing components selected for each budget.The extracted parameters form nested models corresponding to the desired parameter counts.Source: Elastic Model Deployment · pp. 10–11
  • Zero-shot slicing is computationally negligible and produces inference-ready variants without retraining, fine-tuning, or additional distillation.Any budget used during training, including intermediate sizes, can be deployed directly from the single full-model checkpoint.Source: Elastic Model Deployment · p. 11
  • One full-size checkpoint can serve multiple trained variants on demand, enabling per-request selection based on latency or resource constraints.The variants share learned representations and architectural decisions, avoiding separate fine-tuning or calibration.Source: Elastic Model Deployment · p. 11

Experiments and Results

Nemotron Elastic trains nested 6B, 9B, and 12B variants from one 12B parent using two-stage, multi-budget distillation. The approach matches the parent’s performance closely, reduces training and deployment costs, and benefits smaller reasoning models most from extended-context training.

Experimental Setup

  • The evaluation spans knowledge, language, mathematical, algorithmic, and code-reasoning benchmarks using pass@1 with reasoning enabled.The benchmark suite includes MMLU-Pro, GPQA, MATH-500, AIME-2024, AIME-2025, and LiveCodeBench v5.Source: Experimental Setup · p. 11
  • Training proceeds in two stages: an initial short-context phase with sequence length 8192, followed by an extended-context phase with sequence length 49152.Source: Experimental Setup · p. 11

Results

  • 77.41 average score for Nemotron-Elastic-12B nearly matches NanoV2-12B’s 77.38 while also producing 9B and 6B deployments from one run.All three variants use a frozen 12B teacher and share a common foundation.Source: Experimental Setup · p. 11, Results · p. 12, Table 1 · p. 13
  • Around 7X fewer training tokens are required than Minitron-SSM for deriving 6B and 9B models.The saving comes from eliminating separate exploratory distillation runs through end-to-end simultaneous budget optimization.Source: Results · p. 12, Table 2 · p. 13
  • 43% less deployment memory is used than NanoV2’s two models despite storing three Nemotron Elastic variants.Nested weight sharing makes total memory equivalent to the largest model plus routing metadata rather than the sum of separate checkpoints.Source: Results · pp. 12–13, Table 3 · p. 13

Effects of Two-Stage Training

Related Work

Prior work combines pruning, distillation, hybrid SSM-Transformer architectures, and nested weight sharing, but does not jointly address hybrid elastification, reasoning-specific long-context training, and learned heterogeneous routing.

Model Compression and Pruning

  • Structured pruning with knowledge distillation improves compressed Transformer models, while group-aware SSM pruning preserves sequence-modeling constraints but still requires separate distillation per target size.These approaches reduce model size but do not provide one jointly optimized family.Source: Related Work · p. 14

Hybrid SSM-Transformer Models

  • Hybrid Transformer–SSM architectures show promise for efficient long-context modeling; Nemotron-H replaces 92% of attention layers with Mamba2 blocks, achieving 3× inference speedup, while concurrent compression work [15] requires separate distillation per model size.Source: Related Work · p. 14

Elastic and Nested Architectures

  • MatFormer, Flextron, and MatMamba establish nested or elastic architectures, yet prior methods lack joint support for hybrid models, reasoning-focused extended-context training, and end-to-end heterogeneous layer selection.Gemma 3n provides a practical example of nested models with conditional parameter loading.Source: Related Work · p. 14

Reasoning Model Training

  • Reasoning-model research emphasizes extended thought chains and reasoning traces, but does not address architectural efficiency or elastic deployment.Nemotron Elastic identifies 49K-token extended-context training as critical for compressed reasoning variants.Source: Related Work · pp. 14–15

Conclusions

Nemotron Elastic makes elastic reasoning-model training practical by deriving multiple zero-shot submodels from one parent with substantially lower training cost and constant deployment memory.

  • 110B training tokens derive a model family from one 12B parent, reducing cost 360x versus training from scratch and 7x versus sequential compression.The nested submodels preserve accuracy, can be extracted by zero-shot slicing, and add no deployment memory overhead as the family grows.Source: Conclusions · p. 15
  • Extended-context training is critical for preserving reasoning performance across deployment scales.The framework identifies extended-context training as a key requirement for elastic reasoning models.Source: Conclusions · p. 15
  • Zero-shot slicing extracts all nested submodels from the largest model during deployment.This supports a many-in-one model with constant deployment memory against the number of models in the family.Source: Conclusions · p. 15

Related papers

Loading 2511.16664v1…