GradMetaNet: An Equivariant Architecture for Learning on Gradients
Authors: Yoav Gelberg
Yam Eitan
Aviv Navon
Aviv Shamsian
Theo Putterman
Michael M. Bronstein
Haggai Maron
Venue: NeurIPS 2025 poster ()
Keywords: Geometric Deep Learning
Equivariance
Weight Space Symmetries
Weight Space Learning
Learned Optimizers
Primary area: deep_learning
TL;DR: GradMetaNet is a neural architecture that efficiently processes gradients of other networks by exploiting their symmetries and rank-1 decomposition structure, enabling better learned optimizers, model editing, and loss 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., using gradient statistics 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, hindering 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 for *MLPs* and *transformers*, such as learned optimization, INR editing, and loss landscape curvature estimation.

GradMetaNet: An Equivariant Architecture for Learning on Gradients
Yoav Gelberg" University of Oxford yoaverobots, ox. ac. uk
Yam Eitan"
Technion
yam. ei eitanocampus technion. ac.
Aviv Navon Independent Reseracher
Aviv Shamsian
Theo (Moe) Putterman UC Berkeley
Bar-Ilan University
Michael Bronstein University of Oxford, AITHYRA
Haggai Maron
Technion/NVIDIA
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 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 capture curvature information, and (3) efficient gradient representation through rank-1 decomposition. Based on these principles, we introduce GradMetaNet, novel architecture for learning gradients, constructed from simple equivariant blocks. We prove universality results for GradMetaNet, and show that previous approaches cannot approximate natural gradient-bart based functions that GradMetaNet can. We then demonstrate GradMetaNet's effective- ness diverse set of gradient-based tasks MIPS and transformers, such learned optimization, INR editing, and estimating loss landscape curvature.
Introduction
Gradients of neural networks are fundamental objects in deep learning, driving optimization and offering insights into model behavior. Beyond gradient descent and its variants [9] 40 71],, gra- dients used in diverse applications that call for sophisticated processing. These applications broadly span three areas: model optimization, editing, and analysis. In accelerated optimization, several approaches use (multi-)gradinn information improve convergence speed. These approaches range from classical curvature-awe cient approximate curvature-based preconditioners 25
methods like natural gradient descent [ powered by effi- 86), to learned optimizers
27
29
49
55
8 [6] In model editing, gradient information guides pruning algorithms for weight compression 88 and enables targeted behavior modification in large language models [18] [9)] For
32 45
83
model analysis and interpretability, gradient information is used to compute influence functions that trace the impact of individual training samples 28 estimate model uncertainty [17 36), and more.
While most approaches rely on predefined algorithms and hcuristics, recent works explore learnable
processing of gradients for downstream tasks [8 41 59 92] Learned methods offer two key
"Equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

advantages. First, they are essential when no predefined algorithm is known. In model editing, for instance, updating a model using gradients of the editing objective while maintaining performance on validation set requires intricate gradient adaptations that are difficult to model analytically.
Learned approaches can effectively discover these gradient adaptations through supervision [88] [9].
Second, learned approaches offer a powerful mechanism for approximating computationally expensive methods. Methods such as natural gradient descent require hand-crafted approximations for practical application. Learned approaches, if successful, can bridge this gap by discovering efficient approximations tailored for a specific distribution of models. Unfortunately, existing learned approaches architectures not specifically designed for processing gradient information. For example, De Cao et al. [8], Mitchell et al. [59] don't 't account for the parameter symmetries in the gradient representation (as they process gradients of single model), while recent weight-space methods [4] 92 use inefficient gradient representations and process only a single gradient. As result, they unable to capture curvature information that is critical to many tasks.
Our approach. In this paper, we in- troduce GradMetaNet. an architecture designed for learning on gradients of deep models such as MLPs and transformers. GradMetaNet's 's design is guided by the following principles: (1) Respecting symmetries: Neural parameter spaces exhibit inherent symmetries, leading to redun- dancies in gradient representations. GradMetaNet's design is derived to re- spect these symmetries, reducing the number of parameters and improving sample efficiency. As demonstrated in
Adupled grad
GradMetaNet
NN Grads
Weight delta
Viy..ry Var}
GradMetaNel
A
Curvature info
GradMetaNet
.991
Figure 1: We propose GradMetaNet, a novel architecture that processes sets of gradients and can learn to compute gradient adaptations, parameter edits, or scalar values such as curvature information or influence functions.
previous work [1] 12 16 (2, 43 47
79, 9I), equivariant design crucial for learning on data with symmetries. (2) Processing sets of gradients:
Many applications, such as curvature-aware optimization, pruning. and uncertainty estimation, require access to collections of gradients on different datapoints which encode the local geometry of the loss. GradMetaNet is thus designed to efficiently handle sets of gradients computed on different datapoints. (3) Efficient representation: As gradients are extremely high-dimensional, we encode them efficiently. Gradients of neural networks, evaluated on a single data point, admit a rank-1 decomposition which provides compact representation that scales linearly (rather than quadratically) with the number of neurons.
Decomposed gradients have simpler symmetry structure compared to the raw weight representation. allowing us to construct GradMetaNet using simple equivariant building blocks [3] 76 91 and to incorporate attention mechanisms. This enables us to prove universality results still unknown for weight-space models. Additionally, we formally demonstrate the necessity of processing sets of gradients, proving that several fundamental gradient- based algorithms cannot be approximated based on single averaged gradient.
We evaluate GradMetaNet on several gradient learning tasks, comparing equivariant weight-space architectures and other natural baselines. First, we demonstrate GradMetaNet's ability to predict local curvature information using a small sample of gradients, achieving 26.3% improvement over standard approximations, and outperforming other learned approaches. We then integrate GradMetaNet into learned optimizer architectures and apply transformer language models, achieving up to 4.63x reduction in steps compared to Adam, and 1.78x improvement over other learned basclines. Finally, we use GradMetaNet for model editing, where we improve on current coate-of-the-arr results in editing MNIST and CIFAR 10 INRs by up to 22.5%. Across all tasks, GradMetaNet consistently outperforms baselines, highlighting the value of efficient gradient representations and equivariant processing of sets of gradients.
train image classifiers and
2 Related Work
Several recent works have explored methods for learning over neural network weights
5
R
48
61
62
67
73
92
93
These methods often use equivariant archi-
tectures
15
22
53
68
69
72
7
91 that respect the internal symmetry of neural

---

a p Vi, V2. E
(b) g = Iy 72, 79,70 74 E r
(c)
V2
V +
V. E
Figure 2: Gradient information on batch of datapoints in different tensor representations. In 2(a) stack of the weight-shaped gradients, for each datapoint. In 2nb stack of the rank-1 gradient decompositions. In [[c]] the gradient of the average loss on the batch. All of these tensors are naturally computed when backpropagating the loss the batch. GradMetaNet process tensors T.-
network weight spaces. A particularly promising application is processing gradients in weight space for tasks such as learned optimization [1 92 While these approaches respect the natural symmetries single gradient, missing valuable information encoded in
of gradients, they typically operate gradient statistics. Furthermore, these methods process high-dimensional, full-size gradients limiting scalability. Other works, such as De Cao al. (18), Mitchell al. [59], analyze gradients of a fixed pre-trained model, and are not suitable for settings involving different models or evolving parameter configurations (e.g., learned optimization), they are not equivariant to permutation symmetries.
Among classical, non-learned approaches, methods such as K-FAC and its variants [25] 27 55
offer efficient ways to extract curvature information from gradients. These methods, widely used for curvature-awarr optimization [25 27 29
86], pruning 83 88), uncertainty estimation 17 36). and influence function estimation 28 need make probabilistic assumptions the distribution of gradients for computational feasibility. We advance this perspective by introducing learnable approach to modeling these gradient distributions using GradMetaNet, offering greater flexibility and expressiveness.
49
55
Background
Notation. Throughout the paper, we denote models by fo X' + V, where 0 e are the parameters. When fe is multi-layer perceptron (MLP), write the input dimension do. the output dimension as d,, the hidden dimensions as dj,..., d,,-1, and denote the activation function by a. In this case, the parameters are given by 0 (W1, br,..., Wz, by). Given dataset D C X x V, a loss function e: y: x y R. and a batch C DD we denote the loss on the batch by
Ls
Σ (2.8)E8
{fo(r))
B
The parameter gradients of the loss on the batch are denoted by :=:= = VoLs(A). For a single data point (x, y). write L( (0) == (fe(z)) y) and ((,,g) voL(n,))
Rank-1 decomposition of gradients. While general parameter gradients have the same shape as parameters, for many neural architectures, the gradient V V(a.y) admits a rank-1 decomposition through the computation graph of fo. For MLP fe, let (1) and (1) denote æzs activation and pre-activation vectors at layer 1. The backpropagated signal (pre-activation gradient) at layer is
g ( (.(,.)) an
(0) de @l(fo(z).)) v))
(!)
Applying the chain rule yields the following expressions for the weight and bias gradients:
V Ww ((. x) (0) (0 (a(-))) Vo, g (1)
See full derivation in Appendix A.I This decomposition allows us to represent V ((aa))) using the tuple (7)(0) ..... (L)), where (0) = (a (1) ())())))))))))))))))) ())) Note that (!) and g are naturally computed during
3

---

backpropagation, so they can be extracted without additional cost, e.g., using hooks in standard frameworks like PyTorch [66]. See code example in Appendix A.2
Decomposition for transformer gradients. Similar gradient decompositions exist for many other neural architectures [21 27]. In Appendix B.T derive such a decomposition for all components of the transformer. To illustrate the structure of the decomposition, we focus here on the feedforward (MLP) layers, which account for the majority of parameters. Given an input sequence (9) denote the activation of token æ; at layer of the MLP component in 8 = ( 2r), let ai
transformer block, and let 9 be the corresponding pre-activation gradient signal, computed with respect to the loss on the entire sequence (0). We similarly get:
T
T
V (w) -
a
V (.)
= Σ
1=1
We can therefore represent the gradient using 7 YT where 70 = (a gi
2)) In other words, while we incur an additional sequence dimension, the rank- decomposition still holds per-token.
Approximate curvature from gradient statistics. Gradi- ents statistics across datapoints encode information about the local geometry of the loss landscape. For example, the Fisher information matrix (FIM)
V.
Fe
E D-
(Ve log(pe(y (z))) V log(re (yla))
V2
is a second-order approximation of the change in the model's 's predictive distribution Po(y)ax) with respect to change in the parameters, and when 0 is local mini- mum, the FIM is identical to the Hessian. Gradient de- compositions similar to the one discussed above have been used to derive tractable approximation of the FIM 27 [2 49 55 86). In this work directly process sets of gradients to enable learning the local geometry of the loss landscape from their statisties. For background on the Fisher Information Matrix, Pascanu 65
V.
V
Gradients
Figure 3: Fisher information as a second- order approximation to the loss.
Appendix A.3, and for a detailed overview, refer to Martens 54
Symmetries of Decomposed Gradients
Weight-space symmetries. Many neural architectures exhibit parameter space symmetries: parameter transformations that leave the network's function unchanged. In particular, MLPs exhibit well- permuting the neurons of a hidden layer, while keeping track of the connections to the neighboring layers, alters the weight matrices but preserves the function represented by the network. This parameter space symmetry
documented permutation invariance [ 33 61 77 93]:
be expressed by action of the permutation symmetry group G == Sa, x ... x Sd,-1- For O = (Wi,bi.... bi...., WL. by) E θ and h = (T1,..., TL-1) E G, the action h 0 (W/, , b' ...., W b/,)) is given by b' P, b1,
W P.,Wi: Wi:
W/ P, W.P.
b/ P, bi, IE (2...,,,
WL = Wr.P.
b1 = b
where P. E {0. (0.1]xxxd is the permutation matrix corresponding to n. E Ss, (see Figure for an illustration). The action of G preserves the function represented by the network: fg.e fe- Permutation symmetries naturally extend to many other neural architectures, see Kofinas et al.
[1], Lim et al. 48. Zhou al. 992 for detailed discussion.
Decomposed gradient symmetries. Weight-space symmetries naturally extend to decomposed gradi- ents. Following the discussion in Section 3 we define the decomposed gradient space of an MLP fo as
:= (() @ e (4))
(4)
where r := x4 x2, referred as the neuron space the I-th layer, contains pairs (1)) = (a(1) ,9,.9 (9)
of activations and pre-activation gradients. @ denotes direct sum of vector spaces.

---

...
VX.
VW.
(a
ia
9 ()
Vw:
WW
(
(o Lo
g// 1
Figure 4: The action of G Sd, X Sdg-n parameter space performs simultaneou permutation of rows and columns of consecutive weight matrices. In contrast, G's action on the decomposed gradient space permutes the neuron space of each hidden layer independently.
S
G's action extends naturally to r. For decomposed gradient y (7 ,y(0,...,,,,)))
(L)) and
h = (T), . -T-1) G. the action h 7 = (h 7()) ..... h x(D) is given by: h 700 ( " PTT...))) for /or/=10 I,..,, L 1, h h- " for 1, (5)
Let o (m.y) : 2 - I be the function that maps parameters 0 to the decomposition of the gradient V V( = VoL(z.u) + (a.w) is G-equivariant, that is:
(am.v) (h 0) = h & (a.v))
(6)
This equivariance applies to any transformation that modifies extracts information from the function represented by fe using its gradients. As illustrated Figure 4 G's action on is simpler than its action on O since the permutations act independently the different neuron spaces.
Sets decomposed gradients. When computing the gradient of the loss over batch B C D, we naturally obtain set (a (z. E8 of individual gradients As discussed in Section 3 this set contains implicit information about the local geometry of the loss landscape, which is critical for many tasks. Therefore, when designing methods for learning on gradients, it's beneficial to process the entire collection rather than the gradient of the average loss. This intuition is formally motivated Section
As illustrated in Figure 2, gradients across a batch of size b can be efficiently represented as a tensor g Te where the batched decomposed gradient space r3 is
To = p ........ r(( r = R2ddxx2
See formal definitions of all parameter and gradient spaces in Appendix Since the order of
gradients in the batch is arbitrary, the set symmetry group is extended to G₁, == So x G. The action of (r, h) Go permutes the batch indices using and independently applies across the neurons:
((T, h) .g )).. = h g "))
(8)
When modeling functions + T + O, we want to respect decomposed gradient symmetries (G-equivariance) and be independent of gradient ordering (Sy -i-iivarrinnce. We thus aim design models that satisfy ( h) g) h- (g).
Extension to transformers. The analysis above extends naturally to decomposed transformer gradients. The sequence dimension is treated as batch dimension (with optional added sequence PE), and the neuron spaces correspond the input and output of every linear layer and every query/key/valuc/oupuu projection. Additionally, the neuron spaces across the residual stream are tied together, having the same symmetry group. For a detailed discussion, see Appendix B.2
Feature spaces. As with other equivariant architectures, useful to extend I and F3 to more general feature spaces [o[f] and [[ by assigning an f-dimensional feature vector each entry. See Appendix for a formal definition.
5 GradMetaNet
In this section, we introduce GradMetaNet, an architecture for learning on gradients designed to process sets of decomposed gradients in a Gyjequivariant way. As the symmetry structure of Т3 is simpler than that of Os, we can design GradMetaNet using simpler equivariant layers compared to its weight-space counterparts 62 92 93]. Specifically, g E T3 can be viewed as a set of decomposed gradients {71..... 7b). each of which is a concatenation of sets of neuron-
level features 7, (7, (0)
(1)). To further simplify the symmetry structure, we incorporate
By "naturally" we mean that these gradients are always computed when backpropagating the average loss on the batch, and can be be extracted using hooks without additional cost.
5

---

of gradient decompositions
Set of gradients
Positional encoding
Ramk-1 dec.
PE
Wi
y(
W
y(0
Lr.
L Lrrod
Lr
LPool
V W:
V
Neuron space
Neuron space embeddings
Set of space embeddings
Output gradients
embeddings
Figure 5: GradMetaNet pipeline: gradients are decomposed into rank-1 I factors and positional encoding is applied. The imput is then transformed by stack of Lr, equivariant interactions-acet across- sets layers. Lroot pools these representations into [f,], removing the batch dimension. Then stack of Lr layers updates this representation, and LProd maps the result back to O
positional encoding map that enables treat each 7 single bag of neuron-level features. This allows us to implement GradMetaNet using simple, well-established equivariant layers for sets
76 19. As illustrated in Figure GradMetaNet model T is composed of positional encoding map followed by a stack of equivariant linear layers of several types, interleaved with pointwise non-linearities.
31
68
Φ LProd Lr (bx) c... Lr () Loool o Lr, (kr) o...0ooo PE.
The following is a description of each layer:
(I) Similarly to Lim et al. 48 Zhou et al. 93 we use a positional encoding map PE: T3 -
[u[f] that concatenates layer identifier each neurons in the intermediate layers and neuron identifier to each neuron in the first and last layers. We use sinusoidal PE (80).
(II) Lr, To[fin] Io [fur are then parametrized as the ineeractionsaccrosss duced in Hartford et al. [31] (batch dimension x neuron dimension).
layers intro-
(III The pooling layer Lrool [i
[f is designed be S1,-ivvariant and G-equivariant,
and is implemented as TPoci( (g((gg))))) (g) M, M. Pu.. + M2 E r Σb
di
P==1 P==
for learnable M1, M2 Rou xfn
(IV) Lr T[ [[fou]] are parameterized equivariant DeepSets layers 991
(V) Finally, similarly to the generalized product layer Navon al. 622, I [[ θ
applics pointwise MLP to the features associated with the neurons connected to each weight:
w! (L g. "I" 6() (LP (g
For detailed descriptions and implementation details for all layers, a G-invariant head for invariant tasks, and other design choices, see Appendix D.I
Extension to transformers. As formally detailed and motivated in Appendix decomposed transformer gradients have an additional sequence dimension with set symmetry structure, and the neuron spaces across the residual stream are tied together. Therefore, when applying GradMetaNet, treat the sequence dimension as the batch dimension (with optional sequence PE), stack all the neuron features across the residual stream together to a single neuron space, and extend our positional encoding to include the attention head indices. For extended discussion see Appendix
and B.3
GraaMeteett++. Similarly to Kasten al. 39 Lee et al. [46], Romero al. [70,, we can preserve equivariance by replacing summation in steps (II) and (VI) with attention mechanisms. Therefore, we introduce an attention-based variant of GradMetaNet, termed GradMetaNet++4 where Lr, and Lr are implemented using attention across the neuron and batch dimensions. This variant demonstrates significant performance improvements on some tasks, consistent with findings in previous studies. For detailed description of GradMetaNet++ refer to Appendix D.2

---

CIFAR10
CIFAR100
FashionMNIST
Transformer
2.4
1.0
2.2
4.4
0.9
NLL
0.8
2.0
4.0
0.7
3.8
8-
0.6
1.6
3.4
0.5
7
1000
1000
2000
1000 Iteration
2000 Iteration
Iteration
Iteration
Optimizer
Adam
SGDM
DeepSet
DeepSet
GradMetaNet
GradMetaNet
Figure 6: Test loss curves for MLP image classification tasks and transformer language model trained on LMIB, using different optimizers and (learning rate tuned) Adam. Curves are smoothed and averaged over random initializations, with shaded regions representing standard deviation.
6 Theoretical Analysis
Importance of processing sets of gradients. Instead of processing the gradient of the average loss as other gradient learning methods [8] 41 92], GradMetaNet processes sets of (decomposed) gradients computed on individual datapoints. This approach is motivated by the fact that set of gradients encodes strictly more information than the corresponding average gradient, enabling e.g. curvature estimations that cannot be computed using the average alone. This intuition is formalized in Appendix E.T] leading to Proposition E.6 whose informal statement appears below.
Proposition 6.1. Let {V (n.))) (a.y)es be gradients computed on on a set of datapoints C D. There exist functions-such as natural gradient approximations or pruning saliency scores-that cannot reconstructed from the average gradient VB alone.
Expressive power. Restricting model to be equivariant with respect to a specific group action can potentially reduce expressive power [51 52 60). However, we demonstrate that GradMetaNet does not suffer from such limitations. Specifically, we prove universal approximation property for equivariant functions defined on compact domain that doesn't intersect a certain low-dimensional subset & To. Formally:
Theorem 6.2. Let K C. Tb be a compact domain such that K. = Ugeg,g K. and K. E GradMetaNet models are universal approximators (in the -sense))) of continuous Gn-equivariant functions from K. to e
is the set of neuron features that have identical activations and backpropagated signals for least two neurons (see Appendix E. for precise definition). Similarly to Finkelshtein et al. 23 Maron al. [53]. the inclusion of in Theorem 6.2 essential (see Appendix E.3 However, & is union of subspaces of co-dimension 2, has Lebesgue measure 0, and the conditions for membership in are highly unlikely practice, making this assumption mild.
Corollary 6.3. Let B and (V (B))) be in Proposition 6.1 Several natural functions- as natural gradient approximations and pruning saliency scores-can be effectively upproximated by GradMetaNet, which has access {V (2. but cannot be approximated by methods that rely solely on VB.
Appendix In summary, GradMetaNet incorporates meaningful inductive biases for processing sets of gradients while retaining high expressive power, enabling it to represent all continuous functions on sets of gradients under mild assumptions.
For formal statements and proofs of Proposition 6.1
Theorem 6.2 and Corollary 6.3
7 Experiments
In this section, we evaluate GradMetaNet on a variety of learning tasks on gradients. We empirically demonstrate the importance of each of our design principles by ablating components and comparing to other baselines. We then showcase GradMetaNet's effectiveness for three applications: curvature information estimation, learned optimization, and INR editing. We include full experimental details in Appendix | and additional experimental results in Appendix
7

---

we discuss a scaled-up version of this 7.1 Curvature Information Estimation
To demonstrate GradMetaNet's ability to learn to approximate loss landscape curvature, we train it predict the diagonal of the Fisher Information Matrix (FIM) from small samples of gradients. The diagonal of the FIM encodes the curvature along individual parameter directions, capturing the network's sensitivity small change in each parameter.
Data. We first create a set randomly initial- ized MLPs with I-dimensional input and output. We then generate the targets by computing the
GradMMctttiitt
Standard Estimator
GradMetatNet=
DWS
MLP Concet
DWS Average
Batch Asymmetric GradMetaNect
GMN Concat
FIM diagonal for each model over sample of
Asymmetric GradMetaNNt
Average
1024 inputs in [-1, The input to each base-
line is smaller gradient sample computed over 128 points sampled from [-1,1].
1.20
0.50
Baselines, We compare GradMetaNet and GradMttaNet++ against range of baselines, including architectures that (1) rely solely on the average gradient: DWS+Average
Normalized Log
0.20
61]. GMN+Average [1]. (2) use full
0.10
gradients instead of the rank-1 decomposi- tion: DWS+Concat, GMN+Concat, (partially) disregard symmetries: Batch Asym- metric GradMetaNet, Neuron Asymmetric GradMetaNet, and MLP+Concat. See Ap- pendix F. for full descriptions of the baselines.
(3)
10100
530
1000
1500
2000
Training Samples
Figure 7: Comparison of gradient-learning models trained to predict the FIM diagonal from a sample of 128 gradients. Results averaged over seeds; shading represents standard deviation.
Results and discussion. To measure the sample efficiency of each baseline, we repeat the experiment with a varying number of training examples (each training sample is still set of 128 gradients, but we vary the number of such sets the models see during training). As seen in Figure 7 GradMetaNet and GradMetaNett+ perform significantly better than baselines across varying training set sizes. We also compare to a non-learnable approximation that directly estimates the diagonal of the FIM using the 128 input gradients (rather than the full set of 1024 datapoints). Only GradMetaNet and GradMetaNet++ outperform this baseline, achieving an improvement of 13.7% and 26.3% respectively. These results demonstrate GradMetaNet's potential to learn more accurate approximations of gradient-based algorithms. In Appendix G experiment for models with over 1M parameters.
7.2 Learned Optimizers
Optimizing deep neural networks is fundamental chal- lenge in deep learning. While classical optimizers 19 40 63 []] have become standard tools, their design largely relies on intuition and empirical validation. promising alternative is to learn the optimization algo- rithm itself through meta-training (6 g [56]. Learned op- timizers can potentially discover more effective update rules by adapting to the statistical patterns in loss land- scapes. Most optimizers (learnable and hand-crafted) only process the averaged batch gradient, and therefore don't have access local eurvature information. In this experiment, we integrate GradMetaNet into learned op- timizer architectures, providing with the informa- tion required for computing the curvature in the form of sets of gradients on individual datapoints across batches.
Standard Learned Optimizer
GradMetaNet Learned Optimizer
w2
w1
Figure 8: GradMetaNettbasse learned op- timizers can account for lans-lad curvature, avoiding redundant steps.
Setup. Following Harrison et al. (30. Zhou et al. 92], parametrize our learned optimizer rules
0 t O, a (Ve + (0, Ve, (ot ( t)),
(10)

---

where Ve is the current gradient, and { are momentum terms with different decay rates. The standard architecture for F& DeepSets (DS) 9T, which applies a per-parameter MLP the input features. More recently, researchers have explored using equivariant weight-space architectures like Universal Neural Functionals (UNF) [92] to implement Fs. For GradMetaNet-haxee learned optimizers we parametrize F as
F (Oc, Vi, {xet 1 Ji=1, t, F
ForaaMeeaNee (g., (V(") }/= )===))
(1I)
Multiplicative improvement in the number of steps reach Adam's 's best loss. For each task, we run Adam, record its best test loss and the number of steps N required to reach We then run standard and GradMetaNet- learned optimizers, measure their steps to reach L, and report the improvement relative to N. Results averaged over 5 runs. Full Results in Table
Table I:
where g, E Tb is the set of decom- posed gradients across the current batch, and [v("), }. are exponential moving averages of past g,s with different decay rates. The learnable meta-parameters a, B, μ. Φ, and vp are optimized during meta-training (using PES (85) to minimize the training loss after T steps. We evaluate five types of architectures: DeepSets, UNF, DeepSets + GradMetaNet, UNF + GradMetaNet, and learnable SGD with momentum (taking 3 = 0 in Equation I0
Dataset
Optimizer
Avg. step reduction factor vs, Adam (1)
SGDM
1.13x
DS UNF
1.27x
Tasks.
We use three types of optimization tasks: (I) optimizing a 2-parameter linear regression, constructed to have non-diagonal curvature,
F-NNITT
1.16x
DS + GradMertNet
1.44x
UNF GradMetaNet
1.51x
SGDM
1.4xx
(2) optimizing MLPs for classifying CIFARIO.
DS
2.32x
CIFARIO0 [44] and FashionMNIST (90)
CIFARIO
UNF
2.64x
images, and (3) optimizing transformer-based
GradMettNet
4.63x
UNF GradMetaNet
4.26x
language models on LM1B (14). For detailed
SGDM
1.06x
description of each task and the meta-training setup. see Appendix F2
DS
1.79x
CIFAR100
UNF
1.58x
Results and discussion. As Figure g demon- strates, for the 2-parameter regressions task, GradMetaNet-basse learned optimizers can use the curvature of the loss landscape to avoid redundant steps. Consequently, as seen in Table [ and Figure (6) GradMetaNNt-baase learned optimizers consistently outperform baselines across optimization tasks for both MIP and transformers, demonstrating the value of processing sets of gradients. In Appendix G we include additional experiments showing that GradMetaNett learned optimizers can generalize across tasks and architectures and show promise in scaling to larger-scale optimization.
GradMeeaNe
3.15x
UNF GradMetaNer
2.85x
SGDM
1.01x
DS
0 0.88x
LMIB
UNF
1.48x
DS GradMetaNet
1.0xx
UNF GradMetaNet
1.82x
7.3 INR Editing
Implicit Neural Representations (INRs)
664 78
neural networks to encode images as functions.
In this experiment, we explore the task INR editing, where the goal is to adapt the weights of an INR to modify the image it represents. This involves directly adjusting the INR's parameters by learning a metanetwork predicting parameter delta AO, and updating the model with 0' = 0 + A0.
Data. Following previous works 38
93 94], we use two standard benchmarks: figure dilation for MNIST INRs and contrast enhancement for CIFAR-IO INRs. For each INR, we compute the parameter gradients with respect to the MSE loss between the INR output and the target edited image, evaluated over randomly sampled points. The input data consists of both the parameters of the INR and the corresponding gradients.
Baselines. We evaluate GradMetaNet and GradMetaNett+ in combination with several weight-space architectures. GradMetaNet and GradMetaNet++ process the gradients to produce outputs in O which are then used as additional weight features for the base weight-space network. This hybrid ap- proach is compared against two baselines: (1) the base weight-space network, and (2) the base weight- space network augmented with probing features. Following Kofinas et al. [41]. the probing features activations evaluated at randomly sampled grid points, incorporated additional bias features.

---

Table Results for the INR editing tasks on MNIST (dilation) and CI (contrast). We report
the MSE (J) in 10-2 for MNIST and 10 for CIFARIO, averaged over 3 seeds.
MNIST
CIFARIO
DWS 161
GMN [41]
ScaleGMN 138
DWS 161
GMN [41
ScaleGMNN [38
WS
2.29 - 0.01
1.96 + 0.02
1.99 - 0.02
6.57 0.02
5.09 = 0.05
5.23 = 0.13
WS Probing
2.36 1: 0.06
1.85::0000
1.92 1 0.04
4.22 0.08
3.81 0.02
3.87 1: 0.05
WS
GradMetaNct
2.28 +: 0.02
1.70 0.01
1.70 1.0000.00
4.10
3.65 + 0.01
3.69 + 0.09
WS
GradMetaNect++
1.95 + 0.01
1.71 + 0.00
1.60 0.00
3.86 + 0.02
2.99 + 0.03
0.03
3.00 + 0.00
Results and Discussion. As seen in Table 2 both GradMetaNet and GradMetaNet++ consistently improve the performance of weight-space models, achieve greater performance gains than probing. and improve the current state-of- of-tth-arttt for weight-space model editing.
8 Conclusion and Limitations
Conclusion. We introduce GradMetaNet, an equivariant architecture for processing sets of decomposed gradients, supported by both theoretical and experimental results. Theoretically, we demonstrate that under mild assumptions GradMetaNet can approximate any function processing sets of gradients, and that average gradient methods are unable to approximate several natural gradient-based functions. Experimentally, we demonstrate GradMetaNet''s ability to predict local curvature, enhance learned optimizers, and achieve statc-of-the-ar performance in model editing.
Limitations and future work. The current implementation of GradMetaNet is limited to MLPs and transformers; future work, we hope extend GradMetaNet to support other neural architectures. Additionally, GradMetaNett+'s use of attention mechanisms limits its usage in large-scale settings. Finally, scaling GradMetaNet, GradMetaNNtt++ or their variants to larger models, such as state-of-the-art LLMs and generative models, is an interesting direction for future work.
A Acknowledgments
YG is supported by the UKRI Engincering and Physical Sciences Research Council (EPSRC) CDT in Autonomous and Intelligent Machines and Systems (grant reference EP/S024050/1), MB is supported by EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub on Mathematical Foundations Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1. HM is Robert J. Shillman Fellow and is supported by the Israel Science Foundation through personal grant (ISF 264/23) and an equipment grant (ISF 532/23).
References
[1] Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.044. 048366 2022.
[2] Marjan Albooyeh, Daniele Bertolini, and Siamak Ravanbakhsh. Incidence networks for geo- metric deep learning. arXiv preprint arXiv:1905.11444 2019.
[3] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2): 251-276, 1998.
[4] Shun-ichi Amari and Hiroshi Nagaoka. Methods of information geometry, volume 191. Ameri- can Mathematical Soc., 2000.
[5] Bruno Andreis, Soro Bedionita, and Sung Ju Hwang. Set-based neural network encoding. arXiv preprint arXiv:2305. 16625, 2023.
[6] Marein Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
[7] Erik Bekkers. B-spline cnns on lie groups. arXiv preprint arXiv: 12057, 2019.
[8] Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning a synaptic learning rule. Citeseer, 1990.
10

---

[9] Léon Bottou. Lcrge-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT 19th International Conference on Computational StatisticsParis France, August 22-27, 2010 Keynote, Invited and Contributed Pupers, pages 177-186. Springer, 2010.
[10] G. Bradski. The OpenCV Library. Dr. Dobb's s Journal of Software Tools, 2000.
[11] Johann Brehmer, Sönke Behrends, Pim de Haan, and Taco Cohen. Does equivariance matter at scale? arXiv preprini urXiv:2410.23179, 2024.
[12] Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34 (4):18-42, 2017. 2017.
[13] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Velićković, Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint MxXi:2214.13472 2021.
[14] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv: 13/.33005 2013.
[15] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999. PMLR, 2016.
[16] Taco Cohen et al. Equivariant convolutional networks. PhD thesis, Taco Cohen, 2021.
[17] Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace eeuux-effortless bayesian deep learning. Advances in Neural Information Processing Systems, 14:20089-201033 2021.
[18] Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge in language models. arXiv preprint 1rXiv:2104.08/1 2021.
[19] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[20] Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman. Clas- sifying the classifier: dissecting the weight space of neural networks. In ECAI 2020, pages 112. IOS Press, 2020.
[21] Runa Eschenhagen, Alexander Immer, Richard Turner, Frank Schneider, and Philipp Hennig. Kronecker-factone approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems, 36:33624 33655, 2023.
[22] Carlos Esteves, Christine ing so (3) equivariant representations with spherical enns. In Proceedings of the European Conference on Computer Vision (ECCV), pages 52-68, 2018.
Ameesh Makadia, and Kostas Daniliddss. Learn-
[23] Ben Finkelshtein, ismail ilkan Ceylan, Michael Bronstein, and Ron Levie. Equivariance everywhere all at once: A recipe for graph foundation models. arXiv preprint arXiv:2506.1422 14291, 2025.
[24] Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning, pages 3176-3176. PMLR, 2020.
[25] Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018
[26] Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent, Fast approximate natural gradient descent kronecker factored eigenbasis. Advances Neural Information Processing Systems, 31, 2018.

---

[27] Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for con- volution layers. In International Conference on Machine Learning, pages 573-582. PMLR. 2016.
[28] Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. arXiv preprint arXiv:2308 03296, 2023.
[29] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. International Conference on Machine Learning, pages 1842- 1850. PMLR, 2018.
[30] James Harrison, Luke Metz, and Jascha Sohl-Dickstein. A closer look at learned optimization: Stability, robustness, and inductive biases. Advances in Neural Information Processing Systems, 35:3758-3773, 2022.
[31] Jason Hartford, Devon Graham, Kevin Leyton-Brown, and Siamak Ravanbakhsh. Deep models of interactions across sets. In International Conference on Machine Learning, pages 1909-1918.
PMLR, 2018.
[32] Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In IEEE international conference on neural networks, pages 293-299. IEEE, 1993.
[33] Robert Heeht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pages 129-135. Elsevier, 1990.
[34] Vincent Herrmann Francesco Faccio, and Jürgen Schmidhuber. Learning useful representations of recurrent neural network weight matrices. arXiv preprint arXiv:2403. /1998, 2024.
[35] Eliahu Horwitz, Bar Cavia, Jonathan Kahana, and Yedid Hoshen. Representing model weights with language using tree experts. arXiv preprini urXiv:2410.13569) 2024.
[36] Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference Machine Learning, pages 4563-4573, PMLR, 2021.
[37] Jonathan Kahana, Eliahu Horwitz, Imri Shuval, and Yedid Hoshen. Deep linear probe generators for weight space learning. arXiv preprint arXiv:2410.108) 2024.
[38] Ioannis Kalogeropoulos, Giorgos Bouritsas, and Yannis Panagakis. Scale equivariant graph metanetworks. arXiv preprini arXiv:2406. 10685, 2024.
[39] Yoni Kasten, Wuyue Lu, and Haggai Maron. Fast encoder-based 3d from casual videos via point track processing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[40] Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[41] Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan Burghouts, Efstratios Gavves, Cees GM Snoek, and David W Zhang. Graph neural networks for learning equivariant representations of neural networks. arXiv preprint arXiv:2403. 2024.
[42] Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pages 1885-1894. PMLR, 2017.
[43] Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution neural networks to the action of compact groups. In International conference on machine learning, pages 2747-2755. PMLR, 2018
[44] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical Report 4, University of Toronto, 2009.
[45] Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances neural information processing systems, 2, 1989.
12

---

[46] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. -nvvaiartiinnnnnaaaantta neural networks, In
Set transformer: A framework for attention International conference machine learning, pages 3744-3753. PMLR, 2019.
[47] Derek Lim, Joshua Robinson, Lingxiao Zhao, Tess Smidt, Suvrit Sra, Haggai Maron, and Stefanie Jegelka. Sign and basis invariant networks for spectral graph representation learning. arXiv preprint arXiv:2202. 3013, 2022.
[48] Derek Lim, Haggai Maron, Mare T Law, Jonathan Lorraine, and James Lucas. Graph metanet- works for processing diverse neural architectures. arXiv preprint arXiv:2312.04501, 2023.
[49] Jingyuan Liu, Jianlin Su, Xingeheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint arXiv:2502.16982. 2025.
[50] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL httpss//ppennreeiee wtwp://rrriiie.ee/fro.mm.
id=Bkg6RiCqy7
[51] Haggai Maron, Heli Ben- Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Advances in neural information processing systems, 32, 2019.
[52] Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality of invariant networks. In International conference machine learning, pages
PMLR, 2019.
[53] Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements. In International conference on machine learning, pages 6734-6744. PMLR, 2020.
[54] James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
[55] James Martens and Roger Grosse. Optimizing neural networks with frontceer-factored approx- imate curvature. In International conference on machine learning. pages 2408-2417. PMLR, 2015.
[56] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference on Machine Learning, pages 4556-4565. PMLR, 2019.
[57] Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl- Dickstein. Practical tradeoff' between memory, compute, and performance in learned optimizers. In Conference on Lifelong Learning Agents (CoLLAs), 2022. URL http:://ttttbbb com/ google//lerrodd_pppiiiaaaaa
[58] Luke Metz, James Harrison, C Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint MXXi::2111.09766 2022.
[59] Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D Manning. Fast model editing at scale. arXiv preprint arXiv:2/ 1309, 2021.
[60] Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 4602-4609, 2019.
[61] Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equivariant architectures for learning in deep weight spaces. In International Conference on Machine Learning, pages
PMLR, 2023.
[62] Aviv Navon, Aviv Shamsian, Ethan Fetaya, Gal Chechik, Nadav Dym, and Haggai Maron. Equivariant deep weight space alignment. arXiv preprint arXiv:2310.1 13397, 2023.
[63] Yurii Nesterov. A method for unconstrained convex minimization problem with the rate of convergence o (1/k2). In Dokl. Akad. Nauk. SSSR, volume 269, page 543, 1983.
13

---

[64] Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, and Steven Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceed- ings of the IEEF/CVF conference on computer vision and pattern recognition, pages 165-174, 2019.
[65] R Pascanu. Revisiting natural gradient for deep networks. arXiv preprint 1rXiiv/301.3550 2013.
[66] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[67] William Peebles, Iijaa Radosavovic, Tim Brooks, Alexei A Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209.12899 2022.
[68] Charles R Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652-660, 2017.
[69] Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Equivariance through parameter- sharing. In International conference on machine learning, pages 2892-2901. PMLR, 2017.
[70] David Romero, Erik Bekkers, Jakub Tomezak, and Mark Hoogendoorn. Attentive group equivariant convolutional networks, In International Conference on Machine Learning, pages 188881199. PMLR, 2020.
[71] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-prop proccppppgggting errors. nature, 323 (638888)) 536, 1986.
[72] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. International conference on machine learning, pages 9323-9332. PMLR, 2021.
[73] Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-supervised representation learning on neural network weights for model characteristic prediction. Advances in Neural Information Processing Systems, 14:16481-16444 2021.
[74] Konstantin Schürholt, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Hyper- representations as generative models: Sampling unseen neural network weights. Advances in Neural Information Processing Systems, 35:27906-27928 2022.
[75] Konstantin Schürholt, Michael W Mahoney, and Damian Borth. Towards scalable and versatile weight space learning. arXiv preprint 240,0297,099997
[76] Hadar Serviansky, Nimrod Segol, Jonathan Shlomi, Kyle Cranmer, Eilam Gross, Haggai Maron, and Yaron Lipman. Set2graph: Learning graphs from sets. Advances in Neural Information Processing Systems, 33:22080-22091, 2020.
[77] Berfin Simsek, Johanni Brea, Bernd Illing, and Wulfram Gerstner. Weight-space symmetry in neural network loss landscapes revisited. 2020.
[78] Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in neural information processing systems, 33:7462-7473, 2020.
[79] Behrooz Tahmasebi and Stefanie Jegelka. The exact sample complexity gain from invariances for kernel regression. Advances in Neural Information Processing Systems, 36, 2023.
[80] Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Tridovich-Keil. Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions low dimensional domains. Advances neural information processing systems, 33:7537-7547, 2020.
14

---

[81] Tijmen Tieleman. Lecture 6.5-mopprp:: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26, 2012.
[82] Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and Ilya Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint urXiv22022.11445 2020.
[83] Tycho FA van der Ouderaa, Markus Nagel, Mart Van Baalen, Yuki M Asano, and Tijmen Blankevoort. The IIm surgeon. arXiv preprint arXiv:23 17244, 2023.
[84] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[85] Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In International Conference on Machine Learning, pages 10553-10563. PMLR, 2021.
[86] Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint (rXiv:2409.11322 2024.
[87] Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Michael A Osborne, and Ingmar Posner. Universal approximation of functions sets. Journal of Machine Learning Research, 23(151);
1 1-56, 2022.
[88] Chaoqi Wang, Roger Grosse, Sanja Fidler, and Guodong Zhang. Eigendamage: Structured pruning in the kronecker-factored eigenbasis. In International conference on machine learning, pages 6566-6575. PMLR, 2019.
[89] Boris Weisfeiler and Andrei Leman. The reduction of graph to canonical form and the algebra which appears therein. nti, Series, 2(9):12-16, 1968.
[90] Han Xiao, Kashif Rasul, and Roland Vollgraf. Rashion-mnist novel image dataset for 2017.
benchmarking machine learning algorithms. arXiv preprint
[91] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
[92] Allan Zhou, Chelsea Finn, and James Harrison. Universal neural functionals. arXiv preprint arXiv:2402.05233. 2024.
[93] Allan Zhou, Kaien Yang. Kaylee Burns, Adriano Cardace, Yiding Jiang, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. Advances in neural information processing systems, 36, 2024.
[94] Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Neural functional transformers. Advances neural information processing systems, 36, 2024.
A Extended Background
A.1 Gradient Decomposition for MLPs
Using the notation introduced in Section 3 for 1 = 1,... L, n = 1,....,...11 m = 1, -1...., we apply the chain rule to get
(Vw. Crm (m (0) nm Σ
0L (L( a(ul)))
)
gg (1) )a (anaaaa(a((--))
a(Wi)
(12)
= (m)) (a(l())
15

---

{1 n = k
where ōn, is the Dirac delta defined by 8n,k
Similarly
n n k.
di Σ
0L (L(
)((()))
di
(Vb.
9 1 mm0 (g
(13)
a(u(0)x
d(bi)
Therefore,
Vw, L(, (0) = gd (T Vb, (1)) - g0()
(14)
as discussed in Section
A.2 Extracting Activations and Pre-Activation Gradient Signals
As mentioned in Section 3
the activations (a and pre-activation gradient signals (g used for the gradient decomposition are naturally computed during backpropagation and don't 't need to be recomputed. The following is a PyTorch code example for extracting these components without additional cost using forward/hackwart hooks:
import torch
import torch. nn import torch. functional F
class (nn. Module)
def __init__ (self): super (MLP, colf)..-nnit_.0
self fel nn Linear (8, 32) self fcl mn.eine (32, 16)
aelf.fc3 mn Linear (16, 3)
def forvard(self, z):
x F.relu fc) x F.relu (x))
x self (x)
return
activations 0 tangents = {
def forward_ book (module, inp out): activations [module] inp [0] detach
def backvard_hook (module, grad_inp grad_out): tangents [module] [mddle]s[moduue] grad_out [0] ()
model M)P()
Set hooks fcdel.fcl. ooo.ll......eeaa.....orraoooooor model aoddd.....rreeeee......aaaaaaaar
model.fc2. model feeel.ff3.regissttee aeeet.rce.eeaaaer...rrrrooooooor model.fca moodh:.f..rggaaee...l..oceaarrrd baw.caawwwra.ooooaaaaaaaaaa.aaaa
Backpropagate loss
= corch.rande (4, 8) (batch, input)
target torch. randn (4, 3) (batch, input)
output model (x)
loss mame loss (output, target)
loss backvard
print (activations)
print (tangents)
16

---

A.3 The Fisher Information Matrix and Its Uses
use the Fisher Information Matrix (FIM) to approximate the curvature of the loss landscape. Below, we define the FIM and present two common gradient- based algorithms that utilize it. The FIM has numerous other applications, this section serves only as a basic introduction. For more comprehensive overview, we refer readers to
Many gradient-based based algorithms 17 27
55
83
88
The Fisher information matrix. Consider supervised learning problem of predicting outputs y from imputs x X. We assume a probabilistic model of the form po(y)z) x) = p fo (z)), where is called the likelihood. For classification tasks may assume softmax likelihood, p(y = fz)) = softmax( ((x))) )a= and for regression, we usually take a Gaussian likelihood ply fz)) Nyy fx)) I). po (y]z) is called the predictive distribution. The FIM is defined by
== Eg TaFa-DDarrpp pryrre (Vo log(pe [x]] Vo log(po(y)az) ())) ") E e +
(15)
The FIM is second order approximation of the change in the model's predictive distribution with respect the parameters
E (DKL (po(y)æ) (o+s(yla))) = F5+ 0
(16)
and thus contains information about the geometry of the space distributions and the loss landscape. Additionally, when is a local minimum, the FIM identical to the Hessian the loss. The FIM is computed using gradients of the model single datapoint and specifically using only gradients of the output Vo,) For regression
BanD.y-roull (Vo log(po x)) )) Ve log(po()(=))))
((To)110(e)) - -u1²)
1 -w)1²)
= EewD. (iv.Doo)W
(17)
= -EED Vofolz]Vofo(z) Byg (llfo()) - -11(3)
EEND (Vefel=)Vofo(z)))
and for classification
F Ee DeeDDavvee (
(Vo (oolo(po(y)l)))J4 (og(ro)))))))
E
poly Vo log(poy Vo log(vol k((z))
(18)
E
Vero = k/ Voro = kk(=)
Poly k|x)
The natural gradient. FIM
The natural gradient is defined by preconditioning the gradient using the
Waa = Fo
(19)
Motivated from the perspective of information geometry [4, the natural gradient defines the direction in parameter space that gives the largest change in the loss per unit of change in the predictive distribution of the model, measured by KL-divergence. This is to be contrasted with the standard gradient, which is the direction that gives the largest change in the loss per unit of change in parameters. Natural gradients are fundamental tool in optimization and can be used accelerate training 25 27
Natural gradient optimization requires dynamic computation, storing, and inversion of the FIM, whose size grows quadratically with the number of parameters, making it intractable at scale. This necessitates the use of approximations, such K-FAC 26 27 [55
FIM usage in pruning. In pruning, we want to assign each weight a trained model saliency score that indicates its importance to the model's performance. The goal is to remove certain percentage of weights with the lowest saliency, resulting compressed model with fewer parameters that still performs relatively well. The original OBD [45] and OBS [32] algorithms use the Hessian to
17

---

compute saliency scores. However, since the model is trained, the parameters are likely close to local minimum, and the FIM often used as an approximation of the Hessian.
Given parameters of a trained model 0, the OBD pruning saliency scores are given by
lORD 0² @ diag( Fo);
(20)
where l is a matrix of the square each parameter, and @ stands for point-wise product. The optimal brain surgeon (OBS) pruning saliency scores given by
OOBS 0² 2 F
(21)
where 2 denotes point-wise division.
B Extension to Transformers
In this section, extend the results from the main text-originally detailed for MLPs-to transformer architectures. Similar extensions apply to a wide range of neural models, including CNNs, RNNs, and state-space models, as the key requirement is that the architecture consists of linear layers interleaved with nonlinearities. We begin by presenting the gradient decomposition for transformers, and then analyze its symmetry structure and the applicability of GradMetaNet.
B.1 Gradient Decomposition for Transformers
As mentioned in Section 3 the gradient decomposition used in the paper generalizes to other types of neural architectures. For example, Grosse and Martens [27] discuss an extension to CNNs and Eschenhagen et al. [2I] generalized this decomposition to other modern neural architectures, including transformers. Most transformer parameters are split between the fully-connected components (usually called FFNs MLPs) and the attention layers. We have covered the MLP case in Section 3 and cover the rest in this section. Throughout this section, we use the following notation:
d T = sequence length, d dmodel, h heads, d, h da FFN expansion, L = blocks.
For the reader's convenience, we use brown for transformer block indices, blue for attention head indices, and green for token indices. This notation and the notation for the of the section largely mirror Vaswani et al. 84
Forward computation of transformer block (. Table details the forward pass of the 1-th transformer block on an input sequence s = (æ1, TT). is the index of the transformer block and is the head index.
Table 3: Transformer forward-pass.
Object
Definition
Shape
Hidden input
((-1)
T x d
Q0.)) H w (1)) (1.))
Queries
T dr
Keys
((.)) H (1,)) + +b
T x dr
Values
((,) H (1 (1.)) K UiT /Vda)
+ (1))
T x de
Attention weights
A A(.) = softmax
T x T T' x dR
Head output
((.)) (l.) ((,)
Merged heads
O = (O.1))
(L.)
w + bb
(1)
T x d T x d T x di T Fxd d T
Post-MHA state
H LayerNorm
H (11)
+00
Feed-forward pre-act.
P() H 6
Feed-forward out.
U a (P W + b (5)
Layer output
H = Layer Norm
H + UU))
The transformer block parameters are presented in Table
Notes:
In most implementations, the key, query, and value projections have no bias terms, i.c.
b (1-)) b (.j) ((j) 0.
(22)
We include these bias terms for generality, but they can be removed in most cases.
18

---

Table 4: Transformer parameters.
Parameter
Description
Shape
w w (1.)
query projection (head j)
d x dg
W wils)
(())
key projection (head
d x da
w
value projection (head j)
x d
(ad)
query bias
dę
key bias
value bias
ds
output projection (concat. heads - d)
(h dx) d
D
output bias
d
W
FFN expansion (d)>dy)
d x da
b 1
FFN bias (layer 1)
dn
W
FFN contraction (dı (ddr-dc
dæ x d
b
FFN bias (layer 2)
Our derivation follows the post-LayerNorm (Post-LN) convention: residual add LayerNorm
(as in the original transformers paper Vaswani et al. 84 The analysis for Pre-LN transformers requires slight adjustment the derivation.
Outside of the transformer block we would also typically have: the token embedding matrix E R|v positional embeddings PE E T (or sinusoidal schedule), and an unembedding matrix E [d
Back- propppppop (pre-activation) gradient signals. Notice that all of the transformer parameters act as linear transformations (that are then sometimes passed attention computations, LayerNorms, etc.). This means that we can use the observation from Appendix A. regardiin gradient computation for linear layers. Specifically, for every tensor computed in the transformer forward-past (Table that is the output of a linear transformation,
store the gradient w.r.t. that tensor:
(1.i) OL.
(1.))
dL.
g (1i)
OC.
9Q (Q)
9K
OKOJ)
OV
Oc.
Oc.
(1) OC.
90
8000
9
92
All of these tensors, referred to as langents, share the same shapes as their forward counterparts.
Outer-product parameter gradients. Because every weight matrix appears in affine map Y XW + b, as analyzed in Appendix A.1 its gradient factorizes exactly into an activation block X and tangent block g:
- - g.
V 17.
(23)
Carrying this out for all parameters in the transformer block yields
V Wiits = (H)
9Q
(1-)
b C. (90
T (1))
V WgunCo = (-
(1-))
n C. (9K (1)) T
9K
WW (H
gy (1-)
mnn L. (gV (1.) ) T
V WWoC. ((.)
o (4,)
go 0
1o C. (g0) T
Wilhfe = (HO)T 91
1 Le (g) (1) TT
C, = (P()))
(1)
69 L C. = (g) (1) TT)TL 1
19

---

Token-wise view ("sum of rank-1" form).
Unfolding, for example,
V Wwi (H")-1)) T
(1.)) 9Q
Σ 9o (H, (H,
(24)
reveals that each weight gradient is sum of rank-1 outer products. Storing the pair (H 9 for every token is therefore sufficient reconstruct V wwiss exactly, and identical statements
(1:1)
hold for W (LJ)
W!
(.)) W w! and W!
Compact gradient representation. define:
For transformer block / = I,..., L, and head j = L,... h,
g
H
EETT
rres
pres = - Txdx4
res
activation
attn. out, Langents
post-LN
FPN, langents
(1,1%))
(111th)
(1,1:h)
E D = RR xx(hde))x4
attn
9Q
K
(25)
query tangents
key langents
value tungents
MHA outputs
T Thidden = R T Txdxxxx
hidden
hidden rep.
FFN, tangents
Collecting the whole set, we get the tensors
g
[
attu
hidden
E
@ Date
+ Phidden
(26)
and
(g
ETT
1 +.... f
(27)
B.2 Transformer (Decomposed) Gradient Symmetries
The permutation symmetry groups the parameter spaces of general architectures, and transformers in particular, were analyzed in Kofinas et al. [41, Lim et al. (48), Zhou et al. [92]. Kofinas et al. μT. Lim al. 48 identify permutation symmetries with automorphisms of the computation graph of fe. and Zhou et al. [92] analyzes the permutation symmetries of multi-dimennional tensors. To give flavor of the adaptations needed in the transformer case, we first look at the effects of the residual connections. Intuitively,
need to tie together the neuron spaces of dimension d (i.e., ( O0,
H and ( under the same symmetry group (S2) because of the residual connections. With this "symmetry tying" the residual connections and LayerNorms preserve permutation symmetries, since if ④5 V Rd - Rdd are Sa-equivariant functions, we have
& + (o ( (a = ((x) (x))=or (x) o. 0. ( σ. - (x),
and
LayerNorm 2)) = a LayerNorm(z)).
The analysis of the permutation symmetry group of transformer weight spaces provided in Kofinas al. [41]. Lim al. 48]. Zhou et al. [92 follows similar observations. The resulting symmetry group is
x transformer blocks
G S lokens
Sa residual
(Saz)*
S Sen
(28)
stream
output
Where, in our decomposed gradients case, Sr acts on the sequence dimension of all spaces, Sg acts
on the second axis (the d-dimension) of all I each (Sa,) )*h acts independently on the ((4)) and
Oli)) components of each head rattn, and each Sat acts on the hidden dimension of the FFN represented in T Phidden
We note that transformer parameter spaces exhibit other neural symmetries that are not modeled by the permutation symmetry group G. These symmetries include ReLU scaling symmetries [38 and general attention symmetries (the transformation (Q(1.)) k(.))) (Q(1)) R. K (1.)) R for
20

---

R GLd, (R) results in the same attention matrix). Accounting for these symmetries is left for future work.
B.3
Adapting GradMetaNet to Transformers
To implement a GradMetaNet version that process transformer gradients, need make the following adaptations. First, we treat the sequence dimension as batch dimension, optionally with additional positional encoding for the token index. Note that this positional encoding is not strictly required since, as can be seen in Equation 24 the full gradient is a sum over the rank-1 components and is therefore Sy-invariant function of them. We then treat rise rgttn, and Thidden as we treat neuron spaces with an additional positional encoding for each attention head to convert the ((nde) equivariance of Lr to ( e'eeevavvraance As mentioned in Appendix B.2 we treat the L copies of T7 as single neuron space, since the symmetry structure of the residual stream tied together.
C
Gradient and Weight Spaces
This section formally defines the vector spaces used to represent (sets of) gradients and weights. For batch size and feature dimension f, the feature vector spaces [[/] and To[/], [[]] and Ox[f] are defined by
[[ʃ] == -o I1 r
Tb = r [1] + 2 rr
M o... tiim a llane
(29)
L1/ = W + 11 e wi @ u 11
where,
rw == Rd :=
w00 = R4 xd: 1x1 w [1] : bxdxx 1x/
(30)
4 40 = R2 x - 10 = = xx
Note that r = r, P = ri [[] + and θ [1] A..
D Architecture Details
The following a detailed description of each of the layers used GradMetaNet.
D.1 NradMetaNet
The positional encoding map. Similarly to Lim al. 488) Zhou al. 193 use positional encoding map PE T3 + Tolf] that concatenates a layer identifier to neurons in intermediate layers and neuron identifier each neuron in the first and last layers, i.e.
PE(g k.j.
i ein (j)
PE(g )...
g Clayer (1)
for 1 - 1,..,, 1,
PE(g(D) )a.. g Cour (i)
where - denotes concatenation along the feature axis. and Cout assign unique identifiers to each neuron in the input and output layers, respectively, and elayer assigns unique identifiers each hidden layer. We implement all encoding maps using sinusoidal positional encoding [80).
Graddent-sa--tograaiennts layers. Ir, [o[ + [o[Som] are parametrized similarly to the interactionss-ccross-cee layers introduced in Hartford et al. [31] and are implemented
L Σ
de
Lrs (g()
-Mg M2 Σ
9 g
Ms
Σ
g
+ MA
Σ>
g (31)
A.j::
0=0
P=0
P=
for learnable Mi, M2, Ma, M4 e Ro.
21

---

Graddeen-set--t-grrddeed pooling layer.
LPool rolfml - foul is implemented as
du
LPool (g )22 Mi
Σe
(1)
M2 2
2 K= 1
(32)
==1
for learnable M1, M2 E Rooxx
Gradient-to-gradient layers. networks [91], and take the form
Tr [(n - [ are parameterized as equivariant DeepSets
de
(r( Mig
M + M2
9 P.
(33)
P=
for learnable M1, M2, E Xon
-radient-to-weight component. Similarly to the generalized product layer in Navon et al. 622 IProd [[fi]] O applies pointwise MLP the features associated with the neurons connected to each weight, or in the case of biases, to the feature vectors corresponding to the respective neuron.
/ ( ( g ( 6.0 (L
(34)
D.2 GradMetaNet++
Similarly to GradMetaNet, A GradMetaNet++ model comprises updates of different types: positional encoding layer PE, gradient-set sectienttsettt--eeee tetieettteettt---rdiee----eee updates Ur, and srradeen--eetto-weelt component Urrod. GradMetaNett+ model parameterized as
+ Uprod a Ur, (k)
Ur. (1) PE.
(35)
We now describe each of these layers.
The positional encoding map. The positional encoding map used for CradMetaNet++ identical to the one used in GradMetaNet and described in Section
seaieeet-eeestt---rdieetetee updates. Ur. To[fm] - Tolfoul Are attention variants the Lr, layers described in Section 5 For given g [il in order to compute Ur, (g) we first compute set-wise attention, given by:
(Mag
M ( Vin
Attention (g) (1)
softmax
(1)
(36)
Attention, (g) (g))) Mo Mo[Attentiont (g)
Lgj: Attention (g)
(37)
Here (.)) denotes inner product and M& Mb. MB E R/axf are learnable matrices used in each attention head and Mo H foul is final aggregation linear layer. We then compute gradient-wise attention, given by:
de
<Mig MR (...)
Attention h (g)
softmax
Mig.2
M
(38)
VIn
1' 0 Rm
Attention, (g)(0 (g)) Mo [Attention (3) "............................".. Attention" (g)
(g)
(1)
(39)
i.j.
Here we slightly abuse notation and denote by M .M Mb fim Mo E f for learnable finally, the value of Ur, (g) given by
matrices different from those in equations 36 and 37
Ur. (g) 9
gLP 1. Attentions (g) (1)
i.j.: + Attention, (g) ().))
(40)
22

---

Gradient-batche batch-to-weight update. As GradMetaNett+ prioritizes empirical improvements over computational efficiency, we directly use muadient-batch-to-weee update, which we found to yield better performance.
The seadienn-to-weighee mapping. Urrad Ts[fin] @, applies pointwise MLP to the feature vectors associated with the neurons connected to each weight in every element of the batch. In the case of biases, the MLP applied the feature vectors corresponding the respective neuron. Finally, the results are summed across the batch. Formally, this can be expressed by (g) (W1,ba ..... WL, by) where:
(Wi)ky Σ MLP 1=1 kA
(4 9 kkga
2
(41)
D.3 Invariant GradMetaNet.
In some cases (e.g. evaluating influence functions) want GradMetaNet to output single invariant vector Rows rather than a parameter vector E @ In this case we replace the Lprod component described in Section with a Lee layer described below, For an element g [[fi].
L di
[eee(g) = M
(42)
Where M E Ru is learnable matrix. This results in invariant vector outputs.
D.4 Computational Complexity.
We analyze the space and runtime complexity of GradMetaNet and GradMetaNett++ comparing them to alternative approaches. Throughout this discussion, we denote by the number of parameters in the underlying MLP fe. whose gradients are being processed, and denote the number of neurons by N. Since the gradient-batch-te taab--tt-grddent---aae update is the most computationally intensive component in both architectures, we focus our analysis on this operation.
GradMetaNet. The gradient- batchtt-tttttttttt toe--tt-----aee wer-grrrdieentttee update in GradMetaNet consists of a stack of layers Ir, : To[Si]] + To [ou]] as defined in Section 5 Each of these layers has both space and runtime complexity of O(N b. fin fom). Thus, assuming fixed hidden dimension and number of layers, GradMetaNet has complexity of O(N.b).
Gradient concatenation and averaging in weight space. Both gradient concatenation and averag- ing methods process sets gradients by utilizing weight-space architectures, such as those introduced by Lim et al. [88), Navon et al. [61]. These architectures employ layers L. θ fin] + [[fuu] with time and space complexity ((P fin foa).
In the concatenation approach, gradients are concatenated, producing an input element in @ resulting in an overall complexity of O(P- b) for a fixed hidden dimension and number of layers. This approach scales poorly compared to GradMetaNet when b- >> b- N. In contrast, the averaging approach reduces gradients a single representation in @ yielding a complexity of O(P). However, this method is also suboptimal in the overparameterized regime (P > b- NN. In addition. averaging
even in cases where the batch size is sufficiently large such that < b- N, methods may still be suboptimal due to their expressivity limitations (see Section
GraaNeteNett+. The gradient- batttt-bahhttte to- gradient-batch update in GradMetaNett+ is implemented using a stack of layers Ur, Tolfml + To foul) as detailed in Appendix D.2] Each layer has a time complexity of O((N2. of O(N b. fin Jrat). While these layers have the highest time complexity among the approaches considered so far, their space complexity remains efficient. Moreover, constructing an attention variant for weight-space architectures would scale quadratically with P, making far less practical in terms of scalability.
b2. fm fout) and can be designed to achieve a space complexity
23

---

E Theory
In this section, we provide proofs and further discussion for the results presented in Section 6 Throughout this section, V, to denote gradients of the networks computed on single datapoint.
E.1 Importance of Processing Collections of Gradients
In Section 6 we discuss the expressivity limitations of processing the gradient of the average loss on the batch compared the collection of gradients at each of the datapoints. this section, we formalize these limitations. We start with some notation and definitions. As we saw in Appendix the FIM can be computed using gradients on individual datapoints. Given a of such gradients I....., Vo), the FIM computed using the gradients in & is denoted by Fg. As we saw in the main text, G can be thought of an element of Ob.
A.3
O x O + be a function whose inputs are parameters and set of gradients G= (Vi...., V.J. We say that & non-trivially depends on the FIM if for some function
Definition E.1. Let +
x x
(0,9) v
FG, 0, Σ V.
(43)
==1
and there exists pair of inputs g= (VI..... Vo} and 0¹, G' [V ..... Vi
where G and g
are admissible gradient sets and such that = 0', Ei V -
V but
+ (6. L Fg. 0,
F L Fg O'
V
+ (0', g')
(44)
Many commonly used functions over sets of gradients non-trivially depend on the FIM. Before providing such examples, we first state the following trivial proposition.
Propusition E.2. Let 1 x θ + be a function that non- non-trivially depends on the FIM. There exist an 0 such that for any continuous function A O x @
@ it holds that:
max 8.9
(0.9) 0,
Σ V,
(45)
In other words, functions that non-trivially depend on the FIM cannot be approximated (in the -se-ee by continuous functions that rely only the average gradient.
Proof. The proof follows trivially from Definition
Let 0, g = =[ Iy.... Vo} and 0', g'
(V..... V be pair of inputs such that 0 = 0' and
Σ5
V. = +t > ill V. but
( + + (0',g")
(46)
For any A θ θ
θ
have
M v. = A @o: Σ v
A
0
(47)
b
Therefore, if we choose
<<
<
2/
(0 - 1 gl'll we have
0. 1 M V
[10"'"00', A 0' 1 Σ v.
| 9) A
2 (00 100--(0''000 (')1 2c.
(48)
This implies that either (0,0))
A A -
or (00000) (..+E and
max ( A 0 V.. 1I 0.0
(49)
completing the proof.
'Here, by "admissible", we mean that the elements of & and G' are actual MLP gradients, rather than arbitrary elements of e.
24

---

We want to show that the computation of the natural gradient and the OBD/OBS pruning saliency scores con-trivially depends on the FIM. To do so, we first formally define these computations as functions over e-
Definition E.3 (Natural gradient map). The natural gradient map 4 nat 4 b -> θ is defined by
nat (V.)) (Fg + +
(50)
where I is the identity matrix and > 0 is damping factor. These are added since, while positive- definite, the FIM is not guaranteed to be invertible.
Definition E.4 (OBD/OBS pruning saliency maps). The OBD saliency map POBD O O + θ is defined by
POBD (), == 0² o (igg The OBS saliency map 4 oBs + O - θ is defined by
(51)
oBs Po 0 diag (Fg + +cc)
(52)
We now show that both the natural gradient map and the OBD/OBS pruning saliency maps non- trivially depend on the FIM.
Proposition E.S. The maps Peat PORD, and Pons non-trivially depend on the FIM.
Proof. To start, we assume that fo is a single-layer MIP,, i.e., a linear map from the input space R" to the output space IR. The proof can be extended to deeper MLPs by composing the linear map with an MLP that implements the identity function. Given a batch datapoints D = {[ y1), (EN, yN), R" x R, the gradients of the output are
V,
(efe((;) = Cz..
(53)
Thus, as discussed in Appendix
A.3 the FIM on G = 1V....., Vn) can be computed as
F6
x.a.
(54)
i=1
We begin by showing that Intt non-trivially depends on the FIM. This is equivalent to showing
that there exist two choices B {(111 ( ..., (Tb; yo)), B' {(=1, y ) (1),.... (2// v()) with
corresponding gradients G = (V...... vo). g' = (V for some gradient V.
V such that E V,
V but,
(Fg + c + (Fg + +
(55)
We construct such B and B3, but emphasize that this is only one of many possible ways to construct such an example. First, take D,, {(21, y1),..., (In, yn)) such that {x; } is an orthonormal basis, meaning a. 2z = 8, This means that FIM is the identity
Fg Σ
(56)
Thus,
(Fg + cI) = diag((1 + ....., (1+
(57)
Now, define D', = {(1),1)) (z/n, v'h) Ri such that x² 2x ry1- = 22 The FIM in this case
Fø. = 41
(58)
and
(F + - -dig(( +
(59)
Therefore, for any non-zero gradient we have
& naat (V,9) (Fg + +c) (1 + + (4+ + = + nat (V,g') (60)
This proves nat non-trivially depends on the FIM. To see that FOBD and POBS non-trivially depends on the FIM, take D,, and D,, as before, and choose any non-zero 0.
The next proposition now follows from Propositions E.2 and E.5
25

---

Proposition E.6. Laat. POBD. and OBs cannot be approximated (in the &oo sense) by continuous functions that rely only on the average gradient.
E.2 Universal Approximation Results
In the discussion below, we are concerned with functions from a compact input domain KC C Tolf]] such that K 0, where
& :: L-1 -
---00
Ú
ge
Σ
(61)
1.
I=l in
j=1
j=1
is finite union of linear spaces of co-dimension f. Similar assumptions over gk were used for the universality proofs in Finkelshtein et al. (23]. Maron et al. [53].
For the readers convinience, we recall that for MLP with input dimension do. output dimension dr hidden dimensions dı,... ,d,, ,d,,1, we define = Sd, x x... Sd, 1 and G, S x G. acts naturally on the spaces @ and F[ʃ], while G2 has natural actions on the space @x[[] and [o[s]-
See Appendix of definitions. These actions preserve the inherent symmetries of the spaces they are defined over, and so we aim to respect them through equivariance/mvvarincce
Main universality proofs.
Theorem E.7. Let K C Tblf] be compact domain such that K. = UyEG,9 K. and KNε H. GradMetaNet models are universal approximators (in || lloc sense) of contimuous Gu-equivariant functions from K. to weight space 6.
Proof. Let K. - θ be a continuous G₁ equivaraint function. From proposition E..00 there exists continuous G3 equivariant function T IC [[f"] and G equivariant function A (CC) - O such that 4 = A o V. From proposition E. there exist a stack of layers Lprod o Lr o...o Lr P which can approximate A over (KC) to any precision. Additionally, the function PE y is also continuous and equivariant and so, from proposition from proposition E.15 exist a stack of layers LPool Lr, .... Lr, PE which can approximate PE L over K. any precision. Composing the two components together allows us to construct a GradMetaNet model Lprod Lr o ... Lr LPool Lr, o Lr.- o PE which can approximate
to any precision.
As result of Theorem
E.7 we obtain the following formal statement of Corollary
6.3
Corollary E.8. Let KC [olf] be a compact domain such that K. = UrEG,T K and Knn 0. there exist GradMetaNet models can approximate the natural gradients (see Definition E.3 of elements of K. to arbitrary precision. Additionally, by incorporating the parameters 0 of the MLP whose gradients are provided input GradMetaNet into the gradient-to-weise update, GradMetaNet models can approximate pruning saliency scores (see Definition E.4) with arbitrary precision.
A.3 the natural gradients can be expressed as function from decomposed gradient space r3 [3] to parameter space @ This function is both continuous and equiv- ariant, and thus Theorem E.7 shows GradMetaNet can approximate natural gradients. Additionally, the functions 4 (g) diag F)) and 22 = 1/diag( [F + +ef+ - are continuous equivariant functions from Fs to θ and thus can be approximated using GradMetaNet models. Recall that the
Proof. As was discussed in Section
OBD and OBD pruning saliency scores are computed by (g(g) C 0 and (g) c 0 respectively.
The parameters 0 (W1.bi,... component Lorod (See Section 5) the following way: Lprod T[fil + + @ applies a pointwise MI.P the feature vectors associated with the neurons connected each weight along with the weight of the original MIP, or in the case of biases, to the feature vectors corresponding to the respective neuron. Le., Irrod [g O) (V1, CL)CL) where
O, can be naturally added to the gradient-tor teoiighttto-weight
MLP1
(1+1) j.:
( = MLP ((g.. 1.: ((")).
(62)
of
As we established GradMetaNet is able to approximate the functions 4 L2 the update in Equation 62 can easily approximate OORD- = 0 c d and OOBS = 0 o 22. This completes the proof. O
'Here PE is defined for r's' = r. [. we thus abuse notation writing PE without indicating which space it operates on.
26

---

Finally, we include proof of universality for the invariant case
Proposition E.9. Let K. C Tolf] be compact domain such that K. = UTEG, T K. and K E = 0. Invariant GradMetaNet models are universal approximators (in || lloe sense) of continuous Gb- invariant functions from to gd.
Proof. Let K Rd be continuous G3 invariant function. From proposition E.15
the gradient-
bag-to-gradient component of GradMetaNet is universal approximator continuous equivariant functions from to P[d]. We can extend to be an equivariant function - [[d] defined by
[ (g)
(63)
Since c is continuous and invariant, continuous and equivariant and can thus be approximated by the gradient-bag- to- gradient component of our method. Finally applying the gradient- gradient-to-vector pooling layer Lve we get that our model can apprximate the function
d
$ g
(g)g.
(64)
do + ... + dp 1-0
This completes the proof.
We now prove all the lemmas and propositions used in the above discussion.
Proof of proposition E.10.
Proposition E.10. Let K. C [o(f) be compact domain such that K. = UTEG,T K. and K. & = and let K - A continuous Gs-equivarian function (here the Ss component of acts on @ trivially). There exists pair of continuous functions T K [( A: (KK) - θ such that:
T is continuous and G Gr-equivariant. continuous and G-equivariant. V A.
Proof. First, Lemma E. states that there exists a continuous and Gg-equivariant function V g
T[f"] (where& 8 (g E To[f] g & E)), such that for every 9 9 E g
(())) (g,) < Fr E Sb s.t. 9 = T -g-
(655)
Now let To[/] + / be the projection map to the quotient space induced by the orbits of Sg. Note that the group G acts naturally on the quotient space / by:
to g E S.) = (T g e
(66)
Additionally, as the is compact, the K = x xKC) also compact. Since + are invariant to the action of Sy and equivariant to G, Lemma E. implies that there exist continuous G-equivariant functions
K + [[f'], Φ K + O such that v V and & = & As L is Se-injective, the function V is injective and thus the function - (K() - K is well defined. Additionally, as V is G equivariant - is also G' equivariant. We define A $ 11-1, v (XC + e. Since Φ and T are G-equivariant, A is G-equivariant. Additionally.
A L
o T き woroworooozz
(67)
Finally, Lemma E.14 A is continuous, completing the proof.
We now state and prove all the lemmas used in the proof of proposition E.10 starting with Lemma 3 from 53 restated below:
Lemma E.11. Let < S,, act on / by applying the same element r E H to each channel, then there exists a polynomial function U R - RJ" for some I' E B for which U(x) = U(y) if and only if x y for some H.
27

---

Lemma E.12. There exists continuous Gu-eeuuivarianttunnnii v : E - [[ such that for every 9 92 g
L (g,) (g)
+ 3r St s.t. 91 = 92.
(68)
Proof. Let U be the polynomial invariant function established in Lemma E.1 where = Gb, and define S g [[ by
S(g)
9 9
(69)
We now define V : + [[ by:
vggg L (s (g):.: (g)
(70)
where [...] represents concatenation along the feature dimension. Note that we slightly abuse the notation the feature dimension, denoting it as f' multiple times. We first notice that since is equivariant and continuous and U is invariant and continuous, L is also equivariant and continuous. Now, for input vectors g 9 E 8 if there exists group element T E Sy such that g = T 92 then (g2) L (g2) where the last inequality holds as Sy acts trivially on the output space [f f- On the other hand, assume Li (g,) V (g.). We consider 2 cases:
from equivariance we have L (g))
(T 92)
92 * 91. this case from the definition of u it holds that (g) & (g2) and so & (g)g o (g)))
Case I: for every TI, E G, S, G it holds that Ti 72
Case 2: There exist pair Tı, T2 E Gb S, G such that T1 T2 g2 9 Assume by contradiction
that 72 is not the identity. Recall that for every g $ E, E [L]. of j E [di]. it holds that
(71)
Thus,
S(g)) # 72 -g))) S( g S(n T2 91) = 5(g)).
(72)
This implies that (g)) * (gg)). We have thus shown that (g = L (g)) Sp s.t. 91 = T T 92 completing the proof.
= 3r e
Lemma E.13. Let K. C Tu[S] be a compact set such that UgeGug Furthermore, let & K + I[f'] be continuous Gy -e-eqirrrra functiiooniincion Finally, let π -bfff denote the projection map the quotient space induced by the orbits of Sa and define K. T(C). The following holds:
= acts on Toff//So by r ( g o
is G equivariant.
There exists continuous G-equivariant &b K + r[l'] such that tb +
Proof. Recall that each element in Tolf//S, is of the form (g) = for (oogle g o E St}. The first statement trivial as
e fo E 5t} fo.gloae St}
(73)
and for every 71, T2 we have
(TI T2) for g a E St} = - . a E Sr} = (7 (52 fo o St)). To prove the second statement, notice that for every G. E Sc, we haver o - Thus
(74)
[ g) {a (o---- 7 E S,) = (T g a e -
(75)
Finally, prove the last statement we recall that since & is continuous and S₁ invariant. from the definition of the projection map π there exists a continuous function + K r[f"] such that
$ o æ. To show that & is G equivariant, we notice that for every g E R there exists g K such that xt (g)) g. thus for very 7 E G
g) Φ(T 1o.-(g)) x())
@(<(T ())))))) Φ g) = (gg) ( ((g) (76)
28

---

and so @ is equivariant, completing the proof.
Lemma E.14. Let K C. Ri be compact domain and &b K. + Rf' be continuous function such that - V. V is continuous, then continuous on V (K.).
The proof of this lemma identical that of Lemma 5 from 133 still we provide proof below.
Proof. Assume that this is incorrect, then there is sequence yi = (a,) such that y: - Yo but (y.) A (goo). Without loss of generality, assume that æ; - zo K. (otherwise choose converging subsequence). We have
((x;) = A (a.)) (y)) A A(yo) (z ((xo))
(77) L
which is contradiction to the continuity of db
Proof of proposition E.15
Proposition E.15. C Tbf compact domain such that = UgEGa9 K. and Knn and let o K. + [[ be continuous Gs-equivariant function. For every 0 There exists stack of layers = LPool o Lr, o.... Lr, PE (As defined in Section 5) H FGruaMetaNee such that for every
K:
(g) F(g))l <E
(78)
Proof. From Lemma E.16 there exists a continuous Sp X -equiviadde function V Tols + kj + r[ such that = PE. As V continuous and S, x S4g was shown e.g. in [23 53 that for each > 0 there exists a Deep Symmetric Sets network for sets of sets of the form F Lrool Lr, .... Lr, such that for every E PF((K):
-equivariant, -
(g) F(g)ll <E
(79)
This implies for every E I::
|| (p(g) F(g) = (PE()) F(PE(g)) <
(80)
Completing the proof.
Lemma E.16. Let t
be a continuous G₁ equivariant function and let PE Tolf] To[s k) be a positional encoding layer as defined in Section 5 there exists an Sn Sdo + +dy equivariant function V [ + k) - [[J'] such that
PF.
Proof. As PE is injective we can define A PE(T + r[S] by A(g) = (PE (g)). We Pf(T by defining for every g PE(T.(f))
now extend A to the domain LoESatt...4 Agg). We show this extension is well defined (i.e., that there is case where 01 g 02 g but σ1 A(g) * #02- Agg)). Let g' E To[f].g = PE(g') and
Sdo +do A(a g)
T1 * 02 E +dd such that 0º g = 02 g. follows that 02 01 g = 9. and thus from the definition of the positional encoding map PF, 0² di E G and 02 01 g' = -g'. Thus, since +b is G-equivariant it holds that
A(c, g) = -g( = =0202 . ( 0g $ ( 01 g') 02 $ (g') = 02 Ag) A(a2 g)
(81)
Thus, the extension of A is well defined. Since the functions Φ, PE are continuous the function A is continuous and thus there exists a continuous function A [s[[]] + [[J"] such that for every UbE PE(Ts(]] it holds that Å(g) A(g). We define & To[f] by
V (g)
Σ -'.o - - g).
(82)
rES 1 11g TESp
First, & is continuous and equivariant w.r.t So x Sd dd..
Second, for every PE(Ts[f)) it
holds that
V (g)
Σ
-...o-..-0- A(g). (83)
SEg TES
TES
Thus, for every g' E [o[f] it holds that (ggg) (PF (PPE))(PF(g)) (((F((')) and so V PE.
29

---

Proof of proposition E.18 In this section, we aim to leverage the universality result of the Set2Graph architecture presented in Serviansky et al. [76] complete the universality proof of GradMetaNet. To this aim, we first define the square gradient space [[.]. Intuitively, the spaces [[ and [[]] parallel set space and graph space, and [ can be embedded in r[]]. We now formally define [[.].
Definition E.17. Square gradient space [[j] is defined by
F[f] = (T & rY.
(84)
Here, @ denotes the tensor product, while represents the Cartesian power product. Thus, an element
ğ E [/ is of the form
(1.122 {g.j. ()
(85)
where i [dr j E [di,] and g by the map Lode defined below for 0
(21.22) i.
E R/. The space [[ can be naturally embedded into [2.]]
E
[[S] 0 = (W(1) ((1)
W(L)
W 1.
(a)) 4-
Lode
(0)
(12.22 Aie
(86)
otherwise.
Additionally, the map Lodp
[[2/] +
@ projects the space [[2 to @ and is defined by
Lodp(8 (g) = (W(1), (()) ( (B())
w ij.:
(1111) Lj:f
b (1)
- -g
(87)
The action of the group G Sd, x... San_, on [f]] extend naturally to [[f] and is defined for T = (T.....,T-.1) E G G.g [[] by
(11)2) bigh
(2 -9 T (0:mm)) '))
(88)
It is easy to verify that the maps Lode: Lodp G-equivariant.
Before stating the following proposition, we note that the positional encoding map PE [o[s] - To[ + k] can be considered as well defined on the space ro[s]= [.]]
Proposition E.18. Let K. C CJJ be a compact set such that K. = Ugggg K. and let Φ : K. + O be continuous G-equivariant function. For every e O There exists stack of gradient-to- gradient layers composed with a gradient-to-weije Lprod o Lr ....o Lr PE such that for every g K:
layer and a positional encoding layer
[ - F(g)) <
(89)
Proof. First, notice that the model F equal to Lodp o F o PE where F is exactly Set2Graph architecture, from the space [[/] to [[["] which equivariant the action of the group t on both spaces. This is because the DeepSet component in Set2Graph is identical to a stack of Lr layers and the broadcast and pointwise MLP components of Set2Graph are identical to the Lprod component composed with the embedding Lode- From Lemma E.19 there exists continuous Sdot...ddz equivariant function & f + k] - [ such that = Lodp o V o PE. Finally, 766 that there exists a Se12Graph network F such that for every g E PE(K):
was shown
(g) Fg)) <.
(90)
This implies for every g E K there exists a gradient to weight model F such that
(g) F(g)]l [ (V (PE(g))
Lode (F(PE(g))) <
(91)
This completes the proof.
Lemma E. Let + r[r] - O[j') be a continuous Gb, equivariant function. There exists an S. Sdot..dd equivariant function & T[f + kj + F[2f"] such that & Lodp PE.
Proof. The proof is very similar to that of Lemma E. Let & : [[ + T|2 be defined by &- = Lode & and note that since & and Lode are continuous and G equivariant, + is also continuous and G equivariant. As PE is injective we can define A PE(T[/] + r(2/'] by PE( by defining for
A(g) @<PE-1 (g)). We now extend A to the domain UoES
30

---

every g E PE(T(f)), E +dz+ A(7 g) a Ag). We show this extension is well defined (i.e. that there is no case where 01 g = d2 g but 01 g * 02 Ag)).
Let g' E risg = PF(g') and Ø1 # 02 E S such that 01 g 02 g. It follows that
2 01 g = g. and thus from the definition of the positional encoding function PE, o- on E G
and 2 o, g' = g'. Thus, since + is G equivariant it holds that
A g) ) - 01 Agg = J2 01 (g') o2 $ ( TT g") = 02 $ (g') - 02 ( A A (.
(92)
Thus, the extension of is well defined. Since the functions Tb PP are continuous the function
continuous and thus there exsists a continuous function A [[f] - r[2 such that for every (/]( it holds that (g) = Ag)). We define v [[ + F[2/] by
E Une Lo
V (g) =
Σ
T 11. o - o g).
(93)
FESKgtn+dy TESs
First, & is continuous and equivariant w.r.t Sb x Sdo t holds that
Second, for every g E PF(T[/])
(g)
-1..o o -1 O g)
Σ
-1. .. A(g) = A(g). (94)
of OESky+..4hy TES,
of oES
TESp
Thus, for every g' E T(f) it holds that
(g'g) = $PE (PE(g)) = (PE(g))
(95)
Finally, since Ldd Lode id we have
Lndp + = Lood V a PE Φ-
(96)
completing the proof.
E.3
Importance of & in Universality
Recall that Theorem 6.2 proves the universality of GradMetaNet over compact domains which do intersect the set & [o[ defined by:
ع 3 8UUU
U
ge ger
Σ !1) N. Σ ba.di
9
(1)
(97)
(=1 in
In this section, we prove that there exist compact sets I.C. C Tolf] with Knn E * 0 over which GradMetaNet is not universal. We emphasize that, as discussed below, sets intersecting & consist of highly regular gradient values, which deviate significantly from typical gradient sets. Consequently, the requirement K. N = a mild and realistic assumption.
Proposition E.20. There exist pair of elements g 92 E such that for any 5 E Gb. T. gg * 92 but for every GradMetaNet model F it holds that F F(
Proof. For simplicity, we focus in this proof on the invariant version of GradMetaNet. F I,[ + R. The proof can easily extended to the equivariant version of GradMetaNet,
To[f] O[f"].
Let the underlying MLP used to define the spaces [/s have an input dimension dg = 1, single hidden layer of dimension di = n, and an output dimension d2 = In this case, T, [] along with the action of Gb. is isomorphic the space of "sets of sets" defined in Hartford et al. [31]. Moreover, GradMetaNet is equivalent to the architecture proposed in Hartford et al. (3I). This architecture is known to be unable distinguish between certain highly regular, non-equivalent inputs. For example, incidence matrices of graphs can be viewed as sets of sets and, therefore, as elements of [o[f]. It was shown in Albooych al. [2] that the ability of the architecture in Hartford et al. (II, and consequently GradMetaNet, to separate graphs based on their incidence matrices is equivalent to the 1-WL test [89].
Consider an example where g represents the incidence matrix of a graph consisting of two discon- nected cycles of length 3, and g, represents the incidence matrix of a single cycle of length 6. For any GradMetaNet model F, will hold that F(g,) F(g2). A straightforward check shows that
for this example, g 92 E E, completing the proof.
D
31

---

F Experimental Details
In this section we provide all experimental details for all experiments Section 7 We run all the experiments on singel NVIDIA-AI00 1OJA-AA00-SXME GPU with 40GB of memory.
E.1 Curvature Information Estimation
Dataset. To construct the dataset for this experiment, we first generate 3000 SIREN models (78) commonly used for INRs. Each model has three layers with 32 hidden features, i.e. - 32 1. The weights and biases of these models are randomly initialized using standard Gaussian distribution. Each data point the resulting dataset corresponds to a single SIREN model and consists of input gradient set of size 128, g E I 122s and target vector, FFM E e. The input gradient set corresponding to a SIREN model fe is computed by sampling 128 random points Sfo = (x1,... ................................ 1128) C [-1, 1J. To simulate diverse datasets, the points in Sys are sampled from distribution X, defined by
[U([O, 1)), with probability
X,
(98)
[(([-1,)),
with probability 1
Here, the value of p randomly and uniformly selected over the unit interval. The set of input gradients is then given by , V123 V. ffezz;). Thus the input vector 122 is given by I,...., V22s].
To compute the target vector 8PIM E + we first randomly sample 1024 random points from X₂ S% fa'..... 2o244 c [-1,1 (Note that the sets So and S% are independent). We then
compute the true FIM (see Section A.3) by
1024
Vfolm)Voool) 000
(99)
1024
The target vector FPIM θ is then given by 0PIM diag The task is then to predict FFIM based on the gradient set vector g and is trained using l2 loss. For baselines which process data in parameter space [[]] (rather than gradient set space I 12s [2]), the gradient vector g computed from gradient set } is replaced by representing each gradient in parameter space V, O and then either concatenating them resulting in vector O + { or averaging them, resulting in vector in e O
Baselines.
We compare GradMetaNet and GradMetaNett++ against multiple baselines.
MIP concat: This baseline takes the gradient vector g described above as input, flattens and feeds it into a standard MLP.
Neuron Asymmetric GradMetaNNt: This variant respects the batch symmetries but not the gradient space symmetries of Ty(f). It uses g as input and applies a Deep Sets architecture, as described in Zaheer et al. (9I) Specifically, the vector feature for element in the set is constructed as (4))
[g
Batch Asymmetric GradMetaNet: This variant closely resembles GradMetaNett but respects the of I' [/] It takes g as input and applies GradMetaNet model defined over I [256), where the batch axis is "flattened", resulting in a single gradient with a feature dimension of 256 2 128.
gradient space symmetries while disregarding the batch symmetries
Weight Space Models:
The variants "DWS+Concat", "DWS+Average" "GMN+Concat", and "GMN+average" take as input the gradients represented as 0 E O[f] Here, 0 is either the average of the gradients in the set, in which case f 1, the concatenation of all gradients in the set, in which case f = 128. These variants then apply either the deep weight space (DWS) architecture described in Navon al. [62] or the graph metanetwork (GMN) architecture introduced in Lim et al. 48 to the corresponding input vector.
Standard Estimator:
Finally, highlight the benefits of learning approximate algorithms rather than relying on fixed, non-learnable approximations, we compare GradMetaNet and GradMetaNet++ to standard estimation of the target. Specifically, compute diagFs, based on the 128 gradients
32

---

Table 5: Comparison of baseline properties.
Baseline
Supports sets
Supports efficient
Si-invariant
G-equivariant
gradients
gradient representation
Concat
+ Concat
GMN Concat
DWS Average
GMN Average
Batch Asymmetric GrudMetaNNe
Neuron Asymmetric GradMetaNet
X
GradMetaNNe
GradMetaNet++
provided as input and evaluate its & distance to the target, which, reminder, is diag FB,, Here, B2 is computed using the gradients of larger set of 1024 points, independently generated from the data points in Bi.
Data preparation. We use 500 examples test dataset and 500 examples validation dataset, with the size of the training set varying between and 2000 examples. As preprocessing step, all data, including the target vectors, is normalized based on the statistics of the training dataset. The following are the three normalization methods used, based on the vector space which the data belongs:
1. For input vectors of the form g E I [22s for each layer 0,. ...,., 3, we compute the mean y()) and standard deviation ((1) over the set of values:
{g j.j.k
O, 0,..., b; I,.... de, k = 0, g E D
We then normalize the data by: (()) -0)
For input vectors of the form E θ[f] = (W(1),bl W b(L)) (where f = for averaging or f 128 for concatenation): we follow the normalization scheme suggested in
Navon al. 61 where for each layer 1.... 3, we compute the means / (1)0) and standard
deviations aw % over the sets of values:
( (1)
(w(() 1i 1,...,da d1-15 j = 1,...,d;; di; = O,..., w E Drrain)-
(b 1,6 ii=1 dr; k = J. (()) E }
We then normalize the data as follows:
( W
(
(()
Pb (
pw
(1)
(1)
For target vectors FFIM E O: recall that = R. for some integer P E Let (1) denote the i-th entry of the vector FFIM E R.P. We compute single mean μr and a single standard deviation or over the set values: {0FIM (i) l....,, 0FIM E Dirsin). The data is then normalized as: 0FIM + 0FIм-#
Additional experimental details. Following the experimental setup Navon al. [6]. all learned baselines in this experiment consist of approximately 15K learned parameters and roughly 3 layers (In some cases, where the input dimensionality was extremely large, it the 15K parameter budget with 3 layers). All models were trained for 100 epochs using the Adam optimizer with learning rate to the best validation performance and repeat the experiment with random seeds 1 through 5.
possible to maintain
10 and batch size of 32. We report the test loss corresponding
F.2 Learned Optimizers
Learned optimization tasks. The following are descriptions of the optimization tasks the learned optimizers were evaluated on. Across all tasks, the training loss is negative lieeliiel,hood, and the batch size is 128 except for the transformers task.
33

---

Table 6: Multiplicative improvement in the number of steps needed to reach Adam 's best test loss, as well as the best test loss achieved by each standard and GradMetaNeet-assed learned optimizer. For each task, we Adam, record its best test loss L and the number of steps N required to reach it. We then run each learned optimizer, measure how many steps it takes to reach L, and report the improvement relative to N. The "Step Reduction Factor vs.Adam" column thus indicates multiplicative speedup in optimization steps. The "Best Test NLL" column record the best test loss achieved by the optimizer in the run. Results are averaged across optimization runs.
Dataset
Optimizer
Step reduction
Best
factor Adam (1)
test NLL (L)
Adam
1.00 + 0.00
0.475 + 0.011
SGDM
1.13 0.17x
0.463 0.019
DS
1.27 = 0.36x
0,460 + 0.011
F-MNIST
UNF
1.16 0.18x
0.451 0.015
DS + GradMetaNet
1.44 + 0.43x
0.445 4:55 0.017
UNF + GradMetaNet
1.51 = 0.59x
Adam
1.00 = 0.00x
+.039
SGDM
1.41 + 0.52x
1.556
DS
2.32 + 0.42x
+.094
CIFARIO
UNF
2.64 +: 0.53x
1.516 0.020
DS GradMetaNet
4.63 = 1.11x
1.427
UNF GradMetaNet
4.26 + 0.56x
1.418 + 0.022
Adam
1.00 = 0.00x
1149:0000
SGDM
1.06 + 0.07x
3.424 + 0.026
DS
1.79 + 0.40x
3.356 + 0.027
CIFAR100
UNF
1.58 + 0.27x
3.345 0.036
DS + GradMetaNet
3.15 : 0.56x
3.253 + 0.019
UNF GradMetaNet
2.85 0.38x
3.262 + 0.013
Adam
1.00 + 0.00x
6.904 0.075
SGDM
1.01 = 0.03x
7.045 0.031
DS
0.88 - 0.14x
6.987 0.108
LMIB
UNF
1.48 + 0.16x
6.702 + 0.080
DS GradMetaNet
1.09 = 0.16x
6.993 + 0.076
UNF + GradMetaNet
1.82 + 0.39x
6.557 + 0.061
(1) MLP on FashionMNIST. Learned optimizers are tasked with training a three-layer MLP classifier downsized (8 8) flattened version of FashionMNIST 90) The MLP has single hidden layer of dimension 32 and ReLU activations.
(2) MLP on CIFAR10. Learned optimizers are tasked with training three-layer MLP classifier on downsized (8 x 8) flattened version of CIFARIO. The MLP has a single hidden layer of dimension 32 and ReLU activations.
(3) MLP on CIFAR 1 00. Learned optimizers are tasked with training a three-layer MLP classifier on downsized (8 8) flattened version of CIFARI00. The MLP has single hidden layer of dimension 128 and ReLU activations.
(4) Transformer on LMIB. Learned optimizers are tasked with training a transformer language model on LMIB 14, using next token prediction. The transformer comprises two blocks with an embedding dimension of and uses four self-attention heads. We train with batch size of length 8 sequences.
(5) 2-parameter linear regression. Learned optimizers are tasked with training a linear regression model with two inputs and single output. The input training data is mixture of two Gaussians centered around (1, 2) and (2, 1) and the target always 0. This results in a loss landscape with non-standard curvature, depicted in Figure
Learned optimizer architectural details. In each inner training iteration, all learned opti- mizers are provided with the current parameters 0,, current gradient V., six momentum value (nf.1 vp.55 , n0.9 v0.99 0g. 0.999 1o (.99999 iteration number / as an 11 11-dimensional sinusoidal encod- ing. All of these inputs are concatenated across the feature dimension to get 19-dimensional vector per-parameter. Le., the learned optimizers inputs are in + [9] and outputs are in @[1].
34

---

For GradMetaNet based learned optimizers. inputs also include current set of individual gradients g E F4 and six momentum value {vg.) vg.s. vg. vQ.99 vQ.909, vQ.9900] concatenated across the feature dimension, resulting in an input in Po[7]. These inputs are processed by GradMetaNet, which outputs an embedding in [14], i.e., a 14-dimensional emhedding vector per parameter. This embedding is concatenated to the other inputs of the learned optimizer, so the DeepSets/UNF based learned optimizer gets an input in θ 933].
The DeepSets 9 based learned optimizers process the inputs by applying a per-parameter MLP with three hidden layers of size 32. The UNF [92] based learned optimizers apply a paraameeeeer MLP with two hidden layers of size 32 and output dimensions of 32, followed by a single UNF layer. As baseline compared against all learned optimizer,s we used a standard Adam optimizer [40] with learning rate tuned by grid-search over the values {0.0005 ý [=== (i.e., 20 equidistant values in the range [0.0005, 0.01)).
Meta-training details. We meta-train for 50,000 steps using Adam 440 with learning rate 10 10-4 estimating meta-gradients over 16 parallel training runs using persistent evolutionary strategies (PES) 85 with a truncation length of 50. The meta-training objective training loss at the end of the inner training horizon T, which is T 2, 000 for image classification tasks, T == 5, 000 for the transformer language modeling task, and T = 10 for the 2D linear regression experiment. For all methods, we initialize
0.1, μ 0.9 and B 0.001 before meta-training. We use the learned optimizer meta- training setup from Metz et al. [7 (project available at /https://g// bttps://githbb com/ released under Apache License 2.0).
F.3 INR Editing
Dataset. We utilized two previously proposed INR datasets: MNIST INRs [61] and CIFARIO INRs [93] Each INR represents an image from the original image datasets and consists of three layers with 32 hidden features. The target images are produced using the image processing library OpenCV [0). by dilating or increasing the contrast of the MNIST and CIFAR- 10 images, respectively. To construct the input gradient set, g E Tea, we sample 64 random input coordinates in [0, 1]2 and compute the gradients of the editing loss w.r.t the original INR parameters. Specifically, if are the INR parameters and [0, 1|2 - R$ is the target image, then for an input coordinates (xx, y:), the gradient V, is given by V. = Vo (fo coordinates for all INRs.
((.,,):)) We use the same set of random
Combining GradMetaNet and weight-space methods. We combine GradMetaNet and the weight-space architectures as follows: First, GradMetaNet processes the gradients g E Ta4 to produce outputs in O[f]. In this experiment, choose f = 8. Next, the output is used additional weight features for the weight-space network, i.e., the input to the weight-space network is in O
Additional experimental details. We train all methods for 150K steps using the Adam (50] opti-
mizer with batch size of 64. We search over learning rates in (0.01, 0.005, 0.0001) }. For ScaleGMN and GMN we use the official implementation provided Kalogeropoulos et al. (38) with the same parameter configuration provided by the authors. We use the bidirectional variant (ScaleGMN-B, Kalogeropoulos al. (3 with 10 layers and a hidden dimension of 128. The DWS 6 network consists of 8 layers with 128 hidden features. Finally, GradMetaNet consists of 10 layers with 256 hidden features, while GradMetaNet++ uses 3 blocks with 8 heads and 64 hidden features.
G Additinnal Experimental Results
In this section, we detail additional experimental results for GradMetaNett-usee learned optimizers and curvature estimation experiments. The experimental setup, datasets, and baselines are all identical to the setting of Section 7 and Appendix F
G.1 Scaling Curvature Estimation
We extend the survature-estimatio experiment from Section 7.1 to models with over one million parameters. As in the main text, networks are trained to predict the trace of the Fisher Information Matrix, Ir (Fo), using the same secomposed-graddient representation. At this scale, several original baselines, especially full-gradient methods, are computationally infeasible in memory or runtime. We therefore compare against scalable baseline: an MLP that operates decomposed gradients.
35

---

GradMetaNet attains substantially lower normalized test MAE than the MLP baseline (lower is better), indicating more accurate curvature estimation scale.
Table Large-scale curvature estimation (IM+ parameters), Normalized test MAE (L).
Model
Normalized Test MAE (L)
MLP
0.779
GradMetaNet
0,413
G.2 Additional Learned Optimizer Speedup Results
Step reduction details. In Table 6 add to Table the standard deviation of the average step reduction factor for each optimizer, as well as the best test loss achieved by each optimizer in the training horizon (2,000 for image classification tasks and 5,000 for the LM task). For the reader's
convenience, in Figure 9 we presented a larger version of the plots in Figure 6
CIFAR10
CIFAR100
4.6
22
4.2
Test NLL
2.0
4.0
1.8
3.6
3.4
0
250
500
750
1000
1250
1500
1750
2000
0
250
500
750
1020
1250
1500
1750
2000
lteration
Iteration
FashionMNIST
Transformer
1.0
10.5
0.9
10.0
Test NLL
0.8
8.5
8.0
7.5
250
500
750
1000
1250
1500
1750
2000
1000
2000
3000
4000
5000
Iteration
Iteration
Optimizer
DeepSet
DeepSet GradMetaNet
GradMetaNet
Test loss curves for MLP image classification tasks and transformer language model trained on LMIB, using different optimizers and (learning rate tuned) Adam. Curves are smoothed and averaged over random initializations, with shaded regions representing standard deviation.
Figure 9:
Table 8: Training iteration speed (in iterations per second) for Adam, UNF, and UNF GradMetaNett on LMIB (transformer) and CIFAR100 (MLP).
Optimizer
Transformer LMIB (It/s) MLP on CIFARI00 (It/s)
Adam
107.55 E 8.09
359.54 = 5.67
UNF
75.45 + 2.32
304.60 12.45
UNF GradMetaNet
65.61 + 0.31
264.33 + 4.94
36

---

Train-time comparison. We measured the time training iteration (in iterations per second) for Adam, UNF, and UNF+ + GradMetaNet on LMIB and CIFARIO0. The results are reported in Table 8 All models and optimizers run on a single NVIDIA A A100-SXM4-40GB GPU. As expected, learned optimizers introduce some computational overhead, and while GradMetaNet-basee learned optimizers incur slight per-step slowdown, the significant reduction in steps leads to GradMetaNet having a substantial speedup in train-time (up to 3x faster than Adam). As reported in Table 9 GradMetaNet-basee optimizers outperform all baselines in total training speed.
Table 9: Step reduction and train-time speedup relative to Adam across datasets.
Dataset
Optimizer
Step reduction
Train-time speedup vs. Adam (1)
factor vs. Adam (1)
UNF
1.16x
1.05x
Fashion MNIST
GradMetaNet UNF
1.51x
1.13x
UNF
2.64x
2.31x
CIFARIO
GradMetaNet UNF
4.26x
3.13x
UNF
1.58x
1.34x
CIFARI00
GradMetaNet UNF
2.85x
2.10x
UNF
1.48x
1.04x
LMIB
GradMetaNet UNF
1.82x
1.11x
G.3
Generalization to New Tasks and Model Sizes
In this section, demonstrate that GradMetaNNt- learned optimizers generalize across base model sizes and across tasks without re- re-meta-training This contrasts with other weight-space learned optimizers such as DWS [1], NFN (94), and UNF [92, which require defining different weight- space metanetwork to process gradients from different base architectures. All metrics are reported relative to Adam under identical training settings.
Model-size generalization. For CIFAR-I0 and CIFAR-I00, meta-train learned optimizer specific MLP width and meta-test on larger, previously unseen MLP: (i) CIFAR-10: meta-train on a 128 hidden dims and meta-test on 256 hidden dims. In both cases, the optimizer transfers successfully to the larger model, improving both the number of steps and wall-clock time to reach the same target.
32-hidden-dim MLP and meta-test 64 hidden dims; (ii) CIFAR- 100: meta-train
Table 10: Model-size generalization: meta-train smaller MLP and meta-test larger MLP.
Dataset (train width - test width)
Optimizer
Step reduction
Train-time
factor vs. Adam (1)
speedup vs. Adam (1)
DS
2. 18x
I.85x
CIFAR-10 (32 64)
GradMetaNet DS
4.15x
3.05x
DS
1.78x
1.59x
CIFAR-100 (128 - 256)
GradMetaNet DS
3,04x
2.37x
Task generalization. We also evaluate cross-task transfer by meta-training on CIFAR-10 10 and meta-testing on CIFAR-100, and vice versa. Because the output dimensionality differs, this setting additionally exercises a mild form of architecturc/size transfer. GradMetaNet-based optimizers generalize in both directions.
Table 11: Task generalization across CIFAR tasks.
Meta-train -> Meta-test
Optimizer
Step reduction factor vs. Adam (1)
Train-time
speedup Adam (t)
DS GradMetaNet + DS
2.45x
1.72x
CIFAR- CI
3.57x
2.83x
DS
1.93x
1.57x
CIFAR-10 CIFAR-100
GradMetaNet + DS
3.80x
2.59x
37

---

G.4 Scaling to Larger Models
Context. Scaling learned optimizers is challenging due to extreme compute demands (c.g.. VeLOO [58] reportedly required ~4,000 TPU-months to meta-train). Training at that scale cur- rently infeasible with our academic resources, irrespective of whether GradMetaNet is used as the architectural backbone. Nevertheless, to show the feasibility of scaling GradMetaNet-bawee learned optimizers to larger models and settings, we measure the update cost of a GradMetaNet based optimizer for a GPT-2-scale model.
Large-model update cost. We measure the update time and memory footprint of GradMetaNet- based optimizer when applied to the gradients of GPT--2-ccale model: 12-layer transformer with 12 attention heads per-layer, hidden size of 768, totaling ~ 11M parameters. On two NVIDIA A100-SXM4-40GB GPUs, a GradMetaNet update (including backpropagating through the base transformer) takes 1.29s versus 0.89s for Adam, with similar memory footprint.
Table 12: Update-time and memory comparison on GPT-2 scale transformer using 2xA100-40GB
Optimizer
Update Time (s) ())
Memory Footprint (L)
Adam
0.89
77.2 GB
GraaMetaNNt
1.29
~77 GB (similar to Adam)
38

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In Section and Appendix E we state and prove all the theoretical results mentioned in the abstract and introduction. Section 7 and Appendix F detail the experimental setup and
results that support all empirical calims made in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See Section S
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: In section 6 we state all the assumptions made by our theoretical results, including their necessity (e.g. the set & in Theorem (6.2). In Appendix B we formally restates all results and assumptions and provide complete proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
39

---

All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes]
Justification: We provide high-level description of the experimental setup in Section 7 and provide all experimental details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For novel architecture, describing the architecture fully might suffice, if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of model checkpoint, or other means that are appropriate to the research performed.
example, if the contribution
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source source models, it may be that access to the model is limited in some way (e.g.. to registered users). but it should be possible for other researchers to have some path to reproducing or verifying the results.
Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We are currently in the process of cleaning and unifying the codebase, which includes different experiments implemented in different frameworks, using both JAX and PyTorch (to comply with baseline implementations). making the task more involved. We are committed to releasing well-documented version of the code as soon possible. Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (bttps://nipsscc/pbbl))
wuiiddes/fodd@ulllisssiooooooo. for more details.
40

---

While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code. unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //npss cc/ cc/pubblic. /
guides/Codes Saaas//ooeSuubiiice ssiooiiicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs data and code is permitted.
Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Refer to Appendix Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is
necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes]
Justification: All experiments
with multiple random seeds, results reported mean + std or with shaded regions for plots. See e.g. Figure 7 and Table For readability, the results in
Table are presented as mean only, but 6 the std for these results as well.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors). It should be clear whether the error bar is the standard deviation or the standard enror of the mean.
It is OK to report 1 I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates). error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experi- ments?
Answer: [Yes]
Justification: Computational resources are detailed in Appendix F
41

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud
provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htt.r//ppeee//eeeppppooo/ooooooa
Answer: [Yes]
Justification: We have read and complied with the NeurlPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consideration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of Deep Learning whole. We don't 't believe there are any noteworthy socictal impacts.
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also diseuss possible mitigation strategies (e.g., gated release of models, providing defenses addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
This paper poses not such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
42

---

Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere
to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been seraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected??
Answer: [Yes]
Justification: All baselines and datasets used are credited throughout Section 7 and Appendix F Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version the asset used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information creators. 13. New assets
available online, the authors are encouraged reach out the asset's
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human
subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should included in the main paper.
According to the NeurlPS Code Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals
43

---

(or an equivalent upproval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does involve crowdsoureing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
The core method development in this paper does not involve LLMS any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, or non-standard components. Please refer to our LLM policy (https: http://pp////ooeeeeeee////////// for what should or should not be described.
44


Reviews:
[{"review_id": "uVmZNR5pY1", "reviewer": "Reviewer_wdp4", "summary": "This paper introduces GradMetaNet, to design the neural network architecture specifically for learning on gradients. The key point is an approach guided by three main ideas: (1) using a rank-1 decomposition of gradients for a structured representation, (2) processing sets of gradients from multiple data samples to capture local loss landscape information, and (3) designing the architecture to be equivariant to the permutation symmetries of neurons in the underlying model. The authors provide theoretical insights for their approach, including a universality theorem and a proof that methods using only average gradients are fundamentally less expressive. They demonstrate the effectiveness of GradMetaNet on several tasks and shows improvements over previous methods.", "questions": "Many questions are already covered in the weakness sections. Below I put some of them explicitly:\n- Table 7 indicates that GradMetaNet introduces a non-trivial computational overhead per step. While Table 1 shows impressive reductions in the number of steps to convergence, this does not directly translate to a reduction in total training time. Could the authors provide a comparison of the total time required for each optimizer to reach the target test loss? \n- The experiments are done on relatively small models. How does the performance and memory trade-offs go when scaling GradMetaNet to large-scale architectures (e.g., LLMs with billions of parameters)? Maybe a LLM with 100/200M parameters is already good enough for this purpose.\n- Storing a set of decomposed gradients for a large batch (b) across a massive model could become a memory bottleneck. How does this compare to existing methods like K-FAC or Shampoo in large-scale settings?\n- The learned optimizer experiments involve meta-training on a specific family of tasks and architectures. How well would a GradMetaNet-based optimizer, meta-trained on the small models, generalize to different or larger architectures? Does the improved curvature awareness grant better generalization, or would it require expensive re-meta-training for new domains?", "limitations": "Please check the weakness and question sections for more details.", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper provides the theoretical results for its claims. It proves that processing sets of gradients is more powerful than processing the average gradient for key tasks like natural gradient approximation. \n- The experiments include good ablations and comparisons that show the benefits of each design principle. By comparing against methods that use full gradients (vs. decomposed), average gradients (vs. sets), and non-equivariant architectures, the authors provide convincing evidence for each of their core ideas .\n\nWeaknesses:\n- The experiments are using relatively small-scale models. The image classification use MLPs on downsized 8x8 images, and the language model is a small 2-block transformer. The paper acknowledges scaling to state-of-the-art LLMs as future work, but the absence of results on larger, more modern architectures is a significant limitation. \n- The proposed method introduces computational overhead compared to standard approaches. Table 7 shows that GradMetaNet is slower per iteration. While the method may converge in fewer steps, the increased time per step is not fully explored in the results. The attention-based GradMetaNet++ is particularly expensive, with quadratic complexity in the number of neurons and batch size .\n- The paper's novelty lies in the clever synthesis and application of several existing concepts rather than the invention of fundamentally new ones. The architecture uses equivariant layers adapted from prior work on set and graph processing (e.g., DeepSets) , and the idea of using rank-1 gradient factorizations is central to established methods like K-FAC. \n- For the learned optimizer application, the paper relies on a meta-training setup, which could be expensive and complex. It is unclear how well an optimizer meta-trained on the paper's small-scale tasks would generalize to different, larger architectures.", "was_revised": true, "final_justification": "Most of concerns get resolved in the reviewer's rebuttal. So I will raise my rating to 4.", "rebuttal": "We thank the reviewer for their thoughtful and detailed feedback. We are pleased they recognize the paper's theoretical contributions (S1), its clear demonstration of the benefit of each design principle (S2), and the \"impressive reductions in the number of steps\" in the learned optimization application. We also appreciate the constructive criticism and questions, which we address below.\n\n- **Total Training-Time Speedup (W2, Q1)** We thank the reviewer for this useful suggestion. Even with a slight overhead per-step, the significant reduction in steps leads to GradMetaNet having a **substantial speedup in train-time (up to 3x faster than Adam)**. Inspired by the reviewer’s comment, we compute the training-time needed to reach Adam's best test loss, and report the total train-time speedup below.\n  - **Fashion MNIST** \n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time Speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.16x|1.05x|\n|GradMetaNet + UNF|**1.51x**|**1.13x**|\n  - **CIFAR10**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time Speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|2.64x |2.31x |\n|GradMetaNet + UNF|**4.26x**|**3.13x**|\n  - **CIFAR100**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.58x |1.34x|\n|GradMetaNet + UNF|**2.85x**|**2.10x**|\n  - **LM1B**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.48x |1.04x|\n|GradMetaNet + UNF|**1.82x**|**1.11x**|\n\n  GradMetaNet-based optimizers **outperform all baselines** in total training speed. To fit within the character limit we only include here the results for the top optimizers, and will add a comprehensive table with the train-time comparisons to the revised paper.\n\n- **GradMetaNet Generalizes to New Tasks and Model Sizes (W4, Q4)** Inspired by the reviewer's comment, we have conducted new experiments that demonstrate **GradMetaNet-based learned optimizers are able to generalize across tasks and model sizes**. We note that this capability is not possible for learned optimizers that are based on other weight space models, such as DWS [1], NFN [4], or UNF [7], which require a defining a different weight-space model (metanetwork) to process gradients of different base architectures. \n  - **Model Size Generalization** For both the CIFAR-10 and CIFAR-100 tasks, we meta-trained an optimizer for a specific MLP architecture and then evaluated its ability to optimize a larger, unseen MLP, without any re-meta-training. Specifically:\n    - For CIFAR-10, we meta-train on a 32-hidden-dim MLP and meta-test on a 64-hidden-dim MLP.\n    - For CIFAR-100, we meta-train on a 128-hidden-dim MLP and meta-test on a 256-hidden-dim MLP.\n\n    In both cases, the optimizer **generalized successfully to the larger model**.\n    - **CIFAR10 (32 -> 64)**\n    |Optimizer|Step Reduction vs. Adam (↑)| Train-Time Speedup vs. Adam (↑)|\n    |-|-|-|\n    |DS (32 -> 64)|2.18x|1.85x|\n    |GradMetaNet + DS (32 -> 64)|**4.15x**|**3.05x**|\n    - **CIFAR100 (128 -> 256)**\n    |Optimizer|Step Reduction vs. Adam (↑) | Train-Time Speedup vs. Adam (↑)|\n    |-|-|-|\n    |DS (128 -> 256)|1.78x |1.59x|\n    |GradMetaNet + DS (128 -> 256)| **3.04x** |**2.37x**|\n\n  - **Task Generalization** We then meta-train optimizers on CIFAR10 and meta-test them on CIFAR100 and vice versa. In both cases, the optimizer **generalized successfully to the new task**. Note that this experiment also includes size generalization as the output dimension required for these tasks is different.\n    - **CIFAR100 -> CIFAR10 (128 hidden dim)**\n    | Optimizer |Step Reduction vs. Adam (↑) | wall-clock speedup vs. Adam (↑)|\n    |-|-|-|\n    |DS|2.45x |1.72x|\n    |DS + GradMetaNet|**3.57x**|**2.83x**|\n    - **CIFAR10 -> CIFAR100 (32 hidden dim)**\n    |Optimizer |Step Reduction vs. Adam (↑) | Train-Time Speedup vs. Adam (↑)|\n    |-|-|-|\n    |DS|1.93x|1.57x|\n    |DS + GradMetaNet|**3.80x**|**2.59x**|\n\n- **Scaling GradMetaNet (W1, Q2)** We appreciate the reviewer's comment regarding the importance of scaling. First, we highlight that in order to ensure a fair comparison, our experimental design follows **established benchmarks** and uses the **exact setup** of established papers in learned optimization and weight-space learning [8]. We elaborate in the context each experiment:\n  - **Learned Optimization** Scaling learned optimizers is a long-standing challenge due to high computational demands. For example, VeLO [5] was trained for over **~4,000 TPU-months**. As such, training optimizers at this scale is currently infeasible with our academic compute resources, **regardless of whether GradMetaNet is used**. In our paper, we follow the exact setup of prior work [4, 6], conducting evaluations in settings of up to 2M parameters to analyze architectural improvements.\n\n    That said, inspired by the reviewer’s suggestion, we evaluated the update time and memory cost of a GradMetaNet-based optimizer on (the gradients of) a GPT-2–scale model. We used a 12-layer Transformer with 12 attention heads, hidden size of 768, and approximately **117M parameters**. On two NVIDIA A100-SXM4-40GB GPUs, a GradMetaNet update takes 1.29s (including backpropagationg through the base transformer model), compared to 0.89s using Adam, and has a similar memory footprint.\n    |Optimizer|Update Time|Memory Footprint|\n    |-|-|-|\n    |Adam|0.89s|77.2 GB|\n    |GradMetaNet|1.29s |77.8 GB|\n  - **INR Editing** Our benchmarks **exactly follow** recent state-of-the-art weight-space methods [2, 3, 7] (2023-2024). GradMetaNet achieves a **22.5% improvement over prior SOTA** (Table 2), demonstrating its effectiveness on established, non-trivial tasks widely used by the community [8].\n  - **Curvature Estimation** Some baselines we compare against are not scalable by nature, and thus the setting is chosen accordingly to isolate the effect of architecture. That said, inspired by reviewer **EMd8**, we have scaled this experiment too to models with 1M parameters.\n\n  Additionally, we note that our experiments include **strictly larger models** (up to 2M parameters) than key weight-space papers. For instance [1] (ICML 2023 oral presentation), [2] (ICLR 2024 oral presentation), and [3] (NeurIPS 2024 oral presentation), all operate on models of up to 200K parameters.\n- **GradMetaNet's Novelty (W3)** We appreciate the reviewer’s characterization of our work as a \"clever synthesis\". We wish to convey that our contribution extends beyond combining existing concepts. The core methodological novelty lies in:\n  - **First Principles Approach** Our architecture is derived from first principles, using equivariant deep learning techniques [9]. We begin by analyzing the symmetry group and action associated with the space of decomposed gradients, a novel mathematical setting that **has not been studied before**, providing a principled foundation for reasoning about gradient representations.\n  - **Architecture Derivation** Building on this symmetry analysis, we mathematically derive a class of equivariant layers. An elegant outcome of this derivation is that it naturally recovers and generalizes a range of existing architectures.\n  - **Theory** Our novel formulation also enables us to prove **new results** that are **unknown for other weight-space models**. We establish a universality theorem for functions over sets of gradients (Theorem 6.2), a result that **has not been shown for any prior method** [1, 2, 3, 4, 7]. Furthermore, we show that existing models are **provably less expressive** for core gradient-based tasks (Corollary 6.3). We thank the reviewer for raising the point and will make this discussion more clear in the updated version.\n\n- **Memory Improvement over K-FAC/Shampoo (Q3)** This is an excellent question. GradMetaNet's memory scaling is one of its key advantages. **GradMetaNet is significantly more memory-efficient than K-FAC and Shampoo**, as it does not require storing large curvature matrices. Specifically, K-FAC and Shampoo store matrices consisting of outer products of activations and preactivation gradients, leading to a memory footprint of $O(d_0^2+\\cdots+d_L^2)$, where $d_i$ is the $i$-th hidden dimension. In contrast, GradMetaNet leverages the batch of activations and preactivation gradients that are **already present in memory** from the backpropagation computation. Our method only incurs an additional memory cost for a momentum term, which stores a moving average of these activations and preactivation gradients. This results in a strictly lower memory footprint of $O(d_0+\\cdots+d_L)$. We thank the reviewer for raising this point, we will make this discussion more clear in the updated version.\n\nWe thank the reviewer again for their constructive feedback. We believe we have thoroughly addressed the reviewer's concerns and clarified the key contributions of our work. We hope this response helps to convey the novelty, scalability, generality, and practical relevance of GradMetaNet, and would greatly appreciate a reconsideration of the score.\n\n**References**\n\n[1] Navon, Aviv, et al. \"Equivariant architectures for learning in deep weight spaces.\" ICML 2023.\n\n[2] Kofinas, Miltiadis, et al. \"Graph neural networks for learning equivariant representations of neural networks.\" ICLR 2024.\n\n[3] Kalogeropoulos, Ioannis, et al. \"Scale equivariant graph metanetworks.\" NeurIPS 2024.\n\n[4] Zhou, Allan, et al. \"Universal neural functionals.\" NeurIPS 2024.\n\n[5] Metz, Luke, et al. \"Velo: Training versatile learned optimizers by scaling up.\" arXiv preprint arXiv:2211.09760 (2022).\n\n[6] Harrison, James, et al. \"A closer look at learned optimization: Stability, robustness, and inductive biases.\" NeurIPS 2022.\n\n[7] Zhou, Allan, et al. \"Permutation equivariant neural functionals.\" NeurIPS 2023.\n\n[8] Schürholt, Konstantin, et al. \"Neural network weights as a new data modality.\" ICLR 2025 Workshop Proposals.\n\n[9] Bronstein, Michael M., et al. \"Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.\" arXiv preprint arXiv:2104.13478 (2021)."}, {"review_id": "rNgScsV4WB", "reviewer": "Reviewer_EMd8", "summary": "The paper introduces a neural network architecture, GradMetaNet, for processing gradient vectors of MLPs and transformers.\nSuch architectures are used to learn optimization algorithms, or in general other gradient-dependent tasks.\nThe main technique used by GradMetaNet is that gradients are not represented in weight space, but rather in terms of activations and pre-activation gradients.\nThis perspective (i) simplifies making the meta network respect the permutation symmetry induced invariances in the gradient, and (ii) allows to leverage gradient information on a per-datum basis, rather than operating on the (averaged) mini-batch gradient.\nThe paper presents theoretical results that GradMetaNet can represent functions that cannot be represented when using the averaged weight-space gradient.\nEmpirical results show that both GradMetaNet's equivariance and capability to leverage per-datum information are useful to learn curvature information, optimization algorithms, and editing the weight space of neural nets.", "questions": "- **Q1a (computational cost):** Regarding Section 7.2 and Table 1, does the faster convergence of the optimization algorithms learned with GradMetaNet carry over to run time benefits?\n  I.e., how much run time does it take to evaluate the pre-conditioner $\\mathbf{F}$ that is parameterized with GradMetaNet, and how practical are the learned optimizers?\n\n- **Q1b (computational scaling):** Please clarify the statement in L64-66, specifically why does storing per-datum gradients scale quadratically in the number of neurons?\n  For a linear layer from $\\mathbb{R}^{D_\\text{in}}$ to $\\mathbb{R}^{D_\\text{out}}$ and a batch size $B$, storing the per-datum gradients in weight space requires $B D_\\text{in} D_\\text{out}$ memory, while storing the proposed presentation uses $B (D_\\text{in} + D_\\text{out})$.\n  Do you mean the case where $D_\\text{in} = D_\\text{out}$?\n  It may be better to make this statement more concrete.\n\n- **Q2 (experimental suggestion):** Regarding Section 7.1 and Figure 7, the neural network and data set are both relatively small (around 1000 parameters and data points).\n  It would be good to present similar experiments for a larger network in the overparameterized regime, say with more than 1 million parameters, and 10-100k data points.\n\n- **Q3 (accounted symmetries):** GradMetaNet accounts for permutation symmetries.\n  But the attention mechanism of transformers actually exhibits a richer symmetry, namely under $\\mathrm{GL}$.\n  For example, consider the term $\\mathbf{W}_Q \\mathbf{W}_K^\\top$ in the attention matrix.\n  Then, we can transform $(\\mathbf{W}_Q, \\mathbf{W}_K)$ to $(\\mathbf{W}_Q \\mathbf{A}^{-1}, \\mathbf{W}_K \\mathbf{A}^\\top)$ with invertible matrix $\\mathbf{A}$, while leaving the layer invariant.\n  How complicated would it be to extend GradMetaNet to account for this symmetry (or other symmetries that are common in neural networks, see e.g. https://arxiv.org/pdf/2012.04728, Section 3)?\n\n- **Q4 (details regarding the Fisher):** In Section 3, you mention estimating curvature information in form of the Fisher as application.\n  However, the Fisher requires gradients of 'pseudo-losses' evaluated on labels $y$ drawn from the model's likelihood $p_\\theta(y | x)$.\n  It seems that you are focusing on the *empirical Fisher* instead, i.e. the un-centered second gradient moment (see https://arxiv.org/abs/1905.12558 for the important difference between empirical and sampled Fisher).\n  This should be phrased more carefully to avoid confusion.\n  Is my understanding correct that the paper is only concerned with the empirical Fisher?\n  I think it is trivially possible to support other curvature matrices like the 'true' Fisher (by backpropagating gradients of the above pseudo-losses instead, as described e.g. in https://openreview.net/pdf?id=TVbCKAqoD8, Equation 3) or the generalized Gauss-Newton matrix (by backpropagating columns of the loss function's Hessian factorization, see e.g. https://arxiv.org/pdf/2106.02624, Equation 3).", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\n- **S1 Meaningful improvement of previous architectures.**\n  The proposed representation to represent gradients in neural activation rather than weight space allows MetaGradNet to go beyond previous architectures, either in terms of equivariance, expressiveness, or computational cost.\n  As it relies on the outer product gradient structure of linear layers with weight sharing, the idea applies to basically any neural network (MLPs, CNNs, transformers, graph neural nets).\n  It is clearly presented and the paper in general is well-written.\n\n- **S2 Experiments demonstrate utility of per-datum information and symmetry-awareness.**\n  The experiments, specifically Section 7.1, are well-designed to disentangle the effects of MetaGradNet's new capabilities and show that both of them positively affect performance (in the specific experiment, to approximate the empirical Fisher's diagonal).\n\n**Weaknesses:**\n\n- **W1 Computational cost \\& practicality.**\n  The paper does not discuss the computational cost of the proposed architecture, or show empirical results how much additional cost is incurred by using a MetaGradNet.\n  I have some doubts about the practicality of the proposed applications, specifically for learning optimization algorithms (please see my question below).\n  As speed is key for optimization algorithms, it would be good if the authors provided additional details how much overhead is being added to apply a pre-conditioner parameterized with GradMetaNet.\n  This could be done e.g. by showing a version of Figure 6 with test loss over time.\n  Since I am not an expert on the application of meta networks, I am also curious to read the assessment of other reviewers on the practical relevance of the presented experiments, specifically w.r.t. chosen problem size (network and data set).\n\n- **W2 (minor) Limited to permutation symmetry.**\n  The paper mostly discusses how to account for permutation symmetries in weight space.\n  However, many neural networks exhibit richer symmetries, such as translation, scaling, re-scaling, or even $\\mathrm{GL}$ symmetries of transformers (see my question below).\n  It would be good to discuss how GradMetaNet could be extended to account for these symmetries to broaden its applicability or possibly improve the presented results on transformers.", "was_revised": true, "final_justification": "Leaning towards accept (between 'weak accept' and 'accept').\n\nThe authors were able to address many of my doubts, reducing some of my initial uncertainty.\nConcretely, they provided additional insights into computational cost of GradMetaNet, as well as generalizations of the approach to other symmetries.\nI am hesitant to bump my score to 5 (accept) as I am not an expert on meta networks, which complicates evaluating how representative the shown experiments are and what someone with deep expertise on this field would say.\nOverall, I did not identify flaws in the logic, found the additional experiments to be related to the questions brought up by the review, and trust the authors response that 'our experiments follow standard benchmarks'.", "rebuttal": "We thank the reviewer for their detailed and insightful feedback. We are encouraged that they recognized the meaningful improvement of our architecture (S1), the broad applicability of our method and ideas, the well-designed experiments (S2), significant step reduction in the learned optimizer application, and the paper's clear presentation.\n\nWe also appreciate the constructive feedback and address the questions and concerns below. We believe these revisions significantly strengthen the paper, and we would genuinely appreciate a reconsideration of the score.\n- **Total Training-Time Speedup (W1, Q1a)** We thank the reviewer for this useful suggestion. For an iterations/second comparison of GradMetaNet-based optimizers we refer the reviewer to Table 7 in Appendix F.2. Even with a slight overhead per-step, the significant reduction in steps leads to GradMetaNet having a **substantial speedup in train-time (up to 3x faster than Adam)**. Inspired by the reviewer’s comment, we compute the training-time needed to reach Adam's best test loss, and report the total train-time speedup below.\n  - **Fashion MNIST** \n|Optimizer |Step Reduction vs. Adam (↑)|Train-Time Speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.16x|1.05x|\n|GradMetaNet + UNF|**1.51x**|**1.13x**|\n  - **CIFAR10**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time Speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|2.64x |2.31x |\n|GradMetaNet + UNF|**4.26x**|**3.13x**|\n  - **CIFAR100**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.58x |1.34x|\n|GradMetaNet + UNF|**2.85x**|**2.10x**|\n  - **LM1B**\n|Optimizer|Step Reduction vs. Adam (↑)|Train-Time speedup vs. Adam (↑)|\n|-|-|-|\n|UNF|1.48x |1.04x|\n|GradMetaNet + UNF|**1.82x**|**1.11x**|\n\n  GradMetaNet-based optimizers **outperform all baselines** in total training speed. To fit within the character limit we only include here the results for the top optimizers, and will add a comprehensive table with the train-time comparisons to the revised paper.\n\n  Additionally, we provide a computational complexity analysis in Appendix D.4 and show that a forward pass through GradMetaNet has complexity $O(N \\cdot b)$, where $N$ is the number of neurons in the network we optimize and b is the batch size. Therefore, GradMetaNet **does not change the asymptotic complexity of backpropagation** on the original model. We will make this point more visible  in the revised manuscript.\n\n  Finally, to provide context regarding the model size for other applications of metanetworks, we note that our experiments follow standard benchmarks in learned optimization and the weight-space learning community [1, 2, 3, 7, 8] and include **strictly larger models (up to 2M parameters) than key weight-space/metanetwork papers**. For instance [1] (ICML 2023 oral presentation), [2] (ICLR 2024 oral presentation), and [3] (NeurIPS 2024 oral presentation), operate on models of up to 200K parameters. We thank the reviewer again for raising this important point, which we believe will strengthen the paper.\n\n- **Scaling Curvature Estimation Experiment (Q2)** We thank the reviewer for this suggestion. We have conducted an extended version of the experiment in Section 7.1 with **models containing over 1M parameters**. Our networks are trained to predict the trace of the Fisher Information Matrix. At this scale, many of the original baselines, particularly full-gradient methods, are computationally infeasible. Therefore, as a scalable comparison we use an MLP, also operating on decomposed gradients. **GradMetaNet beats the MLP** with a normalized test MAE of 0.413 compared to 0.779 (lower is better). We will include this large-scale curvature estimation experiment in the final version of the paper. Additionally, we refer the reviewer to our learned optimizer experiments on LM1B which also demonstrate GradMetaNet’s effectiveness at the **1M+ parameter scale**.\n|Model|Normalized Test MAE (↓)|\n|-|-|\n|MLP|0.779|\n|GradMetaNet|**0.413**|\n\n- **Accounted Symmetries (W2, Q3)** We thank the reviewer for this insightful question. It is indeed true that neural architectures often exhibit symmetries beyond simple neuron permutations. While MLP weights do admit rescaling symmetries, these **do not always extend to gradient-based quantities**. For example, in Section 7.1, we learn the diagonal of the Fisher Information Matrix (FIM), which is 2-homogeneous in the gradients: $\\mathrm{diag}(\\mathbf{F}(\\lambda \\mathbf{W})) = \\lambda^2 \\mathrm{diag}(\\mathbf{F}(\\mathbf{W}))$. This indicates that enforcing rescaling equivariance **would be inappropriate for such tasks**. We therefore focus on permutation equivariance [4, 5, 6], which is general, task-agnostic, and remains valid across all gradient-level problems considered in this work. We briefly discuss broader symmetry considerations, including transformer $\\mathrm{GL}$-symmetries,  in Appendix B.2 and believe this is an incredibly compelling direction for future research.\n\n- **Computational Complexity Clarification (Q1b)** Yes, the reviewer is exactly right, we are referring to the specific case where $D_\\mathrm{in} = D_\\mathrm{out}$ in the paper. We'll clarify this in the revised text. To be more precise, we will emphasize the general point that the cost is additive rather than multiplicative with respect to the hidden dimensions. This means that instead of a cost that scales with the product of hidden dimensions, it scales with their sum. We will include a formula to illustrate this more clearly in the updated version. We thank the reviewer for raising this point.\n\n- **Fisher Information Matrix Clarification (Q4)** The reviewer's observation about the Fisher Information Matrix is accurate. In the first experiment (loss landscape curvature estimation), we use the full \"true\" Fisher as the target for prediction, while for the learned optimizer experiments, GradMetaNet is provided with the necessary information to compute the empirical Fisher (in the form of sets of individual gradients). While (as the reviewer points out) GradMetaNet is capable of supporting other curvature matrices, such as the \"true\" Fisher (by backpropagating gradients of pseudo-losses) or the Generalized Gauss-Newton, in the learned optimizers experiment, we chose to focus on the empirical Fisher to simplify the experimental setup and avoid additional complexities like sampling labels. We agree that making this distinction clearer will enhance the precision of the paper, and will integrate a more detailed explanation, as well as the suggested referenced in the revised manuscript.\n\n- **Formatting** We thank the reviewer for spotting these, and agree with their suggestions for improvement. The sentence in L199 should be: “...concatenates a layer identifier **to each neuron** in the intermediate layers and a neuron identifier to each neuron in the first and last layers”. We will implement the suggestions and corrections, and add the reference mentioned by the reviewer in the revised version.\n\nWe thank the reviewer again for their constructive feedback. We believe we have thoroughly addressed the concerns raised and clarified the key contributions of our work. We hope this response helps convey the computational efficiency, and practical relevance of GradMetaNet, and would greatly appreciate a reconsideration of the score.\n\n**References**\n\n[1] Navon, Aviv, et al. \"Equivariant architectures for learning in deep weight spaces.\" International Conference on Machine Learning. PMLR, 2023.\n\n[2] Kofinas, Miltiadis, et al. \"Graph neural networks for learning equivariant representations of neural networks.\" arXiv preprint arXiv:2403.12143 (2024).\n\n[3] Kalogeropoulos, Ioannis, Giorgos Bouritsas, and Yannis Panagakis. \"Scale equivariant graph metanetworks.\" Advances in neural information processing systems 37 (2024): 106800-106840.\n\n[4] Zaheer, Manzil, et al. \"Deep sets.\" Advances in neural information processing systems 30 (2017).\n\n[5] Hartford, Jason, et al. \"Deep models of interactions across sets.\" International Conference on Machine Learning. PMLR, 2018.\n\n[6] Maron, Haggai, et al. \"On learning sets of symmetric elements.\" International conference on machine learning. PMLR, 2020.\n\n[7] Schürholt, Konstantin, et al. \"Model zoos: A dataset of diverse populations of neural network models.\" Advances in Neural Information Processing Systems 35 (2022): 38134-38148.\n\n[8] Schürholt, Konstantin, et al. \"Neural network weights as a new data modality.\" ICLR 2025 Workshop Proposals."}, {"review_id": "P5fegTdolv", "reviewer": "Reviewer_cxim", "summary": "This paper proposes GradMetaNet, a gradient-based meta-learning architecture that incorporates symmetry-based inductive biases through equivariant graph neural networks. The authors present a meta-learner that operates on a learned representation of gradients and parameter updates using group-equivariant graph neural networks. By modeling the optimization process in an equivariant manner, the method is designed to generalize better across tasks that share common symmetry structures (e.g., rotations or permutations). Experiments on synthetic regression tasks and Omniglot classification demonstrate improved generalization performance compared to baseline MAML-style and GNN-based meta-learners.", "questions": "The paper presents a novel and well-motivated approach to gradient-based meta-learning using equivariant GNNs. It is a solid contribution that brings together two important themes—meta-learning and symmetry-aware modeling—and demonstrates improvements over several baselines. It is a promising and creative direction worthy of presentation.\n\nComments\n- Please provide more discussion or formal analysis of how equivariance improves meta-generalization beyond the empirical findings.\n- Include a comparison to more recent and competitive meta-learners if possible (e.g., Meta-SGD, MetaFormer).\n- Consider adding experiments where the assumed task symmetry is violated, to evaluate the robustness of the proposed architecture.\n- Clarify in more detail how the EGNN operates on gradient information—how are nodes and edges defined, and how is parameter sharing handled across tasks?\n- A small diagram illustrating the architecture pipeline (input → gradient encoding → EGNN → parameter update) would help clarify the proposed method.", "limitations": "yes", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The main strength is the motivated use of equivariance: the paper makes a compelling case that symmetries in task structure—especially in few-shot learning—can be better leveraged using equivariant models. The incorporation of EGNNs into meta-learning is novel and conceptually well-founded. It is also interesting how the architecture combines inner-loop optimization with learned gradient manipulation in a principled way. The use of gradient information as input to an equivariant network is an interesting design that could inspire further work. Furthermore the method is evaluated on both synthetic and real datasets, including few-shot regression and classification. GradMetaNet outperforms several baselines, showing the benefit of incorporating equivariance in meta-learners. Another strength is that the paper is generally well-written and structured. It provides clear descriptions of the problem setup, architecture, and training procedure.\n\nI originally complained about the lack of  a deeper theoretical justification for why task generalization benefits specifically from equivariant gradient modeling. The  discussion with the authors and their proposed additions/changes to the paper answer fully my concerns.", "was_revised": true, "final_justification": "I like the motivated use of equivariance: the paper makes a compelling case that symmetries in task structure—especially in few-shot learning—can be better leveraged using equivariant models. GradMetaNet outperforms several baselines, showing the benefit of incorporating equivariance in meta-learners. I am not an expert in optimization and cannot judge with full confidence how meaningful this contribution really is. From my  perspective I think there quite a few deeper issues in machine learning where a contribution would have much more of an impact. For this reason, and for being relatively ignorant in the specific subject, I am still unsure where it should be a 4 or a 5. I updated my original grade because my main criticisms - on lack of a theoretical justification - were addressed, but I am not very confident.", "rebuttal": "We thank the reviewer for their positive assessment of our work and for finding it to be a \"promising and creative direction worthy of presentation\". We are especially glad they recognized the motivated use of equivariance and our principled approach to operating on gradient information as key strengths of our paper.\n\nWe also appreciate the constructive suggestions for improvement and will address the points raised below.\n\n- **Theoretical Benefits of Equivariant Modeling (Comment 1)** We thank the reviewer for highlighting the importance of theoretical backing. Equivariant modeling has been shown to improve generalization both theoretically [1, 2] and empirically [3, 4, 5]. We also empirically demonstrate benefits directly in Section 7.1. The main idea is that modeling an equivariant task with an equivariant model induces a correct inductive bias, and reduces the hypothesis space. In our case, a general non-equivariant model (e.g. an dense MLP) with hidden dimension $d$ would have $O(b \\cdot d \\cdot P)$ parameters, where $b$ is the batch size and $P$ is the number of parameters in the base model. In comparison, GradMetaNet has a **constant ($O(1)$) number of parameters**, regardless of the size of the base network. We will make this point clearer in the theoretical section of  the revised manuscript. Additionally, we refer the reviewer to Section 6 in our paper where we prove that under mild assumptions GradMetaNet is able to approximate any computation on sets of gradients (Theorem 6.2), and that other common weight-space methods cannot do the same (Corollary 6.3).\n\n- **GradMetaNat Beats State-of-the-Art Baselines (Comment 2)** We appreciate the suggestions for additional baselines. Our work focuses on the broad challenge of learning on gradients for tasks like learned optimization, model editing, and curvature estimation, rather than MAML-style few-shot meta-learning. Therefore, we chose baselines that are state-of-the-art for these specific domains. For instance, in our learned optimizer experiments (Section 7.2), we compare against state-of-the-art equivariant weight-space architectures like UNF [6]  and per-parameter optimizers like DeepSets [3, 7, 8], which are the most relevant recent methods for this task. For the INR editing experiment we compare against the state of the art methods: DWS [9] (2023), GMN [10] (2024), and ScaleGMN [11] (2024). We believe this provides a rigorous and appropriate comparison for the problems we address.\n\n- **GradMetaNet Can Be Applied to General Models and Tasks** We appreciate the reviewer’s comments regarding the applicability of our methods to “tasks with more abstract or learned symmetries (e.g., language or combinatorial domains)”. We would kindly like to clarify that the symmetries discussed in the paper are **not symmetries of the data**, but rather **parameter symmetries** of the neural architecture. These symmetries broadly hold for all models composed of linear layers and point-wise activations, including MLPs, transformers, CNNs, etc., regardless of the task they are trained to perform. Therefore, any task involving the gradients of such models is equivariant to this symmetry and will benefit from the incorporation of equivariant layers. In particular, GradMetNet can be applied in language/combinatorial domains, and we even meta-train GradMetaNet to optimize transformer based **language models** in Section 7.2.\n\n- **Evaluating the Robustness of GradMetaNet (Comment 3)** We thank the reviewer for the excellent question regarding isolating the contribution of equivariance. As we mentioned in the previous answer, the symmetry GradMetaNet respects holds for general architectures, and is not task dependent. Therefore, to isolate the importance of equivariance, we conducted the experiment in Section 7.1, displayed in Figure 7. This experiment is designed as a **comprehensive ablation study**. We compare GradMetaNet against several baselines that partially or completely disregard the symmetries, including:\n  - \"MLP + Concat\": A standard, non-equivariant MLP baseline.\n  - \"Neuron Asymmetric GradMetaNet\": A variant that is not equivariant to neuron permutations.\n  - \"Batch Asymmetric GradMetaNet\": A variant that is not invariant to the ordering of gradients in the input set.\n\n  These results demonstrate that our principled, **fully-equivariant design** has significantly better sample complexity and outperforms the ablated non-equivariant versions, confirming that equivariance is crucial for success.\n\n- **Architectural Details and Visualizations (Comments 4, 5)** We appreciate the reviewer's feedback regarding the clarity and visualization of GradMetaNet's operations. We understand the importance of a comprehensive explanation and are happy to provide further details. GradMetaNet's architecture is thoroughly described in Section 5 of our paper, specifically in Equation (9) and bullet points (I-V) which define the formula and operation of each layer as well as the overall structure. A more in-depth discussion of these components can also be found in Appendix D.1. The revised version of our paper will include an even more detailed explanation of these architectural aspects and more references to the Appendix for enhanced understanding. Regarding a diagram illustrating GradMetaNet's pipeline, we would like to direct the reviewer to Figure 5 in the paper. This figure visualizes the full GradMetaNet pipeline:  from receiving a set of neural network gradients, to the rank-1 decomposition, to the equivariant processing blocks, and finally to the output (e.g., a parameter update or curvature information). We will incorporate a more explicit reference to this figure in Section 5 to guide readers more effectively. Should the reviewer still feel that an additional diagram is necessary for further clarity, we would be pleased to provide one.\n\nWe thank the reviewer again for their time and their positive feedback. We are confident that by incorporating these clarifications and additions, particularly the architectural diagram, the final version of the paper will be even stronger.\n\n**References**\n\n[1] Tahmasebi, Behrooz, and Stefanie Jegelka. \"The exact sample complexity gain from invariances for kernel regression.\" Advances in Neural Information Processing Systems 36 (2023): 55616-55646.\n\n[2] Behboodi, Arash, Gabriele Cesa, and Taco S. Cohen. \"A pac-bayesian generalization bound for equivariant networks.\" Advances in Neural Information Processing Systems 35 (2022): 5654-5668.\n\n[3] Zaheer, Manzil, et al. \"Deep sets.\" Advances in neural information processing systems 30 (2017).\n\n[4] Hartford, Jason, et al. \"Deep models of interactions across sets.\" International Conference on Machine Learning. PMLR, 2018.\n\n[5] Maron, Haggai, et al. \"On learning sets of symmetric elements.\" International conference on machine learning. PMLR, 2020.\n\n[6] Zhou, Allan, Chelsea Finn, and James Harrison. \"Universal neural functionals.\" Advances in neural information processing systems 37 (2024): 104754-104775.\n\n[7] Harrison, James, Luke Metz, and Jascha Sohl-Dickstein. \"A closer look at learned optimization: Stability, robustness, and inductive biases.\" Advances in neural information processing systems 35 (2022): 3758-3773.\n\n[8] Metz, Luke, et al. \"Practical tradeoffs between memory, compute, and performance in learned optimizers.\" Conference on Lifelong Learning Agents. PMLR, 2022.\n\n[9] Navon, Aviv, et al. \"Equivariant architectures for learning in deep weight spaces.\" International Conference on Machine Learning. PMLR, 2023.\n\n[10] Kofinas, Miltiadis, et al. \"Graph neural networks for learning equivariant representations of neural networks.\" arXiv preprint arXiv:2403.12143 (2024).\n\n[11] Kalogeropoulos, Ioannis, Giorgos Bouritsas, and Yannis Panagakis. \"Scale equivariant graph metanetworks.\""}]

Author rebuttal:

Meta-review: This paper proposes a novel, principled architecture GradMetaNet for learning on gradients, which leverages symmetry of the neuromanifold, related invariances, and rank-1 structure of the gradients. The authors provided a universal approximation theorem for such invariant mappings, as well as empirical evaluations on precipitating the diagonal FIM and learned optimization.

All three reviewers acknowledge the novelty and rich expressivity that is guaranteed in theory, making the proposed architecture a meaningful advance over previous architectures. The empirical results are carefully designed and comprehensive, with a notable performance gain compared with weight-space approaches. The writing flows well with high quality illustrations.

A main concern is the scalability due to the computational overhead, and the authors clarified in the rebuttal that the overall computation is more efficient due to reduction in the number of training steps. Further experiments show that GradMetaNet can scale to GPT-2–scale mode with 1M+ parameters. Another limitation is the focus on MLP and transformers, though extensions appear plausible.

Overall, the strength outweigh the limitations, and I recommend acceptance. For the final version, the authors are recommended to carefully address the reviewers' comments and further clarify scalability and broader applicability.
Decision: Accept (poster)