【PIDM】物理的正則化付き拡散モデル
3つの要点
✔️ 拡散モデルに物理的制約を導入する理論的手法を提案
✔️ オーバーフィッティングに対する正則化機能を確認
✔️ 等式だけでなく不等式や最適関数まで展開可能
Physics-Informed Diffusion Models
written byJan-Hendrik Bastek, WaiChing Sun, Dennis M. Kochmann
(Submitted on 21 Mar 2024)
Comments: 15 pages, 4 figures; added further theoretical motivation, new residual estimation mechanism and additional experimental study
Subjects: Machine Learning (cs.LG); Computational Engineering, Finance, and Science (cs.CE)code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
概要
拡散モデルは非常に複雑なデータ分布を近似する上で極めて高い性能と汎用性を備えており、近年、自然科学においてもその応用が急速に広まっています。科学的な文脈における応用を考えると、特定のクラスの問題に対してはデータが従う支配方程式が明示的に分かっており、データがその支配方程式に従うという正則化を与えたいケースが頻繁にあります。これまでの、拡散モデルの多くの科学的応用は、純粋にデータ駆動で行われ、生成されるサンプルが物理的法則に従うかは非自明でした。この研究では、生成されるサンプルが支配方程式に従うように、拡散モデルに物理的正則化を与える理論的方法を検討し、数値実験を通してその有用性を実証しました。
背景
現行のAIにおける重要な基礎となっているものとして拡散モデルです。拡散モデルは、データが従う複雑な分布を近似する上で極めて高い性能を持っており、急速に利用が広まりました。近年では、自然科学においても拡散モデルの応用が期待されています。しかし、拡散モデルの科学的応用の多くが、純粋にデータからそのデータの従う分布を近似することを目指したものであり、科学的な知見との融合に対する検討は十分ではありません。科学的知見としてデータの従う支配方程式が明示的に分かっているような特定のケースでは、生成されるサンプルが支配方程式に従うよう物理的な制約を拡散モデルに与える手法の確立が必要です。これらを背景として、この研究では、拡散モデルに物理的正則化を与えることを理論的に検討しました。
提案手法
ここでは、この研究中で重要な幾つかの要素を説明します。
拡散モデル
拡散モデルは、最先端の生成モデルの一種です。拡散モデルの目的は簡潔に言えば、データ$x_0$の従う分布$q(x_0)$を近似しモデル化することです。そのために、拡散モデルでは$x_0$から$x_T$まで$T$ステップのデータの系列を考え、それぞれのステップでガウシアンノイズを付加していき、$x_T$がピュアなガウシアンノイズになるようにします。このプロセスをforward difussion processと呼び、以下で定義します。
ここで、$\{\beta\in (0,1)\}_{t=1}^T$が拡散過程を決めるパラメータです。また、この逆操作としてガウシアンプロセスからサンプルを生成する系列を以下で定義します。
ここで、未知の分布である$q(x_{t-1}|x_t)$をニューラルネットワークによって$p_θ$$(x_{t-1}|x_t)$として近似します。最終的に、Hoらによる簡略化に従うと、以下の損失関数によって訓練を行います。
つまり、概念的には、ノイズを付加したのちノイズを除去した$\hat{x}_0$と元データ$x_0$との誤差が小さくなるように訓練されます。以上が拡散モデルの簡単な説明です。
支配方程式
一般的に支配方程式は、以下のように抽象的に表現することができます。
加えて、以下の一般的な境界条件を考えます。
ただし、ここで、
は抽象的な微分演算子、
は境界条件、
は支配方程式の解を表しています。この研究中では、拡散モデルによって生成されるサンプル$x_0$が上述の境界条件までを含めた支配方程式を満足することを仮定します。この研究では、$x_0$は画像形式のデータを考えており、例えば、力学の問題における応力の分布を表す画像などが考えられます。この支配方程式に従って、以下の残差を定義します。
つまり、簡単に言えばこの残差によって支配方程式を満たしているかを測ることができます。
Physics-informed diffusion model (物理的正則化付き拡散モデル)
・目的関数の設計
この研究では、拡散モデルの確率的な観点を損なわないため仮想的な残差$\hat{r}$を導入し、以下の分布に従うと仮定しています。
これを用いて、仮想的な尤度$p_θ$($\hat{r}$)を以下のように考えます。
これらを用いることで、物理的な正則化の目的を以下のように表現することができます。
つまり、残差が0になる確率が最大になるようにパラメータを調整するということです。筆者らはこれを、Physics-informed neural networkの誤差関数の確率的な解釈であると指摘しています。
さらに、この研究中では上述の誤差関数に加えて観測データによる誤差関数も考えています。筆者らの指摘している通り、これは推定する関数の崩壊を防ぐ効果があります。つまり、例えば、至る所0になるような関数は、特定のタイプの支配方程式を満たしますが物理的には意味のない解です。階の探索の際にそのような自明の解にはまることの無いようにある種の正則化項として観測データからの誤差が必要です。その観測データからの誤差を追加して、以下のように目的関数を定義しました。
・目的関数の簡略化
ノイズ除去の系列全てに対して残差が0になることを拡散モデルに課すのは正則化としては厳しすぎ、拡散モデルの柔軟性を損なう危険性があります。そのため、筆者らは適切にそれぞれのステップで正則化がスケールするように設計しました。つまり、$T$から0ステップにノイズ除去を行うに従って、正則化が強くなるように設計しました。以下がそのための定義です。また、図１としてノイズ除去過程におけうPIDMの概略図を示しています。
さらに、
が
を近似するために設計されていることを加味して、目的関数を簡略化することで、最終的に以下の目的関数を得ることができます。
第一項が、観測データに対する誤差関数であり、第二項が支配方程式に対する誤差関数です。以上が、PIDMに関する簡単な説明です。
実験結果
この研究中では、PIDMの有効性を示すために幾つかの数値実験を行っています。この解説中では、そのうち物理的な意味合いの強いものとして、多孔質中の流れを表す二次元のDarcy flowに関する例題を取り上げます。対象とする支配方程式は以下です。
ただし、
です。この研究では、比較のために幾つかのモデルを用意しました。(i)データを用いた目的関数から学習する拡散モデル、(ii)データのみを用いた目的関数から学習するが、残差の情報をガイダンスとして考慮した拡散モデル（Phisics-guided diffusion）、(iii) データのみを用いた目的関数から学習する拡散モデルを用いるが、推論の際に残差の一次の修正を加えるモデル（CoCoGen）です。また、この研究では、上述の式中のpermitivityとpressureを生成するモデルを考えています。
図２として、それぞれのモデルとこの研究で提案されたPIDMの学習の過程における誤差関数の履歴を可視化したものを示しています。図２(a)から、PIDMが残差に対する誤差として二桁程度の改善が確認できます。さらに、図２(b)の結果から、従来手法は学習が進んでいくと、テストデータに対する精度が大きくなり、オーバーフィッティングが起こっていることが予測されますが、PIDMは物理的な正則化によってそれが防がれている様子が確認できます。これらの結果は、正則化項によって精度の改善に加えて、オーバーフィッティングを防ぐ効果があることを示唆しています。
また、図３として通常の拡散モデルとPIDMの与えるpermitivityとpressureを比較したものを示しています。また、その際の残差の空間分布も示しています。この結果も、図２の結果と整合して、通常の拡散モデルよりも精度の高い推定が行われていることが確認されます。また、図３の(b)と(c)は異なる物理的状態を表現していることから、PIDMが一つの解に陥ることなく、多様な状態を表現できることを示唆しています。以上の結果は、PIDMの高いポテンシャルを支持するものです。
まとめと結論
この研究では、物理的な正則化項を加えた拡散モデルPIDMの理論的な導出と数値実験によるその性能の実証を行いました。特に、これまでの研究のような、推論時に潜在変数を「修正」することによってではなく、拡散モデルに直接、物理的正則化を課すことを可能にしました。また、数値実験の結果、PIDMは通常の拡散モデルに比べて、残差に関する誤差が二桁程度改善することに加えて、オーバーフィッティングに対しても堅牢であることが示唆されました。今後、自然科学における拡散モデルの需要がますます高まる中で、この研究の与えた理論的基礎は重要なものになることが予想されます。今後、具体的で様々な自然科学の問題に適用され、汎用的なツールの基盤になることが期待されます。
この記事に関するカテゴリー
