Potential bugs and confusion with `attribution.jvp`

[JacksonKaunismaa]

I have been working through the paper trying to understand things and examining the code for computing edge weights and I believe I have discovered some unexpected behavior, as well as some other confusing areas. I would greatly appreciate some clarification on where I'm going wrong here.

Unexpected zeros

We wish to compute the edge weights between two residual layers, as in

feature-circuits/circuit.py

Lines 157 to 208 in c1a9b7b

	for layer in reversed(range(len(resids))):
	resid = resids[layer]
	mlp = mlps[layer]
	attn = attns[layer]
	MR_effect, MR_grad = N(mlp, resid)
	AR_effect, AR_grad = N(attn, resid)
	edges[f'mlp_{layer}'][f'resid_{layer}'] = MR_effect
	edges[f'attn_{layer}'][f'resid_{layer}'] = AR_effect
	if layer > 0:
	prev_resid = resids[layer-1]
	else:
	prev_resid = embed
	RM_effect, _ = N(prev_resid, mlp)
	RA_effect, _ = N(prev_resid, attn)
	MR_grad = MR_grad.coalesce()
	AR_grad = AR_grad.coalesce()
	RMR_effect = jvp(
	clean,
	model,
	dictionaries,
	mlp,
	features_by_submod[resid],
	prev_resid,
	{feat_idx : unflatten(MR_grad[feat_idx].to_dense()) for feat_idx in features_by_submod[resid]},
	deltas[prev_resid],
	)
	RAR_effect = jvp(
	clean,
	model,
	dictionaries,
	attn,
	features_by_submod[resid],
	prev_resid,
	{feat_idx : unflatten(AR_grad[feat_idx].to_dense()) for feat_idx in features_by_submod[resid]},
	deltas[prev_resid],
	)
	RR_effect, _ = N(prev_resid, resid)
	if layer > 0:
	edges[f'resid_{layer-1}'][f'mlp_{layer}'] = RM_effect
	edges[f'resid_{layer-1}'][f'attn_{layer}'] = RA_effect
	edges[f'resid_{layer-1}'][f'resid_{layer}'] = RR_effect - RMR_effect - RAR_effect
	else:
	edges['embed'][f'mlp_{layer}'] = RM_effect
	edges['embed'][f'attn_{layer}'] = RA_effect
	edges['embed'][f'resid_0'] = RR_effect - RMR_effect - RAR_effect

This is done by computing the direct effect of the residual at layer n to the residual at layer n +1 (line 199) and then subtracting off the indirect contributions of the residual in layer n on residual n+1 via both the MLP and the attention submodules (line 204).

However, these indirect contributions (RMR_effect and RAR_effect) appear to be always all zeros in my testing with all the provided supervised datasets in data regardless of edge/node thresholds, batch size, example length, or input. Is this behavior expected?

downstream_feat indexing

It also seems to me that in this section of attribution.jvp, downstream_feat is being used in two unrelated ways.

feature-circuits/attribution.py

Lines 346 to 357 in c1a9b7b

	for downstream_feat in downstream_features:
	if isinstance(left_vec, SparseAct):
	to_backprop = (left_vec @ downstream_act).to_tensor().flatten()
	elif isinstance(left_vec, dict):
	to_backprop = (left_vec[downstream_feat] @ downstream_act).to_tensor().flatten()
	else:
	raise ValueError(f"Unknown type {type(left_vec)}")
	vjv = (upstream_act.grad @ right_vec).to_tensor().flatten()
	if return_without_right:
	jv = (upstream_act.grad @ right_vec).to_tensor().flatten()
	x_res.grad = t.zeros_like(x_res)
	to_backprop[downstream_feat].backward(retain_graph=True)

Using the notation of equation 6 from the paper, on line 350, it indexes into left_vec, selecting the gradient of some particular feature in d with respect to the features in intermediate node m. Then, on line 357, we index into to_backprop with downstream_feat, where to_backprop is an element-wise product of an intermediate node gradient m and the current activation of m. If downstream_feat corresponds to some feature in node d, why can we use it to index into intermediate node m?

vjv vs. jv

Finally, I believe there could be some issue with how vjv and jv are computed inside of attribution.jvp. As far as I can tell (and running the code confirms this), by the return statement here

feature-circuits/attribution.py

Lines 388 to 391 in c1a9b7b

	return (
	t.sparse_coo_tensor(vjv_indices, vjv_values, (d_downstream_contracted, d_upstream_contracted)),
	t.sparse_coo_tensor(jv_indices, jv_values, (d_downstream_contracted, d_upstream))
	)

vjv_indices and vjv_values are identical to jv_indices and jv_values. Therefore, the only way that MR_effect and MR_grad differ in

feature-circuits/circuit.py

Line 162 in c1a9b7b

MR_effect, MR_grad = N(mlp, resid)

is their size, but each having the same underlying values and indices. However, on lines 186 and 196 in

feature-circuits/circuit.py

Lines 179 to 197 in c1a9b7b

	RMR_effect = jvp(
	clean,
	model,
	dictionaries,
	mlp,
	features_by_submod[resid],
	prev_resid,
	{feat_idx : unflatten(MR_grad[feat_idx].to_dense()) for feat_idx in features_by_submod[resid]},
	deltas[prev_resid],
	)
	RAR_effect = jvp(
	clean,
	model,
	dictionaries,
	attn,
	features_by_submod[resid],
	prev_resid,
	{feat_idx : unflatten(AR_grad[feat_idx].to_dense()) for feat_idx in features_by_submod[resid]},
	deltas[prev_resid],

, this differing size matters when we reshape these *_grad variables. If we had instead applied the same reshape to the analogous *_effect variables, we would end up with 2 different tensors that have the same values but rearranged and moved around in a strange permutation that does not seem to correspond to how these variables were computed in the first place. Is this behavior expected?

[saprmarks]

Many, many thanks—you're absolutely right about your "Unexpected zeros" and "vjv vs. jv" points: it looks like the computation of the correction terms for the inter-residual-stream effects was just completely bugged.

I've attempted a bugfix in the new jvp-fix branch. I've tested that it runs and doesn't have the issues you raised above (e.g. there are nonzero values for RMR and RAR). On brief inspection, it also looks like the new implementation coincides with the results of the old implementation for AR, MR, RA, and MA (and I'm relatively confident that the old implementation gave correct results for these). But better testing is evidently needed, which I'm hoping to get to in the next few days.

One note about the use of jvp for computing RAR and RMR.

RMR_effect = jvp(
    input=clean,
    model=model,
    dictionaries=dictionaries,
    downstream_submod=mlp,
    downstream_features=features_by_submod[resid],
    upstream_submod=prev_resid,
    left_vec={feat_idx : unflatten(MR_grad[feat_idx].to_dense()) for feat_idx in features_by_submod[resid]},
    right_vecdeltas[prev_resid],
)

(Argument names added for clarity.) One weird, but intentional, thing that's happening here is that the downstream_submod is not the same submodule that the downstream_features are coming from. What's going on here is that we need to compute, for each $\mathbf{d}$ and $\mathbf{u}$ the RHS of the expression
$$\sum_{\mathbf{m}} \nabla_\mathbf{d}m \nabla_\mathbf{m}\mathbf{d} \nabla_\mathbf{u}\mathbf{m} = \nabla_\mathbf{u}\left(\sum_{\mathbf{m}} \nabla_\mathbf{d}m\nabla_\mathbf{m}\mathbf{d} \mathbf{m}\right)$$
where the gradients inside the parentheses on the RHS are treated as a constant. This requires iterating over downstream features $\mathbf{d}$ (here, from the next layer's residual stream) but backpropping from (a weighted sum of) $\mathbf{m}$ (here, MLP or attention activations).

In other words, when reading the implementation of jvp in this case, you should keep in mind that downstream_feat doesn't correspond to a feature of the downstream_submod—it's actually a feature of a deeper submodule.

(Of course, the old implementation still wasn't correct; for instance, it didn't sum over the $\mathbf{m}$ dimension!)

[tue147]

Although I have not run the code, I think there is a bug in jvj and jv. As in the code, they are computed identically except for the return

return (
        t.sparse_coo_tensor(vjv_indices, vjv_values, (d_downstream_contracted, d_upstream_contracted)),
        t.sparse_coo_tensor(jv_indices, jv_values, (d_downstream_contracted, d_upstream))
    )

where the d_upstream will have the shape (batch x n_ctx x (d_feature + d_model)), and d_upstream_contracted will have (batch x n_ctx x (d_feature + 1)). This will result in a mismatch in shape of jv since it is computed exactly the same as jvj (which has the latter shape).

Another evidence is that, in circuit.py, the return of the jv is grad:

 MR_effect, MR_grad = N(mlp, resid)

which should not be mutiply with the right_vec:

            vjv = (upstream_act.grad @ right_vec).to_tensor().flatten()
            if return_without_right:
                jv = (upstream_act.grad @ right_vec).to_tensor().flatten()

So I think the code here should be:

            vjv = (upstream_act.grad @ right_vec).to_tensor().flatten()
            if return_without_right:
                jv = (upstream_act.grad).to_tensor().flatten()

which also match the shape of (batch x n_ctx x (d_feature + d_model)) later on.

Please notify me if I am wrong, I am planning to run the code soon.
Btw, great work from the author!

[saprmarks]

Sorry, I'm having a hard time figuring out if you're talking about the code in main (known to be bugged, see above discussion) or in the jvp-fix branch (which I never got around to merging into main, sorry).

In either case, nowadays the place to look to find a correct implementation is the gemma branch. Here the edge attributions are computed using the conceptually simpler (and easier to implement) approach discussed in appendix A.1 here: instead of computing the weight of an edge $$\mathbf{u}\rightarrow \mathbf{d}$$ by computing total effects and then subtracting off indirect effects via intermediate nodes, we simply compute

where $$\nabla_{\mathbf{u},\text{stop}(\mathcal{M})}$$ means "compute the gradient, but with a stop-gradient applied to all of the nodes intermediate between $$\mathbf{u}$$ and $$\mathbf{d}$$."

[tue147]

Thank you for your clarification.