Spatially-Enhanced Temporal Fusion Transformer: Interpretable Multi-Output Prediction for Parametric Dynamical Systems with Time-Varying Inputs

あなたは、気象システムや巨大な電力網のような、複雑な機械の未来を予測しようとしていると想像してください。科学の世界では、これらの機械は「微分方程式」と呼ばれる、巨大で複雑な数学の方程式によって記述されることがよくあります。これらの方程式は、その機械がどのように振る舞うかを示す「レシピ」のようなものだと考えてください。問題は、これらのレシピがあまりにも巨大で詳細すぎるため、解くのにスーパーコンピュータでも長い時間がかかるということです。特に、いくつかの材料（温度や圧力など）を調整したり、外部からの力（突然の突風など）を変えたりした場合に、何が起こるかを見たいときにはなおさらです。科学者たちは何十年もの間、これらの方程式のより単純で高速なバージョン、すなわち「サロゲートモデル（代理モデル）」を構築しようとしてきました。

最近、「Transformer（トランスフォーマー）」と呼ばれる人工知能の一種が、長い物語を読み、言葉が時間の経過とともにどのように関連し合うかを理解する能力によって有名になりました。チャットボットや翻訳ツールで見かけることがあるかもしれません。これらは、単一の事象（例えば、ある都市の気温）を予測することに長けています。しかし、現実世界の機械は通常、多くの動くパーツが同時に互いに影響を及ぼし合っています。例えば、扇風機の速度を変えると、温度、圧力、そして騒音レベルがすべて同時に変わるかもしれません。大きな疑問は、「Transformerに、これらすべての異なるパーツを一度に観察し、それらがどのように共に踊っているかを理解させ、混乱することなく、一度にショー全体を予測させることはできるのか？」という点でした。

この論文では、**SE-TFT（Spatially-Enhanced Temporal Fusion Transformer）**と呼ばれる新しいAIモデルを紹介しています。著者であるマックス・プランク研究所の研究者たちは、過去のデータと将来の入力に基づいて単一の結果を予測することに長けていた既存のモデルである「Temporal Fusion Transformer (TFT)」を取り上げ、それに大幅なアップグレードを施しました。彼らは、複雑なシステム（複数の出力を持つもの、例えば温度、圧力、速度を同時に扱うもの）を予測するためには、AIが単に「いつ」物事が起こるか（時間）だけでなく、「どこに」位置しているか（空間）を理解する必要があることに気づいたのです。

元のモデルを、一つの出来事のタイムラインを暗記するのは得意だが、同時に進行する3つの異なるストーリーラインを追わされると圧倒されてしまう学生だと考えてください。SE-TFTはその同じ学生ですが、今度はグリッド（格子）付きの特別なノートを持っています。単なるリストを書く代わりに、彼らはあらゆる瞬間に「温度」のストーリーが「圧力」のストーリーとどのように結びついているかを見ることができます。著者らは、新しい「マスキング」技術を作成することでこれを実現しました。Transformerの世界において、「マスク」とは、AIが何を見ることが許可されているかを伝えるルールのことです。通常、AIは未来を予測するために過去だけを見ることを許可されています。SE-TFTは新しいルールを追加しました。それは、すべての異なる出力の過去を同時に見ることができるというルールです。これにより、モデルは、異なる要素がどのように時間の経過とともに変化するか（時間的関係）を学ぶと同時に、それらが互いにどのように影響し合っているか（空間的関係）を学ぶことができるようになりました。

研究者たちは、この新しいモデルが混沌とした状況を扱えるかどうかを確認するために、3つの非常に異なるタイプのシステムでテストを行いました。第一に、彼らは、混沌とした気象パターンを記述する有名な数学的システムであるLorenz-63モデルを使用しました。これは、蝶の飛行経路の正確な軌跡を予測しようとするようなものです。出発点のわずかな変化が、劇的に異なる結果をもたらします。SE-TFTは、初期条件がランダムであっても、システム内の3つの変数の将来の経路を正確に予測することに成功しました。

次に、彼らは、ニューロン（脳細胞）が電気信号に反応してどのように発火するかをシミュレートするFitzHugh-Nagumoモデルを試しました。これは、ドミノが倒れていく様子を見ているようなものですが、そのドミノはリセットされて再び倒れることもあり、複雑なリズムを生み出します。ここでは、モデルは2つの異なる出力（電圧と回復変数）を、変化する外部信号に対処しながら予測しなければなりませんでした。SE-TFTは単に数値を推測したのではなく、「信頼区間」を提供しました。これは、本物の答えが位置する可能性が高い場所を示す「緑色のゾーン」を予測の周囲に描くようなものです。実際の結果はこのゾーン内に安全に収まっていました。

最後に、彼らはバッテリーのようなセットアップの中で化学物質がどのように移動し反応するかを記述する、結合電気化学キネティクスおよび拡散モデルに取り組みました。これは、回転速度や電気信号の周波数が化学物質の挙動を変える、非常に厄介なシステムです。SE-TFTは、電流と2種類の異なる化学物質の濃度を驚異的な精度で予測し、その誤差はしばしば1%未満でした。

この論文の最も素晴らしい特徴の一つは、モデルが「解釈可能（interpretable）」であることです。通常、ディープラーニングモデルは「ブラックボックス」であり、データを入力すると数値が出力されますが、コンピュータがどのようにその結論に至ったのかは分かりません。しかし、SE-TFTは自身の「アテンション（注目）」の記録を保持しています。著者らは、モデルの内部にある「アテンション・マップ」を見ることで、どの過去の要素が予測に影響を与えたのかを正確に特定できることを示しました。例えば、化学モデルにおいて、電流の予測はそれ自体の履歴に大きく依存しており、一方で化学濃度の予測は、すべての異なる変数の混合に依存していることがマップから明らかになりました。これは、AIに対して「なぜ温度が下がると考えたのですか？」と尋ねると、AIがその結論に至った特定の過去の出来事を指し示してくれるようなものです。

著者らは、SE-TFTが、一度に一秒ずつ予測してそれをフィードバックするという方法（エラーが蓄積しやすい「自己回帰」と呼ばれる手法）を用いることなく、これらの複雑なマルチアウトプット・システムを一段階で予測できることを見出しました。モデルは、非常に遠い未来を予測する際にはわずかに精度が低下しましたが、驚くほど堅牢（ロバスト）であり続けました。論文は、AIに異なる出力間の「空間的」なつながりを見せることを教えることで、気象パターンから化学反応器に至るまで、複雑な物理世界をシミュレートするための、より速く、より正確で、より理解しやすいツールを構築できると結論付けています。これらの実験で使用されたコードとデータは、誰でも確認できるように公開されており、これらの発見が他の人々によって検証され、発展していくことを保証しています。