Calibrating Generative Models to Feature Distributions with MMD Finetuning

技術要約：MMDファインチューニング（kCGM）による生成モデルの特徴分布へのキャリブレーション

問題提起
生成モデルは、個々のサンプルとしては妥当であっても、特定の関心対象となる特徴に関して分布の不一致を示すことがよくあります。例えば、広範な化学空間で事前学習されたモデルは、特定の治療クラス（抗生物質など）の特定の特徴分布とは大きく異なる分子を生成する可能性があります。これをターゲットセットに対して直接ファインチューニングによって修正することは、小規模なデータセットへの過学習のリスク、ターゲットとなる全サンプル（希少な場合がある）の必要性、およびどの特定の特徴を一致させるべきかを指定できないといった問題を伴います。さらに、既存の手法は、非微分的な（例：Morganフィンガープリントのような）離散値またはブラックボックス的な特徴に対して苦戦することが多いです。本論文は、特徴写像やサンプリング手順の微分可能性を必要とせずに、特徴レベルの監督のみを用いて生成モデルをターゲットの特徴分布に適合させる能力を持つアルゴリズムにおけるギャップを特定しています。

手法：kCGM
著者らは、生成された特徴分布とターゲットの特徴分布との間の最大平均偏差（MMD）を最小化し、事前学習済みモデルへのカルバック・ライブラー（KL）ダイバージェンスによって正則化するファインチューニング手法であるkernel Calibrating Generative Models (kCGM) を導入します。

• 目的関数: この手法は、以下の2つの項からなる目的関数 $L_\theta$ を最適化します：

  1. MMD損失: モデルが誘導する特徴分布 $P$（ここで $y = h(x)$ かつ $x \sim p_\theta$）と、経験的なターゲット分布 $Q$（ターゲット特徴 $y^{targ}_{1:N}$ からなる）の間の二乗MMDの、不偏有限サンプル推定値。
  2. KL正則化: 事前学習済み分布への近さを維持し、破滅的忘却やモード崩壊を防ぐための項 $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$。

• カーネル選択: ドット積カーネル（平均のみを一致させるもの）に依存する可能性のある従来のアプローチとは異なり、kCGMは確率分布上の計量を定義するカーネル（同一の分布に対してのみゼロとなるもの）を利用します。

  • 連続的なユークリッド特徴に対しては、エネルギー距離カーネルが使用されます。
  • 離散的なビットベクトル特徴（例：Morganフィンガープリント）に対しては、タニモト/ジャカードカーネルが採用されます。
  • これらの選択により、追加のバンド幅ハイパーパラメータの導入を回避します。

• 勾配推定: 特徴写像 $h$ およびサンプリングプロセスが非微分的である可能性があるため、kCG�Mはスコア関数勾配（REINFORCE形式）を採用します。

  • Leave-One-Out (LOO) ベースライン: 重要な技術的貢献は、MMD損失のための新しいLOOベースラインの導出です。標準的なLOOベースラインはバッチ内の他のサンプルの係数を平均化しますが、MMDの係数はサンプル間で結合されています。著者らは、特定のサンプルがミニバッチに現れなかった場合の平均係数を近似する、特定の「MMD-LOO」推定量を導出し、バイアスを導入することなく分散を低減します。
  • アルゴリズム: プロセスは、現在のモデルからサンプリングし（勾配追跡なし）、特徴ベクトルを計算し、LOOベースラインを用いてMMDおよびKL係数を算出し、スコア関数勾配を介してモデルの重みを更新するという流れになります。

主な貢献

1. 定式化: 本論文は、特徴レベルの監督のみを用いて、ターゲットの特徴サンプルに適合させる特徴分布マッチングを、KL正則化されたMMDファインチューニングとして定式化しました。これにより、元のサンプル空間におけるフルターゲットサンプルの必要性を排除しています。
2. 不偏推定器: 著者らは、非微分的な特徴写像や再パラメータ化不可能なサンプラー（自己回帰モデル、離散拡散モデルなど）の使用を可能にする、MMD目的関数の不偏スコア関数勾配推定量を導出しました。
3. 汎用性: kCGMはCGMを一般化します。CGMはターゲット特徴の「平均」を一致させますが、kCGMはMMDを介してフルな「分布」を一致させます。論文では、kCGMでドット積カーネルを使用すると、CGM-relax目的関数が回収されることが述べられています。

実験結果
著者らは、小分子生成、タンパク質構造生成、および制御DNA生成の3つのドメインでkCGMを評価しています。

• 小分子抗生物質:

  • タスク: 174種類の既知の抗生物質の機能分布に適合するように、G2PT自己回帰グラフモデルをキャリブレーションする。
  • 特徴: Morganフィンガープリント（局所的な部分構造）、Murckoスキャフォールド（トポロジー）、および物理化学的記述子。
  • 結果: kCGMは、直接的なファインチューニングと比較して、ターゲット特徴分布への適合度（MMDで測定）と事前学習済みモデルからの距離との間で優れたトレードオフを達成しました。決定的なことに、抗生物質の小規模セットに対する直接的なファインチューニングは化学的に無効な分子の高い割合を招きましたが、kCGMは特徴のマッチングを改善しながら、化学的妥当性を向上させました。また、kCGMは、フィンガープリントのみで訓練された場合でも、スキャフォールドや記述子のマッチングを改善するというクロス特徴転移を示しました。

• タンパク質構造の多様性:

  • タスク: CATHデータセットの二次構造分布（αヘリックス vs βストランド）に一致するようにGenie 2（タンパク質拡散モデル）をファインチューニングする。
  • 結果: 事前学習済みモデルは主にαヘリックスを生成していました。kCGMは、CGMと比較してターゲット分布との不一致を約2倍減少させました。kCGMはKL重み $\lambda$ に対して滑らかなトレードオフ曲線を示しましたが、CGMの性能は $\lambda$ に対して感度が低く、最適化が困難でした。

• 制御DNA生成:

  • タスク: 細胞タイプ特異的な活性プロファイル（AlphaGenomeによって予測される）を持つ現実的なDNA配列を生成するように、離散拡散モデルをファインチューニングする。
  • 結果: kCGMは、47種類のシス制御トピックに関する活性プロファイルのフル分布を正常に一致させ、平均のみを一致させたCGMを上回りました。ファインチューニングされたモデルは、事前学習済みモデルが生成していた低活性な配列ではなく、実際のエンハンサーに類似した活性プロファイルを持つ配列を生成しました。

意義と主張
本論文は、特にターゲットデータセットが小さい場合や、分布の特定の側面のみを保存する必要がある場合に、kCGMが直接的なファインチューニングに代わる堅牢な選択肢を提供すると主張しています。その主な意義は以下の通りです：

• 特徴レベルの監督のみを用いて、生成モデルをユーザー指定の特徴分布に整列させること。
• 科学分野で一般的でありながら、標準的な勾配ベースのファインチューニングとは互換性のない、ブラックボックスおよび非微分的な特徴（フィンガープリント、予測器の出力など）を扱うこと。
• 直接的なファインチューニングでよく見られる失敗モードである、分布の補正を行いながらサンプル品質（妥当性）を維持または向上させること。
• 自己回帰モデル、連続空間拡散モデル、離散拡散モデルを含む、多様なアーキテクチャに適用可能であること。

著者らは、高次元の特徴におけるMMD仮説検定の呪いによる効果の低下や、トラクタブルなサンプリング・軌跡の対数確率（フローマッチングモデルを除く）の必要性などの限界についても言及しています。また、将来の研究として、元の生成空間におけるフルサンプルが存在せず、分布的な制約のみが存在する設定への拡張を提案しています。