Learning from the Right Rollouts: Data Attribution for PPO-based LLM Post-Training

Dong Shu; Denghui Zhang; Jessica Hullman

cs.LGApr 2, 2026

Learning from the Right Rollouts: Data Attribution for PPO-based LLM Post-Training

Authors: Dong Shu, Denghui Zhang, Jessica Hullman

Organizations: Northwestern University · Stevens Institute of Technology

Abstract

Traditional RL algorithms like Proximal Policy Optimization (PPO) typically train on the entire rollout buffer, operating under the assumption that all generated episodes provide a beneficial optimization signal. However, these episodes frequently contain noisy or unfaithful reasoning, which can degrade model performance and slow down training. In this paper, we propose Influence-Guided PPO (I-PPO), a novel framework that integrates data attribution into the RL post-training loop. By calculating an influence score for each episode using a gradient-based approximation, I-PPO identifies and eliminates episodes that are anti-aligned with a validation gradient. Our experiments demonstrate that I-PPO consistently outperforms SFT and PPO baselines. We show that our filtering process acts as an intrinsic early stopping mechanism, accelerating training efficiency while effectively reducing unfaithful CoT reasoning.

Figures & tables

Figure 1: Overview of I-PPO Framework. (a) Traditional PPO uses the raw rollout buffer for policy updates. (b) I-PPO refines the rollout by calculating influence scores for each generated episode. Negative episodes are removed.

		SFT			GRPO			PPO			I-PPO
		MV $\uparrow$	EM $\uparrow$	PK $\uparrow$	MV $\uparrow$	EM $\uparrow$	PK $\uparrow$	MV $\uparrow$	EM $\uparrow$	PK $\uparrow$	MV $\uparrow$	EM $\uparrow$	PK $\uparrow$
Rho	GSM8K	42.52	39.26	71.48	44.32(±0.82)	39.68(±1.15)	71.86(±0.65)	48.63(±0.42)	40.26(±0.85)	72.65(±0.44)	51.93 ( $\pm 0.56$ )	46.10 ( $\pm 0.73$ )	74.83 ( $\pm 0.28$ )
	Olymp	2.82	2.04	12.59	2.67(±0.45)	2.98(±0.61)	15.39(±0.32)	4.41(±0.35)	4.03(±0.42)	20.01(±0.20)	5.82 ( $\pm 0.23$ )	7.27 ( $\pm 0.44$ )	20.47 ( $\pm 0.18$ )
	ECQA	33.19	41.13	92.78	43.72(±0.91)	47.23(±0.88)	69.82(±0.77)	50.47(±0.51)	54.00(±1.12)	77.69(±0.67)	56.76 ( $\pm 0.62$ )	64.45 ( $\pm 1.35$ )	80.58( $\pm 0.31$ )
Gemma	GSM8K	68.13	59.95	81.02	70.15(±0.58)	61.30(±1.20)	82.53(±0.49)	72.06(±0.88)	62.03(±1.95)	83.57(±0.82)	75.04 ( $\pm 0.55$ )	68.41 ( $\pm 1.28$ )	85.23 ( $\pm 0.26$ )
	Olymp	10.55	6.40	27.52	14.30(±0.74)	9.71(±0.55)	29.29(±0.41)	15.65(±0.63)	11.24(±0.61)	31.79(±0.25)	22.51 ( $\pm 0.31$ )	17.82 ( $\pm 0.52$ )	34.15 ( $\pm 0.29$ )

Table 1: Performance Comparison Across Different Models. We report Majority Vote (MV), Exact Match (EM), and Pass@K (PK). Each row, the highest score for each metric is highlighted (e.g., MV , EM , PK ). Values in parentheses is the standard error, calculated over five independent runs of the trained model. Full performance across all datasets and standard error for SFT baseline is in Appendix Table 2 .

Figure 2: Cost Analysis per Training Step. Comparisons are conducted using the Rho model on (a) GSM8K dataset and (b) MATH dataset. The total training duration is annotated next to the plot.

Figure 3: Ablation Study on the Rho-1B model. We compare the I-PPO framework with reweighting against five baselines: (1) an I-PPO variant without reweighting, (2) a random filtering strategy, (3) I-PPO using only 75% and (4) 50% of the available validation data, and (5) a Process Reward Model (PRM).

Figure 4: Distribution of Unfaithful Reasoning Patterns. The chart illustrates total episode counts for Group C (correct final answers) and Group NC (incorrect final answers). The upper section displays episodes with positive influence scores, while the lower section shows episodes with negative influence scores.

		Baseline			PPO			Ours
		MV $\uparrow$	EM $\uparrow$	PK $\uparrow$	MV $\uparrow$	EM $\uparrow$	PK $\uparrow$	MV $\uparrow$	EM $\uparrow$	PK $\uparrow$
Rho 1B	GSM8K	42.52	39.26	71.48	48.63	40.26	72.65	51.93 ( $\pm 0.56$ )	46.10 ( $\pm 0.73$ )	74.83 ( $\pm 0.28$ )
	CollegeMath	17.03	10.46	29.66	27.46	21.85	45.29	31.25 ( $\pm 0.61$ )	28.17 ( $\pm 0.98$ )	47.89 ( $\pm 0.36$ )
	MATH	12.20	8.41	41.00	20.40	15.33	47.62	23.80 ( $\pm 0.57$ )	19.01 ( $\pm 0.71$ )	49.14 ( $\pm 0.22$ )
	OlympiadBench	2.82	2.04	12.59	4.41	4.03	20.01	5.82 ( $\pm 0.23$ )	7.27 ( $\pm 0.44$ )	20.47 ( $\pm 0.18$ )
	ECQA	33.19	41.13	92.78	50.47	54.00	77.69	56.76 ( $\pm 0.62$ )	64.45 ( $\pm 1.35$ )	80.58( $\pm 0.31$ )

Table 2: Performance Comparison Across Different Models. We report Majority Vote (MV), Exact Match (EM), and Pass@K (PK). For each row, the highest score for each metric is highlighted (e.g., MV , EM , PK ). Values in parentheses indicate the standard error, calculated over five independent runs of the trained model.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

		Math				Science				Social
		Group C			Group NC	Group C			Group NC	Group C			Group NC
Model		FP	NS	SC	NF	FP	NS	SC	NF	FP	NS	SC	NF
Rho	N	9.10%	1.67%	30.43%	3.25%	3.23%	2.86%	36.26%	1.58%	–	2.66%	–	–
Rho	P	3.61%	0.00%	18.79%	1.05%	0.64%	0.00%	24.20%	0.26%	–	0.00%	–	–
Gemma	N	3.53%	0.49%	28.10%	0.73%	1.77%	0.91%	32.98%	0.62%	–	0.47%	–	–
Gemma	P	1.84%	0.00%	15.34%	0.14%	0.31%	0.00%	21.13%	0.11%	–	0.00%	–	–

Appendix

Table 3: Distribution of Unfaithful Patterns in Negative (N) and Positive (P) Influence Episodes across Math, Science, and Social Domains. Group C is subdivided into False Positive (FP), Nonsensical (NS), and Reasoning Shortcuts (SC). Group NC exhibited Unfaithful (NF).

Figure 5: Distribution of Influence Scores Across Different Training Stages for the Rho-1B model on the GSM8K Dataset. The histograms illustrate the progressive shift from predominantly positive, widely distributed scores in early training (steps 0-100) to a narrower distribution dominated by zero or negative scores as the model converges (steps 300-382).

Figure 6: Distribution of Influence Scores Across Different Training Stages for the Rho-1B model on the MATH Dataset. The histograms illustrate the progressive shift from predominantly positive, widely distributed scores in early training (steps 0-200) to a narrower distribution dominated by zero or negative scores as the model converges (steps 600-817).

Figure 7: Architecture of the Batch Top-K Sparse Autoencoder (SAE). The encoder maps an input vector

\mathbf{z}

to a latent representation, where a Top-K activation enforces a sparse feature space

h(\mathbf{z})

. The decoder then uses these active features to reconstruct the original input as

\mathbf{\hat{z}}

Episode	ID	Concept
Positive	Feature 1:	This feature activates when the model gives a step-by-step arithmetic solution that ends with a standalone final numeric answer line formatted as ”#### [number]”.
	Feature 15:	This feature activates when the response introduces a single unknown (e.g., ”Let x be…”), formulates a linear equation by expressing successive transformations/transactions as algebraic terms, manipulates/combine terms step-by-step, and produces a numeric solution.
	Feature 10:	This feature activates when the model produces explicit, step-by-step arithmetic that chains multiplications (e.g., ”AB=C”, then ”CD=E”) to scale a per-unit quantity across units/time/frequency (unit conversions or rate × time × count).
	Feature 22:	This feature activates when the model produces an explicit step‑by‑step, line‑by‑line arithmetic solution — stringing together basic operations (multiply/divide/add/subtract, percentages, ratios) in procedural form and ending with a final numeric answer.
	Feature 17:	This feature activates when the model gives a procedural arithmetic solution: it breaks a word problem into explicit small steps (counting items, multiplying for totals, computing percent discounts/fees, then summing), listing intermediate calculations and a final numeric answer.
Negative	Feature 6:	This feature activates when the model gives short, confident numeric answers that contain elementary reasoning or arithmetic errors — e.g., misapplying percentages/ratios, omitting or double‑counting terms, confusing additive vs. multiplicative relations, or producing internally inconsistent calculations.

Appendix

Table 4: Top 5 Sparse Autoencoder Features Activated by Positive and Negative Influence Episodes. The table details the conceptual meaning behind each latent feature.

Figure 8: Extended Ablation Study. We compare the performance of the I-PPO framework using influence score reweighting against a baseline variant without reweighting. Results are shown for Gemma-2 2B, Qwen2.5 3B, Phi-3 4B, and LLaMA-3 8B models evaluated on Majority Vote, Exact Match, and Pass@K metrics across five reasoning datasets.

RL Post-Training

Reinforcement learning post-training adapts pretrained models with reward or feedback signals to improve reasoning, alignment, and interactive behavior.

See all 422 papers

Towards Better Training Signal: Advantage Clipped Policy Optimization
Sep 29, 2026Ruichuan Huang, Jinghan Liu, Congliang ChenPolicy Optimization RL for Language Model Reasoning
Shenzhen Loop Area Institute
Beyond Uniform Token-Level Trust Region in LLM Reinforcement Learning
Jun 9, 2026Renjie Mao, Xiangxin Zhou, Lvfang Tao +7RL for Language Model Reasoning Reinforcement Learning with Verifiable Rewards
Tencent Hunyuan
Cross-Epoch Adaptive Rollout Optimization for RL Post-Training
Jun 4, 2026Yiming Zong, Yige Wang, Jiashuo JiangSequential Decision Making RL for Language Model Reasoning
Department of Industrial Engineering & Decision Analytics, Hong Kong University of Science and Technology

Sep 29, 2026J/K move · Enter open · S save

Towards Better Training Signal: Advantage Clipped Policy Optimization

Ruichuan Huang, Jinghan Liu, Congliang Chen

Shenzhen Loop Area Institute

Policy Optimization RL for Language Model Reasoning RL Post-Training Off-Policy RL

Open paper AskPDF arXiv

Reinforcement learning (RL) has become a cornerstone for improving the reasoning capabilities of large language models (LLMs), but the need for on-policy data substantially limits training efficiency. Reusing off-policy data through importance sampling (IS) can improve efficiency but introduce considerable instability. Hence, algorithms such as PPO and GRPO widely adopt IS-ratio clipping to stabilize training. However, training stability and gradient estimate are mainly determined by the product of IS ratio and advantage. To further stabilize training, we propose ACPO, which clips the product of the IS ratio and the advantage, leading to more stable gradient estimates. We also establish a connection between ACPO and gradient clipping in policy mirror descent (PMD), which is a standard technique to stabilize optimization process, and prove the convergence of clipped-PMD under the standard RL setting. Experiments on widely used mathematical reasoning benchmarks show that ACPO consistently outperforms PPO and GRPO in both accuracy and training efficiency, delivering 4-6 percentage points gains on standard math benchmarks, with Qwen3-8B+PPO. Hence, ACPO is a practical and effective alternative to conventional IS-ratio clipping for RL post-training of LLMs.

Abstract

RL Post-Training

Explore similar work

Towards Better Training Signal: Advantage Clipped Policy Optimization

Beyond Uniform Token-Level Trust Region in LLM Reinforcement Learning

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training