Thomas Kwa’s research journal
Monday 11⁄13
I did SERI MATS applications and thought about Goodhart, but most of my time was spent on the KataGo project. I might more about it later, but the idea is to characterize the nature of planning in KataGo.
Early training runs had produced promising results—a remarkably sparse mask lets the network output almost the worst possible move as judged by the value network—but I was a bit suspicious that the hooked network was implementing some trivial behavior, like always moving in the corner. I adapted some visualization code previously used for FAR’s adversarial Go attack paper to see what the policy was doing, and well...
Turns out the network is doing the trivial behavior of moving in either the top left or bottom left corner. I wish I had checked this earlier (it took ~14 days of work to get here), but it doesn’t kill this project—I can just redo the next training run to only allow moves on the 3rd line or above, and hopefully the worst behavior here won’t be so trivial.
Tomorrow I’m going to table this project and start implementing edge-level subnetwork probing—estimate is 2 days for the algorithm and maybe lots more effort to run benchmarks.
Wednesday 11⁄15
Updates from the last two days:
I finished the basic edge-level subnetwork probing code over the last two days. This is exciting because it might outperform ACDC and even attribution patching for circuit discovery. The original ACDC paper included a version of subnetwork probing, but that version was severely handicapped because it operated on the node level (structural pruning) rather than edge level.
Adria is now on vacation, so I’m planning to get as far as I can running experiments before getting stuck somewhere and coming back to this after Thanksgiving.
I’m starting to think about applications of circuit discovery to unlearning / task erasure. If we take a sparse circuit for some task found by some circuit discovery method, and ablate the circuit, can it remove the model’s ability to do that task better than other methods like task vectors?
Last night 11⁄14, Aryan Bhatt and I thought of a counterexample to one idea Drake had for extending our Goodhart results. This is disappointing because it means beating Goodhart is not as easy as having light-tailed error.
I’m going to spend the next two days on MATS applications and syncing back up with the interp proofs project.
Saturday 11⁄25
Adria is still on vacation, and I can’t make sense of the benchmark code within an efficient amount of time. I did finally fork the ACDC repo and publish my changes.
Today was a good time to think about strategy. Last night I started making a tech tree for alignment because this seems necessary to compare different research directions against each other. This is a first draft of a fragment centering on circuit discovery and its potential applications, including scoping models down. The key point is that, in my view, better circuit discovery algorithms have multiple applications, including removing capabilities from a model and generally furthering mechinterp. I think this means working on circuit discovery is likely to be useful.
Zooming in to what I’m trying to do now (the “gradient-based monosemantic CD” node), there are three problems with current circuit discovery algorithms that feel live to me:
They fail validity checks on tracr models with known hardcoded circuits.
They are regularized for sparsity and operate on neurons / attention heads, but maybe superposition means they should operate on autoencoder features instead.
They poorly capture redundant circuitry (likely more problems in this class).
The first one seems pretty severe and has the best feedback loops, so I think I’m just going to iterate on methods and try them on the tracr benchmarks until it’s solved. Hopefully I have enough ideas if edge-level subnetwork probing doesn’t work straightaway—I can probably get some by looking at where it goes wrong. After this I’m pretty worried about the different possible metrics for validity. What if a method that does better on tracr does worse on trained models, or if circuits that do well on the metric are different from circuits we need to remove from a model to remove a capability? This is an opportunity to do science, but seems like it could get messy and slow.
Tuesday 12⁄5
I spent much of the last week being confused about some results. Basically, I wanted to test the subnetwork probing code on the simplest benchmark, which was the tracr-reverse task. The circuit has 11 edges. But when I tried subnetwork probing with resampling ablation, it got nearly perfect loss on a circuit with only 2 edges: token_embeds -> a3.0_v -> resid_post! I suspected a bug in my code and wrote more and more tests, but it turned out that the trivial 2-edge circuit was actually correct.
tracr-reverse circuit from IOI paper.It turns out that on any valid input, the output of 2.mlp_out (<m2> in the diagram) is identical. Since 2.mlp_out and pos_embeds are both the same between the original and resampled circuits, resampling doesn’t disrupt the circuit at all. Maybe tomorrow I’ll dig into this further, but I don’t know how the RASP program corresponds to blocks in the compiled tracr model.
Edge-level SP is now running; this is another thing to do tomorrow.
I got into Adria and Alex Turner’s MATS streams, but not Buck Shlegeris’s or Stephen Casper’s. Meeting tomorrow with Alex to discuss logistics, then decide by Friday, then fly out for NeurIPS.
I like the idea of a public research journal a lot, interested to see how this pans out!