Papers
Topics
Authors
Recent
Search
2000 character limit reached

Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control

Published 8 May 2026 in cs.LG | (2605.07182v1)

Abstract: Training a family of LLMs, either from scratch or via iterative compression, is prohibitively expensive and inefficient, requiring separate training runs for each model in the family. In this paper, we introduce Star Elastic, a novel LLM post-training method that adds N nested submodels to a given parent reasoning model using the compute of one run (N-fold savings) via a single post-training job. Beyond reducing training costs, Star Elastic also addresses a fundamental limitation of efficient reasoning: the rigidity of static architectures, which forces the allocation of constant resources regardless of token difficulty. By unlocking elastic budget control, Star Elastic enables a novel inference scheme that uses different submodels for each reasoning phase (thinking and answering). Star Elastic supports (1) nesting along the SSM, embedding channel, MoE, and FFN axes, (2) learning nested submodels via an end-to-end trainable router, and (3) curriculum-based knowledge distillation. Building on the Nemotron Elastic framework, we apply Star Elastic to the NVIDIA Nemotron Nano models, with a particular focus on hybrid Mixture-of-Experts (MoE) architectures: from Nemotron Nano v3 (30B/3.6A), we generate 23B (2.8A) and 12B (2.0A) variants with 160B training tokens. All nested models match or outperform independently trained baselines of comparable size and achieve a 360x reduction versus pretraining from scratch and a 7x reduction over state-of-the-art compression. Crucially, elastic budget control advances the accuracy-latency Pareto frontier, achieving up to 16% higher accuracy and 1.9x lower latency via dynamic per-phase model selection. We further extend Star Elastic to quantized regimes via Quantization-Aware Distillation (QAD), producing nested NVFP4 and FP8 elastic checkpoints that preserve zero-shot slicing while delivering smaller deployment footprints.

Summary

  • The paper introduces an elastic post-training framework that embeds nested submodels, enabling zero-shot extraction and dynamic phase-wise budget control.
  • The methodology leverages importance estimation, a learnable router with Gumbel-Softmax optimization, and Nested Quantization-Aware Distillation to achieve up to 16% higher accuracy and 1.9× lower latency.
  • The framework significantly reduces training and deployment costs, offering up to 360× lower token costs and halved storage requirements compared to separate model training.

Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control

Motivation and Contributions

The "Star Elastic" framework addresses the prohibitively high computational and storage costs associated with training and deploying families of reasoning LLMs, especially when each model variant is independently trained. Existing approaches for model compression primarily utilize knowledge distillation or structured pruning, but these still require linear scaling in token costs per variant and lack unified, nested deployment. Star Elastic introduces an elastic post-training paradigm that embeds NN nested submodels within a single parent, enabling zero-shot extraction and deployment of multiple model sizes without additional fine-tuning or retraining. This extends the Nemotron Elastic framework to hybrid Mamba--Attention--MoE architectures, unlocking elasticity across embedding channel, SSM, MoE, and FFN axes.

Core contributions include:

  • The first elastic post-training solution for hybrid Mamba--Transformer--MoE reasoning LLMs.
  • Dynamic, phase-wise elastic budget control during inference, yielding up to 16% higher accuracy and 1.9×\times lower latency.
  • End-to-end learnable router architecture, optimized by knowledge distillation, that efficiently determines nested submodel structures, supporting heterogeneous per-layer compression (e.g., FFN and expert selection).
  • Significant training and deployment cost reduction: 360×360\times less than pretraining from scratch and 7×7\times compared to state-of-the-art compression.
  • Nested Quantization-Aware Distillation (QAD) for elastic checkpoints in FP8 and NVFP4, preserving zero-shot slicing and further reducing deployment footprint.

This paradigm dismantles the rigid resource allocation in static architectures and systematically advances accuracy--latency trade-offs via dynamic, per-phase model selection.

Methodology: Elastic Model Construction and Training

Star Elastic's pipeline comprises three stages:

  1. Importance Estimation and Model Preparation: Model components (embedding channels, Mamba heads, attention heads, MoE experts, FFN channels) are ranked by activation-based importance scores, extending group-aware constraints for SSM structures. For MoE, Router-Weighted Expert Activation Pruning (REAP) ensures principled expert selection, considering both activation frequency and output magnitude.
  2. Elastic Formulation: The architecture embeds nested, weight-shared submodels that vary in both width (embedding dimension, head/expert/channel counts) and depth (layer count), forming strictly nested hierarchies via importance rankings. This enables resource-aware extraction along width and depth axes, supporting homogeneous and per-layer heterogeneous configurations.
  3. Elastic Training: A learnable router, parameterized by budget specification, produces differentiable masks for each axis. Mask selections are optimized via Gumbel-Softmax and supervised through knowledge distillation from the non-elastified parent. A two-stage curriculum-based sampling strategy is adopted: uniform budget sampling at short context, followed by weighted non-uniform sampling for extended context. The loss combines knowledge distillation and router resource-matching, enabling end-to-end sub-network optimization. Figure 1

    Figure 1: Overview of the Star Elastic pipeline showing joint optimization and zero-shot deployment of nested submodels from a single checkpoint.

This unified post-training method eliminates the need for separate compression/distillation runs per variant, regularizes smaller submodels via training signals from the full budget, and supports immediate extraction of any nested variant.

Elastic Budget Control

Traditional budget control in LLM inference applies a fixed architecture across all reasoning phases, missing opportunities for differentiated resource allocation. Star Elastic introduces elastic budget control—allocating varying model sizes for reasoning ("thinking") and answer synthesis ("answering") phases. Phase-specific nesting enables lower-cost long-horizon reasoning with smaller models and high-fidelity answer generation with larger models.

The optimal configuration empirically is to use a smaller model for the thinking phase and a larger one for the answering phase (e.g., $23$B for reasoning, $30$B for answer synthesis), aligning model capacity with phase-specific requirements. Star Elastic ensures cache compatibility across nested variants, facilitating cache transplantation and potentially eliminating recomputation overhead. Figure 2

Figure 2

Figure 2: Left: Star Elastic achieves superior accuracy compared to parent and Qwen on multiple reasoning benchmarks. Right: Hybrid elastic budget control dominates the accuracy--speed Pareto frontier relative to static control.

Figure 3

Figure 3: Elastic budget control configurations across computational budgets, with the MS→ML\mathcal{M}_{S} \to \mathcal{M}_{L} configuration offering the best accuracy--latency tradeoff.

Empirical Results

Accuracy and Efficiency

Experiments focus on Nemotron Nano v3 (30B/3.6A), compressing to nested 23B/2.8A and 12B/2.0A variants with 160B tokens. All nested models match or outperform independently trained baselines at comparable size, with strong results on reasoning benchmarks: AIME-2025, GPQA, LiveCodeBench v5, MMLU-Pro, IFBench, and Tau Bench.

Elastic budget control advances the Pareto frontier: up to 16% higher accuracy and 1.9×1.9\times lower latency through dynamic phase allocation. Nested models enable higher batch throughput and fit on lower-memory hardware, further reducing deployment friction. Figure 4

Figure 4: Per-benchmark Pareto frontiers showing elastic budget control configurations outperform static baselines across accuracy--speed.

Figure 5

Figure 5: Detailed configuration breakdown; $23$B-$30$B dominates at high accuracy/moderate latency, ×\times0B-×\times1B and ×\times2B-×\times3B optimal for low-latency.

Training and Deployment Cost

Star Elastic achieves ×\times4 reduction in token cost (relative to pretraining) and ×\times5 over sequential compression, delivering all variants from a single training run. Storage requirements scale only with the largest model due to weight sharing; for Nano v3, an elastic checkpoint requires less than half the memory vs. separate models.

Quantized Elastic Models

The framework extends elastification to FP8 and NVFP4 formats. For FP8, post-training quantization maintains accuracy via simple quantization. For NVFP4, QAD is necessary for accuracy recovery. Nested QAD preserves the elastic structure, facilitating zero-shot slicing of quantized variants. Quantized elastic checkpoints achieve 3.4×\times6 throughput improvement and fit larger models on lower-end GPUs, drastically reducing memory and bandwidth requirements without compromising performance. Figure 6

Figure 6: Throughput (tokens/s) of Star Elastic variants under vLLM in BF16, FP8, NVFP4, across RTX 5080/5090/Pro 6000; NVFP4 enables deployment under tighter device constraints.

Star Elastic builds upon the Nemotron Elastic post-training framework, extending it from Mamba--Attention to MoE-hybrid architectures. Prior work (e.g., Flextron, MatFormer) established nested weight sharing and zero-shot slicing in transformer-based models. Structured pruning and knowledge distillation (e.g., Minitron, REAP, group-aware SSM pruning) are outperformed by Star Elastic's single-run, nested approach.

Quantization-aware training and post-training quantization (PTQ/QAD) are widely used for dense transformers, but Star Elastic demonstrates their efficacy in preserving elastic structure across quantized regimes. The phase-wise dynamic allocation of nested model sizes introduces a novel axis of efficiency not previously explored.

Implications and Future Directions

Star Elastic sets a new paradigm in training and deploying model families, coupling flexibility, efficiency, and performance. The practical impact is immediate: rapid deployment of nested models, real-time adaptation to inference resource constraints, and substantial reduction in both cost and device footprint. Theoretical implications include new avenues for dynamic envelope modeling and joint optimization of nested substructures under distillation. Future directions span ultra-low-budget compression, automated task/domain-specific routing, and advances in cache-state transplantation. Further integration with context-aware adaptation and heterogeneous configuration will strengthen both generalization and resource utilization.

Conclusion

Star Elastic delivers a unified post-training solution for constructing elastic, reasoning-capable LLM families under hybrid Mamba--Transformer--MoE designs. The methodology efficiently scales multiple nested variants in a single run, dynamically calibrates inference resources via phase-wise budget control, and produces quantized checkpoints for device-constrained deployment. The nested, many-in-one paradigm shrinks both training and deployment requirements, improves accuracy--latency trade-offs, and generalizes across precision regimes. Star Elastic represents a substantial step towards practical, adaptive LLM deployment for complex reasoning workloads (2605.07182).

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 3 tweets with 13 likes about this paper.