🤖 machine learning

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

Star Elastic is a novel post-training method that efficiently generates multiple nested reasoning submodels from a single parent model via one training run, enabling dynamic, phase-specific budget control that significantly reduces training costs while improving the accuracy-latency trade-off compared to independent training or compression baselines.

Original authors: Ali Taghibakhshi, Ruisi Cai, Saurav Muralidharan, Sharath Turuvekere Sreenivas, Aditya Vavre, Ameya Sunil Mahabaleshwarkar, Bilal Kartal, Sheldon Liang, Marcin Chochowski, Zijia Chen, Akhiad Bercovich

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ali Taghibakhshi, Ruisi Cai, Saurav Muralidharan, Sharath Turuvekere Sreenivas, Aditya Vavre, Ameya Sunil Mahabaleshwarkar, Bilal Kartal, Sheldon Liang, Marcin Chochowski, Zijia Chen, Akhiad Bercovich, Ran Zilberstein, Ran El-Yaniv, Yonatan Geifman, Daniel Korzekwa, Yoshi Suhara, Oluwatobi Olabiyi, Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you own a massive, high-end kitchen designed to cook gourmet meals for a crowd. This kitchen is your Large Language Model (LLM).

Traditionally, if you wanted to serve a small family dinner, a medium-sized party, and a huge banquet, you would need to build three completely separate kitchens. You'd have to buy three sets of ovens, hire three different teams of chefs, and pay for three separate construction projects. This is incredibly expensive and wasteful.

Star Elastic is a new invention that changes the game. Instead of building three kitchens, it builds one super-flexible kitchen that can instantly transform itself to fit any need.

Here is how it works, broken down into simple concepts:

1. The "One Kitchen, Many Sizes" Trick

Usually, to get a smaller, cheaper AI model, researchers have to train a giant one, then try to "shrink" it by cutting parts off (like pruning a tree). This is slow, expensive, and often results in a model that isn't as smart as if it had been trained from scratch.

Star Elastic does something different. It takes one giant "Parent" model (the Nemotron Nano v3) and teaches it a special skill: how to be many different sizes at once.

Think of it like a Russian Nesting Doll (Matryoshka). Inside the big 30-billion-parameter doll, there is a perfect 23-billion-parameter doll, and inside that, a perfect 12-billion-parameter doll.
They all live in the same "house" (the same computer file). You don't need to download three different files; you just open the big one and tell it, "I only need the small version today."

2. The "Smart Router" (The Kitchen Manager)

How does the model know which parts to use? It uses a Learnable Router.

Imagine a kitchen manager who looks at your order.
If you ask for a simple salad (a simple question), the manager says, "Okay, let's just use the small blender and the basic knife set."
If you ask for a complex soufflé (a hard math problem), the manager says, "We need the big oven, the heavy-duty mixer, and all the chefs!"
The paper shows that this manager is trained to know exactly which parts of the "kitchen" are the most important. It keeps the best tools for the small versions and only adds the extra heavy-duty tools when the big version is needed.

3. The "Thinking vs. Answering" Strategy

This is the paper's most clever trick, called Elastic Budget Control.

When an AI solves a hard problem, it usually does two things: Thinking (reasoning through the steps) and Answering (writing the final result).
Old Way: The AI uses the same "brain size" for both thinking and answering. It's like using a giant, fuel-hungry truck to drive to the grocery store and then to the post office, even if the post office is just down the street.
Star Elastic Way: The AI can switch gears!
- Phase 1 (Thinking): It uses the smaller, faster model to brainstorm and think through the problem. This is cheap and fast.
- Phase 2 (Answering): Once the thinking is done, it switches to the larger, smarter model just to write the final answer. This ensures the answer is perfect.
The Result: You get the accuracy of the giant brain but the speed and cost of the small brain. The paper claims this can make the AI 16% more accurate and 1.9 times faster than using a single fixed size.

4. The "Magic Slicing" (Zero-Shot Extraction)

Usually, if you want a smaller model, you have to train it for months. With Star Elastic, the "slicing" happens instantly.

Because the model was trained to be flexible from the start, you can "slice" off the small or medium versions zero-shot.
This means you don't need to retrain them. You just take the big file, apply the "cut," and boom—you have a working, high-quality small model ready to go immediately.
The paper claims this saves 360 times the training cost compared to building models from scratch and 7 times the cost of previous compression methods.

5. The "Tiny Suitcase" (Quantization)

Finally, the paper talks about making these models even smaller for phones or small computers.

They use a technique called Quantization-Aware Distillation. Imagine taking a heavy, high-definition painting and turning it into a digital file that takes up very little space but still looks perfect.
They created versions of their models that fit into 4-bit or 8-bit formats (very small digital footprints).
Even in these tiny formats, the "Russian Nesting Doll" trick still works. You can still slice off the small, medium, and large versions from one tiny file.

Summary of the Wins

Cost: You train once, get many models. It's like buying one ticket to a theme park that lets you ride every rollercoaster, instead of buying a separate ticket for each ride.
Speed: By using a small brain for thinking and a big brain for answering, you save time and money.
Storage: You only need to store one file to have access to three different model sizes.

In short, Star Elastic stops us from building separate, rigid AI models for every job. Instead, it builds one chameleon-like AI that can instantly change its size and power to fit the task, saving massive amounts of money and time while actually getting smarter.

Technical Summary: Star Elastic

Problem Statement

Training families of Large Language Models (LLMs) with varying parameter counts is currently prohibitively expensive and inefficient. Traditional approaches require separate training runs for each model size, multiplying compute and storage costs. While recent compression methods (e.g., structured pruning and knowledge distillation) reduce costs by deriving smaller models from a large parent, they still demand hundreds of billions of tokens per compressed model and often fail to support heterogeneous architectures.

Furthermore, existing efficient reasoning frameworks suffer from architectural rigidity. They allocate static resources regardless of token difficulty, forcing the use of a single model size for both the "thinking" (reasoning) and "answering" phases of inference. This leads to computational inefficiencies when phases have disparate complexity. Finally, current elastic nested models (e.g., Nemotron Elastic) have been limited to Mamba-Attention architectures, lacking support for the increasingly popular hybrid Mamba-Attention-Mixture-of-Experts (MoE) designs.

Methodology: Star Elastic

Star Elastic is a post-training method designed to generate $N$ nested submodels from a single parent reasoning model using the compute of one training run. It targets hybrid Mamba-Attention-MoE architectures and introduces elastic budget control for dynamic inference.

1. Elastic Model Construction

The pipeline converts an existing LLM into an elastic model through three stages:

Importance Estimation: Components are ranked by importance to guide pruning.
- Width: Activation-based scoring is extended to hybrid architectures. Scores are computed for embedding channels, FFN/MoE channels, Mamba heads/channels, and attention heads. For MoE layers, Router-Weighted Expert Activation Pruning (REAP) is used, measuring an expert's contribution to output magnitude based on routing weights and activation norms.
- Depth: Layer importance is estimated iteratively using normalized Mean Squared Error (MSE) between full-model logits and logits with specific layers removed.
Elastic Formulation: The model supports nesting along width and depth axes.
- Width: Subnetworks select subsets of components (embedding dimensions, heads, experts, FFN channels) based on the importance ranking. Heterogeneous configurations allow independent selection per layer.
- Depth: A binary selection vector determines active layers, removing less important ones first.
Elastic Training: A learnable router maps user-provided budget specifications to optimal architectural configurations.
- The router uses Gumbel-Softmax to produce differentiable masks for training.
- Loss Function: Combines Knowledge Distillation (KD) from the non-elastified parent with a router loss that aligns selected configurations with hardware constraints (latency, memory).
- Two-Stage Curriculum:
  1. Stage 1 (Short Context): Uniform budget sampling across all target sizes to prevent gradient competition.
  2. Stage 2 (Extended Context): Non-uniform sampling prioritizing larger budgets (e.g., 50% for the full model) to ensure reasoning capabilities are preserved in larger variants while smaller models adapt to long-context reasoning chains.

2. Elastic Budget Control

Star Elastic enables dynamic model selection during inference phases:

Thinking Phase: Uses a smaller submodel ( $\mathcal{M}_S$ ) to generate extensive reasoning traces with minimal overhead.
Answering Phase: Switches to a larger submodel ( $\mathcal{M}_L$ ) to synthesize the final answer with higher fidelity.
Cache Compatibility: The architecture preserves Mamba and attention layer structures across nested models, allowing KV and SSM cache states to be transplanted between submodels (though current framework support for seamless reuse is limited, the architecture enables it).

3. Quantization-Aware Distillation (QAD)

To support deployment in quantized regimes (FP8 and NVFP4), Star Elastic applies QAD.

FP8: Achieved via Post-Training Quantization (PTQ) on the BF16 elastic checkpoint.
NVFP4: Requires a short QAD phase where a quantized student is distilled from the BF16 teacher using the same nested mask hierarchy. This recovers accuracy lost during PTQ while preserving the zero-shot slicing capability.

Key Contributions

First Elastic Post-Training for Hybrid MoE: Extends the Nemotron Elastic framework from Mamba-Attention to hybrid Mamba-Attention-MoE architectures, introducing REAP for expert ranking and heterogeneous per-layer selection.
Elastic Budget Control: Introduces a dynamic inference scheme that allocates different submodels to thinking and answering phases, optimizing the accuracy-latency Pareto frontier.
Learnable Router: A router that automatically determines nested architectures via end-to-end training, optimizing for both accuracy and resource constraints without decoupling configuration from training objectives.
Significant Cost Reduction: Demonstrates a 360× reduction in training tokens compared to pretraining from scratch and a 7× reduction over state-of-the-art compression methods for generating model families.
Quantized Elasticity: Produces nested NVFP4 and FP8 checkpoints that preserve zero-shot slicing, enabling deployment of multiple model sizes from a single quantized file.

Experimental Results

The method was applied to Nemotron Nano v3 (30B/3.6A), generating 23B (2.8A) and 12B (2.0A) variants using 160B training tokens.

Accuracy: All nested models match or outperform independently trained baselines of comparable size. For example, the 23B Elastic variant outperformed the independently trained Qwen3-30B-A3B on LiveCodeBench and IFBench.
Efficiency:
- Training: 360× token reduction vs. pretraining; 7× reduction vs. sequential compression.
- Deployment Memory: Storing the 12B, 23B, and 30B variants in a single elastic checkpoint requires 58.9 GB (BF16), compared to 126.1 GB for separate checkpoints.
- Throughput: Nested models achieve up to 2.4× speedup in throughput on H100 GPUs.
Elastic Budget Control: The $\mathcal{M}_S \to \mathcal{M}_L$ strategy (Small Thinking $\to$ Large Answering) achieves up to 16% higher accuracy and 1.9× lower latency compared to static budget control, significantly advancing the accuracy-speed Pareto frontier.
Quantization: FP8 and NVFP4 variants recover 97–100% of the BF16 accuracy while enabling deployment on consumer-grade GPUs (e.g., RTX 5080) where BF16 variants would OOM.

Significance and Claims

The paper claims that Star Elastic fundamentally shifts the paradigm for deploying reasoning LLMs by:

Eliminating Redundant Training: Making the creation of model families a single-pass process rather than a series of expensive, independent runs.
Unlocking Dynamic Inference: Moving beyond static architectures to allow resource allocation that matches the specific computational demands of reasoning versus answer synthesis.
Bridging the Gap to Quantization: Proving that elastic, nested structures can be preserved in low-precision formats (NVFP4/FP8), enabling high-performance reasoning on resource-constrained hardware without sacrificing the flexibility of multiple model sizes.

The authors position this work as a step toward "many-in-one" reasoning LLMs that are both state-of-the-art in accuracy and highly cost-effective for diverse deployment scenarios.