DiffusionBlocks：ブロック単位でニューラルネットワークを訓練する

現代の人工知能は、日常的なタスクから数学、コーディングに至るまで顕著な性能を示しています。しかし、最先端モデルは通常数百億以上のパラメータを持ち、訓練には数千のGPUが必要であり、開発できる組織は限られています。このリソース需要の大きな原因の一つは訓練方法そのものにあります。現在のニューラルネットワーク（Transformerを含む）は、エンドツーエンド最適化により全パラメータを同時に学習するため、メモリ要求がモデルサイズとともに増大します。理想としては、利用可能なハードウェアに関わらず任意のサイズのモデルを訓練できることです。これがDiffusionBlocksの動機であり、AI開発のアクセシビリティ向上を目指します。

DiffusionBlocksは、ネットワークを複数のブロックに分割し、一度に一つのブロックだけを独立に訓練する手法です。その結果、訓練に必要なメモリはネットワーク全体ではなく単一ブロック分のみとなります。同時に、近年成功を収めている拡散フレームワークを活用することで、エンドツーエンド最適化と同等の性能を維持します。本研究は、東京大学の小山真典氏との共同研究であり、機械学習のトップ国際会議ICLR 2026で発表されました。

背景：深層学習におけるメモリボトルネック

現代のAIはスケーリングによって推進されてきましたが、メモリコストも増大します。エンドツーエンドのバックプロパゲーションでは、全ての中間状態（活性化）をメモリに保持する必要があるため、メモリ消費はネットワークの深さに比例します。Transformerは層数を増やすことでスケールするため、深さが直接メモリコストを押し上げます。

ブロック単位訓練の実践

ブロック単位訓練は自然な解決策です。ネットワークを小さなブロックに分割し独立に訓練することで、訓練メモリは単一ブロック分に低下します。先行研究は主に画像分類に限られ、生成タスクや現代のTransformerアーキテクチャへの拡張は十分に探求されていません。主な制限は、局所目的関数の理論的基盤の弱さと適用範囲の限定性です。

DiffusionBlocksの手法

DiffusionBlocksは各ブロックに明示的な役割を割り当て、ネットワーク全体が目標を達成するようにします。各ブロックは「ブロックを進むにつれて徐々に目標に近づく」というダイナミクスを担当し、これは拡散モデルの本質と一致します。残差接続が常微分方程式の離散ステップに対応するという観察に基づき、ブロックダイナミクスを拡散モデルの逆（ノイズ除去）過程と解釈し、原理的な独立訓練を可能にします。

既存ネットワークを変換するには以下の3ステップが必要です：

1. 分割：L層をBブロックに分割。
2. ノイズ範囲の割り当て：各ブロックは「目標への近さ」の範囲を担当。
3. 条件付けモジュールの追加：ブロックが自身の割り当て範囲を認識できるようにする。

訓練時にはランダムに一つのブロックをサンプリングするだけでよく、メモリ消費は約1/Bに削減されます。

検証結果

DiffusionBlocksは、画像分類（ViT）、画像生成（DiT）、テキスト生成（マスク拡散、自己回帰Transformer、リカレント深層Transformer）の3つのタスクドメイン、5つのアーキテクチャで検証されました。全てのケースでエンドツーエンド訓練と同等の性能を達成し、メモリ使用量を削減しました。

リカレント深層モデルへの拡張

リカレント深層モデル（ループ型Transformer）では、同じネットワークをK回繰り返し適用します。従来の時間方向バックプロパゲーション（BPTT）は計算ボトルネックですが、DiffusionBlocksの視点からは、訓練中に一回の順伝播で済み、計算コストを大幅に削減できます。

今後の展望

二つの重要な方向性：第一に、なぜメモリ効率と性能が共に向上するのかの理論的分析（自然なカリキュラム学習の仮説）；第二に、既存の大規模事前学習モデルへの適用。ファインチューニングによりDiffusionBlocksに変換できれば、大規模モデルの訓練と利用が、潤沢なリソースを持つ組織だけでなく、個人研究者や学生、小規模ラボにも可能になるでしょう。

DiffusionBlocksは、Sakana AIによるAI効率化の広範な取り組みの一部であり、訓練時のメモリ消費という新たな軸を提供します。詳細は論文をご参照ください。