Transferring Linear Features Across Language Models With Model Stitching

TL;DR

This paper introduces a method for transferring features between language models using affine mappings, enabling efficient training and better understanding of model representations.

Contribution

It presents a novel technique for feature transfer via affine mappings, reducing training costs and analyzing representation similarities across model sizes.

Findings

Small and large models learn similar representation spaces.

Transferring features can reduce SAE training costs by 50%.

Semantic and structural features transfer differently.

Abstract

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 weights of Sparse Autoencoders (SAEs) between models of different sizes to compare their representations. We find that small and large models learn similar representation spaces, which motivates training expensive components like SAEs on a smaller model and transferring to a larger model at a FLOPs savings. In particular, 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.