Decoder cosine similarity struggles to predict model feature interactions
Where Decoder Cosine Similarity Fails for SAE Feature Flow Discovery
Machine Learning
Summary
The paper looks at how parts of large AI models work together when they learn new things but keep old knowledge. The authors study small building blocks inside these models called sparse autoencoders and how they predict changes in the model’s internal signals. They find that a common technique called decoder cosine similarity often fails to spot important interactions between features that help the model update itself. This suggests that new tools are needed to better understand and guide how AI models change and learn over time.
What this means in practice
- •For ai model engineers: Diagnose and improve how model features update internally using feature flow atlases beyond decoder cosine similarity.
- •For machine learning operations teams: Develop better tools to steer and maintain foundation models during fine-tuning or alignment by understanding update feature interactions.
Authors
Hendrik Droste, Christian Medeiros Adriano, Kathrin Korte, Holger Giese
Abstract
Foundation models are increasingly adapted through fine-tuning, model editing, and alignment procedures while retaining previously acquired capabilities. Understanding the internal computations that support these adaptations is therefore becoming increasingly important for continual model evolution. Sparse autoencoders (SAEs) provide interpretable feature dictionaries for residual-stream activations and sublayer outputs, but it remains unclear how state features and update features interact to produce downstream residual features. In this work, we focus on MLP updates as a first test case. We construct a transition atlas of triples $s_k + u_j \rightarrow t_\ell$, where a residual-state feature and an MLP-update feature jointly predict a target residual feature, and validate candidate triples by ablating the decoded update feature. In a 20M-token Pythia-160M $L_7 \rightarrow L_8$ run, we find 38,125 strong ablation-effect transitions, but 88.0% have both state-target and update-target decoder cosine similarity below 0.7. As a preliminary cross-model check, a run of 20M-token Gemma-3-4B $L_{21} \rightarrow L_{22}$ causally validates only the top 30,000 ranked candidate triples by ablating the decoded update feature, and 53.6% of strong-effect triples have both state-target and update-target decoder cosine similarity below 0.7. The Gemma result is directionally consistent with Pythia, but weaker, since update-target cosine recovers many of the strongest Gemma effects and the run is not a full-atlas causal validation. Ultimately, our results suggest that feature flow atlases can serve as diagnostics of representation-update mechanisms and thereby inform tools for steering model updates. Future work will validate more complex patterns across layers, models, and SAE families.