💻 computer science

Where to Begin: Efficient Pretraining via Subnetwork Selection and Distillation

This paper presents an efficient framework for pretraining Small Language Models (SLMs) that combines evolutionary search to identify structurally sparse subnetwork initializations with knowledge distillation from larger models, achieving performance comparable to standard baselines while significantly reducing computational costs.

Original authors: Arjun Krishnakumar, Rhea Sanjay Sukthanker, Hannan Javed Mahadik, Gabriela Kadlecová, Vladyslav Moroshan, Timur Carstensen, Frank Hutter, Aaron Klein

Published 2026-01-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Arjun Krishnakumar, Rhea Sanjay Sukthanker, Hannan Javed Mahadik, Gabriela Kadlecová, Vladyslav Moroshan, Timur Carstensen, Frank Hutter, Aaron Klein

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

The Big Problem: Building a Giant Library is Expensive

Imagine you want to build a massive library of knowledge (a Large Language Model, or LLM) that can answer any question. Traditionally, to build this library, you have to hire millions of people (parameters), give them all blank notebooks, and make them read the entire internet from scratch. This takes an enormous amount of time, money, and electricity.

While "Small Language Models" (SLMs) are cheaper to build because they are smaller, they are still too expensive for most regular researchers or small companies to build from scratch. They need a shortcut.

The Solution: The "Whittle" Framework

The authors of this paper propose a new way to build these smaller libraries. They call their framework Whittle. Instead of hiring new people and starting with blank notebooks, they take an existing, giant, highly knowledgeable library (a large "Teacher" model) and try to find the perfect, smaller team inside it that can do the job just as well.

They do this using three main tricks:

1. Finding the "Golden Sub-Team" (Sub-Network Selection)

Imagine the giant library is a massive orchestra with thousands of musicians. You don't need the whole orchestra to play a specific song; you just need the right section.

The Old Way: Usually, if you want a small model, you just randomly pick a few musicians and hope they can play.
The New Way: The authors use a smart search tool (called Evolutionary Search) to look through the giant orchestra and find the specific group of musicians who already know the song perfectly. They don't just pick random people; they look for the specific "sub-network" of neurons that are already doing the heavy lifting.
The Result: They found that these pre-selected "sub-teams" are much better at learning than random groups of the same size.

2. The "Smart Copycat" Method (Knowledge Distillation)

Once they have their small sub-team, they don't let them learn from the internet alone. Instead, they put them in a classroom with the original giant library (the Teacher).

How it works: The Teacher doesn't just say "The answer is A." The Teacher says, "The answer is A, but it's very likely to be A, slightly likely to be B, and unlikely to be C." This gives the small team a much richer understanding of the world.
The Analogy: It's like a master chef teaching an apprentice. Instead of just showing the final dish, the master explains the subtle flavors and techniques. The apprentice learns faster and makes fewer mistakes.

3. The "Search Space" (Trying Different Combinations)

The authors realized that not all "sub-teams" are created equal. Sometimes you need to cut down the number of layers (like removing floors from a building), and sometimes you need to cut down the width (like removing columns).

They tested four different ways to cut the model:
- Coarse: Cutting big chunks (like removing entire floors).
- Fine-grained: Cutting tiny bits (like removing specific bricks).
- Uniform: Making the whole model smaller evenly.
- Layer-wise: Making different parts of the model smaller in different ways.
The Discovery: They found that for very small models, being "fine-grained" (cutting tiny bits) worked best. But for larger small-models, being "coarse" (cutting big chunks) was actually better.

The Results: Doing More with Less

The paper claims that by using this method, they can build small models that perform just as well as models trained from scratch, but with a massive reduction in cost:

Speed: Their best model reached the same level of intelligence as a standard small model but required 5.16 times fewer calculations (FLOPs) for a specific training size.
Efficiency: It's like building a house that is just as sturdy as a mansion, but you only used 20% of the bricks and labor.
Open Source: Unlike some other methods used by big tech companies that are secret, the authors released all their code and tools (the "Whittle" library) so anyone can use this method.

Summary

Think of this paper as a guide on how to recycle a giant, expensive AI model into a smaller, cheaper, and highly efficient one. Instead of building a new car from scratch, they take a luxury car, carefully remove the parts you don't need, and tune the remaining engine so it runs perfectly on less fuel, all while teaching it using the knowledge of the original luxury car.

Technical Summary: WHERE TO BEGIN: EFFICIENT PRETRAINING VIA SUB-NETWORK SELECTION AND DISTILLATION

Problem Statement

While Large Language Models (LLMs) have achieved state-of-the-art performance through massive scale, their training and deployment costs are prohibitive for many research groups and resource-constrained environments. Small Language Models (SLMs) offer a more accessible alternative but still require substantial computational resources for pretraining. Existing strategies to reduce these costs, such as knowledge distillation from open-weight LLMs or pruning, often rely on closed-source implementations (e.g., Llama 3.2 variants, Minitron) or lack a systematic, reproducible framework for discovering optimal sub-network initializations. Furthermore, the mechanisms by which teacher models improve SLM pretraining efficiency remain poorly understood.

Methodology

The authors propose a three-component framework to make SLM pretraining substantially more efficient:

Sub-Network Initialization: Instead of random initialization, the framework extracts structurally sparse sub-networks from a larger, pretrained teacher model. The authors define four distinct search spaces to explore how weights are selected:
- Coarse vs. Fine-grained: Coarse selection takes the first $n$ entries of a component (e.g., first $n$ attention heads), while fine-grained selection samples arbitrary indices.
- Uniform vs. Layer-wise: Uniform applies the same configuration across all layers, whereas layer-wise allows each layer to have independent configurations.
- This results in four search spaces: Coarse Uniform, Coarse Layer-wise, Fine-grained Uniform, and Fine-grained Layer-wise.
Constrained Evolutionary Search: To identify high-quality sub-networks within specific parameter budgets, the authors employ an evolutionary search algorithm.
- The parameter space is partitioned into disjoint "bins" (e.g., 385M–426M, 961M–1.06B) to ensure balanced coverage of model sizes.
- Within each bin, a population of candidate architectures is evolved using mutation and crossover operators.
- Constraint Enforcement: Rejection sampling is used to ensure candidate architectures strictly adhere to the parameter count of the target bin.
- Fitness Function: Candidates are evaluated based on validation perplexity on a small token budget (2B tokens) to guide the search toward optimal initializations.
Knowledge Distillation: The selected sub-networks are pretrained using knowledge distillation from the same teacher model they were extracted from.
- The loss function combines standard cross-entropy with forward KL divergence between teacher and student logit distributions.
- To improve computational efficiency, the authors utilize Top-k Logit Distillation, truncating the teacher distribution to its $k$ most salient outputs (e.g., $k=1024$ ).

Key Contributions

Systematic Open-Source Framework: The authors introduce Whittle, a fully open-source library that enables the extraction of sub-networks from Hugging Face models, supports flexible search space design, and facilitates the entire pipeline from search to pretraining and distillation.
Sub-Network Initialization Strategy: They demonstrate that extracting sub-networks from pretrained teachers yields significantly better starting points than random initialization. For example, smaller variants (~410M parameters) require 1.71× fewer FLOPs to match the validation perplexity of a comparable Pythia model trained from scratch.
Comprehensive Analysis: The paper provides the first systematic comparison of sub-network initialization under knowledge distillation versus standard cross-entropy training, analyzing the impact of teacher size, search space granularity, and distillation hyperparameters.
Empirical Validation: Experiments on the Pythia family (6.9B and 12B teachers) show that the best models, discovered via evolutionary search and initialized with teacher weights, match the performance of comparable Pythia SLMs while requiring significantly fewer floating-point operations (FLOPs).

Results

Efficiency Gains:
- For a 410M parameter model (Bin 1), the method achieves the same validation perplexity with 1.71× fewer FLOPs compared to a standard Pythia-410M pretraining.
- For a 1B parameter model (Bin 2), the savings increase to 1.75×.
- For a 2.8B parameter model (Bin 3), the method requires 5.16× fewer FLOPs to match the validation perplexity of a comparable Pythia-2.8B model trained with a 10B token budget.
- At a 100B token budget, the 2.8B variant still achieves a 1.26× FLOP reduction while attaining a lower final validation perplexity than the baseline.
Search Space Findings:
- Granularity: The optimal search space depends on the model size. For larger variants (Bin 3), the least granular space (Coarse Uniform) yields the best results. For smaller variants (Bin 1), more granular spaces like Fine-grained Uniform perform better.
- Distillation: Distillation consistently improves perplexity over standard pretraining. Using the full logit distribution generally yields lower perplexity than Top-k truncation, though Top-k is used for efficiency in the main experiments.
- Search Metric: Optimizing directly for perplexity during the search phase outperforms proxy metrics like weight magnitude or activation-based importance scores.
Downstream Performance: Sub-networks initialized from the supernet consistently outperform both randomly initialized versions of the same architecture and original Pythia models of comparable size across common-sense and question-answering benchmarks (COPA, MMLU, HellaSwag, etc.). Notably, sub-networks extracted from the smaller 6.9B teacher often outperformed those from the 12B teacher.

Significance and Claims

The paper claims to present the first systematic, open-source study of warm-starting student models from larger teachers for pretraining. The significance lies in providing a practical and reproducible path toward cost-efficient SLM development. By combining evolutionary search for sub-network discovery with knowledge distillation, the authors demonstrate that SLM pretraining can be accelerated by up to 9.2× (when excluding search costs) or 5.16× (including search costs) compared to baseline SLMs of similar size.

The authors emphasize that their approach sidesteps the expensive supernet pretraining required by traditional Neural Architecture Search (NAS) methods by leveraging existing open-weight LLMs. They conclude that while the method is highly effective, future work is needed to derive scaling laws for these initialization strategies and to investigate the impact of teacher model choice (e.g., multilingual vs. monolingual) on student performance. The authors explicitly state that their findings offer guidelines for compute-optimal SLM pretraining across different scales.