Mechanistic Permutability: Match Features Across Layers

TL;DR

This paper introduces SAE Match, a data-free method for aligning features across layers in neural networks, enhancing understanding of feature evolution and persistence in deep models.

Contribution

SAE Match is a novel, data-free technique that aligns features across layers by minimizing parameter differences, improving interpretability of feature dynamics.

Findings

Effectively captures feature evolution across layers.

Features persist over multiple layers.

Can approximate hidden states across layers.

Abstract

Understanding how features evolve across layers in deep neural networks is a fundamental challenge in mechanistic interpretability, particularly due to polysemanticity and feature superposition. While Sparse Autoencoders (SAEs) have been used to extract interpretable features from individual layers, aligning these features across layers has remained an open problem. In this paper, we introduce SAE Match, a novel, data-free method for aligning SAE features across different layers of a neural network. Our approach involves matching features by minimizing the mean squared error between the folded parameters of SAEs, a technique that incorporates activation thresholds into the encoder and decoder weights to account for differences in feature scales. Through extensive experiments on the Gemma 2 language model, we demonstrate that our method effectively captures feature evolution across layers, improving feature matching quality. We also show that features persist over several layers and that our approach can approximate hidden states across layers. Our work advances the understanding of feature dynamics in neural networks and provides a new tool for mechanistic interpretability studies.