[论文解读] Learning to Learn with Generative Models of Neural Network Checkpoints

William Peebles, Ilija Radosavovic+3 more|arXiv (Cornell University)|Sep 26, 2022

Model Reduction and Neural Networks被引用 7

一句话总结

本文提出 G.pt，一种生成式模型，通过在包含 2300 万个神经网络检查点的大规模数据集上进行训练，学习优化神经网络。该模型采用条件扩散 Transformer 架构，能够以单步预测方式为未见过的网络架构和任务生成最优参数更新，实现损失、误差或回报等目标指标，而无需反向传播或展开优化。

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.

研究动机与目标

• 解决基于梯度的优化器无法从过往经验中改进的局限性。
• 克服通过展开优化或强化学习训练学习优化器时的不稳定性和复杂性。
• 利用现有海量神经网络检查点作为元学习优化策略的数据源。
• 开发一种可扩展、数据驱动的方法，实现对多样化架构与任务的神经网络一步优化。
• 实现对不可微分目标（如分类误差或强化学习回报）的有效优化。

提出的方法

• 从超过 10 万个监督学习（MNIST、CIFAR-10）和强化学习（CartPole）任务的训练运行中，构建包含 2300 万个神经网络检查点的数据集。
• 收集包含参数和任务级指标（测试损失、误差、回报）的检查点，使用 Adam 和 SGD 等标准优化器。
• 应用参数级别的数据增强（如权重置换）以提升数据集多样性，同时保持模型功能不变。
• 训练一个条件扩散 Transformer 模型（G.pt），在给定初始参数向量和目标指标提示（如期望损失）时，预测最优参数更新的分布。
• 采用标准生成建模技术，而非展开优化或强化学习，从而实现稳定且可扩展的训练。
• 在推理阶段，通过 G.pt 的单次前向传播，为新的、未见过的初始化和架构生成优化后的参数。

实验结果

研究问题

• RQ1 在大规模神经网络检查点数据集上进行训练的生成式模型，是否能够学习在单步内生成有效的参数更新？
• RQ2 该模型在面对分布外的权重初始化方式和未见过的架构时，其泛化能力如何？
• RQ3 该模型能否为同一优化提示生成多模态解，反映通往相同指标的多样化路径？
• RQ4 该方法是否对不可微分目标（如分类误差或强化学习回报）也有效？
• RQ5 所学习的优化策略性能在不同网络架构和任务上如何随规模扩展？

主要发现

• G.pt 在多种架构和任务上，对未见过的初始化实现了单次参数更新的神经网络优化。
• 该模型能泛化到分布外的权重初始化方法，表现出超越训练分布的鲁棒性。
• G.pt 可为同一优化提示生成多模态解，表明其学习到了实现最优性能的多样化路径。
• 该模型展现出有利的缩放特性，在不同网络深度和宽度下均保持性能稳定。
• G.pt 有效优化了不可微分目标（如分类误差和强化学习回报），而传统学习优化器则无法做到。
• 尽管在 CIFAR-10 等复杂数据集上存在一定程度的欠拟合，但模型在简单任务上表现优异，并展现出对新架构的良好泛化能力。