Memory Papers
Memory Architecture

Memory Augmentation Unlocks Efficient Chain-of-Thought Reasoning

AuthorsSimeng Zhang, Yilong Chen, Wenyuan Zhang et al.

arXiv 20262026
arXivPDF

TL;DR

Memory-Augmented Compression shifts reasoning from decode-time to prefill-side memories, boosting Chain-of-Draft GSM8K accuracy from 67.9% to 89.3% (+21.4 points).

SharePost on XLinkedIn

Read our summary here, or open the publisher PDF on the next tab.

THE PROBLEM

Compressed CoT Collapses Accuracy Under Aggressive Shortening

Existing Chain-of-Thought compression causes compression collapse, where CoD accuracy on GSM8K drops to 67.9% while standard CoT reaches 91.4%.

This breakdown in complex reasoning increases latency pressure: long CoT traces are too slow, but naive compression discards subgoals and constraints, degrading task performance.

HOW IT WORKS

Memory-Augmented Compression Framework

Memory-Augmented Compression introduces a memory bank, memory retrieval, memory-augmented prefill, and memory-guided compressed inference to shift reasoning support into explicit context.

Think of the memory bank as a reasoning card catalog, where distilled solution patterns act like indexed templates that the LLM can look up before thinking.

This design lets Memory-Augmented Compression replace parts of long working memory with cheaper explicit memory, enabling aggressive CoT compression that a plain context window cannot sustain without accuracy loss.

DIAGRAM

Offline Memory Construction and Online Inference Flow

This diagram shows how Memory-Augmented Compression constructs reasoning memories offline and uses them during online compressed inference.

DIAGRAM

Evaluation Setup Across Datasets and Compression Methods

This diagram shows how Memory-Augmented Compression is evaluated with CoT, CoD, and other compression baselines on multiple benchmarks.

PROCESS

How Memory-Augmented Compression Handles a Reasoning Query

  1. 01

    Memory bank construction

    Memory-Augmented Compression first runs Full CoT Generation on training datasets and applies Pattern Summary Extraction to build the memory bank with reusable reasoning patterns.

  2. 02

    Memory retrieval

    Given a new query, Memory-Augmented Compression uses the Memory Retriever with reasoning tags and semantic similarity to select top k relevant memories from the memory bank.

  3. 03

    Memory-augmented prefill

    Memory-Augmented Compression concatenates retrieved memories with the system prompt and query, forming a memory-augmented prefill context processed in parallel before decoding.

  4. 04

    Memory-guided compressed inference

    During inference, Memory-Augmented Compression applies Chain-of-Draft or other compressors to generate a Short-CoT reasoning trace and final answer, relying on explicit memory scaffolds.

KEY CONTRIBUTIONS

Key Contributions

  • 01

    Context Generation Substitution Law

    Memory-Augmented Compression formalizes the Context Generation Substitution Law with J = |z′| + γ|M| + λLperf, showing how explicit memory can replace decode-time reasoning when γ = τpre/τdec is small.

  • 02

    Memory-Augmented Compression Framework

    Memory-Augmented Compression introduces a training-free memory-prefill framework with a memory bank, memory retrieval, memory-augmented prefill, and memory-guided compressed inference compatible with CoD, TokenSkip, RPC, and Extra-CoT.

  • 03

    Improved Accuracy Latency Trade Off

    Memory-Augmented Compression improves CoD accuracy by 21.4 points on GSM8K, 28.0 on MATH, 29.5 on BBH, and 6.61 on MMLU-Sci, while achieving 1.14–1.49× latency speedup over standard CoT.

RESULTS

By the Numbers

Acc.

89.3%

+21.4 over CoD on GSM8K

Latency (ms)

2515.4 ms

1.49× speedup vs CoT on GSM8K

Acc.

70.5%

+29.5 over CoD on BBH

Acc.

71.0%

+28.0 over CoD on MATH

On GSM8K, MATH, BBH, and MMLU-Sci, Memory-Augmented Compression with CoD consistently raises accuracy while remaining faster than standard CoT. These results show that explicit reasoning memories can compensate for information lost under aggressive compression without sacrificing efficiency.

BENCHMARK

By the Numbers

On GSM8K, MATH, BBH, and MMLU-Sci, Memory-Augmented Compression with CoD consistently raises accuracy while remaining faster than standard CoT. These results show that explicit reasoning memories can compensate for information lost under aggressive compression without sacrificing efficiency.

BENCHMARK

Across-domain CoD results with Qwen2.5-7B on GSM8K

Accuracy (%) on GSM8K comparing Memory-Augmented Compression with CoT and CoD.

BENCHMARK

Memory as a plug in for compression methods on GSM8K

Accuracy (%) on GSM8K for TokenSkip with and without Memory-Augmented Compression.

KEY INSIGHT

The Counterintuitive Finding

Memory-Augmented Compression sometimes surpasses standard CoT, such as 70.5% accuracy on BBH versus 60.8% for CoT, despite using compressed reasoning.

This breaks the intuition that longer Chain-of-Thought traces are always better, showing that targeted explicit memories can outperform verbose generation with fewer decode tokens.

WHY IT MATTERS

What this unlocks for the field

Memory-Augmented Compression unlocks efficient long reasoning by shifting reusable patterns into explicit memory, reducing dependence on slow autoregressive traces.

Builders can now combine aggressive CoT compression with memory-prefill scaffolds, designing systems that are both highly accurate and latency-aware across complex reasoning benchmarks.

~14 min read← Back to papers

Related papers

Memory Architecture

A Control Architecture for Training-Free Memory Use

Yanzhen Lu, Muchen Jiang et al.

· 2026

TAG routes low-confidence steps to uncertainty-based routing, filters them with guarded acceptance with rollback, chooses between bank selection across rule and exemplar memory, and prunes via evidence-based retirement inside a unified control loop. On SVAMP and ASDiv, TAG reaches 81.0% and 85.2% accuracy, improving over the 74.0% and 77.5% no-memory baselines while a compute-matched Retry baseline stays flat.

arXiv:2604.18206Read explainer

Questions about this paper?

Paper: Memory Augmentation Unlocks Efficient Chain-of-Thought Reasoning

Answers use this explainer on Memory Papers.

Checking…