Research map / Architectures and efficiency / RNNs and normalization
GradMetaNet: An Equivariant Architecture for Learning on Gradients
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes GradMetaNet, a permutation-equivariant architecture that processes sets of per-datapoint rank-1 decomposed gradients for curvature estimation, learned optimization, and INR editing.
Problem
Existing learned methods that operate on gradients use architectures not designed for gradient inputs: they ignore neuron permutation symmetries, use inefficient full-size gradient representations, or process only a single averaged gradient, which prevents capturing curvature information.
Approach
Represent per-datapoint gradients via rank-1 decompositions (activations and pre-activation gradients), build G_b-equivariant layers over sets of these neuron-level features (with positional encodings, DeepSets and set-attention layers), and prove that set-of-gradient processing can approximate functions that average-gradient methods cannot.
Key results
- GradMetaNet++ achieves a 26.3% improvement over the non-learnable 128-sample FIM diagonal approximation on curvature estimation (GradMetaNet 13.7%).
- UNF + GradMetaNet reaches a 4.63x (DS + GradMetaNet on CIFAR10) step reduction versus Adam for learned optimizers, with up to 1.78x improvement over other learned baselines.
- On INR editing, WS + GradMetaNet++ reduces CIFAR10 MSE from 5.09 (GMN WS) to 2.99 (x10^-3) and MNIST MSE to 1.60 (ScaleGMN, x10^-2), improving prior state of the art.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use GradMetaNet to predict the FIM diagonal from 128 gradients instead of the direct non-learnable estimate | FIM diagonal estimation · Randomly initialized MLPs | improvement over direct estimate: 13.7% | Helped |
| Use GradMetaNet++ to predict the FIM diagonal from 128 gradients instead of the direct non-learnable estimate | FIM diagonal estimation · Randomly initialized MLPs | improvement over direct estimate: 26.3% | Helped |
| Add GradMetaNet to DS learned optimizer | DS learned optimizer · CIFAR10 | avg. step reduction factor vs. Adam: $2.32$x -> $4.63$x | Helped |
| Add GradMetaNet to UNF learned optimizer | UNF learned optimizer · LM1B | avg. step reduction factor vs. Adam: $1.48$x -> $\mathbf{1.82}$x | Helped |
| Add GradMetaNet++ features to the GMN weight-space model | GMN weight-space model · MNIST INR dilation | MSE ($10^{-2}$): $1.96 \pm 0.02$ -> $1.71 \pm 0.00$ | Helped |
| Add GradMetaNet++ features to the GMN weight-space model | GMN weight-space model · CIFAR10 INR contrast | MSE ($10^{-3}$): $5.09 \pm 0.05$ -> $2.99 \pm 0.03$ | Helped |
Surprising findings
- Plain DeepSets learned optimizer on LM1B performs worse than Adam (0.88x step reduction), while UNF + GradMetaNet reaches 1.82x.
Limitations
- Current implementation limited to MLPs and transformers.
- GradMetaNet++ attention mechanisms limit use in large-scale settings.
- Not yet demonstrated on large LLMs or generative models.
Experimental setup
Datasets: CIFAR10 CIFAR100 FashionMNIST LM1B MNIST synthetic 1D MLP regression
Benchmarks: FIM diagonal estimation (synthetic MLPs) learned optimization on 2-parameter regression, F-MNIST, CIFAR10, CIFAR100, LM1B INR editing on MNIST (dilation) and CIFAR10 (contrast)
Metrics: MSE improvement over baseline approximation avg step reduction factor vs. Adam test loss
Baselines: DWS GMN UNF ScaleGMN DeepSets SGDM Adam K-FAC-style FIM diagonal approximation
Builds on: DeepSets Deep Weight Space networks (DWS) Graph Metanetworks (GMN) Universal Neural Functionals (UNF) PES
Closest papers by meaning
- GradMax: Growing Neural Networks using Gradient Information (ICLR 2022 · similarity 0.88)
- Global curvature for second-order optimization of neural networks (ICML 2025 · similarity 0.86)
- Learning (Approximately) Equivariant Networks via Constrained Optimization (NeurIPS 2025 · similarity 0.86)
- A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Models (NeurIPS 2025 · similarity 0.86)
- On the hardness of learning under symmetries (ICLR 2024 · similarity 0.85)
- Gradient Descent as Loss Landscape Navigation: a Normative Framework for Deriving Learning Rules (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.