TLRS-DPO: Token-Level Risk Scoring for Better Direct Preference Optimization

Suxin Ji; Hungtao Wan

cs.LGNov 8, 2026

TLRS-DPO: Token-Level Risk Scoring for Better Direct Preference Optimization

Authors: Suxin Ji, Hungtao Wan

Organizations: University of Pennsylvania · Independent Researcher

Abstract

Direct Preference Optimization (DPO) is a simple and widely used alternative to RLHF for aligning large language models (LLMs), but its sequence-level objective cannot precisely locate or suppress the local patterns that make an otherwise acceptable response unsafe or undesirable. We propose TLRS-DPO (Token-Level Risk Scoring for Better Direct Preference Optimization), which augments DPO with fine-grained token-level control while remaining fully offline and reference-based. TLRS-DPO introduces a token-level robust compensation term that converts the policy-reference log-ratio on rejected responses into localized risk scores, a curriculum-style schedule that gradually increases the strength of this compensation during training, and a risk-adaptive screening mechanism that selectively penalizes only high-risk tokens instead of entire sequences. This design bridges sequence-level preference fitting and token-level behavior control, and can be implemented as a drop-in modification to standard DPO trainers. Empirically, TLRS-DPO improves alignment quality and safety-relevant behavior with minimal impact on helpfulness compared to strong DPO baselines.

Figures & tables

Figure 1: Conceptual comparison between DPO and TLRS-DPO. TLRS-DPO introduces token-level risk scoring that identifies and down-weights high-risk tokens in rejected responses, leading to finer-grained, more interpretable, and more stable alignment compared to standard DPO.

Figure 2: Overview of TLRS-DPO Framework. The framework integrates sequence-level DPO loss with token-level risk scoring and curriculum weighting. The upper path represents sequence-level preference optimization via the log-ratio contrast $\Delta_{\theta,\mathrm{ref}}$ , while the lower path estimates token-level risk and compensation $\mathcal{C}_{\mathrm{TLRS}}$ . Both pathways are combined in the final loss $\mathcal{L}_{\mathrm{TLRS\text{-}DPO}}$ , which is optimized through standard backpropagation within the policy $\pi_{\theta}$ .

Method	A-Eval	A-Eval (LC)	Arena	Arena (LC)	MT-Bench	Std.
DPO Rafailov et al. (2023)	62.8 $\uparrow$ 0.0	58.7 $\uparrow$ 0.0	51.2 $\uparrow$ 0.0	47.3 $\uparrow$ 0.0	6.72 $\uparrow$ 0.0	0.15 $\uparrow$ 0.29
SimPO Meng et al. (2024)	66.1 $\uparrow$ 3.3	61.9 $\uparrow$ 3.2	54.6 $\uparrow$ 3.4	51.4 $\uparrow$ 4.1	7.01 $\uparrow$ 0.29	0.14
SePO Yang et al. (2025a)	68.5 $\uparrow$ 5.7	65.7 $\uparrow$ 7.0	57.3 $\uparrow$ 6.1	53.8 $\uparrow$ 6.5	7.24 $\uparrow$ 0.52	0.12
TIS-DPO Liu et al. (2025a)	70.2 $\uparrow$ 7.4	67.4 $\uparrow$ 8.7	59.9 $\uparrow$ 8.7	56.7 $\uparrow$ 9.4	7.41 $\uparrow$ 0.69	0.11
T-REG Zhou et al. (2025)	71.0 $\uparrow$ 8.2	68.2 $\uparrow$ 9.5	61.3 $\uparrow$ 10.1	58.0 $\uparrow$ 10.7	7.48 $\uparrow$ 0.76	0.09
TLRS-DPO (ours)	73.4 $\uparrow$ 10.6	70.6 $\uparrow$ 11.9	64.1 $\uparrow$ 12.9	61.2 $\uparrow$ 13.9	7.73 $\uparrow$ 1.01	0.08

Table 1: Comparison of preference optimization methods across alignment benchmarks. Metrics are win rate (%) and length-controlled (LC) win rate (%) on AlpacaEval 2.0 (A-Eval) and Arena-Hard (Arena), and MT-Bench average score (0–10). Relative improvements over DPO are shown in red subscripts.

Method	Safety	Reward	GSM8K	MathInst.
	Rate (%)	(↑)	Acc@1 (%)	Acc@1 (%)
DPO	76.2	0.41	58.3	62.4
SePO	80.5	0.47	62.1	67.2
TIS-DPO	82.9	0.49	63.4	68.6
T-REG	83.6	0.50	64.1	69.3
TLRS-DPO (ours)	86.8	0.54	67.5	72.1

Table 2: Cross-domain generalization results. We report Llama-Guard Safety Rate (%) and Beaver-Reward (↑) for safety, and accuracy (Acc@1, %). All results are evaluated on the full GSM8K and MathInstruct benchmarks (Setting S1).

Figure 3: Performance under 20% noisy preference labels on AlpacaEval 2.0. TLRS-DPO exhibits the smallest win-rate degradation compared with baselines, indicating stronger robustness to noisy supervision.

Figure 4: Training efficiency comparison across four methods.

Figure 5: Visualization of token-level risk distributions on a GSM8K reasoning example.

Figure 6: Cross-task visualization of average token-level risk across five alignment and reasoning benchmarks.

Method Variant	GSM8K		PKU-SafeRLHF
Method Variant	Accuracy↑	Risk Score↓	Accuracy↑	Risk Score↓
Vanilla DPO	72.4	0.41	64.8	0.47
w/o Token-Level Scoring	75.1	0.36	66.5	0.43
w/o Risk Curriculum	76.3	0.32	68.4	0.39
w/o Reference Penalty	77.5	0.30	70.2	0.36
Full TLRS-DPO (ours)	80.7	0.25	73.1	0.31

Table 3: Ablation study on GSM8K and PKU-SafeRLHF datasets. For controlled component-wise analysis, GSM8K results are reported on a representative subset (Setting S2), while full-benchmark results are reported in Table 2 (Setting S1).

Figure 7: Token-level risk visualization across three representative cases.

Direct Preference Optimization

Direct Preference Optimization trains generative models from preferred and rejected outputs without fitting an explicit reward model.

See all 174 papers

Improving Safety Alignment via Balanced Direct Preference Optimization
Mar 24, 2026Shiji Zhao, Mengyang Wang, Shukun Xiong +7LLM Safety Alignment LLM Alignment
Institute of Artificial Intelligence, Beihang University, Beijing, China.
TokenRatio: Principled Token-Level Preference Optimization via Ratio Matching
May 12, 2026Truong Nguyen, Tien-Phat Nguyen, Linh Ngo Van +3Pairwise Preference Learning LLM Alignment
Hanoi University of Science and Technology · University of Stuttgart · German Research Center for Artificial Intelligence +3
Token-weighted Direct Preference Optimization with Attention
May 21, 2026Chengyu Huang, Zhuohang Li, Sheng-Yen Chou +1Pairwise Preference Learning LLM Alignment
Cornell University · Vanderbilt University

Mar 24, 2026J/K move · Enter open · S save

Improving Safety Alignment via Balanced Direct Preference Optimization

Shiji Zhao, Mengyang Wang, Shukun Xiong, Fangzhou Chen+6

Institute of Artificial Intelligence, Beihang University, Beijing, China.

LLM Safety Alignment LLM Alignment Direct Preference Optimization

Open paper AskPDF arXiv

With the rapid development and widespread application of Large Language Models (LLMs), their potential safety risks have attracted widespread attention. Reinforcement Learning from Human Feedback (RLHF) has been adopted to enhance the safety performance of LLMs. As a simple and effective alternative to RLHF, Direct Preference Optimization (DPO) is widely used for safety alignment. However, safety alignment still suffers from severe overfitting, which limits its actual performance. This paper revisits the overfitting phenomenon from the perspective of the model's comprehension of the training data. We find that the Imbalanced Preference Comprehension phenomenon exists between responses in preference pairs, which compromises the model's safety performance. To address this, we propose Balanced Direct Preference Optimization (B-DPO), which adaptively modulates optimization strength between preferred and dispreferred responses based on mutual information. A series of experimental results show that B-DPO can enhance the safety capability while maintaining the competitive general capabilities of LLMs on various mainstream benchmarks compared to state-of-the-art methods. \color{red}{Warning: This paper contains examples of harmful texts, and reader discretion is recommended.

Abstract

Direct Preference Optimization

Explore similar work

Improving Safety Alignment via Balanced Direct Preference Optimization

TokenRatio: Principled Token-Level Preference Optimization via Ratio Matching

Token-weighted Direct Preference Optimization with Attention