GAW-PO: Preference Optimization with Gradient-Aligned Token Weights

1 INTRODUCTION

• DPO applies response-level supervision although autoregressive models operate token by token, so every rejected-response token contributes to the optimization signal.Only a small subset of rejected-response tokens may be responsible for the response being judged lower quality.Source: §1 · p. 1
• Rejected responses can contain useful facts, calculations, formatting, or linguistic patterns, yet reducing their likelihood can suppress those components.The relative-preference objective also does not require the chosen response’s absolute likelihood to increase; prior work observed it can decrease while the preference objective improves (Pal et al., 2024; Pang et al., 2024).Source: §1 · p. 1
• Prior token-level objectives and weighting schemes estimate token importance, but do not fully characterize the consequences of penalizing a token during optimization.The paper instead considers interactions between the parameter updates induced by rejected tokens and preferred responses.Source: §1 · p. 2
• GAW-PO downweights rejected tokens whose penalization would interfere with behavior supported by preferred responses, while retaining stronger penalties for conflicting tokens.The method compares rejected-token updates with preferred update directions rather than weighting tokens only by probability, loss, or gradient magnitude.Source: §1 · p. 2
• GAW-PO improves average downstream performance over standard DPO and competing methods, and remains stable under aggressive optimization regimes where standard DPO degrades sharply.The introduction also notes that the implementation, gradient-weight computation pipeline, training recipes, filtered preference data, and evaluation setup are open-sourced.Source: §1 · p. 2

Preference optimization

• Response-level variants change the loss, feedback format, or reference-model use, while the likelihood of chosen responses can still decrease during training.Prior responses to that side effect add a chosen-response likelihood term or filter pairs with similar representations; GAW-PO keeps every pair and changes rejected-token penalties (Pal et al., 2024; Pang et al., 2024; Razin et al., 2025).Source: §2 · pp. 2–3

Token- and step-level credit assignment

• Token- and step-level methods use different signals: TDPO weights tokens equally, TIS-DPO uses contrastive-model probability gaps, and Step-DPO locates the first incorrect reasoning step.Step-DPO requires locating the error and resampling continuations, while TI-DPO weights tokens by input-embedding gradient norm rather than alignment with the chosen response (Zeng et al., 2024; Liu et al., 2025; Lai et al., 2024; Yang et al., 2026).Source: §2 · p. 3
• AAO lowers weights for rejected tokens matching tokens in the chosen response, but its input-embedding comparison ignores context and can spare an incorrect operand.In the example, the rejected response’s “2” is treated as shared because it also appears as the chosen response’s first operand (Li et al., 2025).Source: §2 · p. 3

Gradient interaction between training signals

• GAW-PO scores rejected tokens by signed inner products between their parameter gradients and chosen-response gradients, using the token’s own pair and dataset-averaged chosen responses.Unlike TI-DPO and AAO, the score depends on response context and is computed once before training without an auxiliary model.Source: §2 · p. 3
• Most surveyed token-level methods were evaluated on dialogue, safety, summarization, or a single domain, whereas this work evaluates 11 benchmarks across mathematics, reasoning, code, and question answering.The related-work comparison also distinguishes GAW-PO’s signed, context-dependent parameter-gradient score from TI-DPO’s gradient norm and AAO’s embedding-based token matches.Source: §2 · p. 3

3.1 PRELIMINARIES

• Rejected-token weights in [0, 1] replace the uniform sequence-level contribution, selectively weakening penalties without exceeding standard DPO’s penalty.A weight of 1 recovers the token’s standard DPO contribution.Source: §3.1 · pp. 3–4

3.2 GRADIENT-BASED TOKEN WEIGHTS

• Each rejected-token gradient is compared with paired and dataset-aggregate chosen-response directions; strong alignment with either yields a weaker penalty.The paired direction captures preference-specific interaction, while the aggregate direction captures behavior supported across the dataset.Source: §3.2 · p. 4

3.3 EXPERIMENTAL AND IMPLEMENTATION DECISIONS

• A rank-64 PiSSA adapter probes gradients, while full-model DPO uses fixed weights; aligned tokens are downweighted and conflicting tokens retain stronger penalties.Chosen-token gradients are normalized to approximate an average over token directions, and scalar projections avoid materializing the full token-by-parameter gradient matrix.Source: §3.3 · p. 5

3.4 ALGORITHM AND TRAINING PROCEDURE

• The procedure precomputes rejected-token weights from pair-specific and dataset-level chosen directions, then holds them fixed during full-model DPO training.Stronger alignment produces a smaller DPO penalty.Source: §3.4 · p. 6

4.1 EXPERIMENTAL DETAILS

• Experiments start from OLMo-3-7B-Instruct-SFT1 and use a filtered English, single-turn Dolci subset of approximately 209K pairs for one-epoch full fine-tuning.The default setup uses a maximum sequence length of 2048, learning rate 1 × 10^-5, effective batch size 128, and β = 0.1; alignment-score precomputation costs approximately 63 H200 GPU-hours and is reused across runs.Source: §4.1 · pp. 6–7

4.2 BASELINES

• Comparisons include the starting SFT checkpoint, standard DPO, three sequence-level variants, and three token-weighting methods, trained for one epoch on the same preference data.The sequence-level baselines are IPO, KTO, and SimPO; token-weighting baselines are TIS-DPO, TI-DPO, and AAO.Source: §4.2 · p. 7

4.3 EVALUATION BENCHMARKS

• Evaluation uses three seeds and OLMES benchmarks spanning mathematics, reasoning, code, and question answering.The selected suite includes AIME 2024 and 2025, MATH, Omega 500, AGI Eval English, BBH, HumanEval+, LiveCodeBench, MBPP+, GPQA, and MMLU.Source: §4.3 · p. 7

5.1 COMPARISON WITH EXISTING METHODS

• BBH reaches 49.81, 4.36 above DPO and 4.31 above the strongest competing baseline; LiveCodeBench reaches 24.78 versus 21.36 for DPO and 21.88 for TIS-DPO.Consistent gains are also reported on HumanEval+ (73.52), MATH (70.35), AGI Eval (61.15), MBPP+ (57.73), and MMLU (67.92).Source: §5.1 · p. 7
• GAW-PO does not lead every benchmark: TIS-DPO performs best on AIME 2025 and GPQA, IPO on Omega 500, and GAW-PO underperforms on GPQA and Omega 500.Source: §5.1 · p. 7
• 41.99 average score, 0.97 points above DPO (41.02) and 0.65 above TIS-DPO (41.34), with best results on 8 of 11 benchmarks.Table 1 compares the evaluated methods across benchmarks.Source: §5.1 · p. 7, Table 1 · p. 8

5.2 ROBUSTNESS TO THE DPO REGULARIZATION STRENGTH

• At β = 0.002, GAW-PO achieves its best aggregate result and surpasses the best standard-DPO setting, while DPO loses almost all performance.The sweep uses β values from 0.2 to 0.002 with the same learning rate; DPO peaks near β = 0.02, whereas GAW-PO remains stable and improves as β decreases.Source: Figure 2 · p. 8, §5.2 · p. 8

5.3 MULTIPLE PREFERRED UPDATE DIRECTIONS

• At β = 0.002, K = 4 reaches 47.94 versus 47.82 for paired+global; K = 2 performs worse, and gains do not rise monotonically with K.The clustered variant partitions chosen gradients into groups and uses each cluster centroid as a preferred direction; rejected tokens are protected when aligned with a paired response or recurring preferred direction.Source: Table 2 · p. 9, §5.3 · p. 9

5.2 ROBUSTNESS TO THE DPO REGULARIZATION STRENGTH

• Random or uniform weights with the same mean lose to gradient-aligned weights, and GAW-PO permits larger deviations from the reference policy without loss in downstream capability.The authors interpret this as consistent with aggressive DPO instability partly arising from penalties on locally correct rejected-response tokens.Source: §5.2 · p. 9

• GAW-PO compares rejected-token gradients with preferred update directions and reduces penalties on tokens whose suppression could interfere with learning chosen responses.A low-rank PiSSA gradient probe makes the computation practical, while preference optimization remains a full-model fine-tune.Source: §6 · p. 9
• Across the evaluation suite, GAW-PO has the highest average performance among considered methods and remains stable or improves as β decreases, unlike sharply degrading standard DPO.Clustered preferred directions yield modest gains for some K values, but increasing K alone is insufficient; direction structure and quality matter.Source: §6 · p. 9
• The authors plan to extend the framework beyond pairwise objectives such as DPO to online reinforcement-learning methods, particularly GRPO.They propose downweighting low-reward token updates aligned with high-advantage responses while preserving stronger penalties for conflicting tokens.Source: §6 · p. 9
• The paper identifies offline weight-computation cost and evaluation scope as limitations and directions for future work.Source: §6 · p. 10

• ChatGPT 5.0+ and Claude Opus 5.0+ assisted with drafting, revision, language polishing, and reducing repetition.Source: AI USE DISCLOSURE · p. 10
• Claude also assisted with suggesting and debugging experimental code and refining the teaser figure.Source: AI USE DISCLOSURE · p. 10
• The authors reviewed and verified generated text and code suggestions, remain responsible for the manuscript and research, and modified suggestions where needed.Source: AI USE DISCLOSURE · p. 10

A.1 RELATION TO THE DPO DERIVATION.

• The weighted rejected term is a contrastive surrogate, not the exact dual of a KL-constrained sequence-level objective; detached token weights scale rejected-token contributions in the standard DPO gradient.The weights can be interpreted as token-dependent regularization coefficients, but not as a consistent state-dependent coefficient across chosen and rejected trajectories.Source: §A.1 · p. 13

A.2 PISSA GRADIENT PROBE

• Because per-token gradients over all parameters of a 7B model are prohibitively expensive, a PiSSA-initialized low-rank adapter serves only as a gradient probe; preference training remains full-model fine-tuning.PiSSA initialization leaves the represented function unchanged while giving both low-rank factors non-zero values and informative gradients.Source: §A.2 · p. 13

A.3 CONSTRUCTION OF THE CHOSEN DIRECTION

• Rescaling chosen-token gradients makes their summed direction approximate an unweighted average over token-gradient directions rather than being dominated by a few surprising tokens.The rescaling gives each token’s logit-level error vector unit L1 norm; it does not normalize the parameter gradient itself.Source: §A.3 · p. 14

A.4 EFFICIENT ALIGNMENT PROJECTION COMPUTATION

• Forward-mode Jacobian-vector products compute alignment scores for all rejected tokens in a directional pass, avoiding materialization of the full rejected-token gradient matrix.The main method projects against paired and dataset-level chosen directions; the number of differentiation passes depends on preferred directions, not rejected-token count.Source: §A.4 · p. 14

B.1 CLUSTERING PROCEDURE

• Spherical k-means groups 1000-dimensional CountSketch projections by direction, then recomputes centroids in PiSSA space for token comparisons.Each rejected token independently selects its most-aligned preferred direction; evaluations use K = 2, 4, and 8.Source: §B.1 · pp. 14–15

B.2 EFFICIENT CLUSTER PROJECTIONS

• The clustered extension compares rejected-token gradients with multiple dataset-level preferred directions, using K + 1 directional passes without materializing per-token gradients.It retains the main method’s memory-efficient projection mechanism.Source: §B.2 · p. 15

C BASELINE HYPERPARAMETERS

• Baseline hyperparameters were specified by method: DPO and IPO use β = 0.1, KTO uses β = 0.1, and SimPO uses β = 10.The token-weighting baselines also use their stated method-specific settings.Source: C BASELINE HYPERPARAMETERS · pp. 15–16

D.1 QUALITATIVE EXAMPLES

• Weights tend to peak near the divergence from preferred behavior, while correctly reproduced content receives low weights and token roles depend on context.For example, the same fence token receives a high weight when it violates an instruction and a low weight when it does not.Source: §D.1 · p. 16

D.2 UNIFORM AND RANDOM WEIGHT CONTROLS

• 2.09 points over uniform weights and 3.36 points over random weights are achieved by gradient-aligned weighting.Uniform weighting recovers much of the stability, but attenuation alone does not explain the full improvement.Source: §D.2 · p. 16
• On HumanEval+, BBH, and MMLU, uniform and random weights trail GAW-PO by 12.95/21.47, 2.83/3.83, and 1.95/2.12 points, respectively.The mathematics benchmarks are generally within one standard deviation across variants; random weighting also performs below uniform weighting overall.Source: §D.2 · p. 16
• Table 3 shows learned weights concentrating on listed errors, while correctly transcribed content receives near-uniform low weights.Shading maps monotonically to token weight, from pale for low to deep for high.Source: Table 3 · p. 17

E.1 CATEGORY BREAKDOWN OF ROBUSTNESS TO THE DPO REGULARIZATION STRENGTH

• Standard DPO collapses across all four categories at smaller β, while GAW-PO remains stable and improves in mathematics, reasoning, and coding.Question answering remains approximately stable for GAW-PO despite increasingly aggressive updates.Source: Figure 3 · p. 18, §E.1 · p. 18

E.2 FULL CLUSTER-VARIANT RESULTS

• K = 4 obtains the strongest aggregate score among clustered variants, K = 8 is close behind, and K = 2 performs worse.No value of K dominates across all benchmarks.Source: Table 5 · p. 19, §E.2 · p. 19
• Cluster directions determine scores for up to 46.82% of rejected tokens, but greater cluster use does not monotonically improve downstream performance.K = 8 uses cluster directions most frequently yet performs slightly below K = 4.Source: §E.2 · p. 19

F LIMITATIONS

• The offline weighting stage costs roughly 1.5 training runs, making a single run less lightweight than standard DPO despite amortization across repeated runs.The stage is performed once per dataset and probe and reused across training runs, sweeps, and ablations.Source: F LIMITATIONS · p. 19
• Gradient alignment uses a low-rank PiSSA probe as a tractable proxy, and the effect of probe rank and its agreement with full-parameter alignment remain unstudied.The authors identify a systematic probe-rank and alignment study as future work.Source: F LIMITATIONS · p. 19
• Weights are computed from the SFT initialization and remain fixed during training, although policy gradient geometry evolves.Dynamic weight estimation is identified as an extension required for the online setting.Source: F LIMITATIONS · pp. 19–20
• Evaluation uses one open model family and preference dataset with verifiable benchmarks, excluding judge-based open-ended, multi-turn, and multilingual settings.The authors identify larger models and those excluded settings as directions for future work.Source: F LIMITATIONS · p. 20