DiScoFormer:分布を超えた密度とスコアを扱う単一のトランスフォーマー
Hugging Face は、異なる確率分布にわたって密度推定とスコア計算の両方を処理できる汎用的なトランスフォーマーモデル「DiScoFormer」を発表した。
キーポイント
双機能アーキテクチャの実現
従来の手法では別々に扱われがちだった密度推定とスコア計算の両方を、単一のトランスフォーマーモデルで同時に処理可能にした。
分布横断的な汎用性
特定のデータ分布に依存せず、多様な確率分布にわたって安定して動作するよう設計された新しいアプローチである。
Hugging Face による発表
この技術は Hugging Face Blog を通じて公式に発表され、オープンソースコミュニティへの提供が期待される。
影響分析・編集コメントを表示
影響分析
この発表は、確率モデリングの分野において、異なるタスクを統合するアーキテクチャの重要性を示すものであり、研究コミュニティに新しい実験の枠組みを提供します。特に、データ分布が複雑で多様な実世界問題に対する汎用性の高い解決策として注目されることが予想されます。
編集コメント
密度推定とスコア計算を統合する画期的なアプローチであり、確率論的生成モデルの分野における重要な技術的進展と言えます。ただし、具体的な実装詳細やベンチマークデータが本文に明記されていないため、実際の性能評価には今後の公開コードへの注目が必要です。
- Training a transformer for density and score estimation
- Performance
📄 Tech report: arxiv.org/abs/2511.05924
機械学習や科学の多くの問題は、同じ課題に帰着します。つまり、データポイントの集合があり、その背後にある分布を復元したいというものです。どの値が一般的で、どの値が稀なのかを特定することです。その分布を特定するには、2 つの量を推定する必要があります。それは分布の密度と、次元数が増大するほど有用となるスコアです。密度はヒストグラムの滑らかなバージョンであり、ポイントが集まる場所では高く、少ない場所では低くなります。スコア(対数密度の勾配)は、密度が最も急激に上昇する方向を示します。ポイントをスコアの方向に沿って移動させると、より確率の高い領域へと向かいます。
拡散ベースの生成モデル(Stable Diffusion や DALL-E などの AI 画像生成器の背後にある技術)は、ランダムなノイズから始まり、スコアを繰り返し追跡することで、そのノイズを実写のような画像に変換します。同じスコアがベイズサンプリングや、プラズマなどのシステムをモデル化するために用いられる粒子シミュレーションを駆動しています。
有限のサンプルから密度とスコアを抽出するのは困難であり、現在のツールは一般化性と精度の間でトレードオフを強いる。古典的なアプローチの一つであるカーネル密度推定(KDE)では、データ点の周囲の情報に基づいて任意の位置での密度を計算する:点が近接しておりかつ多数あるほど、密度は高くなる。これは学習を必要とせずあらゆる分布に適用可能だが、次元数が増大すると精度が急激に低下する。一方、スコアを予測するように訓練されたニューラル・スコアマッチングモデルは高次元でも精度を保つが、各モデルは特定の分布を学習する必要があり、別の分布には最初から再訓練しなければならない。
我々は、DiScoFormer(Density and Score Transformer)と呼ばれる新しい解決策を導入する。これは単一のモデルであり、データ点のセットを与えられた場合、再訓練なしで単一の前向きパスを通じて分布の密度とスコアの両方を推定する。
密度およびスコア推定のためのトランスフォーマーの訓練
DiScoFormer は、積み重ねられたトランスフォーマーブロック層を用いて、サンプル全体を背後にある分布の密度およびスコアにマッピングする。このモデルはクロスアテンション(cross-attention)を利用しており、データが存在する場所だけでなく任意の点での密度とスコアを評価可能にする。スコアと密度には数学的な関係があり、スコアは密度の対数の勾配である。我々はこれを活用し、密度用とスコア用の2つの出力ヘッドを持つ共有バックボーンを採用している。
この結合はパラメータの節約以上の意味を持ちます。スコアヘッドは、すべてのクエリにおいて対数密度ヘッドの勾配と一致しなければならないため、両者の間のギャップはラベルなしの一貫性損失となります。私たちはこれを推論時に利用します——コンテキストを固定し、その一貫性損失に対して数回の勾配ステップを実行するだけで、DiScoFormer は即座に分布外入力に適応します。真の密度やスコアは必要ありません。
トランスフォーマーアーキテクチャがこのタスクに適しているのには数学的な理由があります。カーネル密度推定では、単一のバンド幅(各点の影響が及ぶ範囲)しか存在せず、これは事前に固定され、どこでも同一に適用されます。アテンションはこれの厳密な一般化です:私たちは解析的に示します、単一のアテンションヘッドの重みはデータ上においてほぼガウスカーネルであるため、1 つのクロスアテンションブロックだけで既に KDE の密度とスコアを再現できるのです。そこからモデルはさらに進み、複数のそのようなスケールを同時に学習し、データを適応させます。DiScoFormer は古典的手法をブラックボックスとして捨てるのではなく、代わりに KDE を特殊ケースとして包含し、それを改善するものです。
DiScoFormer の訓練にはどのようなデータを使用したのでしょうか?私たちは、2 つの主要な理由からガウス混合モデル(Gaussian Mixture Models)を利用しました。第一に、GMM は普遍的な密度近似器であり、成分数を十分に増やすことで、任意の滑らかな分布を極めて小さな誤差で近似できます。第二に、GMM は解析的な密度関数とスコア関数を持つため、常に厳密な目標値を教師信号として利用できます。私たちはこれらの特性を活用し、バッチごとに新しい GMM を生成することで、モデルに対して事実上無制限のターゲット分布例を提供し、各サンプルを対応する GMM の厳密な密度およびスコアで監督学習します。
Performance
総合的に見ると、DiScoFormer は密度推定とスコア推定の両方でカーネル密度推定(KDE)を上回っており、特に KDE が苦手とする領域でその差が顕著に広がります。100 次元のデータにおいては、その差は甚だしく、最も手調整された KDE と比較しても、スコア誤差を約 6.5 倍、密度誤差を 37 倍以上削減します。さらにサンプル数を増やすにつれて性能が向上し続ける一方、KDE はメモリ不足に陥ってしまいます。また、学習データの外側にある分布に対しても優れた性能を発揮し、学習時に経験した以上のモード数を持つ混合分布や、ラプラス分布や Student-t 分布のような非ガウス型の形状においても正確な推定を維持します。KDE の主な利点は依然として速度にあり、特にデータセットが小さい場合にその優位性が顕著です。
DiScoFormer の中で最も有望だと私たちが考える点は、スコア推定が生成モデル、ベイズ推論、科学計算など多くの分野に共通する依存関係であることです。高次元でも精度を保ち、問題ごとに再学習する必要をなくす事前学習済みでプラグイン型の推定器があれば、これらすべての分野におけるコストを一度に削減できる可能性があります。つまり、スコアと密度が現れるあらゆる場所で再利用される単一のモデルです。
詳細については、技術報告書をお読みいただくことをお勧めします。
原文を表示
- Training a transformer for density and score estimation
- Performance
📄 Tech report: arxiv.org/abs/2511.05924
Many problems in machine learning and the sciences come down to the same task: you have a collection of data points and want to recover the distribution they came from—which values are common, and which are rare. Pinning down that distribution means estimating two quantities: the distribution's density and, more useful as dimensionality grows, its score. The density is the smooth version of a histogram—high where points cluster and low where they're scarce. The score—the gradient of the log-density—points in the direction the density rises fastest: move a point along the score and it heads toward a more probable region.
Diffusion-based generative models (the technology behind AI image generators like Stable Diffusion and DALL-E) start from random noise and repeatedly follow the score, turning that noise into a realistic image. The same score drives Bayesian sampling and the particle simulations used to model systems such as plasma.
Extracting the density and score from a finite sample is challenging, and today's tools force a trade-off between generalizability and accuracy. One classical approach, kernel density estimation (KDE), computes the density at any location from the data points around it: the closer and more numerous they are, the higher the density. It needs no training and applies to any distribution, but its accuracy falls off sharply as dimensionality grows. Alternatively, neural score-matching models trained to predict the score stay accurate even in high dimensions, but each needs to learn the distribution and must be retrained from scratch for another.
We introduce a new solution called the DiScoFormer (Density and Score Transformer)—one model that, given a set of data points, estimates both the density and the score of the distribution in a single forward pass without retraining.
Training a transformer for density and score estimation
DiScoFormer maps an entire sample to the density and score of the distribution behind it using stacked layers of transformer blocks. The model utilizes cross-attention, which allows it to evaluate density and score at any point—not just where you have data. Score and density share a mathematical relationship: score is the gradient of the logarithm of density. We leverage this by having a shared backbone with two output heads, one for the density and one for the score.
This coupling does more than save parameters. The score head has to match the gradient of the log-density head at every query, so any gap between them is a label-free consistency loss. We use this at inference—hold the context fixed, take a few gradient steps on that consistency loss, and DiScoFormer adapts itself to an out-of-distribution input on the spot, no ground-truth density or score required.
There's a mathematical reason why the transformer architecture fits this task. Kernel density estimation has a single bandwidth—how far each point's influence reaches, fixed in advance and applied identically everywhere. Attention is a strict generalization of it: we analytically show that a single attention head's weights are nearly a Gaussian kernel over the data, so one cross-attention block can already reproduce KDE's density and score. From there the model goes further, learning several such scales at once and adapting them to the data. DiScoFormer doesn't discard the classical method for a black box but instead includes KDE as a special case and improves on it.
What data did we use to train DiScoFormer? We relied on Gaussian Mixture Models for two primary reasons. Firstly, GMMs are universal density approximators—with enough components they match essentially any smooth distribution to arbitrarily small error. Secondly, GMMs have closed-form densities and scores, so we always have an exact target to supervise against. We employ both of these properties by drawing a new GMM for every batch, giving the model virtually unlimited examples of target distributions and supervising each against a given GMM's exact density and score.
Performance
Across the board, DiScoFormer beats KDE at both density and score estimation, and the gap widens exactly where KDE struggles. In 100 dimensions, it isn't close—against the best hand-tuned KDE, it cuts score error by about 6.5x and density error by more than 37x, and it keeps improving as you add samples, while KDE runs out of memory. It also travels far outside its training data, staying accurate on mixtures with more modes than it ever saw during training and on non-Gaussian shapes like the Laplace and Student-t. KDE's main advantage remains speed, especially when datasets are small.
The part about DiScoFormer that we find most promising is that score estimation is a shared dependency across many fields, such as generative modeling, Bayesian inference, and scientific computing. A pretrained, plug-in estimator that stays accurate in high dimensions and removes the need to retrain per problem could cut that cost across all of them at once—one model, reused everywhere score and density show up.
We encourage you to read our technical report for more details.
関連記事
今日のまとめ
AI日報で今日の重要ニュースをまとめ読み