Hello! One way we (a.k.a. some folks on Anthropic interp) have found it useful to think about causal localization on arithmetic prompts is to capture residual stream values over some big distribution of prompts like
"For example, {a}{op}{b}={c}",
and then see what happens when you do residual stream patching experiments with various marginals from that big table of values.
Here, I took all pairs of
The probe is about as confident at layer 5 as at layer 20, and i'm pretty sure the model hasn't done active computation that early. I suspect this is just that you can probe for information that 'spilled over' via attn well before the model uses its MLPs to do anything. So while the NLA is very powerful (it can do a lot of computation itself), a linear probe has a bit of power on its own!
see https://arxiv.org/abs/2005.00719 for example. that's why i like the patching experiments for estimating what the model has computed -- it should be legible to the model