PERK: 長文脈推論をパラメータ効率的なテスト時学習として

要旨

長文脈推論では、広範でノイズの多い入力文脈から関連情報を正確に特定する必要がある。これまでの研究では、テスト時にモデルパラメータに直接文脈をエンコードするテストタイム学習を用いることで、ノイズの多い情報に対する推論を効果的に可能にすることが示されている。しかし、テストタイム学習を可能にするメタ学習手法は、メモリ使用量が非常に大きく、長文脈設定への適用が困難である。本研究では、PERK（Parameter Efficient Reasoning over Knowledge）を提案する。これは、テスト時に軽量なモデルアダプタに対する勾配更新を用いて長い入力文脈をエンコードするためのスケーラブルなアプローチである。具体的には、PERKはメタトレーニングフェーズにおいて2つのネスト化された最適化ループを採用する。内側のループは、ベースモデルのためのパラメータ効率的なメモリモジュールとして機能する低ランクアダプタ（LoRA）に迅速に文脈をエンコードする。同時に、外側のループは、更新されたアダプタを使用して、エンコードされた長文脈から関連情報を正確に想起し、推論する方法を学習する。いくつかの長文脈推論タスクにおける評価では、PERKは標準的なプロンプトベースの長文脈ベースラインを大幅に上回り、小規模モデル（GPT-2）では最大90%、最大評価モデル（Qwen-2.5-0.5B）では最大27%の平均絶対性能向上を達成した。一般的に、PERKは推論の複雑さ、長さの外挿、および文脈内の関連情報の位置に対してより頑健である。最後に、PERKはトレーニング時にはメモリ使用量が大きいが、推論時にはプロンプトベースの長文脈推論よりも効率的にスケールすることを示す。

English

Long-context reasoning requires accurately identifying relevant information in extensive, noisy input contexts. Previous research shows that using test-time learning to encode context directly into model parameters can effectively enable reasoning over noisy information. However, meta-learning methods for enabling test-time learning are prohibitively memory-intensive, preventing their application to long context settings. In this work, we propose PERK (Parameter Efficient Reasoning over Knowledge), a scalable approach for learning to encode long input contexts using gradient updates to a lightweight model adapter at test time. Specifically, PERK employs two nested optimization loops in a meta-training phase. The inner loop rapidly encodes contexts into a low-rank adapter (LoRA) that serves as a parameter-efficient memory module for the base model. Concurrently, the outer loop learns to use the updated adapter to accurately recall and reason over relevant information from the encoded long context. Our evaluations on several long-context reasoning tasks show that PERK significantly outperforms the standard prompt-based long-context baseline, achieving average absolute performance gains of up to 90% for smaller models (GPT-2) and up to 27% for our largest evaluated model, Qwen-2.5-0.5B. In general, PERK is more robust to reasoning complexity, length extrapolation, and the locations of relevant information in contexts. Finally, we show that while PERK is memory-intensive during training, it scales more efficiently at inference time than prompt-based long-context inference.

PERK: 長文脈推論をパラメータ効率的なテスト時学習として

PERK: Long-Context Reasoning as Parameter-Efficient Test-Time Learning

要旨

Support