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

TL;DR

MemReward introduces a graph-based experience memory that propagates reward signals across similar rollouts, enabling effective reinforcement learning fine-tuning of large language models with limited labeled data.

Contribution

The paper presents a novel graph neural network framework that propagates rewards from labeled to unlabeled rollouts, reducing the need for extensive ground-truth labels in LLM reinforcement learning.

Findings

MemReward achieves over 96% of Oracle performance with only 20% labeled rollouts.

The framework effectively propagates rewards in mathematics, QA, and code generation tasks.

Approaches near Oracle performance on out-of-domain tasks.

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.