Reasoning models maintain information about the status of their reasoning process through in-context memory. Complex problems require larger context windows and memory requirements. In this scenario, context management without caching risks discarding useful information for reasoning. We propose RCMEM, an approach that extends context management with an external memory to maintain additional information without burdening the reasoning model with additional context. RCMEM significantly reduces computation time on problems that require repeated computation of intermediate steps, while achieving better performance on complex problems. We demonstrate our approach on a few mathematics problems using model sizes ranging from 0.4 to 25 million parameters, and GPT-2.
Figures & tables
Figure 1: Illustration of how memory access is introduced on top of reduction rule in RCMEM
Figure 2: (a) Average number of tokens and (b) special tokens for the Knapsack problem training sets. Note the number of tokens and [CALL] (iterations) of PENCIL exceed RCMEM at N=7 and N=6 respectively.
Figure 3: (a) Average number of tokens and (b) special tokens for LCS problem training sets. The number of tokens and iterations are much larger on PENCIL.
Figure 5: Accuracy of PENCIL and RCMEM on Knapsack across 3 model sizes.
Figure 6: Average generation time for a sample at inference, on 256 and 512 context lengths.
Figure 7: Knapsack - Accuracy of 0.4M, 3.17M, 25.2M and GPT-2 on 25k, 50k, 75k and 100k training samples.
Figure 8: Accuracy and sample generation time of PENCIL and RCMEM on LCS tasks. The accuracy are broken down to accuracy per batch of training samples to understand how fast the models are trained.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
# of Item
No memory
Memory
Fix READ
Fix READ b.p.
Fix R&W b.p.
Accuracy
4
0.816
0.74
0.86
0.824
0.832
5
0.742
0.546
0.732
0.758
0.702
6
0.584
0.284
0.68
0.7
0.55
Memory Correct Rate
4
0.762
0.891
0.975
0.976
0.967
Appendix
Table 1: Knapsack - Average accuracy and actual memory correct rate on 0.4M model with 512 content length, with different memory handle methods. R&W - READ and WRITE, b.p. - bypass invalid memory access
Model Size 3.17M
Context 512
# items
CoT
PENCIL
RCMEM
4
0.776
0.984
0.988
5
0.354
0.984
0.976
6
0.222
0.976
0.974
7
0.086
0.956
0.956
Appendix
Table 2: Knapsack - Accuracy of 3.17M on CoT, PENCIL and RCMEM, with 512 and 256 context tokens.
PENCIL
RCMEM
Items
Acc
Time
Trace Rate
Acc
Time
Trace Rate
4
0.984
1.124
0.932
0.988
0.896
0.938
5
0.966
2.303
0.839
0.976
1.322
0.833
6
0.938
2.267
0.734
0.974
2.438
0.670
7
0.844
3.296
0.587
0.956
2.980
0.578
8
0.794
4.238
0.550
0.956
5.986
0.462
Appendix
Table 3: Knapsack - Accuracy, time (in seconds) and trace rate of PENCIL and RCMEM, trained on the same RCMEM data.
Table 4: Pseudocode of generating training traces for Knapsack problem