Neural Network Diffusion 2024/03/01 飯田啄巳

ひとことでいうと 拡散モデル SGD ノイズを加える ランダムノイズから特定の分布推定 拡散モデルと勾配降下法は似ている！ 拡散モデルもパラメータ更新として使えるのでは？ NNのパラメータを生成する拡散モデル（p-diff）を作った 拡散モデルはランダム分布から ある特定の分布に変換する能力がある

前提知識 拡散モデル 拡散過程（Forward Process） 生成過程（Reverse Process） 分散𝛽𝑡 をもつガウスノイズ𝒩をを加える （平均は無視します） ガウスノイズの平均と分散𝜇𝜃 , Σ𝜃 を学習 損失 平均𝜇𝜃 , 分散Σ𝜃 をもつガウスノイズ𝒩を引いていく … T 0 … 𝑡 = 𝑞 𝑥𝑡 𝑥𝑡−1 𝑝𝜃 𝑥𝑡−1 𝑥𝑡

Neural Network Diffusion (p-diff)

STEP 1: Parameter Autoencoder オリジナルのStable Diffusionのように潜在空間を学習 Stable Diffusion = 画像 の潜在空間 p-diff = パラメータ の潜在空間 入力を再構成する普通のAutoencoder データ 学習 モデルパラメータのサブセットを使う ➢ ゼロから学習して、最後のエポックを保存 ➢ 学習済みモデルのサブセットをファインチューン 𝑆 = 𝑠1 , … , 𝑠𝑘 , … , 𝑠𝐾 各𝑆 を平坦化して1dベクトル𝑉𝐾×𝐷 = 𝑣1 , … , 𝑣𝑘 , … , 𝑣𝐾 に このを𝑉入出力としてオートエンコーダを学習 サブセットパラメータの大きさ ➢ 入力側の𝑉と潜在空間𝑍にはランダムノイズでデータ拡張 𝜉𝑉 𝜉𝑍 𝑍 𝑉 𝑣𝑘 ≔ 𝑘番目のモデルのパラメータ K個のモデルがある場合 [Rombach+, “High-Resolution Image Synthesis with Latent Diffusion Models”, CVPR, 2022]

STEP 2: Parameter Generation パラメータをそのまま生成すると膨大なメモリが必要 → 潜在空間で生成 DDPMの更新式をそのまま利用 ガウスノイズ Denoise Net

Neural Network Diffusion (p-diff) それぞれを連結してやれば、ランダムノイズから学習済みパラメータがたくさん生成できる

実験設定 データセット - MNIST - CIFAR-10/100 - ImageNet-1K - STL-10 - Flowers - Pets - F-101 アーキテクチャ - ResNet-18/50 - ViT-Tiny/Base - ConvNeXt-T/B 4-layer 1D CNN Autoencoderで学習 パラメータデータセット作成 - ResNet-18/50: スクラッチ学習 - ViT-Tiny/Base: ファインチューニング - ConvNext-T/B: ファインチューニング それぞれ200モデル保存 推論時 x 100 train setでのベストモデルを選択 Best Model val setで検証 originals ensemble p-diff x 100 最後の2層のBNレイヤだけ

メインの結果 SGDで学習したもの&アンサンブルしたものと同じかそれ以上のAccuracyを達成

Ablation Study bestの 差は小さい 学習データ（学習されるモデル数）は 安定性に影響がある。 拡散モデルの学習原理上、 少数サンプルだと分布の学習が難しい。 - どのレイヤのパラを生成してもbestは良い - 深い層のパラを生成した方が良い ノイズがあると安定的になる 潜在ノイズの方が効果アリ

モデルのパラメータ全体を生成したら？ 小さなモデルを用意して、モデルのパラメータ全体を生成した場合も試してみる ConvNet-3, MLP-3 細かいアーキテクチャの設定 ResNet, ViT, ConvNeXtはGPUのメモリが厳しいので実験不可能らしい Stable Diffusionの1024x1024x3の画像生成で30~40GBくらい使った記憶があるので厳しいのはわかる いい結果（小並感）

どんなシードで学習してもある程度のパターンがある 各レイヤでうまくいくパラメータのパターンがある

p-diffが記憶している可能性を調査 1. オリジナルモデルのパラメータを記憶しているだけでは？ 同じパラメータが生成されているかも。多様性ある？ 2. ノイズの付加やオリジナルモデルのファインチューニングで 違いはあるか？ 各モデルの出力結果を類似度を測る 指標＝間違った結果のIoU (最近傍) p-diffは多様な パラを生成 普通にSGDで 学習すると 似た感じになる 最近傍 +noise finetune original models p-diff model (b) の補足図 ファインチューニングしてp-diffのパラに行き着くか検証 オリジナルモデルを 超えるのは難しい &類似度高い t-SNEで潜在ベクトルの分布を可視化 潜在空間でもp-diffは オリジナルと異なる

パラメータ生成の軌跡 各時間ステップでのパラメータの場所（t-SNE）をプロット オリジナルモデルに 近づく傾向がある 学習モデル数を増やすと 多様化する 学習に使ったモデル数の違いが 生成されたパラメータの多様性に寄与するか

Appendix

1D CNNのかわりにFC層を使った場合

VAEを拡散モデルのかわりに使った場合

各ノイズの強さ

p-diffのtrainとvalデータの性能上の関係 この図の読み方がわかりません💦

p-diffはSGDで学習するよりも超高速 いや、生成してるだけでは…？

他のタスクでは Object Detection Semantic Segmentation Image Generation