DiScoFormer：分布をまたぐ密度とスコア推定のための単一トランスフォーマー
本文の状態
日本語全文を表示中
詳細モードで約6分の本文を読めます。
同じ出来事の情報源
この情報源を基点に整理
Allen AI (AI2)
Allen AI は、有限のサンプルから1回の順伝播で密度とスコアの両方を推論できる「DiScoFormer」というトランスフォーマーベースの手法を発表した。このモデルは古典的なカーネル密度推定を一般化し、再学習なしで高次元や分布外データでも高精度に動作する。
Source Article
元記事を日本語で読む
本文に関係しない購読案内、埋め込み通知、サイト内プロモーションは除いています。
機械学習や科学の多くの問題は、同じ課題に帰着します。つまり、データポイントの集合があり、その背後にある分布を復元したいというものです。どの値が一般的で、どの値が稀なのかを特定することです。その分布を特定するには、2 つの量を推定する必要があります。それは分布の密度と、次元数が増大するほどより有用となるそのスコアです。密度はヒストグラムの滑らかなバージョンであり、ポイントが集まる場所では高く、希少な場所では低くなります。スコア（対数密度の勾配）は、密度が最も急激に上昇する方向を示します。ポイントをスコアの方向に沿って移動させると、より確率の高い領域へと向かいます。
拡散ベースの生成モデル（Stable Diffusion や DALL-E などの AI 画像生成器の背後にある技術）は、ランダムなノイズから始まり、スコアを繰り返し追跡することで、そのノイズを実写のような画像に変換します。同じスコアがベイズサンプリングや、プラズマなどのシステムをモデル化するために用いられる粒子シミュレーションを駆動しています。
有限のサンプルから密度とスコアを抽出するのは困難であり、現在のツールは一般化性と精度の間でトレードオフを強いる。古典的なアプローチの一つであるカーネル密度推定（KDE）では、データ点の周囲の情報に基づいて任意の位置での密度を計算する：点が近接しておりかつ多数あるほど、密度は高くなる。これは学習を必要とせずあらゆる分布に適用可能だが、次元数が増大すると精度が急激に低下する。一方、スコアを予測するように訓練されたニューラル・スコアマッチングモデルは高次元でも精度を保つが、各モデルは特定の分布を学習する必要があり、別の分布に対しては最初から再訓練しなければならない。
我々は、データ点のセットを与えられた際に、再学習なしで単一の順方向パスを通じて分布の密度とスコアの両方を推定する新しい解決策「DiScoFormer（Density and Score Transformer）」を紹介する。このモデルは、Transformer ブロックを積み重ねた構造を用いて、サンプル全体から背後にある分布の密度とスコアをマッピングする。本モデルはクロスアテンションを活用しており、データが存在する点だけでなく任意の点での密度とスコアの評価が可能である。スコアと密度には数学的な関係があり、スコアは密度の対数の勾配として定義される。我々はこれを活用し、密度用とスコア用の 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 はメモリ不足に陥ってしまいます。また、学習データの外側にある分布に対しても優れた性能を発揮し、学習時に経験した以上のモード数を持つ混合分布や、ラプラス分布や学生 t 分布のような非ガウス型の形状においても正確な推定を維持します。KDE の主な利点は依然として速度にあり、特にデータセットが小さい場合にその優位性が際立ちます。
DiScoFormer の中で最も有望だと私たちが考える点は、スコア推定が生成モデル、ベイズ推論、科学計算など多くの分野に共通する依存関係であることです。高次元でも精度を保ち、問題ごとに再学習する必要をなくす事前学習済みでプラグイン型の推定器があれば、これらすべての分野におけるコストを同時に削減できる可能性があります。つまり、スコアと密度が現れるあらゆる場所で再利用される単一のモデルです。
詳細については、技術報告書をお読みいただくことをお勧めします。
最新の Ai2 のニュースに関する月次更新を受け取るには、購読してください。
原文を表示
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.
Subscribe to receive monthly updates about the latest Ai2 news.
関連記事
今日のまとめ
AIデイリーブリーフで今日の重要ニュースをまとめ読み