kefirski.me
ICMLPaper 0920258 min read

Analyze Feature Flow to Enhance Interpretation and Steering

Cosine similarity between SAE decoder directions, computed from weights alone, traces features across layers into flow graphs that explain and improve model steering.

One model, many dictionaries

A sparse autoencoder rewrites a hidden state hRd\mathbf{h}\in\mathbb{R}^d as a sparse combination of learned directions, h^=Wdecσ(Wench+benc)+bdec\hat{\mathbf{h}}=\mathbf{W}_{\mathrm{dec}}\,\sigma(\mathbf{W}_{\mathrm{enc}}\mathbf{h}+\mathbf{b}_{\mathrm{enc}})+\mathbf{b}_{\mathrm{dec}}, with a dictionary far wider than dd. Each column of Wdec\mathbf{W}_{\mathrm{dec}} is one feature, and on a good day one feature means one thing: a token pattern, a topic, a syntactic role. Gemma Scope ships such dictionaries for Gemma 2 2B at three positions per layer: the residual stream, the MLP output, and the attention output.

Those dictionaries are usually studied one at a time. But a transformer is a pipeline, and whatever a feature means at layer 24 was assembled somewhere below. We wanted the assembly record: where a feature first appears, which modules edit it, how long it survives. And we wanted it without running the model over a corpus, because collecting activation statistics is exactly the expensive part.

Match directions, not activations

Our tool is almost embarrassingly simple. Features live in the same Rd\mathbb{R}^d as the states they reconstruct, so we compare decoder directions directly: feature f\mathbf{f} from one dictionary matches index j=argmaxkfWdec,:,k(B)j=\arg\max_k\,\mathbf{f}\cdot\mathbf{W}^{(B)}_{\mathrm{dec},:,k} in another, with unit norms assumed. No data, no forward passes, just a matrix product between two sets of weights.

We ask the question backward: where did this feature come from? For a feature at a layer's output we compute three scores, s(R)s^{(R)}, s(M)s^{(M)}, and s(A)s^{(A)}, the best cosine against the previous residual, the MLP output, and the attention output. The pattern of highs and lows sorts features into translated, processed, newborn, or unexplained, and the split is not cosmetic: groups defined this way differ in actual activation behavior in 67 to 100 percent of our pairwise statistical tests.

Where did this feature come from?

Data-free matching
Every comparison is a cosine between decoder columns, so the whole classification runs from weights alone. Thresholds for high and low are set per layer.

A feature's biography, layer by layer

Chaining top-1 matches through consecutive layers turns a single feature into a flow graph: its ancestry, its edits, its descendants. The graphs read like biographies. The one below starts from a layer-24 feature for Standard Model parameters and walks back to astronomy vocabulary near layer 1, picking up measurement semantics from an MLP feature at layer 7 with cosine 0.820.82, and gauge-theory semantics from another at layer 17 with 0.790.79.

A real flow graph, abridged

Gemma 2 2B
Interpretations come from Neuronpedia; the graph is built backward from feature 24/res/14548. Grey hops are residual persistence; the vertical arrows are the two module injections our scores flagged as processing events.

Across the whole model the demographics are consistent: layers 0 to 5 are a high-entropy nursery, 6 to 15 a stable middle where modules keep editing features, and past 16 the residual track dominates, with MLP births peaking around layer 18 and fading after. Llama Scope shows the same three-act structure, so this is not one model's quirk.

Do the matches carry cause?

Cosine agreement could still be correlation without cause, so we intervene. Deactivating a predecessor means subtracting it from the hidden state, hh+a(r1)v\mathbf{h}\leftarrow\mathbf{h}+a\,(r-1)\,\mathbf{v} with the rescaling coefficient r=0r=0, then checking whether the descendant feature dies with it. Picking a random feature from the top-5 matches instead of the top-1 collapses the success rate, so the best-matching direction really is the load-bearing one.

Silencing a feature through its predecessor

successful deactivation, %
Silencing a feature through its predecessor 1,894 features across two layers. Our data-free matching ties correlation-based Pearson matching, and trying every active predecessor individually, the practical ceiling, adds only eight points. top-1 cosine (ours): top-1 cosine 65% top-1 Pearson: top-1 Pearson 65% exhaustive search: exhaustive search 73% 0255075100 65%top-1 cosine: top-1 cosine (ours), 65%top-1 cosine (ours)65%top-1 Pearson: top-1 Pearson, 65%top-1 Pearson73%exhaustive search: exhaustive search, 73%exhaustive search successful deactivation, %
1,894 features across two layers. Our data-free matching ties correlation-based Pearson matching, and trying every active predecessor individually, the practical ceiling, adds only eight points.

The same experiment, both metrics

MethodMean activation changeSuccess rate
Top-1 cosine (ours)0.7565%
Top-1 Pearson correlation0.7465%
Exhaustive search ceiling0.8373%
Activation change of 1 means the descendant was fully silenced. Cosine matching needs no data at all; Pearson needs activation statistics over a corpus.

The near-tie with Pearson is the point: a matrix product between weight matrices recovers what correlation mining recovers. The failures are informative too. Features fed by residual and MLP predecessors together often survive a single deactivation, which is exactly what circuit-like redundancy should look like, and residual predecessors matter most, consistent with the residual stream being the model's main communication channel.

Steering along the flow

Flow graphs turn steering into a multi-layer operation. Instead of one large edit at one layer, we amplify or suppress a feature's whole lineage, decaying the per-layer coefficient as s=seαls'=s\,e^{\alpha l} with α=0.05\alpha=-0.05. In the low-strength regime the cumulative intervention beats the best single layer on our judged topic scores, is less sensitive to hyperparameter choice, and perturbs each hidden state less to get the same behavioral shift.

The graphs also explain steering surprises before they happen. A feature labeled as London drifts generations toward fashion when pushed hard; its flow graph's early layers turn out to be full of fashion-and-exhibition semantics, so the drift was written in the ancestry all along. Reading the graph first tells you what you are actually about to amplify.

What the method cannot see yet

The method inherits its blind spots honestly. High and low thresholds are per-layer judgment calls, an unexplained remainder group persists, and matching requires dictionaries that share the hidden dimension, which ruled out Gemma Scope's attention SAEs and forced us to train our own; those weaker probes likely understate attention's role in feature birth.

What survives the caveats is a cheap, causally tested map of feature evolution. A dictionary entry is not a fact about one layer. It is a point on a trajectory, and the trajectory is where both the interpretation and the steering leverage live.

A feature's meaning is a trajectory rather than a point, and the flow graph is how we read it, test it, and steer it.