Learning to Learn with Generative Models of Neural Network Checkpoints
Authors: William Peebles
Ilija Radosavovic
Tim Brooks
Alexei A Efros
Jitendra Malik
Venue: Submitted to ICLR 2023 (2023)
Keywords: diffusion
DDPMs
learning to learn
generative models
transformers

Abstract: We explore a data-driven approach for learning to optimize neural networks. We construct a dataset of neural network checkpoints and train a generative model on the parameters. In particular, our model is a conditional diffusion transformer that, given an initial input parameter vector and a prompted loss, error, or return, predicts the distribution over parameter updates that achieve the desired metric. At test time, it can optimize neural networks with unseen parameters for downstream tasks in just one update. We find that our approach successfully generates parameters for a wide range of loss prompts. Moreover, it can sample multimodal parameter solutions and has favorable scaling properties. We apply our method to different neural network architectures and tasks in supervised and reinforcement learning.


Under review as a conference paper at ICLR 2023
LEARNING TO LEARN WITH GENERATIVE MODELS OF NEURAL NETWORK CHECKPOINTS
Anonymous authors Paper under double-blind review
ABSTRACT
We explore a data-driven approach for learning to optimize neural networks, We construct dataset of neural network checkpoints and train a generative model on the parameters, In particular, our model is a conditional diffusion transformer that, given an initial input parameter vector and a prompted loss, error, or return, predicts the distribution over parameter updates that achieve the desired metric. At test time, it can optimize neural networks with unseen parameters for downstream tasks in just update. We find that our approach successfully generates parameters for a wide range loss prompts. Moreover, it can sample multimodal parameter solutions and has favorable scaling properties. We apply method different neural network architectures and tasks in supervised and reinforcement learning.
INTRODUCTION
Gradient-basse optimization is the fuel of modern deep learning. Techniques of this class, such as SGD (Robbins & Monro, 1951) and Adam (Kingma & Ba, 2015), are easy to implement, scale reasonably well and converge to surprisingly good solutions-even in high-dimensional, non-convex neural network loss landscapes. Over the past decade, they have enabled impressive results in computer vision (Krizhevsky et al., 2012; Girshick et al., 2014), natural language processing (Vaswani al., 2017; Radford et al., 2018) and audio generation (Van Den Oord al., 2016).
While these manual optimization techniques have led to large advances, they suffer from an important limitation: they are unable to improve from past experience. For example, SGD will not converge any faster when used to optimize the same neural network architecture from the same initialization the 100th time versus the first time. Learned optimizers capable leveraging their past experiences have the potential to overcome this limitation and may accelerate future progress in deep learning.
Of course, the concept of learning improved optimizers is not new and dates back to the 1980s, not earlier, following early work from Schmidhuber (1987) and Bengio et al. (1991). In recent years, sig- nificant effort has been spent on designing algorithms that learn via nested meta-optimization, where the inner loop optimizes the task-level objective and the outer loop learns the optimizer (Andrychow- icz et al., 2016; Li & Malik, 2016; Finn et al.. 2017). In some instances, these approaches outperform manual optimizers. However, they are challenging to train in practice due to a reliance on unrolled optimization and reinforcement learning.
Taking modern deep learning perspective suggests simple, scalable and data-driven approach to this problem. Over the past decade, our community has trained a massive number of checkpoints. These checkpoints contain wealth of information: diverse parameter configurations and rich metries such as test losses, classification errors and RL returns that describe the quality of the checkpoint. Instead of leveraging large-scat datasets of images or text, we propose learning from large-scaal scale datasets of checkpoints recorded over the course of many training runs.
To this end, we create dataset of neural network checkpoints (Figure left). Our dataset consists of 23 million checkpoints from over hundred thousand training runs. We collect data from supervised learning tasks (MNIST, CIFAR-10) as well as reinforcement learning tasks (Cartpole), and across different neural network architectures (MLPs, CNNs). In addition to parameters, we record relevant task-level metrics in each checkpoint, such as test losses and classification errors.
Please our project website in supplementary materials for additional results and visualizations.

---

Under review as a conference paper at ICLR 2023
Io
Loss 2.3
Error 90%
Updated 0
014
Loss 1.6
Error 63%
G. G.pt
NN
Loss = 0.18
Error 5%
Starting 0
Prompted Loss (or Error)
Pre-training dataset of neural network checkpoints
A generative model of checkpoints
Figure 1: Generative pre-training from checkpoints. Left: We build dataset of neural network checkpoints from many training runs. Each checkpoint includes the neural network's parameters and relevant metadata (test losses and test errors for supervised learning tasks, returns for RL tasks). Right: G.pt, a generative model of checkpoints. G.pt takes a parameter vector and loss/errro//eturm prompt as input and predicts the distribution over updated parameters that achieve the prompt.
Given this data, we explore generative pre-training directly in parameter space (Figure right). Specifically, train transformer-butt diffusion models of neural network parameters. Given initial input parameter vector and a target loss, error or return, these models are trained to predict the distribution updated parameter vectors for single network architecture that achieve the target metric. Our method is trained using standard generative modeling techniques instead of unrolled optimization and reinforcement learning algorithms. We call our model G.pt'.
We show that our approach has a number of favorable properties. First, it is able to rapidly train neural networks from unseen initializations with just one parameter update (Figure 3). Second, it can generate parameters that achieve a wide range of prompted losses, errors and returns (Figure 5). Third, it able generalize ouitt--disti weight initialization algorithms (Figure 6). Fourth, as a generative model, it is able to sample diverse solutions (Figure 8). Finally, it can optimize non dindeifferentialle objectives, such as RL returns or classification errors.
GENERATIVE PRE-TRAINING FROM NEURAL NETWORK CHECKPOINTS
We pre-train generative model G. pt on neural network checkpoints. At test time, we use it to generate parameters for neural networks that solve a downstream task.
2.1 A DATASET OF NEURAL NETWORK CHECKPOINTS
In order to train G. pt, we build dataset of neural network checkpoints. Each checkpoint contains neural network parameters and relevant task-level metrics like train losses, test use standard optimizers like Adam and SGD with momentum to generate the parameters, and we randomly subset of checkpoints from each training run. Our methodology for generating each individual training run is explained in detail in Algorithm See Section 3 for additional details.
returns. We
Augmenting datasets of neural networks. To offset the computational cost of collecting checkpoints, we use data augmentation in neural network parameter space. Given a checkpoint (0, E). we construct augmented tuples (T(0), E), where T(-) is the parameter-level augmentation. In order for these augmented tuples to be valid, we need (T( (z) fe for all parameter vectors 0 and all inputs to the neural network x. One type augmentation that meets this criteria permutation augmentation. Consider an MI.P. If we apply some permutation to the outgoing weights (and biases) of the input layer and the incoming weights of the next layer, the output of the neural network will be preserved (Roeder al., 2021; Schürholt al., 2021). Different permutations can be sampled for each layer up to the output layer. This technique is generic and can be applied to MLPs and CNNs alike. We apply the same permutation to both the input and target parameters during pre-training.
l and pt refer generative models and checkpoint extensions, respectively.
2

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Checkpoint Data Generation
Algorithm 2 Pre-training from Checkpoints
1: Input: Dataset or simulator D. neural network f.
1: Input: Number of training runs K, checkpoint dataset runs {Sx) G.pt, diffusion process length J, diffusion cumulative variance schedule
loss function L, task metric. meta data store S.
Initialize: Learnable parameters for
3: for t = 1, 1,2, Nu do
2: Initialize: Learnable parameters @ for G
4: # Saampll a mini-batch of data
3: for io 1,2, Nuer do
5: (inputs, labels) ~ D
4: # SSSamlle mini-batch of data
6: Compute the predictions predictions <- (inputs)
5:
{0,,, 0rg: (4,2 Sk
7:
# Noise future parameters
8:
Compute the loss
7:
j n~ Uff,,...,
9:
loss + (prediccioos,, labels)
O2g ~ . NNVaaa022 (1 ā2)
10:
# Update the model's parameters
9:
# Compute the predictions
11:
0.41 < update(loss; 0)
10:
d. Go bt, lear les
12:
Compute the task metric
11:
# Compute the loss
13:
& + metrii(predictions, lahels)
12:
loss + -100, 00211
14:
# Save the checkpoint
13:
+Update G. pt 's parameters
15:
S S E.)
14: di+1 <- update の)
end for
15:
end for
2.2 GENERATIVE MODELS OF NEURAL NETWORK CHECKPOINTS
Using our dataset of checkpoints, we train a generative model that learns to rapidly train other neural networks. Specifically, G predicts the distribution of updated parameters Pc 10, L* E), where 0 is the starting (potentially random) neural network parameters, is the starting loss/crror/retum and l is user's s prompted loss/error/return. Conditioning on l' allows G. pt to learn from checkpoints with good and bad performance alike. In this section, we describe an instantiation of our approach based on diffusion and transformers.
2.2.1 PRE-TRAINING OBJECTIVE: DIFFUSION OF NEURAL NETWORK CHECKPOINTS
We use diffusion (Sohl-Dickstein et al., 2015) good generative modeling framework for neural network parameters since the number of forward passes required to sample novel parameter vector is set by the length of the diffusion process J as opposed to the dimensionality of the data. This instantiation of G. pt samples parameters by gradually denoising the future (updated) parameters 0
generative pre-training task. Diffusion is
Parameterization. Given an input corrupted with noise, diffusion models can be parameterized to predict either the signal or the noise (Ho et al., 2020; Nichol & Dhariwal, 2021). Prior work in the image domain has shown that noise prediction outperforms signal prediction. We find that signal prediction works better in our setting empirically, and so we parameterize G output parameters. We use fixed variances as in Ho et al. (2020).
Training. Our model takes two parameter vectors as input: a starting 0 and a noised future parameter vector 07. where denotes the timestep in the diffusion forward noising process. We minimize the simplified variational lower bound, which reduces to predicting the denoised future parameters:
LGG) = E [110 -
Algorithm 2 details our full training procedure. Note that we need tuples of data (0", ,)) to compute L. We sample these tuples from checkpoint dataset. First, sample training run uniformly at random. Then, 0 and 0* are sampled uniformly at random from the checkpoints saved within the selected training run. We enforce that 0 is always from earlier training step than d*. Note that 0 and 8" can be arbitrarily distant, even the initial and final checkpoints from a run.
Sampling. After pre-training, parameter vector 6, its loss/error/retumn and a prompted loss/crror/retum É Sampling begins by feeding-in Gaussian noise as the 6* input and gradually denoising it. We use DDPM sampling.
sample updated parameters 0" by querying G with an input
3

---

Under review as a conference paper at ICLR 2023
Predicted Future
Per-token Decoding
Diffusion Transformer
Per-token
Encoding
Starting (or Retum)
Diffusion Timestop
Retum)
Noise and Tokenize
Tokenize
Future
Starting 0
The G.pt architecture. During training, we sample two checkpoints from the same run "starting" network's parameters and "future" network's parameters from later in the run as well as their losses/errors/returns.. Each layer's parameters are flattened and linearly encoded. The future network's parameters noised via diffusion forward process prior encoding.
Figure 2:
2.2.2 ARCHITECTURE
Our generative model is transformer (Vaswani et al., 2017) that operates over parameter tokens from both 9 and 6* (Figure 2). It uses few domain-specifie inductive biases beyond tokenization.
Parameter tokenizers. Before being processed by G.pt, the two imput parameter vectors 0 and each need to be decomposed into several tokens. In general, a task-level network fo will contain many unique layers, each with potentially different number of parameters. We define the i-th token as the flattened parameter vector of the i-th layer. Layers with multiple parameter groups (e.g., layers with both a weight and a bias) are decomposed into separate tokens. Note that these tokens will usually be of different dimensionality. We call this layer-by-layer tokenization.
Parameter tokenizers for big neural networks. For larger networks, we find that it is beneficial decompose single layer's parameters into multiple tokens. We do this with layer chunking. We define a hyperparameter M, the maximum number parameters a single token can have. Layers containing more than M total parameters at-most M parameters. For example, if M 1000, a weight matrix with 10 x 768 parameters will be containing 1000 parameters and containing 680 parameters.
flattened and chunked into multiple tokens, each with
decomposed into eight tokens, We set M to be smaller than the hidden size of the Transformer to avoid lossy compression.
Metric tokenizers. We also feed the scalar input metrics é and f* (loss, error, return, etc.) and diffusion timestep j as individual tokens to the transformer. We project each scalar to vector representation using a frequency- based encoding scheme (Mildenhall et al., 2020).
Per-token encoders. After tokenizing fø layers and the input scalars, we project each token to the hidden size of the transformer. We explored more complicated encoders, but find that simple linear layer works well. Each token's encoder has unique set of weights.
Transformer. The core of the G. pt architecture is transformer which operates on the set of input parameters and metrics, linearly-encoder into tokens. Our transformer is version of GPT-2 (Radford al., 2019). We omit causal masking as our model not autoregressive across tokens.
Per-token decoders. The final layer of G.pt is decoder from the transformer's output to the future parameter vector. The i-th token linearly decoded from the transformer's 's hidden size back to the original size of the i-th layer's flattened parameter vector. Note that only the output tokens for the noised future parameter vector 0; are decoded to predictions. Our decoders do not share weights.
Global residual connection. Finally, find that it is beneficial to add residual connection (He et al., 2016) to the input at the very end of G. pt. This amounts to predicting the parameter update 8" 8 instead of directly predicting 8 itself. This residual connection also allows to initialize G to perform the identity function by initializing the decoder weights to zero. Empirically, the global residual connection in conjunction with the identity initialization significantly accelerates training.
4

---

Under review as a conference paper at ICLR 2023
500
400
Retum
100
200
100
0.0
Adamiy
Bemontum
SGD
MAg RMSSprog
Ours
Adamiy
Momonttu
SGD
NAG
RMGpraG
Ours
Abamny
Momontom
SGD
NAG
AMSpeap
Our
Cartpole
MNIST
CIFAR-10
Figure 3: G.pt optimizes unseen network parameters in one step. We compare performance after single update from pt versus single step of gradient-based optimizers. Error bars computed over five input parameter vectors, all of which are randomly-initialized
500
0.6
90
Cartpole
CIFAR-10
scn
0.5
Adam
SGD
Adam
SCD
Return
300
Loss
Error
0.4
60
0.3
100
40
10K
Steps
Steps
Steps
Figure 4: Optimization curves. We compare one step of G. pt optimization training curves produced by SGD and Adam. Error bars are computed over five initializations.
IMPLEMENTATION DETAILS
We consider both supervised and reinforcement learning tasks. In all cases, we generate a large collection of training runs in order to pre-train generative model.
Pre-training data for supervised learning. We create datasets of MNIST and CIFAR-10 network checkpoints, For MNIST, the task-level model is two-layer MLP with 10 hidden units; for CIFAR- 10, the model has two conv layers followed by global average pooling and a fully-connected layer. Both models use ReLU activations. We train the MNIST models for 25 epochs and CIFAR-10 1- models for 50 epochs, each with half-period cosine annealing. We use SGD with momentum 0.9, learning rate of 0.1 and weight decay of 5c-4. We train approximately 10K MNIST models and 55K CIFAR-I0 10 models from different random initializations. We select 200 checkpoints to save each run: the initial checkpoint (before training), the final checkpoint and intermediate checkpoints random iterations. total, this results in 2M trained MNIST MLPs and 1 11M trained CIFAR- 10 CNNs.
Pre-training data for reinforcement learning. For reinforcement learning (RL) experiments, we train policies for the Cartpole task using the IsaacGym simulator (Makoviychuk et al., 2021) I). Our policy is three-layer MI.P with 32 hidden units and SeL.U activations. We also train separate critic network with the same architecture as the policy; we only model the policy's 's parameters in our G pt experiments. We train for 500 iterations using PPO (Schulman et al., 2017) and Adam (Kingma Ba, 2015) with 31 = 0.9 and B₂ = 0.999. We train 50K models and record 200 checkpoints in each. This results in a dataset of 10M trained policies.
Model pre-training. We train G. with Adam (Loshchilov & Hutter, 2017). We maintain an exponential moving average (EMA) of G. pt weights over the course of training. Our transformer uses hidden dimension between 1536 and 2048 depending on dataset and has 12 hidden layers with 12-16 heads for self-attention. We use learned positional embeddings across all tokens, initialized to zero. We show full G. Gtpt hyperparameters in Table and dataset information in Table 2. We train one G. G.pt model per-metric, dataset and architecture (e.g., an crror-conditional MNIST MLP model).
Parameter normalization. We follow DALL-E 2's (Ramesh et al., 2022) normalization scheme, where the data is scaled such that the variance the marginal distribution matches the variance of ImageNet pixels scaled to [-1, 1], for which diffusion hyperparameters have been tuned. We find that this normalization ensures the forward noising process destroys nearly all signal in 87; the KL divergence against standard normal is roughly 10-6 bits/dim across experiments.
5

---

Under review as a conference paper at ICLR 2023
500
Achieved Return
Cartpole
2.0
90
400
Generated Ideniity
Achieved Loss
1.5
Identity
Generated
Achieved Error
20
Idernity
Generated
300
Data Bee!
Data Best
200
80
30
10
0.0
200
0.0
1.0
1.5
Prompted Return
Prompted Loss
Prompted Error
Figure 5: Achieved returns, losses and errors across a range input G. pt prompts. G.pt can train unseen neural network parameters to range of desired values in one update. Each blue curve corresponds to a different randomly-initaalzze input parameter vector. We also show the best value of each metric present in the training split of the checkpoint dataset.
EXPERIMENTS
We compare our method hand-designed optimizers and study the properties of our approach. In all experiments, we report optimization of G. pt unseen neural network parameters.
4.1 COMPARISON
OPTIMIZERS
Training in one step. Figure 3 demonstrates G. G.pt's 's ability to train unseen neural network parameters in one update. This property is unique compared gradient-based optimizers like SGD and Adam which usually require thousands, if not millions, of updates to achieve good performance. We compare against several of these traditional optimizers with tuned learning rates and weight decays'. Note that we did not systematically tune training hyperparameters for checkpoints in our dataset. For each method, we measure performance after applying one update to aandonly-iittilizeed network parameters and average results over five seeds. We prompt G.pt by setting t near the best return/loss/error in our dataset (for some tasks, asking for a value slightly above or below the best value in the dataset works better). G.pt outperforms gradient-based optimizers in this regime across tasks (control, image classification), datasets (Cartpole, MNIST, CIFAR-1O) and conditioning metrics (return, test loss, test error). Additionally, G. pt successfully optimizes non-differentiable metrics (CIFAR-10 test error) whereas baseline optimizers must use smoothed surrogates.
Training in multiple steps. We compare one step of G. pt to multiple steps of SGD and Adam in Figure 4. SGD and Adam use tuned learning rates and weight decays. The baseline optimizers require thousands of iterations match the performance one step of G.pt. With tuned hyperparameters and a sufficiently large number of updates, gradient- based optimizers supersede one-step G. optimization. Our model can also be used an iterative optimizer with recursive prompting. In this setting, we repeatedly feed G. pt's predicted f" back in as its input 0 and ask for low loss/error or high returns. Interestingly, we find that the best performance is usually realized with one-step prompting (recursive prompting usually brings only minor improvements). However, we find that recursive prompting leads to considerably better results when the input neural network comes from an out-of-distribution initialization algorithm not present in our checkpoint dataset (see Figure 6 below).
4.2 PROMPTING FOR LOSSES, ERRORS RETURNS
By prompting for various desired losses, errors, or returns, pt can sample different parameter updates that achieve a range of performance levels. Figure we show that G. pt successfully learns to generate parameters corresponding to a large range of prompted values. We pass G.pt randomly-inittalized neural network parameters and ask to optimize them one step to range of losses/errors/returns. We show results for several different starting parameters. Across different tasks and metrics, G.pt generates parameter updates that are well-correlated with the prompted value. While model is able achieve range of prompted values, note that it currently shows limited ability to extrapolate to values beyond the limits the pre-training dataset.
We perform a grid search over three learning rates (the PyTorch default and 10x ahove/helow) and three weight decay values (0. 5 10- 10 10-*) for each haseline optimizer.

---

Under review as a conference paper at ICLR 2023
D--Disrribttin Initalization
Xavier Nommal
Xovier Uniform
Orthogonal
Error
60
40
20
100
Kaiming Normal (Fan-In)
Kaiming Normal (Fan-Out)
Kaiming Uniform (Fan-in)
Kaiming Uniform [Fan-Outt
80
Error
20
2
4
5
8
9
10
5 6
9
10
2
3
6
10
0
22
3
10
Recursive Prompts
Recursive Prompts
Recursive Prompts
Recursive Prompts
Figure 6: G.pt generalizes to dnttof-distribution parameter initializations. We query G. pt with seditallyintiaalizeed weights sampled from different distribution than those in our MNIST checkpoint dataset. By recursively applying G. G.pt to its own output and prompting for low test error, we rapidly optimize out-of-distribution random initializations.
4.3 GENERALIZATION TO OUT-OF-DISTRIBUTION INITIALIZATIONS
The networks in our checkpoint dataset are initialized with single weight initialization scheme. For MNIST, they are sampled 0 ~ UI- Vn J. where is the fan-in of a layer. In Figure 6, we evaluate randomly-initialize input parameter vectors 6, where the weights
G. 6.pt's ability to generalize are sampled from different distributions (Glorot & Bengio, 2010; Saxe et al., 2013; He et al., 2015) not present in our dataset. While one step prompting performance is degraded, recursive prompting significantly improves results. G. pt is able to rapidly optimize -ittr---istriuttion weights in ten or fewer parameter updates.
4.4 SCALING MODEL AND DATA SIZE
Performance metric. We use prompt alignment measure scaling performance. We define it as the R² coefficient of determination between set of input loss/error/return prompts and the actual loss/error/return achieved by the parameters sampled from pt. We compute R² values over 20 regularly-sample prompts and average results over 128 neural networks. The optimal score is +1, which indicates that G. pt perfectly listens to loss prompts. around - 2.7. We use unseen, randomly-initializz capabilities. Empirically, we find that prompt alignment is a more reliable quality metric than diffusion mean-squared error unseen parameter vectors (see Appendix C for additional details).
handomly-initiaizeenn G.pt score input networks in order to gauge generalization
Model scale. We analyze the impact increasing the number of G. parameters in Figure (top). We train six models with transformer hidden sizes in [64, 128, 256, 512, 1024, 2048]; the smallest model is approximately 2M parameters while the largest
Prompt Alignment
0.5
-0.3
-1.0
is 858M parameters. We evaluate the G. pt checkpoint that attains the highest prompt alignment score over training. We find that larger models generalize much more effectively than smaller mod-
-1.5
-2.9
-2.5
els. Small models (<60M parameters) largely fail generalize
10"
10'
10f
10'
G.pt Size (Parameters)
to unseen parameter vectors. Even at roughly 10" parameters, we find that G. pt has not saturated its model scaling curve.
Prompt Alignment
0.5
0.0
Data scale. Next, we analyze the impact of increasing the num- ber of training checkpoints in Figure (bottom). We train our largest 858M parameter model on [500, 5K, 10K, 25K, 55K] runs, with each run containing 200 checkpoints. Performance improves substantially as the number of training checkpoints is scaled from 100K to 5M. We do not observe significant improvement when further increasing from 5M to 10M checkpoints. This may be result of G. pt requiring additional model scale to benefit from a larger pre-training dataset.
-0.5
-1.0
-1.5
-2.
10a
101
100
10'
Dataset Size (Checkpoints)
Figure 7: Scaling studies.
7

---

Under review as a conference paper at ICLR 2023
Error
30
20
Figure 8: G. learns a multimodal distribution over local error minima. We visualize the test error landscape for an MNIST MLP via parameter space PCA directions (Li al., 2018). The dots samples from G. pt when prompted for low test error; the two plots use different MLP initializations. With fixed inputs, G.pt samples diverse solutions that cover distinct positive-curvatury regions of the етror landscape. We show G. pt samples that reconstruct accurately from PCA encoding.
4.5 DIVERSITY OF GENERATED PARAMETERS
The mapping of loss values to neural network parameters is one-to-many, As generative model, G. G.pt is able sample diverse parameter solutions given loss prompt. By fixing pt inputs (including 0) and varying sampling noise, we can sample multimodal solutions that cover distinet error minima (Figure 8). Visual inspection generated first-layer weights suggests that sampling noise controls subtle variations in individual filters well the specific ordering of filters.
Intuitively, conditioning on a starting 0 should narrow the space of possible parameter solutions (in other words, initialization should have some bearing where optimization converge). Indeed, find that the most significant variation is obtained by re-sampling the starting See Appendix A for more thorough discussion and additional visualizations.
4.6 DATASET DESIGN DECISIONS
Parameter augmentation aids generalization. In the absence of permutation augmentation, we
observe that G.pt can aggressively overfit the training set and fail to generalize to new networks (i.e., it exhibits poor prompt alignment when taking unseen networks as input). Figure 11 (Appendix C) shows that training with parameter augmentation alleviates overfitting our Cartpole G. pt model.
Training on intermediate checkpoints improves one step training, Given the redundancy in neural network parameters over single training run, is worth asking there is value in training G. pt on 200 intermediate checkpoints per-run. Instead, we could train G pt exclusively on the initial and final checkpoint from each run. We find that this setup degrades one step training capabilities by over 50%: average test loss when prompting with f = worsens from 0.2 to 0.32. G. pt significantly benefits from training on a large number checkpoints, even those from the same run.
RELATED WORK
5.1 PRE-TRAINING FROM LARGE-SCALE DATA
Transformers for X. Transformers (Vaswani et al., 2017) were initially developed for language but have been shown to well-suited for a wide range of domains. They have achieved strong results in vision (Dosovitskiy et al., 2020), language modeling (Radford et al., 2018; 2019; Brown et al., 2020), coding (Chen et al., 2021b; Li et al., 2022). reinforcement learning (Chen et al., 2021a; Janner et al., 2021), image synthesis (Esser et al., 2020; Ramesh al., 2021; Yu et al., 2022) and protein folding (Jumper et al., 2021). Likewise, we show that transformers can be used for learning to learn by generative pre-training from neural network parameters.
8

---

Under review as a conference paper at ICLR 2023
Diffusion. Diffusion models (Sohl-Dicksted et al., 2015) have recently been shown to be highly effective for images (Ho et al., 2020; Nichol & Dhariwal, 2021; Dhariwal & Nichol, 2021; Ho et al., 2021; Ramesh et al., 2022; Saharia et al., 2022: Song & Ermon, 2019). In this paper, we show that diffusion models can be used for meta-learning by generating neural network parameters.
Pre-training. Large scale pre-training has led to significant advances in vision (Krizhevsky et al., 2012; Girshick et al., 2014), natural language processing (Devlin et al., 2019; Radford et al., 2018; 2019; Brown et al., 2020) and audio understanding (Van Den Oord al., 2016; Dhariwal al., 2020). We explore pre-training from datasets of neural networks instead of datasets of images and text.
Datasets of neural networks. Past works have constructed datasets of neural networks and used them in various settings: analyzing population-leve trends (Radosavovic et al., 2019; 2020). benchmarking neural architecture search (Ying et al., 2019). training hypernetworks (Knyazev al., 2021), predict- ing model properties (Schürholt et al., 2021) and dataset distillation (Wang et al., 2018; Cazenavette et al., 2022). We share the goal of using datasets of neural networks, but for the novel meta-learning approach of pre-training generative model from trained neural network checkpoints.
5.2 LEARNING TO LEARN
Learning optimizers. Past works have explored parameterizing optimization update rules with neural networks in place of hand-designed rules like Adam. These rules can be parameterized implicitly as neural networks that take gradients as input and output an improved parameter update. They are typically trained with unrolled optimization (Hochreiter et al., 2001; Younger et al., 2001; Gregor & LeCun, 2010; Andrychowicz et al., 60 Ravi & Larochelle, 2017; Wichrowska et al., 2017; Lv al., 2017; Metz et al., 2019; 2022) or reinforcement learning (Li & Malik, 2016: 2017).
Hypernetworks. Rather than parameterizing update rules, neural networks can be used to directly output or modify other neural networks' parameters (Schmidhuber, 1992; 1993; Hinton & Plaut, 1987). For example. hypernetworks (Ha et al., 2017) train parameter regression networks end-to- end with the task objective. Hypemetworks have subsequently been extended to support sampling different parameter solutions (Krueger et al., 2017; Deutsch et al., 2019; Ratzlaff & Fuxin, 2019).
Model-agnostic meta-learning. MAMI learns a parameter initialization that is rapidly adaptable to new tasks (Finn et al., 2017). Subsequent work has built simple probabilistic models over learned MAML initializations (Finn et al., 2018). These methods possess similar characteristics as Icarned optimizers- they rely on unrolled optimization and require differrntiable task-level objectives.
Learning hyperparameters. A large body of prior work has explored learning hyperparameters of standard optimizers (Bergstra et al., 2011; Feurer & Hutter, 2019). For example, learning rates, weight decays and weight initializations can all be learned via hypergradient descent (Maclaurin et al., 2015; Baydin al., 2017: Drucker & Le Cun, 1992), Bayesian optimization (Snoek et al., 2012) and reinforcement learning (Daniel et al., 2016; Xu et al., 2017; 2019; Almeida et al., 2021).
Learning to learn as pre-training. In contrast to learned optimizers, hypemneworks and MAML, G.pt pre-trains from vast amounts trained neural network checkpoints. Our method does not backpropagate through task-level losses and, as a result, does not require the task metric being optimized for to be differeentiabe. This allow us to train with standard generative modeling techniques instead of reinforcement learning or unrolled optimization which can be unstable (Metz et al., 2021).
DISCUSSION
Limitations. The current instantiation of our method has several limitations. First, the model sometimes exhibits signs of underfitting the full loss/error landscape, such as with CIFAR-10. Second, our current G. pt models struggle to extrapolate to losses and errors not present in the pre-training data. Third, this paper only pre-trains from imgle-architecture and single-task data. Finally, this paper considers relatively simple datasets of neural networks with static optimizer hyperparameters.
Conclusion. We propose generative pre-training from neural network checkpoints. We show that our approach enables rapid optimization of neural networks across tasks (supervised and reinforcement learning) and metrics (losses, errors, returns). Learning algorithms designed by humans have led to large advancements across different areas of artificial intelligence. We hope that our work serves as a step towards learning learning algorithms from data using modern deep learning techniques.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. A generalizable approach to learning optimizers. arXiv:2 100.00988 2021.
Marcin Andrychowicz, Misha Denil, Sergio Gómez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas. Learning learn by gradient descent by gradient descent. In NeurlPS, 2016.
Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv:1703.04788 2017.
Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning synaptic learning rule. In IJCNN, 1991.
James Bergstra, Rémi Bardenet, Yoshua Bengio, and Balázs Kégl. Algorithms for hyper-parameter optimization. NeurIPS, 2011.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A. Efros, and Jun-Yan Zhu. Dataset distillation by matching training trajectories. In CVPR, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In NeurIPS, 2021a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv:2107.03374 2021b.
Christian Daniel, Jonathan Taylor, and Sebastian Nowozin. Learning step size controllers for robust neural network training. 2016.
Lior Deutsch, Erik Nijkamp, and Yu Yang. A generative model for sampling sigh-performance and diverse weights for neural networks. arXiv preprint arXiv:1905.02880 2019.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HCT, 2019.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In NeurlPS, 2021.
Prafulla Dhariwal, Heewoo Jun, Christine Payne, Jong Wook Kim, Alec Radford, and Ilya Sutskever. Jukebox: A generative model for music. 20ki::2000.00:444 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x words: Transformers for image recognition scale. In ICLR, 2020.
Harris Drucker and Yann Le Cun. Improving generalization performance using double backpropaga- tion. IEEE Transactions on Neural Networks, 1992.
Patrick Esser, Robin Rombach, and Björn Ommer. Taming transformers for high-resolution image synthesis, 2020.
Matthias Feurer and Frank Hutter. Hyperparameter optimization. In Automated machine learning. 2019.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-leame meta-learnigg for fast adaptation of deep networks. In ICML. 2017.
10

---

Under review as a conference paper at ICLR 2023
Chelsea Finn, Kelvin Xu, and Sergey Levine. Probabilistic model-agnostic meta-learning. In NeurIPS, 2018.
Ross Girshick, Jeff Donahue, Trevor Darrell, and Jitendra Malik. Rich feature hierarchies for accurate object detection and semantic segmentation. In CVPR, 2014.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In AISTATS, 2010.
Karol Gregor and Yann LeCun. Learning fast approximations of sparse coding. In ICML, 2010.
David Ha, Andrew M. Dai, and Quoc V. Le. Hypernetworks. In ICLR, 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In ICCV, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
Geoffrey E. Hinton and David C. Plaut. Using fast weights to deblur old memories, In In Proceedings of the 9th Annual Conference of the Cognitive Science Society, 1987.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. NeurlPS, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. arXiv:2106. 10:::222001222, 2021.
Sepp Hochreiter, A Steven Younger, and Peter R Conwell. Learning to learn using gradient descent. In ICANN, 2001.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In NeurIPS, 2021.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero-Soriano, Parameter prediction for unseen deep architectures. NeurIPS, 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In NeurlPS, 2012.
David Krueger, Chin-Wei Huang, Riashat Islam, Ryan Turner, Alexandre Lacoste, and Aaron Courville. Bayesian hypernetworks. arXiv:1710.04759, 2017.
Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In NeurIPS, 2018.
Ke Li and Jitendra Malik. Learning to optimize. arXiv:1606.0/888 2016.
Ke Li and Jitendra Malik. Learning to optimize neural nets.
2017.
Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-level code generation with alphacode. S.Si::2003008811 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arX1v:1711.0510:1 2017.
Kaifeng Lv, Shunhua Jiang, and Jian Li. Learning gradient descent: Better generalization and longer horizons. In ICML, 2017.
11

---

Under review as a conference paper at ICLR 2023
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-basse based hyperparameter optimization through reversible learning. In ICML, 2015.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, al. Isaac gym: High performance gpu-based physics simulation for robot learning. In NeurlPS, 2021.
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In ICML, 2019.
Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv:2111.05800 2021.
Luke Metz, Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl- Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. vrXiv:2203.11860, 2022.
Ben Mildenhall, Pratul P. Srinivasan, Matthew Tancik, Jonathan T. Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. In ECCV, 2020.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In ICML, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. 2019.
Ilija Radosavovic, Justin Johnson, Saining Xie, Wan-Ye Lo, and Piotr Dollár. On network design spaces for visual recognition. In ICCV, 2019.
Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár. Designing network design spaces. In CVPR, 2020.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In ICML, 2021.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv:2204.6//25, 2022.
Neale Ratzlaff and Li Fuxin. Hypergan: A generative model for diverse, performant neural networks. In ICML, 2019.
Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. In ICLR, 2017.
Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar-10 classifiers generalize to cifar-10? 2018.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, 1951.
Geoffrey Roeder, Luke Metz, and Durk Kingma. On linear identifiability of learned representations. In ICML, 2021.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. arXiv:2205.1 / 2022.
Andrew M Saxe, James I. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv: 2013.
Jurgen Schmidhuber. Evolutionary principles in self-referential learning. 1987.
12

---

Under review as a conference paper at ICLR 2023
Jürgen Schmidhuber. Learning to control fast-weight memories: An alternatve to dynamic recurrent networks. Neural Computation, 1992.
Jürgen Schmidhuber. A 'self-referential weight matrix. In ICANN, 1993.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. (eXiv:/707.6635 2017.
Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-supervised representation learning neural network weights for model characteristic prediction. NeurIPS, 2021.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. NeurlPS, 2012.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In (CML,2015,
Jiaming Song, Chenlin Meng, and Stefano Ermon. 2IXi0:2010002502, 2020.
Denoising diffusion implicit models.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NeurlPS, 2019.
Aäron Van Den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew W Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. SSW, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Tongzhou Wang, Jun- Yan Zhu, Antonio Torralba, and Alexei A Efros. Dataset distillation. 1089v99811.10959, 2018.
Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffman, Sergio Gómez Colmenarejo. Misha Denil, Nando de Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In ICML, 2017.
Chang Xu, Tao Qin, Gang Wang, and Tie-Yan Liu. Reinforcement learning for learning rate control. arXiv: 11159, 2017.
Zhen Xu, Andrew M Dai, Jonas Kemp, and Luke Metz. Learning adaptive learning rate schedule. 190.00799. 2019.
Chris Ying, Aaron Klein, Eric Christiansen, Esteban Real, Kevin Murphy, and Frank Hutter. Nas- bench-101. Towards reproducible neural architecture search. In ICML, 2019.
Steven Younger, Sepp Hochreiter, and Peter R Conwell. Meta-learning with backpropagation. In IJCNN, 2001.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Ziru Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, et al. Scaling autoregressive models for content- rich text-to-image generation. 2rXiv: 10789, 2022.
13

---

Under review as a conference paper at ICLR 2023
Generated Parameters
Real Parameters
Figure 9: Visualizing synthetic parameters generated by G. pt versus real parameters. Each row visualizes the first layer weights of single MNIST MLP with 10 hidden neurons. The j-th column shows the x 28 weights incoming to the j-th hidden neuron. Left: We show G. pt samples when querying for low error. In each row, we feed different unseen, randomly- indimlly-nitiaalizedd input network to G. G.pt. Right: The "ground truth" parameters obtained by training the same andomly-initialized nets with gradient-based optimization instead.
APPENDIX A: ANALYZING MULTIMODALITY
Random noise in the reverse diffusion process and random initialization of the initial network parameters provide two different sources of multimodality in the output parameters for a given target loss/error/return. In our main paper, and also in Supplement Figure 12, we see that pt is capable of producing multimodal predictions given fixed inputs (the input parameters A, its loss/error/retum & and the desired loss/error/retum E*). this section, we show that more variation in G.pt is captured by re-sampling input 0. This result suggest that G.p is doing more than memorizing single network for every input loss/error/return.
Motivation. Let's begin by considering the joint distribution over all inputs to G. G.pt:
p(0' 10,00,E" 0) ( (0010,, p(0 p((1)))) p(() P0 e) p((()) p((7le) p(e).
The above follows directly from the probability chain rule and the assumption that 0 is conditionally- independent of E" given Writing the joint distribution in this way makes clear that there are least two ways we may obtain variation in generative procedure with fixed inputs et and é: we can re-sample 0* ~ PG or re-sample ~ p(O(e). Sampling from the latter distribution somewhat complicated by the fact that, for general É, we do not have model of p(0)2)². However, for the common case where the imput = Lrandmm corresponds to the performance of a random neural network prior to training, with whatever initialization algorithm we are using- e.g. Kaiming (He et al., 2015) or Xavier (Glorot & Bengio, 2010) initialization.
draw reasonable samples from p(0 (raddom)) by instantiating
networks
In short, given inputs L" and f, we can produce diverse G. pt samples in two ways. First, by re- sampling the noise part of the reverse diffusion process; second-if & is the performance of random net-by re-sampling random initializations of the input 0.
Assessing variation. To disentangle the two sources of variation, use DDIM sampling (Song et al.. 2020) in the reverse process with = O. This sampling scheme forces G pt to produce deterministic mapping of noise to parameters. Hence, by varying input 0 but holding reverse process noise constant, we can now assess how much variation in G pt's 's distribution over 0" is controlled by We evaluate the degree of variability induced by the two sources in our MNIST error model:
IIn principle, we could train second diffusion model learn this distribution. We leave this future work.
14

---

Under review as a conference paper at ICLR 2023
Varying Input Parameters
Variance
Varying DDIM Noise
6a2
Parameters of
Figure 10: Varying input yields significant variation in synthesized 0* We compute the per- parameter variance of 0 as we re-sample the input & vector and hold diffusion reverse process noise fixed (blue), And, we compute the per-parameter variance as we re-sample reverse process noise and hold the input 0 fixed (yellow). The 750 parameters with highest O-variance are visualized. Input appear to affect G. pt's predictions more strongly than reverse process noise.
1. Diffusion noise re-sampling: Ee-p(ole) (1ans-Wo.,
2. Input re-sampling: (Z~A((.))) Vare-ptele)
Where lr Pa (f, f), and both the expectation and variance are taken elementwise per- parameter. We the inputs ! and Both the expectations and variances are evaluated with 64 samples, requiring total of 64 64 = 4096 samples. The results are in Figure 10-indeed, 0 appears have a stronger effect on G. G.pt's output than reverse process noise. This result suggests that G. does not merely memorize a single solution for every corresponding target loss/error/return e"; its predictions appear to strongly depend on the input 0.
APPENDIX B: ADDITIONAL RESULTS
Parameter visualizations. In Figure 9, we visualize G. pt's predicted weights for MNIST MLPs and compare them to real weights from our checkpoint dataset. Varying both input and DDPM sampling noise yields visually noticeable variation predicted 6*.
Error surface visualizations. In Figure 12, we show nine additional visualizations of the MNIST test error surface and where G.pt samples lie on it. Note that the results are uncurated. All of the plots suggest that G. pt can sample multimodal solutions given fixed inputs 0, l and E.
APPENDIX C: PRE-TRAINING CURVES AND HYPERPARAMETERS
8.030
81025
Test
Frompt Aligmmen
Trate
Teat
8000
vithing
8.6000
0.0000
12Ck
103C
43c
20C
335C
20K
1256
2054
<0K
130K
10K
Training Steps
Training Steps
Taainin Steps
Traiiigg Steps.
Figure 11: G.pt training curves with and without parameter augmentation. Left: We show G.pt's loss the course of training different splits from our Cartpole checkpoint dataset. Right: Prompt alignment scores over the course of training (+1 indicates perfect alignment).
Training curves. In Figure 11, we plot various G. G. evaluation metrics over the course of training. Parameter augmentation increases training stability and yields significantly improved generalization as measured by prompt alignment on unseen neural networks in the test set. Interestingly. we sometimes observe that G. pt 's prompt alignment the test set improves while its test loss worsens. Similar trends have been found in autoregressive generative models; e.g., AlphaCode's (Li et al., 2022) solve rate for unseen coding problems improves its validation loss increases.
15

---

Under review as a conference paper at ICLR 2023
Error
10
20
30
40
50
Figure 12: Uncurated error surface plots. All nine surface plots visualize the error landscape of an MNIST MLP. Different plots correspond to different random initializations of the MLP that are input to pt. The dots are samples produced by G. pt when prompted for low error. We visualize the surface by moving in the top-two PCA directions computed over 256 G pt samples. We show G.pt samples that reconstruct accurately from two-dimensional PCA encoding (less than 10-3 reconstruction distance in parameter space). The plots suggest that G pt has learned a multimodal distribution over parameters, even when given fixed inputs.
16

---

Under review as a conference paper at ICLR 2023
G.pt Hyperparameters
Caripole
MNIST (Loss)
MINSTT (Error)
CIFAR-IO (Loss)
CIFAR-10 (Error)
Diffusion steps (J)
1000
1000
1000
Noise schedule
linear
linear
linear
lincar
linear
Transformer layers
12
12
12
12
Transformer dim
1536
1536
1536
2048
2048
Self-attention heads
12
16
16
16
16
Model parameters
347M
378M
378M
639M
639M
Layer tokenizer
layer-hy-layer
layer chunk
layer layer chunk
layer chunk
layer chunk
Max parameters per token (M)
1000
1000
576
Scalar encoder: number of frequencies
128
128
128
128
128
Scalar encoder: max frequency (logg)
14
14
14
Data scale factor
2.06
4.185
4. 4.185
1.646
1.646
AdamW 32
0.999
0.999
0.95
0.999
0,999
Base learning rate
x
10-4 lincar
10-4 linear
10-4
10-4
Learning rute warmup
linear
linear
lincar
Learning rate decay
cosine
cosine
cosine
cosine
cosine
Weight decay
0.1
0.1
0.1
0.1
0.1
Batch size
8192
1024
512
550
550
Gradient clip
0.75
0.75
0.1
EMA decay
0,9999
0.9999
0,9999
0,9999
0.9999
Training iterations best performance
128K
110K
288K
425K
Prompted E" for one-step optimization
500.0
0
5.0
1.2
35,0
Table I: G. pt Hyperparameters. When prompting G. with t for one-step optimization,
choose a value close to the smallest loss/error or highest return present in our checkpoint dataset.
MNIST
CIFAR-10
CARTPOLE
#RUNS
10728
56840
50026
#TRAIN RUNS
10228
54790
47976
#TEST RUNS
500
2050
2050
CHECKPOINTS/RUN
200
200
200
#CHECKPOINTS
2.1M
11.3M
10M
ARCHITECTURE
MLP
CNN
MLP
#PARAMETERS
7960
5370
1250
Table 2: Checkpoint dataset statistics.
APPENDIX D: OVERFITTING TEST SET METRICS
2.50
Achieved Loss
2.25
2.00
1.75
1.50
CIFAR-10
11..
1.00
1.2
1.6
1.8
2.0
Prompted Loss
Figure 13: G. pt pre-trained on CIFAR-10 test losses generalize to CIFAR-10.1. We examine if G.pt overfits to CIFAR-10 test losses by evaluating its generated networks on CIFAR-10.1. Prompting for smaller CIFAR-10 test losses yields networks that perform better on CIFAR-10.1.
In our supervised learning experiments, we condition G-pt on test set metrics, like losses and errors. A reasonable concern is that conditioning on test losses might cause G. to overfit to the task's test set. To investigate this, we evaluate networks generated by our test loss-conditional CIFAR- 10 model on the 10. dataset (Recht et al., 2018). Results are in Figure 13. Asking for smaller CIFAR-10 test loss yields a smaller CIFAR- 10.1 loss, indicating G. pt is not overfitting. As with past works that evaluate explained by the observation that checkpoints in our dataset also have worse losses on 10.1 (the best checkpoint in our dataset has CIFAR-10 10 test loss of 1.07 versus CIFAR- 10.1 loss of 1.51).
CIFAR-10.1, generated networks perform worse 10.1 than 10. This is
17

---

Under review as a conference paper at ICLR 2023
APPENDIX E: MEMORIZATION VERSUS GENERALIZATION
Cartpole
MNIST
MNIST
100
NN Score (%6
80
400
40
Prompted Return
Prompted Loss
Prompted Error
CIFAR-10
CIFAR-10
NN Score (%)
1.2
1.4
1.8 2.0
2.2
50
70
90
Prompted Loss
Prompted Error
Figure 14: G.pt predictions on held-out (unseen) random initializations tend to lie closer to the ground truth outcome SGD/Adam than any parameter vector from our checkpoint dataset's training split. For each test run in our dataset, we feed the initial parameters and a metric prompt G. stpt, and we sample a prediction. We count the percentage of runs for which the prediction closer to one of the 200 checkpoints in that same test run than all checkpoints in the training split (Cartpole has 1OM training split checkpoints, CIFAR-10 has 1..3 and MNIST has M.) Each plot corresponds to different G.pt model, and repeat the test for wide range of prompts.
In this section, investigate the extent which Gt memorizes solutions from the training set. This is challenging topic to address for any generative model, and there is no siiversally-acceptee methodology to measure it. For generative models of images, popular methodology is to visualize the nearest neighbors of generated images in the training set. Visualizing parameters is challenging for deep networks beyond the first layer, we instead provide basic way quantify memorization.
Experimental setup. Our approach also based on nearest neighbors. We feed pt an unseen, -initaal-nnitiaaizzed parameter vector from a test run and sample a corresponding solution 0* from our model. If G. pt is memorizing parameters from the training set, then the sampled 0 should be closer one of the millions of parameter vectors across all runs in the training split than the 200 "ground truth" parameter vectors in the same test run from which we took the randomly -nnitaiiiee input parameters. On the other hand, if 6* is closer one of these 200 held-out checkpoints, suggests that it is accurately predicting the outcome of gradient- based optimization (i.e., there is some level of meaningful generalization). For simplicity, we compute distances in Euclidean space. We count the percentage of test runs for which G pt generates solution closer any of the 200 checkpoints in the test run than all checkpoints in the training split our dataset (Cartpole has 10M training split checkpoints, CIFAR- 10 has 11 1.33M and MNIST has 2. M).. We call this percentage the nearest neighbor score. A score of 100% suggests G. pt is perfectly generalizing. We repeat this test for range of loss, eтror and return prompts.
Figure 14 shows nearest neighbor scores for all five of our G. G.pt models. Our models appear to accurately generalize under large number of loss, error and return prompts. Our Cartpole and CIFAR-10 10 models exhibit perfect scores (100%) ) for all input prompts. Interestingly, while our MNIST models also have perfect scores for the majority of loss/error prompts, they have lower scores for smaller prompts (decreasing to about 15-20%). A speculative explanation is that our MNIST models were trained with about fifth of the number of training runs compared to our Cartpole and CIFAR- 10 models; this could possibly degrade generalization capabilities. Overall, this test provides some initial evidence that G. G.pt is generalizing and not just memorizing training set parameters.
Results.
18


Reviews:
[{"review_id": "y5oMYd30R8", "reviewer": "Reviewer_rREn", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposed a generative pre-training approach, termed G.pt, that learns to optimize the parameters of neural networks. Specifically, this approach first constructs a dataset of network checkpoints (along with their corresponding metrics) sampled from a hundred thousand training runs, and then trains a conditional diffusion model on this dataset to learn the distribution of network parameters that fulfill the desired metric. After training, given the target metric and the randomly initialized (or unseen) parameters, we can find the optimized network parameters by directly sampling from the conditional diffusion model. Compared to the conventional SGD type optimization methods that relies on at least thousands of parameter updates, the proposed method is claimed to perform well with “a single parameter update”.\n", "summary_of_the_review": "Overall, I liked the idea but I have major concerns about the practical significance of the proposed method. \n", "strength_and_weaknesses": "Strengths:\n- This work is clearly written, and the presentations of the proposed idea and experiment settings are overall very easy to follow. \n- The idea of training a conditional diffusion model in the (parameter, metric)-space to perform the learning to learn tasks looks novel to me.\n- Experiments on a fairly reasonable range of settings (MNIST, CIFAR-10, and Cartpole) demonstrated that the proposed method can learn the distribution of network parameters conditioned on the desired metric (note that its value cannot surpass the best from the dataset). Also, experiments showed some nice properties of the proposed method, including scaling the model and data size improves the performance, measured by the prompt alignment score, and sampling diverse outputs for the same metric.\n\nWeaknesses:\n- Although I liked the idea of generative pre-training that learns to optimize the parameters, given the claimed properties and current experimental results, my major concern is whether the proposed method is really useful for solving downstream problems. For instance, if the proposed method cannot extrapolate to the metric values beyond the limits of the checkpoint dataset, what extra value does it add on top of the pre-trained checkpoints? Because in practice, especially the classification and RL tasks, we mostly just care about when and how we achieve the better scores (e.g., accuracies and rewards). The claimed favorable properties also seem not add significant value in this regard. Instead of converging to a better local minima, why do we care about how the network ends up with multiple given “mediocre” local minima from unseen initializations? Can we find some use cases to show the practical significance of the proposed method?\n- I realized the authors emphasized many times about “from unseen neural network parameters”. I’m a little confused: Does it just mean we initialize the network parameters that are not from the training set?\n- Another confusing part is that the proposed method is claimed to optimize the parameters “in just one update”. I’m not sure what the one update exactly means. If we use the DDPM sampling, which is an iterative process with usually hundreds of function evaluations, how can we say “one update”? When comparing with the conventional SGD methods, I think we should consider the DDPM sampling steps.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "YKhIFZreINF", "reviewer": "Reviewer_bACh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "First, the authors built a dataset consisting of neural network checkpoints that perform specific tasks, such as MNIST classification, including its loss and error. The proposed data-driven optimizer, G.pt, is a neural network model trained using this dataset. Specifically, if the parameters of the initial checkpoint, the initial loss, and the target loss are given as inputs, then G.pt outputs a new parameter set that can achieve the target loss. Finally, by using the MNIST, CIFAR-10, and Cartpole datasets with a small architecture (e.g., two-layer MLP with 10 hidden units), the authors validated the effect of G.pt compared to an existing gradient-based optimizer (e.g., Adam optimizer).", "summary_of_the_review": "The ideas and research topics of the paper are interesting, and the experimental results are also impressive. However, there are obvious limitations in terms of applicability, so I believe that further discussion is necessary in this regard.", "strength_and_weaknesses": "Strength\n1. I am not sure whether this research topic is promising, but it is definitely interesting. Also, I believe this paper can serve as the first step towards this research direction.\n2. The proposed method is straightforward, and the experimental results are surprising.\n3. Unlike the existing gradient-based optimization algorithms, G.pt can optimize non-differentiable objectives, such as classification errors.\n\nWeakness\n1. The proposed method can be used only for small models, and such a model has limitations in minimizing generalization errors. Therefore, when deploying a deep learning model, I do not know when it is better to use G.pt instead of an existing optimizer. Could the authors provide an additional discussion about when the proposed optimizer is beneficial compared to existing optimizers?\n2. It is interesting that G.pt can easily find multiple parameter sets with the same loss, but there is no discussion of why they are necessary.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "XOQon8_R5sO", "reviewer": "Reviewer_Wayk", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper uses a diffusion Transformer to map a pair (initial parameter, target loss) to a distribution over parameters that achieve the specified performance. The model is trained with a pre-training dataset of neural network checkpoints.", "summary_of_the_review": "My main concern with this paper is that I don't see a practical benefit for the method (see \"Strength and Weaknesses\" above). I do think the idea is creative and interesting.", "strength_and_weaknesses": "Pros\n- The overall idea is interesting: posing performance-conditional parameters as a conditional generative modeling problem lets us leverage advances in diffusion models and is an interesting alternative to e.g. learned optimizers.\n\nCons\n- There doesn’t seem to be any practical benefit for the method, because your pre-training dataset already includes a lot of checkpoints for your specific setting, and your experiment in figure 5 shows that the method does not extrapolate to unseen losses (no runs improve upon “Data Best”). For raw performance, the result of running G.pt has no benefit over simply using the best checkpoint in the dataset.\n- I may be misunderstanding the method, but is anything explicitly preventing this model from simply memorizing checkpoints?\n\nThis isn’t a request for additional experiments, but I think evidence of generalization across network architecture or dataset would address my concerns about the method.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "NPKxDrGLT1y", "reviewer": "Reviewer_iYN2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work trains a generative model (conditional diffusion Transformer) called G.pt over NN checkpoints of supervised + reinforcement learning (RL) tasks such that given a prompt of (initial input parameter vector, target loss/error/return), the generative model outputs an updated parameter vector that achieves the desired metric. The paper explores classification on MNIST and CIFAR-10 for MLPs and CNNs, as well as an MLP on a Cartpole task. This approach performs favorably to existing optimizers – G.pt acts as a learned optimizer such that initially, each step of G.pt corresponds to thousands of gradient update steps of traditional optimizers such as Adam and SGD.", "summary_of_the_review": "The authors introduce G.pt, a generative model of NN checkpoints as a way to efficiently learn how to optimize NNs to target a desired loss/return value. They propose how to make this idea work well in practice with different approaches of tokenizing each NN layer, encoding the loss/model parameters, performing data augmentation, adapting the underlying diffusion model, and outlining the pretraining procedure. They then empirically explore G.pt’s performance on MNIST and CIFAR-10 classification as well as learning a policy to perform Cartpole, and also demonstrate that the model is not straightforwardly memorizing models it has seen before. The submission is high quality, interesting, and novel. I anticipate this being of broad interest to the community, and therefore recommend acceptance.", "strength_and_weaknesses": "Strengths:\n- The idea is really interesting and the paper was fun to read! \n- There were a lot of empirical design decisions that had to be made in order to make the approach work (e.g. how to tokenize each layer of the NN, how to modify the conditional Transformer-based diffusion model, permutation augmentation in weight space to facilitate generative pre-training, etc.) and the paper did a good job outlining each step, as well as justifying the reasons behind each choice.\n- The results are impressive. Though one could argue that the tasks are relatively simple, I thought the paper did a great proof of concept that G.pt is a promising approach for learning to learn.\n\nWeaknesses:\n- The authors did a good job addressing the limitations of the work in Section 6, where they outlined how G.pt does not quite handle extrapolation of loss/error values not present in the training data very well.\n- This is not quite a weakness per se, but I thought Appendix A was interesting in terms of constructing an alternate view of the generative process of the joint over (model parameters, losses). It would’ve been helpful for me for the authors to elaborate upon this discussion point a bit more, either in the main text or in the Appendix. Intuitively, it seems as though the dependence between the loss (\\ell) and model parameters (\\theta) should go in the opposite direction (e.g. p(\\ell \\vert \\theta), with a prior over model parameters p(\\theta)) rather than the other way around, since the loss is a function of the model parameters and the input. Things are a bit different here since the diffusion model is being \\emph{prompted} with fixed loss values and initial model parameters. Would the authors comment on this?\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject