Token-level weighting improves preference optimization in language models
GAW-PO: Preference Optimization with Gradient-Aligned Token Weights
Computation and Language
Summary
When teaching language models which responses are better, it’s tricky to know which parts of an answer to reward or punish. Usually, if a whole answer is disliked, every word in it gets penalized—even those that might be helpful. The authors suggest a new method that looks at each word’s effect more carefully, penalizing words only if they really conflict with the preferred behavior. This approach helps the model learn better and handle stronger training signals more reliably across tasks like math, coding, and question answering.
What this means in practice
- •For machine learning engineers: Improve language model training by applying token-level preference signals to achieve higher performance and robustness across reasoning and coding benchmarks.
- •For ai system integrators: Enhance response quality in automated question answering by using gradient-aligned token weighting for finer control of output preferences.
Authors
Andreea Dutulescu, Stefan Ruseti, Mihai Masala, Traian Rebedea, Mihai Dascalu
Abstract
Most preference optimization methods, such as Direct Preference Optimization (DPO), apply preference supervision at the response level, although autoregressive language models are optimized token by token. As a result, all tokens in a rejected response contribute to the negative training signal, including tokens that may encode behavior that is useful for the preferred response. We introduce GAW-PO, a gradient-aligned token reweighting method for DPO that estimates, for each rejected token, whether penalizing it would interfere with the preferred update directions. Tokens whose gradients are strongly aligned with the preferred behavior receive a weaker negative contribution, while conflicting tokens retain a stronger penalty. Our method achieves the highest average performance among the evaluated preference-optimization methods, improving by 0.97 points over standard DPO and 0.65 points over the strongest competing baseline across 11 benchmarks spanning mathematics, reasoning, coding, and question answering. We further show that gradient-aligned weighting is substantially more robust to aggressive preference optimization: as the DPO regularization parameter $β$ decreases, standard DPO degrades sharply, whereas GAW-PO continues to improve. These results suggest that accounting for the interaction between rejected-token updates and preferred behavior provides an effective form of token-level credit assignment for preference optimization.