Log-Augmented Generation: Scaling Test-Time Reasoning with Reusable Computation¶
🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link
Why this paper was selected
MIT (Roth, Madden, Andreas); reusable reasoning logs for test-time scaling, broadly applicable
Ask a follow-up
Open an assistant pre-loaded with this paper's context.
💬 Ask ChatGPT✦ Ask Claude
TL;DR¶
Log-Augmented Generation (LAG) enables LLMs to reuse prior task reasoning at test time by encoding past multi-turn reasoning traces into KV caches and retrieving them for new, similar tasks. The key insight is that a token's KV value encodes attention over the entire preceding context, so storing only last-round tokens—computed over the full reasoning trace—compactly preserves the complete reasoning context. LAG outperforms reflection-based and KV-cache baselines on both knowledge- and reasoning-intensive benchmarks while also reducing the number of inference iterations.
Problem¶
LLMs process each task independently and discard all intermediate reasoning, causing redundant computation when tasks share overlapping subproblems (e.g., two multi-hop QA queries that pass through the same intermediate entity). Reflection-based approaches (e.g., Dynamic Cheatsheet) distill logs into abstract summaries, losing the original reasoning richness. Existing KV-cache methods (e.g., Block Attention) encode and store the same tokens, failing to capture the full multi-turn reasoning context, and are designed for efficiency rather than accuracy improvement.
Method¶
LAG augments a ReAct-style agentic loop with three components:
Encoding & Storage. After each completed task, the model encodes the concatenated sequence of all multi-turn model responses to produce KV values. Only the KV values corresponding to tokens from the last model response are stored in a log store \(\mathcal{L}\). This decouples the encoding context (full reasoning trace) from storage (last-round tokens), capturing rich semantics cheaply.
Retrieval. Given the next action \(a\) (a sub-query or follow-up), a text embedding model computes its embedding and retrieves the top-\(k\) logs by cosine similarity against pre-computed log embeddings.
Generation with RoPE re-embedding. Retrieved KV values carry positional IDs from their original encoding context. LAG strips the original RoPE rotation and reapplies it with new positional IDs matching the current context: $\(\text{RoPE}(x, \theta') = \begin{pmatrix}\cos\theta' & -\sin\theta' \\ \sin\theta' & \cos\theta'\end{pmatrix} \underbrace{\begin{pmatrix}\cos\theta & \sin\theta \\ -\sin\theta & \cos\theta\end{pmatrix}\text{RoPE}(x,\theta)}_{=x}\)$ The re-embedded KV caches are prepended to the prompt for generation.
Key Contributions¶
- LAG framework: a general store→retrieve→augment loop that integrates minimally with any sequential LLM generation workflow.
- Encode-vs-store decoupling: encodes the full multi-turn reasoning trace but stores only last-round KV values, achieving full-context semantics at reduced storage cost.
- RoPE re-embedding: a method to re-anchor retrieved KV caches to new positional IDs, enabling correct multi-log concatenation.
- Empirical evidence that direct computation reuse outperforms reflection and KV-cache efficiency baselines across knowledge- and reasoning-intensive tasks.
Results¶
All numbers are for \(\text{LAGKV}\) (Llama-3.1-8B-Instruct + Snowflake-arctic-embed-m-v2.0, top-3 log retrieval, 70/30 seen/unseen split) vs. standard agentic (ReAct):
- Musique (unseen): LAGKV 32.2 EM / 45.0 F1 / 2.68 iters vs. standard 27.0 EM / 37.3 F1 / 3.90 iters
- 2WikiMultiHop (unseen): LAGKV 55.2 EM / 64.7 F1 / 1.84 iters vs. standard 51.6 EM / 60.2 F1 / 2.49 iters
- GPQA (unseen): LAGKV 30.4 EM / 1.62 iters vs. standard 18.5 EM / 1.87 iters (+11.9 EM, largest margin)
- MMLU-Pro (unseen): LAGKV 42.3 EM / 1.36 iters vs. standard 41.3 EM / 1.58 iters
- LAGKV outperforms Block Attention (KV cache baseline) and Dynamic Cheatsheet (reflection baseline) on all datasets
- \(\text{LAGtext}\)-last (stores text from last round only) consistently ranks second among all methods, confirming both the log-reuse principle and the added value of KV representation
Limitations¶
- Requires open-weight LLMs with direct KV cache access; incompatible with closed-source API-only models.
- Storage overhead grows with the number of stored tokens per log; performance vs. storage tradeoff is confirmed but not fully characterized in the main paper.
- Main results use only Llama-3.1-8B-Instruct; larger-model evaluations (Qwen3-30B, Llama-3.1-70B) are relegated to appendices.
- No filtering of incorrect logs at storage time in the main setup; noisy logs from wrong reasoning traces may degrade retrieval quality.
- Static log store setting (offline indexing) is the primary evaluation; the dynamic (online, growing) setting is in the appendix only.
Relevance to Agentic AI / LLM Agents¶
LAG directly addresses the cross-task memory deficit in agentic systems: rather than treating each task as isolated, it provides a lightweight mechanism for agents to reuse past reasoning, effectively implementing episodic memory at the KV-cache level. The demonstration that direct computation reuse outperforms reflection/distillation is practically significant for long-horizon agent deployments where reflection adds latency and abstraction degrades precision. LAG's integration as a minimal wrapper around ReAct-style loops makes it immediately applicable to existing agentic frameworks, and its efficiency gains (fewer iterations) compound with accuracy improvements—a rare combination in test-time scaling work.