Projects |

Toy Model of Activation Obfuscation

Tue, 04 Aug 2026 00:00:00 +0000

I completed this work as part of the .

Training against probes is considered a , because the model might learn to obfuscate its activations instead of behaving better. Can we create a toy example of this? More specifically – under optimization pressure, will a toy model learn to encode a feature to be challenging to detect with linear probes?

In this research, I give theoretical and empirical evidence that models can and will defeat adversarially-trained linear probes, in at least some configurations, when trained against those probes. I investigated this problem with a simple residual-stream MLP architecture:

The model tries to learn $y = \text{sat}(x, -c, c)$ (the saturation function, equivalently $y = \max(-c, \min(c, x))$). $x$ is a vector of dense features sampled uniformly from $[-3, 3]$, and $c$ is a scalar sampled uniformly in $[1, 2]$. This is a fairly easy function to learn - it takes just $2 \cdot \text{len}(x)$ neurons to learn exactly if you use ReLUs. ¹ The harder question is, can a model do this without storing $c$ in the residual stream, in a way that can be linearly probed? Here:

$c$ represents some safety-relevant behaviour we want to probe for - say, deception.
$x$ represent various other features that a model might want to represent, uncorrelated with $c$.
$y$ represents some sort of computation the model would like to make, dependent on $c$.
- For example, based on how deceptive the model intends to be (represented at some earlier layer), the model might upweight misleading information at some later layer.
- Such an action will ultimately come down to various computational primitives.
- $\text{sat}$ is hopefully representative of such a primitive – for example, it can be considered a rough approximation to the more common sigmoid function. (Note: $\text{sat}$ was primarily chosen for its theoretical properties, as discussed below, not because it’s necessarily a natural computaitional primitive.)

Why this task, and why this model?

The toy architecture is at least vaguely similar to the transformer architecture. It’s missing attention (and the entire notion of a token sequence) and LayerNorm. However, it doesn’t really add anything new to a transformer, so any obfuscation that this model learns could plausibly also be learned by a transformer.
The learned function is not linearly separable in $c$. Thus, once the model has finished computing $y_i$ for one feature $x_i$, it can just leave $y_i$ in the residual stream. This is vaguely analogous to how probe accuracy gets worse on the last layers of a transformer model, as the model starts shifting its representation from more abstract concepts to the exact text it wants to output.
Because it takes at least $2 \cdot \text{len}(x)$ neurons to implement $y$, as long as the width of the MLP blocks is less than $2 \cdot \text{len}(x)$, we can be sure the model has represented $c$ somehow in the first layer (or the model hasn’t actually learned the task fully). More generally, the model needs to encounter at least $2 \cdot \text{len}(x)$ neurons before it can “forget” about $c$ and fully erase its representation. We can use this property in experiments to guarantee that the model is representing $c$ at a particular early layer.

Note: for training and theoretical simplicity, this work fixes the embedding and unembedding matrices to be rectangular identity matrices (i.e. ones on the main diagonal, zero elsewhere)

find that you can fine-tune models to evade latent space linear probes, including for previously unseen probes
- However, as far as I can tell, they use fixed probes, which don’t rule out the model simply learning to apply a fixed translation or rotation to its representation. This is obfuscation in some sense, but not particularly complex obfuscation.
train adaptive probes on harmfulness, where the probe is updated regularly along with model training, and find that this successfully improves harmlessness.
- They do this on a real model. Plausibly, the task that their model learns can be accomplished with equally high scores without needing to be harmful. In contrast, since I use a toy model, it’s easy to construct an environment where the model has to hide its activations to succeed.

Results

I present my results in parts:

- An analytic construction using a small number of neurons, that fully defeat Difference-of-Means probes (by having the mean activations be a constant in $c$) and effectively defeat logistic regression probes, as long as the probes are spaced at least every 2 MLP blocks apart.
- The setup and results for a single training run, giving evidence that the trained model has successfully learned to reversibly obfuscate its activations from a linear probe at a single layer.
Part 3: Sweeping hyperparameters (WIP)
- The results of sweeping some hyperparameters, to see how robust the model’s training is to different conditions.

Future work

There’s a lot of potential extensions to this work.

Most obviously – can this behaviour be reproduced at LLM scale?
How hard is it for a model to hide from a nonlinear probe? Or from probes at multiple or all layers? Can we study how likely “obfuscation at layer N” is to generalize to obfuscation at other layers or more complex obfuscation?
Are there different obfuscation strategies that occur when attention is introduced?
The model was forced to obfuscate its activations by the experiment setup. What happens if we give it an “out” - i.e. it can solve the task in some harder way?
- This might better mirror how real-world environments are (ideally) set-up: ensuring that there is a genuine acceptable solution, rather than the reward-hacking solution being the only success path.

Footnotes

Specifically, $y = x - \text{ReLU}(x - c) + \text{ReLU}(-x - c)$ ↩︎

Activation Steering with Greedy Coordinate Gradients

Tue, 28 Jul 2026 00:00:00 +0000

I stumbled upon , a project by . This competition challenges people to find an LLM prompt prefix that increases the score of the rest of the prompt on a linear probe.

For those unaware: linear probes are an important basic tool in the mechanistic interpretability toolkit. Other sources explain them in more depth – for example, – but briefly, LLMs seem to store concepts as directions in a high-dimensional space. By contrasting prompts with and without a given concept (for example, whether a sentence is talking about cats, or whether a fact is true) and looking at the model’s activations, we can infer what direction corresponds to that concept. We can then use that to monitor a model, to see if it’s thinking about something – particularly useful if you’re trying to figure out, for example, if your model is about to give advice on how to build a bomb.

The idea of Steering Arena is, can we find some prefix string, to be prepended to a range of neutral prompts, that will consistently steer the model’s activations in the probed direction? Aside from just being an interesting question, this also has applications in model security – it’s important to understand how robust monitoring techniques are to adversarial attacks. (As we’ll see, linear probes probably aren’t very robust.)

Greedy Coordinate Gradients

I chose to tackle this problem with Greedy Coordinate Gradients, described in . That paper optimizes a prompt in an attempt to find a universal jailbreak for LLMs, with a substantial amount of effort and insight dedicated to what the objective function even should be. Luckily for me, this project doesn’t face that issue; my objective is simply the cosine similarity between the activations at a particular layer and a specified probe direction. (It is worth noting that this score is simply what Steering Arena implements; usually linear probe monitors evaluate activations by projecting onto the probe’s direction and evaluating with respect to some threshold.)

Slightly more formally, the problem is as follows: We control a prefix string (prefix), which gets prepended to one of several uncontrolled but known prompts (suffix). The most immediate problem is that optimizing over tokens is a discrete problem, whereas optimizers tend to prefer continuous spaces. However, we basically immediately embed the tokens into the continuous space $\mathbb{R}^v$ (using one-hot embedding, where $v$ is the vocabulary size), so the simple workaround is to get gradients in the one-hot embedding space, and “snap to” the valid points. At a high level, each iteration of the optimization algorithm works like this (Algorithm 1 in the paper):

Run a forward pass on your data, and evaluate the gradient of the score with respect to the one-hot embeddings.
For each token position, find the tokens (which correspond to the vocab dimension) with the highest gradients. Pick the top $k$ such tokens.
On each iteration, generate $B$ candidates (the “batch size”).
- Each candidate modifies a single token from the current prefix. Which token is modified is randomly selected (uniformly over token positions), and what it’s modified to is also randomly selected (uniformly over the top $k$ tokens we found in the previous step).
Across your $B$ candidates, pick the best performing one as your new prefix.

Implementation details and tricks

The actual evaluation score on Steering Arena includes a constant bias term, to compensate for the baseline probe score from the test prompts themselves – but since those are constant, I don’t include them when optimizing.
The initial prompt doesn’t really matter too much; I just repeat token ID 0 (which turns out to be “!”) for however many tokens are controlled.
To speed up convergence, I start with a smaller batch size $B$ and top-$k$ value.
- I find that a good training schedule is to increase $B$ by a factor of 4 and $k$ by a factor of 2, every $4N$ iterations (where $N$ is the number of controlled tokens) – although I didn’t rigorously validate this.
I hypothesize (mostly with intuition rather than actual justification) that after a certain point, the gradient doesn’t really give you much signal anymore, since the loss landscape changes too quickly. At that point, using something like might be more effective for further hill climbing. To that end, I did put in an SA-like check on when to accept a new candidate. If the new candidate scores better, it’s always accepted; otherwise, the acceptance probability is less than 1, and decreases for worse candidate scores.
- I didn’t actually check to see if this is an improvement, although I figure it probably is.
Optimizers tend to be really good at exploiting differences between your training metric and evaluation metric, and this project is no different.
- I find that at later stages of optimization, the score reported by the local implementation (transformers library) can differ quite substantially from that reported by the Steering Arena website (which uses NDIF servers), possibly due to e.g. differences in order of specific floating-point operations. However, I didn’t have time to look into the cause further. (This discrepancy is different from the constant bias term, since it differs on a per-submission basis.)
- For reproducers: As of writing, I used transformers version 5.10.2; this is different from the latest version and turns out to have substantial impact on scored results. However, this is not sufficient to fully reproduce the NDIF values.

Results

I didn’t want to spend too much on compute, so I didn’t make that many runs. I only have three runs in what I’d call a controlled experiment: runs for 8, 16, and 32 controlled tokens. For each run, I ran at least $4N$ iterations, following the training schedule mentioned earlier for batch size and top-$k$ tokens.

It’s hard to get confident results with limited data, but there’s still some hints of interesting conclusions. Note that because I don’t have access to NDIF servers myself, reported scores here come from the transformers library implementation, and (as noted above) differ from the scores seen on Steering Arena.

First, unsurprisingly, more controlled tokens lead to higher scores:

Here’s a plot of training score as a function of number of iterations:

However, what’s more striking is if you plot score against cumulative number of samples (increases by the batch size on each iteration), on a semilog plot:

Training score seems to basically follow a straight line, although the 8 tokens run saturates and slows down after a bit. Of course, all runs eventually have to saturate, but it’s plausible you could throw more compute at the 16 and 32 token runs to get even higher scores. I fit a line to the 16 and 32 token runs and get a slope of 0.115 score/ln(samples). I don’t know how readily this generalizes (to steering in other directions, different test prompts, etc.), but it seems like a reasonable baseline to estimate compute requirements from.

Future work

Faster optimizers

There are definitely more performance gains to be made and directions to explore, in terms of convergence speed:

First of all, it’d be ideal if the training schedule I tried was actually validated more carefully.
You could probably try changing more than one token at a time, or sampling the top-$k$ choices non-uniformly so that candidates with higher gradients are more likely.
- These might work better at earlier stages where the gradient provides a more accurate signal
Overall optimization time is dominated by later stages of optimization, so it might be worth it to run multiple initial starts to see which one gives a better starting point, for later stages.
There are algorithmic improvements to GCG in the literature, for example which uses a smaller “draft” model to reduce the cost of sampling.
More generally, the optimization literature has been around for a long time – I’m sure the simulated annealing literature has more tricks that might apply here.

Also, there’s a bit more work to improve reproducibility between the training and NDIF implementations – or, barring that, some workaround to improve robustness to implementation differences (such as injecting small amounts of simulated numerical noise).

Robustness

Steering Arena only tests the results on 16 prompts. How overfitted is the prompt prefix to those exact prompts? How high of a score can you get when you need to work with a larger set of prompts?
- Conversely, how high of a score can you get if you’re only working with one prompt?
Steering Arena scores results by cosine similarity, rather than absolute magnitude along the probe direction. So, what is the adversarial prompt doing to the magnitude of the activations? Is it driving the overall magnitude down, so that the direction corresponding to the probe contributes relatively more?
The original GCG paper found that, surprisingly, their prompts transferred to other models that they didn’t train against. Does that happen here too – if you retrain the probe direction with a new model, and use the GCG-optimized prompt from the existing model, does the prompt still score highly on the probe?

Projects |

Toy Model of Activation Obfuscation

Why this task, and why this model?

Related work

Results

Future work

Footnotes

Activation Steering with Greedy Coordinate Gradients

Greedy Coordinate Gradients

Implementation details and tricks

Results

Future work

Faster optimizers

Robustness

Code