The problem with J-Lens: we want the contribution of $h_\ell$ to the output $h_L$, which is not exactly the same as the Jacobian $\frac{\partial h_L}{\partial h_\ell}$., which is how a delta to $h_\ell$ would get propagated to $h_L$, not how the entirety of $h_\ell$ contributes to $h_L$

Notably, the two concepts are aligned for linear layers, but not for other nonlinear layers, which look of the form $y = C(x) x$. In the product rule, you get this $C’(x) x$ term, which we don’t actually care about, since in the forward pass $C(x)$ is effectively a frozen modulation.

• so we just add these three stopgrad changes to components when calculating the Jacobian for J-lens

Further justification of why the Jacobian is wrong

Note 1: Unlike the tangent (Jacobian), the secant is not unique. Tuned Lens is essentially another secant, that least squares predicts $h_L$ using $h_\ell$ but it suffers from ignoring the model’s computational pathways and thus being correlational and not causal

Note 2: attributing “contribution of $h_l$ to $h_L$” is a separate problem from estimating the secant. Tuned Lens doesn’t try to estimate it at all and just uses $h_L$ directly, R-Lens uses this frozen gradient method. One could try doing this secant estimation with Integrated Gradients as a more principled R-Lens

Experiments

• matters more as the models get bigger. It doesn’t matter at all at the 4B scale

• seems more effective at causal intervention even at 4B though

The fact that they get reasonable outputs in early layers kinda calls into question the J-space paper’s localizing the workspace into middle layers thing.

• the CKP similarity matrix (only look at bottom left square and upper right square) for J-Lens vs R-Lens.