Understanding and Controlling a Maze-Solving Policy Network

Dear Reviewer,

Thank you. We’re glad you found our work “well-written” and addressing “an important topic”. We’ve made a number of clarifications below, which we hope will allow you to improve your evaluation score.

The scope of the paper is far too narrow. The study is exhaustive, but it is focused on a particular architecture in a particular environment. Is there any chance of applying those results to other agents in other domains? [...] Besides, the presented analysis is based on the 11 layers found by visual inspection. While suitable for a single experiment, this method can't be applied broadly.

We believe the purpose of our work is to develop an understanding of the goals and goal representations of neural networks. Carefully studying a particular network is one way of doing this, similar to studying “model organisms” in biology.

In particular:

• Finding that the goals of AI systems are redundant and retargetable is an important finding. We find this holds across 15 different agents. 

• There is further evidence of retargetable goals in other agents, like LLMs [1,2,3]. We will add these to the paper to contextualize the work. We show these findings also hold here.

the numerical results you obtain are far from convincing. If the logistic regression can predict the goal reachability in 82% compared to naive 71%, then quite few hard cases were actually explained.

Our aim is not to perfectly predict the generalization of the network. Instead, we show that it is possible to generalize depending on environmental factors. We think the interesting part of these results is showing that the pursued goal of the network depends on the environment. One might have thought that this would be random. Moreover, we never claim to perfectly predict the generalization.

Also, when you analyze the interventions in Section 3.1 (and others as well), I see quite little difference between no intervention and intervening on all 11 channels.

In Fig. 83 (Appendix E), we show that intervening on all 11 channels results in increased retargetability in all cases, even as we grow far away from the top-right path. Therefore we find substantial differences in the behavior when intervening. We admit the differences in the heatmaps can be hard to see; we will improve this. Additionally, there is a notable difference between no intervention and intervening on the 11 channels.

Furthermore, I think this analysis would be useful if you exceed the impact of moving the cheese. 

This evidence shows that the goals of a network can be partially retargeted via online activation editing. To our knowledge, we are the first to give a proof of concept of this ability. We think this is an important finding because it develops insight into alignment methods that reuse existing network circuits.

Regarding defining many features and selecting those which contribute the most, we do initially regress upon ten features. Appendix B contains further details. 

Only then you can confidently claim that you can control the policy and convince me.

We agree that we cannot perfectly control the network in this way. That is why we describe our network as “partially” controlling the policy. We do not claim to be able to fully control the policy, only that we provide first steps along a pathway to greater control through intervention techniques based on network activations.

Can the results that you present be generalized to other architectures and environments?

There is some evidence showing these results generalize to LLMS [1,2,3].

Can all those steps you describe be automated?

Potentially! We’ll add this to the paper under future work.

Are there any reasons for the network to learn the goal location as a separate feature?

Yes, the convolutional layers do present a “privileged basis” (the channel activations). However, visual inspection did not reveal an analogous “top-right” feature. Even though our results strongly suggest that the network tracks the location of the top-right accessible location. Therefore, it is interesting that the cheese was so clearly and redundantly represented.

Why did you choose this specific layer to inspect? Are those observations valid for many layers and you've just chosen an arbitrary one, or was it a careful decision?

We first developed the cheese vector technique, and found it worked best at this layer. So we looked into it in more detail.

Thank you for the review. We hope our clarifications are useful, and are looking forward to hearing what you think.

Best, The Authors

1 - Turner, A., et. al. (2023). Activation addition: Steering language models without optimization. arXiv preprint arXiv:2308.10248.

2 - Li, Kenneth, et al. "Inference-Time Intervention: Eliciting Truthful Answers from a Language Model." arXiv preprint arXiv:2306.03341 (2023).

3 - Rimsky, N. Reducing sycophancy and improving honesty via activation steering (2023), AI Alignment Forum

Dear Reviewer,

Thank you for your careful analysis and feedback. We’d like to make some clarifications, which we hope will allow you to increase your evaluation score and vote for acceptance.

The paper lacks further validation of whether the discovered intrinsic representation of the goal in the pre-training policy can be generalized to different policies

We find evidence that the cheese vector technique transfers across 15 different policies. This shows that our findings generalize across policies. Moreover, we perform our behavioral analysis across the same 15 policies, and repeatedly find evidence that the goals of these policies are context-dependent. We believe these are valuable findings. There is other evidence that suggests that we can align policies by modifying activations during forward passes, a technique sometimes known as activation engineering [1, 2].

The experimental phenomena and conclusions of the author cannot fully support their core contribution. We cannot demonstrate from the experiments that the 'intrinsic representation' of the goal in the policy can be represented by the activations of these 11 channels that are selected by human visual inspection

We would like to clarify that we have visual and quantitative evidence that these channel activations depend primarily on the location of the cheese. We do not mean to claim that all cheese information passes through these channels, but these channels are used by the network to track cheese. The visual evidence we include in the appendix, but we also include resampling results (Pg. 4), where we state “Across 200 mazes, resampling the cheese-tracking channels from mazes with a different cheese location changes the most probable action at a decision square in 40% of cases, which is much more than when resampling from mazes with the same cheese location (11%)”.

3: The presentation of some of the experiments is confusing, and it may be helpful to detail the setup of these experiments in the appendix to help understand the work.

Thank you. We’ll improve this. Are there any specific things you’d like to see improved?

Question 1: Are the activations influenced by the agent's location or different steps? When we aim to control the policy by adjusting the activations, should we modify the activation of the initial state or all other states?

The activations depend on the agent's location. We adjust the activations at all time steps (i.e., the policy network, as consistently used by the network), which defines a new policy.

Question 2: Does the number of the most effective channels change when the size of the maze is altered? Or will the features found in 11 selected channels maintain consistency with respect to such changes in experimental settings

The mazes as seen by the network are all the same size. We show the “human-friendly view”, which might be smaller than the full game grid if there is padding at the edges. As a result, the network always sees RGB images of the same size, and the number of effective channels thus do not change as the maze size is altered.

Question 3: Thanks. We’ll clarify this in the camera ready. By the geometric mean, we mean the following equation:

$

P_\text{path}(s_t\mid \pi):= \sqrt[t]{\prod_{i=0}^{t-1}\pi(a_i\mid s_i)},

$

where $s_0,s_1,\ldots, s_t$ is the unique shortest path between $s_0$ and $s_t$, navigated by actions $a_i$. If $s_0=s_t$, then $P_\text{path}(s_t\mid \pi):=\pi(`no-op`\mid s_0)$. The probability the agent goes through a path depends on the probabilities of each step multiplied together, so we used the geometric mean.

Question 4: Empirically we found that the cheese vector worked best there. We’ll clarify this in the paper.

Question 5: Thank you. It is possible to compose some activation additions. We can subtract the cheese vector and add the top-corner vector together, which leads to an increase in top-right corner seeking and a decrease in cheese-seeking. We’ll add this example to the paper, but we leave better understanding of composition to future work.

[1] Turner, A., Thiergart, L., Udell, D., Leech, G., Mini, U., & MacDiarmid, M. (2023). Activation addition: Steering language models without optimization. arXiv preprint arXiv:2308.10248.

[2] Rimsky, N. Reducing sycophancy and improving honesty via activation steering (2023), AI Alignment Forum

We would first like to thank the reviewer for their careful consideration and analysis of our submitted work. We are pleased that you thought our work addresses “a very important problem”. We’ve made a number of clarifications below, which we hope will allow you to increase your evaluation score.

if these findings hold true for more realistic, complicated and relevant tasks, and 2) if the particular methodology used here, can be applied to other RL agents.

Thank you. We believe carefully understanding one particular architecture shed light on generic properties of deep learning systems, much like “model organisms” used in biology. To that end, we would like to clarify that: A number of our findings hold for multiple agents. We perform our behavioral analysis across 15 different trained networks (with substantially varying historical goal regions). We repeatedly find evidence of contextually activated goals. We believe this finding is an important result, because it helps understand the behavior of AI assistants. Some of our methods apply to other architectures. We show that we can re-use existing circuits within the network to control network behavior. The technique of combining network forward passes can be used on other networks, such as used in Turner et. al (2023) with large language models.

As such, we believe better understanding the maze-solving policy network allows us to understand the goals of other RL systems, including LLMs. Applying our approach to multiple seeds provides some additional evidence that the goals of AI systems are redundant, context-dependent, and retargetable. We believe these findings have significant implications for AI alignment. Of course, we acknowledge that only studying one environment is a limitation of the work. We will update the work accordingly for the camera ready, including better contextualisation with the related work, including discussions of how other related work provides evidence for our hypotheses.

Top-right corner motivational vector - is the definition ordering swapped? Figure 7 - the columns seem to be out of order.

Yes, we’ll fix this. Thank you.

Please let us know what you think, and whether these changes are enough for you to vote for acceptance. We think that a revised version of the manuscript will be of great interest to the ICLR community.

Thank you, The AuthorsWe would first like to thank the reviewer for their careful consideration and analysis of our submitted work. We are pleased that you thought our work addresses “a very important problem”. We’ve made a number of clarifications below, which we hope will allow you to increase your evaluation score.

if these findings hold true for more realistic, complicated and relevant tasks, and 2) if the particular methodology used here, can be applied to other RL agents.

Thank you. We believe carefully understanding one particular architecture shed light on generic properties of deep learning systems, much like “model organisms” used in biology. To that end, we would like to clarify that:

1. A number of our findings hold for multiple agents. We perform our behavioral analysis across 15 different trained networks (with substantially varying historical goal regions). We repeatedly find evidence of contextually activated goals. We believe this finding is an important result, because it helps understand the behavior of AI assistants.

2. Some of our methods apply to other architectures. We show that we can re-use existing circuits within the network to control network behavior. The technique of combining network forward passes can be used on other networks, such as used in Turner et. al (2023) with large language models.

As such, we believe better understanding the maze-solving policy network allows us to understand the goals of other RL systems, including LLMs. Applying our approach to multiple seeds provides some additional evidence that the goals of AI systems are redundant, context-dependent, and retargetable. We believe these findings have significant implications for AI alignment. Of course, we acknowledge that only studying one environment is a limitation of the work. We will update the work accordingly for the camera ready, including better contextualisation with the related work, including discussions of how other related work provides evidence for our hypotheses.

Top-right corner motivational vector - is the definition ordering swapped? Figure 7 - the columns seem to be out of order.

Yes, we’ll fix this. Thank you.

Please let us know what you think, and whether these changes are enough for you to vote for acceptance. We think that a revised version of the manuscript will be of great interest to the ICLR community.

Thank you, The Authors