技術的概要：分布制約への生成モデルの較正

問題定義

生成モデルは頻繁に較正不良（miscalibration）に陥り、サンプリング分布の統計量（例えば、クラス頻度、二次構造の割合、または性表現）が望ましい値から体系的に逸脱する。この問題は分野を跨いで現れる：画像モデルはモード崩壊を示し、言語モデルは社会的バイアスを強化し、タンパク質/DNA モデルは非典型的なサブ構造頻度を持つ構造を生成する。これらの誤差は、データセットの偏り、最適化されていない訓練ダイナミクス、または低温サンプリングのような事後調整に起因することが多い。

著者らは、較正を制約付き最適化問題として定式化する。基底生成モデル $p_{\theta_{base}}(x)$、統計量 $h(x)$、および目標期待値 $h^*$ が与えられたとき、制約 $E_{p_\theta}[h(x)] = h^*$ を満たしつつ、基底モデルとの間のカルバック・ライブラー（KL）ダイバージェンスを最小化する新しい分布 $p_\theta$ を見出すことが目標である：
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
この定式化は、較正制約を満たす中で基底分布に最も近いモデルを追求する。しかし、目的関数と制約の両方が高次元の生成モデルにおいて計算が困難な期待値を含むため、これを厳密に解くことは不可能である。

手法

この非計算可能性に対処するため、本論文は微調整のための 2 つの代理目的関数、すなわちCGM-relaxとCGM-rewardを導入する。両方の手法は、サンプルから導出された不偏勾配推定量を用いた確率的最適化に依存する。

1. CGM-relax（Relax Loss）

このアプローチは、硬い制約をペナルティ項に置き換える。目的関数は以下の通りである：
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• 制約違反（$L_{viol}$）: 経験統計量と目標との差の二乗 $\ell_2$ ノルム。著者らは、バイアス補正されたサンプル分散の式を用いて、この項の不偏推定量を導出する。
• KL ダイバージェンス（$L_{KL}$）: モンテカルロサンプリングを通じて測定される。
• 最適化: この手法は、不偏勾配を計算するために重要度サンプリングを使用する。具体的には、現在のモデル $p_\theta$（ストップグラディエント付き）からサンプルを抽出し、サンプルを再重み付けして KL 項と制約違反の勾配を推定する。ハイパーパラメータ $\lambda$ は、制約の満足と基底モデルへの近接性の間でトレードオフを行う。

2. CGM-reward（Reward Loss）

このアプローチは、較正問題を最大エントロピー原理と結びつける。制約付きモーメント制約のもとで基底モデルへの KL を最小化する最大エントロピー問題の解は、基底モデルの指数関数的傾き（exponential tilt）である：$p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$。

• $\alpha^*$ の推定: 真の最大エントロピー分布は計算不可能であるため、この手法はまず、基底モデルからの $N$ サンプルを用いて双対最適化問題を解くことで、最適な報酬パラメータ $\alpha^*$ を推定する。
• 損失関数: 微調整の目的関数は、現在のモデル $p_\theta$ と推定された最大エントロピー分布 $p_{\hat{\alpha}_N}$ との間の KL ダイバージェンスである：
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  これは、「報酬」が制約から導かれる報酬微調整問題と同等である。
• 最適化: CGM-relax と同様に、分散を低減するためにリーフ・ワン・アウト（leave-one-out）ベースラインを用いたスコア関数勾配推定量を使用する。

勾配推定

両方の手法は、重要度サンプリングと組み合わせたスコア関数勾配推定量（REINFORCE 型）を採用する。期待値が $\theta$ に依存する問題に対処するため、著者らは「ストップグラディエント」のトリックを使用する。これは、現在のモデルの固定バージョンからサンプルを抽出し、密度比（スコア関数）を通じて勾配を計算するものである。分散低減は、リーフ・ワン・アウトベースラインを用いて達成される。

主要な貢献

1. 定式化: 生成モデルの較正を、制約付き KL 最小化問題として定式化すること。
2. アルゴリズム: 確率的最適化を用いて制約付き問題の解を近似する、2 つの実用的な微調整アルゴリズム（CGM-relax および CGM-reward）を提案すること。
3. スケーラビリティ: これらの手法が、数百の同時制約および最大 90 億パラメータを持つモデルを処理できることを実証すること。
4. 理論的関連性: 較正、最大エントロピー、および報酬微調整の間のリンクを確立し、CGM-reward が最大エントロピー解を、CGM-relax が正則化されたバージョンをそれぞれ目指すことを示すこと。

結果

著者らは、3 つの多様な応用分野においてこれらの手法を評価した：

1. タンパク質設計（Genie2 および ESM3）:

   • 目標: 二次構造の含有量（$\alpha$ヘリックスと$\beta$ストランド）を、天然タンパク質の統計量（CATH データセット）と一致するように較正すること。
   • 制約: 最大 99 の同時制約（二変数 CDF）。
   • 結果: CGM-relax は、Genie2 において天然タンパク質分布との距離を 4 倍以上、ESM3 において 2 倍以上削減した。CGM-reward は、高次元の制約を用いた $\alpha^*$ の推定の難しさもあって、より控えめな改善にとどまった。両手法とも生物物理学的妥当性を維持したが、「設計失敗」がわずかに増加した。

2. 画像生成（TarFlow）:

   • 目標: AFHQ データセットの「野生動物」カテゴリ内のクラス割合を調整すること（例えば、ライオン、トラ、キツネ、オオカミを均等化すること）。
   • 制約: 5 つの同時制約。
   • 結果: CGM-relax は、クラス割合の総変動距離を 0.306 から 0.101 に削減した。サンプルのリアリズム（FID）はわずかに低下したが、較正されたモデルはより多様でバランスの取れた生成物を生み出した。CGM-reward は制約違反を有意に削減できなかった。

3. 言語モデリング（Gemma-2-9B-IT）:

   • 目標: さまざまな職業に関する子供向け物語の生成における性バイアスを削減すること（例えば、男性/女性の弁護士が等しい確率で登場するようにすること）。
   • 制約: 8 つの職業条件付き制約。
   • 結果: 両手法とも性不均衡を削減した。CGM-relax はベースラインに対して 3 倍以上の不均衡削減を実現し、CGM-reward は統計的に有意ではあるがより小さな改善を示した。微調整されたモデルは物語の品質を維持した（対数確率の劣化なし）。

合成実験:

• 稀な事象: CGM は、基底確率が $10^{-3}$ まで低くても、ガウス混合モデルにおける稀なモードを成功裡にアップウェイトした。
• 高次元: CGM-relax は、最大 $10^3$ の制約を効果的に処理した。対照的に、経験的最大エントロピー問題（式 7）の解の難しさにより、制約が 30 を超えると CGM-reward は実行不可能となった。

意義と主張

本論文は、モデルアーキテクチャの変更やゼロからの再訓練を必要とせずに、高次元の生成モデルを分布制約を満たすように較正するための実用的な枠組みとして CGM を提供するとしている。

• ロバスト性: CGM-relax は、最大エントロピーパラメータの推定が非計算不可能になる高次元制約設定において、特に CGM-reward よりもロバストであることが示された。
• トレードオフ: これらの手法は、較正精度と基底モデルへの忠実度（$\lambda$ によって制御される）の間で、調整可能なトレードオフを提供する。
• 限界: 著者らは、特に非常に稀な事象において残差誤差が残ることを認めている。さらに、この枠組みは現在、拡散モデル、正規化フロー、自己回帰モデルなど、尤度が計算可能なモデルに依存しており、GAN のような暗黙的モデルへの直接適用を制限している。

この研究は、分布整合のための駆動信号として制約そのものを使用することで、標準的な報酬微調整とは区別される、独立かつ有望な研究分野として較正を位置づけている。