Transferring Linear Features Across Language Models With Model Stitching

Alan Chen (Brown University) · Jack Merullo (Goodfire) · Alessandro Stolfo (ETH Zürich) · Ellie Pavlick (Brown University)
affine mappingsfeature transferabilityflops savingsfunctional featuresinitializationmodel size comparisonmodel transferprobesrepresentation learningrepresentation spacesresidual streamssemantic featuressparse autoencoderssteering vectorsstructural featurestraining efficiency

In this work, we demonstrate that affine mappings between residual streams of language models is a cheap way to effectively transfer represented features between models. We apply this technique to transfer the \textit{weights} of Sparse Autoencoders (SAEs) between models of different sizes to compare their representations. We find that small and large models learn highly similar representation spaces, which motivates training expensive components like SAEs on a smaller model and transferring to a larger model at a FLOPs savings. For example, using a small-to-large transferred SAE as initialization can lead to 50% cheaper training runs when training SAEs on larger models. Next, we show that transferred probes and steering vectors can effectively recover ground truth performance. Finally, we dive deeper into feature-level transferability, finding that semantic and structural features transfer noticeably differently while specific classes of functional features have their roles faithfully mapped. Overall, our findings illustrate similarities and differences in the linear representation spaces of small and large models and demonstrate a method for improving the training efficiency of SAEs.