要約
近年の多くの研究では、言語モデルに検索機能を組み込むことで、入力コンテキストに取得したデータを追加するアプローチが採用されている。このアプローチが成功するためには、学習時とテスト時の両方で取得データを入力に含める必要がある。しかし、入力長が取得データのサイズに比例して増加するため、現代のTransformerモデルでは計算コストおよびメモリ使用量が二次的に増大するという問題がある。これを回避するため、我々はテスト時に標準的な学習設定に基づき、取得データ上でモデルのファインチューニングを行うシンプルな手法を採用した。本研究では、Pileデータセットのテキスト埋め込みに基づいて大規模な分散インデックスを構築した。テスト入力ごとに、その近傍データを検索し、そのテキスト上でモデルをファインチューニングする。驚くべきことに、わずか20個の近傍データ（それぞれ1回の勾配更新で）を用いてテスト時学習を行うだけで、Pileに含まれる20以上の言語モデリングタスクにおいて性能が著しく向上した。例えば、最近傍データを用いたテスト時学習は、GPT-2（小規模）とGPT-Neo（10倍以上大規模）の間の性能差を、10倍以上も縮小する効果を発揮した。ただし、十分なインデックス品質と規模が前提となる。本研究は、言語モデリングにおけるテスト時学習の初のベースラインを確立した。