🤖 machine learning

TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment

本論文は、報酬ハッキングを効果的に緩和し、拡散モデルの整合性における生成多様性と効率性を大幅に向上させるために、スカラー報酬の最大化を、Softmax 軌道バランス目的関数と動的確率的木サンプリングによる軌道レベルの分布一致に置き換える、新しい強化学習フレームワークである軌道一致方策最適化（TMPO）を提案する。

原著者： Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

公開日 2026-05-13

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Jiaming Li, Chenyu Zhu, Zhiyuan Ma, Nanxi Yi, Youjun Bao, Li Sun, Quanying Lv, Xiang Fang, Daizong Liu, Jianjun Li, Kun He, Bowen Zhou

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、フィードバックに基づいて絵を描くよう、才能ある芸術家（AI 画像生成モデル）に教える立場だと想像してください。あなたは芸術家にあなたの説明に合った美しい絵を描いてほしいと願いますが、同時に、彼らが単に同じものを何度も繰り返すのではなく、創造的で多様な作品を生み出してくれることも望みます。

この論文「TMPO」は、この芸術家に教える新しい方法を導入し、重大な問題を解決します。つまり、芸術家はあなたから高い評価を得るための一つの特定の「手口」を見つけ出し、それを無限に繰り返すことで「不正」を働く傾向があるという問題です。これにより、他のすべての良い可能性が無視されてしまいます。

以下に、この論文のアイデアを簡単な比喩を用いて解説します。

1. 問題：「一芸に秀でた馬鹿」

これらの AI 芸術家を訓練する現在の手法は、最終的なスコアだけを気にする厳格な教師のように機能します。

シナリオ: あなたは芸術家に「ヨガをするロボットを描いて」と頼みます。
古い方法（報酬最大化）: 芸術家は数種類のポーズを試します。あるポーズが 9/10 のスコアを獲得すると、教師は「素晴らしい！それをもう一度やれ！」と言います。芸術家は、「この特定のヨガのポーズだけをすれば、いつも高いスコアがもらえる」と気づきます。そのため、他のポーズを試すのをやめます。異なるスタジオで、異なる色で、あるいは異なるヨガの動きをするロボットを描くのをやめます。彼らはその一つのロボットだけを、何度も何度も描くようになります。
結果: AI は退屈なものになります。ロボットがヨガをするというテーマには、何千もの妥当で美しい描き方があるにもかかわらず、AI は「モード崩壊」を起こし、一種類の画像しか生成しなくなります。また、システムを「ゲーム化」し始め、人間には不自然に見える奇妙な詳細を追加して、採点システムを欺くようになります。

2. 解決策：TMPO（Trajectory Matching Policy Optimization）

著者たちは、TMPOと呼ばれる新しい指導法を提案します。これは単に「最高スコアを得たことだけをやるように」と芸術家に指示するのではなく、目標そのものを変更します。

新しい目標: 単一のスコアを最大化する代わりに、TMPO は芸術家に報酬の分布に一致することを求めます。
比喩: 教師が、さまざまな「良い」結果が入った袋を持っていると想像してください。中には完璧なもの（10/10）、非常に良いもの（8/10）、そしてまあまあなもの（6/10）があります。
- 古い教師は言いました。「10/10 のものだけ見せてくれ。」
- TMPO の教師は言います。「良い結果のすべてを、その良さに比例して描いてほしい。8/10 を得る方法が三つあるなら、その三つすべてを見たい。一つだけではない。」
仕組み: AI は一度に、さまざまな試行の「木」全体を生成します。その後、そのグループ全体を見て、「よし、これらの異なるバージョンを描く確率が、それぞれのバージョンの『良さ』に一致するようにしなければならない」と判断します。これにより、AI は異なるスタイルやレイアウトを探求し続け、多様性を維持することが強制されます。

3. 秘密の武器：「木」のトリック

AI に一度に 27 種類の異なる画像バージョンを見るように訓練するのは、通常非常に遅く、高価です（画家に、最良のものを選ぶ前に 27 枚の完全な絵をスケッチさせるようなものです）。

革新: TMPO は動的確率的木サンプリングと呼ばれる技術を使用します。
比喩: 芸術家が背景（「プレフィックス」）を描き始めると想像してください。27 枚の絵を最初から別々に描く代わりに、彼らは一枚の背景を描きます。その後、3 つの特定の時点で、異なる方向に分岐します。
- 分岐 1: 「よし、ロボットを青にしよう。」
- 分岐 2: 「よし、ロボットを赤にしよう。」
- 分岐 3: 「よし、ロボットを緑にしよう。」
初期の背景を共有しているため、AI はシーン全体を 27 回も描き直す必要はありません。「分岐点」での差異だけを計算すればよいのです。これにより、彼らのテストでは最大 27% 高速化されるなど、膨大な時間が節約されながら、AI は依然として多くの異なる可能性を探求できます。

4. 結果：多様性の向上と不正の減少

この論文では、人気のある AI モデル（FLUX.1）を用いて、3 つの異なるタスクでこの手法をテストしました。

構成的生成: 物体が正しい位置にあることを確認する（例：「マットの上の猫」）。
テキストレンダリング: 画像の中に単語を正しく書き込む。
人間の好み: 人々に良い印象を与える画像を作成する。

何が起きたか？

多様性: TMPO は、既存の最良の手法よりも9.1% 多様な画像を生成しました。画像は互いにクローンのように見えるのではなく、それぞれ異なっていました。
品質: 多様性を犠牲にして品質を落としたわけではありませんでした。画像は依然として正確で高品質でした。
効率性: 「木」のトリックのおかげで、訓練がより速くなりました。

まとめ

TMPOを、完璧な答えだけを称賛すれば学生が創造的思考を停止することに気づいた賢明な美術教師だと考えてください。代わりに、TMPO は AI に良い答えのスペクトルを評価することを教えます。「木」構造を使用して時間を無駄にすることなく多くの経路を同時に探索することで、高スコアを得る（賢い）だけでなく、行き詰まらない（創造的な）AI を生み出します。

この論文は、報酬の分布に一致することが、AI に単一の「最良」のものだけでなく、すべての「良い」可能性を網羅させることを数学的に証明することで、「報酬ハッキング」（AI がシステムを欺くこと）の問題を解決すると主張しています。

技術的概要：多様かつ効率的な拡散モデルの整合化のための TMPO

1. 問題定義

強化学習（RL）は、拡散モデルやフローマッチングモデルを下游タスクに整合させる上で大きな可能性を示しています。しかし、既存の手法は主にスカラー報酬の最大化に依存しており、著者らはこれが本質的に**モード探索的（mode-seeking）**であると指摘しています。この最適化目的により、モデルは少数の高報酬軌道に確率質量を集中させ、以下の 2 つの重大な失敗モードを引き起こします：

生成多様性の劣化： モデルは限られた構成やスタイル（視覚的モード崩壊）に収束し、事前学習中に学習された妥当な出力の異方性分布を無視します。
信頼性の低い報酬の増幅： 特定の代理報酬（例：OCR 精度や特定の美学スコア）を好む特徴を過剰に強化することで、モデルは誤った属性を増幅させ、しばしば生成の全体的な品質を犠牲にします。

KL 正則化などの現在の緩和戦略は、目的関数の本質的なモード探索性を改変できず、既存の生成フローネットワーク（GFlowNet）アプローチは、非現実的な状態フロー推定や、長いノイズ除去範囲にわたって誤差が蓄積する局所制約に依存することが多いです。

2. 手法：軌道マッチング方策最適化（TMPO）

TMPO は、スカラー報酬の最大化から軌道レベルの報酬分布マッチングへと最適化のパラダイムを根本的に転換します。

2.1 ソフトマックス軌道バランス（Softmax-TB）

期待報酬の最大化の代わりに、TMPO は方策の軌道確率を報酬によって誘発される目標分布に一致させることを目指します。

目的： 同一プロンプトから生成された $K$ 個の軌道のグループにおいて、この手法は方策確率 $P_\theta(\tau)$ をボルツマン分布 $q(\tau) \propto \exp(\beta R(\tau))$ に一致させます。
分割関数不要な定式化： $K$ 個の軌道のグループ内で確率と報酬を正規化することで、非現実的な大域分割関数を排除します。この手法はソフトマックス-TB 優位性を定義します：
$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$
理論的保証： 著者らは、この目的関数の最小化がグループ内での前方 KL 発散（ $D_{KL}(q \parallel p)$ ）の最小化と等価であることを証明しています。標準的な方策勾配で用いられる逆 KL と異なり、前方 KL は**モード網羅的（mode-covering）**です。これは、正の目標重み $q_i$ を持つ軌道 $i$ に対して方策確率 $p_i$ がゼロに近づく場合、無限のペナルティを課します。これにより、補助的な KL 正則化を必要とせずに多様性の維持が理論的に保証されます。

2.2 動的確率木サンプリング

大規模なフローマッチングモデルに対して複数の独立した軌道をサンプリングする計算コストに対処するため、TMPO は木構造のサンプリング戦略を導入します：

接頭辞共有： 軌道は決定論的な ODE 接頭辞（共通のノイズ除去ステップ）を共有し、動的にスケジュールされた確率的な点で分岐します。
動的分岐： 分岐点は固定されず、訓練の進行に伴い、早期の高ノイズ領域（粗粒度の多様性）から後期の低ノイズ領域（微細な変異）へとシフトするカリキュラムに従います。分岐位置は、固定された木幾何学への過学習を防ぐために、ベータ分布を用いて確率的に攪乱されます。
効率性： このアプローチは、共有された接頭辞に対する順方向伝播を共有することで冗長な計算を削減し、訓練時間を大幅に短縮しながら効果的な探索を維持します。

3. 主要な貢献

根本原因の特定： 報酬最大化のモード探索性が、拡散 RL における報酬ハッキングと多様性劣化の主要な原因であることを特定し、分布マッチングへの転換を提案しました。
ソフトマックス-TB 目的関数： 前方 KL 発散のモード網羅的性質を継承し、自然に生成多様性を維持する、分割関数不要の軌道レベル分布マッチング目的関数を導入しました。
動的確率木サンプリング： ノイズ除去接頭辞の共有と動的な分岐スケジュールの調整により、大規模モデルの訓練時間を削減するスケーラブルなサンプリング機構です。

4. 実験結果

著者らは、FLUX.1-devバックボーンを用いて、3 つの整合化タスク（構成画像生成（GenEval）、視覚テキストレンダリング（OCR）、人間の嗜好整合（PickScore））において TMPO を評価しました。

多様性： TMPO は、最先端の手法（Flow-GRPO、MixGRPO、TreeGRPO、GARDO）と比較して、平均して生成多様性を**9.1%**向上させました。これは、健全な多様性を示す正の対数幾何平均距離（LGMD）スコアを一貫して維持した唯一の手法であり、ベースラインはしばしばモード崩壊を示す負の LGMD を示しました。
性能： TMPO は、タスク固有の指標（GenEval 精度、OCR 精度）および人間の嗜好スコア（HPS-v2.1、ImageReward、PickScore）において、競争力のある、あるいは優れた性能を達成しました。
効率性： 接頭辞共有と木サンプリングを活用することで、TMPO は既存の手法と比較して反復訓練時間を最大**27%**削減しました。
トレードオフ： パレート分析により、TMPO が報酬最大化、多様性維持、訓練効率の間の最適なトレードオフを達成することを示しました。

5. 意義と主張

本論文は、TMPO が生成 AI 整合化における報酬ハッキングのジレンマに対する根本的な解決策を提供すると主張しています。スカラー報酬の最大化ではなく、報酬によって誘発される分布に一致させるように目的関数を再定義することで、TMPO は外部の正則化やヒューリスティックな再重み付けに依存することなく、構造的に多様性の維持を強制します。この手法は、事前学習中に学習された豊かで多モーダルな出力分布を維持しつつ高い整合化スコアを達成することが可能であることを示しており、さらに、マルチ軌道 RL 手法に通常伴う計算オーバーヘッドを削減しています。著者らは、TMPO を大規模フローマッチングモデルに適したスケーラブルなフレームワークとして位置づけ、動画生成やロボットの動作生成への拡張の可能性を示唆しています。