他のモデルの重みを使用してニューラルネットワークをトレーニングすると、ニューロンを変更しても出力には影響しません。NVIDIAの魔法の研究

コンピューター ビジョンであれ NLP であれ、ディープ ニューラル ネットワーク (DNN) は、今日の機械学習タスクを完了するための推奨される方法です。これに基づいて構築されたモデルでは、モデルの重みに対して何らかの変換を実行する必要がありますが、そのプロセスを実行する最適な方法は何でしょうか?

最近、NVIDIA は ICML 2023 研究 DWSNet を解釈しました。 DWSNet は、ニューラル ネットワークの配置の対称性の等価性を維持しながら、ニューラル ネットワークの重みを効率的に処理できる新しいアーキテクチャを提供します。

このアプローチにより、他のネットワークの重みに基づいてネットワークをトレーニングできるため、継続的な学習に適したアプローチとなる可能性があります。同様に興味深いことに、DWSNet の調査に基づいて、ネットワークの重みには順列対称性があることがわかりました。つまり、出力を変更せずにニューロンの順序を変更できるということです。

• 論文リンク: https://arxiv.org/abs/2301.12780
• 公式 GitHub: https://github.com/AvivNavon/DWSNets

このアプローチは幅広い可能性を秘めていると考えられており、事前トレーニング済みのネットワークを新しいドメインに適応させるなど、さまざまな興味深いタスクを可能にすることができます。おそらく、別のネットワークから情報を抽出、編集、または削除するようにネットワークをトレーニングできるでしょう。

生成モデルを使用すると、さまざまなタスク用にトレーニングされた多数のネットワークをトレーニングし、実行時に特定のタスク用のネットワークを生成できるようになります。これは、高速重みネットワークの最新バージョンとまったく同じです。

DWSNet がどのようにそれを実行するか見てみましょう:

暗黙的ニューラル表現 (INR) またはニューラル放射場 (NeRF) を使用して表現された 3D オブジェクトのデータセットを操作する場合、カップのハンドルを削除したり、ホイールをより対称的にするなど、オブジェクトを「編集」してジオメトリを変更したり、エラーを修正したりする必要があることがよくあります。ただし、INR と NeRF を使用する際の大きな課題は、編集前にレンダリングする必要があり、編集ツールは INR または NeRF パラメータのレンダリングと微調整に依存していることです。

図 1. データ固有のアーキテクチャの例。

Nvidia の研究チームは、ニューラル ネットワークを一種のプロセッサとして使用し、他のニューラル ネットワークの重みを操作しようとしています。

ディープ ネットワークのパラメーターを特徴付ける最も簡単な方法は、すべての重み (およびバイアス) を単純なフラット ベクトルにベクトル化し、完全に接続されたネットワーク (多層パーセプトロン (MLP)) を適用することです。この方法はニューラル ネットワークのパフォーマンスを予測できます。

しかし、このアプローチには欠点があります。ニューラル ネットワークの重みの空間は複雑な構造を持ち、すべてのパラメーターのベクトル化されたバージョンに MLP を適用すると、この構造が無視され、一般化能力が低下します。

図 2. 2 つの隠れ層 (下) を持つ多層パーセプトロン (MLP) の重みの対称性 (上)。

MLP のこの問題は、幾何学的深層学習 (GDL) の分野で広く研究されてきました。

多くの場合、学習タスクはいくつかの変換に対して不変です。たとえば、ポイント クラウド クラスの検索は、ポイントがネットワークに供給される順序とは無関係です。ただし、ポイント クラウド セグメンテーションなど、ポイント クラウド内の各ポイントにクラスが割り当てられる場合、入力の順序が変わると出力も変わります。

入力が変化すると出力が変化する関数を同変関数と呼びます。対称性を考慮したアーキテクチャには、意味のある帰納的バイアスによる複数の利点があります。たとえば、通常、サンプルの複雑性が向上し、パラメーターが少なくなるため、一般化機能を大幅に向上できる要因となります。

重み空間の対称性

では、MLP によって表される基礎関数が変化しないように、MLP の重みにどのような変換を適用できるでしょうか?

これには、ニューロンの配置という特定の種類の変換が関係します。図 2 に示すように、直感的には、MLP の特定の中間層でニューロンの順序を変更しても機能は変わりません。さらに、各内部層の並べ替えプロセスは独立して実行できます。

MLP は次の方程式系で表すことができます。

このアーキテクチャの重み空間は、ベクトル化された重みとバイアスのすべての連結を含む (線形) 空間として定義されます。

重要なのは、この方法では重みの空間が（定義される）ニューラル ネットワークの入力空間であるということです。

では、重み空間の対称性とは何でしょうか?ニューロンの並べ替えは、1 つの層の出力に順列行列を適用し、同じ順列行列を次の層に適用するものとして正式にモデル化できます。正式には、新しいパラメータ セットは次の式で定義できます。

新しいパラメータ セットは異なりますが、この変換によって MLP によって表される関数が変更されないことは容易にわかります。これは、2 つの順列行列 P と P^t が互いに打ち消し合うためです (ReLU のような要素ごとの活性化関数を想定)。

より一般的には、前述のように、MLP の各レイヤーに異なる順列を独立して適用できます。これは、次のより一般的な変換セットでは基礎となる関数が変更されないことを意味します。これらを重み空間の「対称性」と考えます。

ここで、Pi は順列行列を表します。この観察は、30 年以上前に Hecht-Nielsen が論文「フィードフォワード ネットワークの重み空間の代数構造について」で行ったものです。同様の変換を MLP のバイアスにも適用できます。

深層重み付け空間ネットワークの構築

文献に記載されているほとんどの等価アーキテクチャは同じアプローチに従います。つまり、単純な等価レイヤーを定義し、そのような単純なレイヤーの構成としてアーキテクチャを定義します。レイヤー間にはポイント単位の非線形性が存在する可能性があります。

CNN アーキテクチャは、このような構造の良い例です。この場合、単純な等変層が畳み込み演算を実行し、CNN は複数の畳み込みの組み合わせとして定義されます。 DeepSets と多くの GNN アーキテクチャは同様のアプローチに従います。詳細については、Weisfeiler と Leman の論文「Go Neural: Higher-Order Graph Neural Networks」および「Invariant and Equivariant Graph Networks」を参照してください。

ターゲット タスクが定数である場合、図 3 に示すように、MLP を使用して等価層の上に不変層を追加できます。

図 3: 典型的な等変アーキテクチャは、いくつかの単純な等変層と、それに続く不変層および完全接続層で構成されます。

論文「Deep Weight Spaces での学習のための等変アーキテクチャ」では、NVIDIA の研究者がこの考え方を採用しました。私たちの主な目標は、上で定義した重み空間対称性に対して、シンプルでありながら効果的な等変層を特定することです。残念ながら、一般的な同変関数の空間を特徴付けるのは難しい場合があります。これまでの研究（例えば、クロスアンサンブル相互作用の深層モデル）と同様に、私たちの目標はすべての線形等変層の空間を特徴付けることです。

そのため、研究者らは、重み空間 V が各重み行列 V=⊕Wi を表すより単純な空間の連結であるという観察に基づいて、線形等変層を特徴付ける新しい方法を開発しました。 (簡潔にするためバイアス用語は省略します)。

この観察は非常に重要です。なぜなら、(i,j)番目のブロックがWjとWi Lij : Wj→Wiの間の線形等変層であるブロック行列に、任意の線形層L:V→Vを書き込むことができるからです。ブロック構造を図4に示します。

しかし、Lij のすべてのインスタンスを見つけるにはどうすればよいでしょうか?この論文では、考えられるすべてのシナリオをリストし、これらのレイヤーのいくつかが以前の研究ですでに特徴付けられていることを示しています。たとえば、内部層の Lii は、クロスセット相互作用の深層モデルで特徴付けられます。

この場合、最も一般的な等変線形層が共通であり、これは 4 つのパラメータのみを使用してディープ アンサンブル層を一般化したものであることは注目に値します。他のレイヤーについては、新しい研究では、プーリング、ブロードキャスト、小さな完全接続レイヤーなどの単純な等変操作に基づくパラメーター化を提案し、それらがすべての線形等変レイヤーを表現できることを示しています。

図4は、特定の重み空間間の分割行列であるLの構造を示しています。各色は異なるタイプのレイヤーを表します。李は赤です。各ブロックは、特定の重みマトリックスを別の重みマトリックスにマッピングします。このマッピングは、ネットワーク内の重みマトリックスの位置に応じてパラメータ化されます。

図 4: 線形等価層のブロック構造。

このレイヤーは、各ブロックを個別に計算し、各行の結果を合計することによって実装されます。 Nvidia は、バイアス項の処理や複数の入力および出力関数のサポートなど、いくつかの追加の技術的詳細を論文で説明しています。

これらのレイヤーを Deep Weight Spatial Layers (DWS Layers) と呼び、それによって構築されるネットワークを Deep Weight Spatial Network (DWSNet) と呼びます。ここでは、MLP を入力とする DWSNet に焦点を当てます。

ディープウェイト空間ネットワークの表現力

仮定のクラスを単純な同変関数の組み合わせに制限すると、上記のグラフニューラルネットワークの文献で広範に研究されてきた同変ネットワークの表現力を意図せず損なう可能性があります。 Nvidia の論文では、DWSNet が入力ネットワーク上のフィードフォワード操作を近似できることが示されており、これはその表現力を理解するための一歩となります。その後、新たな研究により、DWS ネットワークは MLP 関数空間で定義された特定の「適切に動作する」関数を近似できることが実証されました。

実験

DWSNet は 2 つのタスク ファミリで評価されます。まず、INR などのデータを表す入力ネットワークが使用されます。次に、標準の I/O マップ (画像分類など) を表す入力ネットワークが使用されます。

実験1: INR分類

この構成では、INR が表す画像に基づいて INR を分類します。具体的には、MNIST と Fashion-MNIST からの画像を表現するために INR をトレーニングします。タスクは、DWSNet がこれらの INR 重みを入力として使用して、MNIST の数字などの画像コンテンツを認識できるようにすることです。結果は、NVIDIA が提案した DWSNet アーキテクチャが他のベースラインを大幅に上回っていることを示しています。

表1: INR分類。 INR のクラスは、それが表す画像によって定義されます (平均テスト精度)。

重要なのは、INR をそれが表す画像カテゴリに分類することは、基礎となる画像を分類することよりも難しいということです。 MNIST 画像でトレーニングされた MLP は、ほぼ完璧なテスト精度を達成できます。しかし、MNIST INR でトレーニングされた MLP は、悪い結果を達成しました。

実験2: INRの自己教師学習

ここでの目標は、ニューラル ネットワーク (特に INR) を意味的に一貫した低次元空間に埋め込むことです。優れた低次元表現は多くの下流タスクにとって非常に重要であるため、これは重要なタスクです。

ここでのデータは、a\sin(bx) の形式の正弦波に適合された INR で構成されます。ここで、a、b は区間 [0,10] の一様分布からサンプリングされます。データはこれら 2 つのパラメータによって制御されるため、密な表現によって基礎となる構造が抽出されるはずです。

図 5: 自己教師ありトレーニングを使用して取得した入力 MLP の TSNE 埋め込み。

SimCLR のようなトレーニング プロセスと目的を使用して、ガウス ノイズとランダム マスクを追加することで、各 INR からランダム ビューを生成します。図 4 は、結果として得られた空間の 2D TSNE プロットを示しています。 NVIDIA が提案した DWSNet 方式はデータの潜在的な特性を非常にうまく捉えますが、同等の方法はより困難です。

実験3: 事前学習済みネットワークを新しいドメインに適応させる

この実験では、再トレーニングなしで事前トレーニング済みの MLP を新しいデータ分布に適応させる方法 (ゼロショット ドメイン適応) を示します。画像分類器の入力重みが与えられた場合、タスクはその重みを、新しい画像分布 (ターゲット ドメイン) で適切に機能する新しい重みセットに変換することです。

テスト時に、DWSnet は分類器を取り込み、単一のフォワードパスでそれを新しいドメインに適応させます。 CIFAR10 データセットがソース ドメインであり、その破損バージョンがターゲット ドメインです (図 6)。

図 6: DWSNet を使用したドメイン適応。

結果を表2に示す。注意: テスト時には、モデルは未知の画像だけでなく、未知の画像分類器にも一般化される必要があります。

表 2: 新しいドメインへのネットワークの適応。

今後の研究の方向性

Nvidia は、学習技術をディープウェイト空間に適用する能力が、多くの新しい研究の方向性を提供すると考えています。まず、重み空間上の関数をトレーニングするための効果的なデータ拡張スキームを見つけることで、DWSNet の一般化能力が向上する可能性があります。第二に、スキップ接続や正規化レイヤーなど、他の種類の入力アーキテクチャやレイヤーに順列対称性を組み込む方法を研究するのは自然なことです。

最後に、DWSNet を変形、NeRF 編集、モデル プルーニングなどの実際のアプリケーションに拡張すると便利です。 ICML 2023 論文「Deep Weight Spaces での学習のための等変アーキテクチャ」を参照してください。