Gradient Routing: Masking Gradients to Localize Computation in Neural Networks

We thank the reviewer for acknowledging the implications of gradient routing for the safe deployment of powerful AI problems. We also thank the reviewer for identifying some major opportunities for improvement for the paper, which we have now incorporated. We hope the reviewer finds that we have clarified how masks are generated and established the general applicability of gradient routing.

Response to weaknesses:

1./2. We have added a paragraph to the end of section 3 explaining the general strategy for choosing gradient masks throughout the paper and referencing a newly-added appendix K, which gives precise specifications for masks used for all experiments. We also remind the reviewer that each section contains descriptions of how gradients are routed; for example Figure 2.a. and Figure 3. However, these explanations were not posed in terms of masks, so we are grateful that the reviewer gave us the opportunity to make this improvement.

3. The reviewer is right that MNIST MLP autoencoding is a limited setting. We have strengthened this part of the paper by adding analogous experiments on a CIFAR100 ResNet classifier to the appendix, section B.1. These results match the experimental design of MNIST and show comparably strong results. We also encourage the reviewer to consider our results on transformer LMs and reinforcement learning, showing the diversity of gradient routing applications.

Response to questions:

Questions 1. and 2. are answered above. We again thank the reviewer for helping us make these improvements.

3. Could you add a Certificate using the top half encoding for comparison in Section 4.1?

This is a great idea. We have added a Certificate to the top half encoding and changed the histogram of results to show how each certificate is able to reproduce differents sets of images from different halves of the encoding. We also include reconstructions from the top half encoding in figure 6 in appendix B.

Summary:

We are grateful for the reviewer's feedback and hope that we have addressed all the reviewer's concerns here and in the top-level comment clarifying the contributions of our work. We ask the reviewer to consider raising their scores to reflect these clarifications and improvements, or to let us know if they have futher concerns about or suggestions for the paper.

Response to questions:

1. We have clarified contributions in the top-level comment and in Response to Dmtv (1/2).
2. Regarding the choice of routing strategy: this is a great point. We have added a paragraph to the end of section 3, and an extended discussion in appendix K regarding how to choose gradient routes. Regarding whether the method could involve dynamic routing decisions depending on the input, the answer is that it can, and does! The MNIST autoencoder could be seen as one example. We also applying routing with continuous mask weights based on relative token frequencies in our Tinystories experiments (section 4.2; appendix C has details). However, the gradient routes we've used so far have been relatively simple. In the discussion (section 5), we state:

   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.

3. We appreciate the reviewer's insightful idea about the Lottery Ticket Hypothesis and potential cost of applying gradient routing. We agree. We have had similar concerns backed by intuitions from Neural Tangent Kernel [1], particularly regarding the fact that in wide networks, individual parameters don't move much during training (which seems contrary to the idea of localizing to small subspaces). We do think this means that certain kinds of gradient routing schemes are "unnatural" with respect to neural net learning dynamics and will fail. So far, we have been pleasantly surprised by the success of different routing schemes, though.

Summary: We understand that the reviewer finds our work sound, clearly presented, and very important, but has concerns about the strength of contribution. We have explained why we believe our contribution is significant: (a) it is a novel idea and method for localization in neural networks, and (b) it presents many new ideas for how to apply localization methods to open problems in ML. If the reviewer has further concerns we would be grateful for the opportunity to address them. Otherwise, we hope the reviewer will increase their scores to reflect the potential importance of our work.

[1] Jacot, A., Gabriel, F. and Hongler, C., 2018. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31.

We thank the reviewer for raising important questions about novelty so that we can clarify them for future readers, and for the insightful connection to the Lottery Ticket Hypothesis. We are also delighted to hear that the reviewer thought the paper was "well written" and that the method is "simple and effective" and "really important" for understanding NNs and representation learning. This is great to hear!

We understand the reviewer has concerns about the novelty of the idea. We have addressed these concerns in the top-level comment (above), explaining that gradient routing has unique properties which enable the novel applications given in the paper. We elaborate below.

Response to weaknesses:

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.

We thank the reviewer for bringing Piggyback and PackNet to our attention. The training steps of these algorithms can be viewed as special cases of gradient routing, which is a neat connection! (Similarly, we note that methods which train a subset of weights, like transfer learning, or low rank-adapters, can be viewed as special cases of gradient routing too.) As mentioned in the top-level comment, we have added detailed comparisons of gradient routing to Piggyback and PackNet in appendix K. (If the point is not clear, we are happy to add more comparisons.)

However, we emphatically disagree with the suggestion that our contributions (see top-level comment) are redundant with ideas from fine-tuning and continual learning. Although the ideas are superficially similar ("apply gradient updates only to specific parameters"), they are quite different. The methods mentioned by the reviewer are concerned with efficient training to achieve low loss in the fine-tuning regime. In contrast, our contribution is about giving the ML practitioner flexible control over model internals; we offer novel ideas for how this can be achieved and what problems it can be applied to. Furthermore, our method is more general, allowing for kinds of localization not achievable by prior methods. For example, we demonstrate a variety of ways of localizing parameter updates for individual tokens in language models (section 4.2).

We anticipate a possible objection: "yes, the aforementioned methods were designed for fine-tuning and continual learning, but you could have studied them as localization methods." Our answer is twofold. (i) In a sense, we did study them, because our formalization of gradient routing generalizes these methods. (ii) To the best of our knowledge, we are the first to propose localization as a solution to the problems in the paper. We see this as a very important aspect of our contribution, so to some extent it is irrelevant which method for localization we chose to apply (as long as the method works). We have identified novel ways to bring localization to bear on open problems in machine learning, and done so via a simple and general framework.

So the contributions of this paper are not enough.

We have attempted to provide a thorough and compelling response to this statement. We welcome further questions if anything is unclear!

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

We believe our idea is quite new, for the reasons explained above. We also hope the reviewer will consider the breadth and design of our experiments (including newly-added experiments on a ResNet classifier), which we believe convey a lot of useful evidence about gradient routing. We also include all of our code with the paper so that others can replicate and build on the results, further strengthening the contribution.

Thank you for your detailed responses. I appreciate the effort in addressing the questions, but I still have a few concerns:

I am particularly interested in the initialization process. In the MNIST generation experiments, the gradients flow through two distinct parts of the network. How are these parts initialized? Are they initialized together as in a default network, or are they initialized independently? In other words, are they initialized dependently or independently?

If Gradient Routing (GR) is used, how should the network be partitioned? Should it be divided across layers, or can different parts of a single layer also be partitioned? I would appreciate further clarification on how GR handles these scenarios.

Answer these questions could help me better understand the method and help readers know how to utilize the proposed method.

Response to 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?

We provide detailed experiment descriptions in section 4 and the appendices corresponding to each subsection (MNIST representation control: appendix B, TinyStories unlearning: appendix C, steering scalar: appendix D, virology unlearning: appendix E, and scalable oversight: appendix F). However, we thank the reviewer for raising the important practical question of how we chose these settings in the first place.

Choosing a mask usually starts with a concrete idea of what capabilities are to be localized, and a motivation for doing so. We find that this starting point constraints the search space and is helpful as a guide. From that point, choosing masks is like choosing a neural net architecture: we find that a combination of intuition about neural net learning dynamics, plus trial-and-error with small-scale experiments, is effective. The newly-added appendix K gives more details on this process and may of interest.

Regarding data: For unlearning, we were surprised to find that localizing just a few tokens that were prevalent in the forget set was sufficient to damage performance significantly (and this damage was not explained by bad performance on merely those tokens; see the second row of table 1 in section 4.2.3). On the other hand, attempting to localize updates for all tokens in the forget set led to terrible performance (high retain set training loss). This led us to the belief that for unlearning, it is important to pick a minimal subset of the task that is necessary for good performance, and localize only that.

Regarding architectures and regions of localization: In terms of architecture-specific challenges: we have yet to encounter an architecture that doesn't admit gradient routing. The MNIST MLP experiments worked right away and were easy to set up. The transformer experiments required more trial-and-error. In transformers, we found that localizing to MLPs was most effective (rather than attention heads or dimensions of the hidden activations). This seems consistent with evidence that MLPs serve as stores of factual information in transformer LMs [1].

Summary: The reviewer's apt comments have given us an opportunity to improve the paper and clarify our findings. The paper now includes an explanation of mask selection (main body + appendix), detailed comparisons to other localization methods (appendix), and an extended discussion of gradient routing's relevance to problems in AI safety (appendix). We eagerly await other suggestions for improvements. Also, we respectfully ask the reviewer to consider raising their scores if we have addressed their concerns.

[1] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 36, 2022. arXiv:2202.05262.

We are very grateful to the reviewer for their thoughtful and positive review, as well as their constructive comments which have helped to improve the paper. We thank the reviewer for enumerating the strengths of the contribution, including the novelty, generality, usefulness, and interesting questions raised.

The reviewer raises important points, which we address below.

Response to weaknesses:

GR requires careful selection of mask weights, data subsets, and regions to localize. How should these be chosen?

We have updated the paper to add a paragraph at the end of section 3, and a detailed appendix K, offering general guidance for choosing masks and explaining how we do so for our experiments.

At least in its current form, the method faces difficulties scaling up to larger models.

We respectfully disagree with this statement. We have not found any evidence that gradient routing faces difficulties scaling to larger models. We show that the method scales from a 20M to a 0.7B parameter transformer without issue. (As hinted at in the limitations section, we stopped at 0.7B only due to resource limitations.) Intuitively, its simple mechanism of limiting where gradients flow seems like the kind of approach that could scale indefinitely. Of course, if a particular gradient routing application requires extensive hyperparameter tuning, that will make it cumbersome to experiment with large models in practice. This is a limitation of our current approach to robust unlearning.

No comparison to similar methods mentioned that can also achieve localization (DEMix and Interchange Intervention Training, both mentioned in the text).

Thank you for highlighting this opportunity for improvement. We have added appendix J, which includes detailed comparisons to other methods and explains the unique properties of gradient routing that make it suitable to our applications. We are also in the process of re-running our experiments on TinyStories unlearning, substituting DEMix for gradient routing as a localization method. We do not plan to experiment with IIT because it requires a causal model and counterfactual versions of data points, such that it doesn't translate naturally to our applications. (However, we remain excited about IIT's potential applicability to controlling neural network internals.)

Unclear relevance of GR to safety-critical applications mentioned in the paper. Demonstrating benefits in a high-stakes scenario would strengthen the paper.

We give the example of removing bioweapon-related or dual-use capabilities of AI models in section 4.2.3. We also believe that the relevance of our work to preventing Goodharting (discussed in section 4.3 and 5) is potentially a major conceptual contribution to safety in high-stakes situations, given that imperfectly-specified training objectives for a powerful AI could lead to harmful outcomes. Of course, this is a conceptual contribution only. We have also added an appendix L, in which we enumerate foundational challenges to assuring safe use of LLMs that gradient routing may help address.

We strongly agree with the reviewer that the inclusion of more discussion here would strengthen the paper. Does the reviewer think that the elaboration we have added to the appendix is suitable?

Response to questions:

Why does "absorption" happen? Is the absorption effect reliable across multiple settings?

As mentioned in the first paragraph of the discussion section (section 5), we see evidence for absorption in multiple settings: "... 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)." After, we hypothesize about why this happens.

Beyond the explanation given in the paper, there is a simple dynamic that likely explains part of the absorption effect: if a feature exists (in a hidden layer) that is predictive for the downstream task, then it is "easier" for the neural net to learn to utilize this feature than to re-learn the feature in another location. In other words, there may a general inductive bias in neural nets away from feature duplication. If so, inducing a feature to be learned in one location makes it less likely to be learned in another. This is especially true when explicit regularization is applied.

On robust unlearning:

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)?

The validation loss increases because the model overfits to the small number of training points. It will continue to overfit if trained longer. Our approach of training on a few samples is standard in the literature: we cite a variety of works that do this in the discussion of robust unlearning in section 2, for example, [1] and [2]. This is a natural thing to do because, as it turns out, conventional unlearning methods can be undone by fine-tuning on very small samples of forget sets. These sets can be as small as 2 (see [1], table 5)! In contrast, our experiments show that ERA is approximately as robust to retraining as a gold-standard "pure" model that was never trained on forget data. This is evidence that we have truly removed capabilities from the model.

Why does RMU appear to outperform ERA on the 4-stories task?

Like many unlearning methods, RMU works by directly degrading the model's performance on the forget set. In contrast, ERA doesn't target high forget loss directly. Instead, high forget loss for ERA is a side effect of ablating network components responsible for forget set performance. This difference is reflected in the forget set retrainability of RMU vs. ERA, showing that the latter is much harder to retrain (on par with a gold-standard "pure" model never trained on forget data). However, it happens that in this case, fine-tuning with only 4 samples isn't enough to overcome RMU's degraded forget set performance. (16 is sufficient.)

In fig 4(b) it appears that post-finetuning retain set loss is also increased. How much of the validation forget loss in ERA and RMU is just due to overall slightly worse models?

This is a great question! We were concerned about this as well. We quantify this effect by comparison to two other models: a "base" model, trained conventionally (see "Alignment tax" in 4.2.2), and a "control" model (trained identically to ERA, but without gradient routing) (see Figure 4b). The answer is that some of the effect (0.1 cross entropy) is due to the model simply being worse at fitting the data, but the majority of the effect comes from ablation (0.4 cross entropy).

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?

We think this is answered by our clarification that ERA is a pre-training method: ERA works by controlling how representations form in the first place, not by modifying existing representations. It is thus applicable to any scale of LLM.

Summary:

We are grateful for the reviewer's attention to detail in our unlearning results. We hope we have satisfactorily demonstrated the soundness of the work by clarifying the interpretation of results and the extent of our experiments. If the reviewer has any suggestions for improving the clarity here (while respecting the page limit), we would be grateful. We also ask that the reviewer consider increasing their score if they feel we have addressed their points adequately.

[1]: Sheshadri, A., Ewart, A., Guo, P., Lynch, A., Wu, C., Hebbar, V., ... & Casper, S. (2024). Latent Adversarial Training Improves Robustness to Persistent Harmful Behaviors in LLMs. arXiv preprint arXiv:2407.15549.

[2]: Deeb, A., & Roger, F. (2024). Do Unlearning Methods Remove Information from Language Model Weights?. arXiv preprint arXiv:2410.08827.

We thank the reviewer for their thoughtful engagement with the entirety of the paper. We are glad that they found gradient routing "novel and interesting" and relevant to "multiple important problems in safety and interpretability." We are also happy that they acknowledge gradient routing's usefulness in the absence of perfect data labeling. These are the reasons we are excited about gradient routing, too.

The reviewer rightfully identifies limitations with our study of robust unlearning, which we address below.

Response to 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.

This is correct: conventional post-hoc unlearning methods do not require data labeling during pre-training. However, as explained in section 2 (and the top-level comment), conventional unlearning methods are "shallow", acting merely as "safety wrappers" rather than robustly removing harmful capabilities. In this case, our approach sacrifices generality and ease of use in exchange for actually removing the capabilities. As mentioned in section 4.2.3, one such application of this method would be the removal of high-stakes information, like bioweapon-related capabilities.

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?

There are many settings where data filtering is ineffective or infeasible, either because of imperfect labeling (discussed throughout the paper, including section 4.3 and the first part of the discussion section; the effect is also quantified at the end of section 4.2 and in figure 10 in appendix C.1.) or because of entangled capabilities (when some capabilities imply or are implied by other capabilities, discussed in second part of section 5).

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?

We present ERA as a pre-training technique: the model is trained from scratch, with the forget and retain gradients being routed to different parts of the model from the beginning. This method should work in general because we are constraining the learning updates for forget data to occur in parameters which we will then delete.

In sections 4.2.2. and 4.2.3, we see that the method works when applied across two model sizes (small 28M, large 0.7B), and two datasets (TinyStories, FineWeb-Edu / WMDP). In section 4.2.2., we provide extensive comparisons with a variety of baselines, including a gold-standard pure model trained on retain data only, a control model, which is identical to ERA but does not apply gradient routing, etc. These comparisons show that routing is responsible for localization and provide comprehensive metrics.

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.

This is a good point and a genuine challenge to potential gradient routing applications. As mentioned in the introduction and section 3, we imagine gradient routing being used to localize a limited set of capabilities in order to provide targeted guarantees about model behavior, for safety purposes. We share the reviewer's perspective that localizing many capabilities will be challenging. In particular, it would likely harm model performance too much to be useful.

That said, we are excited by the results on absorption, which suggest that exhaustive examples may not be required to localize capabilities (see, e.g., the results of applying gradient routing to only the _California token in section 4.2.1). Instead, a set of limited examples may be sufficient.

Thanks to the authors for their thoughtful answers to all my questions, and for clarifying that ERA is a pre-training technique and not one that can be applied to already trained models. I have two major remaining concerns that prevent me from raising my score.

My first concern is that the value of the contribution may be very limited due to the proposed method only being applicable in contrived settings: where we know what the undesirable capabilities are before training, and yet have only imperfect labeling (or "entangled capabilities") of data that correspond to those capabilities. The paper needs more motivation for why studying this setting is important.

My second (and bigger) concern is that the experimental results do not appear to suggest that gradient routing meaningfully outperforms the baseline of data filtering in these settings (RMU is unfair to use as a baseline since it can be applied to models that have already been trained). Thanks to the authors for highlighting the experiment in Appendix C comparing robust unlearning performance of ERA to data filtering when only limited filtering is possible. While Figs 10 and 11 show that ERA does outperform filtering when a low proportion of forget stories are labeled, the performance difference appears to be extremely marginal, and presumably the data filtering model has lower overall validation loss that may explain even that difference (and lower implementation complexity - no need for potentially expensive modifications to the pre-training method). As a result, I am unconvinced of the utility of gradient routing for the robust unlearning problem from the experimental results in the paper. I urge the authors to find an experimental setting where the data has imperfect labeling or "entangled capabilities" and show that gradient routing can actually meaningfully outperform data filtering in that setting.

Dear Reviewers,

Thanks for the reviews. The authors have uploaded their responses to your comments, please check if the rebuttal address your concerns and if you have further questions/comments to discuss with the authors. If the authors have addressed your concerns, please adjust your rating accordingly.

AC

Hello to future readers!

Thank you for your interest in gradient routing. To prevent possible confusion arising from the review, I clarify a few points here.

Gradient routing is not always sensitive to hyperparameters. We did not observe noteworthy sensitivity to hyperparameters in our experiments on MNIST, CIFAR, or reinforcement learning.

Gradient routing outperforms data filtering. A key finding of the paper is that gradient routing outperforms data filtering when data is partially labeled. This held for the two settings we tested: language model unlearning (section 4.2.2) and scalable oversight of reinforcement learning (section 4.3). These results are discussed in the introduction and discussion sections as one of the primary motivations for the work.

We suspect that gradient routing will scale gracefully to larger models and datasets, although we can't be certain without further experiments. In the appendix, we state that "we suspect that naive attempts to localize large numbers of concepts to unique regions will lead to high training loss." This is not a limitation of gradient routing. Rather, it is a hypothesis about a highly "unnatural" learning target. Put differently, we expect that any method that tries to modularize large numbers of capabilities to distinct subregions will incur high training loss. Of course, such a feat is not necessary for a method to be impactful. Our robust unlearning application requires localizing merely a single data source (the forget set) to a single location.

Gradient routing and its applications are novel, as explained here, in extensive comparisons to other methods in Appendix J, and in the clarification of our contributions here. Gradient routing generalizes previous methods, and furthermore, our safety- and control-oriented applications are, to the best of our knowledge, previously unimagined.