[論文レビュー] The Surprising Effectiveness of Test-Time Training for Few-Shot Learning

Ekin Akyürek, Mehul Damani+6 more|arXiv (Cornell University)|Nov 11, 2024

Intelligent Tutoring Systems and Adaptive Learning被引用数 4

ひとこと要約

本稿では、テスト時学習（TTT）が、ARCベンチマークにおける大規模言語モデルの少样本推論を顕著に向上させることを示している。8Bパラメータモデルを用いた場合、53%の精度を達成し、プログラム合成とアンサンブルした場合に61.875%にまで上昇する。これは平均的人間の性能と一致する。本手法は、タスク固有のLoRAアダプタ、幾何的データ拡張、可逆変換に基づく自己一貫性推論を用い、未知の抽象的推論タスクにおける一般化性能を向上させている。

ABSTRACT

Language models (LMs) have shown impressive performance on tasks within their training distribution, but often struggle with structurally novel tasks even when given a small number of in-context task examples. We investigate the effectiveness of test-time training (TTT) -- temporarily updating model parameters during inference using a loss derived from input data -- as a mechanism for improving LMs' reasoning and few-shot learning capabilities. On the Abstraction and Reasoning Corpus (ARC), performing TTT with in-context examples yields up to $6 imes$ higher accuracy compared to fine-tuned baselines -- reaching $53.0\%$ on the public validation set with an 8B-parameter LM and $61.9\%$ when ensembled with program-synthesis methods, matching average human performance. On BIG-Bench Hard (BBH), TTT on in-context examples surpasses standard few-shot prompting in the $10$-shot setting by $7.3$ percentage points ($50.5\%$ to $57.8\%$). Our findings highlight the limitations of in-context learning for novel tasks and demonstrate the potential of test-time training to enhance language model adaptability.

研究の動機と目的

• テスト時学習（TTT）が、未知の抽象的推論タスクにおける大規模言語モデルの少样本推論を向上させられるかどうかを調査すること。
• Abstraction and Reasoning Corpus（ARC）における効果的なTTTを可能にする主要な設計要因を同定すること。
• ARCにおける高い性能を達成するには、記号的プログラム合成が必要であるという仮定に挑戦し、完全にニューラル的でTTTを拡張したモデルで最先端の結果を示すこと。
• TTTと既存のニューラル的またはプログラム合成アプローチとの相乗効果をARCで評価すること。
• ARCの公開バリデーションセットにおいて、完全にニューラル的で公開可能なモデルの新たなSOTAを確立すること。

提案手法

• 推論時に、ラベル付きのコンテキスト例から導出される教師あり損失関数を用いて、1つの入力または少々の例に対して勾配更新を適用することで、テスト時学習（TTT）を実装する。
• テスト時におけるパラメータ効率的なファインチューニングを可能にするために、タスク固有のLoRAアダプタを用いる。これにより、各入力インスタンスごとに高速な適応が可能になる。
• 幾何的変換（例：回転、反転）を用いて、テスト時データセットを拡張することで、インダクティブバイアスを高め、一般化性能を向上させる。
• 可逆変換に基づく推論パイプラインを採用し、複数の予測を生成し、自己一貫性投票を適用してより頑健な出力を選択する。
• 階層的自己一貫性スキームを用いて、各インスタンスごとの学習を統合することで、多様なARCタスク全体における予測の信頼性を向上させる。
• TTTパイプラインを、既存のニューラルモデル（例：BARC）やプログラム合成手法と統合し、アンサンブル性能を達成する。

実験結果

リサーチクエスチョン

• RQ1 テスト時学習におけるどの要素が、ARCベンチマークにおける言語モデルの少样本推論を向上させる上で最も重要か？
• RQ2 記号的プログラム合成を一切用いずに、純粋にニューラル的でテスト時適応を行うモデルが、ARCで人間水準の性能を達成できるか？
• RQ3 テスト時学習が、コンテキスト内学習とデータ拡張とどのように作用し合い、未知の抽象的推論タスクにおける一般化性能を向上させるか？
• RQ4 TTTは、ニューラルモデルと神経記号的アプローチの間の性能格差をどの程度埋めることができるか？
• RQ5 適切に構造化されたテスト時計算は、複雑な推論タスクにおいて、記号的探索の代替手段として実用的であると見なせるか？

主な発見

• タスク固有のLoRAアダプタと幾何的データ拡張を用いたテスト時学習（TTT）により、標準的なファインチューニング済みモデルと比較して、ARCにおける少样本推論精度が最大6倍向上した。
• 8Bパラメータの言語モデルが、TTTを用いることでARCの公開バリデーションセットで53%の精度を達成した。これは、完全にニューラル的で公開済みの手法における、前回のSOTAから約25%の向上に相当する。
• プログラム合成手法とアンサンブルした場合、TTT拡張モデルは61.875%の精度を達成し、データセット全体の平均的人間性能と一致した。
• TTTパイプラインにより、完全にニューラルなモデルが、従来はプログラム合成モデルでのみ解けるとされていたタスクの73.5%を解けるようになった。これは、体系的な推論パターンの強力な転送を示している。
• BARCのニューラルモデルにTTTパイプラインを適用したところ、精度が54.375%から53%に向上した。これは、元のTTT手法に比べて相対的に35%の向上であり、本手法の堅牢性とスケーラビリティを示している。
• TTTと可逆変換に基づく推論、自己一貫性投票の組み合わせにより、多様なARCタスク全体における予測の信頼性と一般化性能が顕著に向上した。