トランスフォーマー・ニューラル・プロセス（Kernel Regression）

要旨: ニューラル・プロセス (NP) は、確率過程の事後予測分布を直接モデル化するように設計された、急速に進化しているモデル群である。元々、計算量が O(n^3) となるために制約のあるガウス過程 (GP) の、スケーラブルな代替として開発された。最も正確な現代的な NP はしばしば GP に匹敵する性能を発揮するが、それでも注意機構によって O(n^2) のボトルネックを抱えている。私たちは、Transformer Neural Process - Kernel Regression (TNP-KR) を導入する。これはスケーラブルな NP であり、以下を特徴とする: (1) Kernel Regression Block (KRBlock)。単純で拡張可能で、パラメータ効率に優れたトランスフォーマブロックであり、その計算量は O(n_c^2 + n_c n_t)、ここで n_c と n_t はそれぞれ文脈点 (context) とテスト点の数である; (2) カーネルに基づく注意のバイアス; (3) 2 つの新しい注意メカニズム: scan attention (SA) は、メモリ効率の高い走査 (scan) ベースの注意であり、カーネルに基づくバイアスと組み合わせることで TNP-KR を翻訳不変にできる。また deep kernel attention (DKA) は Performer 風の注意で、距離バイアスを暗黙に取り込み、計算量をさらに O(n_c) にまで削減する。これらの改良により、TNP-KR の両バリアントは、単一の 24GB GPU 上で 1 分未満に 100K 個の文脈点を用いた推論を、100 万を超えるテスト点に対して実行できる。メタ回帰、ベイズ最適化、画像補完、疫学といったベンチマークにおいて、DKA を備えた TNP-KR はほぼすべてのベンチマークで Performer 版の対応手法を上回り、SA を備えた TNP-KR は最先端の結果を達成する。