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

arXiv: https://arxiv.org/abs/2504.02273 · HTML: https://mlsanity.com/paper/2504.02273/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 26 of its 36 references that are in the map.
- First posted on arXiv 3 Apr 2025.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)

## Builds on (most cited first, 26 of 26)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [Do NOT Think That Much for 2+3=? On the Overthinking of o1-Like LLMs](https://mlsanity.com/paper/2412.21187/) (2024)
- [ReST-MCTS*: LLM Self-Training via Process Reward Guided Tree Search](https://mlsanity.com/paper/2406.03816/) (2024)
- [Demystifying Long Chain-of-Thought Reasoning in LLMs](https://mlsanity.com/paper/2502.03373/) (2025)
- [Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training](https://mlsanity.com/paper/2309.17179/) (2023)
- [O1-Pruner: Length-Harmonizing Fine-Tuning for O1-Like Reasoning Pruning](https://mlsanity.com/paper/2501.12570/) (2025)
- [Training Language Models to Reason Efficiently](https://mlsanity.com/paper/2502.04463/) (2025)
- [Neural Episodic Control](https://mlsanity.com/paper/1703.01988/) (2017)
- [LLM Post-Training: A Deep Dive into Reasoning Large Language Models](https://mlsanity.com/paper/2502.21321/) (2025)
- [Training Trajectories of Language Models Across Scales](https://mlsanity.com/paper/2212.09803/) (2022)
- [Can large language models explore in-context?](https://mlsanity.com/paper/2403.15371/) (2024)
- [Large Language Models Are Semi-Parametric Reinforcement Learning Agents](https://mlsanity.com/paper/2306.07929/) (2023)
- [The Diminishing Returns of Masked Language Models to Science](https://mlsanity.com/paper/2205.11342/) (2022)
- [Model-Based Episodic Memory Induces Dynamic Hybrid Controls](https://mlsanity.com/paper/2111.02104/) (2021)
- [Selective Self-to-Supervised Fine-Tuning for Generalization in Large Language Models](https://mlsanity.com/paper/2502.08130/) (2025)
- [Episodic Policy Gradient Training](https://mlsanity.com/paper/2112.01853/) (2021)
- [Large Language Models Prompting With Episodic Memory](https://mlsanity.com/paper/2408.07465/) (2024)

## Built on by (most cited first, 6 of 6)

- [Cognitive Foundations for Reasoning and Their Manifestation in LLMs](https://mlsanity.com/paper/2511.16660/) (2025)
- [Memory-Augmented Transformers: A Systematic Review from Neuroscience Principles to Enhanced Model Architectures](https://mlsanity.com/paper/2508.10824/) (2025)
- [UFO-RL: Uncertainty-Focused Optimization for Efficient Reinforcement Learning Data Selection](https://mlsanity.com/paper/2505.12457/) (2025)
- [Hear Both Sides: Efficient Multi-Agent Debate via Diversity-Aware Message Retention](https://mlsanity.com/paper/2603.20640/) (2026)
- [Uncertainty-Guided Checkpoint Selection for Reinforcement Finetuning of Large Language Models](https://mlsanity.com/paper/2511.09864/) (2025)
- [SenTSR-Bench: Thinking with Injected Knowledge for Time-Series Reasoning](https://mlsanity.com/paper/2602.19455/) (2026)

## Abstract (arXiv)

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.