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

arXiv: https://arxiv.org/abs/2410.04332 · HTML: https://mlsanity.com/paper/2410.04332/

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 93 of its 140 references that are in the map.
- First posted on arXiv 6 Oct 2024.

## Builds on (most cited first, 30 of 93)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Constitutional AI: Harmlessness from AI Feedback](https://mlsanity.com/paper/2212.08073/) (2022)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Man is to Computer Programmer as Woman is to Homemaker? Debiasing Word Embeddings](https://mlsanity.com/paper/1607.06520/) (2016)
- [RealToxicityPrompts: Evaluating Neural Toxic Degeneration in Language Models](https://mlsanity.com/paper/2009.11462/) (2020)
- [BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models](https://mlsanity.com/paper/2106.10199/) (2021)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [Steering Llama 2 via Contrastive Activation Addition](https://mlsanity.com/paper/2312.06681/) (2023)
- [Machine Unlearning](https://mlsanity.com/paper/1912.03817/) (2019)
- [Concept Bottleneck Models](https://mlsanity.com/paper/2007.04612/) (2020)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [PackNet: Adding Multiple Tasks to a Single Network by Iterative Pruning](https://mlsanity.com/paper/1711.05769/) (2017)
- [InfoGAN: Interpretable Representation Learning by Information Maximizing Generative Adversarial Nets](https://mlsanity.com/paper/1606.03657/) (2016)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [Beyond English-Centric Multilingual Machine Translation](https://mlsanity.com/paper/2010.11125/) (2020)

## Built on by (most cited first, 30 of 34)

- [Inoculation Prompting: Eliciting traits from LLMs during training can suppress them at test-time](https://mlsanity.com/paper/2510.04340/) (2025)
- [Inoculation Prompting: Instructing LLMs to misbehave at train-time improves test-time alignment](https://mlsanity.com/paper/2510.05024/) (2025)
- [Model Spec Midtraining: Improving How Alignment Training Generalizes](https://mlsanity.com/paper/2605.02087/) (2026)
- [Shaping capabilities with token-level data filtering](https://mlsanity.com/paper/2601.21571/) (2026)
- [Distillation Robustifies Unlearning](https://mlsanity.com/paper/2506.06278/) (2025)
- [Beyond Data Filtering: Knowledge Localization for Capability Removal in LLMs](https://mlsanity.com/paper/2512.05648/) (2025)
- [Physics of Skill Learning](https://mlsanity.com/paper/2501.12391/) (2025)
- [Memorization Sinks: Isolating Memorization during LLM Training](https://mlsanity.com/paper/2507.09937/) (2025)
- [On the creation of narrow AI: hierarchy and nonlocality of neural network skills](https://mlsanity.com/paper/2505.15811/) (2025)
- [Access Controls Will Solve the Dual-Use Dilemma](https://mlsanity.com/paper/2505.09341/) (2025)
- [Position: Don't Just "Fix it in Post": A Science of AI Must Study Training Dynamics](https://mlsanity.com/paper/2606.06533/) (2026)
- [Modular Pretraining Enables Access Control](https://mlsanity.com/paper/2607.08077/) (2026)
- [Studying Cross-cluster Modularity in Neural Networks](https://mlsanity.com/paper/2502.02470/) (2025)
- [Position: Capability Control Should be a Separate Goal From Alignment](https://mlsanity.com/paper/2602.05164/) (2026)
- [Anatomy of Post-Training: Using Interpretability to Characterize Data and Shape the Learning Signal](https://mlsanity.com/paper/2606.12360/) (2026)
- [Natively Unlearnable Large Language Models](https://mlsanity.com/paper/2606.13873/) (2026)
- [Prototype Language Models](https://mlsanity.com/paper/2607.00510/) (2026)
- [LACUNA: A Testbed for Evaluating Localization Precision for LLM Unlearning](https://mlsanity.com/paper/2607.02513/) (2026)
- [Train the Model, Not the Reader: Decodability Supervision for Verifiable Activation Explanations](https://mlsanity.com/paper/2607.20379/) (2026)
- [Inducing, Detecting and Characterising Neural Modules: A Pipeline for Functional Interpretability in Reinforcement Learning](https://mlsanity.com/paper/2501.17077/) (2025)
- [Wide Neural Networks as a Baseline for the Computational No-Coincidence Conjecture](https://mlsanity.com/paper/2510.06527/) (2025)
- [ModalImmune: Immunity Driven Unlearning via Self Destructive Training](https://mlsanity.com/paper/2602.16197/) (2026)
- [Self-Ablating Transformers: More Interpretability, Less Sparsity](https://mlsanity.com/paper/2505.00509/) (2025)
- [Preserving Bilinear Weight Spectra with a Signed and Shrunk Quadratic Activation Function](https://mlsanity.com/paper/2509.01874/) (2025)
- [PerNodeDrop: A Method Balancing Specialized Subnets and Regularization in Deep Neural Networks](https://mlsanity.com/paper/2512.12663/) (2025)
- [The Persistent Vulnerability of Aligned AI Systems](https://mlsanity.com/paper/2604.00324/) (2026)
- [Latent Instruction Representation Alignment: defending against jailbreaks, backdoors and undesired knowledge in LLMs](https://mlsanity.com/paper/2604.10403/) (2026)
- [Safety Anchor: Defending Harmful Fine-tuning via Geometric Bottlenecks](https://mlsanity.com/paper/2605.05995/) (2026)
- [Policy-Masked Private Experts: Auditable and Reversible Capability Access Control in Sparse MoE Models](https://mlsanity.com/paper/2608.06690/) (2026)
- [Sparse Competition during Training For the Emergence of Specialized Modules](https://mlsanity.com/paper/2608.30978/) (2026)

## Abstract (arXiv)

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.