Skip to main content
Back to timeline
arXivSource publication:

Why the Jacobian lens favors slow features: a theory writes its choice as three coherence factors and predicts their ordering in toy models

Synopsis

The Jacobian lens (J-lens), an interpretability tool that decodes a model's hidden states into words, had been used to find working-memory-like linear directions, but what makes a feature detectable to it was unclear. This work gives a mathematical theory: the lens preferentially selects directions that grow across layers, stay smooth across token positions, and remain consistent across contexts, that is, coherent Lyapunov modes. On a Bayes filter and a two-layer toy transformer, changing task parameters to modulate coherence shifts the ordering of latent states the lens selects, as predicted.

Source-provided article image: A Mathematical Theory of the Jacobian Lens
arXiv

Interpretation

The Jacobian lens selects not the highest-variance directions but directions that jointly have three coherences: growth along layers, smoothness across token positions, and consistency across contexts. The lens had been used to identify working-memory-like linear directions in frontier language models, but which mechanistic properties make features detectable to it was unclear. The theory writes the lens singular value as a product of three coherence factors and uses it to explain the original finding that J-space covers no more than 10% of activation variance.

For each token-position direction the lens averages over lags, which is equivalent to taking the zero-frequency component over a time window, so rapidly decaying or rapidly oscillating responses are suppressed. A single-point Jacobian describes perturbation response only in a local neighborhood, while the lens averages over prompts, token positions and future positions, and the averaged matrix need not resemble any of its elements. On the Bayes filter, three latent states with velocities 0, 1 and 4 on a ring of 8 states are identified in the order predicted by lag coherence, and the transformer result holds across seeds.

Changing task parameters independently modulates the three coherences, shifting the ordering of latent states the lens selects, and the point where two latent states swap order is also predicted. There was no manipulable comparison for judging whether the lens preference comes from task structure or from the model itself. On the Bayes filter the theory exactly predicts the crossing point of two latent states; on the transformer the prediction carries a small bias, in opposite directions at the two layers.

Perspective

Validation covers a Bayes filter and a two-layer decoder-only transformer with residual dimension 128 and 4 heads per layer on a hidden Markov model task that is relatively easy to learn and whose coherence parameters are directly manipulable. The theory itself targets general spatiotemporal dynamical systems, but the numerical comparison does not go beyond this toy setting.

The theory explicitly does not decide whether these properties correspond to a Global Workspace. The lens is noisier in the context-coherence sweep, both transformer layers compress the range of context coherence, and some lens directions align with reveal and key states rather than latent belief states. The transformer crossing prediction carries a small bias, in opposite directions at the two layers.