What progress have you made since your last update?
Finished deliverable 2.1. It was scoped as:
"Does SCA work in a transformer at all? Anchor a concept such as red across the residual stream in the color-mixing task; probe each layer for the anchored concept, and confirm that completion accuracy (predicting the correct result color) matches an un-anchored baseline."
tl;dr: Yes, it works.
I tested the method piece by piece over six preregistered experiments (2.1.6-2.1.11; see Appendix A below).
Anchoring. Residual stream probes show that the target concept (red) aligns with the anchor as expected. A repulsive term reserves the subspace, keeping the alignment graded (less-red colors are less aligned; non-red colors are nearly orthogonal).
Labelling. Sparse, noisy labels suffice, as they did in M1. Labels indicate "there is something red in this equation" (but not which tokens). Smooth minimum pooling allows the regularizer to target the "right" token, attracting it to the anchor position.
Accuracy. Guiding the concept during training left the model's accuracy on its actual task essentially unchanged (it mixes colors accurately). I searched for the point where that stops being true, and the main experiments are operating well below it.
I have not yet tested whether the anchored concept can be intervened on (e.g. to prevent the model from completing red equations in particular). That is scheduled for D2.2.
Delivered as source code and experiment reports (on GitHub).
What are your next steps?
I will begin work on D2.2: anchoring operations, and intervention (selective suppression of capabilities).
I also plan to post on LessWrong about the work so far.
Is there anything others could help you with?
Same as last time: feedback from other researchers. I'm happy to discuss if that's easier than reading the reports.
I plan to seek feedback directly from some people. I said that last time too, but I didn't feel like I had enough results to talk about. I think I do now.
Appendix A. Anchored experiments
Every anchored experiment was preregistered in Git, with the method, hypotheses and gates frozen before running. Several have hypotheses that failed, which informed the next design.
2.1.6: Our first anchored transformer. We anchor red as we did in autoencoders, but this time the noisy labels are applied to whole equations (according to the first operand), and the model/optimizer has to figure out which tokens that applies to on its own. It worked: red was decodable from the residual stream where we expected it to be. But without the repulsive terms (omitted on purpose), the whole color space is dragged toward the anchor.
2.1.7: Reserving the subspace, and narrowing the anchor. 2x2 factorial: 1. We restore the anti-subspace term from M1, which pushes all samples away from the anchor to reserve the space for the labelled concept. 2. We also try limiting the anchor term to only affect the first operand (even though this would be unrealistic for language model application). Both work, but 2. worked better, and their effects stack somewhat. Figure: red is the most aligned with the anchor, but the other colors come along for the ride when the pull is not limited to the first operand.
2.1.8: Repulsion tuning. We tested anti-subspace schedules (trailing timing and strength) to find an operating point that contains the cube-wide drift without reducing selectivity. Holding the repulsion high for longer contains the drift and keeps the margin, which nothing before this did. It's unclear if the response is well-graded.
2.1.9: Softmin sequence pooling. Anchoring on op1 alone worked best so far, but in natural language we won't know which tokens hold the concept. So we pool over sequences with a soft minimum (mellowmax) and let the pull choose its own position. At the embedding it picks op1 unaided, the operating point stays healthy, and grading improves. Only the softest pooling stays graded in every run.
2.1.10: A label that doesn't point. Either operand can trigger the label, and the pooled pull finds the red operand line by line: the weight profiles track the label groups, and selectivity matches the slot oracle. Figure: mellowmax allocates the pull to the red operand automatically, whichever that is.
2.1.11: A hyperparameter survey to close out D2.1. The schedules and weights in the recipe were inherited piece by piece, some from M1, and it was unclear which were needed and what the best values were. We ablate first by replacing schedules with constants, and dropping training epochs. Then we run a search over what is left (anchor weight, pooling temperature, repulsion dose). The anchor schedule and half the epochs could go. The anti-subspace schedule must be kept: a constant delivering the same total dose loses grading.
Appendix B. One more unanchored experiment
2.1.5: Do disjoint vocabularies learn the same representations? A model trained on named and hex colors, but not both in any one equation, learns both tasks well but uses different latent representations. Curiously, a variant that included both types could mix colors between the vocabularies, but the representations remain separate. There's more to discover down this path, if we want to learn more about the latent space. Figure: residual stream probes at all depths show that the named colors are represented as a unified concept, while hex colors are represented (and computed) one channel at a time.