AI News HubLIVE
站内改写

Sakana AIが提案するDiffusionBlocks：残差ネットワークを独立して訓練可能なデノイジングモジュールに変換するブロック単位学習フレームワーク

Sakana AIと東京大学の研究者らは、Transformerベースのネットワークをブロックごとに訓練するDiffusionBlocksを提案。訓練メモリをブロック数B分の1に削減しつつ、多様なアーキテクチャで性能を維持する。残差接続を拡散モデルのオイラー法ステップと解釈し、スコアマッチングによる原理的な局所目的を実現する。

記事インテリジェンス

エンジニア上級

要点

  • DiffusionBlocksはネットワークをB個の独立訓練可能なブロックに分割し、メモリをB倍削減。
  • 残差ネットワークと拡散モデルの関連性を活用し、理論的に裏付けられた局所訓練目標を提供。
  • ViT、DiT、マスク拡散、自己回帰、再帰的深層Transformerの実験で、エンドツーエンド訓練と同等の性能と大幅なメモリ削減を達成。
  • 拡散モデルでは、推論時に各デノイジングステップで1ブロックのみ活性化され、計算量も削減。

重要な理由

このニュースが重要なのは、DiffusionBlocksはネットワークをB個の独立訓練可能なブロックに分割し、メモリをB倍削減ためです。

技術的影響

モデル選定、推論コスト、プロダクト能力、評価基準に影響する可能性があります。

Sakana AIと東京大学の研究チームは、深層ニューラルネットワーク訓練におけるメモリボトルネックを解決する新たなフレームワーク「DiffusionBlocks」を提案した。このフレームワークはTransformerネットワークを複数の独立したブロックに分割し、各ブロックを独立に訓練することで、訓練メモリをブロック数B分の1に削減する。

従来のエンドツーエンドのバックプロパゲーションでは、全層の中間活性化を保存する必要があり、メモリ消費はネットワークの深さに比例して増加する。アクティベーションチェックポインティングは活性化メモリを削減できるが、パラメータ、勾配、オプティマイザ状態のメモリは削減できない。例えばAdamオプティマイザでは、各層にパラメータ、勾配、モメンタムと分散の2つの状態を保存する必要があり、合計でパラメータサイズの約4倍のメモリが必要となる。ブロック単位訓練はネットワークをB個のブロックに分割し独立に訓練することで、メモリ要件を約1/Bに削減する。しかし、各ブロックに原理的な局所目標を定義することが課題であった。従来のForward-Forwardアルゴリズムや欲張り層単位訓練はアドホックな局所目標に依存し、エンドツーエンド訓練に劣り、主に分類タスクに限定されていた。

DiffusionBlocksの核心は、残差ネットワークと拡散モデルの間の深い関係を利用することにある。残差ネットワークの更新z_ℓ = z_{ℓ-1} + f_θℓ(z_{ℓ-1})は常微分方程式のオイラー離散化に対応する。研究チームは、この更新が特にスコアベース拡散モデルにおける確率流ODE（分散爆発（VE）定式化）に対応することを示した。逆拡散過程のオイラー離散化は残差接続更新と構造的に一致する。したがって、残差ブロックのスタックは離散化されたデノイジングステップとして解釈できる。スコアベース拡散モデルでは、スコアマッチング目標は各ノイズレベルで独立に最適化できるため、各ブロックは自身の局所目標のみで独立に訓練でき、ブロック間通信は不要となる。

標準的な残差ネットワークをDiffusionBlocksに変換するには3つの修正が必要である。(1) L層ネットワークをB個の連続ブロックに分割、(2) ノイズ分布（推奨：対数正規分布）を定義し、各ブロックにノイズ範囲の区間を割り当て、(3) 各ブロックの入力にターゲットのノイズ付きバージョンを追加し、AdaLNによるノイズレベル条件付けを追加。訓練時は毎イテレーション1ブロックのみサンプリングされ、他のブロックは計算されない。メモリ消費は全L層ではなくL/B層となる。

分割戦略として、DiffusionBlocksは等確率分割（equi-probability partitioning）を採用する。一様分割はノイズレベル間のデノイジング難易度の違いを無視するが、等確率分割は各ブロックが総確率質量の1/Bを処理するように境界を設定する。CIFAR-10でのアブレーション実験では、等確率分割がFID 38.03を達成し、一様分割の43.53を大きく上回った。

実験評価は5つのアーキテクチャにわたって行われた。CIFAR-100のViTではDiffusionBlocksが59.30%の精度（エンドツーエンド：60.25%）、CIFAR-10のDiT-S/2ではFID 37.20（ベースライン：39.83）、ImageNet 256×256のDiT-L/2ではFID 10.63（12.09）、text8のマスク拡散モデル（MDM）ではBPC 1.45（1.56）、LM1Bの自己回帰TransformerではMAUVE 0.71（0.50）、OpenWebTextではMAUVE 0.82（0.85）を達成した。また、再帰的深層モデルHuginnでは、約10倍の計算削減でMAUVE 0.70（ベースライン：0.49）を実現した。これらの結果は、DiffusionBlocksがメモリを大幅に削減しながら、エンドツーエンド訓練と同等またはそれ以上の性能を発揮することを示している。

同時期の研究NoPropとの比較では、DiffusionBlocksは連続時間定式化とブロック単位訓練を組み合わせた唯一の手法であり、CIFAR-100で46.88%の精度を達成し、エンドツーエンドのバックプロパゲーションと1%ポイント以内の差に留まった。

DiffusionBlocksの強みは、スコアマッチングに基づく原理的な理論的基盤、タスク固有の修正なしで多様なアーキテクチャに適用可能、B倍の訓練メモリ削減、拡散モデルにおける推論計算のB倍削減、等確率分割の有効性、再帰的深層モデルにおけるBPTTの置き換え、GPU間の並列訓練とゼロ通信オーバーヘッドなどである。

一方、入力と出力の次元が一致する必要があるためU-Net型アーキテクチャには適用できない、スクラッチからの訓練のみ検証済み、最適なブロック数を選択する原理的方法がない、ノイズ条件付けによる時間オーバーヘッド（0.0543秒 vs 0.0507秒）、OpenWebTextでの一部指標がわずかに劣るなどの限界もある。

総じて、DiffusionBlocksは深層Transformerネットワークの訓練に対するメモリ効率が高く理論的に妥当な代替手段を提供し、特に拡散モデルや再帰的深層モデルに有望である。