[論文レビュー] $k$NN Prompting: Beyond-Context Learning with Calibration-Free Nearest Neighbor Inference

Benfeng Xu, Quan Wang+4 more|arXiv (Cornell University)|Mar 24, 2023

Domain Adaptation and Few-Shot Learning被引用数 12

ひとこと要約

論文は kNN Prompting を導入します。これは勾配不要・キャリブレーション不要の手法で、近傍ベースのアプローチを用いて LM の分布上で学習データを無制限に活用し、文脈長を超える学習を可能にするfew-shot および fully supervised の性能を複数の LLM で高めます。

ABSTRACT

In-Context Learning (ICL), which formulates target tasks as prompt completion conditioned on in-context demonstrations, has become the prevailing utilization of LLMs. In this paper, we first disclose an actual predicament for this typical usage that it can not scale up with training data due to context length restriction. Besides, existing works have shown that ICL also suffers from various biases and requires delicate calibration treatment. To address both challenges, we advocate a simple and effective solution, $k$NN Prompting, which first queries LLM with training data for distributed representations, then predicts test instances by simply referring to nearest neighbors. We conduct comprehensive experiments to demonstrate its two-fold superiority: 1) Calibration-Free: $k$NN Prompting does not directly align LLM output distribution with task-specific label space, instead leverages such distribution to align test and training instances. It significantly outperforms state-of-the-art calibration-based methods under comparable few-shot scenario. 2) Beyond-Context: $k$NN Prompting can further scale up effectively with as many training data as are available, continually bringing substantial improvements. The scaling trend holds across 10 orders of magnitude ranging from 2 shots to 1024 shots as well as different LLMs scales ranging from 0.8B to 30B. It successfully bridges data scaling into model scaling, and brings new potentials for the gradient-free paradigm of LLM deployment. Code is publicly available.

研究の動機と目的

• In-context learning (ICL) の文脈長制限を、プロンプト長を超える訓練データを再利用することで解決する。
• キャリブレーション技術への依存を排除し、直接的なラベルトークンの整列ではなく LM の分布を用いて照合する。
• プロンプト長と訓練データ利用をデカップリングすることで、より大規模なデータセットでの学習を可能にし、スケーラブルなデータ活用を実現する。
• さまざまなデータセットとモデルスケールに対して、アンカー基盤の表現の堅牢性、分割戦略、効果を調査する。

提案手法

• 訓練データをデモンストレーションセット D とアンカーセット A に分割する。
• D を含むプロンプトと各アンカー x^a を組み合わせて LLM にクエリを投げ、完全な LM 分布 p(v|P,θ) を取得し、ラベル y^a を値としてデータストアにキー k_i としてキャッシュする。
• 推論時には、テストインスタンスの p(v|P_test,θ) を計算し、KL ダイバージェンス D_KL(p_test || k_i) を用いてデータストア内の最近傍を検索する。
• その後、k 最近傍のラベルを多数決風の投票で集約して y を予測する：argmax_y sum_{i in NN^k(p_test, K)} 1(y_i^a = y)。
• このフレームワークは LM の出力を固定タスク固有のラベル空間と揃えることを避けるためキャリブレーション不要であり、アンカーセットを活用して無限のデータを利用できるBeyond-context 学習をサポートする。
• 訓練は不要で、既存の LLM と訓練データを用いて機能するように設計されている。

実験結果

リサーチクエスチョン

• RQ1 calibration-free nearest-neighbor inference は few-shot 設定下で calibration-based ICL 手法と比較して性能を改善し得るか？
• RQ2 文脈長を超える訓練データが増えると、モデルサイズやタスクを越えて kNN Prompting はどの程度スケールするのか？
• RQ3 デモンストレーション-アンカー分割が性能に与える影響は何か、訓練分割の変更に対して手法はどれだけ頑健か？
• RQ4 retrieval においてラベル語だけでなく全ての LM 分布を活用することは、精度と頑健性を向上させるか、それとも低下させるか？

主な発見

• 最先端の calibration-based 手法を厳密な few-shot 設定で上回る（いくつかのタスクで最大 +7.07 ）、。
• 文脈を越えたスケーリングを強く示し、訓練データが 2 から 1024-shot に増えるとともに、0.8B から 30B パラメータの LLM で着実に改善。
• わずか 32 shots でも、ピーク時に +13.58 の平均ゲインを達成（0.8B モデル）、完全監視学習設定では最大で +18.84 まで。
• 既定の ICL に対して頑健性の改善を示し、出力の不安定さが低下し、シード間の標準偏差が減少。
• 様々なデータスケールで ICL Ensemble および標準のファインチューニングを上回り、0.8B–30B を含むモデルサイズでも競争力を維持。