ReasonCACHE: Efficient Multi-Step Reasoning

• ReasonCACHE is a mechanism that enables large language models to engage in complex multi-step reasoning by injecting trainable prefix key-value pairs without modifying the backbone.
• It combines the benefits of in-context learning and in-weight learning to overcome context window limitations and computational scaling issues.
• Empirical outcomes show that ReasonCACHE improves accuracy and efficiency, reducing data requirements and inference cost compared to traditional methods.

ReasonCACHE is a mechanism for teaching LLMs to reason without any weight updates, by leveraging prefix tuning to inject new skills into the model through a compact, trainable key-value (KV) cache. Developed as a "middle path" between in-context learning (ICL) and in-weight learning (IWL), ReasonCACHE overcomes the scaling and expressivity limitations of ICL without the drawbacks associated with parameter updates in IWL. Empirical and theoretical results demonstrate that ReasonCACHE matches or surpasses strong weight-update-based baselines in challenging reasoning domains, while being more efficient in terms of data, inference cost, and trainable parameter budget (Gupta et al., 2 Feb 2026).

1. Motivation and Problem Statement

LLM adaptation traditionally occurs via either ICL or IWL. In ICL, task demonstrations are concatenated to the prompt at inference time, leveraging the model’s frozen weights to process the input and build a KV-cache that subsequent tokens can attend to. While ICL is highly sample-efficient—often requiring only a few demonstrations—it faces fundamental challenges for reasoning-heavy tasks:

• Context window bottleneck: Only O(1–10k) tokens fit within typical model limits.
• Computational scaling: Attention computation grows quadratically with context length (in prefill), and decoding remains linear.
• Performance saturation: Adding more demonstrations often leads to either no improvement or degradation (“context rot”).
• Limited reasoning depth: ICL is often shallow, failing to generalize beyond pattern mimicry on multi-step reasoning.

IWL, involving weight updates via fine-tuning or adapter-based methods such as LoRA, internalizes new skills without context-size limitations and can, in principle, support complex reasoning. However, it is data-hungry, computationally expensive, and susceptible to catastrophic forgetting upon learning new tasks. Each skill increases parameter storage and distribution overhead.

ReasonCACHE addresses these limitations by enabling the acquisition of new reasoning capabilities—specifically, multi-step reasoning—without modifying pretrained weights and without overloading the context window (Gupta et al., 2 Feb 2026).

2. Technical Methodology and Formulation

2.1 Prefix Tuning in Transformer Architectures

ReasonCACHE utilizes prefix tuning within the Transformer attention mechanism. For each layer $\ell$, with hidden representations $H^{(\ell-1)}\in\mathbb{R}^{n\times d}$, standard attention queries, keys, and values are obtained via:

$$Q^{(\ell)} = H^{(\ell-1)}W_Q^{(\ell)},\;\; K^{(\ell)} = H^{(\ell-1)}W_K^{(\ell)},\;\; V^{(\ell)} = H^{(\ell-1)}W_V^{(\ell)}$$

and attention output at this layer is: $$Y^{(\ell)} = \operatorname{softmax}\left(\frac{Q^{(\ell)} (K^{(\ell)})^{\top}}{\sqrt{d}}\right)V^{(\ell)}$$

ReasonCACHE introduces, for each layer $\ell$, $m$ trainable prefix key-value pairs $$P_K^{(\ell)}, P_V^{(\ell)}\in\mathbb{R}^{m\times d}$$, which are concatenated to the token-derived keys and values. The modified attention becomes:

$$\widetilde{K}^{(\ell)} = \begin{bmatrix}P_K^{(\ell)} \ K^{(\ell)}\end{bmatrix},\;\; \widetilde{V}^{(\ell)} = \begin{bmatrix}P_V^{(\ell)} \ V^{(\ell)}\end{bmatrix}$$

$$\widetilde{Y}^{(\ell)} = \operatorname{softmax}\left(\frac{Q^{(\ell)} (\widetilde{K}^{(\ell)})^{\top}}{\sqrt{d}}\right) \widetilde{V}^{(\ell)}$$

All backbone weights ($W_Q$, $W_K$, $W_V$, etc.) are kept frozen. Only the prefix tensors $\{P_K^{(\ell)},P_V^{(\ell)}\}_{\ell=1}^L$ are trained.

2.2 Relationship to Prior Methods

• ICL: Special case with $P_K^{(\ell)}=K_{\mathrm{demo}}^{(\ell)}$, $P_V^{(\ell)}=V_{\mathrm{demo}}^{(\ell)}$, no learning.
• Prompt tuning: Prefix is output of an embedding generator; less expressive than free parameter prefix tuning.
• Prefix tuning in ReasonCACHE: Free parameters at every layer, strictly increasing expressivity beyond prompt tuning.

3. Training Regime and Implementation

3.1 Data Selection and Preprocessing

• MetaMathQA: $\sim$400K math problems with chain-of-thought solutions.
• OpenThoughts-3: Reasoning traces filtered to length $\leq4096$ tokens.

Demonstrations requiring more than 4096 tokens are removed. This reflects deployment constraints.

3.2 Objective and Hyperparameters

Training minimizes next-token log-likelihood:

$$\min_{\mathcal{P}}\;\sum_{j=1}^{M} -\sum_{t=1}^{|D_j|} \log p_\theta \left(x_t^{(j)} \mid x_{<t}^{(j)}, \mathcal{P} \right)$$

Key hyperparameters include:

• Optimizer: AdamW (no weight decay), cosine LR schedule, 5% warmup.
• Learning rates: $1\times10^{-3}$ for prefix/prompt tuning and LoRA; $2\times10^{-5}$ for full fine-tuning.
• Prefix lengths: $m\in \{2^k: k=0,\dots,10\}$ (i.e., 1–1024).
• LoRA ranks: $r\in \{2^k: k=0,\dots,7\}$.
• Batch sizes: MetaMathQA (128, seq-len 2048); OpenThoughts-3 (32, seq-len 8192).

4. Theoretical Expressivity Comparison

The expressivity of ReasonCACHE (prefix tuning) is contrasted with low-rank weight update methods such as LoRA at a single attention layer.

Let $X\in\mathbb{R}^{n\times d}$, with input span $S_X = \operatorname{span}\{v_1,\dots,v_n\}$ of token values, input rank $t_X = \operatorname{rank}(X)$, and novelty capacity $\nu_X = d - \dim S_X$. A subspace $U\subseteq S_X^\perp$ cannot be produced by the base model.

• LoRA Bottleneck: A subspace $U\subseteq S_X^\perp$ is realizable via a rank-$r$ update iff $\dim U \le \min\{t_X,r\}$.
• Prefix Tuning Bottleneck: A subspace $U\subseteq S_X^\perp$ is realizable via prefix tuning with $m$ key-value vectors iff $\dim U \le m$.

Define the families: $$\mathcal{L}(r) = \{U \subseteq S_X^\perp : \dim U \le \min(t_X,r)\}$$

$$\mathcal{P}(m) = \{U \subseteq S_X^\perp : \dim U \le m\}$$

Key Theorem: $$\mathcal{L}(r)\subseteq\mathcal{P}(m)\iff \min(t_X,r)\le m$$; ${m} > \min\{t_X,r\}$ implies prefix tuning strictly surpasses any rank-$r$ weight update in expressivity. For LoRA applied only to $W_Q, W_K$ ($r=0$), prefix tuning with any $m\geq1$ is strictly more expressive, since $\mathcal{L}(0)=\{\{0\}\}\subset \mathcal{P}(m)$.

This demonstrates that ReasonCACHE, via prefix tuning, bypasses the carrier bottleneck imposed by input rank in low-rank adapter methods.

5. Empirical Outcomes and Comparative Analysis

Experiments freeze the backbone weights of LLMs such as LLaMA-2 7B and Qwen-2.5-7B; methods compared include ICL, prompt tuning, LoRA, full supervised fine-tuning (SFT), and ReasonCACHE.

5.1 Benchmarks and Evaluation Metrics

Benchmarks: GSM8K, MATH (high-school math), GPQA-Diamond, AIME (graduate-level math/physics). Metrics: Exact-match accuracy (pass@1), inference TFLOPs (pre-fill + decoding), generation length, trainable parameter count, data-efficiency.

5.2 Main Results

Method	GPQA-Diamond Accuracy (%)
ICL	31.8
Prompt Tuning	34.2
LoRA	38.5
Full SFT	39.1
ReasonCACHE	41.92

ReasonCACHE achieves the highest accuracy, outperforming all competitive baselines. On GSM8K and MATH, it similarly surpasses ICL and matches or exceeds LoRA and SFT.

5.3 Efficiency and Parameter Analysis

• Data Efficiency: To reach 50% accuracy on GSM8K, ReasonCACHE uses 59% less data than LoRA.
• Inference Cost: On GSM8K, replacing long prompts (∼2K tokens) with short prefixes ($m\leq512$) reduces prefill cost from $O(S^2)$ to $O(m^2)$; ReasonCACHE achieves 44.8 percentage points higher accuracy than ICL with 90% less inference compute.
• Reasoning Chain Length: On GPQA-Diamond, ReasonCACHE yields 34% shorter generation chains than SFT and 11 percentage points higher accuracy.
• Parameter Efficiency: For target accuracy 50% on GSM8K, ReasonCACHE requires only 54% as many trainable parameters as LoRA.

6. Implications, Limitations, and Open Directions

6.1 Implications

• Modularity: Prefixes constitute a pluggable “skill cache” that does not alter the backbone, allowing independent composition and removal of learned skills.
• No Catastrophic Forgetting: ReasonCACHE preserves all pretrained model capabilities; removing prefixes restores original behavior.
• Inference Scalability: Fixed-size prefixes decouple adaptation cost from context length, allowing efficient reasoning skill transfer without context constraints.
• Enhanced Expressivity: Unlike low-rank adapters, prefix tuning in ReasonCACHE is not limited by the span of input tokens, enabling strictly higher expressivity for a given parameter budget.

6.2 Limitations and Prospective Research

• Prefixes are trained offline and lack continual adaptation at test time.
• Runtime engine support for prefix injection (e.g., native support in vLLM) is rare.
• Open research directions include:
  • Compositionality: Can multiple prefixes be combined or arbitrated for multi-task transfer?
  • Memory hierarchy: Where should prefixes be placed (across layers) to encode abstract versus concrete knowledge?
  • Continual learning: Mechanisms for dynamic prefix adaptation during inference.

ReasonCACHE establishes that LLMs can acquire multi-step reasoning capabilities, with competitive or superior accuracy and efficiency, by distilling task demonstrations into a trainable, fixed-sized KV cache—without tuning or augmenting the backbone parameters (Gupta et al., 2 Feb 2026).