Gradient Routing: Masking Gradients to Localize Computation in Neural Networks

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Cloud, Alex, Goldman-Wetzler, Jacob, Wybitul, Evžen, Miller, Joseph, Turner, Alexander Matt
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929608772812800
author Cloud, Alex
Goldman-Wetzler, Jacob
Wybitul, Evžen
Miller, Joseph
Turner, Alexander Matt
author_facet Cloud, Alex
Goldman-Wetzler, Jacob
Wybitul, Evžen
Miller, Joseph
Turner, Alexander Matt
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_04332
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Gradient Routing: Masking Gradients to Localize Computation in Neural Networks
Cloud, Alex
Goldman-Wetzler, Jacob
Wybitul, Evžen
Miller, Joseph
Turner, Alexander Matt
Machine Learning
Artificial Intelligence
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.
title Gradient Routing: Masking Gradients to Localize Computation in Neural Networks
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2410.04332