Selective Preference Optimization via Token-Level Reward Function Estimation
Kailai Yang, Zhiwei Liu, Qianqian Xie, Jimin Huang, Erxue Min, Sophia Ananiadou
TL;DR
SePO reframes LLM alignment as selective, token-level optimization by first training a small DPO-based oracle to estimate a token-level reward function. The oracle scores all tokens on large-scale data, enabling selection of top-performing tokens from chosen responses and bottom-performing tokens from rejected ones, which then supervise the target policy via a contrastive objective. The approach achieves strong improvements across multiple benchmarks while reducing training costs, and further demonstrates strong potential for weak-to-strong generalization by leveraging weak or out-of-distribution supervision signals. Theoretical results link token-level rewards to optimal policies under DPO and establish that training on random subsets yields pessimistic but consistent reward estimates, supporting the practicality and efficiency of SePO. Overall, SePO offers a flexible, data-efficient pathway to high-quality preference alignment with broad applicability to existing alignment datasets.
Abstract
Recent advancements in large language model alignment leverage token-level supervisions to perform fine-grained preference optimization. However, existing token-level alignment methods either optimize on all available tokens, which can be noisy and inefficient, or perform selective training with complex and expensive key token selection strategies. In this work, we propose Selective Preference Optimization (SePO), a novel selective alignment strategy that centers on efficient key token selection. SePO proposes the first token selection method based on Direct Preference Optimization (DPO), which trains an oracle model to estimate a token-level reward function on the target data. This method applies to any existing alignment datasets with response-level annotations and enables cost-efficient token selection with small-scale oracle models and training data. The estimated reward function is then utilized to score all tokens within the target dataset, where only the key tokens are selected to supervise the target policy model with a reference model-free contrastive objective function. Extensive experiments on three public evaluation benchmarks show that SePO significantly outperforms competitive baseline methods by only optimizing 30% key tokens on the target dataset. SePO applications on weak-to-strong generalization show that weak oracle models effectively supervise strong policy models with up to 16.8x more parameters. SePO also effectively selects key tokens from out-of-distribution data to enhance strong policy models and alleviate the over-optimization problem.