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

TL;DR

PERK reframes long-context reasoning as test-time learning by encoding lengthy inputs into a lightweight parameter-efficient memory adapter that is updated via gradient steps. It employs two nested meta-learning loops, with the inner loop updating a LoRA-based memory scratchpad and the outer loop optimizing reasoning performance using the encoded context. Across Needle-in-the-Haystack, open-domain multi-document QA, and Drops-in-the-Ocean benchmarks, PERK yields substantial performance gains over long-context finetuning baselines and demonstrates robust length extrapolation and positional bias resilience, while offering improved inference efficiency compared with prompt-based approaches. The work provides a scalable, practical solution for deploying long-context reasoning in real-world LLM applications, across diverse model scales and architectures.

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.

Figures (8)

• Figure 1: Meta-learning PERK for long-context reasoning. The training procedure involves a nested inner and outer loop. The inner loop optimizes the likelihood of a batch of long context segments with respect to the parameters of the LoRA-based memory scratchpad. In the outer loop, the model uses the encoded information in the memory scratchpad to answer questions. In both cases, only the memory scratchpad parameters are updated while the base LLM's parameters are frozen.
• Figure 2: Performance on Long-context Reasoning. We show the evaluation results on NIAH with BabiLong, Multi-Doc with HotpotQA & TriviaQA, and DIO with Student Records. All PERK and FT-ICR models (including Mamba) are trained on contexts with 8K tokens. When evaluated on the out-of-distribution contexts with 32K tokens, they must extrapolate to a new context length. Note that we use a substring exact match as the Accuracy metric for HotpotQA & TriviaQA.
• Figure 3: PERK with diverse model scales and families on BabiLong QA2 with an 8K train & test context length. PERK's performance generalizes across models and scales, consistently outperforms FT-ICR baselines.
• Figure 4: Test-time context length robustness evaluation on BabiLong QA1 (a) and QA2 (b) between PERK and FT-ICR on the Qwen-2.5-0.5B model. The y-axis represents the training context lengths, while the x-axis indicates various test-time context lengths. We test for both test lengths shorter than the training length, and test lengths longer than the training length. Bordered cells denote evaluation on context lengths equal to those in training. PERK shows stronger robustness across both settings.
• Figure 5: Positional Bias. Comparison of PERK and FT‑ICR on 4K and 8K contexts, on Qwen‑2.5‑0.5B. We train on problems where the relevant information appears in the beginning (Pre), middle (Mid), or end (Post) of the context, and evaluate on all three positional settings. We also train models on contexts where the relevant information is randomly located (Rnd), testing these on all four positional distributions (Pre, Post, Mid, Rnd). Bordered cells show in-distribution performances. PERK demonstrates strong positional robustness.