Research map / Large language models / Reasoning and chain of thought

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

NeurIPS 2026 · Large language models › Reasoning and chain of thought · approach: RL with verifiable rewards for reasoning

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Memory-R+, a kNN episodic-memory intrinsic reward (exploit and explore terms) added to GRPO to train tiny (≤1B) LLMs on math CoT reasoning, with reduced training collapse.

Problem

RL fine-tuning of tiny LLMs (≤1B parameters) for chain-of-thought reasoning suffers from sparse outcome rewards, ineffective exploration, and training collapse when mixed with format rewards.

Approach

Two episodic memories (successful and failed responses) indexed by Sentence Transformer embeddings produce a kNN-based exploitation reward (negative distance to success centroid) and exploration reward (1 minus max cosine similarity to failures), both min-max normalized over a sliding window and added to the outcome reward in GRPO.

Key results

• On Llama3.2-1B, Memory-R+ reaches 40.7 best GSM8K accuracy versus 37.2 for R1 and 38.1 for Cosine.
• Memory-R+ with K=20 on AI-MO-trained Qwen2.5-0.5B outperforms R1 and Cosine by about 5% on GSM8K and 4% on MATH-500 (as stated in text).
• Memory-R+ avoided training collapse across all runs, seeds and settings, while R1 and Cosine collapsed in some cases.

Surprising findings

• Cosine length-based reward causes response length collapse (maximum-length generation or shortened responses) in several settings.
• Base Falcon3-1B outperforms R1 on GSM8K (32.9 vs 10.9 last checkpoint), indicating RL with format rewards can degrade tiny models.

Limitations

• AIME24 results are near zero for all methods, so gains on the hardest benchmark are not meaningful.
• Evaluation uses strict extractive-match format requirements, and the paper's gains on MATH-500 are modest (about 1-2 points in several settings).

Experimental setup

Datasets: GSM8K AI-MO NuminaMath-TIR

Benchmarks: GSM8K MATH-500 AIME24

Metrics: extractive match accuracy response diversity (1 - similarity)

Baselines: R1 Cosine reward Base model

Builds on: GRPO DeepSeek-R1 Open-R1

Closest papers by meaning

• RL of Thoughts: Navigating LLM Reasoning with Inference-time Reinforcement Learning (ICLR 2026 · similarity 0.91)
• Training Large Reasoning Models Efficiently via Progressive Thought Encoding (ICLR 2026 · similarity 0.91)
• ProRL: Prolonged Reinforcement Learning Expands Reasoning Boundaries in Large Language Models (NeurIPS 2025 · similarity 0.91)
• On the Interplay of Pre-Training, Mid-Training, and RL on Reasoning Language Models (ICML 2026 · similarity 0.91)
• Demystifying Long Chain-of-Thought Reasoning (ICML 2025 · similarity 0.91)
• Count Counts: Motivating Exploration in LLM Reasoning with Count-based Intrinsic Rewards (ICLR 2026 · similarity 0.91)