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

We construct an episodic memory module $\mathcal{M}$ to store past reasoning trajectories, facilitating efficient retrieval of relevant experiences. To ensure efficient reasoning retrieval, both queries and responses are encoded into a shared high-dimensional vector space using Enc, implemented as a pre-trained Sentence Transformer Reimers and Gurevych (2019):

$$\mathbf{q}_{i}=\text{Enc}(q_{i})\in\mathbb{R}^{d},\quad\mathbf{a}_{i,j}=\text{Enc}(a_{i,j})\in\mathbb{R}^{d}$$

(1)

Here, each entry in the memory consists of embeddings of a query $q_{i}$ and a set of associated responses $\{a_{i,j}\}$: $\mathcal{M}=\{(\mathbf{q}_{i},\{\mathbf{a}_{i,j}\}_{j=1}^{L})\}_{i=1}^{N}$ where $N$ is the maximum number of stored queries, and each query $q_{i}$ maintains at most $L$ associated responses. Memory retrieval is denoted as: $\mathcal{M}[q_{i}]=\{\mathbf{a}_{i,j}\}_{j=1}^{L}$.

During training, we sample $G$ responses from the LLM for a given query. New query-response pairs are incorporated into memory following an update rule: (1) If $q$ is a novel query (i.e., not present in $\mathcal{M}$) and $|\mathcal{M}|\geq 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 $q$ already exists in $\mathcal{M}$, the new response set $\{a_{j}\}$ is merged with the existing responses. If the total number 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, $\mathcal{M}^{s}$, and another for failed responses, $\mathcal{M}^{f}$. Given a query $q$ and a set of generated responses $\{a_{j}\}_{j=1}^{m}$, we update the memories as follows:

$$\mathcal{M}^{s}[q]\leftarrow\mathcal{M}^{s}[q]\cup\{\mathbf{a}_{j}\mid R(q,a_{j})>\tau^{s}\},\ \mathcal{M}^{f}[q]\leftarrow\mathcal{M}^{f}[q]\cup\{\mathbf{a}_{j}\mid R(q,a_{j})\leq\tau^{f}\}$$

(2)

where $\tau^{s}$ and $\tau^{f}$ are reward thresholds for classifying successful and failed responses, respectively. For instance, in mathematical problem-solving, where the outcome reward is defined as $R(q,a)=1$ for a correct final answer and $0$ otherwise, we can set the thresholds as $\tau^{s}=\tau^{f}=0.5$.

Given a new query $q$ and a response $a$, we compute their embeddings $\mathbf{q}=\text{Enc}(q)$, $\mathbf{a}=\text{Enc}(a)$ and retrieve the top-$k$ nearest queries from a memory $\mathcal{M}$ based on cosine similarity (CS): $\{{q}_{k}^{\prime}\}_{k=1}^{K}=\text{top-k}\left(\arg\max_{\textbf{q}^{\prime}\in\mathcal{M}}\text{CS}(\mathbf{q},\mathbf{q^{\prime}})\right)$. Then, the set of relevant responses from the memory is computed as:

$$B\left(\mathcal{M},q\right)=\bigcup_{k=1}^{K}\mathcal{M}[q_{k}^{\prime}]$$

(3)

where $K$ is the number of nearest neighbors considered in the memory retrieval. For simplicity, the same $K$ is used for both $\mathcal{M}^{s}$ and $\mathcal{M}^{f}$.

To reinforce successful reasoning patterns, we compute the exploitation reward using responses stored in the success memory $\mathcal{M}^{s}$. The model is rewarded for generating responses similar to those that previously led to correct final answers. We first compute the centroid of retrieved response embeddings from the success memory: $\mathbf{c}(\mathcal{M}^{s},q)=\frac{1}{|B\left(\mathcal{M}^{s},q\right)|}\sum_{\mathbf{a}_{j}\in B\left(\mathcal{M}^{s},q\right)}\mathbf{a}_{j}$. The Euclidean distance between the response $a$ and this centroid derives the exploit reward as:

$$R_{\text{exploit}}(q,a)=-\|\mathbf{a}-\mathbf{c}(\mathcal{M}^{s},q)\|$$

(4)

By measuring the distance to the centroid, we encourage the model to align with the general distribution of successful reasoning patterns rather than overfitting specific past answers. This provides a smoother optimization signal, overcoming the reward sparsity problem and capturing structural commonalities in effective reasoning paths.

To encourage novel reasoning paths, we compute the exploration reward using responses stored in the failure memory $\mathcal{M}^{f}$, ensuring that the model avoids repeating past mistakes. Specifically, novelty is measured as the minus of the cosine similarity between the generated response embedding $\mathbf{a}$ and its closest retrieved embedding from the failure memory:

$$R_{\text{explore}}(q,a)=1-\max_{\mathbf{a}_{j}\in B\left(\mathcal{M}^{f},q\right)}\text{CS}(\mathbf{a},\mathbf{a}_{j})$$

(5)

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: early in training, when most outputs are wrong, this encourages broad exploration—generating anything unlike previous attempts. As correct responses accumulate in the success memory $\mathcal{M}^{s}$, 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 a natural progression from broad to focused exploration. In practice, we can set a warm-up period (e.g., 50 training steps) to collect the initial data before applying the exploration reward.

The intrinsic rewards are normalized using a running min-max scaling to ensure they are evaluated relative to recent performance trends. For example, given a sliding window of past intrinsic rewards $\{R_{\text{explore},i}\}_{t-w}^{t}$, where $w$ is the window size, the normalized intrinsic reward is computed as:

$$\hat{R}_{\text{exploit/explore}}(q,a)=\frac{R_{\text{exploit/explore}}(q,a)-\min_{t-w\leq i\leq t}R_{\text{exploit/explore},i}}{\max_{t-w\leq i\leq t}R_{\text{exploit/explore},i}-\min_{t-w\leq i\leq t}R_{\text{exploit/explore},i}+\epsilon}$$

(6)

where $\epsilon$ is a small constant to prevent division by zero. Here, the rewards are interpreted relative to recent performance, allowing the model to adapt dynamically. We argue that since this is 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.

The final memory-based intrinsic reward $R_{mem}$ is computed as a weighted sum of the normalized components:

$$R_{mem}=\beta_{s}\hat{R}_{\text{exploit}}+\beta_{e}\hat{R}_{\text{explore}},$$

(7)

where $\hat{R}_{\text{exploit}}$ and $\hat{R}_{\text{explore}}$ are the normalized rewards, ensuring that the model evaluates improvements relative to its recent history. The weighting factors $\beta_{s}$ and $\beta_{e}$ determine the balance between reinforcing past successes and encouraging novel reasoning, providing explicit control over exploitation and exploration trade-offs. For simplicity in this paper, we do not tune them and set $\beta_{s}=\beta_{e}=1$ throughout experiments. With $R_{mem}$, dense performance-driven signals are incorporated into the training rewards. We hypothesize that this facilitates learning performance-based rewards, narrows the difficulty gap between performance-based and format-based learning, and mitigates training collapse issues (see more in Sec. 4.3).

We train the model using a reinforcement learning objective, where the task outcome reward $R$ and the intrinsic reward $R_{mem}$ are used to update the policy $\pi_{\theta}(a\mid q)$. Specifically, we adopt the Group Relative Policy Optimization (GRPO Shao et al. (2024)), a variant of policy gradient methods designed for improved stability and efficiency in language model RL fine-tuning. The training objective maximizes the expected total reward:

$$\max_{\theta}\,\mathbb{E}_{q\sim D,a\sim\pi_{\theta}}[R+R_{mem}]$$

(8)

where $D$ is the training dataset, and $\pi_{\theta}$ is the LLM with its trainable parameters $\theta$.