Research map / Large language models / Reasoning and chain of thought
Puzzle: Distillation-Based NAS for Inference-Optimized LLMs
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.
Puzzle applies hardware-aware decomposed NAS (blockwise local distillation, MIP-based architecture search, then global KD) to compress Llama-3.1-70B/3.3-70B into non-uniform 49-51B models that fit a single H100 with ~2.17x throughput while retaining ~98.4% accuracy.
Problem
High inference cost of large LLMs limits deployment; standard LLM architectures are uniform and not designed for inference efficiency on specific hardware, and conventional NAS or distillation is too costly at LLM scale.
Approach
Decompose NAS into per-block distillation trained in parallel (decoupled BLD over attention and FFN variants), score each block by KL divergence replace-1-block, then select a non-uniform architecture via mixed-integer programming under memory, latency and throughput constraints, followed by global KD with cosine and KLD losses.
Key results
- Nemotron-51B retains 98.4% of Llama-3.1-70B-Instruct benchmark accuracy and reaches up to 2.17x throughput speedup on a single H100 (FP8, TensorRT-LLM).
- Nemotron-49B-Base retains over 98% of parent RULER performance at 64K and about 95% at 128K after uptraining to 128K context.
- After only 2.9B GKD tokens, Nemotron-49B-Base recovers 98.47% of parent accuracy (MMLU/MT-Bench).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply Puzzle distillation to Llama-3.1-70B-Instruct | Llama-3.1-70B-Instruct · GSM8K Chat | accuracy: 81.58 -> 81.88 | Helped |
| Apply Puzzle distillation to Llama-3.1-70B-Instruct | Llama-3.1-70B-Instruct · MT-Bench | score: 8.93 -> 8.99 | Helped |
| Apply Puzzle distillation to Llama-3.1-70B-Instruct | Llama-3.1-70B-Instruct · Text Generation 128/1024 | throughput (tokens/sec/GPU): 2975 (TP4) -> 6472 (TP1) | Helped |
| Add cosine and KLD losses, without LM loss, during GKD (ablation) | Nemotron-51B after BLD · MMLU | accuracy: 78.39 -> 79.61 | Helped |
| Add cosine and KLD losses, without LM loss, during GKD (ablation) | Nemotron-51B after BLD · MT-Bench | score: 8.67 -> 8.87 | Helped |
| Add LM loss to cosine and KLD GKD losses (ablation) | Nemotron-51B after BLD · MT-Bench | score: 8.87 -> 8.66 | Hurt |
| Use coupled BLD refinement rather than decoupled BLD alone (ablation) | Puzzle on Llama-3.1-8B-Instruct · Llama-3.1-8B-Instruct derivative | accuracy: 73.10 -> 73.98 | Helped |
| Use Distillation Mix rather than Project Gutenberg for training, pruning, and block scoring (ablation) | Puzzle on Llama-3.1-70B-Instruct · MT-Bench | score: 7.98 -> 8.61 | Helped |
| Increase BLD training budget from 0.25B to 1.0B tokens (ablation) | Puzzle-optimized architecture · MT-Bench | score: 8.51 -> 8.98 | Helped |
| Score blocks with KL divergence rather than LM loss (ablation) | Puzzle on Llama-3.1-8B-Instruct · Accuracy vs. throughput frontier | accuracy | Helped |
| Score blocks with Half-MMLU accuracy rather than KL divergence (ablation) | Puzzle on Llama-3.1-8B-Instruct · Half-MMLU test set | accuracy: 64.94 -> 66.24 | Helped |
| Score blocks with Half-MMLU accuracy rather than KL divergence (ablation) | Puzzle on Llama-3.1-8B-Instruct · MT-Bench | score: 8.06 -> 7.57 | Hurt |
Surprising findings
- Nemotron-49B-Base after lightweight alignment surpasses its parent on Arena Hard (82.1 vs 71.70) and MT-Bench (9.21 vs 8.84).
- Adding an LM loss term to GKD harms downstream performance in most cases, with KLD plus cosine losses performing best.
Limitations
- Each parent layer is replaced by exactly one child block; no cross-layer block grouping or embedding-dimension changes.
- Long-context performance degrades beyond 64K for Nemotron-51B trained only to 8K.
- Inter-block incompatibility from BLD requires the GKD stage to fix error propagation.
Experimental setup
Datasets: Distillation Mix (FineWeb, Dolma, Buzz-V1.2; 224B tokens pool) Project Gutenberg
Benchmarks: MMLU MT-Bench Winogrande ARC Challenge HellaSwag GSM8K TruthfulQA XLSum HumanEval RULER Arena Hard MMLU-STEM
Metrics: accuracy preserved (%) throughput (tokens/s/GPU) speedup MMLU MT-Bench human preference
Baselines: Llama-3.1-70B-Instruct Llama-3.3-70B-Instruct Llama-3.2-3B-Instruct Llama-3.1-8B-Instruct greedy search random architectures
Builds on: LANA DONNA DNA grouped-query attention Minitron
Closest papers by meaning
- Search for Efficient Large Language Models (NeurIPS 2024 · similarity 0.90)
- REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving (NeurIPS 2025 · similarity 0.90)
- xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference (ICML 2025 · similarity 0.90)
- EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.89)
- Understanding and Mitigating Numerical Sources of Nondeterminism in LLM Inference (NeurIPS 2025 · similarity 0.89)
- Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (ICML 2024 · similarity 0.89)
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.