[論文レビュー] 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

ひとこと要約

この論文では、2300万件のニューラルネットワークのチェックポイントを含む大規模データセットを用いて学習することで、バックプロパゲーションやアンロールド最適化を経ずに、損失、誤差、または報酬などの望ましい指標に到達するための最適なパラメータ更新を1ステップで予測できる生成モデルG.ptを提案する。

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.

研究の動機と目的

• 勾配ベースの最適化手法が過去の経験から改善できないという制限を解決すること。
• アンロールド最適化や強化学習による学習最適化の不安定さと複雑さを克服すること。
• 既存のニューラルネットワークのチェックポイントの豊富な蓄積を、メタラーニングによる最適化戦略のデータソースとして活用すること。
• 多様なアーキテクチャとタスクにわたるニューラルネットワークの1ステップ最適化のスケーラブルでデータ駆動型の手法を開発すること。
• 分類誤差や強化学習の報酬など、微分不可能な目的関数の最適化を可能にすること。

提案手法

• 教師あり（MNIST, CIFAR-10）および強化学習（CartPole）タスクにおける10万件以上のトレーニング実行から、2300万件のニューラルネットワークのチェックポイントのデータセットを構築する。
• 標準的な最適化手法（Adam, SGD など）を用いて、パラメータとタスクレベルの指標（テスト損失、誤差、報酬）を含むチェックポイントを収集する。
• 重みの順列変更などのパラメータレベルのデータオーグメンテーションを適用し、モデルの関数を保持したままデータセットの多様性を向上させる。
• 初期パラメータベクトルとターゲット指標プロンプト（例：望ましい損失）を入力として受け取り、最適なパラメータ更新の確率分布を予測する条件付き拡散トランスフォーマー（G.pt）を訓練する。
• アンロールド最適化や強化学習ではなく、標準的な生成モデリング技術を用いることで、安定的かつスケーラブルな訓練を可能にする。
• 推論時において、G.ptの1回の順伝播によって、新しい未学習の初期化やアーキテクチャに対応する最適化済みパラメータを生成する。

実験結果

リサーチクエスチョン

• RQ1 大規模なニューラルネットワークのチェックポイントデータセットを学習データとして用いた生成モデルは、1ステップで効果的なパラメータ更新を生成できるか？
• RQ2 このようなモデルは、分布外の重み初期化や未学習のアーキテクチャへも一般化できるか？
• RQ3 同じ最適化プロンプトに対して、多様な最適化経路を反映するマルチモーダルな解を生成できるか？
• RQ4 分類誤差や強化学習の報酬など、微分不可能な目的関数に対しても、この手法は効果的に機能するか？
• RQ5 異なるネットワークアーキテクチャとタスクにおいて、学習最適化戦略の性能はどのようにスケーリングするか？

主な発見

• G.ptは、複数のアーキテクチャとタスクにおいて、未学習の初期化に対しても1回のパラメータ更新でニューラルネットワークを最適化することに成功した。
• モデルは分布外の重み初期化法に対しても一般化できており、学習分布を超えたロバストネスを示した。
• G.ptは、同じ最適化プロンプトに対してマルチモーダルな解を生成でき、最適性能への多様な経路を学習していることが示された。
• モデルは良好なスケーリング特性を示し、ネットワークの深さや幅に関わらず性能を維持した。
• 従来の学習最適化手法とは異なり、G.ptは分類誤差や強化学習の報酬など、微分不可能な目的関数に対しても効果的に最適化を実行した。
• CIFAR-10のような複雑なデータセットでは若干のアンダーフィッティングが見られたが、シンプルなタスクでは優れた性能を示し、新しいアーキテクチャへの一般化も成功した。