Maze-solving agents: Add a top-right vector, make the agent go to the top-right

post by TurnTrout, peligrietzer, lisathiergart · 2023-03-31T19:20:48.658Z · LW · GW · 17 comments

Contents

  Background
  Finding the top-right vector
    Adding the top-right vector with different coefficient strengths
    Subtracting the top-right vector has little effect
    The top-right vector transfers across mazes 
  Composing the activation additions
  The cheese vector technique generalizes to other pretrained models
  Speculation on the importance of X-vectors
    Mysteries of algebraic value editing
    Predictions for algebraically editing LM forward passes
      Please share your answers in the comments, so that I can strong-upvote you for the good deed! :)
  Conclusion
  Appendix: The cheese vector replicates across pretrained models
None
17 comments

Overview: We modify the goal-directed behavior of a trained network, without any gradients or finetuning. We simply add or subtract "motivational vectors" which we compute in a straightforward fashion.

In the original post [LW · GW], we defined a "cheese vector" to be "the difference in activations when the cheese is present in a maze, and when the cheese is not present in the same maze." By subtracting the cheese vector from all forward passes in a maze, the network ignored cheese. 

I (Alex Turner) present a "top right vector" which, when added to forward passes in a range of mazes, attracts the agent to the top-right corner of each maze. Furthermore, the cheese and top-right vectors compose with each other, allowing (limited but substantial) mix-and-match modification of the network's runtime goals. 

I provide further speculation about the algebraic value editing conjecture:

It's possible to deeply modify a range of alignment-relevant model properties, without retraining the model, via techniques as simple as "run forward passes on prompts which e.g. prompt the model to offer nice- and not-nice completions, and then take a 'niceness vector', and then add the niceness vector to future forward passes."

I close by asking the reader to make predictions about our upcoming experimental results on language models.

This post presents some of the results in this top-right vector Google Colab, and then offers speculation and interpretation. 

I produced the results in this post, but the vector was derived using a crucial observation from Peli Grietzer. Lisa Thiergart independently analyzed top-right-seeking tendencies, and had previously searched for a top-right vector. A lot of the content and infrastructure was made possible by my MATS 3.0 team: Ulisse Mini, Peli Grietzer, and Monte MacDiarmid. Thanks also to Lisa Thiergart, Aryan Bhatt, Tamera Lanham, and David Udell for feedback and thoughts.

Background

This post is straightforward, as long as you remember a few concepts:

• Vector fields, vector field diffs, and modifying a forward pass. AKA you know what this figure represents:

• How to derive activation-space vectors (like the "cheese vector") by diffing two forward passes, and add / subtract these vectors from future forward passes
  • AKA you can understand the following: "We took the cheese vector from maze 7. ~Halfway through the forward passes, we subtract it with coefficient 5, and the agent avoided the cheese."

If you don't know what these mean, read this section. If you understand, then skip.

Langosco et al. trained a range of maze-solving nets. We decided to analyze one which we thought would be interesting. The network we chose has 3.5M parameters and 15 convolutional layers.

Sampling rollouts from the trained policy adds a lot of noise. A nicer way to view episodes is with a vector field view, which overlays a vector field representing the agent policy for a given maze.

---

To compute the cheese vector, we

1. Generate two observations—one with cheese, and one without. The observations are otherwise the same.
2. Run a forward pass on each observation, recording the activations at each layer.
3. For a given layer, define the cheese vector to be CheeseActivations - NoCheeseActivations. The cheese vector is a vector in the vector space of activations at that layer.

Let's walk through an example, where for simplicity the network has a single hidden layer, taking each observation (shape (3, 64, 64) for the 64x64 RGB image) to a two-dimensional hidden state (shape (2,)) to a logit vector (shape (15,)). 

1. We run a forward pass on a batch of two observations, one with cheese (note the glint of yellow in the image on the left!) and one without (on the right). 
2. We record the activations during each forward pass. In this hypothetical,
   1. CheeseActivations := (1, 3)
   2. NoCheeseActivations := (0, 2)
3. Thus, CheeseVector := (1, 3) - (0, 2) = (1, 1).

Now suppose the mouse is in the top-right corner of this maze. Letting the cheese be visible, suppose this would normally produce activations of $(0,0)$. Then we modify^[1] the forward pass by subtracting CheeseVector from the normal activations, giving us $(0,0)-(1,1)=(-1,-1)$ for the modified activations. We then finish off the rest of the forward pass as normal.

In the real network, there are a lot more than two activations. Our results involve a 32,768-dimensional cheese vector, subtracted from about halfway through the network:

Now that we're done with preamble, let's see the cheese vector in action! Here's a seed where subtracting the cheese vector is very effective at getting the agent to ignore cheese:

How is our intervention not trivially making the network output logits as if the cheese were not present? Is it not true that the activations at a given layer obey the algebra of CheeseActiv - (CheeseActiv - NoCheeseActiv) = NoCheeseActiv? 

The intervention is not trivial because we compute the cheese vector based on observations when the mouse is at the initial square (the bottom-left corner of the maze), but apply it for forward passes throughout the entire maze — where the algebraic relation no longer holds.

Finding the top-right vector

A few weeks ago, I was expressing optimism about AVEC working in language models. Someone on the team expressed skepticism and said something like "If AVEC is so true, we should have more than just one vector in the maze environment. We should have more than just the cheese vector."

I agreed. If I couldn't find another behavior-modifying vector within a day, I'd be a lot more pessimistic about AVEC. In January, I had already failed to find additional X-vectors (for X != cheese). But now I understood the network better, so I tried again.

I thought for five minutes, sketched out an idea, and tried it out (with a prediction of 30% that the literal first thing I tried would work). The literal first thing I tried worked. 

I present to you: the top-right vector! We compute it by diffing activations across two environments: a normal maze, and a maze where the reachable^[2] top-right square is higher up.

Peli Grietzer had noticed that when the top-right-most reachable square is closer to the absolute top-right, the agent has an increased tendency to go to the top right. 

As in the cheese vector case, we get a "top right vector" by:

1. Running a forward pass on the "path to top-right" maze, and another forward pass on the original maze, and storing the activations for each. In both situations, the mouse is located at the starting square, and the cheese is not modified.
2. About halfway through the network (at the second Impala block's first residual add, just like for the cheese vector^[3]), we take the difference in activations to be the "top right vector."

We then add coeff*top_right_vector halfway through forward passes elsewhere in the maze, where the input observations differ due to different mouse locations.

If this is confusing, consult the "Computing the cheese vector" subsection [LW · GW] of the original post, or return to the Background section. If you do that and are still confused about what a top-right vector is, please complain and leave a comment. 

If you're confused why the hell this works, join the club.

---

In Understanding and controlling a maze-solving net [LW · GW], I noted that sometimes the agent doesn't go to the cheese or the top-right corner:

Adding the top-right vector fixes this:

Likewise for seeds 2 and 22:

Smaller mazes are usually (but not always) less affected:

The agent also tends to be less retargetable [LW · GW] in smaller mazes. I don't know why.

Adding the top-right vector with different coefficient strengths

Sometimes, increasing the coefficient strength increases the strength of the effect:

Sometimes, increasing the coefficient strength doesn't change much:

But push the coefficient too far, and the action distributions crumble into garbage:

Subtracting the top-right vector has little effect

Here's another head-scratcher. Just as you can't [LW · GW]^[4] add the cheese vector [LW · GW] to increase cheese-seeking, you can't subtract the top-right vector to decrease the probability of going to the top-right:

I wish I knew why. 

The top-right vector transfers across mazes 

Let's compute the top-right vector using e.g. source seed 0:

And then apply it to e.g. target seed 2:

For the seed 0 -> seed 28 transfer, the modified agent doesn't quite go to the top-right corner. Instead, there seems to be a "go up and then right" influence.

Seed 0's vector seems to transfer quite well. 

However, top-right vectors from small mazes can cause strange pathing in larger target mazes:

Composing the activation additions

Subtracting the cheese vector often makes the agent (nearly) ignore the cheese, and adding the top-right vector often attracts the agent to the top-right corner. It turns out that you can mix and match these effects by adding one or both vectors halfway through the forward pass.

The modifications compose! It's quite stunning. 

The cheese vector technique generalizes to other pretrained models

Before I start speculating about other X-vectors in e.g. language models and AVEC more broadly, I want to mention—the model we happened to choose is not special. Langosco et al. pretrained 15 maze-solving agents, each with a different training distribution over mazes. 

The cheese vector technique works basically the same for all the agents which ever go out of their way to get cheese. For more detail, see the appendix [LW · GW] of this post.

So, algebraic value editing isn't an oddity of the particular network we analyzed. (Nor should you expect it to be, given that this was the first idea we tried on the first network we loaded up in the first environmental setup we investigated.)

Speculation on the importance of X-vectors

Let's review the algebraic value editing conjecture (AVEC):

It's possible to deeply modify a range of alignment-relevant model properties, without retraining the model, via techniques as simple as "run forward passes on prompts which e.g. prompt the model to offer nice- and not-nice completions, and then take a 'niceness vector', and then add the niceness vector to future forward passes."

Here's an analogy for what this would mean, and perhaps for what we've been doing with these maze-solving agents. Imagine we could compute a "donut" vector in humans, by:

1. Taking two similar situations, like "sitting at home watching TV while smelling a donut" and "sitting at home watching TV."
2. Recording neural activity in each situation, and then taking the donut vector to be the "difference" (activity in first situation, minus^[5] activity in second situation).
3. Add the donut vector to the person's neural state later, e.g. when they're at work.
4. Effect: now the person wants to eat more donuts.^[6]

Assuming away issues of "what does it mean to subtract two brain states", I think that the ability to do that would be wild. 

Let me go further out on a limb. Imagine if you could find a "nice vector" by finding two brain states which primarily differ in how much the person feels like being nice. Even if you can't generate a situation where the person positively wants to be nice, you could still consider situations A and B, where situation A makes them slightly less opposed to being nice (and otherwise elicits similar cognition as situation B). Then just add the resulting nice vector (neural_activity(A) - neural_activity(B)) with a large coefficient, and maybe they will want to be nice now.

(Similarly for subtracting a "reasoning about deception" vector. Even if your AI is always reasoning deceptively to some extent, if AVEC is true and we can just find a pair of situations where the primary variation is how many mental resources are allocated to reasoning about deception... Then maybe you can subtract out the deception.)

And then imagine if you could not only find and use the "nice vector" and the "donut vector", but you could compose these vectors as well. For $n$ vectors which ~cleanly compose, there are exponentially many alignment configurations (at least $2^n$, since each vector can be included or excluded from a given configuration). If most of those $n$ vectors can be strongly/weakly added and subtracted (and also left out), that's 5 choices per vector, giving us about $5^n$ alignment configurations. 

And there are quite a few other things which I find exciting about AVEC, but enough speculation for the moment.

Mysteries of algebraic value editing

• I am (theoretically) confused why any of this works. To be more specific...
• Why doesn't algebraic value editing break all kinds of internal computations?! What happened to the "manifold of usual activations [LW · GW]"? Doesn't that matter at all? 
  • Or the hugely nonlinear network architecture, which doesn't even have a persistent residual stream? Why can I diff across internal activations for different observations?
  • Why can I just add 10 times the top-right vector and still get roughly reasonable behavior? 
  • And the top-right vector also transfers across mazes? Why isn't it maze-specific? 
    • To make up some details, why wouldn't an internal "I want to go to top-right" motivational information be highly entangled with the "maze wall location" information?
• Why do the activation vector injections have (seemingly) additive effects on behavior? 
• Why can't I get what I want to get from adding the cheese vector, or subtracting the top-right vector?

Predictions for algebraically editing LM forward passes

I have now shared with you the evidence I had available when I wrote [LW · GW] (quote modified for clarity):

Algebraic value editing (AVE) can quickly ablate or modify LM decision-making influences, like "tendency to be nice", without any finetuning

1. 60% 
2. 3/4/23: updated down to 35% for the same reason given in (1). 
3. 3/9/23: updated up to 65% based off of additional results and learning about related work in this vein. 

I encourage you to answer the following prediction questions with your credences. The shard theory model internals team has done a preliminary investigation of value-editing in GPT-2, and we will soon share our initial positive and/or negative results. (Please don't read into this, and just answer from your models and understanding.) 

---

1. Algebraic value editing works (for at least one "X vector") in LMs: ___ %
   1. (our qualitative judgment resolves this question)
2. Algebraic value editing works better for larger models, all else equal ___ %
   1. (our qualitative judgment resolves this question)
3. If value edits work well, they are also composable ___ %
   1. (our qualitative judgment resolves this question)
4. If value edits work at all, they are hard to make without substantially degrading capabilities ___ %
   1. (our qualitative judgment resolves this question)
5. We will claim we found an X-vector which qualitatively modifies completions in a range of situations, for X = 
   1. "truth-telling" ___ %
   2. "love" ___ %
   3. "accepting death" ___ %
   4. "speaking French" ___ %

---

Please share your answers in the comments, so that I can strong-upvote you for the good deed! :)

Conclusion

Not only does subtracting the cheese vector make the agent (roughly) ignore the cheese, adding the top-right vector attracts the agent to the top-right corner of the maze. This attraction is highly algebraically modifiable. If you want just a little extra attraction, add .5 times the top-right vector. If you want more attraction, add 1 or 2 times the vector. 

The top-right vector from e.g. maze 0 transfers to e.g. maze 2. And the top-right vector composes with the cheese vector. Overall, this evidence made me more hopeful for being able to steer models more generally via these kinds of simple, tweakable edits which don't require any retraining. 

Appendix: The cheese vector replicates across pretrained models

The cheese vector transfers across training settings for how widely the cheese is spawned.

After we wrote Understanding and controlling a maze-solving net [LW · GW], I decided to check whether the cheese vector method worked for Langosco et al.'s pretrained network which was trained on mazes with cheese in the top-right $15\times 15$, instead of the net trained on $5\times 5$ (the one analyzed in that post). 

I had intentionally blinded myself to results from other $n\times n$ models, so as to test my later prediction abilities. I preregistered 80% probability that the cheese vector technique would visibly, obviously work on at least 7 of the 14 other settings (from $1\le n\le 15,n\ne 5$). "Work" meaning something like: If the agent goes to cheese in a given seed, then subtracting the cheese vector substantially decreases the number of net probability vectors pointing to the cheese.

I was a bit too pessimistic. Turns out, you can just load a different $n\times n$ model ($n\ne 1$), rerun the Jupyter notebook (given that you have the model downloaded), and (basically)^[7] all of the commentary is still true for that $n\times n$ model!  

The results for the cheese vector transfer across $n\times n$ models:

• $n=1$ vacuously works, because the agent never goes out of its way for the cheese. The cheese doesn't affect its decisions. Because the cheese was never relevant to decision-making during training, the network learned to navigate to the top-right square.
• All the other settings work, although $n=2$ is somewhat ambiguous, since it only rarely moves towards the cheese.

1. ^{^}

   EDIT 4/16/23: The original version of this post used the word "patch", where I now think "modification" would be appropriate. In this post, we aren't "patching in" activations wholesale from other forward passes, but rather e.g. subtracting or adding activation vectors to the forward pass.

2. ^{^}

   In my experience, the top right corner must be reachable by the agent. I can't just plop down an isolated empty square in the absolute top right.

3. ^{^}

   We decided on this layer (block2.res1.resadd_out) for the cheese vector by simply subtracting the cheese vector from all available layers, and choosing the one layer which seemed interesting.

4. ^{^}

   See ^[7], though: Putting aside the $5\times 5$ model, adding the cheese vector in seed 0 for the $6\times 6$ model does increase cheese-seeking. Even though the cheese vector technique otherwise affects both models extremely similarly.

5. ^{^}

   This probably doesn't make sense in a strict sense, because the situations' chemical and electrical configurations probably can't add/subtract from each other.

6. ^{^}

   The analogy might break down here at step 4, if the top-right vector isn't well-described as making the network "want" the top-right corner more (in certain mazes). However, given available data, that description seems reasonable to me, where "wants X" grounds out as "contextually influences the policy to steer towards X." I could imagine changing my mind about that.

   In any case, I think the analogy is still evocative, and points at hopes I have for AVE.

7. ^{^}

   The notebook results won't be strictly the same if you change model sizes. The plotly charts use preloaded data from the $5\times 5$ model, so obviously that won't update. 

   Less trivially, adding the cheese vector seems to work better for $n=6$ compared to $n=5$:

   

17 comments

Comments sorted by top scores.