GradMetaNet: An Equivariant Architecture for Learning on Gradients
This work addresses a bottleneck in gradient processing for machine learning practitioners, offering a novel architecture that enhances efficiency and applicability in tasks like model editing and optimization.
The paper tackled the problem of designing architectures for processing gradients in neural networks, which are often used in tasks like pruning and optimization, by introducing GradMetaNet, an equivariant architecture that achieved improved performance on gradient-based tasks such as learned optimization and loss landscape curvature estimation.
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.