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

TL;DR

PERK introduces a scalable, parameter-efficient method for long-context reasoning that uses test-time learning with lightweight adapters, significantly improving performance over prompt-based methods while maintaining inference efficiency.

Contribution

PERK proposes a novel nested optimization approach with low-rank adapters for scalable, memory-efficient long-context reasoning during test time.

Findings

PERK achieves up to 90% performance improvement on small models.

PERK outperforms prompt-based baselines in robustness and reasoning complexity.

It scales more efficiently at inference despite memory-intensive training.

Abstract

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.