A Mathematical Theory of the Jacobian Lens

1 INTRODUCTION

• The illustrative model assumes shared input and output directions across contexts, separable depth and lag factors, and independent depth and token-position exponents.It also assumes the context-specific response is no larger than the shared response.Source: §1 · p. 3
• Averaging across token positions suppresses rapid decay or oscillation, while depth weights modes by e^Tχ and favors growing or slowly decaying responses.Slow modes retain their response under the averaging window; the depth weighting may be weak near the output as the Jacobian aligns with the identity.Source: §1 · p. 3
• When responses are weakly overlapping across contexts, averaging suppresses incoherent responses and the leading singular directions approach those of the shared component.The convergence is stated as the overlap parameter η tends to zero.Source: §1 · p. 4
• The J-lens favors directions that grow across layers and remain coherent across token lags and contexts.The general construction ranks modes by magnitude a and lag and context coherence, ρτ and ρc.Source: §1 · p. 4
• The variance explained by a direction selected through Jacobian averaging can be arbitrarily small or large, so singular values do not determine it.The paper gives a two-dimensional counterexample in Appendix D.Source: §1 · p. 5

1.1 MAPPING COHERENT MODES ONTO TOKENS

• The singular modes underlying J-space become token-indexed directions through unembedding alignment and sparse non-negative combinations.Singular values weight coherent response channels, while alignment with a token’s unembedding vector weights each channel’s contribution.Source: §1.1 · p. 5

1.2 FAMILIES OF JACOBIAN LENSES

• Changing the Fourier projection frequency yields ω-J-lenses that can emphasize different temporal scales, with band-pass filtering also proposed.The paper demonstrates their selectivity on a toy task.Source: §1.2 · p. 5

2 TOY MODEL

• Independent latent states on rings form an HMM task whose parameters selectively modulate the J-lens coherence factors.The task is designed for next-token prediction and has three latent states in the experiments.Source: §2 · pp. 5–6
• Changing persistence λ modulates response magnitude, while ring velocity changes angular velocity and modulates lag coherence.Latents are revealed with probability ε so the model must integrate over time; otherwise the Jacobian would be zero across timepoints.Source: §2 · p. 6
• The setup is intended to let researchers predict the ordering and approximate coherence parameters of transformer J-lens modes in advance.The task is described as easily solvable by a transformer, whose embeddings learn to represent the latent states.Source: §2 · p. 6

2.1 INTRODUCING CONTEXTUAL DEPENDENCY

• A fully observed key state W_t flips the encoded observation according to key bias ψ, making read-in and read-out directions context-dependent.Each keyed state has an independent key.Source: §2.1 · p. 6

3 THE BAYES FILTER AND ITS JACOBIAN LENS

• The Bayes filter optimally integrates evidence into belief states, and Fourier modes account for rotational symmetry of the ring latents.Unkeyed evidence updates each mode using the expected evolution operator when evidence is absent, or replaces the state with the observation when it is present.Source: §3 · p. 6

3.1 KEYED FILTERS

• Keyed observations make the Bayes filter’s read-in and prediction directions context-dependent, modulating contextual coherence.The observed Fourier mode is related to the latent mode through the directly observable key; the filter accounts for unchanged and flipped states.Source: §3.1 · p. 7

3.2 THE JACOBIAN LENS

• When angular velocity is zero, the terms accumulated before the next reveal are real and positive; nonzero rotation makes them partly cancel and lowers lag coherence.Persistence λ changes magnitude, while context coherence is |ψ| when θ_j = 0.Source: §3.2 · p. 8

3.3 CONNECTION TO THE TRANSFORMER

• The transformer J-lens reweights the Bayes-filter Jacobian by a positive factor, preserving the ordering of magnitude and coherence factors.Because it maps through the unembedding and softmax, the transformer lens concerns perturbations to log probabilities rather than probabilities.Source: §3.3 · p. 8

4 NUMERICAL RESULTS

• The transformer learns to predict next tokens as well as the optimal Bayes filter, and the J-lens unsupervisedly identifies independent latent modes ordered by the selectively modulated coherence factor.The experiments modulate magnitude, lag coherence, or context coherence across three independent latent states.Source: §4 · p. 8, §4.1 · p. 8

4.1 SELECTIVELY MODULATING MAGNITUDE

• The J-lens ranks latents by magnitude a as persistence λ changes, while lag and context coherence stay fixed, in both the Bayes filter and transformer.The tested persistence values are λ = 0.8, 0.9, and 0.95; the transformer ranking holds across seeds and for input- and output-side directions.Source: §4.1 · p. 8, Figure 4 · p. 9

4.2 SELECTIVELY MODULATING LAG COHERENCE

• Changing ring velocity selectively changes lag coherence ρτ, and both models’ J-lenses rank latents in the predicted order.For transformer latents, ρτ varies most with velocity, while magnitude a and context coherence ρc are less sensitive.Source: §4.2 · p. 9, Figure 5 · p. 10

4.2.1 ω-LENS FAMILY

• Sweeping ω makes ω-J-lenses preferentially upweight latent directions whose lag-coherence parameters align with ω.The authors suggest sweeping ω to identify information that changes rapidly or slowly across tokens.Source: §4.2.1 · p. 9, Figure 6 · p. 10

4.3 SELECTIVELY MODULATING CONTEXT COHERENCE

• When key bias ψ changes context coherence ρc independently, the Bayes filter and transformer rank latents in the predicted order.In the Bayes filter, ρc tracks ψ; transformer lens factors show a compressed range, with higher ρc for the incoherent latent and lower values for coherent latents.Source: §4.3 · pp. 9–10, Figure 7 · p. 11

4.4 THEORY PREDICTS THE CROSSING OF TWO LATENTS

• Theory exactly predicts the Bayes-filter persistence λ*₀ at which two latents swap order, and predicts transformer crossings with O(0.01) error and opposite-direction bias across layers.The transformer prediction is informative but does not perfectly match the Bayes-filter theory.Source: §4.4 · pp. 10–11

5 DISCUSSION

• J-lens directions jointly favor growth across layers, spatial coherence across token positions, and context coherence, connecting the lens to Lyapunov modes.The discussion describes magnitude as layer-wise growth and lag coherence as limited spatial oscillation; transformers showed higher context coherence than the Bayes filter for a context-incoherent latent.Source: §5 · pp. 11–12
• These preferences explain prior findings that J-space captures at most 10% of activation variance and identifies features prevalent across token positions or broadly verbalizable.The authors relate lag coherence to identifying a feature present across positions and context coherence to concepts not specific to one prompt.Source: §5 · p. 12
• The analysis motivates Jacobian lenses for different information scales, including high-frequency modes for fast-changing syntax and modes minimizing context coherence.The authors also suggest that ω-J-lenses could be applied to biological neural data, given recent work estimating Jacobians from dynamical data.Source: §5 · pp. 12–13
• The discussion's evidence is limited to toy models with relatively easy tasks and coherence parameters that can be manipulated.Source: §5 · p. 13

A JACOBIAN SINGULAR VALUES UNDER AVERAGING

• Averaging Jacobians with shared singular vectors but permuted singular values turns each sample's decaying spectrum into a flat mean spectrum.The mean singular values converge to the same value.Source: A JACOBIAN SINGULAR VALUES UNDER… · p. 15

B REAL JACOBIANS OF LLMs

• The LLM analysis computed per-token Jacobians from 96 prompts across 16 public datasets, truncated texts to 57 tokens, and produced 768 matrices per model.The authors visualized top-100 singular-value distributions by projecting individual Jacobians and the mean into a PCA space.Source: B REAL JACOBIANS OF LLMS · p. 16

C THE SHARED-MODE MODEL

• After averaging, shared singular strengths are |g_i|, while the assumptions bound accumulated depth growth and exclude growing lag dynamics.Source: C THE SHARED-MODE MODEL · p. 17
• Spectral dominance requires shared strengths to exceed the averaged context-specific remainder; the assumptions prevent that remainder from becoming arbitrarily larger than the shared response scale.Source: C THE SHARED-MODE MODEL · p. 18
• At fixed response scale, a direction's contribution to the lens is capped by the product of lag and context coherences.The response scale is the product of ρc(v), ρτ(v), and a(v).Source: C THE SHARED-MODE MODEL · p. 18

D PROOF OF COROLLARY 1.1

• The example's first direction is perfectly coherent across contexts and constant across future lags, while its second changes sign across contexts and oscillates rapidly across lags.Source: D PROOF OF COROLLARY 1.1 · p. 18
• The unique nonzero singular value is s, and the corresponding input direction is e1.Source: D PROOF OF COROLLARY 1.1 · p. 18
• The selected direction can capture arbitrarily little variance, or approach a fraction of 1 when covariance eigenvalues are interchanged, without changing the Jacobians.The example's J-lens directions are proportional to e1; coherent, slowly varying responses are weighted by response magnitude, not activation variance.Source: D PROOF OF COROLLARY 1.1 · p. 18

E BAYES FILTER DERIVATIONS

• A distribution p on Zn is defined by its modes E_pϕ_j(X).Source: E BAYES FILTER DERIVATIONS · p. 19
• The mode values are expectations under the distribution p.Source: E BAYES FILTER DERIVATIONS · p. 19
• The passage uses the collection of modes to characterize the distribution on Zn.Source: E BAYES FILTER DERIVATIONS · p. 19

• The reveal-time process gives E_c A_τ=(1−ε)^τ because A_τ remains one until the next reveal and is zero thereafter.Source: §E.1 · p. 19
• For real α_j, static modes have ρ_c=1, while alternating modes have ρ_τ=(1−r_j)/(1+r_j) and perfect context coherence.For intermediate rotations, both coherences are below one, while their product is fixed by Eq. (44).Source: §E.1 · p. 20
• Independent latents make the Bayes-filter lens block diagonal, so their singular values cross and the top direction switches latents when their values are equal.Source: §E.1 · p. 20
• With latent 0 static and its persistence swept, its singular value rises monotonically and overtakes latent 1 above a critical persistence λ*_0.A faster-turning latent 1 is overtaken at a lower λ*_0; the closed-form threshold applies when latent 0 is static.Source: §E.1 · p. 20

• Each model is a decoder-only transformer with causal self-attention and rotary position embeddings applied to queries and keys.The HMM determines the vocabulary size.Source: F TRANSFORMER TRAINING METHODS · pp. 20–21
• The two-block residual stream is indexed at positions before each block, and the J-lens is read at layers ℓ=0 and ℓ=1.Source: F TRANSFORMER TRAINING METHODS · p. 21
• Models are trained for 20,000 steps with AdamW on next-token cross-entropy and compared against the optimal Bayes filter and a stationary-distribution memoryless predictor.Source: F TRANSFORMER TRAINING METHODS · p. 21

F.1 DATA FOR ALL SWEEPS

• The sweeps vary persistence, angular velocity, or key probability to selectively change magnitude, lag coherence, or context coherence, respectively.The reveal probability is ε=0.2 across sweeps.Source: §F.1 · p. 21

G COMPUTING ALIGNMENT OF J-LENS VECTORS AND LATENT STATE

• The J-lens is computed on revealed tokens after burn-in, averaged over 128 lags, and decomposed so right singular vectors are read directions and left vectors are written directions.For the Bayes filter, inputs and outputs are separated by latent.Source: G COMPUTING ALIGNMENT OF J-LENS… · p. 21

H ALIGNMENT BETWEEN J-LENS VECTORS AND LATENT VECTORS

• Alignment is measured by projecting lens vectors onto encoding or decoding subspaces for latent Fourier modes, with each lens direction assigned to its highest-overlap latent.A random direction’s expected squared projection is k_i/d; the stated chance overlap is p k_i/d.Source: H ALIGNMENT BETWEEN J-LENS VECTORS AND… · p. 22

I VALIDATING SINGULAR VALUE SWITCHING THEORY

• The switching test predicts each latent’s first-mode singular value and a critical persistence where the top direction changes, using two unkeyed ring latents.The swept latent is static; the other has fixed persistence and one of four nonzero angular velocities.Source: I VALIDATING SINGULAR VALUE SWITCHING… · p. 22

I.1 BAYES FILTER

• For the Bayes filter, singular values are measured across a persistence sweep and the crossing is interpolated near the predicted switch point.The singular-value sweep reuses sampled sequences across persistence values.Source: §I.1 · p. 23

I.2 TRANSFORMER

• For transformers, models are trained around the predicted switch point, and the crossing is estimated from smoothed seed-averaged singular-value differences.The comparison concerns crossing location because transformer and filter singular values use different scales.Source: §I.2 · p. 23

• Figure 11 compares predicted and measured magnitude, lag coherence, and context coherence across parameter sweeps.Each sweep has one visible point because the HMM fixes the measured metrics for its fixed parameters.Source: Figure 11 · p. 24
• Figure 12 compares predicted singular-vector rank order with measured J-lens rank for the Bayes filter and transformer.Source: Figure 12 · p. 24
• Figures 13, 15, and 17 compare J-lens directions with encoding directions across magnitude, lag-coherence, and context-coherence sweeps.The panels distinguish Bayes-filter belief states and transformer layers; Figures 14, 16, and 18 show decoding accuracy and alignment with each latent’s decoding directions.Source: Figure 13 · p. 25, Figure 15 · p. 25, Figure 14 · p. 25
• Figure 19 shows J-lens alignment with non-latent HMM features: the reveal state and key state across the context sweep.Source: Figure 19 · p. 27