⚡ electrical engineering

AGG: Jacobian-Aggregated Group Gradient for Efficient GRPO Training of Diffusion Models

本論文は、拡散モデルの軌道の近線形な変化を利用して中間ヤコビアンを近似し、グループ勾配を集約する手法であるJAGG（Jacobian-Aggregated Group Gradient）を導入するものであり、これによりDiffusion TransformerのGroup Relative Policy Optimization（GRPO）学習における計算コストを、品質の大幅な低下を招くことなく約2倍削減する。

原著者： Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

公開日 2026-07-21

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

CC BY 4.0

原著者： Ruiyi Ding, Jie Li, He Kang, Ziyan Liu, Chengru Song, Yuan chen

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

あなたは、デジタルアーティストに、より優れた絵を描く方法を教えようとしていると想像してください。このアーティストは人間ではなく、「拡散モデル」と呼ばれる複雑なコンピュータプログラムです。このモデルを、ノイズや静止画の詰まった粘土の塊から、一歩ずつノイズを削り取っていくことで、完璧な彫像を浮かび上がらせる彫刻家だと考えてみてください。この彫刻家に、人間が実際に好む芸術を作れるよう教えるために、私たちは「強化学習」と呼ばれる手法を用います。これは、彫刻家がいくつかの異なるノミの動きを試し、審判（報酬モデル）からスコアを受け取り、そのスコアをもとに次により良くするための学習を行う、というゲームのようなものです。

問題は、この学習プロセスが非常に高価であることです。彫刻家がひとつの小さなチップ（「ステップ」と呼ばれる工程）を作るたびに、コンピュータは彫刻家の脳をどのように調整すべきかを正確に判断するために、膨大なエネルギーを消費する計算を実行しなければなりません。もし彫像を完成させるのに20ステップかかるなら、コンピュータは一つのレッスンごとに、この重い数学的計算を20回行わなければなりません。それは、ダンスの新しいステップを学ぶたびに、動作の途中で立ち止まってコーチに詳細な講評を求め、筋肉の記憶をすべて書き換え、そして最初からやり直すようなものです。高解像度の画像の場合、これには膨大な時間と電力がかかるため、これらのモデルを効果的に訓練することはほぼ不可能です。

ここで、「JAGG」と呼ばれる新しいアイデアが登場します。この論文の著者たちは、賢い問いを投げかけました。「本当に、一歩進むたびに立ち止まってすべてを再計算する必要があるのだろうか？」彼らは、彫刻プロセスの初期の乱雑な段階では、変化が非常に滑らかで予測可能な方法で行われること、つまり、ほとんど直線的な動きであることを発見しました。もし、彫刻がどこから始まり、数ステップ後にどこに到達するかを知っていれば、重い数学的計算をすべてのステップに対して行うことなく、その中間で何が起きたのかを実際に推測することができるのです。

この論文は、JAGG（Jacobian-Aggregated Group Gradient）と呼ばれる手法を紹介しています。JAGGは、ステップごとに停止して計算を行う代わりに、モデルに一連のステップ（例えば4回連続の動き）を行わせ、重い「脳の更新」のための数学的計算を、グループの最初と最後だけ、計2回だけ行います。中間のステップについては、スマートなショートカットを使用します。それは、始点と終点の間の経路が直線であると仮定し、その隙間を埋める方法です。

研究者たちは、もし彫刻プロセスが完全に滑らかで線形であれば、このショートカットは実際には完璧であり、毎回重い作業を行う場合と全く同じ答えを導き出すことを数学的に証明しました。現実の世界では、物事は必ずしも真っ直ぐではありませんが、彼らはこのショートカットが、画像生成の「ノイズが多い」初期段階において驚くほどうまく機能することを見出しました。この手法を用いることで、彼らはこれらのモデルの訓練にかかる時間を、1ステップあたり約17%から18%削減することができました。

最も素晴らしい点は、画像の品質が損なわれなかったことです。彼らがFluxやQwenImageといった異なるモデルでJAGGをテストしたところ、生成された画像は、低速で伝統的な手法で作られたものと同等の品質でした。実際、ケースによっては、このショートカットがモデルの学習をより安定させる助けにもなりました。この論文は、重い作業を行う「タイミング」について少し賢くなることで、最終的な芸術作品の美しさを犠牲にすることなく、これら強力なAIアーティストの訓練をより速く、より安価にできることを示しています。

技術要約: JAGG (Jacobian-Aggregated Group Gradient)

問題提起

Group Relative Policy Optimization (GRPO) は、大規模言語モデル（LLM）において、個別の価値ネットワークを必要とせずに生成モデルを人間の好みに適合させるための主要な強化学習（RL）アルゴッチズムとして台頭しており、成功を収めています。しかし、GRPOを拡散モデルやフローマッチングモデル（具体的にはDiffusion TransformerであるDiT）へ拡張すると、深刻な計算上のボトルネックが生じます。

拡散ベースのGRPOでは、生成プロセスは $W$ ステップのサンプリング・トラジェトリにわたる多段階の意思決定問題として扱われます。標準的な学習では、方策勾配を計算するために、全ステップにおいて高容量のDiTバックボーンのフル・フォワードパスおよびバックワードパスを実行する必要があります。高解像度のテキスト・トゥ・イメージ（T2I）生成では、トラジェトリが数十ステップに及ぶことが多いため、これは極めて高価な計算コストをもたらします。トークンステップが軽量な言語モデルとは異なり、現代的なDiTにおける各バックワードパスは、トランスフォーマーのフォワードパスの全コストを伴うため、この「ステップごとのバックワード要求」が支配的な学習コストとなります。

推論加速のためのトレーニングフリー手法（ $\Delta$ -DiT、ScalingCacheなど）は、DiTの隠れ状態と速度予測がデノイジング・トラジェトリに沿って滑らかかつほぼ線形に変化するという観察結果を利用して中間評価をスキップすることに成功していますが、この「線形性」をRL学習におけるバックワードパスのコスト削減に活用できるかどうかは、未解決の課題でした。

手法: JAGG

著者らは、連続する $W$ 個のデノイジング・ステップのグループに対して、ネットワークのバックワードパスの回数を $W$ 回から 2 回へと削減する、数学的に原理に基づいた加速手法である JAGG (Jacobian-Aggregated Group Gradient) を提案しています。

コアとなる洞察

この手法は、デノイジング速度場 $v_\theta(z, t)$ のパラメータ $\theta$ に関するヤコビアンが、特にデノイジングプロセスの高ノイズ領域において、入力 $(z, t)$ に対して近似的な線形性を示すという洞察に基づいています。これにより、エンドポイントのヤコビアン（ $J_0$ および $J_{W-1}$ ）を用いることで、中間ステップのヤコビアンを線形補間によって正確に近似することが可能になります。

アルゴリズムのワークフロー

JAGGは、 $W$ ステップのタイムステップを以下の4つのフェーズで処理します：

フォワードパス: モデルは $W$ ステップのフォワードパスを実行します。ただし、勾配追跡はエンドポイントの2点のみ（ステップ0とステップ $W-1$ ）に対して有効にします。中間 $W-2$ ステップは、メモリと計算量を節約するために torch.no_grad() を用いて実行されます。
アップストリーム勾配の計算: 各ステップ $j$ について、モデル出力のデタッチされたプロキシが作成されます。ステップごとの損失 $L_j$ （PPOスタイルの目的関数から導出）が計算され、アップストリーム勾配 $s_j = \partial L_j / \partial \hat{v}_j$ が自動微分によって算出されます。極めて重要な点として、このステップにはトランスフォーマーのバックワードパスは含まれません。勾配は軽量なSDEステップとPPO目的関数のみを流れます。
タイムステップ重み付き集約: アップストリーム勾配 $s_j$ は、タイムステップ $t_j$ から導出された補間重み $\alpha_j$ を用いて、2つの合成信号 $s_{base}$ と $s_{corr}$ に集約されます：
$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$
$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$
この集約により、線形補間の仮定を用いて完全な勾配信号を再構成します。
結合バックワード: 2つのエンドポイント出力 $[\hat{v}_0, \hat{v}_{W-1}]$ に対し、集約された勾配 $[s_{base}, s_{corr}]$ を用いて、単一の torch.autograd.backward コールが発行されます。これにより、正確に1回のFSDP all-gatherがモジュールごとにトリガーされ、標準的な学習で必要とされる $W$ 回の個別のバックワードパスが置き換えられます。

理論的保証

著者らは、もし速度場 $v_\theta(z, t)$ が $(z, t)$ に対して厳密に線形であれば、 $t$ 重み付きヤコビアン補間は近似誤差ゼロの厳密なものとなることを示す形式的な証明（定理1）を提供しています。これらの条件下では、2エンドポイント・バックワードパスによって、正確なグループ勾配が復元されます。

適応型ルーティング ( $jagg\_frac$ )

デノイジングの低ノイズ（後半）ステージにおいて線形性の仮定が劣化することを認識し、著者らはルーティングパラメータ $jagg\_frac$ を導入しました。このパラメータは、最初の一定割合のデノイジング・グループ（高ノイズの線形性仮定が成り立つことが勾配のコサイン類似性により検証された領域）にのみJAGGを適用し、残りのグループには正確なステップごとのバックワードパスを使用します。テストされたT2Iモデルにおいて、 $jagg\_frac$ は0.6に設定されました。

主な貢献

ボトルネックの特定: 本論文は、サンプリング・トラジェトリにわたるステップごとのバックワードパスによって引き起こされる、拡散モデルのGRPO学習における計算上のボトルネックを正式に特定し、特徴付けました。
JAGGフレームワーク: 速度ヤコビアンの線形性に関する厳密な理論分析に基づき、バックワードパスをグループあたり $W$ 回から2回に削減するヤコビアン集約フレームワークを提案しました。
理論的証明: 線形速度条件下において、提案された補間が厳密であることを証明し、近似の強固な理論的基礎を提供しました。
実証的検証: 複数の拡散バックボーン（Flux, QwenImage）およびRLアルゴリズム（DanceGRPO, Flow-GRPO-Fast）を用いた広範な実験により、本手法の汎用性と有効性を実証しました。

実験結果

著者らは、HPSv2.1、PickScore、およびCLIPScoreを指標として、T2Iベンチマーク上でJAGGを評価しました。

効率性: JAGGは、標準的なDanceGRPOの設定において、学習ステップあたり**17.3%から18.2%の一貫したウォークロック・スピードアップを達成しました（QwenImageでは約835秒から約691秒へ、Fluxでは約279秒から約228秒へ短縮）。Flow-GRPO-Fastに適用した場合、更新ステージ単体で29.9%から30.0%**のスピードアップが見られました。
品質: この手法は、生成品質の低下をほとんど伴わずにこれらのスピードアップを実現しました。多くの場合、JAGGはベースラインの報酬曲線に密接に追従し、収束時にはベースラインと同等またはそれをわずかに上回る結果を示しました。
アブレーション: エンドポイントに対してバックワードパスを行うものの、ヤコビアン補間を行わない（エンドポイントに独立してフルロスを割り当てる）アブレーション・バリアントは、一貫してベースラインを下回りました。これは、ヤコビアン補間による中間勾配情報の集約が、安定した学習に不可欠であることを裏付けています。
勾配の質: 実証的分析により、高ノイズ・グループ（0–2）において、JAGG集約勾配と厳密な勾配との間のコサイン類似性が高く（0.74–0.85）、線形性の仮定が妥当であることが示されました。後半のグループでは類似性が低下しており、 $jagg\_frac$ ルーティング戦略の正当性が示されました。

意義と主張

本論文は、JAGGが大規模な拡散GRPO学習を実用化するための重要な一歩であることを主張しています。高ノイズ領域における速度ヤコビアンの近線形性を利用することで、生成品質を損なうことなく、RL学習の計算負荷を軽減します。

著者らは、JAGGは単なるヒューリスティックなショートカットではなく、デノイジング・トラジェトリの数学的特性に基づいていることを強調しています。本手法は、既存の推論加速技術とは直交するものであり、特に拡散RLに特有のバックワードパスのボトルネックを標的にしています。これらの結果は、JAGGが計算コストの障壁を大幅に下げることで、視覚的生成モデルに対するRLベースの整列（アライメント）のより広範な採用を可能にすることを示唆しています。