Can TI‑DPO Fix DPO’s Blind Spot? Token‑Importance Guided Direct Preference Optimization for Better LLM Alignment

Research Background

Post‑training alignment methods face two core challenges: (1) sequence‑level binary supervision hides harmful tokens inside otherwise good responses, causing distribution shift; (2) token‑level importance estimates inherit the model’s inherent “U‑shaped” attention bias, over‑emphasizing start and end tokens while neglecting central semantic content.

Core Mechanism of TI‑DPO

Hybrid Weighting

TI‑DPO computes a gradient‑attribution weight for each token by taking the norm of the loss gradient with respect to the token embedding, assigning higher weight to tokens that contribute more to the output. A Gaussian prior centered on the middle of the sequence counteracts the U‑shaped bias by encouraging focus on middle tokens where semantic cores often reside. The final token weight is a convex combination of the gradient‑attribution signal and the Gaussian prior.

Triplet Loss

TI‑DPO replaces binary contrast with a metric‑learning triplet loss. During training three roles are constructed: Anchor (the model’s current response), Positive (a high‑quality human‑preferred answer), and Negative (a low‑quality rejected answer). The loss pulls the anchor toward the positive while pushing it away from the negative, forming a structured geometric objective in semantic space.

Experimental Results

Overall Capability Assessment

On the Llama‑3.1‑8B‑Instruct base, TI‑DPO achieves an average score of 62.3, surpassing GRPO (62.1) and DPO (60.8).

Fine‑Grained Task Performance

On three detail‑sensitive benchmarks—IFEval (instruction following), TruthfulQA (truthfulness), and HumanEval (code generation)—TI‑DPO markedly outperforms DPO, SimPO, and GRPO.

Ablation Study

Removing any of the three core components—hybrid weighting, Gaussian prior, or triplet loss—causes a noticeable drop across general ability, mathematical reasoning, and code generation metrics, confirming that each component is essential.

Case Study: Medical Consultation

A token‑weight heatmap for the query “What should I do for a headache?” shows that TI‑DPO assigns high weight to safety‑critical tokens such as “seek medical attention” and “promptly” in the preferred response, while penalizing risky advice like “painkillers casually” in the non‑preferred response. An intermediate response illustrates the model’s self‑assessment before alignment.

Conclusion

TI‑DPO shifts large‑model alignment from coarse sequence‑level optimization to fine‑grained token‑level control, explicitly modeling each token’s contribution to value alignment. Experiments confirm stable improvements on instruction following, truthfulness, and code generation, validating that finer granularity in data utilization is an effective path for enhancing model capabilities.

Paper: https://arxiv.org/abs/2505.19653

Code: https://github.com/gracefulning/TIDPO