[論文レビュー] Transformers Can Do Bayesian Inference

Samuel Müller, Noah Hollmann+3 more|arXiv (Cornell University)|Dec 20, 2021

Gaussian Processes and Bayesian Inference被引用数 24

ひとこと要約

この論文では、事前分布から抽出されたデータポイントの集合値入力を用いて、コンテキスト内学習を活用することでベイズ推論を実行する、Prior-Data Fitted Networks (PFNs) を提案する。PFNs は MCMC や NUTS よりも最大 200 倍の高速化を達成し、ガウス過程のほぼ完全な模倣を実現するとともに、テーブル回帰、少数のサンプルでの画像分類、ベイズニューラルネットワークなど多様なタスクで優れた性能を発揮する。

ABSTRACT

Currently, it is hard to reap the benefits of deep learning for Bayesian methods, which allow the explicit specification of prior knowledge and accurately capture model uncertainty. We present Prior-Data Fitted Networks (PFNs). PFNs leverage in-context learning in large-scale machine learning techniques to approximate a large set of posteriors. The only requirement for PFNs to work is the ability to sample from a prior distribution over supervised learning tasks (or functions). Our method restates the objective of posterior approximation as a supervised classification problem with a set-valued input: it repeatedly draws a task (or function) from the prior, draws a set of data points and their labels from it, masks one of the labels and learns to make probabilistic predictions for it based on the set-valued input of the rest of the data points. Presented with a set of samples from a new supervised learning task as input, PFNs make probabilistic predictions for arbitrary other data points in a single forward propagation, having learned to approximate Bayesian inference. We demonstrate that PFNs can near-perfectly mimic Gaussian processes and also enable efficient Bayesian inference for intractable problems, with over 200-fold speedups in multiple setups compared to current methods. We obtain strong results in very diverse areas such as Gaussian process regression, Bayesian neural networks, classification for small tabular data sets, and few-shot image classification, demonstrating the generality of PFNs. Code and trained PFNs are released at https://github.com/automl/TransformersCanDoBayesianInference.

研究の動機と目的

• 低データ環境におけるベイズ手法の適用という課題に取り組み、ディープラーニングを用いて効率的な事後分布の近似を可能にする。
• ガウス過程やベイズニューラルネットワークのような複雑なモデルにおける正確なベイズ推論の非可解性を克服する。
• データをサンプリング可能な任意の事前分布を用いて、事後分布の近似が可能な汎用フレームワークを構築する。
• 解析的表現や複雑な近似を必要とせず、事前のサンプリングメカニズムのみを用いて、スケーラブルで微分可能かつ柔軟なベイズ推論を実現する。
• 小規模なテーブルデータ、少数のサンプルでの学習、回帰など多様なタスクにおいて、強力なキャリブレーションと不確実性の評価を実現する。

提案手法

• PFNs は、集合値入力を備えた教師あり分類問題として事後分布の近似を再定式化する：各訓練タスクに対して、事前分布から関数をサンプリングし、(x, y) ペair の集合を収集し、1つのラベルをマスクして、モデルがマスクされたラベルを確率的に予測するように学習する。
• PFN の入力は、1つのラベルがマスクされた (x, y) ペアの集合である。モデルはトランスフォーマー構造におけるアテンションメカニズムを用いて、欠落したラベルの分布を学習する。
• PFN は、複数のサンプルタスクにおけるマスクされた予測目的に基づき、最大尤度によりエンドツーエンドで訓練される：−∑ log qθ(y_test|x_test, D_train) を最小化する。
• 推論時、実際のデータセット D_train とテストポイント x_test が訓練済みの PFN に与えられ、1回の順伝播で全予測分布 qθ*(y_test|x_test, D_train) が出力される。
• 回帰タスク用の新しい予測分布が導入され、連続出力における適切な不確実性推定が可能になる。
• この手法は柔軟な事前分布をサポートする：サンプリング可能な任意の分布が使用可能であり、例えば非可解または複雑な事前分布（ベイズニューラルネットワークやガウス過程など）も含む。

実験結果

リサーチクエスチョン

• RQ1 解析的表現を必要とせず、トランスフォーマーを用いてベイズ推論における事後予測分布の近似に効果的に利用できるか？
• RQ2 トランスフォーマーにおけるコンテキスト内学習は、多様で低データの機械学習タスクにおいて、どの程度ベイズ推論に活用可能か？
• RQ3 MCMC に NUTS を用いたベースラインや、Bayes-by-Backprop を用いた SVI と比較して、PFNs の正確性、速度、キャリブレーション性能はどの程度か？
• RQ4 PFNs は、非可解または標準的手法では近似が難しい事前分布に対しても一般化可能か？
• RQ5 アーキテクチャの選択（例：アテンションヘッド、位置エンコーディング、活性化関数）は、PFNs の事後分布近似性能にどのような影響を及えるか？

主な発見

• PFNs はガウス過程の予測をほぼ完全に模倣し、Dionis データセットでは平均 AUC 0.981、jannis データセットでは 0.996 を達成し、すべてのベースラインを上回った。
• PFN-BNN モデルは、30 個の訓練サンプルを用いた 21 個のテーブルデータセットで平均 AUC 0.855 を達成し、T-BNN (0.654) や他のベースラインと比較して、正確性とキャリブレーションの両面で顕著に優れた性能を示した。
• PFNs は MCMC に NUTS を用いた場合と比較して、推論時間で最大 200 倍の高速化を達成し、同じベンチマークで GPU 上で 13 秒の推論時間であったのに対し、NUTS では 12 時間以上かかっていた。
• この手法は強力なキャリブレーションを達成し、PFN-BNN の期待キャリブレーション誤差 (ECE) は 0.025 であり、ベースラインのロジスティック回帰モデルの 0.157 よりも顕著に低い値を示した。
• Omniglot における少数のサンプル画像分類タスクで、PFNs は平均 AUC 0.981 を達成し、KNN (0.871) や CatBoost (0.945) を上回り、少数のサンプル設定への強い一般化能力を示した。
• PFN フレームワークは多様な事前分布に一般化可能である：ガウス過程、ベイズニューラルネットワーク、その他の非可解モデルの事後分布を、事前のサンプリングメカニズムのみを用いて効果的に近似できた。