GradMetaNet: An Equivariant Architecture for Learning on Gradients

TL;DR

GradMetaNet is a new neural network architecture designed specifically for processing gradients, leveraging symmetry-preserving design, set-based gradient processing, and efficient representation to improve tasks like optimization and model editing.

Contribution

The paper introduces GradMetaNet, an architecture built on principles of equivariance and set processing, with proven universality and superior approximation capabilities for gradient-based functions.

Findings

Effective on gradient-based tasks for MLPs and transformers

Outperforms previous methods in approximation of gradient functions

Demonstrates versatility in optimization, editing, and curvature estimation

Abstract

Gradients of neural networks encode valuable information for optimization, editing, and analysis of models. Therefore, practitioners often treat gradients as inputs to task-specific algorithms, e.g. for pruning or optimization. Recent works explore learning algorithms that operate directly on gradients but use architectures that are not specifically designed for gradient processing, limiting their applicability. In this paper, we present a principled approach for designing architectures that process gradients. Our approach is guided by three principles: (1) equivariant design that preserves neuron permutation symmetries, (2) processing sets of gradients across multiple data points to capture curvature information, and (3) efficient gradient representation through rank-1 decomposition. Based on these principles, we introduce GradMetaNet, a novel architecture for learning on gradients, constructed from simple equivariant blocks. We prove universality results for GradMetaNet, and show that previous approaches cannot approximate natural gradient-based functions that GradMetaNet can. We then demonstrate GradMetaNet's effectiveness on a diverse set of gradient-based tasks on MLPs and transformers, such as learned optimization, INR editing, and estimating loss landscape curvature.