PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Augmentation Unlocks Efficient Chain-of-Thought Reasoning

Simeng Zhang, Yilong Chen, Wenyuan Zhang, Zhenyu Zhang, Yao Chen, Junyuan Shang, Tingwen Liu

arXiv:2608.21265v1cs.CL

TL;DR

Aggressive Chain-of-Thought compression reduces decode overhead but can discard reasoning support needed for accurate problem solving. Memory-Augmented Compression retrieves reusable reasoning memories as prefill scaffolds, improving compressed reasoning accuracy while lowering latency than standard CoT.

  • Problem

    Generation-side reasoning compression can discard useful subgoals, constraints, or logical dependencies, while long autoregressive traces remain a critical latency bottleneck.

  • Method

    Memory-Augmented Compression is a training-free framework that retrieves abstract memories from historical traces and injects them into the prefill context as reasoning scaffolds.

  • Results

    Memory improves accuracy over compressed baselines across diverse reasoning tasks while achieving lower latency than standard CoT.

  • Takeaways & Limitations

    Relevant prefill-side reasoning memories can compensate for information lost under aggressive compression and remain compatible with multiple efficiency mechanisms.

  • Takeaways & Limitations

    The method depends on memory-bank and retrieval quality, and irrelevant or poorly matched memories may distract the model and degrade reasoning quality.

Abstract

from arXiv · show

Large language models often rely on Chain-of-Thought (CoT) reasoning to solve complex tasks, but verbose reasoning traces introduce substantial inference overhead. CoT compression shortens generation, yet aggressive compression may disrupt logical coherence and degrade performance. We formalize this trade-off as the \textit{Context-Generation Substitution Law}, where explicit reasoning context substitutes for part of decode-time generation. Based on this principle, we propose \textit{Memory-Augmented Compression}, a training-free framework that constructs reusable reasoning memories from historical traces and retrieves them as prefill-side scaffolds. Rather than using raw demonstrations, these memories summarize reusable reasoning patterns, key constraints, and critical operations to compensate for information lost during compression. Experiments show that Memory consistently improves prompt-based Chain-of-Draft (CoD) compression across mathematical reasoning, complex reasoning, and science question answering tasks, yielding accuracy gains of 21.4, 28.0, 29.5, and 6.61 points over CoD on GSM8K, MATH, BBH, and MMLU-Sci, while achieving a 1.14--1.49$\times$ latency speedup over standard CoT. Memory is also compatible with token-level, reasoning-trace-level, and inference-state compression mechanisms. Further analyzes show that the gains come from relevant reasoning memories rather than simply increasing context length.

1 Introduction

The introduction frames verbose Chain-of-Thought generation as an efficiency problem and proposes shifting reusable reasoning support into prefill context. It presents Memory-Augmented Compression as a training-free approach that retrieves abstract memories to preserve accuracy under shorter reasoning traces.

  • Verbose reasoning traces improve difficult-task performance but incur substantial token-by-token generation costs.
  • Aggressive compression can discard useful subgoals, constraints, and logical dependencies, threatening reasoning quality.
  • The Context-Generation Substitution Law formalizes how prefill-side reasoning context can substitute for decode-time generation.The motivation is that prefill computation is more parallelizable than autoregressive decoding.
  • Memory-Augmented Compression is training-free and retrieves abstract memories from historical traces as prefill-side scaffolds for compressed reasoning.These memories capture reusable problem-solving patterns, key constraints, and critical operations.
  • Experiments show improved accuracy-latency trade-offs across domains, models, and compression mechanisms, with relevant memories compensating for information lost during aggressive compression.

2 Preliminaries

CoT inference generates an intermediate reasoning sequence before producing the final answer, improving accuracy but increasing autoregressive decoding cost. Memory-augmented compression instead adds retrieved reasoning context to shorten decoding, framing the trade-off between memory prefill and compressed generation.

  • CoT inference: CoT decomposes inference into reasoning-sequence generation followed by answer generation conditioned on that sequence.The joint probability is Pθ(y, z | x) = Pθ(z | x) · Pθ(y | x, z).
  • Decoding bottleneck: Long reasoning traces create a serving-latency bottleneck because autoregressive decoding grows with sequence length and proceeds sequentially.Ldec = |z| can reach thousands of tokens in modern reasoning models.
  • Memory-augmented compression: External memory C stores reusable reasoning information distilled from historical examples and supplies retrieved context instead of relying solely on fully generated reasoning.At inference, memories are injected into the prefill context to support compressed Short-CoT reasoning.
  • Efficiency–accuracy trade-off: Memory conditioning produces a compressed chain z′ with ideally |z′| ≪ |z| while preserving task performance, balancing added prefill context against decode-time savings.The trade-off includes memory length |M|, compressed reasoning length |z′|, compression penalty Lperf, and relative cost γ = τpre/τdec.

3 Memory-Augmented Compression

Memory-Augmented Compression shifts reusable reasoning support from costly decode-time working memory into explicit prefill-side memory. By retrieving compact reasoning structures, it compensates for aggressive compression while improving the accuracy–latency trade-off.

  • Cognitive memory view: Generated CoT forms autoregressive working memory and becomes the main inference bottleneck when verbose reasoning is decoded token by token.Explicit context is processed in parallel during prefill, giving it lower marginal cost than autoregressive generation.
  • Motivation and trade-off: Explicit memory compensates for compression collapse by recovering accuracy lost when decode-time reasoning is aggressively shortened.Injecting prior reasoning examples substantially recovers degradation caused by reduced reasoning length.
  • Motivation and trade-off: A net latency reduction occurs when added prefill cost is smaller than the decode-time savings from reducing required reasoning length.Relevant memory acts as a scaffold that bypasses redundant or recoverable intermediate steps.
  • Framework: Memory-Augmented Compression is a plug-in framework with memory-bank construction, retrieval, memory-augmented prefill, and memory-guided compressed inference.It shifts reasoning support from decode-time generation to prefill-side memory.
  • Framework: The memory bank transforms solved examples into compact representations of problem types, constraints, subgoals, strategies, and critical operations rather than raw demonstrations.Retrieved memories are selected using reasoning tags and semantic similarity, then concatenated with the prompt and query to guide concise reasoning.

4 Experiments

Experiments across mathematical, general reasoning, and science benchmarks show that Memory-Augmented Compression recovers accuracy lost by CoD while retaining an end-to-end latency advantage. Ablations and analytical studies attribute the gains to compact, relevant reasoning structure that supports compressed generation, with compatibility across retrievers, compression methods, models, and compression ratios.

  • General evaluation across domains: Memory-Augmented Compression improves CoD by 21.4, 28.0, 29.5, and 6.61 points on GSM8K, MATH, BBH, and MMLU-Sci, respectively.CoD+Memory approaches standard CoT on GSM8K and MATH and surpasses it on BBH and MMLU-Sci.
  • Accuracy–latency trade-off: 1.14–1.49× latency speedup is achieved across domains despite increased prefill tokens, and end-to-end latency remains below standard CoT after retrieval overhead.Memory reduces decode-side reasoning compared with standard CoT while preserving the end-to-end speed advantage.
  • Plug-in compatibility: Memory augmentation improves TokenSkip across compression ratios and models and yields positive gains with RPC and Extra-CoT under otherwise identical settings.The compatible methods include training-based approaches such as TokenSkip and Extra-CoT.
  • Impact of retrieval strategy: Reasoning Tag achieves the highest accuracy at ≈683 ms/query, while BM25 costs ≈1 ms and retains most accuracy; all retrievers outperform no-memory CoD.These results indicate robustness to retriever choice and adaptation to different deployment constraints.
  • Impact of memory representation: Summary Memory gains +6.67 points, whereas fixed demonstrations and CoT-trace memories degrade performance, favoring compact reasoning abstractions.The ablation compares fixed few-shot demonstrations, truncated/full CoT traces, and abstract reasoning memories.
  • Scalability to reasoning models: CoD+Memory gains +6.67, +9.17, and +4.17 points over CoD on AIME 2024 for DeepSeek-V3.2, Qwen3.5-plus, and o4-mini, respectively.The gains are reported as pass@8 and show generalization beyond a specific open-weight backbone.
  • Compression dependence and memory size: The memory benefit is largest under aggressive compression, diminishes as the baseline recovers, and can decline at k=20 as prefill cost rises.This pattern supports compensation for information lost during compression while revealing a coverage–cost trade-off and over-injection risk.

5 Related Work

Prior work improves inference through context augmentation, context compression, reasoning-memory reuse, concise generation, and decoding-efficiency techniques. MAC differs by retrieving reasoning scaffolds as a pluggable prefill-side compensation layer for content omitted during CoT compression.

  • Prior methods augment or optimize input context through retrieved demonstrations, rationale-based examples, long-context approaches, and context compression.
  • Reasoning-memory methods reuse prior reasoning, whereas MAC supplies retrieved reasoning as a pluggable compensation layer for intermediate content omitted by CoT compression.MAC decouples the memory, retriever, and compressor, shifting part of reasoning from decode-time generation to prefill context.
  • CoT-shortening methods prune or skip tokens, allocate reasoning budgets, self-compress, or use reinforcement learning, but aggressive compression can remove information needed for correct reasoning.MAC instead injects retrieved reasoning scaffolds to support compressed reasoning.
  • Speculative decoding, model sparsification, and KV-cache compression make outputs cheaper or faster, whereas MAC reduces explicit decode-time reasoning through prefill-context support.The two mechanisms are orthogonal and can in principle be combined.

6 Conclusion

The conclusion presents Memory-Augmented Compression as a training-free framework that retrieves reusable reasoning memories as prefill-side scaffolds. It compensates for information lost under compressed decoding while improving accuracy over compressed baselines and reducing latency relative to standard CoT.

  • Memory-Augmented Compression retrieves reusable reasoning memories as prefill-side scaffolds to improve compressed reasoning without additional training.The method shifts part of reasoning support from decode-time generation to the input context.
  • Experiments across diverse reasoning tasks show higher accuracy than compressed baselines such as CoD and lower latency than standard CoT.

Limitations · Appendix

The paper identifies limitations in latency measurement and MAC’s runtime benefits. API interfaces hinder direct comparisons, while MAC’s added processing overhead may reduce its advantage on short reasoning queries.

  • Limitations: API-based models lack a consistent interface for measuring internal prefill and decode latency.This limits direct runtime comparisons across models.
  • Limitations: Direct runtime comparisons across API-based models are therefore limited.The limitation arises from inconsistent access to model-internal latency measurements.
  • Limitations: MAC adds overhead through query-tag generation.This overhead is part of MAC’s additional processing cost.
  • Limitations: MAC also incurs embedding, retrieval, and memory-prefilling overhead.These operations contribute to the framework’s additional latency.
  • Limitations: MAC’s latency advantage may be smaller for queries requiring only short reasoning traces.The added overhead is less easily offset when little reasoning is needed.
  • Limitations: The reported limitations concern both cross-model latency evaluation and the practical runtime cost of MAC.The passage highlights measurement constraints alongside framework-induced overhead.

A Method and Implementation Details … B Experimental Configuration

MAC is a training-free, modular system that builds a fixed memory bank offline and retrieves tagged reasoning memories online for prefill-side augmentation during compressed inference. Its configuration specifies memory formats, disjoint historical sources, tag-based retrieval, prompt integration, and evaluation procedures.

  • A.1 Overall Pipeline: MAC has offline memory-bank construction and online memory-augmented compressed inference stages.The backbone model remains unchanged, while memory representation, retrieval, and compression components are independently replaceable.
  • A.2 Memory Representation: Each memory record stores the problem, answer, verified trace, memory content, structured tags, and tag embedding; inference injects only selected memory content.Supported formats are Long-CoT, Short-CoT, and Summary memory.
  • A.3 Offline Memory-Bank Construction: Memory banks use historical examples disjoint from evaluation sets and remain fixed, excluding evaluation answers and reasoning traces from construction or modification.Sources may include official training splits, prompt exemplars, or historical competition problems.
  • A.3 Offline Memory-Bank Construction: Offline, MAC validates historical reasoning traces, formats retained examples, generates tags, and computes embeddings for a reusable memory bank.Only examples with correct final predictions and complete solutions are retained.
  • A.4 Tag Generation and Query Representation: An LLM generates structured tags covering knowledge domains, solution strategies, problem type and context, and complexity for memories and input queries.Query tags use only the current question; gold answers and reasoning traces are withheld.
  • A.5 Retrieval and Memory Injection: MAC encodes query and memory tags, ranks entries by cosine similarity, selects top-k memories, and injects serialized results into model prefill.The model then generates a compressed reasoning trace and answer under the selected compression method without parameter updates.
  • A.6 Prompt Templates: Prompt templates insert retrieved memories before the unchanged current question as auxiliary reasoning scaffolds, while specifying concise or minimum-draft reasoning and boxed final answers.The memory block supports independent problem solving rather than replacing it.
  • B Experimental Configuration: The experimental configuration covers datasets and memory sources, models and compression baselines, hardware and software, decoding and retrieval settings, answer evaluation, and latency measurement.These components define the evaluation protocol for MAC and its baselines.

B.1 Datasets, Splits, and Memory Banks … B.6 Latency and Cost Measurement

The evaluation spans diverse reasoning benchmarks and model families, using disjoint memory and evaluation data with fixed paired-comparison settings. Accuracy, token usage, latency, and API cost are measured under standardized protocols that separate reusable offline work from per-query processing.

  • B.1 Datasets, Splits, and Memory Banks: Experiments cover GSM8K, MATH-500, BBH, MMLU Science, and AIME 2024, with memory sources disjoint from corresponding evaluation sets.Memory banks are filtered for correctness, completeness, and format where applicable.
  • B.2 Models and Baselines: The study evaluates instruction-tuned Qwen2.5-7B, LLaMA-3.1-8B, Qwen2.5-72B, DeepSeek-V3.2, Qwen3.5-plus, and o4-mini without parameter updates.Baselines include standard CoT, CoD, TokenSkip, RPC, and Extra-CoT; memory variants add retrieved reasoning memories while preserving compression instructions.
  • B.3 Hardware and Software Environment: Open-weight inference uses eight NVIDIA H20 GPUs with fixed tensor parallelism across methods, while memory banks, retrieval indexes, and tag embeddings are built before evaluation.No model parameters are updated during memory construction, retrieval, or inference.
  • B.4 Decoding and Hyperparameter Settings: Paired compressed and memory-augmented methods use the same maximum output length, so memory does not provide an additional decode-token budget.Main open-weight experiments use deterministic greedy decoding with temperature T = 0 and ptop = 1.0, with sampling disabled.
  • B.5 Evaluation Protocol: Evaluation uses exact-match accuracy with benchmark-specific answer extraction, and malformed, empty, or unparsable outputs count as incorrect.Reported efficiency measures include average prefill tokens, average decode tokens, total processed tokens, model latency, and end-to-end latency when available.
  • B.6 Latency and Cost Measurement: Model latency separates input-context processing from autoregressive generation, while MAC end-to-end latency additionally includes online query processing, retrieval, serialization, and prompt construction.Offline memory construction, memory-side tag generation, and memory-index construction are excluded because they are performed once and reused.
  • B.6 Latency and Cost Measurement: Latency comparisons hold model, hardware, batch size, tensor-parallel configuration, and serving process constant, reporting batch size one by default and batch-size-eight measurements additionally.Prefill–decode cost analysis averages per-token latencies over the complete evaluation set.
  • B.6 Latency and Cost Measurement: API usage cost uses provider-reported input and output tokens and applicable prices, counting memory tokens as input and generated reasoning and final-answer tokens as output.API experiments do not report separately measured prefill and decode latency because provider-side execution details are unavailable.

C Additional Experimental Results … C.3 End-to-End Latency Breakdown

Additional experiments show that memory-augmented compression scales to larger models, improves TokenSkip more reliably than Extra-CoT, and preserves an accuracy–latency advantage despite online retrieval costs. End-to-end analysis identifies query-side semantic processing, rather than vector search, as the main retrieval overhead.

  • C Additional Experimental Results: The complete results evaluate multiple reasoning domains and analyze plug-in compatibility, memory format, retrieval size, compression strength, and decoding robustness.The reported plug-in variants use top-k = 5 Long-CoT memories, with Base denoting the original compressor and +Mem. its memory-augmented counterpart.
  • C.1 Model Scaling: On MATH-500 with Qwen2.5-72B, CoD+Memory recovers accuracy to 78.20% and remains 1.62× faster than full CoT.CoD reduces decode tokens from 597.4 to 122.9 and latency from 22592.1 ms to 4393.3 ms, but loses 14.8 percentage points in accuracy; CoD+Memory uses 1772.6 prefill and 373.0 decode tokens on average.
  • C.2.1 TokenSkip: Memory consistently improves TokenSkip across most compression ratios and both backbone models, with larger gains under strong or intermediate compression.On GSM8K, gains occur across all reported ratios and backbones; on MATH-500, gains are consistent except in the most aggressive LLaMA-3.1-8B setting at γ = 0.1.
  • C.2.2 Extra-CoT: Memory improves Extra-CoT only on MATH at the two highest reported ratios, revealing a compressor- and strength-dependent boundary.The remaining configurations do not improve, so retrieved reasoning is not uniformly beneficial across compressors and compression strengths.
  • C.3 End-to-End Latency Breakdown: Latency analysis separates post-retrieval model inference from end-to-end latency, which additionally includes online memory retrieval operations.Post-retrieval inference consists of model prefill and autoregressive decoding; retrieval comprises query-tag generation, query embedding, and top-k vector search.
  • C.3 End-to-End Latency Breakdown: On MATH, additional compression-ratio curves compare no-memory baselines with Long-CoT and Short-CoT memories across retrieval sizes k ∈{3, 5, 10}.The comparisons vary retrieval size and memory format under different compression ratios.
  • C.3 End-to-End Latency Breakdown: 683.0 ms per query is spent on query-tag generation and embedding, whereas top-k vector search requires only 0.10 ms.Vector search contributes negligible overhead, while query-side semantic processing dominates retrieval cost.
  • C.3 End-to-End Latency Breakdown: 3198.5 ms end-to-end latency for CoD+Memory versus 3742.5 ms for standard CoT yields a 1.17× end-to-end speedup.Although query-tag generation and embedding add overhead, MAC remains faster overall; optimizing query-side representation is more important than optimizing vector search.

C.4 Compression Ratio Analysis … C.7 Decoding Robustness

Across compression, memory format, memory size, and decoding temperature, the appendices show that memory benefits depend on compression strength, retrieval breadth, memory format, and decoding conditions. Retrieved memories improve coverage but can add noise, while memory augmentation remains comparatively stable across temperatures.

  • C.4 Compression Ratio Analysis: On MATH, memory gains often peak under intermediate rather than strongest compression across Qwen2.5-7B and LLaMA-3.1-8B.Figure 6 compares Long-CoT and Short-CoT memories with retrieval sizes k ∈ {3, 5, 10}.
  • C.4 Compression Ratio Analysis: Overly aggressive compression may remove the internal reasoning capacity needed to interpret and apply retrieved scaffolds on harder MATH problems.Retrieved memories therefore do not fully replace model-side reasoning.
  • C.5 Memory Format Analysis: Under TokenSkip compression on GSM8K with top-k = 5 retrieval, complete per-γ results report accuracy gains for each compression ratio.The appendix expands aggregated strong, medium, and mild compression results into individual ratios.
  • C.5 Memory Format Analysis: Memory gains are generally largest under stronger compression, while Long-CoT memories remain more stable than Short-CoT memories as compression becomes milder.These complete memory-format results follow the main-paper trend.
  • C.6 Memory Size Analysis: Increasing retrieved-memory count k substantially lengthens prefill while leaving decode length comparatively stable on GSM8K and MATH.Table 13 reports accuracy, average prefill tokens, and average decode tokens across k values on Qwen2.5-7B.
  • C.6 Memory Size Analysis: Accuracy initially improves with broader retrieved reasoning coverage, then saturates or decreases when irrelevant or redundant memories introduce noise.This forms the reported coverage–noise trade-off.
  • C.7 Decoding Robustness: Across temperatures T ∈ [0, 1], memory augmentation shows smaller variation in accuracy and decode length than full CoT.Table 15 reports accuracy and average decode tokens for both methods.
  • C.7 Decoding Robustness: The temperature analysis measures sensitivity across settings and is not a repeated-seed estimate of statistical uncertainty.This qualification limits how the observed robustness should be interpreted.

E Data and Code Availability

The work uses publicly available academic benchmarks under their respective licenses and terms of use. Baselines and MAC resources are provided through official implementations, supplementary materials, and a planned MIT-licensed release.

  • Datasets: The evaluation uses GSM8K, MATH, BBH, MMLU Science, and AIME 2024 solely for research purposes.All listed datasets are publicly available academic benchmarks.
  • Data Use: The authors follow the respective licenses and terms of use for the datasets.
  • Code: Official implementations are used for TokenSkip, RPC, and Extra-CoT when available, while MAC resources are included in the supplement and planned for MIT-licensed release.The supplementary package contains the MAC implementation, experimental configurations, and evaluation utilities.
Loading 2608.21265v1…