Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Owen-Shapley Policy Optimization (OSPO): A Principled RL Algorithm for Generative Search LLMs

Abhijnan Nath, Alireza Bagheri Garakani, Tianchen Zhou, Fan Yang, Nikhil Krishnaswamy

TL;DR

OSPO reframes credit assignment for generative search LLMs by distributing sequence-level rewards to semantically coherent token segments using Owen-Shapley values. By restricting coalitions to contiguous segments and mapping segment attributions to token-level advantages, OSPO achieves fine-grained credit without external value models and stabilizes training via potential-based reward shaping. Empirical results on ESCI and H&M show OSPO, particularly in its proportional form, delivering superior or competitive performance compared with online and offline baselines, while generalizing across retrievers and showing robustness to distribution shifts. The work highlights coalition structure as a key hyperparameter, with moderate-width, contiguous coalitions and sufficient sampling delivering the best generalization. Overall, OSPO offers a principled, efficient pathway to more interpretable and transferable RL-based optimization for retrieval-aligned generation in LLMs, with clear avenues for extension to multi-turn and multimodal settings.

Abstract

Large language models are increasingly trained via reinforcement learning for personalized recommendation tasks, but standard methods like GRPO rely on sparse, sequence-level rewards that create a credit assignment gap, obscuring which tokens drive success. This gap is especially problematic when models must infer latent user intent from under-specified language without ground truth labels, a reasoning pattern rarely seen during pretraining. We introduce Owen-Shapley Policy Optimization (OSPO), a framework that redistributes sequence-level advantages based on tokens' marginal contributions to outcomes. Unlike value-model-based methods requiring additional computation, OSPO employs potential-based reward shaping via Shapley-Owen attributions to assign segment-level credit while preserving the optimal policy, learning directly from task feedback without parametric value models. By forming coalitions of semantically coherent units (phrases describing product attributes or sentences capturing preferences), OSPO identifies which response parts drive performance. Experiments on Amazon ESCI and H&M Fashion datasets show consistent gains over baselines, with notable test-time robustness to out-of-distribution retrievers unseen during training.

Owen-Shapley Policy Optimization (OSPO): A Principled RL Algorithm for Generative Search LLMs

TL;DR

OSPO reframes credit assignment for generative search LLMs by distributing sequence-level rewards to semantically coherent token segments using Owen-Shapley values. By restricting coalitions to contiguous segments and mapping segment attributions to token-level advantages, OSPO achieves fine-grained credit without external value models and stabilizes training via potential-based reward shaping. Empirical results on ESCI and H&M show OSPO, particularly in its proportional form, delivering superior or competitive performance compared with online and offline baselines, while generalizing across retrievers and showing robustness to distribution shifts. The work highlights coalition structure as a key hyperparameter, with moderate-width, contiguous coalitions and sufficient sampling delivering the best generalization. Overall, OSPO offers a principled, efficient pathway to more interpretable and transferable RL-based optimization for retrieval-aligned generation in LLMs, with clear avenues for extension to multi-turn and multimodal settings.

Abstract

Large language models are increasingly trained via reinforcement learning for personalized recommendation tasks, but standard methods like GRPO rely on sparse, sequence-level rewards that create a credit assignment gap, obscuring which tokens drive success. This gap is especially problematic when models must infer latent user intent from under-specified language without ground truth labels, a reasoning pattern rarely seen during pretraining. We introduce Owen-Shapley Policy Optimization (OSPO), a framework that redistributes sequence-level advantages based on tokens' marginal contributions to outcomes. Unlike value-model-based methods requiring additional computation, OSPO employs potential-based reward shaping via Shapley-Owen attributions to assign segment-level credit while preserving the optimal policy, learning directly from task feedback without parametric value models. By forming coalitions of semantically coherent units (phrases describing product attributes or sentences capturing preferences), OSPO identifies which response parts drive performance. Experiments on Amazon ESCI and H&M Fashion datasets show consistent gains over baselines, with notable test-time robustness to out-of-distribution retrievers unseen during training.
Paper Structure (57 sections, 1 theorem, 21 equations, 7 figures, 11 tables, 1 algorithm)

This paper contains 57 sections, 1 theorem, 21 equations, 7 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Generative Recommendation Systems.
  4. Densifying Sparse Rewards.
  5. Problem Formulation and Background
  6. Owen-Shapley Policy Optimization: Segment-Level Credit Assignment
  7. From Segment Attributions to Token Advantages.
  8. Proportional vs. Rank-Based Redistribution.
  9. Rank-Based Owen Redistribution.
  10. Experimental Setup
  11. Expert Data Generation
  12. Baselines
  13. Evaluation Settings
  14. Results
  15. Comparison with Online RL and Offline Methods.
  16. ...and 42 more sections

Key Result

Lemma 1.1

Let $y^{(g)}$ be a generated sequence of length $T$ from a Monte Carlo rollout with reward $R^{(g)}$. Following the REINFORCE algorithm williams92kool2019buyshao2024deepseekmath of sampling rollouts without replacement with baseline subtraction, define the sequence-level advantage as: is the mean reward across $G$ sampled rollouts. Let $\{\phi_t^{(g)}\}_{t=1}^T$ denote Owen values satisfying the

Figures (7)

  • Figure 1: OSPO overview: fine-grained credit assignment via Owen-Shapley values. Standard value-model-free RL (e.g., GRPO shao2024deepseekmath) assigns uniform advantages (via a single terminal reward) to all tokens (grey bars), ignoring segment-level contributions. OSPO evaluates contiguous coalitions by querying a retriever (or a reward model) with partial sequences, computing each segment's marginal contribution to retrieval quality. For example, the table shows coalition sampling for segment $\mathcal{P}_3$ ("coats"): adding $\mathcal{P}_3$ to coalition $\{\mathcal{P}_1, \mathcal{P}_2\}$ increases the reward from 0.50 to 0.72, yielding marginal $\Delta_3 = +0.22$; averaging such marginals across all coalitions produces the Owen value $\phi_3^{\text{Owen}} = +0.26$, indicating that "coats" consistently improves retrieval by providing explicit category information. These Owen values are then used to redistribute advantages proportionally to impact (center-right, teal bars) or “causally”, focusing gradient updates on truly high-contribution segments. Training dynamics (bottom-right) show OSPO reaches target performance in 50% fewer steps than GRPO, demonstrating improved sample efficiency.
  • Figure 2: OSPO-Prop ablation results on varying coalition structure on ESCI product search task. w denotes the maximum coalition span ($w_{max}$) and p the number of coalitions sampled ($M$) for Owen value estimation in OSPO (\ref{['alg:ospo']}).
  • Figure 3: Left and middle: Chain-of-thought (CoT) lengths within <think> fields and refined query lengths within <answer> fields during RL training of OSPO variants and GRPO on the H&M product search task. Right: evaluation performance on 500 randomly sampled H&M test queries, measured every 200 training steps.
  • Figure 4: Full prompt used for training and evaluation on the ESCI product search dataset reddy2022shopping. Following lin2025rec, we simplify the format by directly requesting the refined query within <answer> fields instead of JSON-style <query> tags. Only the text within <answer> tags is used for Owen value computations in OSPO, while the <think> section supports intermediate CoT reasoning. The example query shown in purple is drawn from the dataset.
  • Figure 5: Full prompt used for training and evaluation on the H&M Fashion dataset. Unlike ESCI, this setup grounds the LLM’s instructions in the user’s purchase history for contextualized query refinement. Only the text within <answer> tags is used for Owen value computations in OSPO, while the <think> section supports intermediate CoT reasoning. The example query shown in purple is drawn from the dataset.
  • ...and 2 more figures

Theorems & Definitions (4)

  • Lemma 1.1: Length-Invariant Owen Redistribution
  • proof
  • Remark 1.2
  • Remark 1.3: Connection to the Efficiency Axiom