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

## Abstract

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 ( P arameter E fficient R easoning over K nowledge), 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.

## 1 Introduction

Large language models (LLMs) struggle with identifying and reasoning over relevant information in long, potentially noisy, input contexts (i.e., long-context reasoning [2, 34, 36, 38, 65]), including modern models with substantial context windows [32, 35, 23] such as GPT-4 (claimed context window of 128k tokens [23]).
These challenges are not surprising, as longer contexts contain more irrelevant distractor information, and may also reflect more complex reasoning problems (i.e., multiple hops), both of which increase the challenge of reasoning tasks and degrade LLM performance [13]. Furthermore, positional biases in long-context models focus disproportionate attention on the beginning or end of contexts, neglecting intermediate information [39].

In this work, we address these challenges of long-context reasoning by reframing it as test-time learning, enabling models to encode context at inference time using gradient-based parameter adaptation. Our approach processes and compresses lengthy sequences as a batch (or multiple batches) of shorter segments of the long context, rather than processing the full sequence token-by-token. The updated model can then use the newly stored parametric information (and its initial parametric knowledge) to answer relevant questions. During training, we operate two nested optimization loops: the inner loop, which is also used at test time, encodes the long context segment batches into model parameters, while the outer loop learns to retrieve and reason over relevant information from the updated parameters that now encode the long context.

However, meta-learning algorithms for training test-time learning methods entail prohibitive memory overhead [17, 13, 25], particularly when updating the complete set of model parameters, because they require backpropagating through a complete multi-step optimization trajectory.
To reduce the memory overhead, we introduce Parameter-Efficient Reasoning over Knowledge (PERK), a lighter-weight meta-learning algorithm that keeps training-time memory costs low by:
(1) internalizing the context into a Low-Rank Adapter (LoRA [24]), and (2) backpropagating the outer loop over only a truncated trajectory from the final few inner-loop adaptation steps [54]. By lowering the memory cost, PERK scales more efficiently at training, supporting larger models and longer contexts.

We benchmark PERK relative to models finetuned for classical in-context reasoning (FT-ICR) on multiple Needle-in-a-Haystack benchmarks, demonstrating PERK’s advantages. A small GPT-2 model (with 127M parameters and a native 1k context window) can be scaled by PERK to reason effectively over 8k token contexts, achieving a 60% absolute performance gain over fine-tuning GPT-2 for in-context reasoning. For larger models (Qwen-2.5-0.5B) with a native 32k context window, PERK surpasses FT-ICR by a 27% absolute improvement.
As relevant facts in Needle-in-a-Haystack problems are often easily distinguishable from distractor text, we also propose a more challenging evaluation setting where relevant facts are surrounded by distributionally-similar distractor facts (i.e., Drops-in-the-Ocean). In this more complex setting, PERK models again outperform all baselines, including FT-ICR across the GPT-2, Qwen, and Mamba model families.

Further analyses reveal PERK’s strong test-time length extrapolation ability. Models trained with PERK on 1k-token contexts successfully extrapolate performance to unseen 32k-token contexts with only a 42% degradation, far better than FT-ICR’s 52% decline. In addition, unlike FT-ICR’s dramatic drops in performance (up to 90%) when relevant information shifts positions at test time, PERK robustly generalizes across contexts where relevant information is variably distributed along context positions. Finally, our inference complexity analysis shows that PERK’s memory footprint scales more efficiently than FT-ICR, as does its wall-clock runtime on extremely long contexts (> 32k tokens).

## 2 Background and Preliminaries

Causal Language Models
In this work, we focus on *causal language models* (CLM) $f_{\bm{\theta}}$ with a decoder-only transformer [51] architecture $f$ and parameters ${\bm{\theta}}$. The common objective in autoregressive-style training of a CLM is the *token-level negative log-likelihood* (NLL). For a collection of text sequences ${\mathcal{D}}=\{x\}$ with total length $N_{\mathcal{D}}$, this is defined as:

$$ $\mathcal{L}_{\text{NLL}}({\bm{\theta}},{\mathcal{D}})\;=\;-\;\frac{1}{N_{ \mathcal{D}}}\sum_{x\in{\mathcal{D}}}\;\;\sum_{t=1}^{\mathcal{T}}\log p_{\bm{ \theta}}\bigl{(}x_{t}\;\big{|}\;{\bm{x}}_{<t}\bigr{)}.$ (2) $$

For training, ${\bm{\theta}}$ is optimized iteratively via an optimization algorithm $\mathcal{A}lg:(\mathcal{L},{\mathcal{D}},{\bm{\theta}},\bm{h})\mapsto{\bm{
\theta}}^{\prime}$ where $\mathcal{L}$ is the objective function, ${\mathcal{D}}$ is the training corpus, ${\bm{\theta}}$ are the initial parameters, ${\bm{\theta}}^{\prime}$ are the obtained parameter values, and $\bm{h}$ are the hyperparameters, such as optimizer choice and training length.

Reasoning with CLMs
We denote *reasoning* problems as tuples $r=({\mathcal{K}},{\bm{q}},y)$ of context ${\mathcal{K}}$, question ${\bm{q}}$, and target response $y$. With the pre-trained large CLMs’s ability to reason over facts in input contexts with reasonable success [14], prompt-based *in-context reasoning* (ICR) has become popular for automating reasoning tasks. In this setting, ${\mathcal{K}}$ and ${\bm{q}}$ are combined via a template into a single prompt ${\bm{x}}(r)$. The model’s parameters are optimized to predict the tokens of the correct answer:

$$ $\mathcal{L}_{\text{reason}}({\bm{\theta}},{\mathcal{D}})\;=\;-\;\frac{1}{N_{{ \mathcal{D}}_{y}}}\sum_{r=({\mathcal{K}},{\bm{q}},y)\in{\mathcal{D}}}\;\sum_{t =1}^{\mathcal{T}_{y}}\log p_{\bm{\theta}}\bigl{(}y_{t}\,\big{|}\,{\bm{x}}(r),y _{<t}\bigr{)}$ (3) $$

where ${\mathcal{D}}$ is the training corpus, $\mathcal{T}_{y}$ is the length of an example answer, and $N_{{\mathcal{D}}_{y}}$ is the total length of all answers in the dataset.

Reasoning as Test-time Learning
In this setting, models simulate reasoning by first encoding the given context ${\mathcal{K}}$ into a CLM’s parameters ${\bm{\theta}}$, and then using the updated parameters $\phi({\bm{\theta}},{\mathcal{K}})$ to generate a response to the given question ${\bm{q}}$ [13]. The context is encoded by adapting the model’s parameters using a CLM objective over the context:
$\phi({\bm{\theta}},{\mathcal{K}})=\mathcal{A}lg(\mathcal{L}_{\text{NLL}},{
\mathcal{K}},{\bm{\theta}},\bm{h})$.
We then use $f_{\phi}$ to generate answers for any relevant questions ${\bm{q}}$.

Test-time learning methods for reasoning stem from a key observation: CLMs can store and retrieve diverse information in their parameters from pre-training [52, 5]. Yet these same models struggle with long-context problems [23, 32, 65, 35], despite the contexts containing far less information than what they’ve seen in pre-training. CLMs may reason better with knowledge in their parameters than in context, motivating methods that encode new context parametrically for reasoning.

Training a CLM for Test-time Learning
Given a distribution $R$ of possible reasoning problems, the optimal parameters ${\bm{\theta}}^{*}$ for a test-time learning CLM minimize the expected reasoning loss of the CLM’s adaptation to each reasoning problem $r\sim R$:

$$ ${\bm{\theta}}^{*}=\operatorname*{arg\,min}_{{\bm{\theta}}}\;\mathop{\mathbb{E} }_{r\sim R}\bigg{[}\mathcal{L}_{\text{reason}}\Bigl{(}\phi({\bm{\theta}},{ \mathcal{K}}),\{({\bm{q}},y)\}\Bigr{)}\bigg{]}$ $$

where $\phi({\bm{\theta}},{\mathcal{K}})$ is a CLM adaptation of ${\bm{\theta}}$ to the context ${\mathcal{K}}$: $\phi({\bm{\theta}},{\mathcal{K}})=\mathcal{A}lg(\mathcal{L}_{\text{NLL}},{
\mathcal{K}},{\bm{\theta}},\bm{h})$.

Meta-learning algorithms, such as MAML [17], can be used for training effective test-time learning methods. In MAML, we obtain a meta-model $f_{{\bm{\theta}}}$ via bi-level optimization: the *outer loop* samples batches of reasoning problems from a training dataset of reasoning problems, the *inner loop* obtains the adaptations of ${\bm{\theta}}$ to each problem’s context ${\mathcal{K}}$, and the outer loop then optimizes the meta-parameters based on the performance of each $f_{\phi({\bm{\theta}},{\mathcal{K}})}$ on the corresponding questions. The outer loop optimization backpropagates through the entire inner loop, an operation that becomes increasingly expensive with more parameters and inner loop update steps. This approach—adapting a model with a CLM objective on context in the inner loop while meta-optimizing a reasoning objective in the outer loop—was first introduced by Reckoning [13], which demonstrated its potential.

## 3 PERK : Meta-Learning Parameter-Efficient Reasoning over Knowledge

While prior work has demonstrated the potential of test-time learning as a reasoning paradigm [13], these approaches are hindered by high memory overhead during training. The underlying meta-learning approach, MAML [17], requires backpropagating through the expected test-time learning phase (the *inner loop*). This gradient unrolling process involves computing higher-order derivatives, which are known to have substantial memory requirements [57], thereby limiting the number of model parameters that can be optimized and the effective sequence length that can be encoded.

To scale the training of larger and more capable test-time learning models, we propose PERK
(Parameter Efficient Reasoning over Knowledge), which uses a parameter-efficient test-time adaptation to reduce both the size and amount of gradient unrolling required during optimization. Specifically, in the inner and outer loops, we only optimize a parameter-efficient adapter, LoRA [24], to encode the contexts, and we apply truncated gradient unrolling [54] to reduce the memory cost of backpropagating through the inner loop. We visualize the PERK training process in Figure [1](https://arxiv.org/html/2507.06415v1#S3.F1).

Figure: 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.
Refer to caption: extracted/6607134/figures/training.png

### 3.1 Parameter-Efficient Meta-Learning with Low-Rank Adaptation

Similar to Reckoning [13], PERK uses bi-level optimization to prepare a model for reasoning via test-time adaptation to any given context. At inference time, a model is adapted to the context via CLM optimization on that context, after which the adapted model is used to answer any relevant questions.
However, unlike Reckoning, which directly finetunes the full model at training and test time, PERK only adapts a low-rank adapter [24]. During training, PERK encodes a context into the adapter through gradient-based updates, leaving the base model parameters unchanged.
We denote by ${\bm{\theta}}_{\text{base}}$ and ${\bm{\theta}}_{\text{adapter}}$ the parameters of the base model and the LoRA adapter. ${\bm{\theta}}_{\text{base}}$ is never optimized, while ${\bm{\theta}}_{\text{adapter}}$ is meta-learned to a good initial state for its eventual test-time adaptations.

Test-Time Learning and Inner Loop Training
We set the test time learning objective (thus also the inner loop objective) to causal language modeling of the context: $\mathcal{L}_{\text{NLL}}({\mathcal{K}},({\bm{\theta}}_{\text{base}},{\bm{
\theta}}_{\text{adapter}}))$, and set the adaptation algorithm $\mathcal{A}lg(\mathcal{L}_{\text{NLL}},{\mathcal{K}},({\bm{\theta}}_{\text{
base}},{\bm{\theta}}_{\text{adapter}}),\bm{h})$ to only update ${\bm{\theta}}_{\text{adapter}}$. The adaptation algorithm computes the gradient $\nabla\mathcal{L}_{\text{NLL}}$ on a batch of sub-sequences from the full context sequence ${\mathcal{K}}$. This compression as a parallel batch (or multiple batches) allows us to process lengthy sequences beyond the model’s context window. The exact manner by which ${\mathcal{K}}$ is broken into batches is a hyperparameter of the algorithm, and for efficiency reasons, we set it to one batch during training.

Outer Loop Training    In the outer loop, we optimize the meta parameters ${\bm{\theta}}_{\text{adapter}}$ to learn from $R$, the corresponding distribution of reasoning problems. The optimal ${\bm{\theta}}_{\text{adapter}}$ minimize the expected reasoning loss of the adapter’s adaptation to each reasoning problem $r=({\mathcal{K}},{\bm{q}},y)\sim R$:

$$ $\displaystyle{\bm{\theta}}_{\text{adapter}}^{*}=\operatorname*{arg\,min}_{{\bm {\theta}}_{\text{adapter}}}\;\mathop{\mathbb{E}}_{r\sim R}\bigg{[}\mathcal{L}_ {\text{reason}}\Bigl{(}{\bm{\theta}}_{\text{base}},\,\phi_{\text{adapter}}(({ \bm{\theta}}_{\text{base}},{\bm{\theta}}_{\text{adapter}}),{\mathcal{K}}),\,\{ ({\bm{q}},y)\}\Bigr{)}\bigg{]},$ (4) $$

where $\phi_{\text{adapter}}(({\bm{\theta}}_{\text{base}},{\bm{\theta}}_{\text{
adapter}}),{\mathcal{K}})=\mathcal{A}lg(\mathcal{L}_{\text{NLL}},{\mathcal{K}}
,({\bm{\theta}}_{\text{base}},{\bm{\theta}}_{\text{adapter}}),\bm{h})$ is the parameter-efficient CLM adaptation of ${\bm{\theta}}_{\text{adapter}}$ to the context ${\mathcal{K}}$ obtained via the test-time adaptation. We may use any optimization algorithm to optimize ${\bm{\theta}}_{\text{adapter}}$, and in our experiments, we employ AdamW [40].

### 3.2 Scalable Meta-Learning with Truncated Gradient Unrolling

To optimize Equation ([4](https://arxiv.org/html/2507.06415v1#S3.E4)) with gradient-based methods, we need to differentiate through $\mathcal{A}lg$, which involves higher-order derivatives and saving the complete trajectory to compute the meta-gradient (i.e., the gradient of $\mathcal{L}_{\text{reason}}$).
However, this creates high memory costs, capping the method’s applicability. To reduce memory costs, we truncate the backpropagation for the meta-gradient computation. We briefly describe this approach here, presenting the case where $\mathcal{A}lg$ optimizes ${\bm{\theta}}_{\text{adapter}}$ via direct gradient descent for simplicity (though we use more complicated optimizers for the inner loop in practice).

Gradient Unrolling (GU)    We denote an N-step inner-loop optimization’s intermediate states by:

$$ $\phi_{\text{adapter}}^{(0)}\;=\;{\bm{\theta}}_{\text{adapter}},\quad\phi_{ \text{adapter}}^{(n+1)}=\phi_{\text{adapter}}^{(n)}-\alpha\,g^{(n)},\quad\phi_ {\text{adapter}}^{*}\;=\;\phi_{\text{adapter}}^{(N)},$ $$

where
$g^{(n)}\;=\;\nabla_{\phi_{\text{adapter}}^{(n)}}\!\mathcal{L}_{\text{NLL}}
\bigl{(}{\bm{\theta}}_{\text{base}},\,\phi_{\text{adapter}}^{(n)},\,{\mathcal{
K}}\bigr{)}$ is the gradient at step $n$, and $\alpha$ is the learning rate.

In vanilla MAML, to compute the gradient of the outer-loop loss $\mathcal{L}_{\text{reason}}$ with respect to the meta-parameters
${\bm{\theta}}_{\text{adapter}}$, one backpropagates through all $N$ inner steps. By the chain rule, this equates to accumulating the Jacobian $J^{(n)}=\frac{\partial\phi_{\text{adapter}}^{(n+1)}}{\partial\phi_{\text{
adapter}}^{(n)}}$ from each step:

$$ $\nabla_{{\bm{\theta}}_{\text{adapter}}}\mathcal{L}_{\text{reason}}\;=\;\frac{ \partial\mathcal{L}_{\text{reason}}}{\partial\phi_{\text{adapter}}^{(N)}}\;J^{ (N-1)}J^{(N-2)}\dotsm J^{(0)}\;\frac{\partial\phi_{\text{adapter}}^{(0)}}{ \partial{\bm{\theta}}_{\text{adapter}}}\;=\;\frac{\partial\mathcal{L}_{\text{ reason}}}{\partial\phi_{\text{adapter}}^{(N)}}\;\prod_{n=0}^{N-1}\!J^{(n)},$ (5) $$

where $J^{(n)}$ can be computed via:
$J^{(n)}=I-\alpha\,H^{(n)},\;H^{(n)}=\nabla^{2}_{\phi_{\text{adapter}}^{(n)}}\!
\mathcal{L}_{\text{NLL}}.$
Each Jacobian $J^{(n)}$ requires storing the computation graph of step $n$ in memory, so the cost of the backward pass scales linearly with $N$, creating a high memory requirement for gradient unrolling.

Truncated Gradient Unrolling (TGU)
To reduce memory, we follow [54] and run the inner loop for all $N$ specified update steps, but only store the computational graph for the last $T\ll N$ steps.
This substantially reduces the computational memory requirements of the method, at the cost of slightly less accurate meta-gradients for the outer-loop optimization of ${\bm{\theta}}_{\text{adapter}}$.
To implement this approximation, we do not retain the inner loop optimization’s computation graph until the start of the *retention window*: the last $T$ inner-loop optimization steps.
Since all Jacobians before step $n_{\text{start}}$ are treated as constants, the Jacobian product in Equation ([5](https://arxiv.org/html/2507.06415v1#S3.E5)) is effectively cut:

$$ $\nabla_{{\bm{\theta}}_{\text{adapter}}}\mathcal{L}_{\text{reason}}\approx\frac {\partial\mathcal{L}_{\text{reason}}}{\partial\phi_{\text{adapter}}^{(N)}}\; \underbrace{\prod_{n=N-T}^{N-1}(J^{(n)})}_{\text{last $T$ steps}}\;\;{\color[ rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5} \pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}{\cancel{\prod_{n=0}^{ N-T-1}(J^{(n)})}}}.$ (6) $$

## 4 Experiments

In this section, we empirically evaluate PERK’s long-context reasoning capabilities in two challenging scenarios: reasoning over (1) Needles-in-a-Haystack (NIAH, [32]) and (2) Drops-in-an-Ocean (DIO), a novel long context evaluation setting we propose.

Baselines   We compare our method to two long context reasoning adaptations applied to multiple models. In particular, we finetune transformer [60] and Mamba [21] models for prompting-based in-context reasoning (FT-ICR), as in conventional long-context finetuning [18, 19]. We apply FT-ICR to GPT-2 (127M parameters) [51], a small model with a limited context window (1024 tokens) and no long-context pretraining, Qwen-2.5-0.5B [50] (with a 32k token context window), and to two Mamba [21] models with sizes of 130M and 1.4B parameters.
We refer to these baselines as FT-ICR (GPT-2), FT-ICR (Qwen), FT-ICR (Mamba-130M), and FT-ICR (Mamba-1.4B).
Whenever evaluating tasks with contexts longer than the 1024 token limit of GPT-2, in order to effectively fine-tune it, we first provide GPT-2 with a series of additional positional embeddings instantiated using the interpolation method proposed in [29]. The Qwen models, meanwhile, all have sufficient context length for our experiments, and the Mamba models are designed to work with sequences of any length.
In our experiments, we refer to models trained with PERK as PERK (GPT-2) and PERK (Qwen) depending on their base model. With only the LoRA parameters being trainable, PERK (GPT-2) contains 38M trainable parameters, and PERK (Qwen) has 140M trainable parameters.

### 4.1 Reasoning Over Needles-in-a-Haystack

Figure: Figure 2: Performance on NIAH with BabiLong and DIO with Student Records. All models are trained and tested on sequences ranging from 1K to 8K tokens. All PERK models are trained to first generate the supporting facts relevant to a query, followed by the final answer prediction. In contrast, baseline models directly generate the answer, as this approach yields better performance for them in this setting. The number of trainable parameters for each method is indicated in the legend.
Refer to caption: extracted/6607134/figures/results.png

#### Setup

We first evaluate PERK on reasoning over Needles-in-a-Haystack using the BabiLong framework [32] for assessing models’ reasoning abilities across relevant and distractor facts scattered in very long documents. We mainly focus on tasks involving single-hop (QA1), two-hop (QA2), and three-hop (QA3) reasoning. We use the framework to construct the tasks with long natural documents sampled from FineWeb-Edu [41], ensuring the content postdates our base model’s knowledge cutoff.

#### Results

In the top row of Figure [2](https://arxiv.org/html/2507.06415v1#S4.F2), we show the evaluation results on BabiLong.
PERK outperforms all baselines in all cases we test (across all models), including when compared to larger models. In particular, PERK (GPT-2) outperforms even FT-ICR (Qwen) and FT-ICR (Mamba-1.4B) for all task complexities and context lengths, with an average 20% performance gain. PERK (Qwen) achieves the highest accuracy across all task complexities (QA1, QA2, QA3) and context lengths (1K to 8K). PERK is also more robust to increasing context length and task complexity than the baselines. At best, PERK (Qwen) loses almost no accuracy with growing context length (for QA1 and QA2), and at worst (for QA3), it loses ${\sim}10\%$. When averaging across context lengths, PERK (Qwen)’s performance gain over FT-ICR (Qwen) increases from 10% to 41% when the number of reasoning hops increases from 1-hop (QA1) to 3-hop (QA3).

### 4.2 Reasoning Over Drops-in-the-Ocean

Needle-in-a-Haystack (NIAH) datasets for long-context evaluation typically have a fundamental limitation: the target information often exhibits stylistic differences from surrounding irrelevant text, creating an artificially simple test for identifying relevant information. In practical scenarios, relevant and irrelevant information will likely be distributionally similar, increasing the difficulty of identifying critical facts. As a result, we propose Drops-in-the-Ocean (DIO), a new evaluation setting that forms long contexts from structurally similar documents.

#### Setup

We construct a synthetic dataset, Student Records, where each context simulates a database containing multiple student records. Each record includes several attributes: ID, name, school, major, and grade. Unique IDs and names are generated for each database instance. Schools and majors are sampled with replacement from predefined pools, and grades are integers sampled uniformly from the range [0, 100]. Context length scales directly with the number of student records included. Inspired by BabiLong, we define evaluation tasks of varying complexity: Recall (retrieving attributes for a specific ID), Relation (comparing attributes between two IDs), and Aggregate (calculating the maximum, minimum, and average grade across all students). Crucially, to test generalization, specific entities (e.g., IDs, names) and attribute value combinations encountered in the test set are disjoint from those used during training. Examples are available in the Appendix.

#### Results

Results on the Student Records in Figure [2](https://arxiv.org/html/2507.06415v1#S4.F2) mirror the trends seen in BabiLong. FT-ICR (GPT-2) fails to reason in the DIO setting with 0% accuracy or random performance (for Relation, the random performance is 50%). FT-ICR (Qwen) appears to be the strongest baseline, especially on recalling with a 1K context, where it briefly matches PERK (Qwen)’s near-perfect accuracy. PERK models again outperform all baselines. PERK (Qwen) and PERK (GPT-2) consistently achieve high accuracy (often >85%) and are more robust to increasing context length, even on the more challenging aggregation task with the 8k length, which requires the model to aggregate information across all 8k tokens. With Qwen as a base model, PERK achieves a 9% average gain over FT-ICR.

## 5 Analysis

### 5.1 Test-Time Length Generalization

We evaluate PERK’s ability to generalize to longer inputs than those seen in training [48].

#### Setup

We evaluate using the BabiLong tasks, which allow us to generate contexts of arbitrarily long lengths, and train a Qwen-2.5-0.5B base model using PERK on fixed-length contexts ranging from 1K to 8K tokens (as well as an FT-ICR baseline). Subsequently, we evaluate each model on contexts of up to 32K tokens, measuring each approach’s ability to generalize to new, longer lengths (i.e., extrapolation), as well as shorter contexts (i.e., interpolation). Overall, we train on contexts of increasing lengths (1K to 8K tokens) and test on contexts ranging from 1K to 32K tokens.

#### PERK generalizes to new context lengths at test time.

Figure [5](https://arxiv.org/html/2507.06415v1#S5.F5) shows that PERK outperforms FT-ICR on the BabiLong QA1 (single-hop) and QA2 (two-hop) tasks for all generalization settings.
For extrapolation, FT-ICR’s performance drops drastically as the inference context length exceeds the training context length, especially for shorter training contexts.

Figure: Figure 3: Test-time length extrapolation beyond 32K on BabiLong QA1 (a) and QA2 (b). Both PERK and FT-ICR are trained on 8K-token sequences. The context length for inference grows from 64K to 128K. PERK extrapolates substantially better than FT-ICR.
Refer to caption: extracted/6607134/figures/length-generalization_v3.png

While PERK’s accuracy also decreases (from a higher starting point), the degradation is dampened at longer test lengths. For QA1 (Figure [5](https://arxiv.org/html/2507.06415v1#S5.F5)a), PERK trained on 1K sequences shows only a 42% drop on 32K sequences while FT-ICR exhibits a 52% decline. When trained on 8K sequences, PERK drops only 5% on 32K sequences, while FT-ICR still exhibits a 32% drop. Similar dynamics are observed on the more challenging QA2 task (Figure [5](https://arxiv.org/html/2507.06415v1#S5.F5)b).
To stress-test the limit of PERK models’ length extrapolation ability, we test on sequences with context lengths beyond the context window of the Qwen-2.5 model – 32K tokens in Figure [3](https://arxiv.org/html/2507.06415v1#S5.F3). Specifically, we test on sequences with tokens increasing from 64K to 128K, and demonstrate that PERK consistently extrapolates to longer contexts than the FT-ICR baseline.
We acknowledge that with 128K token context windows, PERK’s accuracy experiences a noticeable drop to 61.4% for QA1 accuracy and 44.4% for QA2 accuracy. However, these scores remain substantially better than the 0% QA1 and QA2 absolute performance of the FT-ICR baselines at this same context length.

Meanwhile, for interpolation, FT-ICR performance drops on shorter contexts when finetuned on longer ones, matching observations from prior work [19]. In contrast, PERK maintains its performance and even increases performance for the QA2 task.
Overall, PERK demonstrates strong length generalization, extrapolating to longer contexts while preserving high performance on shorter ones.

Figure: Figure 4: Test-time context length generalization evaluation on BabiLong QA1 (a) and QA2 (b): comparison 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 interpolation (test lengths shorter than the training length) and extrapolation (test lengths longer than the training length). Bordered cells denote the boundary: evaluation on context lengths equal to those in training. PERK shows stronger generalization across both settings.
Refer to caption: extracted/6607134/figures/length-generalization_updated.png

### 5.2 Robustness to Positional Biases

Prior work has shown that the position of relevant information in long contexts affects a language model’s ability to utilize that information [39]. In this experiment, we test PERK’s robustness to positional biases by evaluating its test-time performance on contexts where the position of relevant information is distributed differently from those seen in training.

#### Setup

We use API documents from [45], where each document has an associated API call, and create long contexts containing multiple API-call-to-document pairs. In this task, the model must retrieve the correct API call based on user instructions and documents. Using fixed 4k and 8k token contexts, we train PERK and FT-ICR models on contexts with the target document at varying positions and test them on contexts distributed across all positions.

#### PERK generalizes regardless of information position in the context.

Figure [5](https://arxiv.org/html/2507.06415v1#S5.F5) compares PERK to FT-ICR under different position distributions (at train and test time) and context lengths. When trained on contexts with relevant documents randomly located (Rnd) throughout the context, PERK outperforms FT-ICR by 24%. Furthermore, FT-ICR shows large performance drops when the relevant document consistently appears at different locations (Pre, Mid, and Post) at test time. Meanwhile, shifting relevant positions at test-time shows minimal effect (within 1-2%) on PERK.

To further stress-test positional generalization, we also force relevant documents into specific positions during training and testing, spanning the full context. We find that FT-ICR easily overfits to the position pattern, completely failing to generalize to test-time position changes (performance drops to close to 0% when the position shifts at test time).

Although PERK underperforms FT-ICR when the train-test positions are the same and the beginning (Pre) or end (Post) of the context, performance remains consistent regardless of the positional distribution shift between train and test time.
We attribute PERK’s robustness in this setting to the fact that documents in the long context are encoded as members of a batch into the parameter space, limiting the overall impact of absolute position in the sequence. Meanwhile, FT-ICR directly aligns specific positions with answers during training.

Figure: Figure 6: (a): Peak GPU memory usage during training with context lengths ranging from 1K to 8K tokens for Reckoning and PERK. While PERK successfully scales to 8K tokens, Reckoning encounters out-of-memory (OOM) errors at shorter context lengths. (b): Memory footprint and wall-clock runtime during inference as context length increases (up to 128K tokens), comparing PERK with FT-ICR. Curves that terminate before 128K indicate that the method failed with an OOM error at longer context lengths, preventing further measurement. PERK demonstrates more efficient scaling in both memory and runtime, particularly for extremely long sequences.
Refer to caption: extracted/6607134/figures/efficiency_results.png

### 5.3 Training and Inference Complexity Analysis

We first show PERK’s training scalability compared to prior test-time learning method Reckoning [13] by measuring each method’s peak training memory across different context lengths. We then evaluate PERK’s training efficiency on long contexts by measuring its inference memory cost and run-time, compared to in-context reasoning with finetuned models.

#### Setup

We complete all runs on a single Nvidia H100 GPU with 93 GB of VRAM. We perform four gradient update steps for the test-time and inner loop adaptations. For training, we have four truncation levels: truncate 0 (basically gradient unrolling without truncation), 1, 2, and 3 steps. For inference, both methods generate 64 tokens using greedy decoding with Huggingface’s text generation pipeline. We measure runtime complexity in seconds (wall-clock time) and memory complexity in gigabytes. Both methods undergo 10 hardware warm-up iterations.

As explained in Section [3](https://arxiv.org/html/2507.06415v1#S3), PERK splits long sequences into batches of shorter sequences. Here, each batch has an effective context length defined as 128 tokens. Since PERK performs gradient-based encoding over these batches, we can apply gradient accumulation to trade runtime for reduced memory consumption. We thus analyze PERK across varying gradient accumulation steps, ranging from 2 to 16. Note that for the minimum context length of 1024 tokens, which yields a batch of 8 sequences, gradient accumulation with 16 steps is infeasible as it exceeds the batch size.

#### PERK scales more efficiently in memory and runtime on extremely long contexts.

Figure [6](https://arxiv.org/html/2507.06415v1#S5.F6)a shows that, as the training context length grows from 1K to 8K tokens, Reckoning quickly runs into Out-Of-Memory (OOM) errors at 2K, while PERK with increasing truncation steps continues to fit the maximum GPU memory. We successfully validate PERK’s stronger scalability on long contexts.

Figure [6](https://arxiv.org/html/2507.06415v1#S5.F6)b for inference complexity shows that PERK provides more efficient memory and runtime scaling for extremely long contexts compared to FT-ICR. While FT-ICR is initially more efficient, its memory and runtime grow rapidly, leading to OOM errors at a context length of 128K. In contrast, PERK can manage the long sequences through gradient accumulation, which, while increasing runtime, reduces the memory footprint. For instance, at 8K tokens, increasing the accumulation steps from 1 to 16 reduces memory usage from 35.2GB to 5.9GB, although the runtime increases from 1.9s to 8.5s. Crucially, even with this runtime trade-off, PERK runs faster than FT-ICR at very long contexts where both can operate. At 64K tokens, FT-ICR takes 32.6s and 55.7GB, whereas PERK with 16 steps of accumulation completes in 11.4s, using only 19.6GB. With 8 steps, it is even faster at 9.7s, using 35.2 GB. Ultimately, at 128K tokens, where FT-ICR fails, PERK with 16 steps successfully processes the context using 35.2GB in 20.9s, showing that PERK provides a practical path to handle extreme context lengths efficiently in both memory and runtime when compared to standard approaches.

## 6 Related Work

#### Test-Time Learning

MAML [17] introduces optimization-based model-agnostic meta-learning for few-shot adaptation. [13] applies MAML to GPT models, integrating knowledge into model parameters through test-time gradient updates. Some methods achieve test-time learning in the model architecture. [58] proposes RNN layers with learnable hidden states updated via self-supervised gradients on test inputs. Titans [8] incorporates differentiable memory with transformers for rapid inference-time memorization. Building on the Titan memory, ATLAS [7] proposes a new transformer architecture that contains the differentiable memory module as part of the network to support large-scale distributed pretraining. Other approaches adapt through context alone. MetaICL [43] meta-learns from in-context examples without gradient updates. Our method, PERK, introduces parameter-efficient meta-learning using lightweight adapters, improving scalability for long contexts.

#### Long-Context Language Models

Developing long-context models has become a major research topic of its own.
Many works explore extending transformer-based language models’ context windows with minimal training, either by position extrapolation [12, 47, 63, 69] or improving the transformer’s attention mechanism [9, 44, 56, 27]. Other works show that using the default attention, applying simple position extrapolation, and fine-tuning the model on long documents [18, 42] or synthetically-generated long data [1, 64] yields much stronger results. The FT-ICR baseline in our work follows this paradigm. Other LLM works [37, 20, 19] propose achieving long-context capabilities through a long-context continued training stage between standard pre-training and supervised fine-tuning. In contrast to all the approaches above, PERK achieves long-context reasoning through the ability to encode long sequences using test-time gradient updates on a lightweight model adapter.

#### Long-Context Architectures

In parallel, there have been many efforts in designing more efficient architectures, for example, linear attention with RNNs [11, 46, 21, 15, 66, 6] and alternative attention architectures [53, 59, 22]. However, these new architectures often require training from scratch, and many have inherent limitations in terms of long-context recall [26, 3]. Recent works also explore hybrid models [37, 16] or distilling existing LLMs into hybrid models [61] and show promising results. In comparison, PERK directly augments off-the-shelf pre-trained LLMs without requiring additional architecture modifications or extensive pre-training, while naturally being parameter-efficient.

#### Long-Context Evaluation

Many benchmarks have been proposed to evaluate language models’ long-context abilities [31, 55, 4, 71, 23, 70, 68, 62]. Some works focus on the impact of extending input lengths with irrelevant contexts [32, 35] or changing relevant information positions [39] on LLM reasoning. Many works also reframe natural-language-processing tasks, such as retrieval [34, 49], summarization [30], and in-context learning [36, 65, 10], as long-context challenges. In this work, we define Drops-in-the-Ocean as a new long-context evaluation paradigm that makes the long context as distributionally similar to relevant information as possible.

## 7 Conclusion

In this work, we introduce PERK, a parameter-efficient meta-learning approach that learns to perform long-context reasoning by encoding information into a lightweight adapter through test-time gradient updates. Our experiments on multiple long-context benchmarks (Needle-in-the-Haystack, Drops-in-the-Ocean, and API retrieval) demonstrate substantial performance gains over models finetuned for classical in-context reasoning across reasoning complexity and context lengths. Our generalization analysis highlights PERK’s strong test-time length extrapolation ability, where the model can generalize to unseen contexts 32 times longer than the training data. When relevant information shifts positions at test-time, PERK robustly generalizes across different positions. Finally, we show that PERK scales efficiently at inference to extremely long contexts. Overall, PERK excels in performance, generalization, and efficiency for long-context reasoning.

## References

- [1]
Shengnan An, Zexiong Ma, Zeqi Lin, Nanning Zheng, Jian-Guang Lou, and Weizhu Chen.
Make your llm fully utilize the context.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 62160–62188. Curran Associates, Inc., 2024.
- [2]
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Ambrose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur.
Exploring length generalization in large language models.
Advances in Neural Information Processing Systems, 35:38546–38556, 2022.
- [3]
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré.
Simple linear attention language models balance the recall-throughput tradeoff.
arXiv preprint arXiv:2402.18668, 2025.
- [4]
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li.
Longbench: A bilingual, multitask benchmark for long context understanding, 2024.
- [5]
Deniz Bayazit, Negar Foroutan, Zeming Chen, Gail Weiss, and Antoine Bosselut.
Discovering knowledge-critical subnetworks in pretrained language models.
In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 6549–6583, Miami, Florida, USA, November 2024. Association for Computational Linguistics.
- [6]
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter.
xlstm: Extended long short-term memory.
arXiv preprint arXiv:2405.04517, 2024.
- [7]
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni.
Atlas: Learning to optimally memorize the context at test time, 2025.
- [8]
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni.
Titans: Learning to memorize at test time.
arXiv preprint arXiv:2501.00663, 2024.
- [9]
Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley.
Unlimiformer: Long-range transformers with unlimited length input.
In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 35522–35543. Curran Associates, Inc., 2023.
- [10]
Amanda Bertsch, Maor Ivgi, Emily Xiao, Uri Alon, Jonathan Berant, Matthew R. Gormley, and Graham Neubig.
In-context learning with long-context models: An in-depth exploration.
arXiv preprint arXiv:2405.00200, 2025.
- [11]
Kedar Chandrayan, Lance Martin, Greg Kamradt, Lazaro Hurtado, Arkady Arkhangorodsky, Lazaro Hurtado, and Pavel Kral.
Needle in a haystack - pressure testing llms.
2023.
- [12]
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian.
Extending context window of large language models via positional interpolation.
arXiv preprint arXiv:2306.15595, 2023.
- [13]
Zeming Chen, Gail Weiss, Eric Mitchell, Asli Celikyilmaz, and Antoine Bosselut.
Reckoning: reasoning through dynamic knowledge encoding.
Advances in Neural Information Processing Systems, 36:62579–62600, 2023.
- [14]
Peter Clark, Oyvind Tafjord, and Kyle Richardson.
Transformers as soft reasoners over language.
In Christian Bessiere, editor, Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020, pages 3882–3890. ijcai.org, 2020.
- [15]
Tri Dao and Albert Gu.
Transformers are ssms: Generalized models and efficient algorithms through structured state space duality.
arXiv preprint arXiv:2405.21060, 2024.
- [16]
Soham De, Samuel L. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, Guillaume Desjardins, Arnaud Doucet, David Budden, Yee Whye Teh, Razvan Pascanu, Nando De Freitas, and Caglar Gulcehre.
Griffin: Mixing gated linear recurrences with local attention for efficient language models.
arXiv preprint arXiv:2402.19427, 2024.
- [17]
Chelsea Finn, Pieter Abbeel, and Sergey Levine.
Model-agnostic meta-learning for fast adaptation of deep networks.
In International conference on machine learning, pages 1126–1135. PMLR, 2017.
- [18]
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng.
Data engineering for scaling language models to 128k context.
arXiv preprint arXiv:2402.10171, 2024.
- [19]
Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen.
How to train long-context language models (effectively).
arXiv preprint arXiv:2410.02660, 2025.
- [20]
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang,
Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur,
Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher,
Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang, Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh
Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna
Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim
Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan,
Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang, Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi
Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma.
The llama 3 herd of models.
arXiv preprint arXiv:2407.21783, 2024.
- [21]
Albert Gu and Tri Dao.
Mamba: Linear-time sequence modeling with selective state spaces.
arXiv preprint arXiv:2312.00752, 2024.
- [22]
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab S. Mirrokni, David P. Woodruff, and Amir Zandieh.
Hyperattention: Long-context attention in near-linear time.
ArXiv, abs/2310.05869, 2023.
- [23]
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg.
Ruler: What’s the real context size of your long-context language models?
arXiv preprint arXiv:2404.06654, 2024.
- [24]
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.
Lora: Low-rank adaptation of large language models.
ICLR, 1(2):3, 2022.
- [25]
Nathan Hu, Eric Mitchell, Christopher D. Manning, and Chelsea Finn.
Meta-learning online adaptation of language models.
arXiv preprint arXiv:2403.04317, 2023.
- [26]
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach.
Repeat after me: Transformers are better than state space models at copying.
arXiv preprint arXiv:2402.01032, 2024.
- [27]
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yuan Chang, Huiyuan Chen, and Xia Hu.
Llm maybe longlm: Self-extend llm context window without tuning.
arXiv preprint arXiv:2401.01325, 2024.
- [28]
Damjan Kalajdzievski.
A rank stabilization scaling factor for fine-tuning with lora.
arXiv preprint arXiv:2312.03732, 2023.
- [29]
Petros Karypis, Julian McAuley, and George Karypis.
Extending input contexts of language models through training on segmented sequences.
In Kevin Duh, Helena Gomez, and Steven Bethard, editors, Findings of the Association for Computational Linguistics: NAACL 2024, pages 3040–3052, Mexico City, Mexico, June 2024. Association for Computational Linguistics.
- [30]
Yekyung Kim, Yapei Chang, Marzena Karpinska, Aparna Garimella, Varun Manjunatha, Kyle Lo, Tanya Goyal, and Mohit Iyyer.
Fables: Evaluating faithfulness and content selection in book-length summarization, 2024.
- [31]
Kalpesh Krishna, Erin Bransom, Bailey Kuehl, Mohit Iyyer, Pradeep Dasigi, Arman Cohan, and Kyle Lo.
Longeval: Guidelines for human evaluation of faithfulness in long-form summarization, 2023.
- [32]
Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev.
Babilong: Testing the limits of llms with long context reasoning-in-a-haystack.
Advances in Neural Information Processing Systems, 37:106519–106554, 2024.
- [33]
Nathan Lambert, Jacob Morrison, Valentina Pyatkin, Shengyi Huang, Hamish Ivison, Faeze Brahman, Lester James V. Miranda, Alisa Liu, Nouha Dziri, Shane Lyu, Yuling Gu, Saumya Malik, Victoria Graf, Jena D. Hwang, Jiangjiang Yang, Ronan Le Bras, Oyvind Tafjord, Chris Wilhelm, Luca Soldaini, Noah A. Smith, Yizhong Wang, Pradeep Dasigi, and Hannaneh Hajishirzi.
Tulu 3: Pushing frontiers in open language model post-training.
arXiv preprint arXiv:2411.15124, 2025.
- [34]
Jinhyuk Lee, Anthony Chen, Zhuyun Dai, Dheeru Dua, Devendra Singh Sachan, Michael Boratko, Yi Luan, Sébastien M. R. Arnold, Vincent Perot, Siddharth Dalmia, Hexiang Hu, Xudong Lin, Panupong Pasupat, Aida Amini, Jeremy R. Cole, Sebastian Riedel, Iftekhar Naim, Ming-Wei Chang, and Kelvin Guu.
Can long-context language models subsume retrieval, rag, sql, and more?
arXiv preprint arXiv:2406.13121, 2024.
- [35]
Mosh Levy, Alon Jacoby, and Yoav Goldberg.
Same task, more tokens: the impact of input length on the reasoning performance of large language models.
arXiv preprint arXiv:2402.14848, 2024.
- [36]
Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen.
Long-context llms struggle with long in-context learning.
arXiv preprint arXiv:2404.02060, 2024.
- [37]
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham.
Jamba: A hybrid transformer-mamba language model.
arXiv preprint arXiv:2403.19887, 2024.
- [38]
Bingbin Liu, Jordan Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang.
Exposing attention glitches with flip-flop language modeling.
Advances in Neural Information Processing Systems, 36:25549–25583, 2023.
- [39]
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang.
Lost in the middle: How language models use long contexts.
Transactions of the Association for Computational Linguistics, 12:157–173, 2024.
- [40]
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
arXiv preprint arXiv:1711.05101, 2017.
- [41]
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf.
Fineweb-edu: the finest collection of educational content.
2024.
- [42]
Yi Lu, Jing Nathan Yan, Songlin Yang, Justin T. Chiu, Siyu Ren, Fei Yuan, Wenting Zhao, Zhiyong Wu, and Alexander M. Rush.
A controlled study on long context extension and generalization in llms.
arXiv preprint arXiv:2409.12181, 2024.
- [43]
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi.
MetaICL: Learning to learn in context.
In Marine Carpuat, Marie-Catherine de Marneffe, and Ivan Vladimir Meza Ruiz, editors, Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2791–2809, Seattle, United States, July 2022. Association for Computational Linguistics.
- [44]
Amirkeivan Mohtashami and Martin Jaggi.
Random-access infinite context length for transformers.
In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- [45]
Shishir G Patil, Tianjun Zhang, Xin Wang, and Joseph E Gonzalez.
Gorilla: Large language model connected with massive apis.
Advances in Neural Information Processing Systems, 37:126544–126565, 2024.
- [46]
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu.
Rwkv: Reinventing rnns for the transformer era.
arXiv preprint arXiv:2305.13048, 2023.
- [47]
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole.
Yarn: Efficient context window extension of large language models.
arXiv preprint arXiv:2309.00071, 2023.
- [48]
Ofir Press, Noah A. Smith, and Mike Lewis.
Train short, test long: Attention with linear biases enables input length extrapolation.
arXiv preprint arXiv:2108.12409, 2022.
- [49]
Yifu Qiu, Varun Embar, Yizhe Zhang, Navdeep Jaitly, Shay B. Cohen, and Benjamin Han.
Eliciting in-context retrieval and reasoning for long-context large language models.
arXiv preprint arXiv:2501.08248, 2025.
- [50]
Qwen, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu.
Qwen2.5 technical report.
arXiv preprint arXiv:2412.15115, 2025.
- [51]
Alec Radford and Karthik Narasimhan.
Improving language understanding by generative pre-training.
2018.
- [52]
Adam Roberts, Colin Raffel, and Noam Shazeer.
How much knowledge can you pack into the parameters of a language model?
In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 5418–5426, Online, November 2020. Association for Computational Linguistics.
- [53]
Ohad Rubin and Jonathan Berant.
Retrieval-pretrained transformer: Long-range language modeling with self-retrieval.
Transactions of the Association for Computational Linguistics, 12:1197–1213, 2024.
- [54]
Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots.
Truncated back-propagation for bilevel optimization.
In Kamalika Chaudhuri and Masashi Sugiyama, editors, Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 1723–1732. PMLR, 16–18 Apr 2019.
- [55]
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy.
ZeroSCROLLS: A zero-shot benchmark for long text understanding.
In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Findings of the Association for Computational Linguistics: EMNLP 2023, pages 7977–7989, Singapore, December 2023. Association for Computational Linguistics.
- [56]
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian.
A study on relu and softmax in transformer.
arXiv preprint arXiv:2302.06461, 2023.
- [57]
Qianli Shen, Yezhen Wang, Zhouhao Yang, Xiang Li, Haonan Wang, Yang Zhang, Jonathan Scarlett, Zhanxing Zhu, and Kenji Kawaguchi.
Memory-efficient gradient unrolling for large-scale bi-level optimization.
In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
- [58]
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin.
Learning to (learn at test time): Rnns with expressive hidden states.
arXiv preprint arXiv:2407.04620, 2025.
- [59]
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei.
You only cache once: Decoder-decoder architectures for language models.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 7339–7361. Curran Associates, Inc., 2024.
- [60]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998–6008, 2017.
- [61]
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, and Tri Dao.
The mamba in the llama: Distilling and accelerating hybrid models.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 62432–62457. Curran Associates, Inc., 2024.
- [62]
Menglin Xia, Victor Ruehle, Saravan Rajmohan, and Reza Shokri.
Minerva: A programmable memory test benchmark for language models.
arXiv preprint arXiv:2502.03358, 2025.
- [63]
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis.
Efficient streaming language models with attention sinks.
arXiv preprint arXiv:2309.17453, 2024.
- [64]
Zheyang Xiong, Vasilis Papageorgiou, Kangwook Lee, and Dimitris Papailiopoulos.
From artificial needles to real haystacks: Improving retrieval capabilities in llms by finetuning on synthetic data.
arXiv preprint arXiv:2406.19292, 2024.
- [65]
Xiaoyue Xu, Qinyuan Ye, and Xiang Ren.
Stress-testing long-context language models with lifelong icl and task haystack.
arXiv preprint arXiv:2407.16695, 2024.
- [66]
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim.
Gated linear attention transformers with hardware-efficient training.
arXiv preprint arXiv:2312.06635, 2024.
- [67]
Han-Jia Ye and Wei-Lun Chao.
How to train your MAML to excel in few-shot classification.
In International Conference on Learning Representations, 2022.
- [68]
Xi Ye, Fangcong Yin, Yinghui He, Joie Zhang, Howard Yen, Tianyu Gao, Greg Durrett, and Danqi Chen.
Longproc: Benchmarking long-context language models on long procedural generation.
arXiv preprint arXiv:2501.05414, 2025.
- [69]
Howard Yen, Tianyu Gao, and Danqi Chen.
Long-context language modeling with parallel context encoding.
In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 2588–2610, Bangkok, Thailand, August 2024. Association for Computational Linguistics.
- [70]
Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izsak, Moshe Wasserblat, and Danqi Chen.
Helmet: How to evaluate long-context language models effectively and thoroughly.
arXiv preprint arXiv:2410.02694, 2025.
- [71]
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun.
$\infty$Bench: Extending long context evaluation beyond 100K tokens.
In Lun-Wei Ku, Andre Martins, and Vivek Srikumar, editors, Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15262–15277, Bangkok, Thailand, August 2024. Association for Computational Linguistics.

## Appendix A Datasets

In this section, we provide details on the datasets used for each task in our experiments and analysis. For each dataset, we visualize an example in Figure [7](https://arxiv.org/html/2507.06415v1#A1.F7).

### A.1 BabiLong

We construct the train and evaluation data using the data generation framework proposed by [32]. As we mentioned in the experiment setup, we mainly focus on tasks involving single-hop (QA1), two-hop (QA2), and three-hop (QA3) reasoning.

#### Tasks

The main task for BabiLong is to reason over facts distributed in a long context. For each QA task, we have a set of facts, including distractors and supporting facts. Each fact describes an agent’s action or movement (which also describes the agent’s location). For example, Mary journeyed to the garden is a movement, where Mary is the agent, and garden is Mary’s current location. As an example for actions, John picked up the apple there describes John’s action of picking up an apple. Together with multiple agents performing multiple movements and actions, the set of facts became a temporal sequence of events.

QA1 focuses on one-hop reasoning, where the problems can be answered by using a single supporting fact. For example, in a sequence of actions and movements, the agent John’s last movement is John moved to the bedroom. Then the one-hop question is “Where is John located now?”. By the last movement alone, we can find the answer “bedroom”.

QA2 focuses on one-hop reasoning, where the problems can be answered by using two supporting facts, where one fact is a movement and the other one is an action. For example, in the event sequence, we have John moved to the garden and John dropped the apple there. Using these two facts, we can answer “Where is the apple now?”, which is “garden”.

QA3 focuses on one-hop reasoning, where the problems can be answered by using three supporting facts, as a combination of movements and actions. For example, we have three supporting facts: Daniel went back to the office, Daniel went back to the bedroom, and Daniel left the milk in order. We can infer the answer to the question “Where was the milk before the bedroom?” is “office”.

#### Build Long Context

To build the long context, we use the framework to construct the tasks with long natural documents sampled from FineWeb-edu [41], specifically from the CC-MAIN-2024-42 subset, ensuring the content postdates our base model’s knowledge cutoff. Given a particular context window, such as 8192, we sample that many tokens from the corpus as the Haystack long context.

We need to convert a long context into a batch of subsequences, such that PERK can encode them in parallel through test-time gradient-based adaptation. In practice, we set the effective context length (i.e., the number of tokens of a subsequence on average) as 256. We divide the long context into chunks where each chunk is approximately 256 tokens long. Then, we randomly distribute the set of facts into these chunks. For each fact, we randomly select a chunk and then insert this fact into the chunk at a random position. To preserve the temporal order of the facts (since they come from a sequence of events with temporal order), we index each fact with a prefix. Finally, we have a batch of subsequences where each sequence might contain some facts at random positions.

### A.2 Student Records

In the main paper, we propose Drops-in-the-Ocean (DIO), a new evaluation setting where long contexts are formed from structurally similar documents. We aim to create problems in which the relevant and irrelevant information will likely be distributionally similar, making it more difficult to identify critical facts. We construct a synthetic dataset, Student Records, where each context simulates a database containing multiple student records. Each record includes several attributes: ID, name, school, major, and grade.

#### Tasks

The task operates on top of a database that consists of multiple records. Each record holds information about a student. The record is structured as follows:

$$ Student Id: 3109998, Student Name: Alison Keith, Year: Sophomore $$

$$ School: School of Computer Information and Data Sciences, Major: Computer Science, Grade: 72 $$

We generate the student IDs and grades using a random number generator. The student names are generated from the Name package of Python. We define a pool of years, schools, and courses. Then, we randomly select each entity for each student. Crucially, to test generalization, specific entities (e.g., IDs, names) and attribute value combinations encountered in the test set are disjoint from those used during training.

Based on the database of student records, we construct three main tasks: Recall (retrieving attributes for a specific ID), Relation (comparing attributes between two IDs), and Aggregate (calculating the maximum, minimum, and average grade across all students).

Recall tests a model’s ability to retrieve an attribute of a student given a specific student ID number. For example, we can ask “What major does student 1778924 study?”. The model should answer “What major does student 1778924 study?”. The questions cover all attributes for testing coverage.

Relation tests a model’s ability to compare the attributes of two students. For example, for the attribute of grade, we can ask “Does student 1778924 have a higher grade than student 7077015?”. The model should be able to compare their grades and infer “Yes”. Or we can also ask “Do student 1778924 and student 7077015 study the same major?”, the model should recognize that their majors are different and answer “No”.

Aggregation is the most difficult task in this dataset. This task requires many-hop reasoning ability, where the model needs to aggregate attributes from all students in a long context to make an accurate conclusion. We focus on three types of numerical aggregation among the students’ grades: (1) Maximization, (2) Minimization, and (3) Averaging. For example, we ask the model “What is the average grade of all students”, the model should be able to compute the sum of all student grades, count the number of total students in the context, and then compute the average, “49.3”.

Figure: Figure 7: Dataset Examples. Here we show an example from the three datasets we used: (Top) BabiLong, (Middle) Student Records, and (Bottom) Lost-in-the-Middle API Retrieval.
Refer to caption: extracted/6607134/figures/data_examples.png

#### Build Long Context

Context length scales directly with the number of student records included. For a given context window, we continuously append student records to the context until the number of tokens exceeds the context window. To convert the long context into a batch of subsequences for PERK, we simply treat each student record as an individual subsequence in the batch, since they are already independent documents by design.

### A.3 Lost-in-the-Middle API Retrieval

To analyze a model’s robustness against relevant information positions, we follow the setting proposed by Lost in the Middle [39]. We use API documents from [45], where each document has an associated API call, and create long contexts containing multiple API-call-to-document pairs. We place the relevant API document in a particular user instruction at different locations in the context (beginning, middle, end, and random position).

#### Task

We formulate the main task as gold API retrieval. In this task, the model must retrieve the correct API call based on user instructions and documents. Using fixed 4k and 8k token contexts, we train PERK and FT-ICR models on contexts with the target document at varying positions and test them on contexts distributed across all positions.

#### Build Long Context

We follow Student Records and scale the context length with the number of API documents included. For a given context window, we continuously append API documents to the context until the number of tokens exceeds the context window. Similarly, since each API document is independent, we treat a document as a subsequence in the batch for PERK. Since the documents are processed in parallel as a batch, they do not preserve any ordering information. Thus, by design, PERK is robust to the position of the relevant information in a context.

## Appendix B Training and Inference Details

#### Hardware Information

All training runs are conducted on a single Nvidia H100 GPU with 93GB of VRAM. The training jobs run on an internal research cluster managed by Slurm. Each job is provided with 64 CPUs and 64GB of RAM. The GPU is completely utilized without sharing with other concurrent jobs.

#### Standard Hyperparameters

Each training (for PERK and baselines) sets 2 epochs as the maximum runtime with early stopping based on validation loss. In practice, most training runs typically stop early, within one epoch of runtime. We set the baseline training and PERK’s outer loop learning rate to 1e-5, and the weight decay to 0.01. We use a cosine learning rate scheduler with 0.03 of the total optimization steps as the warm-up phase. We use PyTorch Lightning (^1^11[https://lightning.ai/docs/pytorch/stable/](https://lightning.ai/docs/pytorch/stable/)) as the trainer for PERK. The Qwen and GPT-2 FT-ICR baselines are trained using the Open-Instruct [33] trainer from AI2, and the Mamba FT-ICR baselines are trained using the Huggingface Trainer (^2^22[https://huggingface.co/docs/transformers/en/main_classes/trainer](https://huggingface.co/docs/transformers/en/main_classes/trainer)). For the stress test of length generalization beyond 64K tokens, we change the default maximum position of 32768 to 131072.

#### LoRA Hyperparameters

We set LoRA rank to 256 for both PERK GPT-2 and PERK Qwen. We use the default values for the other hyperparameters, including setting the alpha to 16 and the dropout to 0.1. We apply LoRA to all modules in the model architecture. For GPT-2, the number of active training parameters of PERK is 38M, 30% of the total parameter number. For Qwen, the number of active training parameters of PERK is 140M, which is also 30% of the total parameter number. We set the scaling factor $\alpha$ to 16 initially and enable rank-stabilized LoRA fine-tuning [28], which results in a target $\alpha$ of 256, corresponding to a rank of 256.

#### Meta-Learning Hyperparameters

For the inner loop adaptation, we use 4 inner loop gradient steps. For the truncated gradient unrolling, we truncate steps 1 and 2 for input context length $\leq$4K, and we truncate steps 1, 2, and 3 for context length $>$ 8K. We follow [67]’s practice on how to train a good MAML model and use per-layer-per-step dynamic learning rate for the inner loop, using 5e-5 as the initial value. For the language modeling loss in the inner loop, we use a weighted version similar to CaMeLS [25]. The token-level weights are predicted via a weighting model implemented as a two-layer MLP network. The parameters of the weighting model are co-meta-learned along with the LoRA parameters and the adaptive per-layer-per-step learning rates. We use the differentiable version of the AdamW [40] optimizer for the inner loop optimization, implemented by the Higher (^3^33[https://github.com/facebookresearch/higher](https://github.com/facebookresearch/higher)) library for higher-order optimization. We compute higher-order derivatives using the automatic differentiation mechanism implemented by PyTorch(^4^44[https://docs.pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html](https://docs.pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html)). The outer loop optimizer is the standard PyTorch implementation of AdamW.

#### Evaluation Hyperparameters

At inference time, we generate the answer for a reasoning problem via greedy decoding (temperature is set to 0). The maximum number of tokens that can be generated is 512. The EOS token is set to the default one defined in the tokenizer. The generation is handled by the text-generation pipeline implemented by Huggingface (^5^55[https://huggingface.co/docs/transformers/v4.43.4/en/main_classes/pipelines](https://huggingface.co/docs/transformers/v4.43.4/en/main_classes/pipelines)).

## Appendix C Additional Analysis

**Table 1: Analyzing LoRA rank’s impact on PERK’s performance.**
| Method | QA1 | QA2 | QA3 |
| --- | --- | --- | --- |
| FT-ICR | 91.1 | 70.5 | 48.6 |
| PERK LoRA-16 | 99.0 | 96.9 | 84.7 |
| PERK LoRA-256 | 100 | 98.1 | 89.1 |

### C.1 Impact of LoRA Rank in PERK

In this section, we analyze the impact of LoRA’s rank on PERK’s long-context reasoning performance. We focus on the BabiLong tasks with a 1024 context window for an efficient ablation study. We apply PERK to Qwen-2.5-0.5B with LoRA rank of 16 and 256 and evaluate on the BabiLong tasks. We show the results in Table [1](https://arxiv.org/html/2507.06415v1#A3.T1). We also include the performance from the FT-ICR Qwen baseline as a reference. The results show that PERK’s performance only decreases slightly when the rank is reduced to 16. However, PERK still maintains a much stronger performance than the FT-ICR Qwen baseline. The analysis shows that PERK can maintain its strong performance with a more parameter-efficient setting, highlighting its potential for more efficient long-context reasoning.

**Table 2: Analyzing the impact of the number of inner loop adaptation steps on PERK’s performance.**
| Method | QA1 | QA2 | QA3 |
| --- | --- | --- | --- |
| FT-ICR | 91.1 | 70.5 | 48.6 |
| PERK 2 Steps | 99.2 | 94.5 | 82.4 |
| PERK 4 Steps | 100 | 98.1 | 89.1 |

### C.2 Inner Loop Adaptation Steps

In this section, we analyze the impact of inner loop adaptation steps on PERK’s long-context reasoning performance. Similarly, we focus on the BabiLong tasks with a 1024 context window. We again apply PERK to Qwen-2.5-0.5B for 2 and 4 inner loop steps. With 4 steps, we truncate the first 2 steps in backpropagation. With 2 steps, we truncate the first step in backpropagation. We show the results in Table [2](https://arxiv.org/html/2507.06415v1#A3.T2). The results show that reducing the number of inner loop adaptation steps decreases the performance more noticeably than reducing LoRA rank. However, PERK still outperforms FT-ICR Qwen on all three tasks. The results here suggest that we should select a higher number of inner loop adaptation steps to maximize the performance. If in settings where GPU memory might be limited, reducing the number of inner loop steps can reduce the performance, but it is still able to outperform the baseline. However, the analysis is done with limited hardware resources, which limits our analysis to 4 steps. To further analyze the impact of the inner loop step counts, more systematic experiments in a scalable hardware setting are needed for future work.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - One of the mentioned strengths is that it encodes context as permutation-invariant representations, but that might not always be a strength. What if there is a query about positional information? The benchmarks do not consider that.
- In figure 6, you show that you encode context + question before the LoRA updates. This means that for each new question, you need to do the update procedure. Is there a way to encode only context so that you can ask multiple questions?
soundness: 3
strengths: - Paper is well written, and the results are nice
- PERK demonstrates strong generalization to long context extrapolation (e.g., training on 8K, testing on 128K) and superior robustness to positional biases in the relevant information.
- The idea of using a LoRA adapter as a differentiable memory module for context is a good alternative to ICL
- The use of LoRA and TGU is nice, and seems well executed.
summary: This paper introduces PERK, a method to improve reasoning over long contexts.

Perk learns from the context at test time by updating a small, efficient memory module (using LoRA) while keeping the main model frozen. It uses a two-step training process: first, rapidly encodes the context, and then optimizes how the model uses that info to answer questions.

Experiments on long context benchmarks show PERK significantly outperforms standard long-context finetuning.
weaknesses: - The biggest concern is the added complexity. The gains are nice, but I'm not sure they justify using this method over FT-ICR.
- Related to the added complexity, in general the inference time is significantly longer than ICR (at least up to 34k according to fig 7)
- PERK is designed for length generalization, as noted in the experiments (train on 8k, eval on 32k). However, FT-ICR is not suited for this kind of evaluation, as is well known. So I'm not sure this is the best baseline to compare against. Especially since we can fine-tune on 32K lengths, I think it would be nice to see an FT-ICR training with 32k length.
- Permutation-invariance can be detrimental for certain tasks (though maybe rare).

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - Will the code be publicly released?

- Reading PERK made me thing about PATHGOOSE [1], where at each LLM token the authors propose to route to a LoRA module that has been finetuned for a specific task. I haven't seen a comprehensive description of any future work you envision for PERK and was wondering if some variant of PATHGOOSE is something you thought about to improve efficiency. 


[1] Mohammed Muqeeth, Haokun Liu, Yufan Liu, Colin Raffel, Learning to Route Among Specialized Experts for Zero-Shot Generalization,
soundness: 3
strengths: - The paper offers a wide-ranging experimental study across multiple long-context reasoning benchmarks, including Needle-in-a-Haystack, HotpotQA, TriviaQA, and the newly introduced Drops-in-the-Ocean (DIO) dataset. The introduction of DIO, which features distributionally similar distractors to better test reasoning precision, strengthens the evaluation by addressing limitations of prior benchmarks. Comparisons against both open-source and commercial LLMs convincingly demonstrate PERK’s superiority.

- PERK exhibits remarkable robustness to reasoning complexity, unseen context lengths (up to 128K tokens), and positional shifts of relevant information. Its consistent performance across model families (GPT, Qwen, LLaMA) and scales indicates that the method learns transferable reasoning strategies rather than overfitting to a particular model or data distribution.

- The approach is built on solid meta-learning foundations, combining bi-level optimization with truncated gradient unrolling to balance effectiveness and efficiency. Its LoRA-based parameter-efficient adaptation enables scalable test-time learning without full model updates, making PERK both methodologically principled and computationally practical.
summary: This paper introduces PERK (Parameter-Efficient Reasoning over Knowledge), a novel meta-learning framework that enhances large language models’ (LLMs) ability to perform long-context reasoning. Instead of relying solely on in-context learning or expensive fine-tuning, PERK allows models to encode long contexts through gradient updates at test time, using lightweight Low-Rank Adapters (LoRA) as memory modules. The approach employs two nested optimization loops: the inner loop encodes context into LoRA parameters, while the outer loop learns how to use these adapted parameters for reasoning. Across multiple benchmarks, including Needle-in-a-Haystack, HotpotQA, and a new Drops-in-the-Ocean dataset, PERK consistently outperforms long-context fine-tuning baselines and specialized long-context models, showing strong robustness to reasoning complexity, length extrapolation, and positional variation in relevant information.
weaknesses: - The authors don't mention any public release of the code

- In section 2 of the appendix I see you only ran you experiments setting temperature to 0. Have you tried other values and what results have you found? I'd be curious to see the robustness of PERK at different decoding parameters

- While I didn't find any major weakness, what prevented me from giving a strong accept to this paper is a clear metric of inference cost for PERK. For this method to be useful to researchers and ML practitioners, data on inference cost is fundamental. While a direct comparison with fine-tuning methods is not fair, at least some stats on time needed by PERK would make the paper contribution much more comprehensive

Overall, I think PERK makes a solid contribution and I vote for its inclusion at the conference as is. With additional inference cost analysis, I'd be happy to further increase my score

---

confidence: 3
contribution: 3
initial_rating: 4
presentation: 3
questions: N/A
soundness: 3
strengths: * The test-time learning problem is relatively useful to further improve the model performance.
* Figure 1 presents the pipeline of this method, and the inner loop is clear.
* The method compares the performance with FT-ICR, proving that the PERK is better.
* The experiments have various models, supporting that the PERK is general.
summary: This work proposes a meta-learning method that learns to encode long contexts through gradient updates to a model adapter at test time. The inner encodes the context to Lora, and the outer loop uses the adapter for the relevant information from the encoded long-context. The experiment proves that the proposed PERK is effective for long-context reasoning.
weaknesses: * Is the method test-time training? According to the definition of the outer loop, Equation 3 needs the label of the question. However, during test time, there is no such label.
* Figure 4 is not related to length extrapolation. The Qwen-2.5-0.5B is trained with 32K by the Qwen Team. However, Figure 4 presents the performance with a maximum 32K, which is NOT longer than the Qwen-2.5-0.5B training length. This is a misclaim.
* It should provide the training cost, such as time cost, for the comparison of PERK and FT-ICR.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Inference Latency Trade-Off: Could you provide a more direct comparison of the end-to-end inference latency (context encoding + question answering) of PERK vs. FT-ICR, particularly in the common 8K-32K context range? How does this latency scale with the number of inner-loop adaptation steps, and where is the practical break-even point where PERK becomes faster than a standard forward pass?

2. Optimality of NLL for Encoding: The inner loop uses a standard NLL objective to "compress" the context into the LoRA. Have you considered alternative objectives that might more explicitly encourage the adapter to store salient facts? For instance, would a contrastive loss or an information-theoretic objective that prioritizes reasoning-relevant information lead to a more effective parametric memory?

3. Comparison to RAG: Your method encodes the entire context into parameters, which is a key difference from RAG systems that first retrieve a small subset of the context. Could you discuss the conceptual trade-offs between PERK and a standard RAG pipeline in terms of performance, latency, and robustness to different types of reasoning tasks?

4. Ablation on Training Components: The PERK training process involves several advanced components (TGU, learned per-layer LR, weighted NLL). How critical is each of these to the final performance? For example, what is the performance drop if you use a fixed inner-loop learning rate or a standard (unweighted) NLL loss?
soundness: 3
strengths: **Novel Conceptual Framing:** The paper's most significant contribution is reframing long-context reasoning as a test-time, gradient-based adaptation problem. This is a fundamental and clever departure from the standard in-context reasoning (ICR) paradigm, moving from "reasoning over context" to "reasoning from context stored in parameters."

**Exceptional Empirical Performance and Robustness:** PERK demonstrates substantial performance gains (up to +20% absolute) over standard fine-tuning. More importantly, the method shows remarkable robustness where standard models fail:
- Length Extrapolation: PERK models trained on 8K contexts maintain strong performance on contexts up to 128K, significantly outperforming baselines and even specialized long-context models.

- Positional Robustness: By processing context as a permutation-invariant batch of segments, PERK largely overcomes the "lost in the middle" problem, showing minimal performance degradation when the position of relevant information changes, unlike FT-ICR models which fail catastrophically.

**Scalability and Model-Agnosticism:** The combination of LoRA and truncated gradient unrolling makes the complex meta-learning approach tractable . The paper convincingly shows this advantage holds across multiple model families (GPT-2, Qwen, LLaMA) and scales (0.5B to 8B), highlighting the general applicability of the technique.

**Strong Evaluation Methodology:** The introduction of the Drops-in-the-Ocean (DIO) benchmark is a valuable contribution, creating a more challenging and realistic test than standard Needle-in-a-Haystack by using distributionally similar distractors.
summary: This paper introduces PERK (Parameter Efficient Reasoning over Knowledge), a novel method that reframes long-context reasoning as a test-time learning problem. Instead of processing a long context in a single forward pass, PERK first encodes the context into the parameters of a Low-Rank Adapter (LoRA) via several gradient-based update steps at inference time. The model then answers a given question using this newly created parameter-efficient memory module, without needing the original long context in its attention window. To enable this, PERK is trained using a bi-level meta-learning optimization scheme: an inner loop learns to rapidly encode context into the LoRA adapter, while an outer loop learns to use the adapted model for reasoning. To make this process scalable, the authors employ truncated gradient unrolling. The paper demonstrates through extensive experiments that PERK significantly outperforms standard long-context fine-tuning, shows superior length extrapolation and robustness to positional biases, and is competitive with or exceeds the performance of specialized proprietary and open-source long-context models.
weaknesses: **Inference Latency Trade-Off:** The paper emphasizes efficiency in terms of training memory, but the trade-off with inference latency is a major practical concern. The method requires performing multiple gradient updates at test time to encode the context before answering a question. While Figure 7b shows PERK can be faster at extreme context lengths (e.g., 64K), it appears slower than FT-ICR for more common lengths. This makes the "efficient" framing potentially misleading for many real-world applications where inference speed is critical.

**Complexity of the Training Framework:** The PERK training pipeline is highly complex, involving bi-level optimization, truncated gradient unrolling, learned per-layer-per-step learning rates, and a weighted NLL loss in the inner loop. This high degree of complexity could pose a significant barrier to reproducibility and adoption compared to the much simpler standard fine-tuning approach.

**Sub-Optimal Encoding Objective:** The inner loop's objective is simply next-token prediction over the context segments. This seems to incentivize the model to learn the surface-level statistics of the text, which may not be the most effective way to distill salient, reasoning-critical information into the LoRA parameters. The paper does not explore or discuss alternative, potentially more targeted, encoding objectives.


Decision: Accept (Poster)