Sparse Autoencoder Attribution Graphs: Boundaries Decide Everything

Sparse autoencoder (SAE) feature attribution graphs are commonly used to predict ablation effects and to compare across different training seeds. But a new study points out that these attribution graphs are actually determined by the feature "margin"—the distance between the pre-activation value and the TopK cutoff threshold. Ignoring the margin leads to significant prediction errors, while a simple margin correction adds almost no cost yet greatly improves accuracy.

Attribution Graphs and the Margin Problem

Attribution graphs predict model behavior by analyzing interactions between SAE features, and are especially used in ablation experiments to determine which features influence the output. A TopK SAE keeps the largest few pre-activation values on each token, and a feature's margin is defined as the distance between its pre-activation value and the cutoff threshold. The researchers conducted experiments on five public TopK SAE pairs from three Pythia models (the largest being pythia-160m).

The key finding is: when a feature is zeroed out, the proportion of next-layer features crossing the cutoff threshold reaches as high as 0.74 or more, and these crossing events carry most of the response energy. This means that a simple ablation operation causes a large number of features to "enter and exit" the active set, while traditional attribution graphs only consider currently active features and ignore this margin effect.

TopK Sparse Autoencoder (TopK SAE) is a type of autoencoder that learns sparse features by forcing the hidden-layer activations of each token to retain at most k non-zero values. Unlike traditional SAEs that use L1 regularization to encourage sparsity, TopK SAE directly sorts pre-activation values and keeps the largest k, zeroing out the rest. This design avoids the shrinkage bias caused by L1 penalties, making the magnitude of active features more faithful to the original activations. The cutoff threshold is the k-th largest pre-activation value, and whether a feature is active is entirely determined by the relative size of its pre-activation value and the threshold. Attribution graphs analyze interactions among active features on this basis, but traditional methods assume the active set remains unchanged under perturbation, ignoring features whose pre-activation values are close to the threshold and may enter or exit the active set due to small changes.

Advantages of the Margin Graph

To solve this problem, the researchers proposed the "margin graph": when constructing the graph, they add a first-order change term to each pre-activation value and then re-sort. This method has the same construction cost as the original graph, but the error is significantly reduced. Experiments show that in a write operation that scales feature encodings by three times, the median relative error of graphs built from previously active features is at least 0.247, while the maximum error of the margin graph is only 0.033.

In addition, in the median case, the decoded change of features crossing the threshold is more than 1.24 times the norm of the model's own change, indicating that these features' contribution to the output cannot be ignored. The margin graph can capture this dynamic more accurately.

First-order change term refers to the linear approximation of the pre-activation value with respect to the input perturbation, i.e., estimating the change in the pre-activation value by multiplying the gradient by the perturbation amount. In the margin graph, the researchers do not simply sort by the original pre-activation values, but add the corresponding first-order change term to each pre-activation value and then re-sort, thereby simulating which features will cross the cutoff threshold after perturbation. This method is essentially a local linearization of the TopK selection process; the computational cost is the same as the original attribution graph (only one forward and backward pass is needed), yet it can capture the dynamics of features entering and exiting the active set. Since the TopK operation itself is non-differentiable, directly attributing to active features misses the influence of features near the boundary, whereas adding the first-order change term is equivalent to a continuous relaxation of TopK's discrete selection, making the attribution graph more accurately predict the actual effects of operations such as ablation or scaling.

Cross-Seed Consistency

The study also compared attribution graphs across different training seeds on the same testbed. It found that the sparse mapping between features can explain 0.85 of the energy of one graph on its active features, but when another seed's graph retains only its own features, this proportion drops to 0.47. This indicates that although features learned by different seeds have some correspondence, the structure of the attribution graph is still affected by the seed, and the margin effect may be one of the causes of the discrepancy.

Training seed refers to the initial state of the random number generator set before model training, which affects random processes such as weight initialization and data order. SAEs trained with different seeds may learn feature dictionaries that are semantically similar but differ in specific direction or magnitude, leading to differences in attribution graph structure. Sparse mapping here refers to mapping the active features of one SAE to the feature space of another SAE through some correspondence (such as cosine similarity of decoder directions), thereby comparing the proportion of energy explained by the two graphs on the same feature basis. When only the active features under each seed are retained, the mapping fails and the explained proportion decreases, indicating that the attribution graph is sensitive to the specific choice of features. The margin effect may amplify this difference: features near the threshold have different active states under different seeds, which in turn affects the graph structure. Understanding cross-seed consistency is crucial for evaluating the reliability of attribution graphs.

Implications for Interpretability

This study reminds us that the accuracy of SAE attribution graphs is highly dependent on how margin conditions are handled. Ignoring the margin leads to prediction bias, especially in operations involving ablation or scaling. The margin graph provides a lightweight correction scheme that adds almost no computational overhead yet significantly improves the fidelity of attribution graphs. For those working in mechanistic interpretability research, this is an improvement worth incorporating into standard workflows.

Future work can explore more complex margin interactions and generalization across different model scales. The largest model studied so far is pythia-160m, and performance on larger models still needs verification.

One-line Summary

New research reveals that SAE attribution graphs are governed by feature margins, and margin graphs can greatly improve prediction accuracy at low cost.

Paragraph Summary

Sparse autoencoder (SAE) attribution graphs are often used to predict ablation effects, but a new study finds that their accuracy is governed by feature "margins"—the distance between pre-activation values and the TopK cutoff threshold. Experiments show that zeroing out a feature causes up to 0.74 or more of next-layer features to cross the threshold, carrying most of the response energy, while traditional attribution graphs ignore this effect. The researchers propose the margin graph, which adds a first-order change term and re-sorts, reducing the median relative error from at least 0.247 to a maximum of 0.033 at the same cost. Cross-seed comparisons show that sparse mapping can explain 0.85 of the energy, but this drops to 0.47 when retaining each seed's own features; the margin effect may be a source of the discrepancy. The margin graph is a lightweight correction and is recommended for inclusion in standard interpretability research workflows.