最適輸送が組み込まれた「生成拡散モデル」の学習則の解明

• 画像AI生成サービスの基盤である生成拡散モデルを、変分オートエンコーダの拡張として定式化することで、複雑な学習則の単純化に成功
• これによりモデル内の2つのニューラルネットワークが果たす役割を明確化し、一方の学習を早期に打ち切ることにより、精度を落とさずに学習コストを削減
• 計算資源の限られた環境における高性能な生成モデルの活用に期待

東京科学大学（Science Tokyo） 理学院 物理学系の蒲健太郎大学院生と大関真之教授、東北大学 情報科学研究科の清水怜央大学院生（当時）と杉山友規特任准教授の研究チームは、最適輸送^[用語1]の理論を取り入れた生成拡散モデル^[用語2]の学習則を単純化し、学習の効率化に成功しました。

生成拡散モデルは、生成AI^[用語3]の一種であり、多くの画像生成AIサービスの基盤となる技術です。近年、その性能向上のために、最適輸送の手法である「シュレーディンガー橋^[用語4]」を組み込んだモデルが注目されていますが、複雑な数理構造ゆえに学習則の直感的な理解が困難でした。

本研究では、生成拡散モデルを変分オートエンコーダ^[用語5]という他の生成AIの拡張として定式化することで、モデル内の2つのニューラルネットワークが果たす役割を明確化しました。さらに、この理論的枠組みに基づき、一方のニューラルネットワークの学習を早期終了することで、精度を犠牲にせずに計算コストが削減可能であることを数値計算で実証しました。これらの成果は、計算資源の限られた環境における高性能な生成モデルの活用を推し進めることが期待されます。

本成果は、9月3日付（現地時間）の「Physical Review Research」誌に掲載されました。

近年、生成AIは多様な分野で社会基盤技術として急速に普及しています。特に、生成拡散モデルは、画像や音声などの高次元データを高精度に出力可能な生成AIとして注目を浴びており、幅広い応用が期待されています。

生成拡散モデルの性能向上のための研究は盛んに行われており、中でも、「シュレーディンガー橋」と呼ばれる最適輸送の手法を用いることで、柔軟かつ高速なサンプリングを可能にすることが報告されています^{[参考文献1]}。従来のモデルでは、生成拡散モデルを構成する2つの拡散過程^[用語6]のうち片方にのみニューラルネットワークが組み込まれていましたが、最適輸送理論を用いることで、もう一方にもニューラルネットワークを組み込むことが可能になります。しかし、このモデルは数理構造が複雑になるために学習則の直感的理解が難しく、また、その性能と引き換えに学習のための計算コストが高くなるという問題点を抱えています。

本研究では、シュレーディンガー橋を用いた生成拡散モデルを、変分オートエンコーダの拡張として再解釈する新たな理論枠組みを提案しました。これにより、モデルに組み込まれた2つのニューラルネットワークの学習の役割が、それぞれ「事前分布の学習」と「分布間を繋ぐダイナミクスの再現」に分けられることを示しました。具体的には、前者はデータをノイズへ変換する最適なダイナミクスを構築し、後者は新たなデータ生成を行うために前者を逆向きに再現します。また、前者は後者よりも学習が素早く安定化していることも数値実験で明らかにしました。これにより、前者の学習を早期に打ち切ることにより、精度を落とさずに学習コストが削減可能かつ過学習^[用語7]も抑制可能であることが示されました。この枠組みは従来のモデルも記述可能であるため、さまざまな生成拡散モデルを統一的な視点から扱うことが可能になります。

生成拡散モデルは、生成AIの基盤技術として画像生成、音声合成、物理シミュレーションなど多様な応用が期待されています。本成果により、シュレーディンガー橋型モデルの学習コスト削減と汎化性能向上が可能となり、計算資源の限られた環境でも高性能な生成モデルの活用が進むことが期待されます。さらに、統計物理学などでのサンプリング手法への応用も見込まれます。

本研究における理論の枠組みは、他の確率過程を用いた生成AIに対しても適用可能です。今後は、非マルコフ過程や他種のダイナミクスへの拡張による新たな生成AIの開発や、統計力学におけるボルツマン分布サンプリングへの応用などを進める予定です。

本研究は、日本学術振興会（JSPS）科学研究費補助金 基盤研究（Ｂ）連続時間量子フィードバック制御によるロバストな量子誤り訂正の開発（23H01432）、戦略的イノベーション創造プログラム（SIP）「先進的量子技術基盤の社会課題への応用促進」量子コンピュータを活用した新事業を共創する研究開発基盤（23836436）、研究開発とSociety5.0との橋渡しプログラム（BRIDGE）量子プロダクト事業化推進プラットフォーム構築事業の支援を受けて実施されました。

[1]

T. Chen, G.-H. Liu, and E. A. Theodorou, Likelihood training of Schrödinger bridge using forward-backward SDEs theory, in International Conference on Learning Representations (2022).

[用語1]

最適輸送：1つの分布を、もう1つの分布に移す際にかかるコストを最小化するような輸送方法を決める問題。

[用語2]

生成拡散モデル：生成AIの一種であり、画像生成AIサービスなどの基盤技術。データにノイズを加える過程と、それを再現する生成過程の2つの拡散過程を用いて記述される。

[用語3]

生成AI：手元にあるデータを元に、それらに似た新たなデータを作り出すことを目的とするAI。

[用語4]

シュレーディンガー橋：拡散過程を用いた輸送を行う場合の最適輸送問題の一種。

[用語5]

変分オートエンコーダ：生成AIの一種であり、隠れ変数を用いたモデルとしては最も初期に提案されたモデル。

[用語6]

拡散過程：各時刻にランダムな力が加わるダイナミクス。ブラウン運動などを記述するために用いられる。

[用語7]

過学習：生成AIが手元にあるデータの再現のみを行い、新たなデータを生成するための汎化能力を失う状態。

掲載誌：

Physical Review Research

タイトル：

Schrödinger bridge-type diffusion models as an extension of variational autoencoders

著者：

Kentaro Kaba, Reo Shimizu, Masayuki Ohzeki, Yuki Sughiyama

DOI：

10.1103/dxp7-4hby

• プレスリリース 最適輸送が組み込まれた「生成拡散モデル」の学習則の解明—計算資源の限られた環境における高性能な生成モデルの活用にも期待—（PDF）
• 2024年ノーベル物理学賞を読み解く：AI時代を切り拓いたニューラルネットワークの革命 | Science Tokyoニュース

• 大関 真之 Masayuki Ohzeki | Science Tokyo研究情報データベース（理工学系）
• 大関研究室
• 理学院 物理学系
• 理学院

• 東北大学