Paper: Understanding and Controlling a Maze-Solving Policy Network (ICLR-2024)

Reviewer comment (Experiments):
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. There is doubt about whether these results can be applied to other agents in other domains, and no simple way is seen to achieve this. Thus, the work would have a very low impact on the community.

Author response:
**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.

Reviewer comment (Reproducibility):
The presented analysis is based on the 11 layers found by visual inspection. While suitable for a single experiment, this method cannot be applied broadly.

Author response:
**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.

Reviewer comment (Evaluation):
The numerical results are far from convincing. If the logistic regression can predict the goal reachability in 82% compared to naive 71%, then few hard cases were actually explained. There are likely more important features missing from the analysis.

Author response:
**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.

Reviewer comment (Experiments):
The reviewer would like to see the definition of around 20 different features, achieve over 95% accuracy, and then identify those that contribute most. The current work still misses important features.

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

Reviewer comment (Evaluation):
In analyzing the interventions in Section 3.1 (and others as well), there seems to be little difference between no intervention and intervening on all 11 channels. This does not convincingly demonstrate control of the policy, only some correlation between intervention and intended behaviour.

Author response:
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.

Reviewer comment (Evaluation):
The analysis would be more useful if it exceeded the impact of moving the cheese, which would allow for a confident claim of controlling the policy. Currently, it seems more like a bias than control.

Author response:
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.

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.

Reviewer comment (Novelty):
Overall, the reviewer agrees that _some_ features contributing to behaviour were identified, but as acknowledged by the authors themselves, there are clearly many more, reducing the contribution of the work.

Author response:
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.

Reviewer comment (Experiments):
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. This would weaken the value attributed by the article to its exploration of the representation of the goal within the policy. This is because the phenomena mentioned in the article are only specific to a particular parameterized strategy rather than a general family of strategies trained on that environment.

Author response:
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].

Reviewer comment (Evaluation):
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.

Author response:
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%)”.

Reviewer comment (Presentation):
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.

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

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
Why did you choose activations after the first residual block of the second IMPALA block as your target? Are there any insights or observations behind it?

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

Reviewer comment (Experiments):
As authors said, we can control the behavior of the policy by combining different forward passes of the network, is it possible to control the policy to exhibit more flexible behavior in this way? For example, let the policy tend to move toward the upper-left corner (neither the position of the goal nor the "upper-right corner" bias introduced in the training phase).

Author response:
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.

Reviewer comment (Experiments):
The paper's deep analysis of the maze-cheese task and policy, while a strength, is also a main weakness because many of the analyses and experiments depend on the authors' knowledge of the task and their design of the policy, raising the question of whether these findings hold true for more realistic, complicated, and relevant tasks.

Author response:
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.

Reviewer comment (Experiments):
It is unclear if the particular methodology used in the paper can be applied to other RL agents beyond the specific case studied.

Author response:
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.

Reviewer comment (Experiments):
There is no mention of seed variance, raising the concern that the findings might only emerge with a specific seed, and leaving the generality of these results across RL agents (even within the same Maze task) in question.

Author response:
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.