技術要約：オフポリシー・ログ分散正則化を用いたボルツマン生成器の効率的な学習

1. 問題設定

物理系のボルツマン分布のような、正規化されていない確率密度からサンプリングすることは、計算物理学および化学における根本的な課題である。ターゲット分布は $p_X(x) = \tilde{p}_X(x)/Z$ と定義される。ここで $\tilde{p}_X(x) = \exp(-E(x)/k_B T)$ は未正規化密度であり、$E(x)$ はエネルギー関数、$Z$ は計算不可能な正規化定数である。

分子動力学（MD）やマルコフ連鎖モンテカルロ（MCMC）といった古典的な手法は漸近的に正しいものの、高次元またはエネルギー的に複雑な系においては、長く相関のある軌道による混合の遅れや高い計算コストに悩まされることが多い。ボルツマン生成器（BG）は、生成モデルとしての代替案を提供し、平衡分布からの独立したサンプリングを可能にする。しかし、その実用性はデータ効率によって制約を受ける。標準的な学習目的関数は通常、以下のいずれかに依存している：

1. ターゲット分布からの平衡サンプル（MDによって生成されることが多く、高コストである）。
2. 変分設定における広範なエネルギー評価。

既存のアプローチは、データセット生成時に利用可能なターゲットエネルギーラベルを十分に活用できていなかったり、この情報を効果的に組み込むために追加のオンポリシーサンプルやエネルギー評価を必要としたりすることが多い。

2. 手法：オフポリシー・ログ分散正則化 (LDR)

著者らは、追加のオンポリシーサンプルや追加のエネルギー評価を必要とせずに、ターゲットエネルギー情報をボルツマン生成器の学習に統合するフレームワークである**オフポリシー・ログ分散正則化（Off-Policy Log-Dispersion Regularization: LDR）**を提案している。

コア概念

本手法は、オンポリシーの変分学習で伝統的に使用される対数分散目的関数に基づいている。著者らは、負のモデル密度と負のターゲット密度の差（未正規化重要度重み）を表す関数 $f_\theta(x)$ を以下のように定義する：
$$f_\theta(x) = -\log q_\theta^X(x) - (-\log \tilde{p}_X(x)) = -\log q_\theta^X(x) - \frac{E(x)}{k_B T}$$
$q_\theta^X = p_X$ となる最適点において、$f_\theta(x)$ は定数（$\log Z$）となる。

ログ分散目的関数

著者らは、ログ分散目的関数を以下のログ分散目的関数のファミリーへと一般化する：
$$\mathcal{L}_\theta^{LD} = \mathbb{E}_{r_X} \left[ \rho(f_\theta(x) - c_\theta) \right], \quad c_\theta = \mathbb{E}_{r_X}[f_\theta(x)]$$
ここで $\rho(u) \geq 0$ は分散尺度（例：$|u|^p$）であり、$r_X$ は参照分布である。

• LDR-L2 ($p=2$): 分散（二次中心モーメント）を最小化し、標準的な対数分散目的関数を回収する。
• LDR-L1 ($p=1$): 第一絶対中心モーメントを最小化する。著者らは、このバリアントがエネルギーの外れ値に対してより鈍感であり、高次元のエネルギーランドスケープにおける重要度重みのヘビーテイルな性質を考慮すると、より安定した最適化を提供すると仮定している。

オフポリシー正則化戦略

重要な理論的洞察（付録A）によれば、もし参照分布 $r_X$（例：固定されたデータセット）が構成空間全体をカバーするサポートを持っていない場合、ログ分散目的関数のみを最小化することはダイバージェンス（発散）にはならない。つまり、データ多様体の外側におけるモデルを制約できず、正規化が不正確になる。

これを解決するために、LDRは標準的なデータベースの目的関数（例：Forward KL、Flow Matching、またはScore Matching）と並んで正則化項として適用される：
$$\mathcal{L}_\theta = \lambda_{data}\mathcal{L}_\theta^{data} + \lambda_{LD}\mathcal{L}_\theta^{LD}$$

• データベースの項は、グローバルな正規化と正しいターゲット分布への収束を保証する。
• LDR項は、形状正則化器として機能し、データ多様体上のログ重要度重みの分散を最小化することで、提案分布をターゲットエネルギーランドスケープに適合させる。

このアプローチにより、バイアスのかかったシミュレーションデータや補助的なエネルギーラベル付きデータを含む、任意の参照分布を使用することが可能になる。

3. 主な貢献

1. オフポリシー設定への適応: 著者らは、オンポリシーの変分学習から、固定されたエネルギーラベル付きデータセットに対する正則化として、対数分散目的関数をオフポリシー設定へと適応させた。
2. ログ分散への一般化: ログ分散損失をログ分散目的関数のファミリーへと一般化し、標準的なL2バリアントと比較して外れ値に対して優れた安定性を提供するL1バリアント（LDR-L1）を導入した。
3. 幅広い適用性: 本フレームワークは、以下の3つの異なる学習レジームにおいて有効であることが示されている：
   • 非バイアス平衡データ上での学習。
   • バイアスのかかったシミュレーションデータ上での学習（重要度サンプリングおよび混合データセットを活用）。
   • ターゲットサンプルへのアクセスがない純粋な変分学習（ターゲットエネルギー評価のみを使用）。
4. 最先端技術との統合: 本手法は、変分学習における現在の最先端技術であるConstrained Mass Transport (CMT) に正常に統合され、その効率を大幅に向上させている。

4. 実験結果

著者らは、2Dガウス混合モデル (GMM)、アラニン・ジペプチド ($d=60$)、およびアラニン・ヘキサペプチド ($d=180$) のベンチマークを用いてLDRを評価した。性能は負の対数尤度 (NLL) および有効サンプルサイズ (ESS) によって測定された。

主な知見：

• 非バイアスデータ: 非バイアスデータセットにおいて、LDRはNLLとESSの両方を大幅に改善した。アラニン・ジペプチドでは、LDRを用いた $1 \times 10^6$ サンプルの性能が、非正則化学習の $5 \times 10^6$ サンプルを上回った。これは5倍のデータ効率の向上を示している。
• バイアスデータ: 短い軌道によって生成されたバイアスのかかったデータセットにおいて、再サンプリングされた平衡データセット ($D_{IS}$) および元のバイアスデータセットにLDRを適用したところ、大幅な利得が得られた。LDRを用いることで、$1 \times 10^5$ の重要度サンプルと初期のバイアスデータセット構築のみで高性能なモデルの学習が可能となった。これに対し、非正則化のベースラインは10倍以上の重要度サンプルを必要とした。
• 変分学習（ターゲットサンプルなし）: 純粋な変分設定（CMTを使用）において、LDRは必要なターゲットエネルギー評価の数を約一桁削減した。アラニン・ジペプチドでは、CMT + LDRは、標準的なCMTが $1 \times 10^8$ 回の評価を必要としたのに対し、$1 \times 10^7$ 回の評価だけで同等の性能を達成した。アラニン・ヘキサペプチドでは、LDRはバニラなCMTが失敗した領域（ESS < 6%）において、モード崩壊を防ぎ（ESS > 30%）、性能を維持した。
• デカルト座標: 訓練をTarFlowを用いてデカルト座標で直接行う場合でも、同様の利得が再現された。これは、本手法の恩恵が内部座標表現に限定されないことを裏付けている。
• 安定性: 特に高次元の系におけるヘビーテイルな重要度重みを考慮すると、LDR-L1はLDR-L2よりも一般的に安定した学習ダイナミクス（勾配ノルムの分散が低い状態）を提供した。

5. 意義と主張

本論文は、LDRがボルツマン生成器の学習にターゲットエネルギー情報を組み込むための、シンプルかつ強力なフレームワークであると主張している。その主な意義は以下の通りである：

• データ効率: 多様な設定において、最終的な性能とデータ効率を一貫して向上させ、多くの場合、必要なサンプル数またはエネルギー評価の数を一桁削減する。
• 柔軟性: 非バイアス、バイアス、および純粋な変分的学習のすべてをサポートしており、追加のオンポリシーサンプルや追加のエネルギー評価を必要としない。
• 補完性: 著者らは、LDRがパス勾配（ターゲット勾配を使用する）のような既存の分散低減技術と補完関係にあることを指摘している。LDRはターゲットエネルギーの値（ゼロ次情報）を利用するためである。

著者らは、LDRがデータセットサイズへの依存を排除するわけではないものの、各サンプルから抽出される情報を増やすことで、大規模なデータセットの必要性を大幅に軽減できることを強調している。このフレームワークは、正規化フローや、将に拡散ベースの生成器への容易な統合を想定して設計されている。