Gradient Routing: Masking Gradients to Localize Computation in Neural Networks

## Abstract

Abstract Neural networks are trained primarily based on their inputs and outputs, without regard for their internal mechanisms. These neglected mechanisms determine properties that are critical for safety, like (i) transparency; (ii) the absence of sensitive information or harmful capabilities; and (iii) reliable generalization of goals beyond the training distribution. To address this shortcoming, we introduce gradient routing , a training method that isolates capabilities to specific subregions of a neural network. Gradient routing applies data-dependent, weighted masks to gradients during backpropagation. These masks are supplied by the user in order to configure which parameters are updated by which data points. We show that gradient routing can be used to (1) learn representations which are partitioned in an interpretable way; (2) enable robust unlearning via ablation of a pre-specified network subregion; and (3) achieve scalable oversight of a reinforcement learner by localizing modules responsible for different behaviors. Throughout, we find that gradient routing localizes capabilities even when applied to a limited, ad-hoc subset of the data. We conclude that the approach holds promise for challenging, real-world applications where quality data are scarce.

## 1 Introduction

As AI systems become more powerful and more prevalent, there is an increasing need to explain and control the inner mechanisms governing their behavior. To address this challenge, some researchers aim to fully understand AI systems,
either by reverse engineering the operations of conventionally trained models (Olah et al., ; Olsson et al., ) or
with inherently interpretable architectures (Koh et al., ; Hewitt et al., ; Xin et al., ). This is not necessary. If we could understand or control the mechanisms underlying a neural network’s computation with respect to a limited set of safety-critical properties, such as hazardous information or the capacity for deception, that might be sufficient to make significant safety guarantees.

To achieve targeted control over neural network internals, we propose gradient routing, a training method for localizing capabilities to chosen subregions of a neural network. Gradient routing is a modification of backpropagation that uses data-dependent, weighted masks to control which network subregions are updated by which data points. By appropriately specifying these masks, a user can configure which parts of the network (parameters, activations, or modules) are updated by which data points (e.g. specific tokens, documents, or based on data labels). The resulting network is similar to a conventionally trained network, but with some additional internal structure.

Our contributions are as follows. In Section 2, we discuss prior work on neural network modularity, unlearning, and scalable oversight. In Section 3, we define gradient routing and comment on its practical implementation. Most of the paper is a tour of gradient routing applications:

We use gradient routing to control the encodings learned by an MNIST autoencoder to split them into two halves, with each half representing different digits.

We apply gradient routing to localize features in language models. First, we train a model that can be steered by a single scalar value, showing that feature localization is possible, even with narrowly-scoped labels. Next, we present Expand, Route, Ablate, an application of gradient routing that enables robust unlearning via ablation of a pre-specified network subregion. This unlearning is nearly as resistant to retraining as a gold-standard model never trained on the task. Finally, we show that this unlearning method scales to a large (0.7B) model.

We apply gradient routing to the problem of scalable oversight (Amodei et al., ), where the aim is to train a performant policy despite limited access to reliable labels. We train a policy network by reinforcement learning to navigate to two kinds of grid squares in a toy environment, Diamond and Ghost. Using gradient routing, we localize modules responsible for these two behaviors. We show that we can steer the policy towards Diamond by ablating the Ghost module. Gradient routing trains steerable networks even when the amount of labeled training data is small, and even when the policy is able to condition on the existence of labels. As a result, our method outperforms baselines, including data filtering.

In Section 5, we discuss themes from our findings, including an observed absorption effect, where gradient routing applied to a narrow subset of data has a broader localizing effect on capabilities related to that data. Absorption provides an answer to the question: “If one has labels that are suitable for localizing undesirable computation, why not simply use those labels to filter the data?” When labels do not encompass all training data from which harmful capabilities arise, absorption means that localization can still occur, whereas filtering may be inadequate. Furthermore, localization does not explicitly influence the learned behavior of a model, a fact we exploit to achieve scalable oversight.

We conclude by noting that black-box training techniques may be inadequate for high-stakes machine learning applications. Localization techniques, like gradient routing, may provide a solution.

## 2 Related work

Training to localize pre-specified capabilities. Akin to gradient routing, work in modular machine learning trains modules to contain concepts or abilities determined in advance of training. Typically, modular architectures involve a routing function that selects modules to apply on a forward pass (Pfeiffer et al., ). Routing functions are often unsupervised, as with a typical mixture of experts setup (Jacobs et al., ; Eigen et al., ; Shazeer et al., ). However, some approaches route inputs based on metadata, creating modules with known specializations (Waibel & II, ). For example, routing has been based on (i) the modality of data in multi-modal models (Pfeiffer et al., ), (ii) language (Pfeiffer et al., ; ; Fan et al., ), and (iii) low- vs. high-level control or task type in robotics (Heess et al., ; Devin et al., ). Gururangan et al. () separate the training data of a language model by domain and assign one expert in each layer to a single domain. By disabling the expert for a domain, they are able to approximate a model that was not trained on the domain.

Other methods freeze the weights of a pre-trained model and train a newly added module, with the aim of localizing the task to the new module (Rebuffi et al., ; ; Houlsby et al., ; Bapna & Firat, ). Zhang et al. () locate capabilities in models by learning a weight mask, transfer the identified sub-network to a randomly initialized model, then train as if from scratch. By choosing a suitable sub-network, they can, for example, induce a vision model to identify ImageNet (Deng et al., ) classes by shape, not texture.

Adversarial representation learning and concept erasure. In order to control the information in learned representations, prior works have trained feature extraction networks adversarially against discriminator networks that predict this information (Goodfellow et al., ; Schmidhuber, ; Ganin & Lempitsky, ; Ganin et al., ; Edwards & Storkey, ). Other works have removed concepts by modifying activations during inference (Ravfogel et al., ; Belrose et al., ; Elazar et al., ; Bolukbasi et al., ).

Robust unlearning. Machine unlearning seeks to remove undesired knowledge or abilities from a pre-trained neural network (Cao & Yang, ; Li et al., ). Typical unlearning methods are brittle in the sense that the unlearned abilities of the model can be recovered by fine-tuning on a tiny number of data points (Henderson et al., ; Sheshadri et al., ; Lynch et al., ; Liu et al., ; Shi et al., ; Patil et al., ; Lo et al., ; Lermen et al., ). Lee et al. (); Łucki et al. () suggest that undesired concepts are more easily “bypassed” than thoroughly removed from model weights. In this paper, we pre-train models with gradient routing such that we can perform *robust* unlearning, which cannot be easily undone by retraining. Tampering Attack Resistance (TAR) (Tamirisa et al., ) also targets robust unlearning in LLMs. While their method does improve robustness to retraining, it degrades general model performance as a side effect. However, we present gradient routing as a training technique rather than a post-hoc modification, so the two methods aren’t directly comparable.

Compared to gradient routing, the most similar approaches prune or mask parts of the network most important for the target behavior. SISA (Bourtoule et al., ) trains multiple models in parallel on partitions of the dataset and takes “votes” from each model at inference time. Similar to ablating a network subregion, a model can be dropped to achieve robust unlearning. The approach of Bayazit et al. () is to learn a mask over parameters in a language model to unlearn specific facts, while Huang et al. () and Pochinkov & Schoots () remove neurons related to harmful behavior in order to restore the alignment of an adversarially fine-tuned language model. Guo et al. () fine-tune the parameters of only the most important components for the task. Lizzo & Heck () instead delete subspaces of the model parameters in order to remove specific knowledge. Unfortunately, Lo et al. () find that models pruned to remove a concept can very quickly relearn the concept with further training. This may be because *identifying* the precise sub-network for a task post-hoc is very challenging, as evidenced by the modest success of “circuit discovery” in mechanistic interpretability thus far (Wang et al., ; Conmy et al., ; Miller et al., ; McGrath et al., ).

Scalable oversight. Scalable oversight is the problem of providing a supervised training signal for behaviors that are difficult or expensive to assess (Amodei et al., ). Semi-supervised reinforcement learning frames scalable oversight in terms of RL on partially labeled data (Zhu et al., ; Finn et al., ; van Engelen & Hoos, ). Another approach is weak-to-strong generalization, in which a less powerful model provides supervision to a more powerful one (Burns et al., ; Kenton et al., ; Radhakrishnan et al., ). Weak-to-strong generalization introduces a potential risk: the stronger model may exploit blind spots in the weaker model’s supervision capabilities.

## 3 Gradient routing controls what is learned where

Figure: Figure 1: Gradient routing applied to three different data points. Gradient routing uses masks (shown here) to zero-out or re-weight gradients as they are backpropagated through the computational graph (denoted by arrow thickness). By supplying different masks for different data points (denoted by colors), the user controls what regions of the network learn to predict different data points.
Refer to caption: x1.png

In order to train models that are interpretable and controllable in targeted ways, we seek to localize specific knowledge or capabilities to network subregions. To do this, gradient routing applies data-dependent, weighted masks to gradients during backpropagation to configure what data (whether it be defined in terms of tokens, documents, or based on other labels) is learned where in the network (e.g. at the level of parameters, activations, or modules). The result is a model with a partially-understandable internal structure, where particular regions correspond to known capabilities. Throughout this paper, we will use “route $X$ to $Y$” to mean “use gradient routing to limit learning updates for data points $X$ to region $Y$ of the neural network.”

Let $(\mathcal{V},\mathcal{E})$ be the nodes and edges of the computational graph corresponding to a neural network and loss function, with $v(z)$ taken to be the output of node $v$ if $z$ is input to the network. Given a dataset $\mathcal{D}=\{z_{i}\}_{i=1}^{n}$, for each data point $z_{i}$, gradient routing requires the specification of a gradient route given by $\widetilde{\mathcal{E}}_{i}=\{\alpha^{i}_{e}\in\mathbb{R}:e\in\mathcal{E}\}$ and visualized in fig. 1. Define $\frac{\partial L(z)}{\partial{v}}\triangleq\frac{\partial L(\zeta)}{\partial{v}(\zeta)}|_{\zeta=z}$, the partial derivative of the loss $L$ with respect to the output of node $v$ when evaluated at input $z$. The routed derivative (denoted with a tilde) of the loss over a batch $\mathcal{B}\subseteq[n]$ is then defined recursively as $\frac{\widetilde{\partial}L(z_{i})}{\widetilde{\partial}L}\triangleq 1\text{ for all }i\in\mathcal{B}$, and

$$ $\displaystyle\frac{\widetilde{\partial}L(z_{i})}{\widetilde{\partial}v}\triangleq\sum_{u\in\mathrm{child}(v)}{\alpha^{i}_{(v,u)}}\frac{\widetilde{\partial}L(z_{i})}{\widetilde{\partial}u}\frac{\partial u(z_{i})}{\partial v},$ $$

for all non-terminal nodes $v\in\mathcal{V}\setminus\{L\}$ and $i\in\mathcal{B}$. Choosing $\alpha_{e}^{i}\equiv 1$ recovers standard backpropagation. This weighting is only applied in the backward pass; the forward pass is left unchanged. Any gradient-based optimizer, like SGD or Adam (Kingma, ), can then be used to train with these modified gradients.

It is costly to define gradient routes over every data point and edge in the computational graph. In practice, we limit masks to a small set of edges, like the outputs of specific MLP neurons or the outputs of specific layers. Also, we typically assign gradient routes to data points based on membership in a coarse partition, like the forget set or retain set in an unlearning problem. Implementation is straightforward and efficient: fig. 2 gives sample Pytorch (Paszke et al., ) code in which masking is applied to the outputs of sequential layers.

Figure: Figure 2: Example of gradient routing implemented in PyTorch. For each batch of training data points x, a batch of gradient_masks corresponding to those data points is passed as well. The detach() method applies the stop-gradient operator, preventing gradients from being backpropagated through act but leaving its value unchanged.

## 4 Applications

### 4.1 Routing gradients to partition MNIST representations

As a first example of feature localization via gradient routing, we train a simple MLP autoencoder on the MNIST handwritten digit dataset (LeCun et al., ) and use label-dependent stop-gradients to control where features for different digits are encoded. The goal is to obtain an autoencoder that reconstructs all digits (0-9) via an encoding that is made up of non-overlapping subcomponents corresponding to distinct subsets of digits. We choose subsets $\{0,1,2,3,4\}$ and $\{5,6,7,8,9\}$. To hint at the potential difficulty of this task, we note the encodings learned by an autoencoder trained on one of these sets admit low-error reconstructions on the other set, despite never being trained on it (details in appendix B).

We use a simple architecture of three-layer MLP modules with ReLU activations: an Encoder, a Decoder, and a Certificate. The Encoder processes a 28$\times$28 image into a vector in $\mathbb{R}^{32}$, and the Decoder processes that vector into a 28$\times$28 reconstruction. The Certificate is a decoder trained only on the bottom half of the encoding, which takes values in $\mathbb{R}^{16}$. Certificate updates do not affect the encoding. If the Decoder can reconstruct a digit that the Certificate cannot, this “certifies” that robust feature localization occurred (into the top half encoding, and away from the bottom half).

We use gradient routing to train an encoding split so that the top half encodes digits 0-4 and the bottom half encodes digits 5-9. While training on all digits, we apply stop-gradients to the bottom half of the encoding for digits 0-4 and stop-gradients to the top half of the encoding for digits 5-9. To induce specialization in the two halves of the encoding, we add the L1 norm of the encoding as a penalty term to the loss. The setup is shown in fig. 3(a). The results, shown in fig. 3(b) and fig. 3(c), are stark: while using the entire encoding allows the Decoder to reproduce all digits with low loss, the Certificate is only able to reproduce 5-9 from the bottom half of the encoding, as desired. Furthermore, the Certificate’s learned predictions for digits 0-4 are approximately constant. This suggests that we have successfully eliminated most information relevant to digits 0-4 from the encoding. We elaborate on experiment details and provide an extensive ablation study in appendix B.

Figure: (a) An autoencoder trained to encode digits 0-4 in the top half encoding and digits 5-9 in the bottom half. The full encoding is processed by a single Decoder module trained with gradient routing; we illustrate this using weight tying. The Certificate is trained to reconstruct all digits using only the bottom half encoding.
Refer to caption: x4.png

### 4.2 Localizing targeted capabilities in language models

In this section, we show that gradient routing applied to a small set of tokens can be used to localize broader features or capabilities in Transformer (Vaswani, ) language models. This is first demonstrated in terms of model activations, then applied to MLP layers for the purpose of robust unlearning.

#### 4.2.1 Steering scalar: localizing concepts to residual stream dimensions

Elhage et al. () frames the inter-block activations of a Transformer, or the residual stream, as the central communication channel of a Transformer, with all layers “reading from” and “writing into” it. Usually, the standard basis of the residual stream is indecipherable, with dimensions not corresponding to interpretable concepts. We pre-train a 303M parameter Transformer on the FineWeb-Edu dataset (Penedo et al., ) while routing the gradients for all _California(^1^11We use a leading _ to represent a leading space before a token.) tokens to the 0^th entry of the residual stream on layers 6-18. On token positions predicting _California, we mask gradients (to zero) on every residual stream dimension except the 0^th in layers 6-18. This masking causes the learning updates for those token positions to be localized to the weights that write into the 0^th dimension of the residual stream. After training, we look at which tokens’ unembedding vectors have the highest cosine similarity with the one hot vector for the 0^th entry of the residual stream. We find that _California has the highest cosine similarity, followed by California, _Californ, _Oregon, _Colorado, _Texas, _Florida, _Arizona, _Sacramento, and _Los; see appendix D for the top 300. These tokens all have semantic similarity to California, but gradient routing was not applied to them. This shows that gradient routing localizesbroader semantic concepts, rather than the narrow set of explicitly-routed tokens.

Past work on activation steering (Turner et al., ; Rimsky et al., ) computed (non-axis aligned) *steering vectors* specified by $d_{\text{model}}$ different values. However, since we localized California-related concepts to the 0^th dimension of the residual stream, we can steer the model to generate text related to California by adding a single scalar value to the 0^th entry of the residual stream during the forward pass. Appendix D provides steered model completions.

#### 4.2.2 Gradient routing enables robust unlearning via ablation

Robust unlearning (Sheshadri et al., ) means training models which lack the internal mechanisms or “knowledge” required for certain tasks, as opposed to merely performing poorly on those tasks. To address this open problem, we show that gradient routing can be used to localize capabilities to a known region of the network, then delete that region, removing those capabilities.

Figure: Figure 4: Backpropagation in the Route step of Expand-Route-Ablate, showing the flow of gradients through a Transformer for tokens in the forget set. Gradients for retain tokens are unmodified. Additional dimensions, shown with dashed outlines, were added to target layers in the MLP and attention blocks, and will be removed after training in the Ablate step. All modules participate in the forward pass.
Refer to caption: x7.png

To enable comprehensive comparisons, our initial study on robust unlearning applies gradient routing to a small (28M parameter) Transformer. This model is trained on an LLM-generated dataset of simple children’s stories based on the TinyStories dataset (Eldan & Li, ; Janiak et al., ). We partition the data into: 1) a forget set made up of any story containing one of the keywords “forest(s)”, “tree(s)”, or “woodland(s)”, and 2) a retain set made up of all other stories. An example story is given in appendix C. The goal is to train a model that performs well on the retain set but poorly on the forget set, and whose forget set performance is not easily recoverable by fine-tuning. To do this, we route specific forget tokens to designated MLP neurons using three-step process termed Expand, Route, Ablate (ERA):

Increase the dimensionality of the model by adding randomly-initialized neurons to particular *target layers*.

Train the model by supervised learning on next-token prediction, but on select tokens in forget stories, reduce the learning rate in the original dimensions of the model at the target layers. Figure 4 illustrates the routing step.

Delete the additional neurons. Post-ablation, apply a very small number of steps of fine-tuning on retain data to correct for degradation caused by ablation.

Does gradient routing localize capabilities that can be robustly ablated? To answer this question, we train five types of models: an *ERA* model that uses gradient routing to localize forget set concepts, a *base* model trained conventionally on all data, a *pure* model trained only on retain data to serve as a gold standard, a *control* model trained equivalently to ERA except without gradient routing(^2^22The control model is expanded, ablated, and fine-tuned. It uses a small L1 penalty (small in the sense that it has no measurable effect on loss; see appendix C) on the MLP activations in the target layers.) and an *RMU* model (Li et al., ) fine-tuned from the base model to serve as an unlearning baseline. Using these models, we obtain the following results. Approximate $95\%$ confidence intervals for the mean ($N=60$ runs) are given in parentheses or highlighted in figures.

- •
Shallow unlearning measures the degradation in forget loss caused by our method by comparing the loss on the forget set for the ERA model vs. the base model.
Result: ERA achieves shallow unlearning, with forget loss of $1.91\,(\pm 0.05)$, vs. base model forget loss $1.47\,(\pm 0.02)$.
- •
Robust unlearning measures the robust removal of capabilities by comparing the forget loss after fine-tuning on forget data for the ERA model vs. the pure model. Figure 5(a) shows that the ERA model is almost as hard to retrain on forget data as the gold-standard pure model. In contrast, RMU’s performance is easily recovered by less than a batch of data. We comment on the choice of RMU as a baseline in appendix C.
- •
Alignment tax measures the cost of gradient routing in terms of retain set performance, by comparing the loss on the retain set for the ERA model vs. the base model. Result: $1.67\,(\pm 0.01)$ ERA, $1.59\,(\pm 0.01)$ base. The reduced performance of ERA is influenced by the prevalence of forget data, which constitutes 21% of the training data. In fig. 11, we show that the performance gap is negligible when the forget set constitutes as much as 5% of the training data.
- •
The differential effect of routing measures the impact of ablation on the ERA model vs. the control model. Figure 5(b) shows that ablation has a large effect on the control model, particularly in terms of increasing forget loss, as compared to a negligible effect on the control model.

Losses are always calculated on held-out validation data. ERA setup and experiment details are given in appendix C.

Figure: (a)
Refer to caption: x8.png

Robustness to missing labels. Finally, we show that ERA achieves robust unlearning even when only a random i.i.d. subset of forget samples are labeled, with unlabeled samples treated as retain data for training purposes. This is meant to model real-world scenarios where comprehensive labeling is infeasible. Compared against a conventionally trained model that does not train on labeled forget data, ERA achieves (on average, across runs) higher retrained validation forget loss for all labeling proportions measured besides 100%. For example, at 60% labeling, ERA achieves a retrained forget loss of 1.53 ($\pm$0.02) as opposed to 1.49 ($\pm$0.02) for the baseline. Full curves are shown in Figure 9 in Section C.1.

#### 4.2.3 Scaling robust unlearning to larger language models

Gradient routing can localize capabilities in larger models. Motivated by the dual-use nature of AI (Urbina et al., ), we would like to train useful models that lack certain harmful capabilities. Here, we seek to localize and remove bioweapon-related capabilities in a 0.7B parameter Transformer. To do this, we route 20 tokens related to virology(^3^33Specifically, we route on COVID, _COVID, RNA, _infections, DNA, _genome, _virus, _gene, _viruses, _mutations, _antibodies, _influenza, _bacteria, PCR, _cell, _herpes, _bacterial, _pathogens, _tumor, and _vaccine.) to the 0^th through 79^th MLP dimensions on layers 0 through 7 of the Transformer. Appendix E provides further details on the model and training.

Table 1 evaluates the model on a validation split of regular FineWeb-Edu data and on some of the WMDP-bio (Li et al., ) forget set. Ablating the target region of the network increases loss greatly on both datasets. We then fine-tune the model on a train split of FineWeb-Edu for 32 steps to restore some performance. Finally, we retrain for twenty steps on a separate split of two WMDP-bio forget set datapoints, as Sheshadri et al. () do, and report the lowest loss on the validation split of the WMDP-bio forget set.

The results are striking: even after retraining on virology data, loss increases much more on the WMDP-bio forget set (+0.182) than on FineWeb-Edu (+0.032), demonstrating successful localization and robust removal of virology capabilities. A natural concern would be that ablation merely decreased probabilities on the routed tokens, without decreasing overall virology capabilities. To test this, we measured cross-entropy loss on the forget set excluding the 20 tokens we routed on. Even after this exclusion, the loss increase is still much higher than the increase on FineWeb-Edu (+0.171 vs. +0.032). This shows that gradient routing generalizes beyond limited labels.

**Table 1: Performance of a language model trained with gradient routing on virology tokens. The final column evaluates the model after finetuning on FineWeb-Edu and then retraining on two examples from the WMDP-bio forget set, choosing the retraining step with the lowest loss. The increase in loss on (the validation split of) the WMDP-bio forget set is much higher than the increase in loss on FineWeb-Edu data, demonstrating successful localization and robust unlearning. Intriguingly, this increase persists even when excluding routed tokens from the loss calculation, showing a broader localizing effect.**
| Dataset | Loss | Ablated loss ($\Delta$) | Retrained loss ($\Delta$) |
| --- | --- | --- | --- |
| WMDP-bio forget set $\uparrow$ | 2.596 | 4.283 (+1.687) | 2.778 (+0.182) |
| WMDP-bio forget set (sans routed toks)$\uparrow$ | 2.567 | 4.205 (+1.638) | 2.738 (+0.171) |
| FineWeb-Edu $\downarrow$ | 2.925 | 4.864 (+1.939) | 2.957 (+0.032) |

### 4.3 Localizing behavioral modules enables scalable oversight in reinforcement learning

In complex settings, reliable data labels are often scarce, especially when labeling requires human input (Stiennon et al., ; Bai et al., ; Baker et al., ). *Scalable oversight* (Amodei et al., ) means effectively utilizing limited labels to obtain a performant policy. In this section, we show that gradient routing’s effective use of limited labels enables scalable oversight, outperforming an outcome-based baseline and a data filtering baseline.

We use gradient routing to train a policy to reach two types of squares in a gridworld, Diamond and Ghost. Access to the type of square reached, or oversight, is only sometimes available. The aim is to produce a policy that reliably navigates to Diamond and not Ghost, even when oversight is lacking. In real-world problems, label availability may be correlated with task type or difficulty, and a policy might not generalize correctly because of this fact (Langosco et al., ).
To model this fundamental difficulty, we allow the policy to condition on the presence or absence of oversight. During training, we route policy gradients through two separate modules in a mixture of experts layer, each responsible for reaching one of the aforementioned squares. During evaluation, by ablating the Ghost module, we can steer the policy toward reaching Diamond even when oversight is lacking.

Figure: (a) Proportion of episodes in which our agent reached Diamond after being steered to it. A diamond indicates the case of 4% oversight level, where Diamond is reached 70% of the time. (Ranges: 95% C.I.)
Refer to caption: x10.png

Environment.
The environment is a $5\times 5$ gridworld Markov Decision Process (Puterman, ) with two randomly placed terminal squares: Diamond and Ghost. At each step, the agent can move north, west, south, or east. Episodes end upon reaching a terminal square or after 256 steps. The environment state includes a boolean indicator variable for every grid square, indicating whether oversight is present. These indicators are randomly sampled at the beginning of each episode. If an episode ends with the agent reaching a terminal grid square with the oversight indicator set to False, then the learning algorithm does not have access to the square type labels (Diamond, Ghost) for that entire episode. The (true) rewards, which are not always observed, are +1 for reaching Diamond, −1 for reaching Ghost, and 0 otherwise, with a discount factor of 0.97.(^4^44For general reinforcement learning terminology, see Sutton & Barto ().) However, our algorithm only uses these labels for gradient routing, and not for RL updates.

Architecture and training. The first layer of the policy network is a modified Mixture of Experts (MoE) layer.
The MoE layer computes a weighted sum of two experts, $\operatorname{E}_{\textsc{Diamond}}(\cdot)$ and $\operatorname{E}_{\textsc{Ghost}}(\cdot)$.
The module weights, which are between 0 and 1, are computed by a multi-layer gating network.
The remaining layers in the network are either convolutional or linear. (See appendix F for details.)

The policy network is trained with Proximal Policy Optimization (PPO) (Schulman et al., ).
We perform gradient routing only when training on episodes with oversight.
In those cases, we route the gradients through one of the modules, $\operatorname{E}_{\textsc{Diamond}}(\cdot)$ or $\operatorname{E}_{\textsc{Ghost}}(\cdot)$, depending on which terminal square the agent reached.
We also add a loss term that incentivizes the gating network to activate only the module through which the gradients were routed.
In episodes without oversight, we perform a normal, non-routed PPO update with full backpropagation.

Results. Our method uses gradient routing based on limited labels to localize a module responsible for going to Ghost, then ablates the module. Below, we show that the method works even in the most challenging setting: instances of the environment where Ghost is the closest terminal square, where we have no oversight (and the agent knows this). Details are given in appendix F.

Our *gradient-routed MoE* agent reaches Diamond 70% of the time (fig. 6(a)) with as low as 4% oversight during training. An *outcome-based* baseline, trained on all episodes and given reward −1 when it reached Ghost under oversight, and +1 when it reached a terminal square in any other case, learns to exploit the limits of the labels, going to Ghost whenever it is closer than Diamond and not under oversight. A *data filtering* baseline was trained on a random subsample of all training episodes, but granted full oversight. Gradient-routed MoE under 4% oversight has comparable performance to the data filtering baseline at a 20% sampling rate, representing 5 times greater data efficiency (fig. 12 in appendix F). So, we achieve scalable oversight: gradient routing avoids the perils of naive outcome-based training while enjoying far greater data efficiency than data filtering.

## 5 Discussion

Absorption. Routing a subset of the data related to some knowledge or capability appears to localize that knowledge or capability more generally. This held for an i.i.d. subset of the data (TinyStories unlearning in section 4.2.2), and for semantically limited data (steering scalar in section 4.2.1, virology unlearning in section 4.2.3, scalable oversight in section 4.3). We hypothesize an absorption effect: routing limited data to a region creates internal features or units of computation in that region which are relevant to a broader task; these units then participate in the model’s predictions on related, non-routed data; the resulting prediction errors are then backpropagated to the same region, creating a positive feedback loop that reinforces those features. To the extent that absorption is true, it has advantages over data filtering methods: if data filtering labels are limited either in quantity or semantically, then harmful capabilities can still be learned where labels are missing, whereas routing that data to a region absorbs those capabilities into that region, which can then be removed.

Benefits of localization vs. suppression. When the ability to label (or score) undesirable behavior is imperfect, attempting to suppress the behavior may be perilous: a model may learn to exploit the limits of the labels, rather than learning the desired behavior (Goodhart, ; Karwowski et al., ). Our study of scalable oversight presents a model of this scenario, demonstrating the advantage of localization as opposed to attenuation of undesirable behavior. This advantage may apply more broadly, for example, to machine learning problems where capabilities are *entangled*, in the sense that there are connections or dependencies between the computation learned to perform different tasks (Arora & Goyal, ; de Chiusole & Stefanutti, ). Entanglement might occur because certain capabilities or behaviors are reinforced by a broad range of training objectives (Omohundro, ; Turner et al., ; Krakovna et al., ). More simply, capabilities required to perform undesired tasks may overlap with those required to perform desired tasks. For example, biological knowledge entails much of the knowledge required to construct biological weapons. For this reason, filtering or training against bioweapon-specific data might not prevent a network from learning enough to create bioweapons from general biology sources.(^5^55Another reason suppression may be insufficient to provide safety guarantees is that poor behavioral performance does not entail the elimination of internal circuitry related to that behavior (Lee et al., ; Casper et al., ).)

Limitations and future work. (a) Gradient routing’s performance is sensitive to its many hyperparameters: what data to route on, what regions to localize to, and what mask weights to use. This makes it hard to balance retain set performance vs. unlearning, for example. We suspect that methodological improvements will reduce this sensitivity. (b) So far, we have studied gradient routing as a pretraining method, making it costly to experiment with large models. (c) In our experiments with language models, we route gradients on a token-by-token basis, ignoring neighboring tokens. This naive strategy is surprisingly effective. However, it is plausible that contextual information will be critical in some problems, necessitating routing strategies that depend on entire sequences. Finding practical ways of choosing what data to route in order to localize broad capabilities is an intriguing open problem. (d) Our empirical results for scalable oversight pertain to a simplistic, narrow setting. Furthermore, our method for scalable oversight requires that the ablated policy produce coherent behavior. This does not hold in general, so scaling oversight via localization may require new ideas. (e) Other methods could be used to achieve similar aims as gradient routing, for example, DEMix Layers (Gururangan et al., ) or Interchange Intervention Training (Geiger et al., ). (f) We elaborate on application-specific limitations in appendix A.

## 6 Conclusion

Gradient routing localizes targeted capabilities in neural networks, creating models with known internal structure. Even when based on simple and limited data labeling schemes, this localization is suitable for robust unlearning of pre-specified capabilities and scalable oversight. Consequently, gradient routing may facilitate the safe deployment of AI systems, particularly in high-stakes scenarios where black-box methods are insufficiently robust.

#### Acknowledgments

We are indebted to the MATS program for facilitating and supporting our collaboration. Bryce Woodworth was especially helpful with planning, team dynamics, and feedback on the final versions of the paper. We are also indebted to Neel Nanda, who provided invaluable feedback, as well as funding (via Manifund) that enabled us to run experiments on larger models. We thank Alice Rigg for help with running our experiments.

We are grateful for the wealth of feedback we received from Kola Ayonrinde, Eric Easley, Addie Foote, Mateusz Piotrowski, Luke Marks, Vincent Cheng, Alessandro Stolfo, Terézia Dorčáková, Louis Jaburi, Adam Venis, Albert Wang, and Adam Gleave.

#### Reproducibility statement

We include detailed descriptions of experiment settings in the appendix. Code to reproduce our results can be found at [https://github.com/kxcloud/gradient-routing](https://github.com/kxcloud/gradient-routing).

## Appendix to Gradient Routing: Masking Gradients to Localize Computation in Neural Networks

## Appendix A Extended discussion of application-specific limitations and future work

MNIST autoencoders. The cleanly separated MNIST autoencoder representations depicted in fig. 3(c) depend on the problem setup (e.g. the choice to *not* use data augmentation, like rotations) and use of heavy L1 regularization on the encoding vector. L1 regularization is required because, by default, a regular MLP autoencoder trained on MNIST digits retains information necessary to decode other digits.

For a wide set of hyperparameters, we find that gradient routing achieves quantitative representation splitting: the Certicate’s reconstruction of digits 0-4 has higher average loss than its reconstructions of digits 5-9 for a wide range of settings, including different partitions of the digits. However, outside the specific hyperparameters chosen for the results in the main body of the paper, the qualitative results are poorer: the visual difference in reconstruction quality between the different digit subsets is less stark than in fig. 3(c). We take this to highlight the problem-dependent characteristics of feature localization. In the case of autoencoding handwritten digits, separation of features for encoding different digits is “unnatural,” so achieving it requires a specific setup and heavy regularization.

Language models. We speculate that gradient routing on particular tokens introduces an “internal tug of war” between the expanded and original dimensions of the model (these dimensions depicted in fig. 4), where parameter updates in the original dimensions consistently decrease the logits for routed tokens and parameter updates in the expanded dimensions increase logits for routed tokens. This effect can be understood as a consequence of the mismatch between the implicit estimands
(learning targets) for the original and expanded dimensions. We were concerned that this effect, rather than localization of capabilities, explained the post-ablation increase in forget loss. However, preliminary measurements suggest that this is not the case. For example, we find that the loss of ERA models is higher on average on non-routed forget tokens than a pure model, whereas it is lower on average on routed tokens. In general, the learning dynamics of gradient routing remain an open question.

If routing one token to a dimension of the residual stream creates an interpretable, axis-aligned feature as discussed in section 4.2.1, then routing many tokens to many neurons could produce a neural network with transparent internal representations. These representations might be made up of “individual neurons…[that] corresponded to cleanly interpretable features of the input,” as imagined in Elhage et al. (), or they could be organized in different ways. In principle, gradient routing provides a straightforward means of achieving this. However, we suspect that naive attempts to localize large numbers of concepts to unique regions will lead to high training loss.

Scalable oversight. Our reinforcement learning results demonstrate the promise of a localization-based strategy for scalable oversight, but further empirical and conceptual work is needed. The toy environment we use is simple, lacking the complexity and asymmetries of real-world problems. Additionally, our proposed solution relies on the fact that ablating an otherwise-active module of a policy network produces a policy with coherent behavior, which may not be true in practice (and isn’t true in general, in principle). We discuss these considerations in appendix G.

## Appendix B MNIST autoencoder details and ablations

Model architecture. The Encoder, Decoder, and Certificate are all three-layer MLPs. The layer sizes for the Encoder produce data with shapes (28 $\times$ 28, 2048, 512, 32) and for the decoder, data with shapes (32, 512, 2048, 28 $\times$ 28). All hidden layers use ReLU activations. The final layer of the Encoder is linear. The final layer of the decoders is affine.

Training. The model was trained for 200 epochs on the 60,000 image training part of the MNIST dataset (LeCun et al., ) with batch size 2048. Images were normalized to have mean and standard deviation 0.5. No data augmentation was used. Optimization was performed with Adam (Kingma, ) with learning rate 1e-3, $\beta=(0.9,0.999)$, and weight decay 5e-5.

The loss used was pixel-wise mean absolute error, with a penalty term for the L1 norm of the encoding and a penalty term for the sum of absolute correlations (across batch elements) between the top and bottom half of the encoding. For a batch of data indexed $i=1,\dots,n$ and encoding size $32$, denote data points by $x_{i}$, encodings as $\widehat{z}_{i}$, and Decoder outputs as $\widehat{x}_{i}$. Then for $\lambda=0.003$ and $\gamma=0.1$, the loss used to train the autoencoder is $\mathcal{L}=\mathcal{L}_{\text{reconstruction}}+\lambda\cdot\mathcal{L}_{\text{L1}}+\gamma\cdot\mathcal{L}_{\text{Correlation}}$, where

$$ $\displaystyle\mathcal{L}_{\text{reconstruction}}$ $\displaystyle=\frac{1}{28^{2}\cdot n}\sum_{i=1}^{n}\|x_{i}-\widehat{x}_{i}\|_{1},$ $\displaystyle\mathcal{L}_{\text{L1}}$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\|\widehat{z}_{i}\|_{1},\text{ and}$ $\displaystyle\mathcal{L}_{\text{Correlation}}$ $\displaystyle=\frac{1}{16^{2}}\sum_{k=1}^{16}\sum_{h=17}^{32}\frac{\sum_{i=1}^{n}|\widehat{z}_{i,k}-\overline{z}_{\star,k}||\widehat{z}_{i,h}-\overline{z}_{\star,h}|}{\sqrt{\sum_{i=1}^{n}(\widehat{z}_{j,k}-\overline{z}_{\star,k})^{2}}\sqrt{\sum_{i=1}^{n}(\widehat{z}_{j,h}-\overline{z}_{\star,h})^{2}}},$ $$

with $\overline{z}_{\star,k}=n^{-1}\sum_{i=1}^{n}\widehat{z}_{i,k}$. Note: this equation does not include gradient routing, which is an intervention applied to gradients when backpropagating $\mathcal{L}_{\textnormal{reconstruction}}$ through $\widehat{z}_{i}$.

Additional results and ablations. Additional findings are given below. Many of them reference table 2, which provides results from ablation experiments.

- •
For a given set of hyperparameters, the run-to-run variability induced by random neural net initialization and data shuffling is small. For our main results (setting 1 in table 2), the 5th and 95th quantiles (across runs) of the average (over digits) final validation loss are (0.31, 0.33) for digits 0-4 and (0.08, 0.09) for 5-9.
- •
We find that training a regular autoencoder on a subset of digits, without regularization or gradient routing, results in an encoding that admits reconstructions of the digits that were not trained on (setting 8 of table 2).
- •
Inclusion of the correlation penalty helps split representations but is not necessary (compare setting 1 and setting 3 of table 2). However, regularization is necessary to achieve splitting (compare settings 1 and 2 to settings 4 and 5 of table 2).
- •
We find that we can learn separate “split” encodings of MNIST digits simply by training autoencoders on subsets of digits with a high L1 penalty, rather than applying gradient routing (setting 7 of table 2). However, gradient routing is still able to produce split encodings even in a more challenging setting where only one of the subsets of digits is routed, while the other has its gradients flow through the whole encoding (setting 6 of table 2, shown in fig. 7 and fig. 8).
- •
(Not presented in this document) For most digit partitions that we tried (other than 0-4 and 5-9), we were able to reproduce results similar to those given in fig. 3 without modifying hyperparameters. Generally, the results were quantitatively comparable to, but less visually striking than, those shown in fig. 3(c). We were even able to split the encoding into 10 parts, one per digit.

**Table 2: The average (over 20 runs) reconstruction losses for the Certificate, a decoder trained only on the bottom half encoding, for different MNIST autoencoder training settings. Approximate 95% confidence intervals are given in parentheses. Default regularization settings are an L1 penalty on the encoding with weight 3e-3, and a penalty on the sum of absolute correlations between the top and bottom half entries with weight 0.1. Gradient routing (Setting 1) is presented in the main body of the paper and uses the default regularization. Settings marked with “separate Decoders” trained a Decoder on digits 0-4 and a different Decoder on digits 5-9 (equivalent to removing weight tying in fig. 3(a)). Setting 6 is the same as Setting 1, with two modifications: no stop gradients are used on the bottom half encoding, and the L1 penalty is increased to 2e-2 on the bottom half encoding. Setting 6 is depicted in fig. 7.**
| Setting | Loss: 0-4 | Loss: 5-9 |
| --- | --- | --- |
| 1. Gradient routing | 0.32 ($\pm$0.02) | 0.08 ($\pm$0.00) |
| 2. Gradient routing, separate Decoders | 0.33 ($\pm$0.02) | 0.07 ($\pm$0.00) |
| 3. Gradient routing, no correlation penalty | 0.28 ($\pm$0.02) | 0.11 ($\pm$0.01) |
| 4. Gradient routing, no regularization | 0.32 ($\pm$0.02) | 0.32 ($\pm$0.01) |
| 5. Gradient routing, no regularization, separate Decoders | 0.09 ($\pm$0.01) | 0.08 ($\pm$0.00) |
| 6. Gradient routing, bottom half encoding trained on 0-9 | 0.23 ($\pm$0.02) | 0.13 ($\pm$0.01) |
| 7. No gradient routing, L1 penalty 1e-3, trained on 5-9 only | 0.27 ($\pm$0.02) | 0.11 ($\pm$0.00) |
| 8. No gradient routing, no regularization, trained on 5-9 only | 0.08 ($\pm$0.01) | 0.08 ($\pm$0.00) |
| 9. No gradient routing, with regularization | 0.13 ($\pm$0.01) | 0.13 ($\pm$0.01) |
| 10. No gradient routing, no regularization | 0.08 ($\pm$0.01) | 0.09 ($\pm$0.00) |

Figure: Figure 7: A variant of the MNIST gradient routing experiment from section 4.1. In this version, gradients from all digits (rather than merely 5-9) are allowed to flow through the bottom half of the encoding. Since the goal is to isolate the representations for digits 0-4 to the top half encoding, the inclusion of digits 0-4 makes the problem more challenging. However, by increasing the strength of the L1 penalty applied to the bottom half encoding, we still achieve splitting, as visualized in fig. 8 and quantified in table 2.
Refer to caption: x12.png

Figure: Figure 8: Certificate reconstructions from the more challenging gradient routing experiment described in fig. 7.
Refer to caption: x13.png

## Appendix C TinyStories unlearning details

Model architecture. We use the TinyStories-28M model from Eldan & Li (), which is an 8-layer Transformer with hidden size 512, 16 attention heads, vocabulary size 50,257, and GELU activations, as found at [https://huggingface.co/roneneldan/TinyStories-28M/tree/main](https://huggingface.co/roneneldan/TinyStories-28M/tree/main).

Training. The model was trained for one epoch on 400,000 stories from the Delphi version of the TinyStories dataset (Janiak et al., ; Eldan & Li, ), with batch size 80, truncating sequences at 256 tokens. The number of times each model was trained (the sample size for confidence intervals) is 60. The Adam optimizer was used with learning rate 5e-4 decaying to 5e-5 over the course of training, $\beta=(0.9,0.999)$, and weight decay 0.1. The forget set was defined as any story containing one of the following strings, separated by spaces or punctuation: “tree”, “trees”, “forest”, “forests”, “woodland”, and “woodlands”.

Unlearning baseline. Our results in fig. 5(a) compare against Representation Misdirection for Unlearning (RMU) (Li et al., ), a post-hoc unlearning method. RMU works by corrupting a model’s internal representations on forget data and preserving its representations on retain data. As is typical of many unlearning methods, much of the degradation to forget set performance caused by RMU is reversible by fine-tuning on a very small number of forget set examples (Sheshadri et al., ). The choice to compare against RMU in particular was arbitrary.

Expand, Route, Ablate settings. The following settings are used for the training process described in section 4.2.2 and depicted in fig. 4.

- •
Target layers: $\{0,1,2,3,4\}$.
- •
Dimensions added: 64 MLP neurons in each of the target layers.
- •
The mask weight for routed forget tokens in the original dimensions of target layers is set to −0.75. All other weights are 1.
- •
Instead of using a binary mask for a small set of tokens, we define a mask weight for each token as a convex combination of two masks: one that lets gradients flow everywhere (1’s everywhere), and one as described in the previous bullet point. The weight in the convex combination is set by the token’s relative frequency in the forget vs. retain set, biased towards retain. So the token “_the”, which has high frequency in both sets, is assigned the mask of 1s. The token “_tree”, which only appeares in the forget set, is given the most “aggressive” mask as defined in the previous bullet. Sample values are shown in table 3.
- •
Additional loss terms: a penalty on the L1 norm of the MLP activations in the target layers, with weight 1e-4. Note: the effect of this penalty is small enough that it is not detectable when comparing the base model to the control model, which have average forget validation set losses 1.47 ($\pm$ 0.02) and 1.47 ($\pm$ 0.02) respectively (not a typo).
- •
Description of post-ablation finetuning: sample 64 random stories from the retain set, and train on those 64 only. Evaluate the retain set training loss at each step and choose the weights with the lowest such loss over the course of retraining. This is usually achieved in two or fewer steps.

### C.1 Additional tables and figures

**Table 3: Mask weights for common tokens from the TinyStories training data. A mask weight of 0 corresponds to “full” routing as described in appendix C, and a mask weight of 1 means gradients will not be modified during the backward pass. In between 0 and 1, these gradient routes are interpolated.**
| Token | Forget set freq.<br>per 10k tokens | Retain set freq.<br>per 10k tokens | Mask weight |
| --- | --- | --- | --- |
| _tree | 99.5 | 0.0 | 0.000 |
| _bird | 73.1 | 18.7 | 0.585 |
| _flew | 10.3 | 3.6 | 0.810 |
| _bear | 10.9 | 3.8 | 0.816 |
| _animals | 10.2 | 3.9 | 0.851 |
| _Bob | 13.2 | 5.9 | 0.901 |
| _walked | 9.7 | 4.5 | 0.909 |
| _find | 19.9 | 9.3 | 0.912 |
| _down | 18.1 | 8.8 | 0.919 |
| _its | 8.4 | 4.2 | 0.922 |
| my | 5.1 | 7.1 | 0.991 |
| _dad | 3.8 | 5.8 | 0.992 |
| _says | 4.3 | 6.7 | 0.993 |
| _box | 6.9 | 10.6 | 0.993 |
| _water | 5.2 | 8.3 | 0.993 |
| _mom | 23.4 | 38.2 | 0.993 |
| _car | 5.3 | 10.9 | 0.996 |
| _toys | 4.3 | 11.2 | 0.998 |
| _room | 1.8 | 8.2 | 1.000 |
| _fish | 1.5 | 6.7 | 1.000 |

Figure: Figure 9: The performance of gradient routing vs. data filtering when a random proportion of the forget data is unlabeled, based on $N=13$ runs per setting (of model and labeling proportion). Gradient routing routes labeled data to expanded dimensions of target layers, and trains on unlabeled data as if it were in the retain set. Data filtering means the model trains on unlabeled forget data but not on labeled forget data. The “Retrained validation forget loss” refers to the lowest forget validation loss achieved when training on a batch of 64 forget tokens from the train set.
Refer to caption: x14.png

Figure: Figure 10: A reproduction of fig. 9, with an additional set of runs included. We present this figure separately because it is not central to the work. The additional runs are equivalent to the ERA runs, except that they sort the forget data such that all labeled forget samples are trained on before all unlabeled forget samples. (Ignoring these labels, the marginal distribution of the data shuffle order is unchanged.) The idea is to measure whether routing can induce features that persist even when routing ceases and training is performed normally. Apparently, it can, although the effect is less strong than i.i.d. training, and run-to-run variability (as reflected in wide confidence intervals) precludes definitive conclusions.
Refer to caption: x15.png

Figure: Figure 11: The influence of forget data prevalence on ERA performance. Left: The retain loss of ERA when the training data constitutes different proportions of forget data. The retain loss diminishes quickly with forget set size, approximately reaching the average loss of the pure model (no forget data) at proportions as large as 0.05. This suggests that in real-world problems where the forget set is relatively small, the alignment tax will be low. Right: the lowest obtained forget loss when retraining on 64 forget data points, showing that the effectiveness of unlearning is robust to the quantity of forget data.
Refer to caption: x16.png

### C.2 Sample story

The following is a story from Janiak et al. () used as part of the forget set in our unlearning experiments.

### C.3 Text completions

We provide TinyStories text completions for a *base*, *pure*, and *ERA* model. These models were trained with the same settings as the models in section 4.2.2, except they trained on 1.6M stories instead of 0.4M stories. We use two input sequences: one that evokes the concept “tree” from the forget set, and one that is unrelated to the forget set.

- •
Forget prompt: *“Once upon a time, there was a big oak tree.”*
- •
Retain prompt: *“Once upon a time, there was a kind girl named Lily.”*

The following completions were generated with a softmax sampling of temperature 0.8 and truncated at 200 tokens. We present the first two completions sampled from each of the models, without cherry picking.

Base model (trained on all data) - forget prompt

Pure model (trained on retain data only) - forget prompt

ERA model (trained on all data with gradient routing, then ablated) - forget prompt

Base model (trained on all data) - retain prompt

Pure model (trained on retain data only) - retain prompt

ERA model (trained on all data with gradient routing, then ablated) - retain prompt

## Appendix D Steering scalar details

Model architecture. We use a modified nanoGPT (Karpathy, ) model with the GPT-2 tokenizer, 20 layers, 16 attention heads, RoPE positional embedding (Su et al., ), and RMSNorm (Zhang & Sennrich, ).

Training. We train on sequences of length 1024 with $589,824$ tokens per step for $10,000$ steps. We use the AdamW optimizer (Loshchilov & Hutter, ) with a learning rate warmup of $2,000$ steps to $1.8\times 10^{-3}$ with cosine decay to $1.8\times 10^{-4}$ after $10,000$ steps, $\beta_{1}=0.9$, $\beta_{2}=0.95$, $0.1$ weight decay, and gradient clipping at $1.0$.

The tokens most similar to the localized dimension. The unembed matrix of a Transformer $U\in\mathbb{R}^{d_{\text{vocab}}\times d_{\text{model}}}$ maps the output of the final hidden layer to logits for the token vocabulary. To find the tokens with the highest cosine similarity to the localized “California dimension” (the 0^th standard basis vector), we sort them according to $U_{i,0}/\|U_{i}\|_{2}$ and take the most negative values. This results in the following 300 tokens, in descending order of cosine similarity.

_California, California, _Californ, _Oregon, _Colorado, _Texas, _Florida, _Arizona, _Sacramento, _Los, _San, _Hawaii, _Nevada, _Utah, _Alaska, _Massachusetts, _Missouri, _CA, _Minnesota, _Illinois, _Hawai, _Southern, _Connecticut, _Kansas, _UC, _Louisiana, _Virginia, _Pacific, _American, _Santa, _Maryland, _Fresno, _Japan, _Mexico, _Maine, _Michigan, _Wisconsin, Calif, _America, _Ohio, _China, _Berkeley, _Washington, _Pennsylvania, _Nebraska, _Kentucky, _New, _Cal, _Americans, _Idaho, _Mexican, _Queensland, _Chicago, _Iowa, _Oakland, _Wyoming, _Oklahoma, _UCLA, _Calif, _Costa, _Hawaiian, _Ventura, Colorado, _US, _Yosemite, _Chile, _Mississippi, _Stanford, _Chinese, _Brazil, _Sierra, _Tokyo, _Indiana, _Alabama, _Arkansas, _Montana, _LA, _Philippines, _United, _Spain, _Ranch, Oregon, _Moj, _Vermont, _Denver, _Carolina, _Peru, _Western, _Alberta, _North, _Hollywood, _Rhode, _Ontario, _Tennessee, _Italy, Texas, _Canada, _Seattle, _Puerto, Florida, _Delaware, _CAL, _Japanese, _Southwest, _Georgia, Los, Arizona, _Marin, _states, _Kenya, _Houston, _statewide, _Pasadena, _Brazilian, _Hong, _Australia, _southern, _UCS, _London, _Italian, _Kerala, America, _European, _U, _Vancouver, _Taiwan, Utah, _Tucson, _Ecuador, _Northern, _Beijing, _Boston, _Honolulu, CA, _Canadian, ornia, Japan, _BC, _Australian, _Coast, _Davis, _South, Ber, _Saudi, _parsed, _Kern, _British, _Silicon, _Palo, _Chilean, _Spanish, _NYC, _Mexicans, _NSW, _Anaheim, _Philippine, _federal, _Texans, _almonds, _Kyoto, _Midwest, _timeout, _States, _Central, _Manhattan, _West, _Proposition, UC, _Miami, Washington, _desert, 688, _Pittsburgh, Mary, _Brooklyn, _Guam, _Colombia, _Bay, _northern, _Riverside, _Philadelphia, _India, _Portland, Virginia, _western, _Panama, _Mediterranean, _Federal, _Angeles, _Mont, _USA, _southwestern, _Cincinnati, orset, _AMERICA, _UK, _Schwarzenegger, _Al, 115, _Per, Santa, _coast, _Berlin, Cal, _Okinawa, Mexico, _Filipino, _cal, apan, _NY, Italy, _Harvard, _nationwide, _Asian, San, _NASA, _Shanghai, _WA, arkable, American, _Victoria, _Saskatchewan, ijuana, _federally, _Honduras, oma, _Argentina, 69, Americans, _Nicaragua, har, _Latino, _Montreal, _Korea, _villain, _Yemen, _climates, _Francisco, _Northwestern, _Northwest, _Cuba, _Europe, _Iceland, asms, _Madrid, Yet, _Las, _Gujarat, Kansas, _cities, _England, _Irvine, erey, China, _Golden, Israel, _Portugal, ohm, _Lincoln, _americ, _Congress, _Kau, _State, _Switzerland, _Honda, _grow, _Paris, _state, Jesus, _ranch, outhern, , _USC, _Indian, _Toronto, !’", _flavors, _Columbia, _Rio, , oming, _Son, _University, _Germany, argument, _Asia, _Bon, _L, _Cannabis, asting, cal, _Israeli, _Singapore, _UAE, 415, assion, Japanese, _college, _Latinos, _Victorian

Many of these tokens are related to California, even though California is the only token that we routed on. This provides evidence for the ability of gradient routing to localize concepts without full data labeling.

### D.1 Steered and unsteered samples from the model

We sample 150 tokens from the model 5 times using temperature 1.0 from the top-200 tokens with a prompt of 
. We then perform the same sampling but add $-30$ to the 0^th embedding dimension after layer 10 to induce steering.

Unsteered

Steered

We can see that the steered text talks about California and states, which is what seemed to get localized to the 0^th residual stream dimension.

## Appendix E Larger model unlearning details

Model architecture and routing settings. We use a modified nanoGPT (Karpathy, ) model with the Qwen-2 tokenizer, 20 layers, 2 key value heads with 8 query heads each, a 1536 dimensional embedding space, and RoPE positional embeddings. We route the specified tokens to the 0^th through 79^th MLP dimensions on layers 0-7. We add additionally set the mask weight for the routed forget tokens in the *original* dimensions of *target* layers to $\text{\textminus}5\times 10^{\text{\textminus}8}$. We also add a $1\times 10^{\text{\textminus}7}$ L1 penalty to the MLP activations of the target layers.

Training. We train on approximately 13B tokens from FineWeb-Edu and add in the approximately one half of the WMDP-bio (Li et al., ) forget set to ensure that the model has seen information about virology. Each step consists of an effective batch size of $1,280$ for a total of $1,310,720$ tokens per step and we train for $10,000$ steps. We use AdamW with a learning rate warmup of $2,000$ steps to $1.8\times 10^{-3}$ with cosine decay to $1.8\times 10^{-4}$ after $60,000$ steps, $\beta_{1}=0.9$, $\beta_{2}=0.95$, and gradient clipping at $1.0$.

Evaluation.
After training, we ablate the 0^th through 79^th MLP dimensions on layers 0 through 7. We then retrain on data from FineWeb-Edu for 32 steps of 128 sequences of 1024 tokens each, while not allowing gradients to flow into the dimensions that had been ablated. After that, we retrain on 2 samples from the WMDP-bio (Li et al., ) forget set for 20 steps and record the lowest loss on FineWeb-Edu and a validation split of the WMDP-bio forget set.

## Appendix F Scalable oversight details

In this section, we provide details on the motivation and setup for our experiments on scalable oversight. Recall that in scalable oversight problems, we seek to train a performant policy despite limited access to reliable labels. We deal with the episodic RL setting. Throughout, we distinguish between:

- •
*Cursory labels*: labels that are available for all episodes, which may lack key information about the episode; and
- •
*Comprehensive labels*: labels that fully characterize the relevant properties of an episode, sufficient to determine its true reward.

For example, in the context of process supervision (Uesato et al., ; Luo et al., ), cursory labels would refer to properties of the outcome of an agent-environment interaction (“did the agent answer the math problem correctly?”), and comprehensive labels would refer to properties of the process used to produce the outcome (“was the agent’s reasoning sound?”).

Partial oversight details. Each episode includes a label $y\in\mathcal{Y}$ that is either cursory (“did the agent reach a terminal grid square at all?”) or comprehensive (“which terminal grid square did the agent reach?”):

$$ $\displaystyle\mathcal{Y}=$ $\displaystyle\left\{\text{not reached},\;\text{reached something},\;\text{reached }\textsc{Diamond},\;\text{reached }\textsc{Ghost}\right\}.$ $$

Define $\mathcal{Y}^{*}=\{\text{reached }\textsc{Diamond},\;\text{reached }\textsc{Ghost}\}$, the set of comprehensive labels.

The problem setting is parameterized by a level of oversight $p\in[0,1]$. At the beginning of an episode, after the agent is randomly placed, Diamond and Ghost are placed uniformly at random on distinct grid squares no closer than 3 steps away. Then, an oversight indicator each for Diamond and Ghost are sampled independently with probability $p$ to determine which terminal squares will be under oversight. Finally, other grid squares are labeled with oversight i.i.d. at random with probability 0.25. This makes the oversight indicator a spurious feature on non-terminal squares; we chose this setup so that an agent “by default” would learn to ignore the labels, absent reinforcement for discriminating based on them. The environment state (which is observed by the agent) comprises a one-hot encoded state of the grid cells (not pixels) and a binary mask that contains the cells’ oversight indicators.

Comprehensive labels are available only for episodes where the agent reached a terminal square with the indicator set to True.
For the remaining episodes, the labels are cursory, i.e. either “not reached” or “reached something”.

Policy network architecture. Our policy network $\pi(s)$ incorporates a mixture of experts (MoE) layer. For a state $s\in\mathcal{S}$,

$$ $\displaystyle\pi(s)$ $\displaystyle=s\triangleright\text{MoE}\triangleright\text{Conv}\triangleright\text{Conv}\triangleright\text{Linear}\triangleright\text{Linear}\triangleright\text{Softmax},$ $$

with ReLU activations before any linear or convolutional layer.

where $\triangleright$ denotes a piping operator, $(x\triangleright f)\triangleq f(x)$. The MoE layer combines outputs from two simple modules $\operatorname{E}_{\textsc{Diamond}},\operatorname{E}_{\textsc{Ghost}}$, using a gating circuit $\Gamma:\mathcal{S}\to[0,1]^{2}$:

$$ $\displaystyle\operatorname{MoE}(s)$ $\displaystyle=\operatorname{E}_{\textsc{Diamond}}(s)\cdot\Gamma(s)_{\textsc{Diamond}}+\operatorname{E}_{\textsc{Ghost}}(s)\cdot\Gamma(s)_{\textsc{Ghost}};$ $\displaystyle\operatorname{E}_{\textsc{Diamond}}(s)$ $\displaystyle=s\triangleright\text{Conv}\triangleright\text{MaxPool};$ $\displaystyle\operatorname{E}_{\textsc{Ghost}}(s)$ $\displaystyle=s\triangleright\text{Conv}\triangleright\text{MaxPool}.$ $$

This architecture allows us to isolate computations responsible for certain behaviors into the modules, and later steer the model by manually manipulating the gating coefficients. We use the same architecture across baselines to enable direct comparisons. Preliminary results suggest that the method is applicable to other architectures, but more work is needed to establish the exact dependency between its performance and model architecture.

Training. We train the policy network using Proximal Policy Optimization (PPO) (Schulman et al., ) with a reward of +1 whenever it reaches any of the two terminal squares, i.e. the training reward for an episode with label $y$ is $r(y)=\mathds{1}\{y\not=\text{not reached}\}$.

Crucially, whenever we have access to a comprehensive label for an episode, we use it to perform gradient routing in the MoE layer. For example, for $y=\text{reached }\textsc{Diamond}$, we use

$$ $\operatorname{MoE}(s)=\operatorname{E}_{\textsc{Diamond}}(s)\cdot\Gamma(s)_{\textsc{Diamond}}+\mathrm{sg}{\left(\operatorname{E}_{\textsc{Ghost}}(s)\right)}\cdot\Gamma(s)_{\textsc{Ghost}},$ $$

where $\mathrm{sg}{(\cdot)}$ is the stop-gradient operator. We apply the inverse for $y=\text{reached }\textsc{Ghost}$, stopping the gradients for $\operatorname{E}_{\textsc{Diamond}}$. We do not use any stop-gradients when oversight is unavailable (i.e. when we only have cursory labels). We only route the gradients in the MoE layer. In all other, non-MoE layers, gradients are propagated without routing in all episodes.

We include an additional loss term to encourage the gating circuit to activate only on one module. We only use the additional loss term when we have access to the comprehensive label for an episode. For a trajectory $\tau=(s_{1},a_{1},\dots,s_{T},r(y))$, the overall loss is

$$ $\mathcal{L}=\mathcal{L}_{\text{PPO}}(\tau)-\alpha\cdot\mathds{1}\{y\in\mathcal{Y}^{*}\}\cdot T^{-1}\sum_{t=1}^{T}\log\Gamma_{y}(s_{t}),$ $$

where $\mathcal{L}_{\text{PPO}}(\tau)$ is the PPO loss, $\alpha>0$ is a tuning parameter, and we abuse notation by writing $\Gamma_{y}$ to mean $\Gamma_{\textsc{Diamond}}$ if $y=\text{reached {Diamond}}$, and similarly for Ghost.

Hyperparameters. We train the model using the CleanRL (Huang et al., ) implementation of PPO for 1.5M environment steps on environments with 4 parallel instances. Each policy rollout consisted of 256 steps. The learning rate was set to 2.5e-4 with annealing enabled. We used a discount factor ($\gamma$) of 0.97 and GAE-$\lambda$ of 0.95. The model was updated for 4 epochs per rollout, using 4 minibatches. Advantages were normalized, and we applied a clipping coefficient of 0.2 for both policy and value function losses. The entropy coefficient was set to 0.01, and the value function coefficient to 0.5. Gradient norms were clipped at 0.5, and we did not use a target KL divergence threshold.

The custom loss term coefficient is $\alpha=0.3$.

Evaluation. To steer the model towards Diamond, we assign gate coefficients $\Gamma(\cdot)_{\textsc{Diamond}}=1$ and $\Gamma(\cdot)_{\textsc{Ghost}}=0$.

Outcome-based baseline. This agent trains on all episodes, with the reward function

$$ $r(y)=\begin{cases}-1&\text{if }y=\text{reached }\textsc{Ghost}{};\\ \mathbin{\phantom{-}}1&\text{if }y\in\left\{\text{reached something},\text{reached }\textsc{Diamond}{}\right\};\\ \mathbin{\phantom{-}}0&\text{otherwise}.\end{cases}$ $$

The outcome-based baseline never reaches Diamond in our test environments (as we verified empirically by testing it at 50% oversight), because it learns the optimal behavior under this reward function, which is to go to Ghost whenever it is closer than Diamond and not under oversight.

Data-filtering baseline. This agent trains with access to comprehensive labels, but only a fraction of the episodes are observed. This fraction is analogous to the oversight level $p$. Its reward function is

$$ $r(y)=\begin{cases}-1&\text{if }y=\text{reached }\textsc{Ghost};\\ \mathbin{\phantom{-}}1&\text{if }y=\text{reached }\textsc{Diamond};\\ \mathbin{\phantom{-}}0&\text{otherwise}.\end{cases}$ $$

The data filtering baseline learns the optimal behavior given this reward function, which is to navigate to Diamond in all cases. However, it requires more data. To learn behavior comparable to a routed policy trained with 4% oversight, the filtering baseline requires about 20% of the training episodes to be kept (fig. 12). This means our method is about $5\times$ more data-efficient than filtering.

Figure: Figure 12: The data efficiency of gradient routing compared to the data-filtering baseline. Dashed blue lines denote the filtering policy. We obtain these curves by training a policy with full oversight, then rescaling the $x$-axis to reflect the curve that would occur had varying proportions of data been filtered (proportions given in parentheses). The solid red line denotes a routed policy. We averaged five full-oversight filtering policies to obtain the displayed curves. The average return is measured on a 100 randomly generated environments.
Refer to caption: x17.png

Note that the highest episodic return for the gradient routing MoE agent is about 0.2 lower than the data filtering baseline. This is because, while the ablated agent virtually never navigates *towards* Ghost alone, it will walk to Ghost if it is on the shortest path to Diamond. Based on the random initial state of the environment, this happens some proportion of the time, leading to reduced reward. We discuss considerations necessary for overcoming this shortcoming in appendix G.

## Appendix G Impacts of localizing capabilities vs. dispositions for scalable oversight

To achieve scalable oversight, our proposed strategy for preventing bad behavior (for example) is to (1) localize a submodule responsible for bad behavior, then (2) ablate the submodule. In this section, we one factor that may complicate this strategy in real-world applications.

We distinguish between two types of processing that might occur within a neural network to cause some behavior, like navigating to a red tile in a gridworld. With respect to a particular behavior, we define:

Processing that is necessary for engaging in the behavior; for example, feature extraction and computation to detect a red tile and compute the shortest path to reach it.

Processing that is not a capability but that determines behavior (as a probability distribution over network outputs). For example, a submodule that processes features representing the shortest path to a red tile and a blue tile and then returns action probabilities corresponding to the red tile.

These definitions are informal. Note: Similar terms have been used in the context of AI evaluations (Beverley et al., ), but, to the best of our knowledge, have not been formalized. See Beverley et al. () for a philosophical treatment of related terms.

Depending on whether capabilities or dispositions are to be localized, the application of gradient routing to scalable oversight faces different challenges, as summarized in table 4.

**Table 4: An overview of the challenges to localizing capabilities vs. dispositions as a means of achieving scalable oversight. A checkmark (✓) indicates a step that we speculate is easy to achieve; a challenge indicates a fundamental difficulty.**
|  | Localization during training | After ablating the target region |
| --- | --- | --- |
| Localizing capabilities | Challenge: entangled capabilities | ✓ |
| Localizing dispositions | ✓ | Challenge: distribution shift |

In the case of capabilities localization, obtaining a performant policy post-ablation is straightforward in principle: by localizing and ablating, one has created an encoding of the state which does not admit any postprocessing which will exhibit the capability (analogous to the MNIST split encoding, whose bottom half did not admit any learned decoding for digits 0-4 as shown in fig. 3). In that case, one can simply train freeze this feature encoder and train on top of it. However, there is a fundamental challenge: in many problems, capabilities may not factor because they are entangled. For example, the skills required to be a cybersecurity researcher vs. a hacker overlap significantly.

On the other hand, we speculate that localizing dispositions is straightforward, and suitable for problems where capabilities are entangled. For example, even if cybersecurity and hacking involve the same capabilities, we expect to be able to localize the disposition for (harmful) hacking. However, localizing dispositions for scalable oversight does not permit post-ablation training, because further training could change the agent’s disposition. Instead, we must either zero-shot ablate, or find another manner of post-training that avoids this issue (e.g. fine-tuning on high-quality labeled data only). The fundamental difficulty to zero-shot ablation is distribution shift: suppose that during the training of a policy, an internal module is learned that governs the policy outputs in some regions of state space but not others. If, upon ablation, that module “becomes responsible” for regions that were previously governed by an ablated component, there is no reason to expect it to perform well in these states which are, with respect to its role in training, off-distribution.

## Appendix H Computational cost of gradient routing

Memory. Storing edge weights for every data point would incur a hefty cost of $O(|\mathcal{B}||\mathcal{E}|)$ memory per batch. In practice, this cost is easily avoided by reducing dependence on the amount of data and the number of edges. First: instead of assigning unique gradient routes to each data point, we assign routes according to membership in parts of a partition $\bm{\mathcal{P}}$ of data points, reducing the $|\mathcal{B}|$ term to $|\bm{\mathcal{P}}|$. For example, in a typical unlearning application, we would use $\bm{\mathcal{P}}=\{\mathcal{P}_{\text{retain}},\mathcal{P}_{\text{forget}}\}$ with a single gradient route assigned to each set. Second: we restrict the set of edges considered. For example, using only edges leaving parameters reduces the $|\mathcal{E}|$ factor to $O(p)$ if the neural net parameters have dimensionality $p$. This amounts to choosing elementwise learning rates for each parameter entry, for each data point.

Runtime. In the general case, gradient routing requires $|\mathcal{B}||\mathcal{E}|$ floating point operations to apply a scalar multiplication to each edge in the computational graph. Since we apply gradient routing to a sparse set of edges, like the $d_{\text{model}}$ entries of a hidden activation of a Transformer, the number of operations is much lower: $|\mathcal{B}|\cdot d_{\text{model}}$, for example. This is negligible compared to the number of operations required for matrix multiplication.

## Appendix I Extended literature review

We start by reviewing further works that, like gradient routing, modify learning rates or backpropagation.

Adjusting learning rates. Discriminative fine-tuning (Howard & Ruder, ) sets the learning rate for each layer independently to improve training efficiency. You et al. () introduce Layer-wise Adaptive Rate Scaling (LARS), which dynamically adjusts learning rates for each layer during training.

Modifying backpropagation. Sun et al. ()’s meProp uses only the top-k dimensions by magnitude of the gradient when updating parameters during training, which improves the accuracy of MNIST classifiers. Panda et al. () and Sung et al. () optimize only a sparse subnetwork of a model during fine-tuning, minimizing catastrophic forgetting and memory usage. Rosenfeld & Tsotsos () go a step further by updating only a small subset of parameters during pre-training, demonstrating competitive performance compared to conventional methods.

The methods above can be framed as multiplying the gradient by a mask vector. Mohtashami et al. () prove the theoretical convergence properties of binary gradient masking methods using a similar notation to our definition of gradient routing in Section 3.

Geiger et al. () train models to respect certain causal structure by applying interventions to the forward pass and minimizing the difference between the actual output and the expected output according to a user-supplied causal model. This method could be used to localize capabilities by ensuring some modules are causally relevant to certain outputs.

Fine-tuning parameter subsets. Many popular fine-tuning methods update only a small subset of parameters with the goal of computational efficiency or minimizing catastrophic forgetting or catastrophic interference (Sun et al., ; Sung et al., ; Rosenfeld & Tsotsos, ; Kaplun et al., ; Lee et al., ; Zhang et al., ; Mallya & Lazebnik, ; Panda et al., ). In some sense this localizes the new capabilities to this small subset of the network (as gradient routing does), although these tuned parameters may be activating latent abilities already present in the network (Ben Zaken et al., ).

Safe LoRA (Hsu et al., ) projects fine-tuned weights into a “safety-aligned subspace’, while subspace-oriented model fusion (SOMF) (Yi et al., ) masks task vectors (Ilharco et al., ) such that they do not interfere with the subspace identified as relevant for safe behavior, before merging them into the model using model fusion (Zhang et al., ; Jin et al., ).

Hierarchical reinforcement learning. Early work in hierarchical reinforcement learning used hand designed sub-behaviors assigned to individual modules to divide and conquer more complex tasks (Maes & Brooks, ; Singh, ; Mahadevan & Connell, ) although later works generally discarded this approach in favor of automatically learned sub-behaviors (Hutsebaut-Buysse et al., ).

Disentangled representations. While gradient routing partitions representations using supervised training, disentangled representation learning attempts to separate representations in an unsupervised manner (Bengio et al., ; Wang et al., ) using methods such as VAEs (Kingma & Welling, ; Mathieu et al., ) and GANs (Goodfellow et al., ; Chen et al., ).

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 4
questions: Why does "absorption" happen? Is the absorption effect reliable across multiple settings?

Questions on the robust unlearning experiment
- In Fig 4(a), why does the validation loss increase with more training steps? Can we see what happens with a more realistic number of finetuning steps (say 1000)?
- Why does RMU appear to outperform ERA on the 4-stories task? Does it just have overall higher loss?
- In fig 4(b) it appears that post-finetuning retain set los is also increased. How much of the validation forget loss in ERA and RMU is just due to overall slightly worse models?

Is ERA expected to work when the pretraining model and dataset are big (such as for modern LLMs) and the representations in the original model have good understanding of the forget set?
soundness: 2
strengths: The proposed method of gradient routing is novel and interesting, and has applications for multiple important problems in safety and interpretability. 

Impressively, the proposed method enables forgetting undesirably capabilities even without labelling full forget sets.
summary: This paper presents a new method for training networks called Gradient Routing (GR) with the goal of isolating capabilities to specific subregions of the network by controlling which network subregions are updated by which data points.

The paper shows applications for gradient routing towards a variety of settings and problems:

- The authors first show that by applying GR to an MNIST-autoencoder, it is possible to isolate the representation of digits 0-4 to the first half of an embedding and digits 5-9 to the second half. Additionally, the authors show that GR can be used for activation steering in language models. 

- The authors then propose to use GR for robust unlearning in already learned language models, via an approach called expand-route-ablate (ERA), where pretrained networks are expanded to include new subregions, some already learned capabilities are re-routed to new subregions, and then new subregions are ablated from the model, unlearning these previously learned capabilities from the network. Experiments show that ERA is successful in unlearning topics for small language models, even if only a fraction of "to-unlearn" samples are labelled.

- The paper then shows that vanilla GR can be successfully applied to learn a 0.7B language model that lacks specific harmful capabilities such as bioweapon related capabilities.

- Finally, the paper demonstrates an application of GR in reinforcement learning, to learn a policy that avoids certain target squares in a grid using a limited oversight signal to route gradients.
weaknesses: - The vanilla GR method requires splitting data into "retain" and "forget" sets before training begins, and then routing gradients during a potentially long and expensive training process. This limits applicability (except of the ERA method) and raises the question if GR is ever necessary for language model unlearning applications. If we knew the split during pretraining, why would we not simply omit the forget set during pretraining, learning the "pure" model considered in the paper?

- The ERA method requires more experiments and analysis/ablations to show why it works -- it is not clear why it is possible to introduce new network units, set low learning rate on old units for forget set samples, ablate the new units, and observe poor performance on forget set samples when running the old units. Is this a one-off result? Does this generalize to more tests and bigger models that may have more significant understanding of forget-set content?

- GR requires isolating each capability to a subset of examples before training and establishing network subregions for each capability, which may be challenging when there are lots of examples or lots of capabilities.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Q1: How the gradient masks are selected / initialized for different tasks / data samples / features? Are there meaningful heuristics that generally work well?

Q2: How to choose the network subregions that gradients should be routed through? Would certain architectures or tasks require a different approach to subregion selection? Have you experienced any achitecture-specific challenges?
soundness: 3
strengths: - The paper makes a new contribution to the field: GR is a new method to control gradient propagation based on data, enabling modular learning within a network.
- GR can make neural networks more interpretable by designing which part of the network is responsible for a specific task;
- GR provides support for targeted unlearning and scalable oversight;
- Broad applicability of GR (and many interesting questions to spark future work);
- GR can reduce interference and improve accuracy on specialized tasks;
- The method helps to control data access in privacy-sensitive scenarios.
summary: The paper introduces gradient routing (GR), which allows localizing computation in neural networks by masking gradients based on data-dependent paths. GR controls gradient flow within specific network subregions, isolating capacities and enabling the model to focus on particular data features / tasks. It promotes modular network design, making it suitable for selective forgetting and scalable oversight in reinforcement learning and language models. GR is evaluated on MNIST, gridworlds, and language tasks to demonstrate that it can improve task-specific performance, support targeted unlearning, and improve model interpretability by maintaining distinct functionalities in different parts of the network.
weaknesses: - GR requires careful selection of mask weights, data subsets, and regions to localize. How should these be chosen? At least in its current form, the method faces difficulties scaling up to larger models.
- No comparison to similar methods mentioned that can also achieve localization (DEMix and Interchange Intervention Training, both mentioned in the text).
- Unclear relevance of GR to safety-critical applications mentioned in the paper. Demonstrating benefits in a high-stakes scenario would strengthen the paper.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Clarification of Contributions

 2. Choosing the Routing Strategy:
Since gradient routing relies on predefined, data-dependent masks, an essential question is how to select an effective routing strategy. Could the method involve dynamic routing decisions based on the characteristics of each input, allowing gradients to flow through particular layers or weights? Further clarification on decision criteria for different types of data would strengthen the approach's versatility and practicality.

3. Impact of Initialization and Routing Interactions:
Gradient routing may be influenced by the initialization of weights, particularly concerning the Lottery Ticket Hypothesis [1]. If certain weights are initially unproductive and specific inputs route gradients exclusively to these weights, it could hinder learning.


[1]. The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks
soundness: 3
strengths: 1. This paper is well written and easy to follow.
2. The idea is simple and effective
3. The topic is really important for helping understand neural networks and learned representations.
summary: This work introduces gradient routing, a training method that enhances neural network safety by isolating specific capabilities within distinct network subregions. By applying data-dependent, weighted masks to gradients during backpropagation, gradient routing allows users to control which parameters are updated for particular data points. This approach achieves three main outcomes: (1) interpretable partitioning of representations, (2) effective unlearning by ablating targeted subregions, and (3) improved oversight of reinforcement learning by confining behaviors to distinct modules. The results indicate that gradient routing can localize capabilities even with limited data, suggesting its potential for complex, real-world applications.
weaknesses: 1. My main concern is about the contributions. Similar ideas has already been proposed by previous papers
(1) Piggyback [1] and PackNet[2] tries to learn different tasks better with different subsets of the weights by leveraging a binary mask.
(2) Parameter efficient tuning method for Large language models(LLMs) such as lora[3], can efficiently learn different downstream tasks using different adaptor modules.

So the contributions of this paper are not enough.

2. Need more large scale experiments. Given the idea is not that new, only small scale experiments may not enough.



[1] Piggyback: Adapting a Single Network to Multiple Tasks by Learning to Mask Weights
[2] PackNet: Adding Multiple Tasks to a Single Network by Iterative Pruning
[3] LoRA: Low-Rank Adaptation of Large Language Models
[4] S-LoRA: Serving Thousands of Concurrent LoRA Adapters

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. For representations learning, the gradient routing method can be applied to more common image classification tasks?
2. Can you provide a more detailed description in Section 3? 
3. Could you add a Certificate using the top half encoding for comparison in Section 4.1?
soundness: 2
strengths: 1. Based on the gradient, a training method is proposed to isolate the ability to a specific sub-region of the neural network.
2. Gradient routing has implications for safely deploying AI systems, especially in high-risk scenarios where black-box methods are not robust enough.
summary: This paper proposes a gradient-routing training method for isolating capabilities to specific subregions of neural networks. We show that gradient routing can be used to (1) learn representations that are partitioned in an interpretable way; (2) achieve robust unlearning by ablating pre-specified subregions of the network; and (3) achieve scalable supervision of reinforcement learners by localizing modules responsible for different behaviors.
weaknesses: 1. In Section 3, authors did not clarify the gradient routing, thus make this paper hard to understand. For example, how gradient mask is generated.

2. Authors did not clarify how gradient mask is generated. I hope you can provide the formula or pseudocode for the mask generation at line 178 of the pseudocode.

3. The MNIST dataset has too little data to prove the generalizability of the method. Experiments on more widely used datasets such as ImageNet and COCO can be more convincing.

Meta-review: (a) summary

This paper addresses some limitations with the current neural network training algorithms relating to safety. It proposes to apply data-dependent masks provided by users for routing backpropagated gradients to the specified subregions in a neural network. Constraining weight updating to isolated subregions help learn partitioned representation, and is suitable for selective forgetting and scalable oversight reinforcement learning.

(b) strengths
+ The gradient routing is useful for deploying AI systems in high-risk scenarios.
+ It is applicable to broad cases including unlearn tasks, protecting privacy.
+ It improves accuracy with reduced interference.

(c) weaknesses
- The proposed method requires careful selection of hyper-parameters which limits its application.
- Gradient routing does not outperforms the baseline of data filtering in the experimental settings.
- There is an issue with scalability when there are lots of examples or lots of capabilities.
- It has limited novelty: similar ideas have already been proposed by previous papers (1) Piggyback [1] and PackNet[2] tries to learn different tasks better with different subsets of the weights by leveraging a binary mask. (2) Parameter efficient tuning method for Large language models(LLMs) such as lora[3], can efficiently learn different downstream tasks using different adaptor modules.

(d) decision

The proposed gradient routing method is interesting and has potential application to broad use cases. However, the issues with sensitivity to hyper-parameters selection and scalability affects its usefulness. The limited novelty also makes it not  ready for publication at ICLR. Please keep the reviewer comments in mind when preparing a future version of the manuscript.

Decision: Reject