Where to Begin: Efficient Pretraining via Subnetwork Selection and Distillation

TL;DR

This work tackles the challenge of efficiently pretraining Small Language Models by combining sub-network initializations extracted from pretrained teachers, an evolutionary search that enforces parameter-bin constraints to find high-quality starting points, and knowledge distillation from larger models to accelerate convergence. The methodology yields substantial gains in token efficiency, with best configurations matching baseline perplexity while using up to $9.2\times$ fewer pretraining tokens, and improves downstream performance when coupled with distillation. The authors provide an open-source Whittle library that enables sub-network extraction, search, pretraining, and distillation, offering practical guidance for compute-efficient SLM development at scale. The study also analyzes how search space granularity, distillation strategy, and teacher size affect performance, delivering actionable recommendations across model scales and tasks.

Abstract

Small Language models (SLMs) offer an efficient and accessible alternative to Large Language Models (LLMs), delivering strong performance while using far fewer resources. We introduce a simple and effective framework for pretraining SLMs that brings together three complementary ideas. First, we identify structurally sparse sub-network initializations that consistently outperform randomly initialized models of similar size under the same compute budget. Second, we use evolutionary search to automatically discover high-quality sub-network initializations, providing better starting points for pretraining. Third, we apply knowledge distillation from larger teacher models to speed up training and improve generalization. Together, these components make SLM pretraining substantially more efficient: our best model, discovered using evolutionary search and initialized with LLM weights, matches the validation perplexity of a comparable Pythia SLM while requiring 9.2x fewer pretraining tokens. We release all code and models at https://github.com/whittle-org/whittle/, offering a practical and reproducible path toward cost-efficient small language model development at scale.

Figures (16)

• Figure 1: Left: Initialization schemes — random weights, sub-network from a pretrained teacher, and our evolutionary search–based sub-network. Right: The same teacher is used for knowledge distillation to train the student.
• Figure 2: Overview of our search spaces and search strategy
• Figure 3: Best perplexity after evolutionary search based on perplexity for different search spaces.
• Figure 4: Pretraining Validation perplexity of the best sub-networks from each search space and the bin-center Pythia models (410M, 1B, 2.8B), all trained for 10B tokens with cross-entropy loss. Sub-networks are extracted from the Pythia-6.9B base model.
• Figure 5: Distillation: Comparison of validation perplexity for models trained with distillation loss v/s cross entropy loss. All sub-networks are extracted from Pythia-6.9B as a base model and trained for 10B tokens