Back to Articles

Kyle Wiggers's avatar

📄 Tech report: arxiv.org/abs/2511.05924

DiScoFormer One Transformer for density and score across distributions - Google -image-1

機械学習や科学分野の多くの問題は、結局同じタスクに帰結します。データポイントの集合があり、それらが従う分布を復元したい—どの値が一般的で、どの値が稀なのかを明らかにすることです。その分布を特定するには、2つの量を推定する必要があります。分布の密度と、次元数が増えるにつれてより有用になるスコアです。密度はヒストグラムの滑らかな版であり、点が密集している場所では高く、希少な場所では低くなります。スコア—対数密度の勾配—は密度が最も速く上昇する方向を示します。点をスコアに沿って動かせば、より確率の高い領域に向かいます。

拡散ベースの生成モデル（Stable DiffusionやDALL-EのようなAI画像生成の基盤技術）は、ランダムノイズから始めてスコアを繰り返し追従することで、そのノイズを現実的な画像に変換します。同じスコアがベイズサンプリングや、プラズマなどの系をモデル化するために用いられる粒子シミュレーションを駆動します。

有限サンプルから密度とスコアを抽出するのは困難であり、現在のツールは汎化性と精度のトレードオフを強います。古典的な手法の1つであるカーネル密度推定（KDE）は、任意の位置における密度をその周囲のデータポイントから計算します。点が近く多数あるほど密度は高くなります。KDEは訓練を必要とせず、任意の分布に適用できますが、その精度は次元数が増えるにつれて急激に低下します。一方、スコアを予測するよう訓練されたニューラルスコアマッチングモデルは、高次元でも精度を保ちますが、それぞれが分布を学習する必要があり、別の分布には一から再訓練しなければなりません。

私たちはDiScoFormer（Density and Score Transformer）と呼ばれる新しい解法を紹介します。これは、データポイントの集合が与えられると、分布の密度とスコアの両方を再訓練なしで単一のフォワードパスで推定する単一のモデルです。

密度とスコア推定のためのTransformerの訓練

DiScoFormer One Transformer for density and score across distributions - Google -image-2

DiScoFormerは、Transformerブロックを積み重ねた層を用いて、データサンプル全体を背後の分布の密度とスコアにマッピングします。このモデルはクロスアテンションを利用しており、データが存在する場所だけでなく、任意の点で密度とスコアを評価できます。スコアと密度には数学的な関係があります。スコアは密度の対数の勾配です。私たちはこれを活用し、密度用とスコア用の2つの出力ヘッドを持つ共有バックボーンを採用しています。

この結合はパラメータを節約するだけでなく、より多くのことを行います。スコアヘッドは、すべてのクエリで対数密度ヘッドの勾配と一致する必要があります。そのため、両者の間のギャップはラベル不要の一貫性損失となります。私たちはこれを推論時に利用します—コンテキストを固定したまま、その一貫性損失に対して数ステップの勾配降下を行うことで、DiScoFormerは分布外の入力にその場で適応します。真値の密度やスコアは必要ありません。

Transformerアーキテクチャがこのタスクに適している数学的な理由があります。カーネル密度推定には単一のバンド幅があり、各点の影響範囲は事前に固定され、どこでも同一に適用されます。アテンションはKDEの厳密な一般化です。単一のアテンションヘッドの重みが、データ上のガウスカーネルにほぼ等しいことを解析的に示しており、1つのクロスアテンションブロックだけでKDEの密度とスコアを再現できます。そこからモデルはさらに進み、複数のスケールを同時に学習し、データに適応させます。DiScoFormerは古典的手法をブラックボックスのために捨てるのではなく、KDEを特殊ケースとして取り込み、それを改善します。

DiScoFormerの訓練にはどのようなデータを使ったのでしょうか？主に2つの理由から、ガウス混合モデル（GMM）を利用しました。第一に、GMMは普遍的な密度近似器であり、十分な成分があれば本質的に任意の滑らかな分布を任意に小さな誤差で近似できます。第二に、GMMは閉形式の密度とスコアを持つため、常に正確な教師信号を得られます。これらの性質を両方活用し、バッチごとに新しいGMMを生成することで、モデルに事実上無制限の対象分布の例を与え、それぞれを指定されたGMMの正確な密度とスコアで監督します。

性能

全体として、DiScoFormerは密度とスコアの両方の推定でKDEを上回り、その差はKDEが苦戦する領域で特に大きくなります。100次元では、最適に手調整されたKDEに対して、スコア誤差を約6.5倍、密度誤差を37倍以上削減し、サンプル数を増やしても精度を向上させ続けますが、KDEはメモリ不足になります。また、訓練データから大きく外れたケース—訓練中に見たことのないモード数を持つ混合分布や、ラプラス分布やStudent-t分布のような非ガウス形状—でも精度を維持します。KDEの主な利点は、特にデータセットが小さい場合の速度です。

DiScoFormerについて最も有望だと考える点は、スコア推定が生成モデル、ベイズ推論、科学計算など、多くの分野で共通の依存関係であることです。高次元でも精度を保ち、問題ごとに再訓練する必要をなくす事前学習済みのプラグイン推定器は、それらすべてでコストを一度に削減できる可能性があります—1つのモデルが、スコアと密度が現れるすべての場所で再利用されるのです。

詳細については、技術レポートをお読みになることをお勧めします。