Gradient Routing: Masking Gradients to Localize Computation in Neural Networks

Authors: Alex Cloud, Jacob Goldman-Wetzler, Evžen Wybitul, Joseph Miller, Alexander Matt Turner
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: representation learning, modularity, unlearning, reinforcement learning, scalable oversight

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.

Decision: reject
Reviewer scores: 5, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1730357567054

thread_type:
official_review

<thread_object>
creator: reviewer_gjlh

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730448945794

thread_type:
official_review

<thread_object>
creator: reviewer_dmtv

<content>
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.

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.

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

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
</content>
</thread_object>

</thread>
<thread>
creation date:
1730473752013

thread_type:
official_review

<thread_object>
creator: reviewer_hlx2

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680474597

thread_type:
official_review

<thread_object>
creator: reviewer_nugs

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731985472838

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Thank you to all reviewers, plus clarification of contributions

comment:
We thank the reviewers for their thoughtful comments and astute questions, which have already helped us make significant improvements to the paper. In addition to revisions to the paper, we include a general response here, as well as detailed responses to each reviewer below.

## Clarifying our scholarly contributions
We respectfully submit that the reviewers have  underestimated the strength of the contributions of the paper, which are as follows:
1. **We establish a novel way to control neural network internals that has unique, desirable properties**: Gradient routing allows the user to flexibly assign learning updates to network regions at any level of granularity and in any training regime. To the best of our knowledge, no other method exists that is relevant to all of the applications areas we present. We have added a new section, appendix J: Extended Comparisons to Other Modularity Methods, to clarify these claims and provide detailed comparisons to existing methods.
3. **Using gradient routing, we demonstrate progress on three major open problems in machine learning**:
    * supervised control of feature learning, an open problem not solved by the subfield of adversarial representation learning (mentioned in section 2 of the paper), nor by disentangled representation learning (an unsupervised approach mentioned in appendix I);
    * robust removal of harmful capabilities, an open problem that has received substantial attention due to a recent flurry of works establishing the "shallowness" of typical unlearning methods, as we discuss in section 2; and
    * scalable oversight, a general open challenge to supervised training of advanced AI systems (discussed in sections 2, 4.3, 5, and at length in the appendix).
4. Our work additionally **contains insights and raises intriguing questions about neural net learning dynamics and AI control**. 
    * We report results from a variety of experiments suggesting the existence of an *absorption* effect of feature localization (as discussed in section 5). 
    * To the best of our knowledge, we are also the first to consider the problem of unlearning that is robust to *imperfect labeling*. 
    * We identify a general strategy for solving the problem of *Goodharting* [1][2], by utilizing imperfect labels to influence behavior without including those labels in the behavioral training objective (as demonstrated in section 4.3, discussed in section 5, and elaborated on in the appendix). 
    * We introduce a simple-yet-challenging toy model for scalable oversight by noting a key challenge to achieving weak-to-strong generalization [3]: that a policy can condition on problem difficulty.

Finally, we urge the reviewers to consider the breadth of our investigations: we have applied gradient routing to feedfoward MLP autoencoders, deep CNN classifiers (newly added per reviewer suggestion), transformer language models of varying sizes, and reinforcement learning agents. We show that gradient routing is applicable across these diverse domains.

We believe that the above represents a significant contribution to the literature. If we are mistaken, we would be grateful for any corrections or feedback reviewers can offer.

## Putting gradient routing in context
Gradient routing is not intended to be competitive with conventional supervised learning methods (which target low loss, without concern for internal mechanisms), nor with conventional post-hoc unlearning methods (which target high forget-set loss, **not** robust removal of capabilities). The purpose of gradient routing is to enable new approaches to tackling foundational challenges presented by the development of powerful AI systems [4]. We reference some of these challenges in the introduction and discussion sections. To clarify the contribution, we have added further discussion in appendix L.

[1] C. A. E. Goodhart, 1984, _Problems of Monetary Management: The UK Experience_, Macmillan Education UK, London, ISBN 978-1-349-17295-5.

[2] Karwowski et al., 2024, _Goodhart's Law in Reinforcement Learning_, The Twelfth International Conference on Learning Representations, [openreview.net/forum?id=5o9G4XF1LI](https://openreview.net/forum?id=5o9G4XF1LI).

[3] Burns et al., 2024. Weak-to-Strong Generalization: Eliciting Strong Capabilities With Weak Supervision. _Proceedings of the 41st International Conference on Machine Learning_, 235:4971-5012.

[4] Anwar et al., 2024, _Foundational Challenges in Assuring Alignment and Safety of Large Language Models_, Transactions on Machine Learning Research, 2835-8856, 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732789283054

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A significant final revision

comment:
We have uploaded a significant revision in response to reviewer feedback. We are indebted to the reviewers whose engagement with the work and thoughtful suggestions enabled these changes.
 
## Changelog

* **Introduction** 
  * Improved the writing to more clearly explain our conceptual and practical contributions.
* **Background and related work**
  *  Added a review of the limitations of data filtering for control of neural network behavior, improving the motivation for gradient routing.
* **TinyStories unlearning (section 4.2.2)**
   * Added _DEMiX plus ablation_ so that ERA can be compared to another localization-based method. Contrasting the two approaches provides strong evidence of _absorption_ (see section 5). (We thank Reviewer hLX2 for the suggestion.)
   * Changed our metrics to highlight the most important comparisons: forget set performance before/after unlearning and retain set performance. This controls for the possibility that a method is simply worse at fitting the data. (We thank Reviewer nUgs for the suggestion.)
   * We report strong results clearly establishing ERA's superior performance in the partial labeling regime.
   * We add a result on ERA plus RMU that outperforms either method alone at robust unlearning. (This combination is natural because RMU is a "light touch" post-hoc method that does negligible damage to retain set performance.)
* **Scalable oversight in reinforcement learning (section 4.3)** We noticed that reviewers have not focused much on this part of the paper, which we feel is the most important part, conceptually.
   * To highlight the contribution, we re-wrote the section, adding clearer motivation and presentation. 
   * We report better results based on improvements to our implementation: gradient routing induces specialized submodules even at 1% data labeling and achieves massive gains in data efficiency as compared with baselines. 
   * We re-wrote appendix F and added learning curves and ablations to provide better context to the results in the main body of the text.
* **Discussion (section 5)**
   * Expanded discussion of absorption by contrasting ERA to DEMix and improved our explanation of the effect. 
   * Added subsections "Mechanistic supervision avoids Goodharting" and "Entangled capabilities motivate gradient routing," further clarifying the strong conceptual motivation for gradient routing.
   * Added further motivation for gradient routing by proposing a concrete way that localization could enable traditional computer security approaches [1] to high-stakes AI deployment. (Discussed further in our newly-improved appendix L, on gradient routing's applicability to problems in AI safety.)


[1] R. S. Sandhu and P. Samarati, "Access control: principle and practice," in IEEE Communications Magazine, vol. 32, no. 9, pp. 40-48, Sept. 1994, doi: 10.1109/35.312842.
</content>
</thread_object>

</thread>


Metareview:
metareview:
(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.