1 Introduction↩︎

Large language models (LLMs) have achieved strong performance across many natural-language tasks, but adapting them to knowledge outside their static pretraining corpus remains difficult. Retrieval-augmented generation (RAG) [1] addresses this limitation by conditioning a frozen LLM on documents retrieved from an external corpus. However, retrieval alone does not solve the full adaptation problem. After relevant documents are retrieved, the model must still decide how to use them for a new task, domain, or query distribution.

Existing RAG systems usually address this problem in one of three ways. The first is to keep both the retriever and generator fixed, and simply prepend retrieved documents to the input. This strategy is efficient, but it treats retrieved documents as static evidence and gives the model no mechanism to adjust how evidence should be used. The second is to fine-tune the retriever, the generator, or both. This can improve task performance, but it requires additional training and can be expensive when the task or domain changes. The third is to use in-context learning (ICL), where a few input-output examples are provided at inference time. ICL is attractive because it avoids full model retraining, but it is still unclear whether these examples merely provide extra demonstrations, or whether they can induce a more systematic update to how a RAG model uses retrieved evidence.

This paper asks a simple question: can retrieved evidence and a few RAG examples act not only as context to read from, but also as a signal for adapting how the model uses evidence? Answering this question requires connecting two views that have mostly been studied separately. On one side, recent theory shows that, under linear self-attention, in-context learning can implement gradient descent on the examples in the context [2]–[4]. This suggests that context can behave like a forward-pass update, rather than only as additional input text. On the other side, RAG introduces structure that is absent from standard ICL theory: a query, retrieved documents, query-evidence interactions, and a generator that must combine them to produce an answer. It remains unclear whether the gradient-descent view of ICL extends to retrieval-augmented prediction, and whether such a view can guide practical adaptation in real RAG systems.

We study RAG from this in-context optimization perspective. Our goal is not to claim that modern retrieval-augmented LLMs literally perform gradient descent during inference. Instead, we use a controlled linear setting to identify where such an update would act. In linear RAG, the relevant update acts on the interaction between the query and retrieved evidence. This gives a simple design principle for LLM-scale RAG: rather than changing which documents are retrieved, we adapt how the frozen generator uses the retrieved documents. Guided by this principle, we propose a forward-only adaptation method for frozen RAG LLMs. The method keeps both the external retriever and the LLM backbone fixed, and adapts only a lightweight generator-side evidence-use interface implemented with LoRA. At inference time, given new few-shot RAG demonstrations, the predictor produces the update in a single forward pass, enabling the generator to adjust how it uses retrieved documents without re-training on the new dataset.

We develop this idea in three steps. First, we prove that one linear self-attention layer can implement one gradient-descent step on a unified linearized RAG objective covering both projection-based and dot-product retrieval interfaces. Second, we test how far this correspondence extends beyond the exact construction. A trained self-attention layer closely matches the constructed gradient-descent predictor under controlled linear shifts, varying document counts, and stacked depths, while nonlinear architectures and real-world regression data reveal a clear dependence on feature distribution. Third, we use the optimization view to guide LLM-scale RAG adaptation. Across seven QA benchmarks, two retrievers, and two frozen LLM backbones, the predicted update improves a shared-interface baseline, transfers to held-out tasks, and approaches test-time gradient adaptation at much lower per-query cost. Our contributions are summarized as follows:

• An in-context optimization view of linear RAG. We extend the ICL-as-gradient-descent perspective from generator-only prediction to retrieval-augmented prediction. We prove that one linear self-attention layer can implement one gradient-descent step on a unified linearized RAG loss covering both linear-projection and dot-product retrieval interfaces. We also show that stacking \(K\) linear self-attention layers gives a multi-step view of in-context optimization for linear RAG.

• A boundary analysis beyond the exact linear setting. We test when the linear construction remains predictive and when it breaks. On synthetic linear regression tasks, a trained self-attention layer closely matches the constructed gradient-descent predictor under distribution shift, varying document counts, and stacked depths. On nonlinear architectures and four real-world regression datasets, the alignment degrades in a structured way and becomes sensitive to feature distribution.

• Adapting evidence use without test-time backpropagation. We use the optimization view to guide adaptation in frozen RAG LLMs. Rather than changing the external retriever, we adapt a generator-side evidence-use interface implemented with Q/K/V LoRA modules. A small context-conditioned predictor amortizes the autograd-defined \(K\)-step update to this interface. Across seven QA benchmarks, two backbones, and two retrievers, the predicted update improves a shared-interface baseline, transfers to held-out domains, and approaches test-time gradient adaptation at much lower per-query cost.

2 Related Work↩︎

Retrieval-augmented generation (RAG) conditions a language model on documents retrieved from an external corpus [1], [5]–[9]. Prior work has improved RAG through better retrieval, prompting, evidence fusion, and joint retriever-generator training [10]–[13]. Our focus is complementary: rather than changing which documents are retrieved, we study how a frozen generator can adapt its use of already-retrieved evidence. We position this contribution relative to three lines of work.

2.0.0.1 In-context learning as gradient descent.

A growing line of theory interprets in-context learning as implicit optimization. Under linear self-attention, a single ICL forward pass can implement one gradient-descent step, and stacked layers can implement multiple steps [2]–[4], [14]. Later work extends this view to preconditioned gradient descent [15], in-context algorithm selection [16], the role of depth [17], [18], and kernel-regression interpretations of attention [19], [20]. These analyses mainly study generator-only settings, often through linear regression or simplified attention. RAG introduces additional structure, including retrieved documents, query-evidence interactions, and evidence-conditioned generation. We extend the gradient-descent view to a linearized RAG setting and use it to identify where evidence-use adaptation should act.

2.0.0.2 Context-conditioned weight prediction.

Another line of work learns auxiliary networks that produce model updates from a small context. HyperTuning [21] predicts soft prompts or low-rank weights from few-shot examples. HyperFlow [22] learns support-conditioned fine-tuning dynamics. MAC [23] maps documents into memory modulations. MEND [24] maps fine-tuning gradients into knowledge edits. RAG-GD follows the broad template of predicting an update from context, but differs in both the target and the adaptation site. The target is not a downstream task loss, a meta-learning objective, a memory objective, or a single editing gradient. Instead, the predictor matches an autograd-defined \(K\)-step SGD update induced by RAG-formatted demonstrations. The adapted parameters are also restricted to a generator-side evidence-use interface, while the retriever and backbone remain fixed.

2.0.0.3 Test-time adaptation.

Standard adaptation either updates model parameters before deployment, as in fine-tuning and LoRA [25], or leaves the model unchanged at inference, as in pure ICL. Test-time training [26] lies between these extremes by updating parameters for each test instance, but this requires per-instance backpropagation and becomes expensive for large LLMs. Recent studies compare ICL, fine-tuning, and trainable RAG as system-level adaptation strategies [27]–[29]. RAG-GD targets the same goal of adapting at inference time, but amortizes the update: a small predictor emits a LoRA update to the generator’s evidence-use interface in one forward pass. Thus, it avoids backpropagation through the LLM at deployment while keeping both the external retriever and the frozen backbone unchanged.

3 A Linear RAG Setting Where Self-Attention Implements Gradient Descent↩︎

This section establishes the linear-regime basis for our in-context optimization view of RAG. We study a controlled setting in which retrieval-augmented prediction and gradient descent can be connected exactly. The goal is not to model modern RAG systems literally: real retrievers involve discrete document selection, and modern LLMs are deep and nonlinear. Instead, we isolate a differentiable retrieval-augmented prediction problem and show that one linear self-attention layer can realize the prediction shift produced by one gradient-descent step. The proof and explicit construction are in Appendix 8, and the derivations for the retrieval variants are in Appendix 9.

3.1 Self-Attention↩︎

We begin with a multi-head self-attention block parameterized by \(\theta = \{P_h, W_{h,V}, W_{h,K}, W_{h,Q}\}_{h=1}^H\). Given tokens \(\{e_1,\ldots,e_N\} \subset \mathbb{R}^d\), the update for token \(e_j\) is \[e_j \gets e_j + \mathrm{SA}_\theta(j, \{e_i\}_{i=1}^N) = e_j + \sum_h P_h V_h \,\mathrm{softmax}(K_h^\top q_{h,j}), \label{eq:sa95update}\tag{1}\] where \(V_h\), \(K_h\), and \(q_{h,j}\) are the value matrix, key matrix, and query vector for head \(h\). Following [2], [17], we remove the softmax and bias terms to obtain the linear self-attention (LSA) update: \[e_j \gets e_j + \mathrm{LSA}_\theta(j, \{e_i\}_{i=1}^N) = e_j + \sum_h P_h V_h K_h^\top q_{h,j}. \label{eq:lsa95update}\tag{2}\]

3.2 A Unified Linearized RAG Predictor↩︎

We use a linearized abstraction of retrieval-augmented prediction. Rather than modeling discrete top-\(k\) selection, this abstraction captures a differentiable interface in which query features and retrieval-derived features jointly determine the prediction. Both a projection-based retrieval interface [1] and a dot-product retrieval interface [6] can be written as \(y = W_1 x_1 + W_2 x_2\) where \(x_1\) denotes the query-side feature and \(x_2\) denotes the retrieval-derived feature. For the projection-based interface, we set \(x_1 = x_q\), \(x_2 = D\), and \(W_2 \triangleq W_1 W_d\), where \(W_d\) projects document embeddings into the prediction space. For the dot-product interface, we set \(x_1 = x_2 = x_q\) and \(W_2 = W_z\left(\sum_i d_i d_i^\top\right)M^\top,\) where \(M\) parameterizes query-document similarity. For tractability, we use the shared-encoder simplification \(M = W_e^\top W_e\), so \(M\) is symmetric. The general DPR formulation [6] allows separate query and document encoders, with \(M = W_q^\top W_d\). Full derivations are provided in Appendix 9.

3.3 Optimization Objective↩︎

Given training examples \(\{(x_1^i, x_2^i, y_i)\}_{i=1}^N\), we consider the squared loss \[L(W_1, W_2) = \frac{1}{2N} \sum_{i=1}^N \left\| W_1 x_1^i + W_2 x_2^i - y_i \right\|^2. \label{eq:rag95loss}\tag{3}\] One gradient-descent step with learning rate \(\eta\) gives \[\Delta W_k = -\eta \nabla_{W_k} L = -\frac{\eta}{N} \sum_{i=1}^N \left( W_1 x_1^i + W_2 x_2^i - y_i \right) (x_k^i)^\top, \qquad k \in \{1,2\}. \label{eq:gd95step}\tag{4}\] For a query token with features \((x_1,x_2)\), the corresponding prediction shift is \(\Delta y \triangleq \Delta W_1 x_1 + \Delta W_2 x_2. \label{eq:prediction95shift}\) Thus, \(\Delta y\) is the change in prediction after updating \(W_k\) to \(W_k' = W_k + \Delta W_k\).

3.4 Linear Self-attention Reproduces one Gradient Step↩︎

The construction is given in Appendix 8. Intuitively, the value projection encodes the residual \(W_1x_1^i + W_2x_2^i - y^i\). The key-query interaction computes the inner products \((x_1^i)^\top x_1^j\) and \((x_2^i)^\top x_2^j\). The output projection then writes the resulting weighted residual sum into the query token’s prediction coordinate.

This construction also gives a controlled multi-step analogue. If each LSA layer represents one gradient-like update, then after \(K\) layers, \[\hat{y}^{(K)}_{N+1} = \hat{y}^{(0)}_{N+1} + \sum_{t=0}^{K-1} \left( \Delta W^{(t)}_1 x^1_{N+1} + \Delta W^{(t)}_2 x^2_{N+1} \right), \label{eq:stacked}\tag{5}\] where \(\Delta W^{(t)}_1\) and \(\Delta W^{(t)}_2\) are the implicit updates represented by layer \(t\). We use this multi-step view as a linear-regime guide rather than as a literal claim about frozen LLM computation. In later sections, this view motivates a forward-only mechanism that adapts how a frozen generator uses retrieved evidence.

4 Testing the Boundary of the Linear Correspondence↩︎

Lemma 1 gives an exact correspondence in a controlled linear setting. We now ask how far this correspondence remains predictive when the setting is varied. The experiments have two goals: first, to verify that a trained linear self-attention layer can reproduce the constructed gradient-descent predictor; second, to identify where the correspondence begins to break beyond the exact regime.

4.1 Linear-Regime Verification↩︎

Each token concatenates an input feature, a retrieval-derived feature, and a target, \(e_i=(x_i,z_i,y_i)\) for \(i=1,\ldots,N\). The auxiliary slot \(z_i\) instantiates the unified RAG view. For the projection-based interface, \(z_i\) is a document-derived feature. For the dot-product interface, \(z_i=x_i\), and document information is injected into the keys and values. We train an LSA layer \(\theta\) to minimize expected squared error across tasks, using minibatch SGD over freshly sampled tasks. Following prior work [2], [30], each task is generated from a teacher with weights \(W_\tau \sim \mathcal{N}(0,I)\). Inputs are sampled as \(x_{\tau,i}\sim\mathcal{U}(-1,1)^{n_I}\), and targets are generated by \(y_{\tau,i}=W_\tau^1x_{\tau,i}^1+W_\tau^2x_{\tau,i}^2\). We set \(N=n_I=10\) and sweep the document count \(k \in \{2,5,10,25\}\). We compare the trained layer \(\theta^*\) with the constructed predictor that exactly realizes one gradient-descent step on the unified RAG loss. On \(T_{\mathrm{val}}=10^4\) held-out tasks, we report the prediction difference \(\|\hat{y}_{\theta^*}-\hat{y}_{\theta,\mathrm{rag}}\|_2\), the cosine similarity between input sensitivities \(\partial \hat{y}/\partial x_{\mathrm{test}}\), and the corresponding sensitivity \(\ell_2\) difference. Details are in Appendix 10.

Figure 1 verifies the construction for both retrieval interfaces. The trained LSA layer closely matches the constructed predictor: the loss difference is small, the sensitivity cosine is close to \(1\), and the sensitivity \(\ell_2\) difference is negligible. This numerically confirms the algebraic correspondence in Lemma 1.

4.2 Controlled Stress Tests↩︎

We next test whether the agreement persists under controlled changes within the linear regime. We vary the document count, shift the test-input distribution, and stack LSA layers with shared parameters. When sweeping \(n \in \{2,5,10,25\}\) and shifting the test-input range to \(\alpha \in \{0.5,1,1.5,2\}\) while keeping training fixed at \(\alpha=1\), the loss difference between the trained Transformer and the gradient predictor remains small (Figure 5, Appendix 10). Stacking LSA layers further supports the multi-step picture in Eq. 5 . At depths \(2\) and \(5\), the loss and prediction differences remain small across document counts. The residual gap at \(\mathrm{Docs}=25\) also shrinks as depth increases (Figure 6, Appendix 10). The projection-based interface shows similar behavior (Appendix 12). These results suggest that the linear correspondence is not a fragile single-step artifact, but remains stable under controlled linear extensions.

4.3 Nonlinear Stress Test↩︎

We then examine where the correspondence begins to break. We add MLP layers after the input embedding and evaluate on four real-world regression datasets: California Housing, Bike Sharing, Wine Quality, and Predict Calorie Expenditure. We focus on the dot-product interface throughout. Dataset details are in Appendix 13.0.0.2. This experiment is diagnostic. We do not claim that normalization solves RAG adaptation. Instead, we use normalization to control the feature geometry that interacts with dot-product retrieval. We compare Z-score [31], Min–Max [31], rank-based normalization [32], and Tanh normalization. The training set is used as the retrieval corpus and is normalized with Z-score throughout. Only the input-side normalization is varied, and alignment is measured using the same metrics as in Section 4.1.

Figure 2 shows two representative cases. On Bike Sharing, Min–Max normalization gives the closest agreement between the trained nonlinear Transformer and the gradient-descent predictor, likely because the features are bounded and not dominated by outliers. On California Housing, the alignment is weaker: skewed and heavy-tailed features make dot-product geometry more sensitive to outliers. The sensitivity cosine drops, the model difference grows, and prediction differences become less stable. The same pattern appears on the remaining datasets. Predict Calorie Expenditure behaves similarly to Bike Sharing, while Wine Quality behaves more like California Housing (Figure 7, Appendix 11). Overall, the linear optimization view remains informative when feature geometry is stable, but becomes less predictive when retrieval-derived dot products are dominated by skewed or heavy-tailed features. This empirical boundary supports our use of the linear construction as a guide for adaptation, rather than as a literal model of LLM computation. Next, we use this view to design a forward-only update to the generator-side evidence-use interface in LLM-scale RAG.

5 LLM-Scale RAG: Amortizing the Gradient-Descent Update↩︎

We now instantiate this view as RAG-GD, a forward-only adaptation method for frozen billion-parameter RAG LLMs. In this setting, we do not assume an exact equivalence between an LLM forward pass and gradient descent. Instead, we use gradient descent as an operational target for adapting how the generator uses retrieval-conditioned information. Given a few RAG-formatted demonstrations, we use autograd during training to compute the update that gradient descent would make to the generator-side retrieval interface. We then train a lightweight predictor to approximate this update. At inference time, the predictor produces the update in a single forward pass, without further training the RAG system or backpropagating through the frozen LLM.

Concretely, we first train a base retrieval adapter \(W_0^{\mathrm{ret}}\) on RAG-formatted examples from the NQ training split. It is a low-rank LoRA perturbation to the Q/K/V projections of every attention layer in a frozen LLM. All inputs are RAG-formatted instances \((x,\mathcal{D}_x,y)\), where \(y\) is the gold answer. In this work, we use a fixed external retriever, either BM25 or E5, to select the retrieved documents \(\mathcal{D}_x\) for each question \(x\) from a fixed corpus. This adapter serves as a generator-side retrieval interface: it does not select documents, but modulates how the generator uses evidence selected by BM25 or E5. The predictor \(g_\phi\) is then meta-trained on few-shot support contexts \(C=\{(x_i,\mathcal{D}_i,y_i)\}_{i=1}^N\) from NQ, TriviaQA, HotpotQA, 2WikiMultiHopQA, and MuSiQue. PopQA and Bamboogle are held out from both stages and used only for evaluation.

5.1 Supervision Target: Autograd-Defined Interface Update↩︎

The supervision target is the update that \(K\) SGD steps would produce on the generator-side retrieval interface using a support context \(C\). Starting from \(W_{\mathrm{ret}}^{(0)}=W_0^{\mathrm{ret}}\), we run \[\Delta W_{\mathrm{GD}}^{(K)}(C) = W_{\mathrm{ret}}^{(K)} - W_0^{\mathrm{ret}}, \qquad W_{\mathrm{ret}}^{(t+1)} = W_{\mathrm{ret}}^{(t)} - \eta \nabla \mathcal{L}(W_{\mathrm{ret}}^{(t)}; C), \label{eq:gd95target}\tag{6}\] for \(t=0,\ldots,K-1\), where \(\mathcal{L}\) is the answer-token cross-entropy conditioned on the question and retrieved documents. We compute \(\Delta W_{\mathrm{GD}}^{(K)}(C)\) with autograd and detach it from the predictor optimizer. Equation 6 is not a theorem-preserving lift of Lemma 1. The setting changes from squared regression with a linear predictor to cross-entropy training of a deep LLM, and the adapted parameter becomes a low-rank Q/K/V LoRA interface. Its role is practical: it provides an optimization-derived target for how RAG demonstrations should adjust generator-side evidence use.

5.2 Predictor Architecture and Matching Objective↩︎

The predictor \(g_\phi\) is a context encoder with per-layer, per-projection update heads. Each demonstration \((x_i,\mathcal{D}_i,y_i)\in C\) is formatted by concatenating the question, retrieved documents, and gold answer. The frozen LLM with the base adapter encodes each sequence, and we use the EOS hidden state \(h_i\in\mathbb{R}^{d_h}\). We aggregate demonstrations by mean pooling, \(\bar h(C)=\frac{1}{N}\sum_{i=1}^{N}h_i\). For each layer \(\ell\) and projection type \(\pi\in\{Q,K,V\}\), the update head outputs \(\widetilde{\Delta W}_{\ell,\pi} = U_{\ell,\pi}V_{\ell,\pi}^{\top} \in \mathbb{R}^{d\times d},\) where \(U_{\ell,\pi},V_{\ell,\pi}\in\mathbb{R}^{d\times r}\) are generated by a two-layer MLP from \(\bar h(C)\). The rank \(r\) matches the base adapter, so \(g_\phi(C)\) has the same shape as \(W_0^{\mathrm{ret}}\). We train \(g_\phi\) to match the autograd-defined target per layer and projection type. Let \(\Delta W_{\ell,\pi}^{\star} \triangleq [\Delta W_{\mathrm{GD}}^{(K)}(C)]_{\ell,\pi}.\) The matching loss is \[\mathcal{L}_{\mathrm{match}}(\phi;C) = \sum_{\ell,\pi} \left[ 1- \frac{ \langle \widetilde{\Delta W}_{\ell,\pi},\Delta W_{\ell,\pi}^{\star}\rangle }{ \|\widetilde{\Delta W}_{\ell,\pi}\|_F \|\Delta W_{\ell,\pi}^{\star}\|_F } + \lambda \left| \log \frac{ \|\widetilde{\Delta W}_{\ell,\pi}\|_F }{ \|\Delta W_{\ell,\pi}^{\star}\|_F } \right| \right], \label{eq:matching}\tag{7}\] where the cosine term matches direction and the log-magnitude term matches scale. We use \(\lambda=0.1\) throughout. At deployment, the predictor emits \(\widetilde{\Delta W}(C)\) in one forward pass, and the frozen LLM answers with the adapted interface \(W_0^{\mathrm{ret}}+\widetilde{\Delta W}(C)\).

5.3 Benchmarks, Baselines, and Metrics↩︎

5.3.0.1 Benchmarks

We evaluate on seven open-domain QA benchmarks: NQ, TriviaQA, PopQA, HotpotQA, 2WikiMultiHopQA, MuSiQue, and Bamboogle. NQ, TriviaQA, and PopQA are single-hop, while the remaining four are multi-hop. We use Qwen 2.5-7B-Instruct [33] and Llama 3.1-8B-Instruct [34] as frozen backbones. Each query is augmented with the top five documents from BM25 or E5-large [35], using the same retrieval cache for all methods. The support size is \(N=3\), and we train predictors with \(K\in\{1,5,10\}\).

5.3.0.2 Baselines and Metrics

We compare RAG-GD with six baselines. Query Only uses no retrieved documents. Vanilla RAG prepends retrieved documents to the prompt. Base adapter applies \(W_0^{\mathrm{ret}}\) without context-conditioned perturbation. + few shot variants concatenate support demonstrations into Vanilla RAG or Base adapter prompts. Prompt tuning [36] learns a soft prefix using the same supervision pool as \(W_0^{\mathrm{ret}}\). HyperTuning [21] uses the same predictor architecture as RAG-GD, but trains through downstream task loss rather than the autograd-defined target. TT-SGD performs \(K\) SGD steps on \(C\) at test time and serves as the non-amortized reference. RAG-GD shares \(W_0^{\mathrm{ret}}\) with Base adapter and adds only \(\widetilde{\Delta W}(C)\). Table 1 reports the headline comparison against no-perturbation baselines. Full results for Prompt tuning, HyperTuning, TT-SGD, and + few shot variants are in Appendix 16. We report SQuAD-style [37] exact match (EM) and token-overlap F1.

5.4 Results↩︎

Table 1 reports the headline comparison between RAG-GD (\(K{=}5\)) and the no-perturbation baselines. Figures 3 and 4 compare against additional context-conditioned methods, including Prompt tuning, HyperTuning, TT-SGD, and + few shot variants. Full per-method and per-benchmark results are in Appendix 16, and Algorithm 9 gives the deployment procedure.

5.4.0.1 The predicted update improves the base retrieval adapter.

Across all backbone and retriever configurations, RAG-GD improves average EM and F1 over Base adapter. This comparison is controlled: both methods share the external retriever, retrieval cache, frozen backbone, and base adapter \(W_0^{\mathrm{ret}}\). The only difference is the predicted perturbation \(\widetilde{\Delta W}(C)\), which isolates the effect of adapting the generator-side retrieval interface from the support context.

5.4.0.2 The learned update transfers to held-out tasks.

PopQA and Bamboogle are held out from training for both \(W_0^{\mathrm{ret}}\) and \(g_\phi\). On PopQA, RAG-GD improves EM over Base adapter in every backbone and retriever setting, with F1 close or improved in most cases. On Bamboogle, gains are strongest on Qwen, while Llama with BM25 stays close to Base adapter. The transfer is consistent but not uniform, suggesting that \(g_\phi\) learns a reusable update rule for generator-side evidence use rather than only fitting the meta-training tasks.

5.4.0.3 Gradient-update supervision matters.

Figure 3 (a) compares RAG-GD with context-conditioned baselines on Qwen-2.5-7B with E5 retrieval. HyperTuning uses the same predictor architecture but trains through downstream task loss rather than matching the autograd-defined update. RAG-GD improves average EM and F1 over HyperTuning and Prompt tuning, indicating that the gradient-derived target contributes beyond the predictor architecture.

5.4.0.4 Performance is largely insensitive to \(K\).

Figure 3 (b) sweeps the inner-loop depth \(K\in\{1,5,10\}\) across both backbones and retrievers. A single amortized step already recovers most of the gain, while additional steps bring only small and configuration-dependent changes. Thus, the \(K{=}5\) setting in Table 1 is representative. Full per-\(K\) results are in Appendix 16.

5.4.0.5 Amortization approaches test-time adaptation at lower cost.

Figure 4 shows the EM-cost tradeoff for Qwen-2.5-7B with E5 retrieval. TT-SGD performs inner-loop backpropagation through the 7B LLM at test time, while RAG-GD moves this computation into training and uses only one forward pass through \(g_\phi\) at inference. As a result, RAG-GD reaches a similar average EM and F1 operating point at substantially lower per-query cost.

5.5 Discussion↩︎

These results complete the theory-to-practice arc. The linear regime gives an exact gradient-descent correspondence, while the nonlinear experiments show where this correspondence becomes feature-dependent. At LLM scale, we do not claim that frozen RAG LLMs implement the linear equivalence internally. Instead, the autograd-defined update provides a practical target for context-conditioned adaptation. Because RAG-GD and Base adapter share the same retriever, retrieval cache, backbone, and \(W_0^{\mathrm{ret}}\), the gains isolate the contribution of the predicted update to the generator-side retrieval interface. The held-out transfer and cost-performance tradeoff suggest that gradient-supervised retrieval-interface adaptation is a promising forward-only alternative to test-time backpropagation for RAG.

6 Conclusion↩︎

We studied retrieval-augmented generation through an in-context optimization lens. In a controlled linear setting, we showed that one linear self-attention layer can implement one gradient-descent step on a unified linearized RAG loss covering projection-based and dot-product retrieval interfaces. Empirical tests verified this construction in the exact regime and revealed a structured boundary under nonlinear architectures and real regression data, where alignment becomes sensitive to feature distribution. At LLM scale, we turned this view into a forward-only adaptation method by using the autograd-defined \(K\)-step update to a generator-side Q/K/V LoRA interface as supervision for a lightweight predictor. Across seven QA benchmarks, two retrievers, and two frozen backbones, the predicted update improved a shared-adapter baseline, transferred to held-out domains, and was largely insensitive to \(K\). Overall, these results suggest that retrieved evidence can be treated not only as external context, but also as a signal for context-induced adaptation in RAG.

7 Limitations↩︎

Our linear construction is an analytical starting point rather than a literal account of modern RAG LLMs, and the nonlinear experiments show that the correspondence depends on architecture and feature distribution. At LLM scale, we instantiate the view with a generator-side Q/K/V LoRA interface while keeping the retriever and backbone fixed. Open questions remain about the inner-loop optimizer, update parameterization, predictor capacity, scaling to larger backbones, and robustness across broader retrieval settings. Future work should characterize when context-induced updates improve evidence use, when they should be suppressed, and how uncertainty-aware gating can make such updates robust to noisy retrieval.

Broader Impacts↩︎

The method we propose lowers the cost of adapting a deployed large language model to a new task or domain. At inference, the LLM still runs a forward pass to generate the answer, but no backward pass through the LLM is required: adapting to the context costs only a single small forward pass through a context-conditional weight predictor. If adopted at scale, this could reduce the energy footprint of test-time adaptation pipelines for retrieval-augmented systems. The flip side is that lower adaptation cost could also accelerate the deployment of models in domains for which the underlying LLM has not been carefully evaluated, including settings where retrieval can amplify biases in the corpus. We encourage practitioners adopting this style of adaptation to retain the same evaluation discipline that would apply to a fine-tuned model. All datasets used in our evaluation are publicly available QA and regression benchmarks that do not contain personally identifiable or sensitive information, and our work does not raise additional ethical concerns beyond those discussed above.

Reproducibility Statement↩︎

All datasets used in our main evaluation are publicly available QA benchmarks (NQ, TriviaQA, PopQA, HotpotQA, 2WikiMultiHopQA, MuSiQue, Bamboogle). The supplementary regression datasets discussed in the appendix are also public. Preprocessing steps, dataset splits, and the training pools used for the reference retrieval adapter \(W_0^{\text{ret}}\) and the predictor \(g_\phi\) are documented in the appendix.

Our implementation uses PyTorch with the Hugging Face Transformers library on top of frozen Qwen 2.5-7B-Instruct [33] and Llama 3.1-8B-Instruct [34] backbones. Hyperparameters for the reference retrieval adapter \(W_0^{\text{ret}}\), the predictor \(g_\phi\), and the inner-loop SGD target (\(\eta\), \(K\), number of demonstrations \(N\)) are listed in Appendix 14. All experiments were conducted on NVIDIA A100 GPUs.

8 Proof of Lemma 1↩︎

8.0.0.1 Statement.

Given a 1-head linear-attention layer and tokens \(e_j = (x_1^j,\, x_2^j,\, y^j)\) for \(j = 1,\ldots,N\), we construct key, query, and value matrices \(W_K\), \(W_Q\), \(W_V\) and a projection \(P\) such that one linear self-attention update on each \(e_j\) matches one gradient-descent step on the unified RAG loss of Section 3. The update modifies only the \(y\)-coordinate of each token: \[e_j \leftarrow e_j + \big(0,\; 0,\; \Delta W_1 x_1^j + \Delta W_2 x_2^j\big) \;=\; e_j + P\, V\, K^\top q_j, \label{eq:appA95update}\tag{8}\] where \(\Delta W_1\) and \(\Delta W_2\) are the gradient-step updates of Eq. 4 in the main text.

8.0.0.2 Setup.

We are given \(N\) context tokens, each of the form \(e^i = (x_1^i,\, x_2^i,\, y^i)\) corresponding to one training pair, plus a query token \(e_{N+1} = (x_1^{N+1},\, x_2^{N+1},\, 0)\) at position \(N{+}1\). The model is asked to predict the updated \(y\)-value at position \(N{+}1\).

8.0.0.3 Step 1: expand the post-update prediction.

Writing the post-update prediction \(y'\) as the original prediction plus the contributions of \(\Delta W_1\) and \(\Delta W_2\), \[\begin{align} y' &= W_1' x_1 + W_2' x_2 \nonumber \\ &= (W_1 + \Delta W_1)\,x_1 + (W_2 + \Delta W_2)\,x_2 \nonumber \\ &= W_1 x_1 + W_2 x_2 \;+\; \Delta W_1 x_1 + \Delta W_2 x_2. \label{eq:appA95yprime} \end{align}\tag{9}\]

8.0.0.4 Step 2: gradient step on the unified RAG loss.

Under the squared loss \(L(W_1, W_2) = \tfrac{1}{2N}\sum_{i=1}^N \|W_1 x_1^i + W_2 x_2^i - y_i\|^2\), one gradient step with learning rate \(\eta\) yields \[\begin{align} \Delta W_1 &= -\eta\,\nabla_{W_1} L = -\frac{\eta}{N}\sum_{i=1}^N \big(W_1 x_1^i + W_2 x_2^i - y_i\big)\,(x_1^i)^\top, \tag{10}\\ \Delta W_2 &= -\eta\,\nabla_{W_2} L = -\frac{\eta}{N}\sum_{i=1}^N \big(W_1 x_1^i + W_2 x_2^i - y_i\big)\,(x_2^i)^\top, \tag{11}\\ \Delta y &= \Delta W_1 x_1 + \Delta W_2 x_2. \tag{12} \end{align}\]

8.0.0.5 Step 3: rewrite \(\Delta y\) as a sum of outer-product contractions.

Substituting Eqs. 10 –11 into Eq. 12 and evaluating at the query token \(j\), \[\begin{align} \Delta y \;=\; -\frac{\eta}{N}\sum_{i=1}^N \big(W_1 x_1^i + W_2 x_2^i - y_i\big)\,(x_1^i)^\top x_1^j \;-\; \frac{\eta}{N}\sum_{i=1}^N \big(W_1 x_1^i + W_2 x_2^i - y_i\big)\,(x_2^i)^\top x_2^j. \label{eq:appA95dy95expanded} \end{align}\tag{13}\] Equivalently, the update applied to the token at position \(j\) is \[\begin{align} \begin{pmatrix} x_1^j \\ x_2^j \\ y^j \end{pmatrix} \;\leftarrow\; \begin{pmatrix} x_1^j \\ x_2^j \\ y^j \end{pmatrix} \;+\; \begin{pmatrix} 0 \\ 0 \\ \Delta y \end{pmatrix}, \qquad \text{with}\qquad \begin{pmatrix} 0 \\ 0 \\ \Delta y \end{pmatrix} = \begin{pmatrix} 0 \\ 0 \\ \Delta W_1 x_1 + \Delta W_2 x_2 \end{pmatrix}. \label{eq:appA95token95update} \end{align}\tag{14}\]

8.0.0.6 Step 4: cast the update as a linear self-attention output.

Using the identity \(a\, b^\top c = (a \otimes b^\top)\, c\) and grouping terms, Eq. 13 can be written as \[\begin{align} \begin{pmatrix} 0 \\ 0 \\ \Delta y \end{pmatrix} \;=\; -\frac{\eta}{N}\sum_{i=1}^N \underbrace{\begin{pmatrix} 0 \\ 0 \\ W_1 x_1^i + W_2 x_2^i - y^i \end{pmatrix}}_{\text{value vector } v_i} \,\otimes\, \underbrace{\begin{pmatrix} x_1^i & x_2^i & 0 \end{pmatrix}}_{\text{key vector } k_i^\top} \;\underbrace{\begin{pmatrix} x_1^j \\ x_2^j \\ 0 \end{pmatrix}}_{\text{query vector } q_j}. \label{eq:appA95outer} \end{align}\tag{15}\] Each factor in Eq. 15 can be obtained by applying a fixed linear projection to the token \(e^i = (x_1^i, x_2^i, y^i)\) or \(e^j\): \[\begin{align} v_i &= \underbrace{\begin{pmatrix} 0 & 0 & 0\\ 0 & 0 & 0\\ W_1 & W_2 & -I_y \end{pmatrix}}_{W_V}\, e^i, \qquad k_i = \underbrace{\begin{pmatrix} I_x & 0 & 0\\ 0 & I_x & 0\\ 0 & 0 & 0 \end{pmatrix}}_{W_K}\, e^i, \qquad q_j = \underbrace{\begin{pmatrix} I_x & 0 & 0\\ 0 & I_x & 0\\ 0 & 0 & 0 \end{pmatrix}}_{W_Q}\, e^j. \label{eq:appA95proj} \end{align}\tag{16}\]

8.0.0.7 Step 5: explicit construction.

Combining Eqs. 15 and 16 , the gradient step of Eqs. 10 –11 is realized as a linear self-attention update with the closed-form projections \[{!}{ \begin{pmatrix} x_1^j\\ x_2^j\\ y^j \end{pmatrix} \leftarrow \begin{pmatrix} x_1^j\\ x_2^j\\ y^j \end{pmatrix} -\frac{\eta}{N}\sum_{i=1}^{N} \Bigg( \underbrace{\begin{pmatrix} 0&0&0\\ 0&0&0\\ W_1&W_2&-I_y \end{pmatrix}}_{W_V} \begin{pmatrix} x_1^i\\ x_2^i\\ y^i \end{pmatrix} \Bigg) \otimes \Bigg( \underbrace{\begin{pmatrix} I_x&0&0\\ 0&I_x&0\\ 0&0&0 \end{pmatrix}}_{W_K} \begin{pmatrix} x_1^i\\ x_2^i\\ y^i \end{pmatrix} \Bigg)^\top \Bigg( \underbrace{\begin{pmatrix} I_x&0&0\\ 0&I_x&0\\ 0&0&0 \end{pmatrix}}_{W_Q} \begin{pmatrix} x_1^j\\ x_2^j\\ y^j \end{pmatrix} \Bigg) } \label{con:main95pro}\tag{17}\] The projection \(P\) is taken to be the identity on the \(y\)-coordinate, so that the value contribution lands in the \(y\)-slot of \(e^j\). Comparing the right-hand side of Eq. 17 to the gradient step in Eqs. 10 –11 , the two sides agree term by term. One linear self-attention update therefore reproduces one gradient-descent step on the unified RAG loss, as claimed. 0◻

9 Linear RAG: derivation of the unified retriever formulation↩︎

9.0.0.1 Main function.

\[y = (W_q, W_z) \begin{bmatrix} x_q \\[3pt] \sum_{i=1}^n (W_e x_q)^\top (W_e d_i)\, d_i \end{bmatrix} = W_q x_q + W_z \sum_{i=1}^n (W_e x_q)^\top (W_e d_i)\, d_i.\]

9.0.0.2 Define \(M\) and rewrite the similarity.

We adopt the shared-encoder simplification, where the same linear encoder \(W_e\) maps both queries and documents into the retrieval space. The general DPR formulation [6] permits separate \(W_q\) and \(W_d\), in which case the analysis below carries through with \(M = W_q^\top W_d\) but \(M\) need not be symmetric. Defining \[M \triangleq W_e^\top W_e \;\;\;\;\Rightarrow\;\;\;\; (W_e x_q)^\top (W_e d_i) = x_q^\top W_e^\top W_e d_i = x_q^\top M d_i.\] Hence, \[y = W_q x_q + W_z \sum_{i=1}^n (x_q^\top M d_i)\, d_i.\]

9.0.0.3 Converting “scalar \(\times\) vector” into “matrix \(\times\) vector.”

Note that \(x_q^\top M d_i\) is a scalar, and the following identity holds: \[(x_q^\top M d_i)\, d_i = d_i (d_i^\top M^\top x_q) = (d_i d_i^\top) M^\top x_q.\] Therefore, \[\sum_{i=1}^n (x_q^\top M d_i)\, d_i = \sum_{i=1}^n (d_i d_i^\top) M^\top x_q = \Big(\sum_{i=1}^n d_i d_i^\top\Big) M^\top x_q.\]

9.0.0.4 Define the document second-moment matrix \(D\).

\[D \triangleq \sum_{i=1}^n d_i d_i^\top \;\;\;\;\Rightarrow\;\;\;\; \sum_{i=1}^n (x_q^\top M d_i)\, d_i = D M^\top x_q.\]

9.0.0.5 Substituting back into \(y\).

\[y = W_q x_q + W_z D M^\top x_q.\]

Then \(M = W_e^\top W_e\) is symmetric, i.e., \(M^\top = M\). Thus the expression simplifies to \[y = W_q x_q + W_z D M x_q.\]

The right-hand side is grouped into an equivalent linear mapping: \[y = (W_q + W_z D M)\, x_q.\]

10 Linear-regime equivalence: full setup and additional figures↩︎

This appendix supplements Section 4.1 of the main paper with a full description of the synthetic-regression setup and the additional robustness and stacked-layer figures referenced there.

10.1 Setup details↩︎

10.1.0.1 Tokens.

Each token concatenates an input vector, a retrieval-derived feature, and a target, \[e_i = (x_i,\, z_i,\, y_i), \qquad i = 1, \ldots, N,\] where \(N\) is the number of in-context examples for a single task \(\tau\). The auxiliary slot \(z_i\) instantiates the unified RAG view of Section 3. Under the linear-projection retriever, \(z_i\) is a document-derived feature; under the dot-product retriever, \(z_i = x_i\) and the document information is injected through the keys and values rather than through the token (see “dot-product injection” below).

10.1.0.2 Pre-training objective.

We train an LSA layer parameterized by \(\theta\) to minimize the expected squared prediction error across tasks: \[\mathcal{L}(\theta) = \frac{1}{B} \sum_{\tau=1}^{B} \big\| \hat{y}_{\theta}\big( \{e_{\tau, i}\}_{i=1}^{N},\; e_{\tau, N+1} \big) - y_{\tau, N+1} \big\|^2,\] where the query token at position \(N+1\) is \(e_{\tau, N+1} = (x_{\text{test}}, z_{\text{test}}, 0)\) and \(y_{\tau, N+1}\) is its target. The objective is optimized with minibatch SGD over a fresh batch of tasks at each iteration. We denote the parameters at convergence by \(\theta^{*}\).

10.1.0.3 Synthetic data.

Following [2], [30], we generate each task \(\tau\) from a teacher with weights \(W_\tau \sim \mathcal{N}(0, I)\). Inputs are drawn from \(x_{\tau, i} \sim \mathcal{U}(-1, 1)^{n_I}\) and targets are constructed as \(y_{\tau, i} = W_\tau^{1}\, x_{\tau, i}^{1} + W_\tau^{2}\, x_{\tau, i}^{2}\). We set \(N = n_I = 10\) with output dimension \(1\), and we sweep the document count \(k \in \{2, 5, 10, 25\}\).

10.1.0.4 Dot-product injection.

Under the linear-projection retriever the document is included in the token, \(e_i = (x_i, \mathcal{D}, y_i)\), and the LSA layer can learn to select relevant documents during pre-training. Under the dot-product retriever the document is not concatenated into the token; instead, document information is injected directly into the key and value matrices, \[K = \begin{bmatrix} K_{\text{ctx}} \\ h_d \end{bmatrix}, \qquad V = \begin{bmatrix} V_{\text{ctx}} \\ h_d \end{bmatrix},\] where \(K_{\text{ctx}}, V_{\text{ctx}}\) are the contextual key/value rows and \(h_d = f(\mathcal{D}) \in \mathbb{R}^{B \times \dim(x)}\) is a fixed projection of the document set \(\mathcal{D} = \{d_1, \ldots, d_n\}\) into the input dimension.

10.1.0.5 Constructed reference predictor.

The trained LSA layer (parameters \(\theta^{*}\)) is compared against a constructed predictor that realizes one gradient-descent step on the unified RAG loss exactly. Following the construction of Eq. 17 , we set the value, key, and query projections so that one LSA update reproduces the gradient step. For the linear-projection retriever, \(W_1\) and \(W_2\) inside \(W_V\) are initialized to zero, following [2]. For the dot-product retriever, \(W_2 = W_z\big(\sum_i d_i d_i^\top\big) M^\top\), with \(W_z \in \mathbb{R}^{d_y \times d_d}\) and \(M \in \mathbb{R}^{d_d \times d_d}\) sampled independently from \(\mathcal{N}(0, \sigma^2)\), and document features \(C \sim \mathcal{U}(-\tfrac{1}{2}, \tfrac{1}{2})^{k \times d_d}\). The inner-loop learning rate \(\eta\) for the constructed predictor is chosen by line search to minimize the constructed model’s loss over \(10^4\) training tasks. We write the resulting prediction \(\hat{y}_{\theta, \mathrm{rag}}(x_{\text{test}})\).

10.1.0.6 Evaluation metrics.

On \(T_{\text{val}} = 10^4\) held-out validation tasks, following [2], we report the mean of three quantities between the trained and constructed predictors: (i) the prediction difference \(\big\| \hat{y}_{\theta^{*}}(x_{\tau, \text{test}}) - \hat{y}_{\theta, \mathrm{rag}}(x_{\tau, \text{test}}) \big\|_2\); (ii) the cosine similarity between the input-sensitivities \(\partial \hat{y}_{\theta, \mathrm{rag}}/\partial x_{\text{test}}\) and \(\partial \hat{y}_{\theta^{*}}/\partial x_{\text{test}}\); and (iii) the corresponding \(\ell_2\) sensitivity difference.

10.2 Robustness to distribution shift and document count↩︎

To probe whether the trained LSA layer captures a generalizable update rule rather than memorizing the training distribution, we vary two factors at test time. First, we sweep the document count \(n \in \{2, 5, 10, 25\}\) and recompute the comparison; second, we sample test inputs from \(\mathcal{U}(-\alpha, \alpha)^{n_I}\) with \(\alpha \in \{0.5, 1, 1.5, 2\}\) while training is fixed at \(\alpha = 1\). Figure 5 reports the resulting loss curves. With the linear-projection retriever, the absolute loss rises with document count (the projected document features carry more variance) but the LSA layer follows the gradient predictor in lockstep. With the dot-product retriever, where the document information enters through the second-moment matrix \(\sum_i d_i d_i^\top\), the loss is largely insensitive to document count. The dot-product variant is also computationally cheaper, since no per-document projection is required.

10.3 Stacked-layer agreement under the dot-product retriever↩︎

Figure 6 reports the dot-product variant at depths 2 and 5. The loss differences between the trained Transformer and the constructed predictor remain small across document counts, and the prediction differences converge to similar values. The number of retrieved documents has a depth-dependent effect on the residual: at depth 2, the prediction difference is smaller for \(\text{Docs}=2\) than for \(\text{Docs}=25\), but this gap narrows at depth 5. The corresponding analysis for the linear-projection retriever is reported in Appendix 12.

11 Normalization analysis: per-dataset extended results↩︎

This appendix supplements Section 4.3 of the main paper. Section 4.3 reports the headline result on Bike Sharing and California Housing; here we cover the remaining two datasets, Predict Calorie Expenditure and Wine Quality. Datasets and normalization methods are as defined in Section 4.3 and Appendix 13.0.0.2. The training set is used as the retrieval corpus and is normalized with Z-score throughout; only the input-side normalization is varied, between Z-score [31], Min–Max [31], rank-based [32], and Tanh [38].

On Predict Calorie Expenditure, the trained Transformer continues to track the gradient-descent predictor closely, mirroring the alignment seen on Bike Sharing in Section 4.3. Wine Quality is the harder case. Under Min–Max normalization, a few outliers dominate the scaling and compress most of the samples near zero. Two effects follow. The sensitivity cosine drops because the sensitivity vectors diverge from those of the gradient-descent predictor. The prediction difference also fluctuates more strongly, indicating instability in the alignment between RAG and ICL dynamics under heavy-tailed feature distributions. This is consistent with the California Housing pattern in Section 4.3: when retrieval-derived dot products are dominated by skewed features, the linear correspondence becomes less predictive.

12 Stacked-layer agreement under the projection-based retriever↩︎

This appendix complements the dot-product analysis of Appendix 10 with the projection-based retriever. Under this interface, the retrieved documents are concatenated with the input tokens, \(e_i = (x_i, \mathcal{D}, y_i)\), so each per-document feature is processed jointly with the query through every stacked layer. As the document count grows, the variance of the per-token features grows with it, which amplifies the discrepancy between the trained Transformer and the constructed gradient-descent predictor at any fixed depth. Stacking layers reduces this discrepancy: at depth 5, the loss and prediction differences are uniformly smaller across document counts than at depth 2, and the sensitivity cosine is closer to 1.

Whereas the dot-product sweep in Appendix 10 reports results at 2, 5, 10, and 25 documents, the projection-based sweep is restricted to 2, 5, 10, and 15. The compute cost of stacking concatenated-document tokens grows substantially with retrieval size, and the growing-residual trend is already clear at 15 documents, so the 25-document run is omitted.

13 Dataset details↩︎

13.0.0.1 QA benchmarks (main evaluation).

The seven question-answering benchmarks used in the main evaluation are all publicly available:

• Natural Questions (NQ) [39]: open-domain factoid QA over Wikipedia.

• TriviaQA [40]: large-scale trivia question answering with evidence documents.

• PopQA [41]: popularity-stratified entity-centric QA (held out from training).

• HotpotQA [42]: multi-hop QA with comparison and bridge questions.

• 2WikiMultiHopQA [43]: multi-hop questions grounded in Wikipedia article pairs.

• MuSiQue [44]: compositional multi-hop QA constructed by composing single-hop pairs.

• Bamboogle [45]: small multi-hop benchmark on long-tail entities (held out from training).

For each benchmark we use the standard FlashRAG [46] corpus and retrieval splits. NQ is used to train the source retrieval adapter \(W_0^{\mathrm{ret}}\). NQ, TriviaQA, HotpotQA, 2WikiMultiHopQA, and MuSiQue are used to meta-train the predictor \(g_\phi\). PopQA and Bamboogle are held out from both adapter training and predictor meta-training, and are used only for evaluation.

13.0.0.2 Synthetic and tabular regression datasets (linear-attention sanity check and normalization analysis).

• California Housing: Given eight features, [’MedInc’, ’HouseAge’, ’AveRooms’, ’AveBedrms’, ’Population’, ’AveOccup’, ’Latitude’, ’Longitude’] , the task is to predict MedHouseVal. The dataset is split into 16,640 training samples and 2,000 test samples.

• Bike Sharing: Using the features [’season’, ’yr’, ’mnth’, ’hr’, ’holiday’, ’weekday’, ’workingday’, ’weathersit’, ’temp’, ’atemp’, ’hum’, ’windspeed’, ’casual’, ’registered’], the task is to predict count. The dataset contains 15,641 training samples and 1,738 test samples.

• Wine Quality: Given eleven physicochemical features, [fixed acidity, volatile acidity, citric acid, residual sugar, chlorides, free sulfur dioxide, total sulfur dioxide, density, pH, sulphates, alcohol] the task is to predict the wine quality (a sensory score ranging from 0 to 10). The dataset is split into 4,408 training samples and 490 test samples.

• Predict Calorie Expenditure: Using the features [Gender, Age, Height, Weight, Duration, Heart_Rate, Body_Temp], the task is to predict the number of Calories expended. The dataset is split into 13,500 training samples and 1,540 test samples.

14 Implementation Details↩︎

14.0.0.1 Source retrieval adapter \(W_0^{\mathrm{ret}}\).

We implement the generator-side retrieval interface with LoRA modules on the \(\{q,k,v\}\) projections of every transformer block. The LoRA rank is \(16\), with \(\alpha{=}32\) and dropout \(0\). The LLM backbone is kept frozen and loaded in 4-bit NF4 precision. We train \(W_0^{\mathrm{ret}}\) with AdamW using learning rate \(10^{-4}\), weight decay \(0.01\), gradient clipping \(1.0\), and gradient accumulation \(4\) for \(3{,}000\) steps on RAG-formatted examples from the NQ training split. Each example is paired with the top-\(K_{\mathrm{ret}}{=}5\) retrieved documents from the fixed external retriever. For each retriever setting, we use the corresponding fixed retrieval cache and keep this cache unchanged across methods. This adapter serves only as a generator-side evidence-use interface: it does not select documents, but modulates how the frozen generator uses the retrieved evidence.

14.0.0.2 Predictor \(g_\phi\).

The predictor maps a support context \(C=\{(x_i,D_i,y_i)\}_{i=1}^{N}\) with \(N{=}3\) demonstrations to a context-conditioned update for the generator-side retrieval interface. Each demonstration is formatted by concatenating the question, retrieved documents, and gold answer, and is encoded by the frozen LLM equipped with \(W_0^{\mathrm{ret}}\). We take the EOS hidden state \(h_i\) of each demonstration and aggregate the support context by mean pooling, \[\bar{h}(C)=\frac{1}{N}\sum_{i=1}^{N}h_i .\] The pooled representation is passed to a two-layer MLP encoder with hidden dimension \(256\) and output dimension \(64\), followed by per-layer and per-projection update heads for the \(\{q,k,v\}\) LoRA modules. The update heads output low-rank perturbations with the same LoRA rank as \(W_0^{\mathrm{ret}}\), so that \(g_\phi(C)\) has the same parameter shape as the base retrieval adapter. We train \(g_\phi\) with AdamW using learning rate \(5\times10^{-4}\), weight decay \(0.01\), gradient clipping \(1.0\), and gradient accumulation \(4\) for \(3{,}000\) steps. The predictor is meta-trained on support contexts sampled from NQ, TriviaQA, HotpotQA, 2WikiMultiHopQA, and MuSiQue. PopQA and Bamboogle are excluded from both adapter training and predictor meta-training, and are used only for held-out evaluation.

14.0.0.3 Matching objective.

The predictor is trained to match the autograd-defined inner-GD target for each layer and projection type. The matching loss uses a cosine term for update direction and a log-magnitude term for update scale, as defined in Section 5.2. We set the magnitude weight to \(\lambda{=}0.1\) throughout. No downstream answer loss is applied when training \(g_\phi\) for the main RAG-GD results.

14.0.0.4 Inner GD target.

For each support context, the supervision target is computed by running \(K\) steps of SGD on the \(N{=}3\) demonstrations, starting from \(W_0^{\mathrm{ret}}\). The inner-loop learning rate is \(\eta{=}10^{-2}\), and we evaluate \(K\in\{1,5,10\}\). Gradients are taken only with respect to the LoRA parameters of the generator-side retrieval interface; the external retriever and the LLM backbone remain fixed.

14.0.0.5 Compute.

All experiments are run on NVIDIA A100 80GB GPUs. A single training run for \(W_0^{\mathrm{ret}}\) takes approximately \(50\)–\(60\) minutes. Training \(g_\phi\) takes approximately \(1.5\) hours for \(K{=}1\), \(4\) hours for \(K{=}5\), and \(7\) hours for \(K{=}10\). The full result table requires approximately \(50\) A100-hours.

15 Algorithm and inference procedure for RAG-GD↩︎

Algorithm 9 summarises the deployment-time procedure of RAG-GD: build a RAG-formatted support context, predict the retrieval-interface update with a single forward pass through \(g_\phi\), and generate with the perturbed interface. No backward pass through the LLM is required at deployment.

16 Full per-method comparison on QA benchmarks↩︎

We split the per-benchmark numbers into two tables. Table 2 reports the methods that we ran on both Qwen-2.5-7B-Instruct and Llama-3.1-8B-Instruct: Query Only, Vanilla RAG, Base adapter, and RAG-GD at \(K\in\{1,5,10\}\). Table 3 reports the additional context-conditioned baselines that we ran only on Qwen due to compute constraints: Vanilla RAG + few shot, Base adapter + few shot, Prompt tuning, HyperTuning, and TT-SGD at \(K{=}5\). Together they complement the slim main-text Table 1 and the per-family aggregates in Figures 3 and 4. HyperTuning [21] uses the same predictor architecture as RAG-GD but supervises against a downstream task loss instead of the SGD-update target, so the contrast between the HyperTuning rows in Table 3 and the RAG-GD (\(K{=}5\)) rows for Qwen in Table 2 isolates the choice of supervision signal. TT-SGD performs \(K{=}5\) inner gradient-descent steps at test time and serves as a reference for what test-time gradient adaptation would achieve at the same backbone and retriever. + few shot variants concatenate the support demonstrations into the prompt under the corresponding base configuration.

2.8pt

2.8pt

17 Use of large language models↩︎

Large language models (LLMs) were only used to assist with language polishing and minor grammatical editing of this manuscript.

References↩︎

---

1. indicates equal contribution↩︎