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.
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.
Δ vs CoD
Prefill
Decode
Total
Latency (ms)
Math
GSM8K
CoT
91.4
–
96.3
299.2
395.5
3742.5 (1.00 × )
CoD
67.9
–
112.3
51.1
163.4
666.6 (5.61 × )
+Memory
89.3
+21.4
1161.9
187.2
1349.1
2515.4 (1.49 × )
MATH
CoT
72.8
–
105.9
569.3
675.2
5117.0 (1.00 × )
CoD
43.0
–
121.9
100.3
222.2
1119.5 (4.57 × )
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
Δ 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.
Method
Acc. (%)
Prefill
Decode
Total
Δ Acc
DeepSeek-V3.2
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)
τdec/τ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 γ 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× speedup over full CoT while improving accuracy from 77.00% to 78.20%.
Model
γ
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
γ
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 ×
CoD+Memory
89.3
2515.4
683.0
0.10
3198.5
1.17 ×
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.
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∈{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 .
γ
LLaMA-3.1-8B
Qwen2.5-7B
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.
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.
Xin Liu, Runsong Zhao, Xinyu Liu +8
School of Computer Science and Engineering, Northeastern University, Shenyang 110819, China · China Unicom Cloud-Link, China · NiuTrans Research, Shenyang, China
Large Reasoning Models (LRMs) achieve strong performance on complex tasks by leveraging long Chain-of-Thought (CoT), but often suffer from overthinking, leading to excessive reasoning steps and high inference latency. Existing CoT compression methods struggle to balance accuracy and efficiency, and lack fine-grained, step-level adaptation to redundancy and reasoning bias. Therefore, we propose State-Aware Reasoning Compression with Knowledge Guidance (STACK), a framework that performs step-wise CoT compression by explicitly modeling stage-specific redundancy sources and integrating with a retrieval-augmented guidance. STACK constructs online long-short contrastive samples and dynamically switches between knowledge-guided compression for uncertain or biased reasoning state and self-prompted compression for overly long but confident state, complemented by an answer-convergence-based early stopping mechanism to suppress redundant verification. We further propose a reward-difference-driven training strategy by combining Proximal Policy Optimization (PPO) and Direct Preference Optimization (DPO), enabling models to learn state-conditioned compression strategies. Experiments on three mathematical reasoning benchmarks show that STACK achieves a superior accuracy-efficiency balance, reducing average response length by 59.9% while improving accuracy by 4.8 points over existing methods.
Yi Sui, Chaozhuo Li, Dawei Song
Beijing Institute of Technology, China · Beijing University of Posts and Telecommunications, China · The Open University, UK
Explicit chain-of-thought (CoT) reasoning substantially improves the reasoning ability of large language models (LLMs), but incurs high inference cost due to lengthy autoregressive traces. Existing latent reasoning methods offer a promis- ing alternative, yet they often treat reasoning as uniformly compressible, causing precision-critical intermediate steps to be overly compressed and thereby degrading reasoning accuracy. In this work, we propose Selective Latent Thinking (SLT), a reasoning paradigm that selectively compresses redundant reasoning spans into latent representations while preserving precision-critical spans as explicit CoT within the same reasoning trajectory. Our primary instantiation uses a lightweight decoder to anticipate a short upcoming reasoning span and applies confidence- based gating to determine the longest span that can be reliably compressed; the resulting selective compression policy is further optimized with trajectory-level reinforcement learning to balance answer correctness against reasoning cost. Ex- periments across four mathematical reasoning benchmarks show that which spans are compressed is critical: with the same compressor, compressing arbitrary spans reduces the GSM-Hard accuracy of Qwen3-4B from 54.7% to 17.3%, whereas keeping numerical tokens explicit retains 51.8%. SLT reduces explicit reasoning length by 58.4% with only a 2.8-point accuracy drop compared to explicit CoT on Llama-3.2-1B. Notably, the principle extends beyond this design: instantiated within existing latent reasoning models via a learned compression gate, it improves accuracy at equal reasoning length (+2.6 points over Latent-SFT) and at shorter length (+5.6 points over Latent-GRPO with 13% fewer tokens).
Hui Xie, Jie Liu, Ziyue Qiao +1
Eindhoven University of Technology, Netherlands · School of Computing and Information Technology, Great Bay University, China