【論文瞬読】AI開発の常識が変わる！？ Stanford発の新フレームワーク「TEXTGRAD」とは？

こんにちは！株式会社AI Nestです。
今回は、スタンフォード大学の研究チームが発表した革新的なAIフレームワーク「TEXTGRAD」について解説します。

タイトル：TextGrad: Automatic "Differentiation" via Text
URL：https://arxiv.org/abs/2406.07496
所属：DEPARTMENT OF COMPUTER SCIENCE, STANFORD UNIVERSITY
, DEPARTMENT OF BIOMEDICAL DATA SCIENCE, STANFORD UNIVERSITY, CHAN ZUCKERBERG BIOHUB
著者：Mert Yuksekgonul, Federico Bianchi, Joseph Boen, Sheng Liu, Zhi Huang, Carlos Guestrin, James Zou

複合AIシステムの最適化という新たな挑戦

最近、ChatGPTやLlama 2といった大規模言語モデル（LLM）の発展が目覚ましいですよね。これらのLLMは、質問応答や文章生成など、単体でもすごい能力を発揮しますが、複数のLLMや他のツールと連携することで、さらに複雑なタスクをこなせる「複合AIシステム」へと進化しています。例えば、数学の問題を解いたり、プログラミングの課題をクリアしたり、さらには新薬の開発や医療の分野にまで応用されているんです。

しかし、こんな複合AIシステム、どうやって開発・最適化すればいいのでしょうか？ これがAI研究における新たな挑戦として注目されているんです。

従来のAI開発とTEXTGRADの違い

従来のAI、特にニューラルネットワークの分野では、「逆伝播」と「自動微分」という手法がモデルの最適化に革命をもたらしました。これは、簡単に言うと、モデルの出力と目標値の誤差を計算し、その誤差を基にモデルのパラメータを調整していく手法です。

TEXTGRADは、この逆伝播の概念をテキストベースのフィードバックに拡張した画期的なフレームワークです。複合AIシステムの各コンポーネントを、あたかもニューラルネットワークの層のように扱い、LLMが生成するテキストフィードバックを逆伝播することで、システム全体を最適化します。図1は、このTEXTGRADの概念図と、様々な分野への応用例を示しています。

Figure1, テキストによる自動「微分」（a、b）。勾配のバックプロパゲーションはディープラーニングの原動力である。ブラックボックスAIシステムの複合システムに対する勾配はないが、テキストベースのフィードバックに対する類似のバックプロパゲーションを構成することができ、これがTEXTGRADの基礎となる。TEXTGRADの抽象化(c)。TEXTGRADはPyTorchと同じ抽象化と構文を共有することで、汎用的で学習しやすいフレームワークを実現している。TEXTGRADの応用(d,e,f,g)。セクション3.4では、ドラッグライクさやタンパク質結合親和性などの特性について分子構造を最適化する(d)。セクション3.1では、コーディング問題の解を最適化する(e)。セクション3.5では、放射線治療計画を最適化し、患者の予後を改善する(f)。セクション3.3では、言語モデルの推論を改善するためにプロンプトを最適化する(g)。

TEXTGRADの仕組みをざっくり解説

1. 計算グラフの作成: まず、複合AIシステムの各コンポーネントの入出力を変数として、計算グラフを作成します。この計算グラフは、各コンポーネントがどのように連携しているかを示すものです。

2. LLMによるフィードバック生成: 次に、LLMは、システム全体の出力に対して自然言語でフィードバックを生成します。例えば、「このコードはもっと効率化できる」とか「この分子の構造はこう変えるべき」といった感じです。

3. テキストフィードバックの逆伝播: TEXTGRADは、このフィードバックを計算グラフに沿って逆伝播します。各コンポーネントは、自分に関連するフィードバックを受け取り、それを基に自分のパラメータを調整します。

4. 最適化の繰り返し: このフィードバック生成と逆伝播のプロセスを繰り返すことで、システム全体が徐々に最適化されていきます。

TEXTGRADのすごいところ

TEXTGRADのすごいところは、その汎用性の高さです。論文では、コードの最適化、複雑な科学的問題の解決、推論タスクの改善、新薬の分子設計、放射線治療計画の最適化など、多岐にわたる分野でTEXTGRADが有効であることが実証されています。しかも、それぞれの分野に合わせてフレームワークを調整する必要はなく、どんなタスクにも適用できるんです。例えば、図2や図3は、TEXTGRADがそれぞれ新薬の分子設計と放射線治療計画の最適化に適用された例を示しています。

Figure2, Text.TEXTGRADによる分子最適化は、タンパク質受容体PPARAに対する薬効（より高いQED）と結合親和性（より低いvinaスコア）を向上させるために、ベンゼネフラグメントを最適化する。最初の3反復のテキスト勾配を(a)に、臨床的に承認されているPPARA標的分子と比較した全10反復のパフォーマンスを(c)に示す。最終反復の分子は、臨床的に承認されている最も類似した分子と構造的類似性が低く、QEDとVinaのスコアが優れており（d）、（e）に示すように、非常に妥当なポーズ形状をしている。29のターゲットと3つの初期フラグメントに おいて、TextGradは臨床的に承認された分子と同程度のVinaスコアと より高いQEDスコアを持つ分子の設計に成功している（b）。

Figure3, 放射線治療計画の最適化。腫瘍の標的化とOARの保護のバランスをとるために、計画標的体積（PTV）と危険臓器（OAR）の重要度の重みを調整するバックプロパゲートテキスト勾配の視覚化。(a)の画像は初期化から反復5までのプランの変化を示す。TEXTGRADは反復的に平均線量を改善し、PTVへの線量のばらつきを減少させ、複数回の反復で臨床目標を達成する(b)。TEXTGRADは膀胱と直腸の被ばく線量を臨床的に許容される最大値以下に抑えている(c)。TEXTGRADで最適化された計画は、臨床的に最適化された計画よりもPTVの線量指標（平均線量とD95）が良好であり（d）、膀胱と直腸の線量が低く、OARの保護が良好であることを示している（e）。

まとめ｜TEXTGRADでAI開発は新たなステージへ

TEXTGRADは、複合AIシステムの開発と最適化を自動化する画期的なフレームワークです。これにより、AI開発の効率が大幅に向上し、これまで以上に複雑で高度なAIシステムを構築できるようになるでしょう。今後のAI研究の方向性を大きく変える可能性を秘めたTEXTGRAD。今後の発展に目が離せませんね！