Building Better Activation Oracles

Work done for our MATS 10.0 Sprint project - mentored by Neel Nanda and Adam Karvonen Huggingface, Github TL;DR: We have improved the original Activation Oracle (AO) training regime by training on on-policy rollouts, improving the conversational dataset, feeding more layers (following the approach by Niclas Luick) and making a small change to the injection formula. We also open source our evals, which we believe are currently the most comprehensive evaluation of AO quality called AObench. The capability improvements are marginal, but quality of life improvements are quite substantial. If you want to play around with the new AOs, we recommend you use this one, if you want to play with our new Activation Oracles live, we will host them for a week on ao.celeste.computer. Alternatively, you can self-host our web interface. Activation Oracles (AOs) by Karvonen et al. are fine-tuned LLMs that can receive the original target LLM’s activations as input and answer natural language questions about them. However, they are plagued by various issues, which limit their usefulness as an off-the-shelf tool for interpretability research. For our MATS Sprint, we set out to work on these issues.  Issues with current Activation Oracles In Current activation oracles are hard to use, Arya Jakkli demonstrates scenarios where AOs are hard to work with. We focus on addressing two of the issues pointed out: Hallucinations: The AO will output false information.  Vagueness: The AO output will be generic (therefore unfalsifiable) and will not answer the user’s question. In addition, they are difficult to evaluate because of the problem of text inversion: the model infers the surrounding text and answers based on that, just as any black box oracle (i.e. a method that only receives text) could, rather than extracting specific info from activations. As part of our evaluations, we focus on some of his specific tasks, you can find details on our evaluations in the Appendix.  Our approaches to improving AO training  A better conversational finetune than LatentQA To make the Activation Oracle be able to answer natural language questions, you need a dataset consisting of questions and answers about activations. The original paper used LatentQA to this end. However, we found that this dataset was of low quality, likely incentivizing vagueness: The model is given a complicated prompt, and then a specific question is asked about this prompt. We think the answers to the questions LatentQA poses are often not easily retrievable from activations, which makes it a difficult task for the AO, not incentivizing much beyond text inversion, and may even directly incentivise hallucinations/guessing if the relevant info is not present. The questions are not about on-policy data, but about specifics of a user prompt: this does not target the model’s internal reasoning.   It was generated by o1, a now outdated model. Our solution: Questions about unarticulated completions  We constructed a new conversational dataset that attempts to address all of these concerns. Because we don’t want the questions learned to be trivially answerable from adjacent tokens (text inversion), we construct QA pairs as follows: To construct this task, a separate LLM (Sonnet 4.6)  is given the target model’s chain-of-thought (CoT), and is instructed to split the chain of thought into a prefix and suffix, and to write a question about the suffix. It is instructed to do this in a way such that the question is hard to answer purely from the text of the prefix (i.e. to avoid text inversion), but plausibly answerable from the prefix’ activations (solvability). You can explore our dataset here.  We ablate the effect of this task by replacing only LatentQA in Adam's recipe (leaving everything else the same) and notice a significant uplift, across the board on our AObench evaluations. We find that the responses are more specific and the resulting model is less vague, and responds better to instructions. Layer choice/feeding multiple layers to the AO As Niclas Luick demonstrated, feeding multiple layers at a time during training and inference increases Activation Oracle capability. Adam originally fed activations randomly selected from either layer 25%, 50% or 75% of total model depth. Since most features live around the 55-80% layer ranges, we suspected a layer sweep could be important. Indeed, we find that AO performance peaks at layer 22 (62%). Feeding 5 contiguous layers from layer 21-25 causes further uplift. Interestingly, the largest uplift is on model diffing tasks. We’d like to point out that training a multi-layer Activation Oracle can cause an increase in training time due to longer context, and that most gains can be had by simply choosing a layer at ~65% depth.  Training on on-policy data To train Activation Oracles, we need scalable unsupervised training tasks. A common way to achieve this is to predict past and/or future tokens from the activations, known as past or fu