Residual Stream Analysis with Multi-Layer SAEs

TL;DR

This paper introduces multi-layer sparse autoencoders (MLSAEs) trained on residual streams from all transformer layers to better understand information flow and representation changes across layers in language models.

Contribution

The paper presents MLSAEs as a novel single autoencoder approach that captures residual stream dynamics across all layers, enabling new insights into layer-specific activation patterns.

Findings

Latents are often active at a single layer per token, but this varies across tokens.

Variance in latent activation distributions over layers is much higher across tokens than within a single token.

Larger models show increased multi-layer latent activity, aligning with more similar residual streams across layers.

Abstract

Sparse autoencoders (SAEs) are a promising approach to interpreting the internal representations of transformer language models. However, SAEs are usually trained separately on each transformer layer, making it difficult to use them to study how information flows across layers. To solve this problem, we introduce the multi-layer SAE (MLSAE): a single SAE trained on the residual stream activation vectors from every transformer layer. Given that the residual stream is understood to preserve information across layers, we expected MLSAE latents to 'switch on' at a token position and remain active at later layers. Interestingly, we find that individual latents are often active at a single layer for a given token or prompt, but the layer at which an individual latent is active may differ for different tokens or prompts. We quantify these phenomena by defining a distribution over layers and considering its variance. We find that the variance of the distributions of latent activations over layers is about two orders of magnitude greater when aggregating over tokens compared with a single token. For larger underlying models, the degree to which latents are active at multiple layers increases, which is consistent with the fact that the residual stream activation vectors at adjacent layers become more similar. Finally, we relax the assumption that the residual stream basis is the same at every layer by applying pre-trained tuned-lens transformations, but our findings remain qualitatively similar. Our results represent a new approach to understanding how representations change as they flow through transformers. We release our code to train and analyze MLSAEs at https://github.com/tim-lawson/mlsae.