[論文レビュー] Equivariant Architectures for Learning in Deep Weight Spaces

Aviv Navon, Aviv Shamsian+4 more|arXiv (Cornell University)|Jan 30, 2023

Machine Learning in Materials Science被引用数 5

ひとこと要約

本稿では、重み行列とバイアス行列から直接事前学習されたMLPを処理する、新しいニューラルアーキテクチャであるDeep Weight-Space Networks (DWSNets)を提案する。この手法は、重みの自然な置換対称性に不変性を課すことで、それらの重みの構造的性質を尊重する。アフィン不変性を持つ層は、プーリング、ブロードキャスト、全結合層の操作によって特徴付けられ、パラメータ効率の高い学習を可能にし、一般化誤差予測やINR分類といったタスクで最先端の性能を達成する。

ABSTRACT

Designing machine learning architectures for processing neural networks in their raw weight matrix form is a newly introduced research direction. Unfortunately, the unique symmetry structure of deep weight spaces makes this design very challenging. If successful, such architectures would be capable of performing a wide range of intriguing tasks, from adapting a pre-trained network to a new domain to editing objects represented as functions (INRs or NeRFs). As a first step towards this goal, we present here a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trained MLP and processes it using a composition of layers that are equivariant to the natural permutation symmetry of the MLP's weights: Changing the order of neurons in intermediate layers of the MLP does not affect the function it represents. We provide a full characterization of all affine equivariant and invariant layers for these symmetries and show how these layers can be implemented using three basic operations: pooling, broadcasting, and fully connected layers applied to the input in an appropriate manner. We demonstrate the effectiveness of our architecture and its advantages over natural baselines in a variety of learning tasks.

研究の動機と目的

• 事前学習されたMLPの重みを、その生の行列形式そのままで処理できる深層学習アーキテクチャの設計を目的とする。
• MLPの重み行列に内在する複雑な対称性構造が、標準的なアーキテクチャの学習を妨えるという課題に対処することを目的とする。
• MLP重み行列に内蔵された置換対称性を尊重する、整合的なフレームワークの構築を目的とする。
• 重み空間上で直接、モデルの適応、プルーニング、表現学習などの下流タスクを可能にすることを目的とする。

提案手法

• アーキテクチャは、MLP重みの置換対称性を尊重するアフィン不変性を持つ線形層（DWS層）に基づいている。
• これらの層は、入力ネットワークの特定の重みおよびバイアス部分空間間のマッピングを表すブロック行列として定式化される。
• 層は、プーリング、ブロードキャスト、および標準的な全結合層の3つの基本操作によって実装される。
• DWS層のスタックに点関数非線形性を組み合わせることで、重み空間処理のための深層アーキテクチャが構築される。
• この手法は不変性を保証する：入力MLP内のニューロンの置換は、出力に対して一意で予測可能な変換を引き起こす。
• 構造的な重み共有のおかげで、完全結合ベースラインと比較して顕著に少ないパラメータ数で実現される、パラメータ効率の高いアーキテクチャである。

実験結果

リサーチクエスチョン

• RQ1 MLP重み行列の置換対称性に対して不変である線形層が満たすべき数学的制約は何か？
• RQ2 基本的なディープラーニング操作を用いて、このような不変層をどのように効率的に実装できるか？
• RQ3 DWSNetは、一般化誤差の予測や密度のある表現の学習といった、事前学習済みMLPの空間上で関数を学習するのに有効に機能するか？
• RQ4 生の重み行列を扱うタスクにおいて、このアーキテクチャは標準的なベースラインと比較してどの程度優れているか？
• RQ5 DWSNetsは、INR や NeRF などの他のアーキテクチャにもどの程度一般化可能か？

主な発見

• DWSNetsは、CIFAR-10およびMNISTデータセットにおいて、重みからモデル性能を予測する一般化誤差の観点で、すべてのベースラインを上回る最低水準の誤差を達成した。
• モデルはINRの周波数と振幅ごとにグループ化される、解釈可能な2次元の密度のある表現を学習した。
• 重み同士、およびバイアス同士のブロックを併用することで、片方のみを用いる場合よりも性能が向上し、対角バージョンではほぼ最適な結果が得られた。
• データ拡張とバッチ正則化を導入することで、MNIST INR分類タスクのテスト精度が77.20%から85.71%に向上した。
• このアーキテクチャは強力な表現力を持ち、学習された層を通じて入力MLPの順方向伝搬を近似可能であることが示された。
• プルーニングやINR処理の面倒さにもかかわらず、適切な初期化とドメインに配慮した正則化を施すことで、有望な結果が得られた。