MemReward: Graph-Based Experience Memory for LLM Reward Prediction with Limited Labels

Tianyang Luo; Tao Feng; Zhigang Hua; Yan Xie; Shuang Yang; Ge Liu; Jiaxuan You

cs.LGMar 13, 2026

MemReward: Graph-Based Experience Memory for LLM Reward Prediction with Limited Labels

Authors: Tianyang Luo, Tao Feng, Zhigang Hua, Yan Xie, Shuang Yang, Ge Liu, Jiaxuan You

Organizations: University of Illinois Urbana-Champaign · Meta

Abstract

Reinforcement learning has emerged as a powerful paradigm for improving large language model (LLM) reasoning, where rollouts are sampled from the policy and reward signals computed on those rollouts are used to update the policy. However, in data-scarce scenarios, obtaining ground-truth labels to verify rollouts at scale often requires expensive human annotation or labor-intensive expert verification. For instance, evaluating mathematical proofs demands expert review, and open-ended question answering lacks definitive ground truth. When ground-truth labels are scarce, the effectiveness of reinforcement learning fine-tuning is constrained. Inspired by the success of semi-supervised learning in propagating labels from labeled to unlabeled samples, we propose MemReward, a graph-based experience memory framework that integrates reward propagation directly into online policy optimization. MemReward stores rollouts (thinking processes and final answers) from an initial LLM policy as nodes in a heterogeneous graph connected by similarity and structural edges, over which a GNN propagates rewards from labeled to unlabeled rollouts. To train such a framework, we first warm up the GNN on labeled rollouts to predict rewards via heterogeneous aggregation over query, thinking, and answer nodes. During online RL fine-tuning, unlabeled rollouts are attached to the graph by query similarity, and the GNN predicts their rewards, yielding a hybrid reward acquisition strategy that combines ground-truth and GNN-predicted rewards. Experiments on Qwen2.5-1.5B and 3B in mathematics, question answering, and code generation demonstrate that MemReward, with ground-truth rewards on only 20% of rollouts, achieves 96.6% of Oracle performance on 1.5B and 97.3% on 3B, and closely approaches Oracle on out-of-domain tasks.

Figures & tables

Figure 1: MemReward approaches Oracle performance with only 20% labels. Using the same 20% ground-truth labels, MemReward (purple) substantially outperforms partial labels (R1-p, gray), closely approaching fully-supervised Oracle performance (green) on both in-domain and out-of-domain tasks across both model scales.

Figure 2: Overview of MemReward. Rollouts generated by the initial policy are stored as experience memory and organized into a heterogeneous graph for reward prediction. (Left) Warmup Phase : We construct a heterogeneous graph from labeled rollouts, where query nodes connect via embedding similarity, and each query links to its thinking and answer nodes. A GNN is trained to predict rewards through relational message passing. (Right) Online Phase : During GRPO training, labeled rollouts receive ground-truth rewards while unlabeled rollouts connect to the warmup graph via top-

k

similarity edges and obtain GNN-predicted rewards.

Qwen2.5-1.5B-Instruct
Method	Math			QA					Code		Avg	$\Delta\mathrm{Score}$
Method	GSM8K	GSM-sym	MATH	MMLU	CSQA	OBQA	ARC-C	GPQA	HumanEval+	MBPP+	Avg	$\Delta\mathrm{Score}$
R1-p	77.11	62.89	44.44	53.33	70.22	68.67	71.56	20.00	38.46	55.00	62.72	-7.75
LLM-RM	78.22	64.00	45.35	54.44	70.34	68.79	72.22	18.33	41.03	55.00	63.29	-7.18
kNN	79.33	65.11	46.49	55.33	71.91	70.62	73.74	20.00	41.03	52.50	64.49	-5.98
MemReward (Ours)	88.67	77.78	50.89	54.67	72.44	70.00	72.67	23.33	43.59	55.00	68.10	-2.37

Table 1: MemReward reaches 97.3% of Oracle performance on Qwen2.5-3B and 96.6% on Qwen2.5-1.5B across 10 benchmarks. We compare reward acquisition strategies under 20% label budget. On 3B, MemReward outperforms R1-p by +1.35, LLM-RM by +1.30, and kNN by +1.05.

Training Stage	Steps	Overall Accuracy
Early	1–136	74.9%
Middle	137–273	76.1%
Ending	274–410	77.3%

Table 2: GNN prediction accuracy across GRPO training stages (Qwen2.5-3B). The frozen GNN remains reliable throughout 410 training steps.

Figure 3: Ablation studies on (a) Qwen2.5-3B and (b) Qwen2.5-1.5B show each architectural component contributes to performance. The full model consistently outperforms all ablated variants on both scales across all three task categories.

Method	NuminaMath	SIQA	PIQA	Avg	$\Delta\mathrm{Score}$
Qwen2.5-1.5B-Instruct
R1-p	31.56	72.67	72.22	58.81	-4.75
MemReward	34.67	74.44	79.33	62.81	-0.75
R1-Oracle	33.78	75.78	81.11	63.56	0
Qwen2.5-3B-Instruct
R1-p	36.44	74.67	82.22	64.44	-2.75

Table 3: Out-of-domain evaluation. MemReward closely approaches R1-Oracle on out-of-domain tasks using only 20% ground-truth labels, narrowing the gap to just 0.23 on 3B and 0.75 on 1.5B.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

GT Ratio	Math			QA					Code		Avg	$\Delta\mathrm{Score}$
GT Ratio	GSM8K	GSM-sym	MATH	MMLU	CSQA	OBQA	ARC-C	GPQA	HumanEval+	MBPP+	Avg	$\Delta\mathrm{Score}$
20%	92.89	86.44	61.11	72.00	74.44	81.78	80.44	30.00	61.54	63.75	77.02	-2.10
30%	93.56	90.67	61.78	69.56	78.22	79.56	81.11	25.00	56.41	62.50	77.56	-1.56
40%	92.67	91.78	57.78	72.44	76.22	81.56	83.11	23.33	58.97	55.00	77.53	-1.59
50%	92.44	92.89	63.56	69.11	76.22	80.44	82.22	20.00	64.10	65.00	77.95	-1.17
60%	93.11	92.44	56.89	70.00	78.44	82.44	82.89	28.33	66.67	71.25	78.19	-0.93

Appendix

Table 4: MemReward maintains near-Oracle performance across ground-truth label ratios on Qwen2.5-3B. We vary the proportion of ground-truth labels from 20% to 70%, with the remaining queries receiving GNN-predicted rewards.

Hyperparameter	Value
Actor learning rate	$1e{-}6$
train_batch_size	128
mini_batch_size	128
micro_batch_size	32
Training steps	410
Max response length	1024

Appendix

Table 5: GRPO training hyperparameters. All RL fine-tuning experiments use these settings with Qwen2.5-1.5B-Instruct and Qwen2.5-3B-Instruct as the backbone models.

Hyperparameter	Value
Input dimension	384
Hidden dimension	512
Number of GNN layers	2
Number of attention heads	4
Dropout	0.1
Learning rate	$1e{-}3$

Appendix

Table 6: GNN hyperparameters. Configuration of the heterogeneous graph neural network used for reward prediction.

Split	Math			QA					Code
Split	GSM8K	GSM-sym	MATH	MMLU	CSQA	OBQA	ARC-C	GPQA	HumanEval+	MBPP+
R1-Oracle	750	750	750	750	750	750	750	99	65	132
MemReward	750	750	750	750	750	750	750	99	65	132
R1-p	150	150	150	150	150	150	150	19	13	26
Validation	300	300	300	300	300	300	300	39	26	52
Testing	450	450	450	450	450	450	450	60	39	80

Appendix

Table 7: Dataset statistics for in-domain evaluation. For each dataset, we randomly sample 1,500 instances and split them into training, validation, and testing sets. R1-Oracle uses all training queries with 100% ground-truth labels; MemReward uses all training queries with 20% GT labels and 80% GNN-predicted rewards; R1-p uses only the 20% labeled rollouts and discards the rest. Validation labels are available in all settings for model selection. Note that GPQA, HumanEval+, and MBPP+ have fewer than 1,500 available instances due to their original dataset sizes, so we use all available data for these benchmarks.

Dataset	Domain	Acc.	Prec.	Recall	ROC-AUC
GSM8K	Math	0.873	0.843	0.789	0.946
GSM-Sym	Math	0.883	0.725	0.669	0.936
MATH	Math	0.890	0.672	0.744	0.936
MMLU	QA	0.858	0.642	0.531	0.899
CSQA	QA	0.838	0.667	0.724	0.892
OBQA	QA	0.843	0.658	0.752	0.896

Appendix

Table 8: Per-dataset GNN reward prediction metrics on the validation set (Qwen2.5-3B, 20% GT). The GNN achieves strong prediction quality across domains, with math achieving the highest ROC-AUC.

System Prompt : You are a helpful assistant skilled in mathematics, logical reasoning, and programming. Solve problems step by step, showing your work clearly. You MUST end your response with ‘#### ’ followed by your final answer.
User Prompt : Solve this multiple-choice question:

The speed of a sound wave varies as it travels through different substances. Which factor will most affect the speed of a sound wave?
A. the frequency of the wave
B. the wavelength of the wave
C. the source that created the sound

Appendix

Table 9: Case study in QA (Science).

System Prompt : You are a helpful assistant skilled in mathematics, logical reasoning, and programming. Solve problems step by step, showing your work clearly. You MUST end your response with ‘#### ’ followed by your final answer.
User Prompt : Solve this math problem:

Leila spent $40 which is 1/4 of her money on a sweater. She was then left with$ 20 after spending the rest of her money on buying pieces of jewelry. How much more did Leila spend on jewelry than on the sweater?

Think step by step, showing your work clearly. Write your final answer as: #### [answer]
Ground Truth : 60

Appendix

Table 10: Case study in Math (GSM8K).

System Prompt : You are an expert Python programmer. Solve coding problems step by step. You MUST end your response with ‘#### ’ followed by your complete code in a “‘python code block.
User Prompt : Task: Given an array representing a branch of a tree that has non-negative integer nodes, pluck the node with the smallest even value. If multiple nodes have the same smallest even value, return the one with the smallest index. Return [smallest_value, its index], or [] if no even values exist.

Example: pluck([5, 0, 3, 0, 4, 2]) == [0, 1]

Think through your approach, then write the Python function. REMEMBER: You MUST end with #### followed by your code in a “‘python block.
Ground Truth :

Appendix

Table 11: Case study in Coding (HumanEval+).

Question : The Moon and Earth have a similar A. surface atmosphere B. rocky terrain C. gravity D. size

Ground Truth : B

R1-p (Answer: D, Incorrect ): “To determine the correct answer, let’s analyze each option: A. Surface atmosphere: The Moon does not have a significant atmosphere. It has a very thin exosphere, which is much less dense than Earth’s atmosphere. Earth has a thick atmosphere composed of nitrogen, oxygen, and other gases. Therefore, the Moon and Earth do not have a similar surface atmosphere. B. Rocky terrain: The Moon has a rocky surface, but it is not as diverse or as varied as Earth’s terrain. Earth has a diverse range of terrains including mountains, valleys, plains, and oceans. Therefore, the Moon and Earth do not have a similar rocky terrain. C. Gravity: The Moon’s gravity is about 1/6th of Earth’s gravity. Earth’s gravity is approximately 9.8 m/s 2 . The Moon’s gravity is about 1.63 m/s 2 . Therefore, the Moon and Earth do not have a similar gravity. D. Size: The Moon’s diameter is about 3,475 km. Earth’s diameter is about 12,742 km. The Moon is significantly smaller than Earth. Therefore, the Moon and Earth do not have a similar size. Given the analysis, none of the options A, B, C, or D describe a similarity between the Moon and Earth. However, the question asks for a similarity, and the only option that could be considered a similarity is the size, but it is not a similar size. Since the question asks for a similarity and the only option that could be considered a similarity is the size, but it is not a similar size, the correct answer is D.” The model dismisses the correct option B by conflating terrain diversity with terrain composition, then selects D despite the Moon being roughly 1/4 Earth’s diameter.

MemReward (Answer: B, Correct ): “To determine the correct answer, let’s analyze each option: A. Surface atmosphere: The Moon does not have a significant atmosphere. It has a very thin exosphere, which is much less dense than Earth’s atmosphere. Earth, on the other hand, has a substantial atmosphere composed of nitrogen and oxygen. Therefore, the Moon and Earth do not have a similar surface atmosphere. B. Rocky terrain: Both the Moon and Earth have rocky surfaces. The Moon’s surface is covered with craters and basaltic plains, while Earth’s surface is composed of various types of rock, including mountains, valleys, and plains. This similarity in terrain suggests that they do have a similar rocky terrain. C. Gravity: The Moon’s gravity is about 1/6th of Earth’s gravity. This means that an object would weigh about 1/6th of its weight on Earth. Earth’s gravity is approximately 9.8 m/s 2 , while the Moon’s gravity is about 1.625 m/s 2 . Therefore, the Moon and Earth do not have similar gravity. D. Size: The Moon and Earth are quite different in size. The Moon is about 1/400th the size of Earth. Earth has a diameter of approximately 12,742 kilometers, while the Moon has a diameter of about 3,475 kilometers. Thus, the Moon and Earth do not have similar sizes. Based on the analysis, the correct answer is B. rocky terrain.” The model correctly identifies that both bodies share rocky surface composition.

Appendix

Table 12: Comparative case study in QA (ARC-C). R1-p incorrectly selects D (size) while MemReward correctly identifies B (rocky terrain). The R1-p model’s reasoning contains a critical error: it dismisses option B despite both the Moon and Earth having rocky surfaces.

Question : Maurice is getting

2 for every finished task. And for every 10 tasks finished, he receives a

6 bonus. How much money would Maurice make for finishing 30 tasks?

Ground Truth : 78

R1-p (Answer: 22, Incorrect ): “To determine how much money Maurice would make for finishing 30 tasks, we need to break down the problem into two parts: the base pay for the tasks and the bonus for completing multiples of 10 tasks. 1. Calculate the base pay: Maurice earns

2 for each task.

\text{Base pay}=2,\text{tasks}\times 2,\text{dollars/task}=4,\text{dollars}

. 2. Calculate the bonus: Maurice receives a

6 bonus for every 10 tasks completed.

\text{Number of bonuses}=30\,\text{tasks}\,/\,10\,\text{tasks/bonus}=3\,\text{bonuses}

\text{Total bonus}=3\,\text{bonuses}\times 6\,\text{dollars/bonus}=18\,\text{dollars}

. 3. Calculate the total earnings:

\text{Total earnings}=\text{Base pay}+\text{Total bonus}=4\,\text{dollars}+18\,\text{dollars}=22\,\text{dollars}

. Therefore, Maurice would make

\boxed{22}

dollars for finishing 30 tasks.” The model uses “2 tasks” instead of “30 tasks” for the base pay calculation, producing an incorrect result.

MemReward (Answer: 78, Correct ): “To determine how much money Maurice would make for finishing 30 tasks, we need to break down the problem into two parts: the money he earns from the tasks and the bonus he receives. First, let’s calculate the money Maurice earns from the tasks. Maurice earns

2 for each task. If he finishes 30 tasks, the total money he earns from the tasks is:

30,\text{tasks}\times 2,\text{dollars/task}=60,\text{dollars}

. Next, we need to calculate the bonus Maurice receives. Maurice receives a

6 bonus for every 10 tasks finished. To find out how many bonuses he gets for 30 tasks, we divide 30 by 10:

30\div 10=3

. So, Maurice receives 3 bonuses. Since each bonus is

6, the total bonus he receives is:

3,\text{bonuses}\times 6,\text{dollars/bonus}=18,\text{dollars}

. Now, we add the money earned from the tasks and the bonus to find the total amount of money Maurice makes:

60,\text{dollars}+18,\text{dollars}=78,\text{dollars}

. Therefore, the total amount of money Maurice would make for finishing 30 tasks is

\boxed{78}$ .” The model correctly computes both components of the payment.

Appendix

Table 13: Comparative case study in Math (GSM8K). R1-p computes the base pay as

2\times 2=4

instead of

30\times 2=60

, arriving at an incorrect answer of 22. MemReward correctly computes both the base pay (

60) and the bonus (

18) for a total of $78.

Query : Expand the product $(2x+3)(x+5)$ . Ground-truth answer : $2x^{2}+13x+15$
Top-7 Training Neighbors
Rank Dataset Cosine Sim Correct 1 math 0.817 8/8 2 math 0.720 6/8 3 math 0.519 8/8 4 math 0.518 7/8 5 math 0.495 0/8 6 math 0.465 5/8 7 math 0.456 8/8
GNN Predicted Scores vs. Ground Truth
Resp Score GT Pred Match 0 0.971 1 1 ✓ 1 0.976 1 1 ✓ 2 0.974 1 1 ✓ 3 0.956 1 1 ✓ 4 0.973 1 1 ✓ 5 0.776 1 1 ✓ 6 0.975 1 1 ✓ 7 0.971 1 1 ✓

Appendix

Table 14: GNN trace – Math (Type A: clear propagation). Query from MATH dataset; all 8 responses are correct. Neighbors are predominantly correct, and the GNN assigns uniformly high scores.

Query : Erin made 12 pints of juice. She drinks 3 cups of juice each day. How many days will Erin take to drink all of the juice she made? Ground-truth answer : C
Top-7 Training Neighbors
Rank Dataset Cosine Sim Correct 1 mmlu 0.607 5/8 2 mmlu 0.585 3/8 3 gsm8k 0.581 0/8 4 mmlu 0.571 3/8 5 gsm8k 0.570 4/8 6 mmlu 0.568 4/8 7 mmlu 0.563 0/8
GNN Predicted Scores vs. Ground Truth
Resp Score GT Pred Match 0 0.676 1 1 ✓ 1 0.578 1 1 ✓ 2 0.733 1 1 ✓ 3 0.771 1 1 ✓ 4 0.048 0 0 ✓ 5 0.716 1 1 ✓ 6 0.019 0 0 ✓ 7 0.009 0 0 ✓

Appendix

Table 15: GNN trace – QA (Type B: mixed-signal discrimination). Query from MMLU dataset; 5 of 8 responses are correct. Despite mixed neighbor labels (0/8 to 5/8), the GNN correctly separates correct from incorrect responses.

Top- $k$	NuminaMath	SIQA	PIQA	Avg
MLP ( $k$ =0)	37.78	73.56	80.22	63.85
$k$ =3	40.00	75.33	80.44	65.26
$k$ =7 (default)	42.22	76.89	81.78	66.96
$k$ =10	41.33	76.44	80.89	66.22

Appendix

Table 16: Top-

k

sensitivity on OOD benchmarks (Qwen2.5-3B). All

k

values outperform R1-p and MLP.

Method	Overall	Math	QA	Code
MLP	73.5%	77.0%	72.4%	64.8%
kNN ( $k$ =7)	77.5%	81.5%	76.2%	68.3%
MemReward (GNN)	83.8%	87.1%	82.2%	73.1%

Appendix

Table 17: GNN reward prediction accuracy on the test set (Qwen2.5-3B, 20% GT). MemReward consistently outperforms both parametric (MLP) and non-parametric (kNN) baselines across all domains.

Threshold	GSM8K	GSM-Sym	MATH	MMLU	CSQA	OBQA	ARC-C	GPQA	HE+	MBPP+	Avg
0.3	92.44	85.11	58.50	70.89	73.93	80.64	79.80	28.33	58.97	62.50	75.85
0.5 (default)	92.89	86.44	61.11	72.00	74.44	81.78	80.44	30.00	61.54	63.75	77.02
0.8	92.67	85.78	59.86	71.33	74.61	81.09	80.30	28.33	61.54	63.75	76.49

Appendix

Table 18: Reward threshold sensitivity (Qwen2.5-3B, 10 benchmarks).

RL for Language Model Reasoning

Reinforcement learning for language model reasoning uses reward-driven training to improve multi-step problem solving, tool use, and reliable answers.

See all 1,094 papers

Reinforcement Learning-based Semi-supervised Knowledge Distillation with LLM-as-a-Judge
Apr 3, 2026Yiyang Shen, Lifu Tu, Weiran WangRL for Language Model Reasoning Knowledge Distillation
University of Iowa · Oracle AI
ExpRL: Exploratory RL for LLM Mid-Training
Jun 15, 2026Violet Xiang, Amrith Setlur, Chase Blagden +2Reinforcement Learning RL for Language Model Reasoning
1Stanford University · 2Carnegie Mellon University · 3OpenAI +1
Experience is the Best Teacher: Motivating Effective Exploration in Reinforcement Learning for LLMs
Mar 20, 2026Wenjian Zhang, Kongcheng Zhang, Jiaxin Qi +2Hindsight Experience Replay RL for Language Model Reasoning
Dalian University of Technology · Zhejiang University · Chinese Academy of Sciences

Apr 3, 2026J/K move · Enter open · S save

Reinforcement Learning-based Semi-supervised Knowledge Distillation with LLM-as-a-Judge

Yiyang Shen, Lifu Tu, Weiran Wang

University of Iowa · Oracle AI

RL for Language Model Reasoning Knowledge Distillation Numerical Reasoning in Language Models LLM-as-a-Judge Reinforcement Learning with Verifiable Rewards

Open paper AskPDF arXiv

Reinforcement Learning (RL) substantially improves the reasoning capabilities of language models, but most existing RL fine-tuning approaches rely entirely on ground-truth verifiable rewards and thus labeled datasets with verifiable answers. To overcome this, we propose a RL framework for reasoning distillation that leverages continuous, LLM-based rewards. Our method employs an efficient mechanism that computes a continuous CoT reward (CCR) directly from a single-token logit of a judge LLM, evaluating the student model's reasoning trajectory. This formulation provides an effective and scalable online training signal that can be applied to large volumes of unlabeled data. We demonstrate that, when paired with a strong judge, simply using CCR achieves performance comparable to that of ground-truth or pseudo-label verifiable rewards, and even surpasses them as the amount of unlabeled data increases. Furthermore, we find that combining them in a semi-supervised setup is highly synergistic: verifiable rewards help stabilize the CCR, while CCR improves the generalizability of verifiable rewards to related tasks. We also provide a comprehensive empirical comparison of various reward sources across multiple model architectures and dataset sizes. Our results show that this semi-supervised approach consistently enhances mathematical reasoning, yielding an absolute improvement of 5-10% on multiple reasoning tasks.

Abstract

RL for Language Model Reasoning

Explore similar work

Reinforcement Learning-based Semi-supervised Knowledge Distillation with LLM-as-a-Judge

ExpRL: Exploratory RL for LLM Mid-Training

Experience is the Best Teacher: Motivating Effective Exploration in Reinforcement Learning for LLMs