Reasoning Under 1 Billion: Memory-Augmented Reinforcement Learning for Large Language Models
Authors: Hung Le
Van Dai Do
Dung Nguyen
Svetha Venkatesh
Venue: Accepted by TMLR ()

Abstract: Recent advances in fine-tuning large language models (LLMs) with reinforcement learning (RL) have shown promising improvements in complex reasoning tasks, particularly when paired with chain-of-thought (CoT) prompting. However, these successes have been largely demonstrated on large-scale models with billions of parameters, where a strong pretraining foundation ensures effective initial exploration. In contrast, RL remains challenging for tiny LLMs with 1 billion parameters or fewer because they lack the necessary pretraining strength to explore effectively, often leading to suboptimal reasoning patterns. This work introduces a novel intrinsic motivation approach, called Memory-R+, that leverages episodic memory to address this challenge, improving tiny LLMs in CoT reasoning tasks. Inspired by human memory-driven learning, our method leverages successful reasoning patterns stored in memory while allowing controlled exploration to generate novel responses. Intrinsic rewards are computed efficiently using a kNN-based episodic memory, allowing the model to discover new reasoning strategies while quickly adapting to effective past solutions. Experiments on three reasoning datasets demonstrate that our approach significantly enhances smaller LLMs' reasoning performance and generalization capability, making RL-based reasoning improvements more accessible in low-resource settings.

Published in Transactions on Machine Learning Research (10/2025)
Reasoning Under 1 Billion: Memory-Augmenter Reinforce- ment Learning for Large Language Models
Hung Le
that. le Bdekkin edu.
Applied Artificial Intelligence Initiative Deakin University, Austrolia
Dai Do
v. v. Adenkin edu.an
Applied Artificial Intelligncee Initiative Deakin University, Austrolia
Dung Nguyen
dung, www.gn.ggnnnnnnueeoom ruguenn &deakin edu.
Applied Artificial Intelligence Initiative Deakin University, Austrolia
Svetha Venkatesh
svetha. www.ateshaeessoo.. &deakin edu.oo
Applied Artificial Intelligence Initiative
Deakin Universily, Austrulia
Reviewed on OpenReview: https: // openrevie.. net/ forum? id= tmdunU2uKs
Abstract
Recent advances in fine-tuning large language models (LLMs) with reinforcement learning (RI) have shown promising improvements complex reasoning tasks, particularly when paired with chain-offthought (CoT) prompting. However, these successes have been largely demonstrated large-scale models with billions of parameters, where strong pretraining foundation ensures effective initial exploration. In contrast, RI. remains challenging for tiny LLMs with 1 billion parameters or fewer because they lack the necessary pretraining strength to explore effectively, often leading to suboptimal reasoning patterns. This work introduces a novel intrinsic motivation approach, called Memory-Rt, that leverages episodic memory to address this challenge, improving tiny LLMs in CoT reasoning tasks. Inspired by human memory-driven learning, our method leverages successful reasoning patterns stored in memory while allowing controlled exploration generate novel responses. Intrinsic rewards computed efficiently using kNN- based episodic memory, allowing the model to discover new reasoning strategics while quickly adapting to effective past solutions. Experiments on three reasoning datasets demonstrate that our approach significantly enhances smaller LLMs reasoning performance and generalization capability, making RL-based reasoning improvements more accessible in low-resource settings.
1 Introduction
Large Language Models (LLMs) have demonstrated remarkable advancements in reasoning and problem- solving, driven by innovations in scaling strategies and training techniques (Google 2024 OpenAl 2024a
Despite its foundational role in defining LLM capability, scaling pro-training is prohibitively expensive and tends to plateau (Xia et. al.) 2023 Hong et. al., 2023). As a result, post-training has become increasingly important, offering improvements in alignment, reasoning depth, and downstream task efficiency (Kumar et al.) ) Among post-training approaches, reinforcement learning (RL) fine-tuning
promising alternative expensive LLM's test-time search methods, such MCTS or Beam Search (Yao et al. 2023;
Feng et. al. 2023 Snell al. 2024h. RT. directly instills chain-of-thought (CoT) reasoning strategies into the

---

Published in Transactions on Machine Learning Research (10/2025)
model, enabling efficient deployment. Recent works like DeepSeek-R1 [uuo et al. 2025 show that RL with
simple outcome rewards can enhance reasoning without relying on heavy inference-time compute methods
(OpenAl 2024a b) and complicated process-based rewards (Lightman et al.) 2023 Zhang et al.).
However, these benefits have been observed mainly in large models (8B-670B) (Guo et al. 2025 Arora &
Zanette 2025 Yeo et al. 2025). In contrast, RT. remains challenging for tiny LLMss which we consider as having S IB porameters. These weak models frequently produce incorrect outputs during training, failing to receive any outcome reward. For example, 0.5B model may repeatedly generate improperly formatted answers to math questions, failing to produce any valid outputs that qualify for a reward. As a result, reward signals are extremely sperse. A common mitigation approach is using hcuristic format-based rewards Guo et al. 2025). However, we will indicate that relying on format-based rewards can cause training collapse in tiny LLMs, as they may overfit to simple format patterns while neglecting the main task. Worse, exploration is ineffective only because small models choose poor actions but also because they lack an explicit exploration mechanism. Unlike RL agents, LLMs do not actively explore or exploit; they passively sample from learned distributions. As noted in KKrshhamurtty et. al. (2024), even large models struggle with exploration and exploitation; this issue becomes acute for small models. Finally, the lack of quality data in downstream tasks poses an additional challenge for training tiny LLMs with RL.
Drawing inspiration from the human brain's episodic memory, which stores and retrieves experiences to guide learning (McClelland et al., [995), , we introduce Memory-R a memory-augmented reinforcement learning framework designed to enhance CoT reasoning in tiny LLMs. To address the challenges of reward sparsity and insufficient exploration, we implement an intrinsic motivation mechanism that emulates the brain's drive to seek successful outcomes (exploit) and avoid repeated errors (explore). This mechanism guides reasoning trajectories by leveraging two distinct episodic memory modules: one dedicated to storing successful reasoning traces and the other to capturing failed attempts. By employing nearest-neighbor estimation within a shared representation space, Memory-R+ derives performance-driver intrinsic rewards from the momory. This process mirrors how humans Icarn from near-correct attempts, allowing LLMs to refine their reasoning by aliguing with successful patterns while avoiding detrimental exploration paths. For example, when asked "There are bags with apples each, how many apples are there?" an LLM might rospond, "Each bag has 4 apples. Three bags mean 3 x 4 = 12. So the answer is 11." Memory-R* still assigns a positive reward for the correct multiplication step when it closely matches past response that reached the correct. final answer. In contrast, if the model repeatedly answers "Each bag has 4 and has 3 more. 3 + 4 it triggers negative rewards. This intrinsic motivation effectively addresses the limitations of sparse external rewards, providing a continuous learning signal based on past experiences.
Unlike traditional episodic control methods that rely on state-action-retum associations for discrete action spaces (Pritzel et al., 2017: Le et al., 2021 2022; Do et al., 2024), Memory-R* simplifies memory storage to input-output pairs, making it more suitable for LLMs' textual reasoning. Upon receiving new query, the framework retrieves outputs from similar past instances by first encoding the query and searching for the top-k most similar queries in memory using cosine similarity. The corresponding response sets from the success and failure memories are then retrieved. The exploitation reward is computed by measuring the Euclidean distance between the generated response and the centroid of the successful response set, encouraging the model to align with generalizable successful patterns rather than memorizing specific past responses. In contrast. the exploration reward is derived from the maximum cosine similarity betweon the generated response and the stored failure responses, ensuring that the model discovers novel outputs differing from incorrect reasoning. To maintain stability in training, both rewards are normalized within sliding window, adapting to the model's recent performance trends.
To evaluate our approach, we conduct extensive experiments on mathematical problem-solving across several tiny LLMs. Our results demonstrate that Memory-R+ significantly improves reasoning accuracy and robust- ness compared to baseline RL and other handcrafted rewards. Moreover, analytic studies provide insights into training collapses and the impact of different memory configurations, highlighting the role of episodic memory in enhancing reasoning performance. In summary, our key contributions are as follows: (1) We pi- oneer an RL fine-tuning approach for tiny LLMs by leveraging memory-basel intrinsic reward mechanism (2) We empirically identify and analyze training collapse issues when fine-tuning tiny LLMs with RL. (3) Extensive experiments on CoT reasoning tasks show that Memory-R*
to teach LLMs to explore and exploit.
2

---

Published in Transactions on Machine Learning Research (10/2025)
Mj
ghbbr
Lookup Memory Retrivaal
Memory- based Intrinsic Reward
Joy pages minutes. How long will it take her to 120 pages?
Rown
22
0g.1
22,2
R3,1
43,2
AS.1
6g.2
R4,1
Novelty (Explore)
Gatasos D
LLM
Similarity (Explott)
M,
First, let's find Joy's reading rate: reads pages in 20 That
2:
A6,1
4G,2
Answer Verifier
R
91
d4.1
4kg,2
4r,1
Figure 1: Memory-Rt Architecture. Left: The LLM receives query from training dataset D, and generates multiple responses. For each response in addition to outcome reward R from an Answer Verifier, Memory-R+ introduces intrinsic reward Rmem based on episodic memory. Right: The query is used to query the failure memory M and success memory Ms using kNN (red arrows), resulting in corresponding retrieved responses. The intrinsic reward Rmem is computed by comparing the current response a to retrieved ones encouraging novelty against failed responses (e.g., 01,1, 03,1, 03,2) and rewarding similarity successful ones (e.g., 05,1, 05,2; 06,1, a6
outperforms other RL methods, enhancing reasoning in small models significantly. (4) We conduct extensive analyses of our method, including hyperparameter selection, intrinsic rewards, exploration behaviors, and the emergence of self-reflection.
The significance of our approach is that our method enables effective RL fine-tuning for models as small as 500M parameters-orders. reasoning research (Guo et al. (025).. This dramatically lowers the barrier to entry for small research labs, academic groups, and companies with limited computing resources, making advanced reasoning capabilities more accessible.
magnitude smaller than current state-of-the-art LLMs used in RL-based
2 Method
2.1
Intrinsic Reward Formulation for CoT Reasoning
When the LLM generates a response to a given qucry, receives two forms of foedback: an outcome reward R from an Answer Verifier that judges the correctness of the final answer extracted from the response, and an intrinsic reward Rmem from memory that reflects how the response aligns with past successes and failures. We note that the Answer Verifier can only assess the final answer and cannot evaluate the quality of the reasoning chains in the response. Therefore, the intrinsic reward is expected to complement the Answer Verifier in providing useful training signals.
Our intrinsic reward balances exploration and exploitation by rewarding responses that resemble past suc- cessful reasoning trajectories while penalizing those similar to previously failed responses. This is achieved through a kNsebased memory system that quantifies the novelty and similarity of generated responses, Fig, illustrates the overall design Mcmory-R+.
2.2 Episodic Memory
Memory Formulation We construct an episodic memory module M to store past reasoning trajectories, facilitating efficient retrieval of relevant experiences. To ensure efficient. reasoning retrieval, both queries and
3

---

Published in Transactions on Machine Learning Research (10/2025)
responses encoded into shared high-dimensional vector space using Enc, implemented as pre-trained Sentence Transformer (Reimers & Gurevych 2019):
q = Enc(g.) E R. a.j = Enc(a.,) E R
(1)
Here, each entry in the memory consists of embeddings of query 9i and set of associated responses {atz): M = {qg (m)ي) where N is the maximum number of stored queries, and each query qi maintains at most L associated responses. Memory retrieval is denoted as: M[a. = {ar.)/=1
Memory Writing During training, we sample G responses from the LLM for given query. New query- response pairs are incorporated into memory following update rule: (1) If q is a novel query (i.e. not present in M) and [M] 2 N, the oldest stored query-response pair is evicted to maintain a fixed memory capacity. The new query and its corresponding responses are then inserted.
(2) If already exists in M, the new response set {ay} is merged with the existing responses. If the total mumber of responses exceeds L, the oldest responses are discarded to preserve memory constraints.
For guiding reinforcement learning, we maintain two episodic memory modules: one for storing successful responses, M and another for failed responses, M Given a query g and set of generated responses {a, {z update the memories follows:
+ [a u (az) R(g, 0z) > r*}, ( + - u R(q, (5),a))s+))
(2)
where 7" and 7 are reward thresholds for classifying successful and failed responses, respectively. For instance, in mathematical problem-solving, where the outcome reward is defined as R(g, a) = 1 for correct final answer and 0 otherwise, can set the thresholds as 7" 7 0.5.
2.3 Memory-based Intrinsic Reward
Memory Read
Given new query q and response a, compute their embeddings q Enc(g).
Enc(a) and retrieve the top-k nearest queries from a memory M based cosine similarity (CS): 1 k=1 = top-k (arg M (5)q. Then, the set of relevant responses from the memory is computed
as:
K
B(M, U M
(3)
k=1
where K is the number of nearest neighbors considered the memory retrieval. For simplicity, the same K is used for both M and M
Exploitation Reward Rexploit To reinforce successful reasoning patterns, compute the exploitation reward using responses stored in the success memory M The model is rewarded for generating responses similar to those that previously led to correct final answers. We first compute the centroid of retricved re- sponse embeddings from the success memory: c(M*, q) = نا EajeB(Mr.s)) aj. The Euclidean distance between the response and this controid determines the exploit reward as:
Rexplont a) = -|a c(M" .))l
(4)
By mcasuring the distance to the centroid, encourage the model to align with the general distribution of successful reasoning patterns rather than overfitting specific past answers. This provides smoother optimization signal, overcoming the reward sparsity problem and capturing structural commonalities in effective reasoning paths.
Exploration Reward Rexplore To encourage novel reasoning paths, we compute the exploration reward using responses stored in the failure memory M ensuring that the model avoids repeating past mistakes. Specifically, novelty is measured as the minus of the cosine similarity between the generated response em- bedding a and its closest retrieved embedding from the failure memory:
4

---

Published in Transactions on Machine Learning Research (10/2025)
Rexplore(q. a) =1
CS(a,
(5)
a.EB(M).
This formulation penalizes responses that closely resemble previously failed attempts while rewarding novel outputs that deviate from incorrect reasoning, Importantly, the design of this intrinsic reward creates a natural curriculum: carly in training, when most outputs are wrong, this encourages broad exploration, generaling anything unlike previous allempts. As correct responses accumulate the success memory M', the failure memory becomes more selective, guiding the model to avoid bad patterns without discouraging similarity to correct ones. As the model improves, this creates natural progression from broad to focused exploration. In practice, can set warm-up period (e.g., 50 training steps) to collect the initial data before applying the exploration reward.
Reward Normalization The intrinsic rewards are normalized using a running min-man max scaling to ensure they are evaluated relative to recent performance trends. For example, given a sliding window of past intrinsic rewards (Rexplore iim where w is the window size, the normalized intrinsic reward is computed as:
dexplolt/exploree a) min, usist exxploot/exxlore.
a) mäxe-wsist Rexploit/explore. - ming-wgist /expoot/ +
(6)
where E a small constant to prevent division by zero. Here, the rewards are interpreted relative to recent performance, allowing the model to adapt dynamically. We arguc that since this an intrinsic reward, its value should be assessed relative to the model's past performance rather than an absolute scale. A response is considered more rewarding if it demonstrates improvement over its recent historical performance, ensuring that the model continuously refines its reasoning rather than converging prematurely.
Final Reward Signal The final memory-based intrinsic reward Rmem is computed as weighted sum of the normalized components:
Rynem = Bg Reseploit + Be Rexplore:
(7)
where Rexploit and Rexplore are the normalized rewards, ensuring that the model evaluates improvements relative to its recent history. The weighting factors 3, and Be determine the balance between reinforcing past successes and encouraging novel reasoning, providing explicit control over-exploitation and exploration trade-offs. With Rmem, dense performance-driven signals are incorporated into the training rewards. We hypothesize that this facilitates learning performance-basee rewards, narrows the difficulty gap between performance-battt based and format-based learning, and mitigates training collapse issues (see more in Sec. 4.4
Training with RL We train the model using a reinforcement learning objective, where the task outcome reward R and the intrinsic reward Ryne used to update the policy To(a q). Specifically, adopt the Group Relative Policy Optimization (GRPO (Shao et al. 2024)), variant of policy gradient methods designed for improved stability and efficiency in language model RL. finc-tuning. The training objective maximizes the expected lotal reward:
max T R.--Damm, [R + Rmem]
(8)
where D is the training dataset, and me is the LLM with its trainable parameters 0.
3 Experimental Setup
Training and Datascts We use three tiny LLMs with most billion parameters: Qwen2.5-0 5-0.5B. Instruct, Falcon3- 1i-is-BB-rre Instructt and Llama3.2 18- Instructt and fine-tune them using single NVIDIA H100 GPU. Training is conducted on two math datasets: (1) the 'easy-math' GSMSK dataset. (Cobee et al. 2021
5

---

Published in Transactions on Machine Learning Research (10/2025)
using the training set of 7,473 samples, and (2) the "hard-math' AI-MO dataset (Jia LI & Polu 2024
where we randomly select only 2,000 samples to reflect real-world high-quality data scarcity. In addition, we employ a logical reasoning dataset: Knights and Knaves (K&K) (Xie et al. 2024), to examine our method on non-math scenarios. The K&K dataset targets logical inference and truth evaluation through procedurally generated puzzles with controllable difficulty and rule-based verification. Despite its simple setup, has proven challenging for both enterprise and open-source LLMs (Xie et al., 2025). We run each training with three seeds to account for the inherent randomness in RL training, ensuring that our results are stable and not dependent on specific initialization. Unless stated otherwise, we implement and execute the training using the Open-R1 codebase HuggiinFFace) 2025 (see more details in Appendix A).
Evaluation We evaluate our approach on three representative mathematical reasoning bench- Dorkks-GMKKK MATH-500, and ATME24 (Cobbe et al., 2021 Lightman et al. 2023) (022))-wichh increase in difficulty in that order; and the logical reasoning K&K test set with novel difficulty levels unseen in the training set. If not stated otherwise, we follow the zero-shot setting for all evaluations. For math datasets, our evaluation framework is based on Lighteval (Fourrier et al., 2023), and we employ its extractive match metric, which rigorously applies regex-based conditions to preciscly extract and parse generated answers. For logical datasels, we follow the evaluation protocols established by the datasel authors. In any case, answers must adhere to a strict, predefined format to be successfully extracted for evaluation; if the model fails to generate an answer that follows the specified format, the answer will not be extracted and will be counted as incorrect.
Basclines We focus on RL fine-tuning approaches. We define R1 as the RL baseline trained using the standard GRPO algorithm (no SFT cold-start), following the DeepSeekRI paper (Guo et al. 2025), with correctness outcome and format-based rewards. Cosine incorporates the response length's property reward signal (Yeo et al. 2025 Our proposed method, Memory-Rt, introduce 2 performance-driven
reward strategies: Rexploit and Rexplore. If not stated otherwise, our method uses 3, = Be = =1 throughout experiments.
4
Experimental Results
4.1 Results with GSM8K Training
For GSMSK training, we use a zero-shot setting for all base LLMs except Llama3.2-1B-Instun which requires a single in-context example per training sample. Without this, it fails to produce any valid correct- ness rewards, preventing learning across all models. All baselines share the outcome correctness rewards and format-based rewards, including the integer reward (rewarding responses that contain integers, tailored for GSM8K task) and the XML reward (ensuring responses match XML structures, c.g., <answer> < </answer>).
We tune GRPO's hyperparameters using baseline R1 and find that in tiny LLMs and low-resource settings, the optimal response length is consistently below 200, There is no significant difference in performance when we increase the training's 's max sequence length to higher values, up to 2048. Thus, we set the maximum length to 200 for all baselines to reduce memory consumption and enable faster training. Regarding hyperparameter G, a higher number of generations per step helps stabilize the training, yet demands more memory resources. To balance resource constraints, we set this value to G 16. Other hyperparameter values are provided in the Appendix A For method, we set the episodic memory capacity equal to the dataset size, i.e., N = [D], ensuring no memory overflow. This design choice is suitable for our setting, where the training dataset is relatively small, aligning with real-world conditions. The maximum number of stored responses per query is fixed at L = 100, and the reward normalization window size is set to w = 100 across all experiments. A key hyperparameter that may require tuning is K, the number of neighbors used for memory retrieval. For simplicity, we set K in this soction as a proof of concopt, without further hyperparameter optimization.
After one training epoch, we evaluate the baselines on mathematical reasoning benchmarks of varying diffi- culty. To analyze the contributions of exploitation and exploration rewards in our method, we also report results for Memory-R, which utilizes only Rexploit as the intrinsic reward for RL training. In contrast, Memory-R+ incorporates both Rexploit and Rexplore; providing a more comprehensive reward structure.
6

---

Published in Transactions on Machine Learning Research (10/2025)
LLM
Baseline
GSMsK
MATH 500
ATME24
Last
Best
Last
Best
Last
Best
Base
27.8
20.0
0.0
R1
27.
28.817.5
18.713.6
18.913.8
0.0.0.00
1.1.1.9
Cosine
29.4=114
31.2:10.7
22.7:11.0
22.7+:1.0
0.0+00.0
0.0=0.0
Owen22550.5B-IInttuuut
Memory-R
33.0:11.1
36.011..6
21.41.99
23.71:1.
0.040.0
0.0.0.0
Memory-R
33.712.55
34.011233
22.3.10.6
2.6
0.0.00.0
1.111.9
Base
32.9
12.2
0.0
R1
10.9=4.6
16.3#1.7
6.5+1.7
10.8+0.4
0.0+0.0
0.0
Cosine
35.3+0.2
37.4+1.3
16.2+1.4
17.0=0.9
0.0+0.0
0.0=0.0
Citkon3-IB-Tnstrucc
Memory-R R
34.6=1.7
36.3#0.6
12.940.3
14.1
0.0
2.2=1.9
Memory-R
34.0=0.6
34.8++0.5
14.0+2.8
16.9-1.3
0.0++0.0
2.2+1.9
Base
26.3
17.4
R1
36.220.3
37.2+1.
15.241.0
19.0.0.5
3.310.0
2.2+1.
Cosine
3..
38.11118
18.510.3
0.3
0.010.0
3.3
Llama3.2-1B- 2-2-1B-lnsstrut
Memory-R
38.7=0.8
39.9:1.22
20.311.4
21.1=0.7
3.340.0
4.4=1.9
Memorr-R"
40.511.1
40.711.1
20.011.2
1.4
1.4
0.040.0
2.211.9
Table 1: Results over different LLMs and datascts. Extractive match (mean = std.) at the last and best training checkpoints, averaged over 3 seeds (for Base, only one seed enough). The average best results are highlighted in bold, and the second- best results underlined; if two best results occur (Cohen effect size < 0.5), all are bold without underlining, and scttings with full zero performance are left unformateed.
statistically identical
Method
GSM8K
Math500
AIME24
Base
73.4
70.0
6.6
R1
73.2 = 0.5
67.9 + 1.7
11.1 - 5.0
Cosine Baseline
74.2 + 0.5
70.2 + 1.1
4.9 + 2.3
Memory-R*
75.4 + 0.3
70.9 + 1.4
13.3 + 3.3
Table 2: Results of Qwen3-0.6B fine-tuned on GSM8K.
Main Results: Table [ presents test accuracy across multiple training checkpoints, reporting the best and last checkpoint's results. Memory-R* emerges as the strongest performer, ranking highest in 10 cases, followed by Memory-R with 8. Cosine and R1 achieve top rankings 7 and 2 times, respectively. Compared to the Base model, our Memory-R variants yield performance improvements ranging from 2% to 14%, depending on the setting. Notably, some RL methods, such as R1, occasionally underperform the Base model due to training collapse. Across all runs, seeds, and settings, we observe that training collapse does not occur with Mcmory-R* whereas does affect the other mothods. We provide a detailed analysis of this phenomenon in Sec.
Scaling with LLM Capability: We additionally examine Memory-R+
stronger small-scale model, Qwen3-0.6B (Thinking Mode), which has recently demonstrated reasoning performance on par with several 8B models (Yang el al. 2025). This serves as robust. testbed to assess whether our method can further enhance the reasoning ability high-performing yet compact LLMs, Table 2 reports the final accuracy on test benchmarks (best checkpoints). Across all benchmarks, Memory-R* consistently outperforms other RL basclines, indicating that its benefits scale with the underlying model's capability and remain effective capable small-scale architectures.

---

Published in Transactions on Machine Learning Research (10/2025)
Cosine
R1
Memory-R+
GSM8K
MATH_50 _50
AIME24
0.26
Test Accuracy
0.04
0.35
0.24
0.22
0.02
0.30
0.20
0.00
100
200
0
100
200
0
100
200
Training Steps
Training Steps
Training Steps
Figure 2: Performance of fine-tuning Qwen2.5--..B-Insttnne on AT-MO data. The test accuracy is evaluated at multiple checkpoints during training (sta over 3 runs). The Base model's performance corresponds to step 0.
Analysis Intrinsic Reward: We also visualize the learning curves of Rexploit and Rexplore over training steps in Appendix Fig. [ Both rewards generally show an upward trend, suggesting that the RL algorithm effectively optimizes them. It is important to note that these rewards are normalized to reflect relative improvements. For exploration reward, there is a warm-up period during which Rexplore remains zero while initial data is collected to estimate novelty. After this phase, a spike in exploration occurs when Rexplore is first applied, followed by stabilization and gradual increase. The rise in Rexplore indicates that the LLM's outputs become: more diverse training progresses.
Effect of Exploration Exploitation Trade-off We varied 3, (exploitation) and 3. (exploration) when training Qwven2.5-0.B-IInsttee on GSM8K (Appendix Table 10 A balanced setting (B, = 33 = 0.5) performed best, showing that neither pure exploration nor pure exploitation alone is optimal and combining both modes more effective than cither alone. For the K&K logic task (sce results in Sec. 4.3 we swept 3. = B, (Appendix Table 11 The optimal scale (10-1) far smaller than in math lasks due procedural generation and highly templated reasoning traces (e.g., 'assume X is knight check consistency - assign roles"). While this allows reusing past reasoning, the high structural similarity means large intrinsic rewards can cause the model to overcommit to familiar template too early. Smaller 3 values mitigate this by letting the model reuse beneficial substructures without locking into an incorrect full template prematurely.
4.2 Results with AI-MO Training
In this task, we focus on Qwen2.5-0.50 5--.5..nsntttrett the smallest LLM in our study. The format-based rewards include the XML reward and a heuristic reward that assesses the clarity of reasoning, dubbed "reasoning step" reward (HuggingFace) 2025). We exclude the integer reward because the answer this task is not limited to integers. We keep the training hyperparameters similar to Soc. except for K, which we vary to study the impact of memory retrieval on the performance of method. Also, the data is limited (only 2,000 samples),
fine-tune the LLMs for epochs to ensure convergence.
Main Results: Fig. 2 reports the test accuracy of 3 RL methods: R1, Cosino, and Mcmory-R* (K = 20) on 3 lest datasets. The results consistently show that Memory-R+ surpasses all baselines by notable margins of approximately 5% on GSM8K, 4% on MATH-500, and 1% on AIME24. In this task, Cosine performs poorly due to response length collapse (sce Sec. 4.4.2p, whereas Memory-R+ and R1 remain unaffocted. However, R1 exhibits significantly slower learning (see Appendix's Fig. 5p) and achieves lower test accuracy compared to M+
8

---

Published in Transactions on Machine Learning Research (10/2025)
Dataset
Model
K
K 10
K 20
K 30
Best
37.6++0.8
38.7+1.2
38.2+0.9
38.4 0.6
GSMSK
Last
36.5+0.2
37.811.4
37.5+0.3
36,8+0.4
Best 25.5+1.1
25.910.3
25.8+0.5
2.3
MATH-500
Last
22.710.11
23.7=0.5
25.3+0.6
25.0=0.4
Best
0.0+0.0
0.7+1.0
2.241.9
0.0
AIME24
Last
0.0+0.00
0.3++088
2.2+1.9
0.0+0.0
Table 3: Memory-R+ test accuracy with different k = 1, 10, 20 and 30 (meandistd. over 3 runs). The best results are highlighted in bold; if two more statistically identical best results occur (Cohen effect size < 0.5), all are bold.
Hyperparameter Selection: Table 3 presents Memory-R*3 performance with different values of K, showing how varying K influences the model's test accuracy across multiple datascts. As observed, incrcasing K generally leads to an improvement in performance. This trend holds for all three datasets, where the highest accuracy is achieved with K 1. However, the improvement varies depending the dataset, indicating that the model's behavior may differ based task complexity or data characteristics. These findings suggest that tuning K can be crucial and further improve Memory-R** performance in downstream tasks. Out of all settings tested, K = 20 demonstrates the highest performance for this task, consistently ranking at the top.
Exploration Analysis: Furthermore, we evaluate the response diversity after fine-tuning LLMs with our method and other approaches, selecting seed that ensures no training collapse occurs in the other methods. We compute diversity scores on 3 randomly sampled responses from the LLMs, given the input from subset of 100 geometry questions Hendrycks et al. (2021). We utilize the Language Model Evaluation Harness library (Gao et al. 2024) to generate model responses using the 'hendrycks, Sendrycks_mat__eeeee
task. The results in
Appendix's s Table 13 demonstrate that Memory-R+ enhances the diversity of the base model, significantly surpassing RI in terms of diversity. Further details and examples are provided in Appendix C.3
Emergence of Self-reflection: Finally, we investigate the outputs of LLMs trained with our method to analyze the self-verification behavior discussed in Guo et al. (2025). To this end, we read through the responses of the LLMs on the geometry task mentioned above. Interestingly, the tiny LLMs trained with Memory-R+ also exhibit self-verification behaviors, as indicated by phrases like "let's re-evaluate' and "let's consider an alternative approach'. Among 100 observed cases, 26 instances demonstrated such behavior, compared to only 6 in the Base model. This highlights the emerging capability of smaller models to perform self-verification, a form of reasoning previously thought to be exclusive to larger, more complex models. These instances suggest that, with the right training and mechanisms, small models can not only generate outputs but also evaluate and refine them. More examples are given in Appendix C.3
4.3 Results with K&K Training
In this task, the difficulty of a question specified by the number of people in the question. We use Qwen2.5- Instrun as the Base LLM. We adhere closely to the training and evaluation protocols outlined in the exclusively
code repository]oof previous studies (Xie al. 2025 Specifically, rewards. The training dataset includes questions from difficulty levels 3 to 7, and evaluate Base, R1, Cosine, and Memory-R* across seven difficulty levels (with the number of people ranging from 2 to 8). In this experiment, our mothod employs = 1 as a proof of concept. Unlike in other experiments, we pick a very small value of B, = B = 0.0001, which yields the best results.
correctness and XML
9

---

Published in Transactions on Machine Learning Research (10/2025)
# People (Difficulty)
Avg.
Model
3
4
5
7
Base
0.0 + 0.0
0.0 = 0.0
0.0 = 0.0
0.0 + 0.0
0.0
0.0
0.0 1 0.0
0.0
1 0.0
0.0 + 0.0
R1
28.0 1.9
20.0
# 22
=
6.3
4.7
2.
1.0
10000
0.0
1.7
12
10.7 2.6
Cosine
28.6 + 5.0
10.0
= 1.7
6.6
= 2.8
3.3
=
1,3
1.5
0.3
0.5
1,3
=
155
7.3 2.5
Memory-R* F 36.0 1 1.0 25.0 1 1.0 22.3 1.1 9.3 1 2.5
2.3
1
2.1
2.0
1.0
2.0
10
14.3 1.6
Accuracy (%) by number of people in K&K puzzles (higher difficulty number increases), with results reported mean + standard deviation over 3 runs. Bold denotes the best mean performance.
Table 4:
R1
Memory-R
Memory-R+
Cosine
Correctness Reward
Integer Reward
XML Reward
0.6
0.4
0.2
Value
0.4
0.2
0.0
0.2
0.0
0.0
-0.2
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Epoch
Epoch
Epoch
Figure 3: Reward Mode Collapse in Matcon3-1B-Instree
As shown in Table 4 the task is challenging as indicated by zero accuracy of the Base model. Among RL baselines, Memory-R+ consistently outperforms both R1 and Cosine across nearly all difficulty levels, especially complex cases involving to 5 people. Notably, Memory-Rt achieves the highest overall average accuracy (14.3), surpassing other baselines by significant margin, with a relative improvement gain of approximately 28.5% over RI and 95.9% over Cosine. While RI demonstrates stronger performance one level of difficulty (6 people), its performance exhibits high instability and the largest variance, resulting in subpar overall performance. Cosine performs reasonably only at the simplest level but struggles with all other settings, showing limited capacity for handling increased complexity. These results emphasize the critical role of our memory-based reasoning in addressing structured logical puzzles.
4.4 Training Collpase in Tiny LLMs
When training tiny LLMs for reasoning. we observe that incorporating multiple reward signals (c.g., format, accuracy, elc.) can enhance performance, particularly boosting specific aspects like accuracy with precise format requirements. However, tiny LLMs can easily converge to local optima (training collapses) when exposed to multiple reward signals, resulting in suboptimal performance. In this section, discuss these collapse cases, highlighting the nuances of reward design and exploring how these challenges can be addressed with our method.
4.4.1 Reward Mode Collapse
In this section, we investigate the reward mode collapse phenomenon, where (1) LLMs prioritize learning simpler, typically format-based reward, (2) LLMs become confused by multiple rewards, struggling to learn any effectively. We observe and report this issue using 1B-omt-t,,, though it is not exclusive to this model.
The values of the main reward (correctness reward) and easier format-based rewards (e.g., integer reward and XML reward) are shown in Fig. Here, Memory-R and Memory-R* enhance both accuracy and integer rewards while trading off the XML reward. We hypothesize that incorporating intrinsic content and performance-based rewards, such Rexploit and Rexplore, facilitates correctness optimization and prevents the
10

---

Published in Transactions on Machine Learning Research (10/2025)
R1
Memory-R
Memory-R+
Cosine
Correctness Reward
Integer Reward
XML Reward
Completion Length
0.8
0.5
0.5
200
0.6
0.4
0.4
150
Value
0.3
0.4
100
0.2
0.2
0.2
0.1
0.1
50
0.0
0.0
0.0
0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 Epoch
1.0
0.0
0.5
1.0
Epoch
Epoch
Epoch
Figure 4: Response Length Collapse in Oxven2.5-0.BB-Instrue
model from overfitting to easy cermat-based rewards. In contrast, using R1, without intrinsic rewards, makes the model immediately focus on easier lormat-based rewards without any improvement in the correctness reward. Additionally, Cosine, which uses length-base based intrinsic rewards, fails to learn any reward, resulting in mediocre performance across all criteria. This suggests that relying on naive intrinsic rewards may hinder learning. Among Memory-R and Memory-R+, the latter shows better performance, likely due to its more diverse intrinsic reward structure, which supports both format-based and correctness bess-baedd rewards.
4.4.2 Response Length Collapse
We observed two distinct types of collapse in response length. In one scenario, depending the setup, the LLM struggles to generate meaningful tokens, resulting in unusually brief responses sometimes as short as just 10 tokens. On the other hand, another form of collapse occurs when the LLM fails to halt its generation process (overthinking problem), leading to output sequences that continually expand until they reach the maximum allowed length. We use Qwen2.5-0.55 Innuttrr trained on GSMSK to illustrate both collapse cases mentioned above. We also note that response length collapse also occurs in other settings (see more in Appendix C.2
We present an example of response length collapse in Fig. In this case, Memory-R and Memory-R* successfully avoid this collapse, achieving high correctness rewards while maintaining reasonable completion lengths and balancod metrics. Interestingly, Cosine, despite focusing on optimizing lengths, leads the model to generate the maximum number of tokens early training, yet none of its corresponding rewards increase. This suggests that the model fails to optimize for any meaningful objective despite the excessive token generation (in some other cases, Memory-R can also suffer similar issues). On the other hand, R1 drastically shortens responses. This results in the integer reward spiking the response length drops significantly, indicating that the model is being guided to generate only minimal number of tokens containing only digits. While this satisfies the integer reward, it is detrimental to correctness, which should be the primary optimization objective. Additionally, the XML reward for RI remains unchanged, indicating a complicated relationship between the response length collapse and the reward mode collapse mentioned above.
5 Related Works
Enhancing LLM Reasoning Recent advancoments in LLM reasoning have focused on scaling test-time computation to improve accuracy in complex tasks. Test-time search stralegies, such as beam search (Gao et al. 2023 and majority vote (Wang et al. 2022). aggregate predictions from multiple inference traces to refine accuracy. While these methods are effective, they come with the drawback of significantly increasing computational costs. More sophisticated techniques, using Monte Carlo Tree Search (Feng al. (2023) and Tree-of-Thoughts (aao et al. 2023), adopt structured search approaches to explore possible reasoning paths
11

---

Published in Transactions on Machine Learning Research (10/2025)
more systematically. However, these methods often require bespoke implementations tailored to specific tasks, and they still lead to high inference costs, making them unsuitable for low-resource devices. Tn addition, alternative approaches, such as process reward models (PRM) (Lightman et al. 2023), aim to address particular aspects of reasoning by modeling rewards during inference. While these methods can improve performance in specilic domains, they face several limitations. For instance, Guo el al. (2025) highlights that process-reward models are costly and not universally applicable. These issues underscore the trade-offs between reasoning accuracy and computational efficiency. Automated annotation often fails to provide satisfactory results, and manual annotation is not scalable. Additionally, introducing model-based PRM leads to reward hacking (Gao et al., 2023) and requires extra resources for retraining, complicating the training pipeline.
Reinforcement Learning for Reasoning Enhancement Recent rescarch, starting with DeepScck-RR (Guo et al. 2025), has shown the effectiveness of pure RL training with outcomee-ased rewards in significantly improving reasoning performance, eliminating the need for costly inference-time searches. However, these methods often depend on verifiable ground truth or domain-spocific hcuristics, such as using response length as rewards (Yeo et al. 2025). For instance, Kimi k1.5 (Team et al.) 2025) introduces a method to shorten chain-of-thought using a length penalty in the reward function during online RI,, while Luo et al. (2025); Arora& & 'Zanette (2025) propose an RL objective aimed at minimizing tokens while maintaining accuracy, Other works, such as Chen et al. (2024), ; explore the overthinking phenomenon and suggest generating data for offline policy optimization using first-correct solutions and diversity criteria. We argue that relying on heuristic reward functions restricts the generalization of LLM reasoning across wide range of datasets. Additionally, small LLMs face challenges in generating long sentences, so emphasizing sentence length may not be beneficial for these models. In contrast, our approach leverages episodic memory to derive intrinsic rewards, making it more adaptable and widely applicable. Furthermore, while most existing methods target large LLMs, our work is the first improve this approach for smaller models ($1B parameters).
Episodic Memory For LLMs Several works have explored episodic memory to enhance LLMs' outputs, but they primarily focus on improving retrieval-based prompting rather than upgrading the model's inherent reasoning capabilities. Experiential learning methods like REMEMBERER (Zhang et al., 2024] store past observation-action pairs and retrieve high-value trajectorios to guide LLMs' actions during inference. Simi- larly, Reflexion (Shinn et al. 2024) and ExpeL (Zhang et al. 2024 use memory to extract insights from past successes and failures, integrating them into prompts to improve decision-making. However, these methods rely on strong LLMs like GPT-4 (OeenAl 2024a) without altering their internal reasoning process, using memory solely for explicit retricval during inference. They still treat memory as an external knowledge base rather than an intrinsic driver of learning. In contrast, our method embeds memory-driven intrinsic moti- vation directly into the learning process. Rather than relying on explicit retrieval for in-context learning, our approach finetunes the model by leveraging intrinsic rewards derived from past successes and failures, enabling adaptive and self iel--improiing reasoning.
6 Conclusion
We present Memory-Rt, a novel memory-augmentted reinforcement learning framework that equips tiny LLMs with intrinsic motivation for effective chain-of-thought reasoning. By leveraging episodic memory to compute exploration and exploitation rewards from past successes and failures, our method mitigates problems such as reward sparsity and poor exploration. Experimental results on math and logical reasoning tasks demonstrate that Memory-R* significantly boosts reasoning performance in small models, making RL fine-tuning more accessible and effective in low-resource settings. Our method is first step toward equipping tiny language models with intrinsic motivation for reasoning, but balancing diverse exploration with efficient exploitation of past successes remains an open direction. Future work will enhance memory-driven rewards to better guide this trade-off across increasingly complex reasoning tasks.
While our Memory-R+ framework shows significant promise, there exist potential failures. A key challenge is the potential for memory corruption. Suppose the episodic memory is initially populated with a large number of flawed or non-generalizable examples. In that case, the model may overfit. to these suboptimal patterns
12

---

Published in Transactions on Machine Learning Research (10/2025)
in the early stage of training, leading to worse performance. This underscores the importance of carefully selecting hyperparameters like 8 to prevent early overfitting to the memory samples. Furthermore, the efficiency of our kNN-based memory system is highly dependent on its size. Although our experiments with a memory size of thousands show fast training, scaling to much larger memory capacitics would introduce substantial computational overhead and latency, necessitating significant engineering efforts to optimize memory retrieval and storage.
Our approach is ideal for applications with limited compute resources, such as edge devices cheap com- puters restricted to single GPU. Here, tiny LLM can be fine-tuned to improve its reasoning capabilities on specialized tasks, thereby achieving rapid inference on less powerful hardware while still capable of solv- ing reasoning problems like maths. In contrast, large models, while lacking specific reasoning fine-tuning, excel at tasks requiring broad, general knowledge and can be better suited for open-ended problems where a diverse knowledge base is essential.
Broader Impact Our work has significant implications for both the accessibility and sustainability of AL By making advanced reasoning capabilities available tiny LLMs, we help democratize access to sophisticated AI for organizations with limited resources. This also has positive environmental implications, as Training smaller models is substantially more energy-effiiint than using massive, iillion-parameter models. However, this accessibility also carries a risk: these models could be deployed in critical applications where their limitations are not fully understood, which could lead to unreliable decisions.
References
Daman Arora and Andrea Zanette. Training language models to reason efficiently. arXiv preprint arXiv:2502.044 04463, 2025.
X. Chen, J. Xu, T. Liang, 7. He, J. Pang, D. Yu, L. Song, Q. Liu, M. Zhou, Z. Zhang, R. Wang, Z. Tu, H. Mi, and D. Yu. Do not think that much for 2+3-? on the overthinking of ol-like llms. 2024. URL
https://arxiv: http///axxv...g//bba.a///. 2412.21187
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL. /tttps://arxii..... abs/ abt///1a./2101..///// 22110.11660
Dai Do, Quan Tran, Svetha Venkatesh, and Hung Lc. Large language model prompting with episodic memory. In ECAI, 2024.
Xidong Feng, Ziyu Wan, Muning Wen, Stephen Marcus McAleer, Ying Wen, Weinan Zhang, and Jun Wang. Alphavero-like tree-search can guide large language model decoding and training. ar preprinl w:Xiv:2309.171197 2023.
Clémentine Fourrier, Nathan Habib, Hynek Kydliček, Thomas Wolf, and Lewis Tunstall. Lighteval: A lightweight framework for 11m evaluation, 2023. URL
L. Gao, J. Schulman, and J. Hilton. Scaling laws for reward model overoplimization. In International Conference Machine Learning, pp. 10835-10866. PMLR, 2023.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Lau- rence Golding, Jeffrey Hsu, Alain Le Noachh, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for fow-shot language model evaluation, 07 2024. URL httppeee/eenoooo.rrreeeood/d//// org/records/1260822
Google. Introducing gemini 2.0: Our new ai model for the agentic era, 2024. URL https://blog.got Accessed: 2025-03-05
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in 1lms via reinforcement learning. Xix preprint arXiv:2001. 12948, 2025.
13

---

Published in Transactions on Machine Learning Research (10/2025)
Sonam Gupta, Yatin Nandwani, Asaf Ychudai, Dinesh Khandelwal, Dinesh Raghu, and Sachindra Joshi. Selective self-to-supervieed fine-tuming for generalization in large language models. In NAACL (Pindings), 2025.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. In Thirty-fifth Con- ference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Zhi Hong, Aswathy Ajith, James Pauloski, Eamon Duede, Kyle Chard, and lan Foster. The diminishing re- turns of masked language models to science. In Findings of the Association for Computational Linguistics: ACL 2023, pp. 1270-1283, 2023.
HuggingFace. Open rl: A fully open reproduction of deepseek-r1, January 2025. URL https://githi: com/huuggingffcccpppen-1l
Lewis Tunstall Ben Lipkin Roman Soletskyi Shengyi Costa Huang Kashif Rasul Longhui Yu Albert Jiang Ziju Shen Zihan Qin Bin Dong Zhou Yann Fleurcan Guillammc Lample Jia LI, Edward Beeching and Stanislas Polu. Numinamath tir.
2024.
Akshay Krishnamurthy, Keegan Harris, Dylan J Foster, Cyril Zhang, and Alcksandrs Slivkins. Can large language models explore in-context? arXie preprinl urx::22403. 15371, 2024.
Komal Kumar, Tajamul Ashraf, Omkar Thawakar, Rao Muhammad Anwer, Hisham Cholakkal, Mubarak Shah, Ming-Hsuar Yang, Phillip HS Torr, Salman Khan, and Fahad Shahbaz Khan. Llm post-training: A deep dive into reasoning large language models. arXiv preprint XXii:2502.21321 2025.
Hung Le, Thommen Karimpanal George, Majid Abdolshah, Truyen Tran, and Svetha Venkatesh. Model- based episodic memory induces dynamic hybrid controls. Advances in Newral Information Processing Systems, 34:30313 30325, 2021.
Hung Le, Majid Abdolshah, Thommen K George, Kien Do, Dung Nguyen, and Svetha Venkatesh. Episodic policy gradient training. In Proceedings of the AAAI Conference on Arlificial Intelligence, volume 36, pp. 7317-7325, 2022.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, llya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
H. Luo, L. Shen, H. He, Y. Wang, S. Liu, W. Li, N. Tan, X. Cao, and D. Tao. O1-pruner: hengthhharmoniziin fine-tuning for ol-like reasoning pruning. 2025. URL https://rrxx wwg. http:/bb/2200...1.5 12570
C. D. Manning, P. Raghavan, and H. Schütze. Introduction to Information Retrieval. Cambridge University Press, 2008
J. L. McClelland, B. L. McNaughton, and R. C. O'Reilly. Why there are complementary learning systems in the hippocampus and neocortex: insights from the successes and failures of connectionist models of learning and memory. Psychological Review, 102(3):419-457, Jul 1995. doi:
OpenAl. Hello gpt-4o, 2024a. URL http::///ppaaaccccoo////eee/////
OpenAI. Learning to reason with 1lms, 2024b.
URL https://openai.coom/neee
11ms/
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adria Puigdomenech Badia, Oriol Vinyals, Demis Has- sabis, Daan Wierstra, and Charles Blundell. Neural episodic control. In International conference on machine learning, pp. 2827-2836. PMLR, 2017.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (FMNIP--ICNI.P), pp. 3982 2992, 2019.
14

---

Published in Transactions on Machine Learning Research (10/2025)
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint 2Xxi:2242.033000 2024.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: En- hancing language agents with verbal reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, 2024.
Charlic Snell, Jachoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can be more effective than scaling model parameters. arXiv preprinl or arXiv:2408.00 03314, 2024.
K. Team, A. Du, B. Gao, B. Xing, C. Jiang, C. Chen, C. Li, C. Xiao, C. Du, C. Liao, C. Tang, C. Wang, D. Zhang, E. Yuan, E. Lu, F. Tang, F. Sung, G. Wei, G. Lai, H. Guo, H. Zhu, H. Ding, H. Hu, H. Yang, H. Zhang, H. Yao, H. Zhao, H. Lu, H. Li, H. Yu, H. Gao, H. Zheng, H. Yuan, J. Chen, J. Guo, J. Su, J. Wang, J. Zhao, J. Zhang, J. Liu, J. Yan, J. Wu, L. Shi, L. Ye, L. Yu, M. Dong, N. Zhang, N. Ma, Q. Pan, Q. Gong, S. Liu, S. Ma, S. Wei, S. Cao, S. Huang, T. Jiang, W. Gao, W. Xiong, W. He, W. Huang, W. Wu, W. He, X. Wei, X. Jia, X. Wu, X. XU, X. Zu, X. Zhou, X. Pan, Y. Charles, Y. Li, Y. Hu, Y. Liu, Y. Chen, Y. Wang, Y. Liu, Y. Qin, Y. Liu, Y. Yang, Y. Bao, Y. Du, Y. Wu, Y. Wang, Z. Zhou, Z. Wang, 7. Li, Z. Zhu, 7. Zhang, Z. Wang, Z. Yang, Z. Huang, 7. Huang, Z. Xu, and Z. Yang. Kimi k1.5: Scaling reinforcement learning with llms. 2025. URL https: /ttp::///xvvvvvoeee orgooog///bo. hwwgg/ab/250.... 2501 12599
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V. Le, Ed H. Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh Internaiional Conference on Learning Representations, 2022.
Mengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasumuru, Danqi Chen, Luke Zettlemoyer, and Veselin Stoyanov. Training trajectories of language models across scales. In The 61st Anmual Meeting Of The Association For Computational Linguistics, 2023.
Chulin Xie, Yangsibo Huang, Chiyuan Zhang, Da Yu, Xinyun Chen, Bill Yuchen Lin, Bo Li, Badih Ghazi, and Ravi Kumar. Large language interpolators can learn logical reasoning: A study on knights and knaves puzzles. In The 4th Workshop on Mathematical Reasoning and AI at NeurIPS244, 2024. URL https://openrevieer net/forummidd=xxX6B
Tian Xie, Zitian Gao, Qingnan Ren, Haoming Luo, Yuqian Hong, Bryan Dai, Joey Zhou, Kai Qiu, Zhirong Wu, and Chong Luo. Logic-rl: Unleashing 11m reasoning with rule-based reinforcement learning, 2025.
URL https: htwrg//bb////22// 1472.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.008 09388, 2025.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Advances in neural information processing systems, 36:11809 18111809-1182,,, 2023.
Edward Yea, Yuxuan Tong, Morry Niu, Graham Neubig, and Xiang Yue. Demystifying long chain-of-thought reasoning in lims. arXiv preprint orXiv:2502. 03377, 2025.
Dan Zhang, Sining Zhoubian, Ziniu Hu, Yisong Yue, Yuxiao Dong, and Jie Tang. Rest-mets*; Lim self- training via process reward guided tree scarch. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems.
Danyang Zhang, Lu Chen, Situo Zhang, Hongshen Xu, Zihan Zhao, and Kai Yu. Large language models are semi- paramerrree reinforcement learning agents. In Advances in Neural Information Processing Systems, volume 36, 2024.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Appendix
A Experiment Details
A.1 System Prompt
Following Guo et al. 2025 HuggingFace (2025 the system prompt is designed as CoT prompling with clear requirement for reasoning and answer format, as shown in Table 5
SYSTEM PROMPT:
A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think>>ansser> answer here </answer>
Table 5: System prompt used in our experiments.
The use of <think> and <answer> tags ensures clear distinction between the internal reasoning process and the final output.
A.2 Training Hyperparameters
The model trained using the GRPO optimization framework with carefully selected hyperparameters ensure stable convergence while being suitable for our computing resources. The key hyperparameters for GSMSK and AI-MO are listed in Table There is slight difference in making the training suitable for each dataset while ensuring efficient training. The hyperparameters for K&K are the same as in the dataset's 's code repository.
Hyperparameter
GSM8K
AI-MO
Learning Rate
5 x 10-6
x 10-6
Adam 31
0.9
0.9
Adam 32
0.99
0.99
Weight Decay
0.1
0.1
Warmup Ratio
0.1
0.1
Learning Rate Scheduler
Cosine
Cosine
Batch Size
2
4
Gradient Accumulation Steps
8
16
Number of GRPO Generations
16
16
Maximum Prompt. Length
256
512
Maximum Completion Length
200
300
Training Epochs
1
Maximum Gradient Norm
0.1
0.1
Mixed Precision
BF16
BF16
Table 6: Key training hyperparameters for GSM8K and AI-MO.
We list the links to the LLM models and datasets in Table
16

---

Published in Transactions on Machine Learning Research (10/2025)
Models/Datasee
URL
Qween2.5-0.B-IInstrruu
Llama3- Maamaa--B--nnsite IB-Instruct
Falcon3- 1193-11-Inssuuu
Sentence Transformer
GSM8K
MATH-500
AIME24
AT-MO
K&K
httpss//ggtabbccc///ullaa//ggggg
Table 7: Models and Datasets Details.
GSMSK
Math500
AIME24
Method
Last
Best
Last
Best
Last
Best
SFT
22.5 + 1.8
22.7 + 1.8
14.5 + 0.8
14.7 1.2
0.0 0.0
0.0 = 0.0
Distill
20.3 + 0.6
21.7 + 1.7
8.9 + 0.3
8.9 0.3
0.0
0.0 0.0
Memory-Ri
33.7 + 2.5
34.0 + 2.3
22.3 + 0.6 24.4 + 0.6
0.0 - 0.0
1.1 - 1.9
Table 8: Performance comparison between standard supervised fine-tuning (SFT), distillation (Distill), and proposed Memory-R. method testing data: GSM8K, Math500, and AIME24 datasets. Results are mean accuracy + standard deviation for the last and best checkpoints over 3 runs.
B Additional Experimental Results
Comparison with Non-Reinforcement Learning Approaches Naive SFT is known to suffer from poor generalization in reasoning tasks, particularly for small models and limited training data. Prior work has demonstrated that SFT can even lead to performance degradation compared to the base model on test sets (Gupta et al.) 2025
Regarding distillation approaches, these methods rely on access to a high-quality teacher model, which may not be fcasible in low-resource settings. Furthermore, using data generated by large models raises concerns about dała quality control and potential ethical and legal issues associated with model-generated content. Therefore, direct comparison of our RT.-based method with distillation is not always appropriate. Nonetheless, to provide comprehensive evaluation, we include baselines using SFT and distillation.
Specifically, we use the Cohere LLM CommandR to generate chain-of-thought reasoning answers for the GSM8k training data, and use these outputs to perform supervised fine-tuning of the Qwwen2.5-0.5B-Insttr model following the procedure similar to (Guo et al. 2025). Table summarizes the final results for SFT, Distillation, and our proposed Memory-R method.
The results demonstrate that both standard SFT and distillation baselines substantially underperform pared to our method. For example, Memory-R+ improves accuracy GSM8K by over 10 percentage points compared to SFT (33.7% vs. 22.7%) and exceeds distillation performance by more than 12 points. These find- ings underscore that supervised fine-tuming and distillation face significant challenges in improving language modols with fower than one billion parameters, making our RL-based approach a more effective solution for enhancing reasoning performance in such tiny LMs.
http:///wagghggaa....ooooaaaaaas www.mmandd--puuseeee 1aammandd--puu-0------ 08-2024 2-222-0
17

---

Published in Transactions on Machine Learning Research (10/2025)
Method
Training Time Memory Usage
Memory-R+
10h19m
50.6 GB
R1
9h02m
50.4 GB
Table 9: Comparison of training time and memory usage between Memory-R+ and RI.
Be Value
pg Value
Accuracy + Std (%)
0.0
1.0
33.0 + 1.1
0.1
0.9
33.2 + 1.2
0.3
0.7
32.6 - 1.8
0.5
0.5
34.5 2.3
0.7
0.3
34.1 1.6
0.9
0.1
32.8 = 1.9
1.0
0.0
32.6 + 2.5
Table 10: Impact of varying exploitation (5,) and exploration (6,) weights on GSM8K accuracy.
Training with Larger Datasets To evaluate the scalability and robustness of our approach, we extend training the larger DeepScaler datased while only using a small memory size (2000). Training on Deep- Scaler also tests whether Memory-R+ continues to provide benefits when the data distribution is richer and more challenging.
Table 12 reports performance on the DeepScaler benchmark. Memory-R+ consistently outperforms both the Base and R1 baselines. On GSM8K, our method achieves 30.8 0.2, an improvement over R1 (30.0 + 0.1) and the Base model (27.8). A similar trend is observed on Math500, where Memory-R+ reaches 23.6 0.1, compared to R1 (23.0 + 0.1) and Base (20.0). These results indicate that Memory-R+ is still effective leveraging larger training corpora.
C Training Collapse Examples
C.1 Reward Mode Collapse
We present the reward mode collapse phenomenon in Figure 3 Here, we show the values of the correctness reward, integer reward (ensuring the output is an integer), and XML reward (ensuring the output contains correctly formatted XML tags for parsing) when optimized simultancously using different reward schemes with same seed. evident that while Memory-R and Memory-R' accuracy rewards increase steadily, those of R1 and cosine do not. Instead, these rewards lead the tiny L.LMs to learn reward patterns easily.
C.2 Collapse in Response Length
In Fig. 1 we present different rewards and corresponding completion lengths when training Qwen2.5-0.5h Instruct with the same random seed. These figures reveal two distinct types of response length collapse. Memory-R and Memory-R+ shows robust behaviors with good correctness rewards. Cosine causes the model to generate excessively, reaching the maximum token limit early in training. In contrast, rewards such as integer and XML remain close to zero. R1 shortens responses while optimizing for the integer reward. Both approaches result in low correctness rewards, highlighting their suboptimal behavior.
Fig. 5 presents additional examples of response length collapse observed during the training of Qwen- 2.5-0.5B-nnstrud on the GSMSK and AI-MO datasets. The Cosine method exhibits severe lengthening collapse GSMSK while experiencing shortening collapse on AI-MO. R1 suffers from shortening collapse
18

---

Published in Transactions on Machine Learning Research (10/2025)
Be Value
Bg Value
Average Accuracy Std (%)
0.0
0.0
10.7 2.6
0.00001
0.00001
11.3 + 3.8
0.0001
0.0001
14.3 + 1.6
0.001
0.001
12.5 = 2.0
0.01
0.01
10.1 = 3.5
0.1
0.1
6.4 444::.9 3.9
Table 11: Effect of 8g and 3e scaling on K&K logic task accuracy. Smaller scales yield better results due high structural similarity between procedurally generated puzzles, where large intrinsic rewards can cause premature overcommitment to incorrect templates.
Method
GSM8K
Math500
AIME24
Base
27.8
20.0
0.0
R1
30.0 + 0.1
23.0 ::0:00 0.1
0.0 0.0
Memory-R,
30.8 = 0.2
23.6 0.1
0.0
Table 12: Results on DeepScaler training. The best results are bolded.
GSM8K. Although it does not exhibit collapse AI-MO, it converges more slowly and underperforms Memory-R+, the only method capable of overcoming training collapse.
C.3 Details on Model Outputs
Memory-based Intrinsic Rewards In Fig. we report the memory-based intrinsic rewards (Rezpton and Rezplore) over training steps while fine-tuning 9wven2.5-0.5B-Instruu with Memory-R+ on GSM8K.
Diversity Evaluation To assess the diversity of responses generated by Qwen22-3-0.5--Isstruu fine-tuned with our method, we employ a pairwise similarity analysis. Specifically, for each question, we sample three responses from the model using a temperature of 0.1, and compute the pairwise cosine similarity between them. This process is repeated for set of 100 questions to obtain comprchensive measure of response diversity.
To capture both lexical and semantic similarities, we utilize two different embedding models: TF-IDF Man- ning et al. 2008) and Sentence Transformer (Reimers & Gurevych 2019). The TF-IDF model provides a measure of lexical overlap, while the Sentence Transformer captures deeper semantic relationships between responses.
For each question, we compute the average pairwise cosine similarity for the three sampled responses using both embedding models. The final results, reported in Table 13 reflect the overall diversity of the model's responses across the dataset. report/ 1 similarity as the diversity score, where higher value indicates greater diversity, suggesting more variation in the generated outputs.
The results confirm that our method effectively encourages the LLM to explore more. Compared to R.I, the final model trained with our approach demonstrates clear improvement in diversity. We present several sampled responses of the Base, R11, and Memory-Ri + in Table 14 these responses listed below:
15 and respectively. The question for
19

---

Published in Transactions on Machine Learning Research (10/2025)
(a) GSM8K
Cosine
R1
Memory-R+
(b) AI-MO
200
Training Accuracy
0.4
Training Length
150
Accuracy
0.08
Training Length
250
200
100
0.006
100
0.1
0.02
50
0.0
0.00
300
100
150
200
150
200
Training Steps
Training Steps
Training Steps
Training Steps
Figure 5: More Training Collapses in Qwen2.5-0.5t 5-0.555-00.B--nstreee during finc-tuning GSM8K (a) and AI-MO
datasets (b). The results have been smoothed to improve clarity and visual appeal.
Memory Explore Reward
Memory Exploit Reward
0.5
0.15
0.4
0.3
0.10
0.2
0..55
0.1
0.00
100
200
300
400
100
200
300
400
Training Steps
Training Steps
Figure 6: Memory-based Intrinsic Reward on GSMSK and Qwen2.5-0.5B. Instrucct. smoothed to improve clarity and visual appeal.
The results have been
Problem: In right triangle ABC' shown below, cos B 10. What is tan C?
draw (0,0) (10,0) (3.6,4.8) cyckebllack (ittr/gggaanleeaaa(00.... (3.6,).)).. black label("C", (abel("C"((10.0), (10,0),E);; label('A ("eee("A,(3.6,,4,,,,"" (3.6,4488),)) label(" D", (O, label("10", (0,0) (1),0))))
Self- Verification Behaviors Examples of our method's output, showcasing self-verification behaviors (highlighted in red), can be found at the end of the paper.
20

---

Published in Transactions on Machine Learning Research (10/2025)
Method
Lexical Diversity
Semantic Diversity
Base
27.3
9.0
R1
8.6 27.6
24.2
Memory-R*
9.3
Table 13: Diversity scores (x100) for different baselines. Bold denotes best results.
Response 1
Response 2
Response
-3
To find tan C' given cos B Step-by-oe bt-step:
10'
Given a right triangle ABC with 10 we need find
solve for tan C, we use:
cos B
Cos B 6
tan C.
cos B 10
Using the Pythagorean theorem:
10
Since
Adjacent side: a, hypotemuse: Using the ratio:
BC = v AB AC
cos B adjacent hypotemuse
We express cosine as:
B adjacent
hypotenuse
B adjacent
we find 6, hypotemuse 10.
hypotenuse
Verifying:
Since cos
cos B AC
10 10
By rross-multiplication:
Next, using sine:
sin
Table 14: Qwven2.5-0.5B-Innsnnnin responses.
21

---

Published in Transactions on Machine Learning Research (10/2025)
Response 1
Response 2
Response 3
To solve for tan C, given cos B
To determine tan C
given
To solve for tan C, analyze
6 in AABC:
cos B = 6 in AABC:
AABC:
Identify sides:
AB = hy-
Trigonometric
definitions:
Given:
potenuse - BC = one leg AC
cos B
adjacent hypotenuse
6
= other leg Step Find sin B:
tan B opposite
COS
adjacent
10
Step 1: Find the hypotenuse
ing cos B:
Step 1: Find sin
cos B adjacent
hypotenuse
cos adjacent
adjacent
cos B
hypotenuse
hypotenuse
For cos B 16'
6
Given:
Since cos B 10'
BC 6
AB
10
cos B
6
adjacent 6
10
hypotemuse 10
Let BC
Let b be the adjacent side and h. be...
Let hypotemuse = h:
x:
AB
10
adjacent 6k
Since AB is the hypotenuse:
where k is a positive number.
10r
5x
Next, find:
AB
6
sin B opposite
Next, find sin B...
hypotenuse
Table 15: Qwen2.55-..BB-IIssrucc fine-tuned with RI's responses.
Response 1
Response 2
Response 3
To solve for tan C, given cos B
To determine tan C, first find sin and cos B for AABC.
To find tan C' in AABC: Given:
6.
10:
Step Recall trigonometric def-
Given:
initions cos B adjacent
cos B
6
3
sin B apposite hyposenuse
B 6
10
5
hypotenuse
10
Step 2: Use Pythagorcan iden-
Trigonometric definition:
tity
Use the Pythagorean theorem to find the hypotenuse:
cos B adjacent
sin B cos B
hypotemuse
AB² BC = 2C
Substituting cos B = To.
6.
Denote sides:
AB c BCC
Substituting known values:
AC
Given:
6
sin? 2 B +
=
10 6 2
10
3
3 5
100 + 2C AC V136 2/34
5
Calculate sin B...
Using Pythagorean theorem:
Next, calculate sin B using:
2 a²
opposite
Solve for remaining sides...
sin hypotemuse
Table 16: 20wen2.5-0.B--nnttmmee fine-tuned with Memory-R+'s responses.
22

---

Published in Transactions on Machine Learning Research (10/2025)
Memory-Ri Rt's Output Example 1
Question: Coplanar squares ABGH and BCDF are adjacent, with CD=10 units and AH=5 units. Point E on segments AD and GB. What the area of triangle ABE, in square units? Express your answer as common fraction.
Response: To find the area of triangle AAB where ABGH and BCDF are coplanar squares with CD = 10 units and AH = 5 units, we need to determine the coordinates of points E and B. First, let's place the squares on the coordinate plane for easier calculation:
Square ABCD has vertices at A(0, 0), B(5, 0), C)5,5), and 5). Square BCDF has vertices at (55, 0), C(15,)) 0), D(15, 10), and F(5, 10).
Next, we need find the coordinates of point E. Since E lies on segment AD and GB, we can the fact that AD and GB are parallel lines (both have slope of 1). Step 1: Determine the coordinates of E Since E lies on AD:
E divides AD into two equal parts.
The equation of line AD can be found using the points A(0, 0) and D(0, 5):
= - 5.
To find the intersection with AD, set æ O:
y = - 5
Thus, E is at 50,5).
Step 2: Determine the coordinates of
Since B lies on BC and DF:
B divides BC into two equal parts.
The distance from A BE simply the absolute value of the y-coordinate
Height = 00 - 0 = 0.
However, this interpretation seems incorrect because A should be above BE. Let's re-evaluate: Correctly, the correct height should be the vertical distance from to the line BE.
Now, calculate the area of AABE:
25
Area x
x5x 2
12.5.
2
Thus, the area of triangle AABE is:
25
23

---

Published in Transactions on Machine Learning Research (10/2025)
Mcmory-R+3s Output Example 2
Question: In the diagram, points X, Y and Z are on the sides of AUVW, as shown. Line segments UY, VZ and WX intersect at Point Y on VW such that VY: YW 4 If APYW has an area of 30 and APZW has an area of 35, determine the area of AUXP.
Model Response:
To solve for the area of AUXP, we need to use the given information about the areas of other triangles in the diagram.
First, let's find the ratio of the lengths VY to VW. We know: The ratio VY YW 3.
Given that 4PZW has area of 35, find the area of AUXP by subtracting the areas of ΔPYW and APZW from the area of AUXW.
First, we find the area of AUXW:
Area of AUXW = Area of AUYW + Area of AUXW
Finally, we find the area of AUXP:
Area of AUXP Area of AUXW (Area of APWW + Area of APZW)
Area of AUXP 30 (15 + 35) 30 50 -20
However, this result seems incorrect because the area cannot be negative. Let's re-evaluate the problem using the correct approach.
Given: Area of APYW = 30 Area 4PZW = 35 The correct calculation should be:
Area of AUXP = Area of AUXW (Area of APYW + Area APZW)
Area of AUXP = 30 (15 + 35) 30 50 = 20
This indicates mistake in the interpretation or calculation. Revisiting the problem, we realize the correct approach involves recognizing the geometric properties and relationships between the triangles.
Given the correct geometric properties and solving step-by-step correctly, we find:
24


Reviews:
[{"review_id": "lQobdxXT7X", "reviewer": "Reviewer_DJYK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The memory-augmented intrinsic reward mechanism is creative and well-motivated, drawing inspiration from human episodic memory while being computationally efficient through kNN retrieval.\n\n2. The paper provides thorough ablation studies, analyzes different failure modes (reward mode collapse, response length collapse), and tests across multiple datasets and model architectures.\n\n3. The identification and visualization of specific training collapse patterns in tiny LLMs (Sections 4.4.1 and 4.4.2) provides valuable insights for the community.\n\n4. The authors provide implementation details and demonstrate that the method works on a single H100 GPU, making it accessible to researchers with limited resources.\n\n5. The discussion of self-reflection behaviors emerging in 0.5B parameter models is particularly interesting and challenges assumptions about model size requirements.\n\n**Weaknesses**\n\n1. While the paper focuses on tiny LLMs, it would be strengthened with direct comparison with larger models (8B-70B), both with and without reasoning enhancements. Understanding how these enhanced tiny models compare to larger base models would provide important context for the results.\n\n2. In addition, the paper primarily compares against other RL methods but doesn't include comparisons with simpler approaches like few-shot prompting or supervised fine-tuning on the same data.\n\n3. While the authors set memory capacity equals dataset size, this approach may not scale well to larger datasets. Discussion on memory management strategies for practical deployment could strengthen the paper.\n\n4. The paper is primarily empirical and lacks theoretical justification for why the proposed intrinsic rewards should lead to better exploration-exploitation trade-offs.\n\n5. The method introduces several hyperparameters, and while some ablation is provided, the sensitivity analysis is limited. The K&K experiment uses very different values without clear justification.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Memory-R+, a novel memory-augmented reinforcement learning framework designed to enhance chain-of-thought (CoT) reasoning in tiny language models. The key contributions are as follows:\n\n1. The authors propose using episodic memory modules to store successful and failed reasoning trajectories, computing intrinsic rewards based on similarity to past successes (exploitation reward) and novelty relative to past failures (exploration reward).\n\n2. The work identifies and addresses critical issues, including reward sparsity, ineffective exploration, and training collapse phenomena (reward mode collapse and response length collapse) that occur when applying RL to small models.\n\n3. Experiments on GSM8K, MATH-500, AIME24, and Knights & Knaves datasets demonstrate that Memory-R+ significantly improves reasoning performance in models as small as 0.5B parameters, achieving 2%-14% improvements over baselines.\n\n4. The paper mentions the emergence of self-reflection and self-verification behaviors in tiny models trained with Memory-R+, previously thought to be exclusive to larger models.\n\nThe core technical innovation lies in the kNN-based episodic memory system, which provides dense, performance-driven intrinsic rewards to complement sparse outcome rewards, thereby enabling effective RL fine-tuning for resource-constrained models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The following changes are critical:\n\n1. Including comparisons of how Memory-R+ enhanced tiny models perform relative to larger models (e.g., 7B, 13B, 70B), both with and without reasoning enhancements, would be essential for understanding where these models sit in the performance spectrum. \n\n2. Adding comparisons with supervised fine-tuning on the same training data could be important to determine whether RL is necessary for the improvements.\n\n3. Addressing how the method would scale to larger datasets where setting memory capacity equal to dataset size becomes impractical would also be practical concerns.\n\n\n\nThe following changes could strengthen the paper:\n\n1. Explicitly mention \"Memory-R+\" in the abstract as the proposed method name.\n\n2. Provide theoretical insights or at least intuitive explanations for why the combination of exploitation and exploration rewards leads to better performance than alternatives. \n\n3. Could the authors also conduct more extensive hyperparameter sensitivity analysis, particularly explaining the drastic difference in $\\beta$ values between math and logic tasks.\n\n4. It is also useful to include a detailed comparison of training time, memory usage, and inference costs compared to baselines. Table 5 provides training hyperparameters but not computational metrics.\n\n5. Including additional discussions on where Memory-R+ fails or underperforms, as well as guidance in choosing small reasoning models vs large models without reasoning, could help readers better understand the use cases.", "broader_impact_concerns": "The paper does not discuss a broader impact, though one would be beneficial given several considerations:\n\n1. The work contributes to making advanced reasoning capabilities accessible to organizations with limited computational resources.\n2. Training smaller models has positive environmental implications compared to large model training.\n3. Small reasoning-capable models could be deployed in contexts where their limitations aren't properly understood, potentially leading to harmful decisions in critical applications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "sRAw5v7Zn9", "reviewer": "Reviewer_aFG5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper clearly identifies the key obstacles in applying RL to tiny LLMs and proposes a well-motivated memory-based solution.\n- The experimental results show that Memory-R+ consistently outperforms traditional RL baselines.\n- The paper includes insightful analyses, such as investigations into training collapse and the effect of memory size, which help to validate and contextualize the proposed method.\n\n### Weaknesses\n- One concern is the lack of comparison with supervised fine-tuning (SFT) or distillation-based baselines. While it is commendable that the authors successfully made RL work for tiny LLMs, it is unclear whether this approach is superior to simpler SFT methods, especially when a teacher model is available.\n    - For instance, distillation often reduces formatting errors, which are the main cause of RL failure in tiny models. In that case, existing RL methods might work just fine. It would be helpful to know whether Memory-R+ still brings significant benefits in such settings.\n- It is also unclear how meaningful RL fine-tuning is for tiny LLMs in practice. The contribution, while technically sound, may be limited in broader impact. Moreover, the proposed method involves additional complexity and hyperparameter tuning. Without a more direct comparison to simpler alternatives (e.g., reward shaping via smarter answer parsing), it is hard to assess whether the added complexity is justified.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses the challenge of applying reinforcement learning (RL) to tiny LLMs (<1B parameters), where existing RL methods often fail due to sparse rewards, mainly caused by these models’ inability to produce properly formatted outputs. To mitigate this issue, the authors propose Memory-R+, a memory-augmented RL framework that stores both successful and failed reasoning trajectories. By measuring semantic similarity between new outputs and those stored in memory, the framework provides dense intrinsic rewards, facilitating more effective learning.\n\nThe proposed approach enables tiny LLMs to be fine-tuned via RL, a previously unreliable setup, and shows consistent improvements across several reasoning benchmarks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Comparison with Supervised Fine-Tuning (SFT) or Distillation Baselines\n\nPlease include empirical comparisons against supervised fine-tuning (e.g., distillation from a larger teacher model). Since formatting errors are a major challenge for RL with tiny LLMs, it is important to assess whether distillation alone can resolve these issues—and whether Memory-R+ still provides benefits in such settings.\n\n2. Clarify Practical Value of RL for Tiny LLMs\n\nThe paper would benefit from a discussion of the practical utility of applying RL to tiny models. Given the added complexity and tuning overhead, please clarify the intended deployment scenarios or concrete advantages over simpler fine-tuning methods.\n\n3. Justify Design Complexity over Simpler Alternatives\n\nConsider comparing Memory-R+ against simpler reward shaping approaches, such as using relaxed or rule-based parsing strategies to extract answers, even from imperfect outputs. This would help establish whether the memory-based method provides a meaningful improvement over lower-cost solutions.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "uBJZRBMspc", "reviewer": "Reviewer_ERN8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n1. Addresses a key limitation of small models: Directly tackles the poor exploration capabilities of tiny LLMs in RL settings, an underexplored but important problem.\n\n2. Efficiency: Designed to work in low-resource settings (only one H100 GPU)\n\n\nWeaknesses:\n\n1. I would like to see more intuitive explanations of why the proposed method works.\n2. Does the proposed method work only for small-scale models, or is it also effective for larger models?\n3. More ablation studies on the exploration–exploitation dynamics would be helpful.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work introduces an intrinsic motivation framework to improve chain-of-thought reasoning in tiny language models (≤1B parameters) using reinforcement learning. Unlike prior approaches that rely on large-scale models with strong pretraining for effective exploration, this method leverages a k-nearest-neighbor-based episodic memory to compute intrinsic rewards. The memory stores successful reasoning traces, encouraging the model to reuse effective patterns while also exploring novel strategies. Inspired by human memory-guided learning, this approach enables a balance between exploitation and exploration, facilitating reasoning improvements even in low-resource settings. Experiments on three reasoning benchmarks demonstrate substantial gains in performance and generalization over standard RL methods, making reinforcement learning-based reasoning enhancements more viable for small models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses", "broader_impact_concerns": "No concerns on the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)