Research map / Large language models / Instruction and fine-tuning
Flextron: Many-in-One Flexible Large Language Model
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Flextron converts a trained LLM (GPT-3 8B, Llama-2 7B) into a nested elastic network with input-adaptive routers, trained via sampling-based continued training and a surrogate-model-guided router training, to flexibly meet latency targets without fine-tuning.
Problem
Deploying LLMs under varying latency, memory, and compute budgets requires training multiple model variants, which is costly; existing elastic/nested approaches are limited in scope and need expensive training.
Approach
Nested elastic MLP and MHA layers with importance-based neuron/head sorting, sampled sub-network continued training, and learned static/input-adaptive routers trained through a surrogate model that predicts LM loss from router outputs, enabling MoE-like token routing with weight sharing.
Key results
- Flextron-Llama2-7B static-0.5x at 3.5B non-embedding params reaches 61.3% average vs. 69.4% for the full 6.5B Llama2-7B model, and dynamic-0.5x reaches 61.8%.
- Single elastic pretraining run uses 89.9B tokens for continued training (7.54% of GPT3 1.1T pretraining tokens, 4.50% of Llama2 2T tokens) plus 1.049B router-tuning tokens.
- Latency on A100 with TensorRT-LLM: Flextron-Llama2-7B drops from 9.64s (full) to 4.91s at 0.5x, and Flextron-8B from 10.43s to 5.48s.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use input-adaptive rather than static routing at 0.7× | Flextron-8B · Downstream tasks (Avg.) | average score: 61.3% -> 62.0% | Helped |
| Use input-adaptive rather than static routing at 0.6× | Flextron-8B · Downstream tasks (Avg.) | average score: 60.8% -> 61.2% | Helped |
| Use input-adaptive rather than static routing at 0.5× | Flextron-8B · Downstream tasks (Avg.) | average score: 59.1% -> 60.0% | Helped |
| Use input-adaptive rather than static routing at 0.7× | Flextron-Llama2-7B · Downstream tasks (Avg.) | average score: 62.8% -> 64.2% | Helped |
| Use input-adaptive rather than static routing at 0.6× | Flextron-Llama2-7B · Downstream tasks (Avg.) | average score: 62.1% -> 62.8% | Helped |
| Use input-adaptive rather than static routing at 0.5× | Flextron-Llama2-7B · Downstream tasks (Avg.) | average score: 61.3% -> 61.8% | Helped |
| Add importance-based sorting before selecting the first half of MHA heads (ablation) | GPT3-2B · Wikipedia (50% MHA heads) | perplexity: 184.5 -> 179.9 | Helped |
| Add importance-based sorting before selecting the first half of MLP neurons (ablation) | GPT3-2B · Wikipedia (50% MLP neurons) | perplexity: 1902.0 -> 193.6 | Helped |
| Reduce the active model budget from full to 0.7× | Flextron-8B · Downstream tasks (Avg.) | average score: 66.8% -> 62.0% | Hurt |
| Reduce the active model budget from full to 0.7× | Flextron-Llama2-7B · Downstream tasks (Avg.) | average score: 69.4% -> 64.2% | Hurt |
Surprising findings
- Embedding layer latency is a large share of the total for the multilingual 8B model (17.4% with a 320K vocabulary), limiting latency gains from elastic layers.
- Un-permuted models cut to 50% neurons yield Wikipedia perplexity above 1000, while importance-based permutation reduces it to 193.6 without training.
Limitations
- Gains relative to full model degrade at lower latency; sub-networks at 0.5x lose several points average accuracy versus the full model.
- Static MHA/MLP latency benefit is limited for the 8B multilingual model due to large embedding latency.
- Dynamic routing gives only modest gains over static routing (e.g., 61.8 vs 61.3 at 0.5x for Llama2-based model).
Experimental setup
Datasets: GPT3 pretraining data (53 languages and code) Nemotron-4 curated continued training dataset Wikipedia Books3 Arxiv
Benchmarks: ARC-easy LAMBADA PIQA WinoGrande MMLU HellaSwag
Metrics: accuracy validation language modeling loss perplexity latency
Baselines: Matformer GPT3-843M GPT3-8B Pythia OpenLLaMA-v2 Sheared-LLaMA Compresso LLM-Pruner SliceGPT LaCo
Builds on: Matformer Once-for-all Mixture-of-Experts
Closest papers by meaning
- Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.88)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.88)
- Star Elastic: Many-in-One Reasoning LLMs with Efficient Budget Control (ICML 2026 · similarity 0.88)
- Puzzle: Distillation-Based NAS for Inference-Optimized LLMs (ICML 2025 · similarity 0.88)
- Universal Model Routing for Efficient LLM Inference (ICLR 2026 · similarity 0.87)
- Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.