PaperlayerView exact revision on arXiv

Source-linked AI summary

Reasoning Under 1 Billion: Memory-Augmented Reinforcement Learning for Large Language Models

Hung Le, Dai Do, Dung Nguyen, Svetha Venkatesh

arXiv:2504.02273v1cs.LG

TL;DR

RL-based reasoning improvements have largely been demonstrated on large models, while tiny LLMs can receive sparse outcome rewards and struggle to explore. Memory-R+ adds intrinsic rewards from episodic memories of successful and failed reasoning. The paper reports improved accuracy and robustness over baseline RL and handcrafted rewards, including strong results across its mathematical benchmarks.

  • Problem

    RL remains challenging for tiny LLMs with ≤1B parameters, which often produce incorrect outputs that receive no outcome reward, leaving sparse training signals.

    Source: §1 · p. 1

  • Method

    Memory-R+ uses episodic memories of successful and failed reasoning traces to compute intrinsic rewards that promote successful patterns and novel responses.

    Source: §1 · p. 2

  • Results

    Memory-R+ improves reasoning accuracy and robustness over baseline RL and handcrafted rewards, ranking highest in 10 cases in GSM8K-training results.

    Source: §1 · p. 2, §4.1 · p. 6

  • Takeaways & Limitations

    The approach supports effective RL fine-tuning down to 500M-parameter models and lowers the barrier for groups with limited computing resources.

    Source: §1 · p. 2

    For binary math outcomes where correct answers receive 1 and incorrect answers 0, the success and failure thresholds are both set to 0.5.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

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 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 for 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 fine-tuning GSM8K and AI-MO datasets demonstrate that our approach significantly enhances smaller LLMs' sample efficiency and generalization capability, making RL-based reasoning improvements more accessible in low-resource settings.

1 Introduction

RL fine-tuning has improved LLM reasoning, but its demonstrated benefits have largely been in large models; tiny LLMs face sparse outcome rewards and weak exploration. Memory-R+ uses episodic memory to provide intrinsic rewards, and experiments report improved reasoning accuracy and robustness for small models.

1 Introduction

  • Tiny LLMs with ≤1B parameters often produce incorrect outputs and receive no outcome reward, creating sparse feedback; format-based rewards can also cause training collapse.The paper also identifies ineffective exploration and limited downstream-task data as challenges.Source: §1 · pp. 1–2
  • Memory-R+ derives intrinsic rewards from stored successful and failed reasoning traces, encouraging alignment with successful patterns and avoidance of repeated errors.The framework uses nearest-neighbor retrieval in a shared representation space to produce performance-driven rewards.Source: §1 · p. 2
  • Memory-R+ improves reasoning accuracy and robustness over baseline RL and handcrafted rewards, with experiments reporting stronger performance on CoT reasoning tasks.The introduction reports extensive experiments across tiny LLMs on mathematical problem-solving.Source: §1 · p. 2
  • Memory-R+ is presented as enabling effective RL fine-tuning for models as small as 500M parameters, lowering the barrier for groups with limited computing resources.The paper contrasts this scale with the much larger models used in RL-based reasoning research [Guo et al., 2025].Source: §1 · p. 2

2 Method

Memory-R+ supplements verifier outcome rewards with intrinsic rewards from episodic memories of successful and failed responses. It retrieves related experience, rewards similarity to successful patterns and novelty from failures, then normalizes and combines these signals for RL training.

2.1 Intrinsic Reward Formulation for CoT Reasoning

  • For each generated response, Memory-R+ combines the verifier’s outcome reward with an episodic-memory reward that complements the verifier’s final-answer-only assessment.The memory reward compares responses with retrieved successes and failures, encouraging similarity to successful responses and novelty from failed ones.Source: §2.1 · pp. 2–3, Figure 1 · p. 3

2.2 Episodic Memory

  • The episodic memory stores query-response pairs encoded by a pretrained Sentence Transformer in a shared high-dimensional space.The memory has a maximum number of queries and a limit on associated responses per query.Source: §2.2 · p. 3
  • Separate success and failure memories classify responses by reward thresholds; for binary math rewards, both thresholds can be set to 0.5.New responses are merged with existing query entries, while oldest queries or responses are discarded when capacity limits are exceeded.Source: §2.2 · pp. 3–4

2.3 Memory-based Intrinsic Reward

  • For a new query, Memory-R+ retrieves responses associated with the top-k nearest memory queries, then rewards proximity to the successful-response centroid and novelty from failure responses.The success reward uses Euclidean distance to a retrieved-response centroid; the failure reward uses response similarity to discourage repeating past mistakes.Source: §2.3 · p. 4
  • The exploration reward encourages broad novelty early, then more focused exploration as successful responses accumulate; a warm-up period can precede its use.The paper gives 50 training steps as an example warm-up period.Source: §2.3 · p. 4
  • The two intrinsic rewards are normalized relative to recent performance and combined as a weighted sum, then added to outcome reward in GRPO training.The weights are set to βs = βe = 1 throughout the experiments; the authors hypothesize that dense signals may mitigate training collapse.Source: §2.3 · pp. 4–5

3 Experimental Setup

The experiments fine-tune three tiny instruction-tuned LLMs on GSM8K or AI-MO and evaluate mathematical reasoning on three benchmarks. Comparisons include standard GRPO and a response-length reward baseline.

3 Experimental Setup

  • Training uses Qwen2.5-0.5B-Instruct, Falcon3-1B-Instruct, and Llama3.2-1B-Instruct on one H100 GPU, with GSM8K’s 7,473 training samples or 2,000 randomly selected AI-MO samples.Training runs use three seeds.Source: §3 · p. 5
  • Evaluation covers GSM8K, MATH-500, and AIME24 in zero-shot settings, using extractive matching that counts answers failing strict format requirements as incorrect.The benchmarks are described as increasing in difficulty in that order.Source: §3 · p. 5
  • The R1 baseline uses standard GRPO with correctness and format rewards, while Cosine adds response length as a reward and Memory-R+ uses exploitation and exploration rewards.The proposed reward strategies are performance-driven.Source: §3 · p. 5

4 Experimental Results

Across GSM8K and AI-MO training, Memory-R+ performs strongly and avoids several training-collapse patterns. The results also show that reward design and memory-retrieval settings shape performance.

4.1 Results with GSM8K Training

  • Memory-R+ ranked highest in 10 cases (Memory-R: 8; Cosine: 7; R1: 2), and Memory-R variants improved performance over Base by 2%–14% across settings.Table 1 reports extractive match at the last and best checkpoints, averaged over three seeds (one for Base).Source: §4.1 · p. 6, Table 1 · p. 6
  • Normalized exploit and explore rewards generally rose during training; exploration began with a zero-reward warm-up before increasing alongside output diversity.The warm-up collects initial data to estimate novelty; the exploration reward then spikes, stabilizes, and gradually increases.Source: §4.1 · p. 7

4.2 Results with AI-MO Training

  • Memory-R+ (K = 20) surpassed baselines by approximately 5% on GSM8K, 4% on MATH-500, and 1% on AIME24 after AI-MO training.Figure 2 reports test accuracy for three RL methods on three test datasets.Source: §4.2 · p. 7, Figure 2 · p. 7
  • Across three datasets, increasing K generally improved test accuracy, with the best tested results at K = 20 for this task.The gains varied by dataset, and the highest accuracy on each dataset occurred with K > 1.Source: §4.2 · p. 7, Table 2 · p. 8
  • Memory-R+ increased response diversity over Base, and self-verification appeared in 26 of 100 observed cases versus 6 for Base.The self-verification examples came from geometry-task responses.Source: §4.2 · p. 8

4.3 Training Collpase in Tiny LLMs

  • With multiple reward signals, tiny LLMs can converge to local optima and achieve suboptimal performance.The section examines training collapse arising in reward design.Source: §4.3 · p. 8

4.3.1 Reward Mode Collapse

  • Reward-mode collapse involves favoring easy format rewards or becoming confused by multiple rewards; Memory-R variants improved correctness and integer rewards while trading off XML reward.R1 focused on easier format rewards without improving correctness, while Cosine failed to learn any reward; Memory-R+ performed better than Memory-R.Source: §4.3.1 · pp. 8–9, Figure 3 · p. 8

4.3.2 Response Length Collapse

  • Response-length collapse ranges from outputs as short as 10 tokens to runaway generation; Memory-R variants avoided the illustrated cases, though Memory-R can also fail in other settings.Cosine reached the maximum token limit, while R1 shortened responses and increased integer reward at the expense of correctness.Source: §4.3.2 · p. 9, Figure 4 · p. 9

5 Related Works

Prior reasoning approaches use inference-time search or reinforcement learning, while episodic-memory work has mainly guided retrieval at inference. Memory-R+ instead uses memory-derived intrinsic rewards during fine-tuning.

Enhancing LLM Reasoning

Reinforcement Learning for Reasoning Enhancement

  • Pure RL with outcome-based rewards can improve reasoning without costly inference-time searches, but existing approaches may depend on verifiable ground truth or domain-specific heuristics.Examples include response-length rewards [Yeo et al., 2025] and length penalties in online RL [Team et al., 2025].Source: §5 · p. 10
  • Memory-R+ derives intrinsic rewards from past successes and failures during fine-tuning, rather than using memory only as an external retrieval base.The paper positions this approach as more adaptable than heuristic rewards and as targeting smaller models (≤1B parameters).Source: §5 · p. 10

Episodic Memory For LLMs

  • Memory-based methods such as REMEMBERER [Zhang et al., 2024], Reflexion, and ExpeL [Zhang et al., 2024] guide inference through retrieved experiences rather than intrinsic learning rewards.The cited methods use strong LLMs such as GPT-4 without altering the model’s internal reasoning process.Source: §5 · p. 10

6 Conclusion

Memory-R+ uses episodic memory to provide tiny LLMs with intrinsic motivation for chain-of-thought reasoning. The paper reports improved math-reasoning performance and greater accessibility of RL fine-tuning in low-resource settings.

6 Conclusion

  • Memory-R+ computes exploration and exploitation rewards from past successes and failures, and experiments report significantly improved reasoning performance in small models.The conclusion states that this mitigates reward sparsity and poor exploration on math reasoning tasks.Source: §6 · p. 10

A Experiment Details

The system prompt requires chain-of-thought reasoning and a specified answer format. It uses separate tags to distinguish internal reasoning from the final answer.

A.1 System Prompt

  • The system prompt uses chain-of-thought prompting with explicit reasoning and answer-format requirements.The prompt is shown in Table 3.Source: §A.1 · p. 14, Table 3 · p. 14

SYSTEM PROMPT:

  • The <think> and <answer> tags separate the model’s internal reasoning from its final output.The prompt text describes the assistant solving the question and giving an answer.Source: SYSTEM PROMPT · p. 14

A.2 Training Hyperparameters

The experiments use GRPO with hyperparameters selected for stable convergence and resource suitability. Training settings vary slightly by dataset, and Table 4 lists key hyperparameters for GSM8K and AI-MO.

A.2 Training Hyperparameters

  • GRPO hyperparameters were selected to support stable convergence while fitting the available computing resources.Table 4 lists the key settings for GSM8K and AI-MO.Source: §A.2 · p. 14, Table 4 · p. 14
  • Table 5 provides links to the language models and datasets used.Source: §A.2 · p. 14

B Training Collapse Examples

The section examines reward mode collapse during tiny-model training and compares how reward schemes affect accuracy. Memory-R and Memory-R+ accuracy rewards increase steadily, unlike R1 and cosine rewards.

B.1 Reward Mode Collapse

  • The reward-mode analysis tracks correctness, integer-format, and XML-format rewards optimized simultaneously.Integer reward checks for an integer output, while XML reward checks for correctly formatted tags for parsing.Source: §B.1 · p. 14
  • Figure 5 presents further examples of collapse with more training on GSM8K and AI-MO.The figure concerns Qwen2.5-0.5B-Instruct; its results are smoothed.Source: Figure 5 · p. 15
  • Memory-R and Memory-R+ accuracy rewards increase steadily, while R1 and cosine rewards do not.The latter schemes instead lead tiny LLMs to learn reward patterns more easily.Source: §B.1 · p. 15

B.2 Collapse in Response Length

Response-length collapse varies by reward scheme and dataset. Memory-R+ is the only method described as overcoming training collapse across the examples discussed.

B.2 Collapse in Response Length

  • Cosine reaches the maximum token limit early, while R1 shortens responses and optimizes integer reward; both yield low correctness rewards.Integer and XML rewards remain close to zero in the reported Qwen2.5-0.5B-Instruct run.Source: §B.2 · p. 15
  • Memory-R+ is the only method described as overcoming training collapse across GSM8K and AI-MO examples.R1 converges more slowly and underperforms Memory-R+ on AI-MO.Source: §B.2 · p. 15
  • Cosine lengthens responses severely on GSM8K but shortens them on AI-MO, whereas R1 shortens responses on GSM8K.R1 does not exhibit collapse on AI-MO in these examples.Source: §B.2 · p. 15

B.3 Details on Model Outputs

The section reports response diversity and self-verification observations for Memory-R+, alongside examples of its outputs and tracked intrinsic rewards. Its diversity evaluation compares sampled responses using lexical and semantic similarity measures.

B.3 Details on Model Outputs

  • Figure 6 reports Memory-R+’s exploitation and exploration intrinsic rewards over training steps on GSM8K with Qwen2.5-0.5B-Instruct.The figure’s results are smoothed.Source: §B.3 · p. 15, Figure 6 · p. 16
  • Memory-R+ shows a clear improvement in response diversity over R1.The diversity scores for different baselines are reported in Table 6.Source: Table 6 · p. 16, §B.3 · p. 16
  • Diversity is measured as 1 minus average pairwise cosine similarity across three sampled responses per question, using TF-IDF and Sentence Transformer embeddings.The analysis repeats this for 100 questions; TF-IDF captures lexical overlap and Sentence Transformer captures semantic relationships.Source: §B.3 · pp. 15–16
  • The paper states that examples of its method’s outputs showcasing self-verification behaviors appear at the end of the paper.Source: §B.3 · p. 16

Related papers

Loading 2504.02273v1…