Memory Augmentation Unlocks Efficient Chain-of-Thought Reasoning

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

cs.CLAug 21, 2026

Memory Augmentation Unlocks Efficient Chain-of-Thought Reasoning

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

Organizations: Institute of Information Engineering, Chinese Academy of Sciences · School of Cyber Security, University of Chinese Academy of Sciences · Tencent Inc. · Baidu Inc.

Abstract

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 Context-Generation Substitution Law, where explicit reasoning context substitutes for part of decode-time generation. Based on this principle, we propose 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.49x latency speedup latency speedup over standard CoT. Memory is also compatible with token-level, reasoning-trace-level, and inference-state compression mechanisms.

Figures & tables

Figure 1: Illustration of memory-augmented compression. Explicit memory shifts part of the reasoning support to the prefill stage, improving short-CoT accuracy while preserving generation efficiency.

Figure 2: Overview of Memory-Augmented Compression . Historical reasoning traces are distilled into reusable memory entries offline. At inference time, relevant memories are retrieved and injected into the prefill context to support compressed Short-CoT reasoning.

Figure 3: Explicit memory mitigates the reasoning collapse caused by aggressive compression.

(a) Across-domain CoD results with Qwen2.5-7B
Domain	Dataset	Method	Acc.	$\Delta$ vs CoD	Prefill	Decode	Total	Latency (ms)
Math	GSM8K	CoT	91.4	–	96.3	299.2	395.5	3742.5 (1.00 $\times$ )
		CoD	67.9	–	112.3	51.1	163.4	666.6 (5.61 $\times$ )
		+Memory	89.3	+21.4	1161.9	187.2	1349.1	2515.4 (1.49 $\times$ )
	MATH	CoT	72.8	–	105.9	569.3	675.2	5117.0 (1.00 $\times$ )
		CoD	43.0	–	121.9	100.3	222.2	1119.5 (4.57 $\times$ )

Table 1: Main results of Memory-Augmented Compression . Panel (a) reports across-domain CoD results with Qwen2.5-7B. Panel (b) evaluates Memory as a plug-in for representative compression settings.

Retriever	Accuracy (%)	Overhead
No Memory (CoD)	43.00	–
Random	50.20	0.1 ms
Query Emb.	53.20	100 ms
BM25	54.60	1 ms
Reasoning Tag	56.20	683 ms

Table 2: Accuracy and online retrieval overhead of different retrieval strategies on MATH-500. Overhead is the approximate per-query online cost.

Memory Content	Acc.	Prefill	Decode	$\Delta$ Acc
CoD (None)	55.00	136.31	2105.50	0.00
Few-shot	46.67	300.79	1406.11	-8.33
Summary	61.67	606.81	2117.04	+6.67
Short-CoT	46.67	605.40	1454.78	-8.33
Long-CoT	43.33	961.37	1531.02	-11.67

Table 3: Impact of memory representation on DeepSeek-V3.2/AIME 2024 with top-

k=1

retrieval and 2048 maximum generation tokens.

	$k{=}1$		Best			$k{=}20$
Dataset	Acc.	Pre.	$k$	Acc.	Pre.	Acc.	Pre.
GSM8K	73.39	259	14	80.52	1881	79.68	2628
MATH	52.20	349	16	62.60	3330	57.40	4120

Table 4: Effect of memory size

k

on Qwen2.5-7B. We report accuracy (Acc.) and prefill tokens (Pre.). Accuracy improves from

k{=}1

to the best

k

, whereas

k{=}20

may reduce accuracy and increases prefill cost.

DeepSeek-V3.2
Method	Acc. (%)	Prefill	Decode	Total	$\Delta$ Acc
CoT	55.00	124.01	2733.94	2857.95	–
CoD	55.00	136.31	2105.50	2241.81	0.00
+Memory	61.67	606.81	2117.04	2723.85	+6.67
Qwen3.5-plus
CoT	75.83	142.30	8657.60	8799.90	–

Table 5: Scalability to API-based reasoning models on AIME 2024. +Memory denotes CoD+Memory. All results are reported as pass@8, following common practice on AIME. Gains are in percentage points over CoD.

Dataset	Batch	Prefill	Decode	Ratio
		(ms/tok)	(ms/tok)	$\tau_{\text{dec}}/\tau_{\text{pre}}$
GSM8K	1	0.136	11.81	87.0
GSM8K	8	0.144	2.11	14.7
MATH-500	1	0.150	11.97	79.4
MATH-500	8	0.134	2.96	22.1

Table 6: Per-token prefill and decode latency for Qwen2.5-7B on an NVIDIA H20 under batch sizes.

Figure 4: Memory compensation across compression ratios on GSM8K with Qwen2.5-7B. Smaller $\gamma$ indicates more aggressive compression. Green: memory-augmented ( $k{=}5$ ); gray: no-memory; shaded: gain.

Figure 5: Case study on GSM8K where retrieved memory guides compressed reasoning to recover the missing subtraction step.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Evaluation Split	Eval. Size	Memory Source	Bank Size
GSM8K	Official test set	1,319	Filtered official training examples	4,307
MATH	MATH-500	500	Filtered official training examples	2,242
BBH	Nine selected tasks	495	Official task-specific CoT exemplars	27
MMLU-Sci	Subject-stratified 20% split	79	Remaining 80% from the same six subjects	317
AIME 2024	Official 2024 problems	30	Historical AIME problems from 1983–2022	287

Appendix

Table 7: Dataset splits and memory-bank statistics. Bank Size denotes the number of retained memories after correctness, completeness, and format filtering. Evaluation examples are excluded from all memory banks and retrieval pools.

Method	Acc.	Prefill	Decode	Latency (ms)
CoT	77.00	105.9	597.4	22592.1
CoD	62.20	121.9	122.9	4393.3
CoD+Memory	78.20	1772.6	373.0	13940.4

Appendix

Table 8: Results on MATH-500 with Qwen2.5-72B. Prefill and Decode denote the average numbers of input and generated tokens, respectively. CoD+Memory achieves a

1.62\times

speedup over full CoT while improving accuracy from 77.00% to 78.20%.

Model	$\gamma$	Base	+Mem.	Gain
Qwen2.5-7B	0.1	52.62	74.60	+21.99
	0.2	60.65	74.53	+13.87
	0.3	66.19	78.77	+12.59
	0.4	68.08	80.74	+12.66
	0.5	71.11	84.00	+12.89
	0.6	77.26	84.99	+7.73

Appendix

Table 9: Full TokenSkip plug-in results on GSM8K with top-

k=5

Long-CoT memories.

Model	$\gamma$	Base	+Mem.	Gain
Qwen2.5-7B	0.1	12.00	13.20	+1.20
	0.2	19.60	22.20	+2.60
	0.3	28.00	33.60	+5.60
	0.4	34.40	44.80	+10.40
	0.5	38.00	52.80	+14.80
	0.6	42.60	58.40	+15.80

Appendix

Table 10: Full TokenSkip plug-in results on MATH-500 with top-

k=5

Long-CoT memories.

Dataset	Ratio	$k$	Base	+Mem.
GSM8K	0.2	5	74.07	67.70
	0.4	5	78.92	71.65
	0.6	5	84.69	80.74
	0.8	5	85.75	80.52
	1.0	5	83.70	81.80
MATH	0.2	5	29.20	24.40

Appendix

Table 11: Full Extra-CoT plug-in results with Qwen2.5-7B and top-

k=5

Long-CoT memories.

Method	Acc. (%)	Prefill + Decode (ms)	Query Tag + Embedding (ms)	Vector Search (ms)	End-to-End Latency (ms)	E2E Speedup
CoT	91.4	3742.5	–	–	3742.5	1.00 $\times$
CoD+Memory	89.3	2515.4	683.0	0.10	3198.5	1.17 $\times$

Appendix

Table 12: End-to-end latency breakdown on GSM8K. Prefill + Decode denotes post-retrieval model inference latency. For CoD+Memory, end-to-end latency additionally includes query-tag generation, query embedding, and top-

k

vector search.

Panel of Figure 6 — Figure 6: Additional compression-ratio curves on MATH. We compare no-memory baselines with Memory using Long-CoT and Short-CoT formats under different retrieval sizes $k\in\{3,5,10\}$ .

GSM8K				MATH				MMLU-Sci
$k$	Acc.	Prefill	Decode	$k$	Acc.	Prefill	Decode	$k$	Acc.	Prefill	Decode
1	73.39	259.0	57.9	1	52.20	348.8	152.8	1	55.07	267.1	51.6
3	76.27	509.4	59.9	3	53.20	751.0	190.6	3	65.64	493.5	68.3
5	77.26	759.8	61.3	5	53.40	1134.4	202.6	5	62.11	493.6	66.5
8	78.15	1132.1	61.7	8	58.60	1729.8	218.1	8	66.52	493.6	67.7
10	78.47	1379.2	61.9	10	58.80	2124.9	227.1	10	65.20	493.6	67.8

Appendix

Table 13: Full memory-size results with Qwen2.5-7B. We report accuracy and average prefill and decode tokens under different values of

k

$\gamma$	LLaMA-3.1-8B		Qwen2.5-7B
$\gamma$	Long	Short	Long	Short
0.1	+30.71	+31.54	+33.66	+29.04
0.2	+14.25	+13.57	+17.97	+16.15
0.3	+7.28	+7.05	+11.75	+8.49
0.4	+6.44	+4.85	+4.40	+1.67
0.5	+4.70	+3.18	+3.56	+1.06

Appendix

Table 14: Full memory-format results under TokenSkip compression on GSM8K. Values denote accuracy gains over the corresponding no-memory compressed baseline in percentage points. Both Long-CoT and Short-CoT memories use top-

k=5

retrieval.

Temperature	Standard CoT		Memory
$T$	Acc.	Decode	Acc.	Decode
0.0	90.1	282	81.8	68.2
0.2	90.8	288	82.1	68.7
0.4	91.2	293	82.3	69.0
0.6	89.8	297	81.9	68.5
0.8	90.5	301	82.0	68.9

Appendix

Table 15: Full temperature results for decoding robustness. Accuracy is reported in percentage points, and Decode denotes the average number of generated tokens. Across-

T

Std. denotes the standard deviation across the six temperature settings rather than across repeated random runs.

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 356 papers

HybridThinker: Efficient Chain-of-Thought Reasoning via Compressed Memory and Transient Thought Steps
Jun 2, 2026Xin Liu, Runsong Zhao, Xinyu Liu +8Memory-Augmented Language Models Efficient Language Model Reasoning
School of Computer Science and Engineering, Northeastern University, Shenyang 110819, China · China Unicom Cloud-Link, China · NiuTrans Research, Shenyang, China
Think Less, Know More: State-Aware Reasoning Compression with Knowledge Guidance for Efficient Reasoning
Apr 10, 2026Yi Sui, Chaozhuo Li, Dawei SongCoT Compression Efficient Language Model Reasoning
Beijing Institute of Technology, China · Beijing University of Posts and Telecommunications, China · The Open University, UK
Selective Latent Thinking: Adaptive Compression of LLM Reasoning Chains
May 25, 2026Hui Xie, Jie Liu, Ziyue Qiao +1Efficient Language Model Reasoning LLM Inference Acceleration
Eindhoven University of Technology, Netherlands · School of Computing and Information Technology, Great Bay University, China

Jun 2, 2026J/K move · Enter open · S save

HybridThinker: Efficient Chain-of-Thought Reasoning via Compressed Memory and Transient Thought Steps

Xin Liu, Runsong Zhao, Xinyu Liu, Junhao Ruan+7

School of Computer Science and Engineering, Northeastern University, Shenyang 110819, China · China Unicom Cloud-Link, China · NiuTrans Research, Shenyang, China

Memory-Augmented Language Models Efficient Language Model Reasoning Efficient Language Model Inference CoT Compression LLM Reasoning

Open paper AskPDF arXiv

Extended chain-of-thought (CoT) traces improve LLM reasoning but incur substantial computational and memory costs. While existing CoT compression methods mitigate this by condensing thought steps into compact representations via memory tokens and retaining only these representations at inference time, the loss of fine-grained information makes subsequent steps more error-prone. To alleviate this, we propose HybridThinker, where in addition to preserved these representations, thought steps are also temporarily retained to provide fine-grained details. However, we observe that naively keeping thought steps accessible to subsequent steps during training lets the model bypass memory tokens by retrieving information directly from these steps, leaving the model's ability to compress and retrieve information through memory tokens insufficiently trained. We therefore introduce a hybrid training scheme, in which only some thought steps are directly accessible through attention to subsequent steps, while the other thought steps are masked, forcing the model to use memory tokens for compression and retrieval. Across 4 reasoning benchmarks, HybridThinker matches the uncompressed baseline, advancing the state of the art in CoT compression by 5.8 points on average accuracy with similar inference time. Ablation studies confirm that both temporary thought-step retention and the hybrid training scheme contribute to these gains.

Abstract

Efficient Language Model Reasoning

Explore similar work

HybridThinker: Efficient Chain-of-Thought Reasoning via Compressed Memory and Transient Thought Steps

Think Less, Know More: State-Aware Reasoning Compression with Knowledge Guidance for Efficient Reasoning

Selective Latent Thinking: Adaptive Compression of LLM Reasoning Chains