Kaiming Heの新作:拡散モデルに正則化を追加、事前学習やデータ拡張不要で、超シンプルに性能向上を実現

拡散モデルが勢いを増す中、Kaiming He氏の最新論文もこれに関連しています。

この研究は、拡散モデルと表現学習をどのように結びつけるかを探るものです。

拡散モデルに「整理整頓」機能を追加し、その内部特徴をより秩序立たせることで、より自然でリアルな画像を生成できるようにします。

画像

具体的には、本論文ではDispersive Loss(分散損失)という、プラグアンドプレイ式の正則化手法を提案しています。

その核心的な考えは、モデル出力の標準的な回帰損失(例えばノイズ除去)に加えて、モデルの中間表現を正則化するための目的関数を導入することです。

これはコントラスト学習における反発効果に似ています。しかし、コントラスト学習と比較して、その独特の利点は以下の通りです:

正例ペアが不要で、コントラスト学習の複雑さを回避できる;

汎用性が高く、既存の拡散モデルに直接適用でき、モデル構造の変更は不要;

計算オーバーヘッドが低く、追加の計算コストがほとんど増加しない;

元の損失と互換性があり、拡散モデル本来の回帰学習目標を妨げず、既存のフレームワークに統合しやすい。

中間表現を隠れ空間に分散させる

論文の詳細を見ていきましょう。

Kaiming He氏と共同研究者のRunqian Wang氏の出発点は3つあります:

拡散モデルの限界

拡散モデルは複雑なデータ分布の生成に優れていますが、その学習は通常、回帰ベースの目的関数に依存しており、中間表現に対する明確な正則化が不足しています。

表現学習からの着想

表現学習(特にコントラスト学習)は、類似サンプルを近づけ、異なるサンプルを分散させることで、汎用的な表現を効率的に学習できます。

コントラスト学習は分類や検出などのタスクで成功を収めていますが、生成タスクにおけるその潜在能力はまだ十分に探求されていません。

既存手法の不足

REPA(表現アラインメント)などの既存手法は、生成モデルの中間表現と事前学習済み表現をアラインすることで生成効果を改善しようと試みていますが、外部データ、追加のモデルパラメータ、事前学習プロセスへの依存という問題があり、コストが高く複雑です。

彼らは、コントラスト自己教師あり学習を参考にし、生成モデルの中間表現が隠れ空間で分散するように促すことで、モデルの汎化能力と生成品質を向上させる方法を検討し始めました。

画像

このような核心的な考えに基づき、彼らはDispersive Loss(分散損失)を設計しました。これは、モデルの中間表現を正則化し、中間表現の分散性を高めることで、隠れ空間により均一に分布させるものです。

コントラスト学習との違いは、コントラスト学習では、正例ペアをデータ拡張などの方法で手動で定義し、損失関数を通じて正例ペアを近づけ、負例ペアを離す必要がある点です。

一方、Dispersive Lossは正例ペアを定義する必要がなく、負例ペア間の分散性を促すことのみで正則化を実現します。

入力サンプル群画像に対し、Dispersive Lossの目的関数は次のように表すことができます:

画像

ここで、画像は単一サンプルの標準拡散損失、画像は分散損失項、すなわち正則化項であり、λは拡散損失と分散損失の重みをバランスさせるための正則化強度です。

見てわかるように、Dispersive Lossの実装は非常に簡潔で、追加のサンプルペアや複雑な操作は不要であり、モデルの中間層表現に直接適用できます。

また、単層への適用だけでなく、多層への重ね合わせもサポートしており、理論的には複数の異なる中間層に同時にDispersive Lossを適用することで、異なるレベルのフィーチャの分散性をさらに高めることができます。

実験結果

著者はImageNet上で、DiTとSiTをベースラインモデルとして使用し、異なるスケールのモデルで広範なテストを実施しました。

結果は、Dispersive Lossがすべてのモデルと設定で生成品質を向上させることを示しています。例えば、SiT-B/2モデルでは、FIDが36.49から32.45に減少しました。

画像

REPA手法と比較して、Dispersive Lossは事前学習済みモデルや外部データに依存せず、生成品質も劣っていません。

SiT-XL/2モデルでは、Dispersive LossのFIDは1.97であり、REPAのFIDは1.80でした。

画像

さらに、多段階拡散モデルであろうと単段階生成モデルであろうと、Dispersive Lossに基づいて顕著な改善が見られました。

画像

著者は、Dispersive Lossは画像生成タスクだけでなく、画像認識などの他のタスクでも可能性を秘めていると考えています。

論文アドレス:https://arxiv.org/abs/2506.09027v1

— 完 —

📪 量子位AIテーマ企画が進行中です!特集「365業種AI実装ソリューション」「千と一つのAI応用」にご参加いただくか、あなたが探しているAI製品や、発見したAIの新しい動きを私たちと共有してください。

💬 量子位の日刊AI交流グループへのご参加も歓迎します。AIについて自由に語り合いましょう〜

画像

ワンクリックでフォロー 👇 スターを点灯

最先端の技術進捗を毎日お届け

ワンクリック三連「いいね」「シェア」「ハート」

コメント欄にあなたの意見を残してください!

メインタグ:機械学習

サブタグ:人工知能正則化生成AI拡散モデルディープラーニング


前の記事:エコノミストの読書リスト | 今年これまでに読むべき40冊の傑作

次の記事:バークレーとスタンフォードが共同で「研究の預言者」を開発:研究アイデアの将来性を77%の精度で予測

短いURLをシェア