Understanding Layer Patching in Model Size Interpolation

Zero-shot model size interpolation allows for the creation of intermediate-sized language models by patching layers of a distilled student model with blocks from a larger teacher model. While this method avoids additional training, the specific order in which layers are patched significantly impacts the performance of the resulting models. This paper investigates how to select the optimal patching order to maximize performance across the interpolation trajectory.

The authors formalize the selection of patching orders as a combinatorial optimization problem. They prove that finding the optimal sequence of layer replacements is equivalent to finding the shortest path in a directed acyclic graph (the interpolation graph), where nodes represent partially patched models and edges are weighted by the KL divergence between the teacher and the intermediate student model. Because the search space for the optimal path is factorial, the authors introduce KLPatch, a greedy algorithm that selects the next layer to patch based on minimizing the KL divergence to the teacher at each step. This reduces the computational complexity from exponential to polynomial time.

Through exhaustive experiments on smaller models (DistilBERT, DistilGPT2) and sampling-based studies on larger LLMs (Qwen, Pythia), the authors demonstrate that patching order is a critical design choice. They find that while simple heuristic strategies—such as patching from the last layer to the first—are often surprisingly effective, they are not always optimal. The proposed KLPatch algorithm consistently identifies near-optimal patching trajectories, outperforming simple sequential heuristics across various model families. The study also highlights that naive patching can sometimes degrade performance, emphasizing the need for a principled approach to layer selection.