Allen AI、密度とスコア同時推定の新モデル公開

DiScoFormerとは

密度とスコアを単一推論で同時推定
再学習なしの汎用モデル
Transformer交差注意を活用

性能と意義

100次元でKDE比誤差大幅減
未学習分布へも高精度に適応
生成AIや科学計算で再利用可能
詳細を読む

Allen Institute for AIは6月29日、データ点の集合から分布の密度とスコアを一度の順伝播で同時推定する新モデル「DiScoFormer」をHugging Face上で発表しました。従来は手法ごとに汎用性と精度のどちらかを犠牲にしていましたが、本モデルは再学習なしで両者を両立する点が特徴です。

機械学習や科学の多くの課題は、観測データから元の分布を復元する作業に帰着します。その鍵となるのが密度と、対数密度の勾配であるスコアです。スコアは確率の高い領域へ向かう方向を示し、Stable DiffusionやDALL-Eといった拡散モデルの画像生成や、ベイズ推論、プラズマなどの粒子シミュレーションを支えています。

DiScoFormerはTransformerブロックを積み重ね、サンプル全体を密度とスコアへ写像します。交差注意によりデータのない点でも評価でき、共有バックボーンに密度用とスコア用の2つの出力ヘッドを持たせました。スコアは対数密度の勾配という関係を利用し、両者のずれをラベル不要の整合性損失として推論時に数ステップ最適化することで、未知の入力へその場で適応します。

注意機構はカーネル密度推定(KDE)の一般化にあたります。1つの注意ヘッドの重みはデータ上のガウスカーネルにほぼ等しく、交差注意ブロック1つでKDEを再現できると数学的に示されました。さらに複数のスケールを同時に学習してデータに適応させ、KDEを特殊例として包含しつつ改善する設計です。

学習にはガウス混合モデル(GMM)を用いました。GMMは万能な密度近似器であり、密度とスコアの厳密な閉形式を持つため、バッチごとに新たなGMMを引いて正確な教師信号として供給でき、事実上無制限の学習例を確保できます。

性能面ではKDEを密度・スコアの双方で上回り、100次元ではスコア誤差を約6.5倍、密度誤差を37倍以上低減しました。学習時より多くのモードを持つ混合分布やラプラス分布などにも高精度を保ちます。スコア推定は生成モデルやベイズ推論、科学計算に共通する依存処理であり、再学習不要の汎用推定器は多分野のコストを一括で削減する可能性があると同社は示しています。