Accelerating Sparse Autoencoder Training via Layer-Wise Transfer Learning in Large Language Models (2024)

Abstract: Sparse AutoEncoders (SAEs) have gained popularity as a tool for enhancing the interpretability of Large Language Models (LLMs). However, training SAEs can be computationally intensive, especially as model complexity grows. In this study, the potential of transfer learning to accelerate SAEs training is explored by capitalizing on the shared representations found across adjacent layers of LLMs. Our experimental results demonstrate that fine-tuning SAEs using pre-trained models from nearby layers not only maintains but often improves the quality of learned representations, while significantly accelerating convergence. These findings indicate that the strategic reuse of pretrained SAEs is a promising approach, particularly in settings where computational resources are constrained.

## introduction transformer-based models have become ubiquitous in a large variety of different application fields @xcite @xcite @xcite . given their tremendous impact on society, concerns about their interpretability have been raised by various stakeholders @xcite . mechanistic interpretability (mi) @xcite , seeks to reverse-engineer how neural networks, and in particular llms, generate outputs by uncovering the circuits they have learned during training, stored inside their parameters, and executed during a forward pass @xcite @xcite . a promising interpretability technique is dictionary learning @xcite @xcite which seeks to capture interpretable and editable features within the internal layers of llms. this method implies training sparse autoencoders (saes) to reconstruct the model's activations using sparse learned features. however, training saes is computationally intensive, particularly when applied across multiple layers in deep networks. this computational burden poses a significant barrier to their widespread application, especially in resource-constrained environments where the cost of training from scratch is prohibitive. recent research has highlighted the potential of transfer learning as a strategy to mitigate these challenges @xcite . in particular, it has been shown in @xcite that adjacent layers in llms are often redundant, suggesting that the knowledge encoded in one layer is also present in neighboring ones and that it can effectively be transferred. this observation forms the basis of our investigation: we hypothesize that saes trained on one set of layers can serve as effective initialization for saes designed for closely related layers. specifically, the forward approach is defined as initializing an sae with the weights of a previous layer sae, and the backward approach as initializing an sae with the weights of a subsequent layer sae. the overall training procedure is summarized in figure 1 . we tested this hypothesis on pythia-160m, a small 12-layer decoder-only transformer from the pythia family @xcite . by reusing the representations learned in earlier layers, computational demands of training can be reduced by at least 25% foot_0 while maintaining, or even improving, the quality of the resulting models. our contributions are as follows: code, data, and trained models will be publicly released after the double-blind review. ## experimental setup to address the questions raised in section 2, we first trained from scratch one sae i for each layer i of pythia-160m, a 12-layer decoder-only transformer model from the pythia family @xcite . each sae was trained using the jumprelu activation function @xcite , with activations taken from the corresponding layer's residual stream after the mlp contribution. the model configuration details are provided in ## conclusions we hypothesized and validated whether sae transfer is an effective method to accelerate and optimize the sae training process. we investigated whether sae weights derived from adjacent layers could maintain efficacy in reconstruction, which our results affirmed. furthermore, we examined whether the transferred saes, when fine-tuned on a layer's activations, could reliably capture monosemantic features comparable to the original sae, which has been also confirmed by our experiments. the transferred saes (both forward and backward) demonstrated comparable and occasionally superior reconstruction loss relative to the original. empirically, we observed frequent overlap in the most strongly activated features across adjacent layers (e.g. figure 23 ). for a given feature index i, the features learned by sae i←i+1 (backward), sae i (no transfer), and sae i←i-1 (forward) appeared to represent similar concepts.