Token-weighted Direct Preference Optimization with Attention

# Token-weighted Direct Preference Optimization with Attention
Source: [https://arxiv.org/html/2605.21883](https://arxiv.org/html/2605.21883)
Chengyu Huang1Zhuohang Li2Sheng\-Yen Chou1Claire Cardie1 1Cornell University2Vanderbilt University \{ch2263,sc3379,ctc9\}@cornell\.edu,zhuohang\.li@vanderbilt\.edu

###### Abstract

Direct Preference Optimization \(DPO\) aligns Large Language Models with human preferences without the need for a separate reward model\. However, DPO treats all tokens in responses equally, neglecting the differing importance of individual tokens\. Existing token\-level PO methods compute the token weights using either token\-position\-based heuristic functions or probability estimates given by a separately trained model, which lacks robustness and incurs extra training cost\. In contrast, we propose Token\-weighted DPO \(TwDPO\)—a novel training objective grounded on token\-weighted RL—and AttentionPO—an instantiation of TwDPO that uses attention from the LLM itself to estimate token weights\. AttentionPO prompts the LLM to serve as a pairwise judge and check where the model attends when comparing the responses\. This design makes AttentionPOcontent\-aware, adjusting weights based on response content, andefficient, incurring only two extra forward passes per example\. Experiment results show that AttentionPO significantly improves performance on AlpacaEval, MT\-Bench, and ArenaHard, surpassing existing Preference Optimization methods\. GitHub:[https://github\.com/HCY123902/AttentionPO](https://github.com/HCY123902/AttentionPO)

Token\-weighted Direct Preference Optimization with Attention

Chengyu Huang1Zhuohang Li2Sheng\-Yen Chou1Claire Cardie11Cornell University2Vanderbilt University\{ch2263,sc3379,ctc9\}@cornell\.edu,zhuohang\.li@vanderbilt\.edu

## 1Introduction

The alignment of Large Language Models \(LLMs\) with human preferences and values is critical for ensuring the quality and safety of generated outputs\. Early efforts in this domain utilized policy\-gradient reinforcement learning \(RL\) algorithms like Proximal Policy Optimization \(PPO\)Schulmanet al\.\([2017](https://arxiv.org/html/2605.21883#bib.bib23)\); Ouyanget al\.\([2022](https://arxiv.org/html/2605.21883#bib.bib24)\)\. However, these traditional RL methods require training a separate reward model on annotated preference pairs to provide feedback during trainingOuyanget al\.\([2022](https://arxiv.org/html/2605.21883#bib.bib24)\)\. To address this complexity, Preference Optimization \(PO\) methods such as Direct Preference Optimization \(DPO\)Rafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\)have emerged\. DPO’s training objective is derived from the policy gradient, but it learns the reward landscape directly from preference pairs, avoiding the need for an external reward model\.

![Refer to caption](https://arxiv.org/html/2605.21883v1/x1.png)Figure 1:AttentionPO weighs each token by attention, surpassing DPO and various other baselines\.Despite the advancements introduced by DPO and its variants—including IPOAzaret al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib16)\), KTOEthayarajhet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib15)\), and SimPOMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)—these methods generally treat every token within a response as equally importantMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\. This uniform weighting neglects the fact that different tokens contribute variably to the overall quality of a responseMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\. Consequently, traditional PO methods fail to provide the fine\-grained credit assignment necessary to account for the specific importance of individual tokens to human preference\.

Recent research has attempted to enable token\-level credit assignment through various estimation techniques\. Grounded in token\-level policy gradient methods, TDPO adds sequence KL terms to the DPO objectiveZenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib12)\)\. Other approaches, such as TIS\-DPOLiuet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib10)\), SePOYanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\), and cDPOLinet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib51)\), estimate token importance by contrasting probabilities between different policies or models\. Alternatively, TI\-DPO utilizes gradient norms and Gaussian priorsYanget al\.\([2026](https://arxiv.org/html/2605.21883#bib.bib8)\), while D2PO employs heuristic temporal decay functions that prioritize earlier tokensShaoet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib13)\)\. However, many of these methods either require the training of additional models to estimate weightsYanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\); Liuet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib10)\); Linet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib51)\)or rely on heuristic functions based on token positions rather than the specific semantic content of the tokensYanget al\.\([2026](https://arxiv.org/html/2605.21883#bib.bib8)\); Shaoet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib13)\)\.

In this work, we propose Token\-weighted DPO \(TwDPO\), a training objective that is theoretically\-grounded on token\-weighted RL, andAttentionPO, an instantiation of TwDPO that utilizes the attentions from the LLM itself to estimate token weights\. By obtaining attention weights for response tokens from a pairwise judge prompt, AttentionPO provides a content\-aware importance metric\. These weights are then normalized and applied during the credit assignment process to better align the model with human preferences\.

Experimental results demonstrate the effectiveness of our proposed approach\. AttentionPO achieves significant performance improvements across several models and standard benchmarks\. On LLaMA\-3\-8B\-Base\-SFTGrattafioriet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib30)\); Menget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\), AttentionPO improves performance by 12% \(win rate against GPT\-4\-1106\-previewHurstet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib31)\)\) on AlpacaEvalLiet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib27)\); Duboiset al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib28)\), 1\.05 \(LLM\-judged score\) on MT\-BenchZhenget al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib29)\), and 40% \(win rate against GPT\-4\-0314\) on ArenaHardLiet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib26),[2025](https://arxiv.org/html/2605.21883#bib.bib25)\)\. On LLaMA\-3\-8B\-Insturct, AttentionPO improves performance by 27% on AlpacaEval, 1\.41 on MT\-Bench, and 14% on ArenaHard\. AttentionPO also surpasses strong baselines such as SimPO \(up to 4% on AlpacaEval; 0\.20 on MT\-Bench\)\. AttentionPO highlights the value of using intrinsic model attentions for precise, token\-level preference optimization\.

## 2Methodology

![Refer to caption](https://arxiv.org/html/2605.21883v1/x2.png)Figure 2:Workflow of AttentionPO\. First, we promptπr​e​f\\pi\_\{ref\}to act as a pairwise judge and indicate which response is better in a single output tokenyv​e​r​d​i​c​ty\_\{verdict\}\. We extractπr​e​f\\pi\_\{ref\}’s attentions at layerℒ\\mathcal\{L\}that attend fromyv​e​r​d​i​c​ty\_\{verdict\}to the response tokens of bothywy\_\{w\}andyly\_\{l\}\. Prompting and weight extraction are repeated two times with the position ofywy\_\{w\}andyly\_\{l\}swapped, and the final attention weightsAwA\_\{w\}andAlA\_\{l\}are averaged over the two rounds\. Second, we post\-process the weights by applying normalization and fixing the attention sink\. Finally, we apply TwDPO with the post\-processed attention weights\.We first introduce the notations \(§[2\.1](https://arxiv.org/html/2605.21883#S2.SS1)\) and preliminaries \(§[2\.2](https://arxiv.org/html/2605.21883#S2.SS2)\), then present our TwDPO objective \(§[2\.3](https://arxiv.org/html/2605.21883#S2.SS3)\) and its math derivations \(§[2\.4](https://arxiv.org/html/2605.21883#S2.SS4)\)\. Finally, we discuss an instantiation of the TwDPO, AttentionPO, which uses attentions as weights\. AttentionPO takes three steps: \(1\) First, we prompt the initial LLM to act as a pairwise judge to extract the attention weights \(§[2\.5](https://arxiv.org/html/2605.21883#S2.SS5)\); \(2\) We post\-process the attention weights to obtain a weight distribution \(§[2\.6](https://arxiv.org/html/2605.21883#S2.SS6)\); \(3\) We apply TwDPO using the attention weight distribution\. See Figure[2](https://arxiv.org/html/2605.21883#S2.F2)for visualization\.

### 2\.1Notations

We initialize our main policyπθ\\pi\_\{\\theta\}with a reference modelπr​e​f\\pi\_\{ref\}\. We trainπθ\\pi\_\{\\theta\}on a preference dataset𝒟=\{x\(i\),yw\(i\),yl\(i\)\}i=1N\\mathcal\{D\}=\\\{x^\{\(i\)\},y\_\{w\}^\{\(i\)\},y\_\{l\}^\{\(i\)\}\\\}\_\{i=1\}^\{N\}, where each example contains a promptx\(i\)x^\{\(i\)\}, a preferred responseyw\(i\)y\_\{w\}^\{\(i\)\}and a dispreferred responseyl\(i\)y\_\{l\}^\{\(i\)\}\. Each responseyyconsists of tokens\{y1,⋯,y\|y\|\}\\\{y^\{1\},\\cdots,y^\{\|y\|\}\\\}\.

### 2\.2Preliminaries

Direct Preference Optimization is an offline variant of traditional policy gradient methods to align LLMs with human preferences\. It avoids the need for a separate reward model\. Instead, it trains LLMs as reward models themselves, on the annotated preference pairs\. In particular, DPO maximizes the model’s predicted probability ofywy\_\{w\}being preferred overyly\_\{l\}\. Formally, given the preference dataset𝒟\\mathcal\{D\}, DPO trainsπθ\\pi\_\{\\theta\}using the following objective:

ℒ=−𝔼\(x,yw,yl\)∼𝒟​\[l​o​g​\(pθ​\(yw≻yl\|x\)\)\]\.\\begin\{split\}\\mathcal\{L\}=\-\\mathbb\{E\}\_\{\(x,y\_\{w\},y\_\{l\}\)\\sim\\mathcal\{D\}\}\\biggl\[log\\biggl\(p\_\{\\theta\}\(y\_\{w\}\\succ y\_\{l\}\|x\)\\biggr\)\\biggr\]\.\\end\{split\}\(1\)wherepθ​\(yw≻yl\|x\)p\_\{\\theta\}\(y\_\{w\}\\succ y\_\{l\}\|x\)is the model’s predicted preference probability\. Following the Bradley\-Terry modelBradley and Terry \([1952](https://arxiv.org/html/2605.21883#bib.bib42)\),

pθ​\(yw≻yl\|x\)=exp⁡\(r​\(x,yw\)\)exp⁡\(r​\(x,yw\)\)\+exp⁡\(r​\(x,yl\)\)=σ​\(rθ​\(x,yw\)−rθ​\(x,yl\)\)\.\\begin\{split\}p\_\{\\theta\}\(y\_\{w\}\\succ y\_\{l\}\|x\)&=\\frac\{\\exp\(r\(x,y\_\{w\}\)\)\}\{\\exp\(r\(x,y\_\{w\}\)\)\+\\exp\(r\(x,y\_\{l\}\)\)\}\\\\ &=\\sigma\(r\_\{\\theta\}\(x,y\_\{w\}\)\-r\_\{\\theta\}\(x,y\_\{l\}\)\)\.\\end\{split\}\(2\)whererθ​\(x,yw\)r\_\{\\theta\}\(x,y\_\{w\}\)is model’s predicted reward andσ\\sigmais the sigmoid function\. Further, DPO derives the model’s predicted reward to berθ​\(x,y\)=β​l​o​g​πθ​\(y\|x\)πr​e​f​\(y\|x\)\+β​l​o​g​\(Z​\(x\)\)r\_\{\\theta\}\(x,y\)=\\beta log\\frac\{\\pi\_\{\\theta\}\(y\|x\)\}\{\\pi\_\{ref\}\(y\|x\)\}\+\\beta log\(Z\(x\)\), whereZ​\(x\)Z\(x\)is a partition function depending only onxx\. Thus, the final DPO objective is to minimize

ℒ=−𝔼\(x,yw,yl\)∼𝒟\[logσ\(βlogπθ​\(yw\|x\)πr​e​f​\(yw\|x\)−βlogπθ​\(yl\|x\)πr​e​f​\(yl\|x\)\)\]\.\\begin\{split\}\\mathcal\{L\}=\-\\mathbb\{E\}\_\{\(x,y\_\{w\},y\_\{l\}\)\\sim\\mathcal\{D\}\}\\biggl\[log\\sigma\\biggl\(\\beta log\\frac\{\\pi\_\{\\theta\}\(y\_\{w\}\|x\)\}\{\\pi\_\{ref\}\(y\_\{w\}\|x\)\}\-\\\\ \\beta log\\frac\{\\pi\_\{\\theta\}\(y\_\{l\}\|x\)\}\{\\pi\_\{ref\}\(y\_\{l\}\|x\)\}\\biggr\)\\biggr\]\.\\end\{split\}\(3\)Equation[3](https://arxiv.org/html/2605.21883#S2.E3)can be rewritten into the token\-level form:

ℒ=−𝔼\(x,yw,yl\)∼𝒟\[logσ\(β∑t=1\|yw\|logπθ​\(ywt\|x,yw<t\)πr​e​f​\(ywt\|x,yw<t\)−β∑t=1\|yl\|logπθ​\(ylt\|x,yl<t\)πr​e​f​\(ylt\|x,yl<t\)\)\]\.\\begin\{split\}\\mathcal\{L\}=\-\\mathbb\{E\}\_\{\(x,y\_\{w\},y\_\{l\}\)\\sim\\mathcal\{D\}\}\\biggl\[log\\sigma\\biggl\(\\beta\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}\\log\\frac\{\\pi\_\{\\theta\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\-\\\\ \\beta\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}log\\frac\{\\pi\_\{\\theta\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\\biggr\)\\biggr\]\.\\end\{split\}\(4\)Note that DPO assigns an equal weight of 1 to all theπθ​\(yt\|x,y<t\)πr​e​f​\(yt\|x,y<t\)\\frac\{\\pi\_\{\\theta\}\(y^\{t\}\|x,y^\{<t\}\)\}\{\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\}terms\. This does not consider the varying degrees of importance of different tokens to the final response quality\.

### 2\.3Token\-weighted DPO

Instead, our objective multiplies each probability ratio term in Equation[4](https://arxiv.org/html/2605.21883#S2.E4)with a token weightata^\{t\}:

ℒ=−𝔼𝒟\[logσ\(β​\|yw\|​∑t=1\|yw\|awt​log⁡πθ​\(ywt\|x,yw<t\)πr​e​f​\(ywt\|x,yw<t\)−β\|yl\|∑t=1\|yl\|altlogπθ​\(ylt\|x,yl<t\)πr​e​f​\(ylt\|x,yl<t\)\)\]\.\\begin\{split\}\\mathcal\{L\}=\-\\mathbb\{E\}\_\{\\mathcal\{D\}\}\\biggl\[log\\sigma\\biggl\(&\\beta\|y\_\{w\}\|\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{t\}\\log\\frac\{\\pi\_\{\\theta\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\\\\ \-&\\beta\|y\_\{l\}\|\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{t\}log\\frac\{\\pi\_\{\\theta\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\\biggr\)\\biggr\]\.\\end\{split\}\(5\)

### 2\.4Mathematical Derivation

We closely follow the proof ofRafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\)to derive the token\-weighted DPO objective from a token\-weighted RL objective\. First, we find the closed\-form expression of the near\-optimal policyπ∗\\pi^\{\*\}and reward functionr∗r^\{\*\}attained by the token\-weighted RL objective \(§[2\.4\.1](https://arxiv.org/html/2605.21883#S2.SS4.SSS1)\)\. Next, we derive the token\-weighted DPO objective \(§[2\.4\.2](https://arxiv.org/html/2605.21883#S2.SS4.SSS2)\)\.

#### 2\.4\.1Deriving the Near Optimum of the Token\-weighted Objective

DPORafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\)derives the equivalence between their optimal policy and the one attained by a sequence\-level RL objective, which maximizes response reward with KL constraints:

maxπ𝔼x∼𝒟,y∼π\(⋅\|x\)\[r\(x,y\)\]−β𝔻𝕂​𝕃\(πθ\(y\|x\)\|\|πr​e​f\(y\|x\)\)=maxπ⁡𝔼x∼𝒟,y∼π\(⋅\|x\)​\[r​\(x,y\)−β​∑t=1\|y\|πθ​\(yt\|x,y<t\)πr​e​f​\(yt\|x,y<t\)\]\.\\begin\{split\}&\\max\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\},y\\sim\\pi\(\\cdot\|x\)\}\\left\[r\(x,y\)\\right\]\-\\beta\\mathbb\{D\_\{KL\}\}\(\\pi\_\{\\theta\}\(y\|x\)\|\|\\pi\_\{ref\}\(y\|x\)\)\\\\ &=\\max\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\},y\\sim\\pi\(\\cdot\|x\)\}\\biggl\[r\(x,y\)\-\\beta\\sum\_\{t=1\}^\{\|y\|\}\\frac\{\\pi\_\{\\theta\}\(y^\{t\}\|x,y^\{<t\}\)\}\{\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\}\\biggr\]\.\\end\{split\}\(6\)
Instead, we show that our objective[5](https://arxiv.org/html/2605.21883#S2.E5)gives a model equivalent to the near\-optimal policy produced by a token\-weighted RL objective\. We start by reformulating objective[6](https://arxiv.org/html/2605.21883#S2.E6)\. For a given promptxx, a policyπ\\pi, and a reference modelπr​e​f\\pi\_\{ref\}, our objective is to find an optimal policy that maximizes the reward under a token\-weighted KL penalty\. As we will show in §[2\.4\.3](https://arxiv.org/html/2605.21883#S2.SS4.SSS3), this objective naturally induces a token\-weighted reward functionrr\. Our KL penalty weighs the log\-probability ratios at each token stepttby the weightata^\{t\}:

maxπ𝔼x∼𝒟,y∼π\(⋅\|x\)\[r\(x,y\)−β\|y\|∑t=1\|y\|atlogπ​\(yt\|x,y<t\)πr​e​f​\(yt\|x,y<t\)\]\.\\begin\{split\}&\\max\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\},y\\sim\\pi\(\\cdot\|x\)\}\\biggl\[r\(x,y\)\-\\\\ &~~~~~~~~~~\\beta\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\frac\{\\pi\(y^\{t\}\|x,y^\{<t\}\)\}\{\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\}\\biggr\]\.\\end\{split\}\(7\)Note thatata^\{t\}is a proper distribution and sums up to 1\. In light of this, we scale the KL term by the sequence length\|y\|\|y\|so that the total weight of the logprob ratio terms\|y\|​∑t=1\|y\|at\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}equals\|y\|\|y\|, which is consistent with the DPO objective[6](https://arxiv.org/html/2605.21883#S2.E6)\.

To solve this mathematically, we define a sequence\-level "weighted" log\-probability for both the policy and the reference model:

log⁡π~​\(y\|x\)=\|y\|​∑t=1\|y\|at​log⁡π​\(yt\|x,y<t\),log⁡π~r​e​f​\(y\|x\)=\|y\|​∑t=1\|y\|at​log⁡πr​e​f​\(yt\|x,y<t\)\.\\begin\{split\}\\log\\tilde\{\\pi\}\(y\|x\)&=\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\pi\(y^\{t\}\|x,y^\{<t\}\),\\\\ \\log\\tilde\{\\pi\}\_\{ref\}\(y\|x\)&=\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\.\\end\{split\}\(8\)Substituting these into our objective, we can follow the standard derivation steps:

maxπ⁡𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)​\[r​\(x,y\)−β​log⁡π~​\(y\|x\)π~r​e​f​\(y\|x\)\]=minπ⁡𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)​\[log⁡π~​\(y\|x\)π~r​e​f​\(y\|x\)−1β​r​\(x,y\)\]=minπ𝔼x∼𝒟𝔼y∼π\(⋅\|x\)\[logπ~​\(y\|x\)1Z​\(x\)​π~r​e​f​\(y\|x\)​exp⁡\(1β​r​\(x,y\)\)−logZ\(x\)\],\\begin\{split\}&\\max\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\left\[r\(x,y\)\-\\beta\\log\\frac\{\\tilde\{\\pi\}\(y\|x\)\}\{\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\}\\right\]\\\\ &=\\min\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\left\[\\log\\frac\{\\tilde\{\\pi\}\(y\|x\)\}\{\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\}\-\\frac\{1\}\{\\beta\}r\(x,y\)\\right\]\\\\ &=\\min\_\{\\pi\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\biggl\[\\log\\frac\{\\tilde\{\\pi\}\(y\|x\)\}\{\\frac\{1\}\{Z\(x\)\}\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\\exp\(\\frac\{1\}\{\\beta\}r\(x,y\)\)\}\\\\ &~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\-\\log Z\(x\)\\biggr\],\\end\{split\}whereZ​\(x\)=∑yπ~r​e​f​\(y\|x\)​exp⁡\(1β​r​\(x,y\)\)Z\(x\)=\\sum\_\{y\}\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\\exp\(\\frac\{1\}\{\\beta\}r\(x,y\)\)is the partition function\.

The near optimal weighted policyπ~∗\\tilde\{\\pi\}^\{\*\}\(we analyze the error bounds betweenπ~∗\\tilde\{\\pi\}^\{\*\}and the intractable true optimal policy in Appendix[B](https://arxiv.org/html/2605.21883#A2)\) that minimizes this expression satisfies:

π~∗​\(y\|x\)=1Z​\(x\)​π~r​e​f​\(y\|x\)​exp⁡\(1β​r​\(x,y\)\)\.\\tilde\{\\pi\}^\{\*\}\(y\|x\)=\\frac\{1\}\{Z\(x\)\}\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\\exp\\left\(\\frac\{1\}\{\\beta\}r\(x,y\)\\right\)\.\(9\)By taking the logarithm of both sides and rearranging the terms, we can express the true rewardr∗​\(x,y\)r^\{\*\}\(x,y\)entirely in terms of the optimal policy:

r∗​\(x,y\)=β​log⁡π~∗​\(y\|x\)π~r​e​f​\(y\|x\)\+β​log⁡Z​\(x\)\.r^\{\*\}\(x,y\)=\\beta\\log\\frac\{\\tilde\{\\pi\}^\{\*\}\(y\|x\)\}\{\\tilde\{\\pi\}\_\{ref\}\(y\|x\)\}\+\\beta\\log Z\(x\)\.\(10\)Substituting the expanded definition of our weighted log\-probabilities back into this reward expression yields:

r∗​\(x,y\)=β​\|y\|​∑t=1\|y\|at​log⁡π∗​\(yt\|x,y<t\)πr​e​f​\(yt\|x,y<t\)\+β​log⁡Z​\(x\)\.r^\{\*\}\(x,y\)=\\beta\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\frac\{\\pi^\{\*\}\(y^\{t\}\|x,y^\{<t\}\)\}\{\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\}\+\\beta\\log Z\(x\)\.\(11\)

#### 2\.4\.2Deriving the Objective Under the Bradley\-Terry Model

We substitute our token\-weighted reparameterization ofr∗​\(x,y\)r^\{\*\}\(x,y\)into the Bradley\-Terry preference model \(Equation[2](https://arxiv.org/html/2605.21883#S2.E2)\):

p∗​\(yw≻yl\|x\)=σ​\(r∗​\(x,yw\)−r∗​θ​\(x,yl\)\)=σ\(\(β\|yw\|∑t=1\|yw\|awtlogπ∗​\(ywt\|x,yw<t\)πr​e​f​\(ywt\|x,yw<t\)\+βlogZ\(x\)\)−\(β\|yl\|∑t=1\|yl\|altlogπ∗​\(ylt\|x,yl<t\)πr​e​f​\(ylt\|x,yl<t\)\+βlogZ\(x\)\)\)\.\\begin\{split\}&p^\{\*\}\(y\_\{w\}\\succ y\_\{l\}\|x\)=\\sigma\(r^\{\*\}\(x,y\_\{w\}\)\-r^\{\*\}\{\\theta\}\(x,y\_\{l\}\)\)\\\\ &=\\sigma\\Bigg\(\\left\(\\beta\|y\_\{w\}\|\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{t\}\\log\\frac\{\\pi^\{\*\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\+\\beta\\log Z\(x\)\\right\)\\\\ &\\quad\-\\left\(\\beta\|y\_\{l\}\|\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{t\}\\log\\frac\{\\pi^\{\*\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\+\\beta\\log Z\(x\)\\right\)\\Bigg\)\.\\end\{split\}Crucially, the partition function termsβ​log⁡Z​\(x\)\\beta\\log Z\(x\)cancel out, leaving:

p∗\(yw≻yl\|x\)=σ\(β\|yw\|∑t=1\|yw\|awtlogπ∗​\(ywt\|x,yw<t\)πr​e​f​\(ywt\|x,yw<t\)−β\|yl\|∑t=1\|yl\|altlogπ∗​\(ylt\|x,yl<t\)πr​e​f​\(ylt\|x,yl<t\)\)\.\\begin\{split\}p^\{\*\}\(y\_\{w\}\\succ y\_\{l\}\|x\)=\\sigma\\biggl\(\\beta\|y\_\{w\}\|\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{t\}\\log\\frac\{\\pi^\{\*\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\\\\ \-\\beta\|y\_\{l\}\|\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{t\}\\log\\frac\{\\pi^\{\*\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\\biggr\)\.\\end\{split\}\(12\)To frame this as a maximum likelihood objective for training a parameterized policyπθ\\pi\_\{\\theta\}, we take the negative log\-likelihood over the preference dataset𝒟\\mathcal\{D\}\. This leads to our final objective in Equation[5](https://arxiv.org/html/2605.21883#S2.E5)\.

#### 2\.4\.3Properties

##### Reward function is token\-weighted\.

Equation[10](https://arxiv.org/html/2605.21883#S2.E10)shows thatrrweighs tokens byata^\{t\}\. For tokens with higherata^\{t\}, their logprob ratio betweenπ∗​\(yt\|x,y<t\)\\pi^\{\*\}\(y^\{t\}\|x,y^\{<t\}\)andπr​e​f​\(yt\|x,y<t\)\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)significantly affectrr, while tokens with lowerata^\{t\}affectrrless\.

##### Token\-weighted KL penalty enables locally\-adaptive trust region\.

Our KL penalty in objective[7](https://arxiv.org/html/2605.21883#S2.E7)scales the KL divergence at the token level viaata^\{t\}\. Consequently, the model is granted greater freedom to deviate from the reference policy on low\-importance tokens, while being strictly constrained on high\-importance ones\. While restricting the most critical tokens may seem counterintuitive, this adaptive weighting is essential because important tokens dominate the implicitly derived reward function, and thus are highly susceptible to reward hacking and over\-optimization\. The token\-weighted KL term acts as a locally adaptive regulator: for high\-weight tokens, the massive incentives to deviate are safely balanced by proportionally stricter penalties, ensuring the model only updates when the expected reward justifies the risk\.

The above suggests that objective[7](https://arxiv.org/html/2605.21883#S2.E7)is indeed a token\-weighted RL objective in terms of both the reward function and KL penalty\. Consequently, both its derived policy and the near\-optimal policy derived from objective[5](https://arxiv.org/html/2605.21883#S2.E5)learn more signals from important tokens, which is what we desire\. We further provide an analysis on the gradient in Appendix[C](https://arxiv.org/html/2605.21883#A3)\. While we use attention as the weights, there are no assumptions about how the weights are derived, and AttentionPO is only one possible instantiation\.

### 2\.5Token Weights Extraction Through Attention Scores

We hypothesize that if we ask the LLM to judge the response quality, the tokens that get more attention are more important\. As such, we promptπr​e​f\\pi\_\{ref\}with a general pairwise judge prompt \(See Appendix[6](https://arxiv.org/html/2605.21883#A9.F6)\) that asks it to predict which ofywy\_\{w\}andyly\_\{l\}has better quality\. We explicitly requireπr​e​f\\pi\_\{ref\}to output the identifier \("A" or "B"\) of the better response, and we denote this single token asyv​e​r​d​i​c​ty\_\{verdict\}\. To avoid bias, we hide the preference labels in the prompt and callπr​e​f\\pi\_\{ref\}two rounds, each round with the positions ofywy\_\{w\}andyly\_\{l\}swapped\. For each round, we do a single forward pass and obtain the attention weights at a fixed layerLL\. We extract the attention weights fromyv​e​r​d​i​c​ty\_\{verdict\}to all the tokens inywy\_\{w\}and denote them asAwr​o​u​n​d​\_​r=\[aw1,⋯,aw\|yw\|\]A\_\{w\}^\{round\\\_r\}=\[a\_\{w\}^\{1\},\\cdots,a\_\{w\}^\{\|y\_\{w\}\|\}\]\. Similarly, we extract the attention weights fromyv​e​r​d​i​c​ty\_\{verdict\}to all the tokens inyly\_\{l\}asAlr​o​u​n​d​\_​r=\[al1,⋯,al\|yl\|\]A\_\{l\}^\{round\\\_r\}=\[a\_\{l\}^\{1\},\\cdots,a\_\{l\}^\{\|y\_\{l\}\|\}\]\. For simplicity, we compute the final attention weight of each token as the mean attention weight for that token across attention heads\.

For roundsr∈\{1,2\}r\\in\\\{1,2\\\},This producesAwr​o​u​n​d​\_​1A\_\{w\}^\{round\\\_1\},Awr​o​u​n​d​\_​2A\_\{w\}^\{round\\\_2\},Alr​o​u​n​d​\_​1A\_\{l\}^\{round\\\_1\},Alr​o​u​n​d​\_​2A\_\{l\}^\{round\\\_2\}\. The final attention is the average of the two rounds:

awt=12​awt,r​o​u​n​d​\_​1\+12​awt,r​o​u​n​d​\_​2,alt=12​alt,r​o​u​n​d​\_​1\+12​alt,r​o​u​n​d​\_​2\.\\begin\{split\}a\_\{w\}^\{t\}&=\\frac\{1\}\{2\}a\_\{w\}^\{t,round\\\_1\}\+\\frac\{1\}\{2\}a\_\{w\}^\{t,round\\\_2\},\\\\ a\_\{l\}^\{t\}&=\\frac\{1\}\{2\}a\_\{l\}^\{t,round\\\_1\}\+\\frac\{1\}\{2\}a\_\{l\}^\{t,round\\\_2\}\.\\end\{split\}\(13\)

### 2\.6Post\-Processing

##### Normalization\.

πr​e​f\\pi\_\{ref\}also attends to other tokens in the pairwise judge prompt and so neitherAwA\_\{w\}norAlA\_\{l\}form a proper distribution \(i\.e\., it does not sum up to 1\.0\)\. Therefore, we normalize the attention weightsAwA\_\{w\}andAlA\_\{l\}:awt←awt/\(∑q=1\|yw\|awq\)a\_\{w\}^\{t\}\\leftarrow a\_\{w\}^\{t\}/\(\\sum\_\{q=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{q\}\),alt←alt/\(∑q=1\|yl\|alq\)a\_\{l\}^\{t\}\\leftarrow a\_\{l\}^\{t\}/\(\\sum\_\{q=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{q\}\)\.

##### Fixing Attention Sink\.

Attention sink is a phenomenon where the LLMs predominantly attend to the initial tokensXiaoet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib43)\)\. Manual inspection on our dataset𝒟\\mathcal\{D\}suggests that these initial tokens are usually of low importance, consisting of starting phrases\. As such, for any response with at leastK′K^\{\\prime\}tokens, we reset the attention weights of the firstKKtokens of each responseyyto the average weight1/\|y\|1/\|y\|and renormalize the remaining tokens:

∀t∈\{K\+1,⋯,\|yw\|\},awt←awt​1−K\|yw\|∑q=K\+1\|yw\|awq∀t∈\{K\+1,⋯,\|yl\|\},alq←alq​1−K\|yl\|∑q=K\+1\|yl\|alq\\begin\{split\}\\forall t\\in\\\{K\+1,\\cdots,\|y\_\{w\}\|\\\},a\_\{w\}^\{t\}&\\leftarrow a\_\{w\}^\{t\}\\frac\{1\-\\frac\{K\}\{\|y\_\{w\}\|\}\}\{\\sum\_\{q=K\+1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{q\}\}\\\\ \\forall t\\in\\\{K\+1,\\cdots,\|y\_\{l\}\|\\\},a\_\{l\}^\{q\}&\\leftarrow a\_\{l\}^\{q\}\\frac\{1\-\\frac\{K\}\{\|y\_\{l\}\|\}\}\{\\sum\_\{q=K\+1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{q\}\}\\end\{split\}\(14\)

## 3Experiment Setup

##### Models and Datasets\.

We follow the settings fromMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\. For the initial modelπr​e​f\\pi\_\{ref\}, we experiment withLlama\-3\-8B\-Base\-SFTfromMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\(princeton\-nlp/Llama\-3\-Base\-8B\-SFT\), which is supervised\-finetuned from the pretrained model LLaMA\-3\-8BGrattafioriet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib30)\)on UltraChat\-200KDinget al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib48)\); andLLaMA\-3\-8B\-InstructGrattafioriet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib30)\)\(meta\-llama/Meta\-Llama\-3\-8B\-Instruct\)\.

For the SFT model \(Llama\-3\-8B\-Base\-SFT\), we apply AttentionPO on the UltraFeedback binarized datasetCuiet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib45)\)\(HuggingFaceH4/ultrafeedback\_binarized\)\. This dataset contains 61,135/2,000 training/validation examples covering diverse instruction\-following tasks such as creative writing, document\-assisted writing, and open\-ended QA\. Each query is accompanied by four responses generated by diverse LLMs and graded by GPT\-4 in terms of instruction\-following, helpfulness, honesty, and truthfulness\.ywy\_\{w\}is set to the response with the highest overall score, andyly\_\{l\}is chosen randomly from the remaining three\.

For the Instruction\-Tuned model \(LLaMA\-3\-8B\-Instruct\), we use the dataset produced byMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\(princeton\-nlp/llama3\-ultrafeedback\)\. It contains 59,876/1,961 training/validation examples\. The questions are subsampled from the binarized datasets, but for each question, five responses are sampled fromπr​e​f\\pi\_\{ref\}and then scored by a reward model PairRMJianget al\.\([2023b](https://arxiv.org/html/2605.21883#bib.bib50)\)\.ywy\_\{w\}is set to the response with the highest score andyly\_\{l\}is set to the one with the lowest score\.

Table 1:Statistics of the attention weights after post\-processing in §[2\.6](https://arxiv.org/html/2605.21883#S2.SS6)\. Std: Standard deviation of token weights within a response\. Max: Maximum weights; Len: Response length; Metrics are averaged across the training set\.Table 2:Most weighted tokens after post\-processing in §[2\.6](https://arxiv.org/html/2605.21883#S2.SS6)\. Weight: Weights averaged across all occurrences of the token in the training set\. We only keep tokens that occur at least 100 times\.
##### Attention Weights\.

We setLLto the index of the last layer for every model \(i\.e\., we take the attention weights from the last layer\)\. To fix the attention sink issue, we useK=1K=1andK′=5K^\{\\prime\}=5\. That is, for responseyythat has at least 5 tokens, we set the attention weight of its first token to1/\|y\|1/\|y\|\. The rationale is that for our models, the attention sink primarily concentrates on the first token\. For LLaMA3\-8B\-Base\-SFT, significant attention weights are on the first token \(12\.55%\), whereas the weights on the second \(3\.35%\) and third \(2\.55%\) tokens are much smaller\. For LLaMA\-3\-8B\-Instruct, the proportion is 8\.25%/2\.11%/1\.55% for the first, second, and third tokens, respectively\.

We show detailed statistics of the attention weights in Table[1](https://arxiv.org/html/2605.21883#S3.T1)\. The standard deviation of the weights from both SFT and Instruct models is above 0\.01, suggesting that the weights vary across tokens\. We also show the ten most weighted tokens in Table[2](https://arxiv.org/html/2605.21883#S3.T2)\. For the SFT model, 7/10 tokens are related to binary judgment words \(e\.g\., "Yes", "True"\) for bothywy\_\{w\}andyly\_\{l\}\. The highly weighted tokens of the Instruct model not only include these but also negation words \(e\.g\., "doesn"\)\.

##### Training Details\.

We set learning rate to1​e−61e\-6with a cosine scheduler and a warmup ratio of 0\.1\. We use the AdamW optimizerLoshchilov and Hutter \([2019](https://arxiv.org/html/2605.21883#bib.bib46)\)\. The KL regularizerβ\\betais set to0\.0050\.005\. See more details in Appendix[E](https://arxiv.org/html/2605.21883#A5)\.

##### Baselines\.

We compare AttentionPO against the initial modelπr​e​f\\pi\_\{ref\}and models trained with common PO methods, includingRRHFYuanet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib20)\),SLiC\-HFZhaoet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib22)\),DPORafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\),IPOAzaret al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib16)\),CPOXuet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib19)\),KTOEthayarajhet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib15)\),ORPOHonget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib18)\),R\-DPOParket al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib17)\), andSimPOMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\. For all the baselines, we adopt the training hyperparameters fromMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\), which have been separately searched to be optimal\.

##### Evaluation\.

Following prior workMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\); Yanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\); Huang and Goyal \([2025](https://arxiv.org/html/2605.21883#bib.bib44)\), we evaluate the models on AlpacaEvalLiet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib27)\), MT\-BenchZhenget al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib29)\), and ArenaHardLiet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib26)\)\. For AlpacaEval, we report the raw and length\-controlled win rate of our model against GPT\-4\-1106\-preview\. For MT\-Bench, we report the LLM\-judged quality score of our model’s generated response on a scale of 1 to 10\. For ArenaHard, we report the win rate of our model against GPT4\-0314\. We use GPT\-4o\-miniAchiamet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib32)\)as the judge model for these benchmarks\. See more evaluation details in Appendix[E](https://arxiv.org/html/2605.21883#A5)\.

## 4Results

We first present the main results \(§[4\.1](https://arxiv.org/html/2605.21883#S4.SS1)\), followed by comparing our method with other choices of token weights \(§[4\.2](https://arxiv.org/html/2605.21883#S4.SS2)\), and finally conducting an ablation study \(§[4\.3](https://arxiv.org/html/2605.21883#S4.SS3)\)\.

### 4\.1Main Results

Table 3:Main results\. Base:πr​e​f\\pi\_\{ref\}; MT: MT\-Bench; AlpacaEval: LR is the length\-controlled win rate, and WR is the raw win rate against GPT\-4\-1106\-Preview; MT\-Bench: LLM\-judged quality score of the model’s responses, on a scale of 1 to 10; ArenaHard: Win rate against GPT\-4\-0314\.As shown in Table[3](https://arxiv.org/html/2605.21883#S4.T3),AttentionPO significantly improves the performance over the base modelπr​e​f\\pi\_\{ref\}in all settings\. On LLaMA\-3\-8B\-Base\-SFT, performance increases by 12% on AlpacaEval, 1\.05 on MT\-Bench, 40% on ArenaHard\. On LLaMA\-3\-8B\-Instruct, performance increases by 27% on AlpacaEval, 0\.41 on MT\-Bench, and 14% on ArenaHard\. Furthermore,AttentionPO outperforms the baselines in most settings\. For the SFT model, AttentionPO outperforms all the baselines on ArenaHard and achieves second place on MT\-Bench and third place on AlpacaEval\. For the Instruct model, AttentionPO outperforms all baselines on AlpacaEval and MT\-Bench\.

### 4\.2Impact of Token Weight Sources

There are alternative ways to obtain token\-level weights\. Prior workLinet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib51)\); Yanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\); Liuet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib10)\); Shaoet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib13)\); Yanget al\.\([2026](https://arxiv.org/html/2605.21883#bib.bib8)\)either requires training a separate model to estimate weights or relies on strong heuristics\. Instead, we stick to our setting where the weights are attained through self\-judging, but we experiment with various changes to the current design\. This includes \(1\) Extracting the attention weights from a different hidden layerLLofπr​e​f\\pi\_\{ref\}; \(2\) Using attention rolloutAbnar and Zuidema \([2020](https://arxiv.org/html/2605.21883#bib.bib41)\)to aggregate the weights across different hidden layers; \(3\) Promptingπr​e​f\\pi\_\{ref\}to verbally assign a weight for each token\. For these experiments, we use LLaMA\-3\-8B\-Instruct as an example\.

![Refer to caption](https://arxiv.org/html/2605.21883v1/x3.png)Figure 3:Results for different layers\. x\-axis: index of the layer from which attention weights are extracted; y\-axis: benchmark performance\.##### Changing layers\.

LLaMA\-3\-8B\-Instruct consists of 32 hidden layers\. We useL=32L=32for our main experiments\. Here, we experiment withL∈\{1,8,16,24,32\}L\\in\\\{1,8,16,24,32\\\}and present the results in Figure[3](https://arxiv.org/html/2605.21883#S4.F3)\. We find thatthe first and last layers perform better than the middle ones, with the last layer achieving the best results\. We leave more investigation on this to future work\.

##### Attention rollout \(w/ attn rollout\)\.

Attention rollout aggregates the attentions across all the layers to estimate the total attentions from tokens in the final output layer to tokens in the initial input layer\. Table[4](https://arxiv.org/html/2605.21883#S4.T4)shows that using their aggregated weights underperforms our method\. This is potentially associated with our previous observation that middle\-layer attentions are not suitable for AttentionPO, since attention rollout incorporates attentions from these layers into the final weights\.

##### Verbalized self\-judged weights \(w/ verbal w\)\.

Instead of using attention, we askπr​e​f\\pi\_\{ref\}to verbally assign an importance score on a scale of 1 to 5 to each token in the response\. We observe that it is difficult for the model to precisely give a score for every token \(e\.g\., malformed scoring report; not being able to recall every token in the response\)\. Therefore, we adapt a coarser granularity and ask the model to first partition the response into parts and then assign a score for each part \(See prompts in Appendix[I](https://arxiv.org/html/2605.21883#A9)\)\. Even with a coarser granularity,πr​e​f\\pi\_\{ref\}still cannot perfectly recall the response tokens\. We filter out examples where fewer than 90% of the tokens are matched with a verbalized weight for eitherywy\_\{w\}oryly\_\{l\}\. After filtering, we get 10,111/343 train/validation examples\. To compensate for a smaller dataset, we increase the number of training epochs from 1 to 5 and validate at each epoch\.

Table 4:Results with different weight sources\.Table[4](https://arxiv.org/html/2605.21883#S4.T4)shows that this method outperforms ours on ArenaHard but underperforms on AlpacaEval and MT\-Bench\. A potential reason is the use of a coarser granularity\. We leave more investigation on the viability of this approach to future work\.

### 4\.3Ablation Study

##### Restoring the attention sink \(w/ attn sink\)\.

We restore the attention sink, in which case the first token occupies roughly 8\.25% of the total weights\. Table[5](https://arxiv.org/html/2605.21883#S4.T5)shows that this variant underperforms our method on all benchmarks, indicating the benefits of averaging out the attention sink\.

##### Adding length normalization \(w/ len norm\)\.

Prior workYanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\); Menget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)argues that DPO is biased toward lengthy responses \(i\.e\., lengthy responses receive more gradient updates than concise ones\)\. They add a length normalizer to the objective\. Similarly, we remove the\|y\|\|y\|multiplier in front of the logprob ratio term forywy\_\{w\}andyly\_\{l\}, giving a weighted, length\-normalized variant:

ℒ′=−𝔼𝒟\[logσ\(β​∑t=1\|yw\|awt​log⁡πθ​\(ywt\|x,yw<t\)πr​e​f​\(ywt\|x,yw<t\)−β∑t=1\|yl\|altlogπθ​\(ylt\|x,yl<t\)πr​e​f​\(ylt\|x,yl<t\)\)\]\\begin\{split\}\\mathcal\{L^\{\\prime\}\}=\-\\mathbb\{E\}\_\{\\mathcal\{D\}\}\\biggl\[log\\sigma\\biggl\(&\\beta\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{t\}\\log\\frac\{\\pi\_\{\\theta\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\}\\\\ \-&\\beta\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{t\}log\\frac\{\\pi\_\{\\theta\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\{\\pi\_\{ref\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\}\\biggr\)\\biggr\]\\end\{split\}\(15\)Adding the length normalization effectively reduces the KL regularization constantβ\\beta\. Therefore, we increaseβ\\betato 2\.0 for this baseline\. For these experiments, we use LLaMA\-3\-8B\-Instruct as an example\.

Table 5:Ablation Results\.Table[5](https://arxiv.org/html/2605.21883#S4.T5)shows that applying length normalization significantly hurts performance on MT\-Bench and AreaHard\. One speculated reason for this is the misalignment with the original RL objective\. Instead of maximizing the total rewards of the entire rolloutyy, the dual\-RL objective of objective[15](https://arxiv.org/html/2605.21883#S4.E15)\(which is the same as objective[7](https://arxiv.org/html/2605.21883#S2.E7)but with\|y\|\|y\|removed in both the reward function and the KL penalty\) maximizes the average per\-token reward ofyy, undermining the total rewards\.

## 5Related Work

##### Preference Optimization\.

Aligning LLMs with human preferences originally relied on RLHF via PPOOuyanget al\.\([2022](https://arxiv.org/html/2605.21883#bib.bib24)\); Schulmanet al\.\([2017](https://arxiv.org/html/2605.21883#bib.bib23)\), which requires a separate reward model\. Modern Preference Optimization \(PO\) methods bypass this by training the policy directly on preference pairs\(yw,yl\)\(y\_\{w\},y\_\{l\}\)using contrastive objectives\. For instance, RRHFYuanet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib20)\)and SLiC\-HFZhaoet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib22)\)apply ranking losses, while DPORafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\)derives a reference\-normalized objective from policy gradients\. Subsequent variants introduce bounded objectives \(IPOAzaret al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib16)\)\), pointwise formulations \(KTOEthayarajhet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib15)\)\), auxiliary SFT objectives \(CPOXuet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib19)\), ORPOHonget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib18)\)\), length penalties \(R\-DPOParket al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib17)\)\), or reference\-free length normalization with reward margins \(SimPOMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)\)\. However, these methods weight all response tokens equally, failing to capture token\-level importance\.

##### Token\-level PO\.

To enable fine\-grained credit assignment, recent work explores token\-level weighting\. TDPOZenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib12)\)incorporates token\-level sequential KL terms\. TIS\-DPOLiuet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib10)\), SePOYanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\), and cDPOLinet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib51)\)estimate token importance by contrasting probabilities between reference/negative models and oracle/positive models\. TI\-DPOYanget al\.\([2026](https://arxiv.org/html/2605.21883#bib.bib8)\)combines Gaussian priors with logit gradient norms\. Alternatively, heuristic approaches like D2POShaoet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib13)\)decay weights based on token position\. Nonetheless, these methods either demand expensive auxiliary models or rely on rigid, position\-based heuristics rather than semantic content\.

## 6Conclusion

We propose TwDPO—a token\-weighted DPO objective, theoretically grounded on token\-weighted RL—and AttentionPO—an instantiation of TwDPO that prompts the model to judge the response qualities and extract attentions from the model itself to weigh the response tokens\. Experiment results show that AttentionPO significantly improves the performance of the base model across AlpacaEval, MT\-Bench, and ArenaHard, outperforming various prior PO methods\. We hope AttentionPO paves the way for highly\-performant and efficient preference alignment with humans\.

## Limitations

Due to computational constraints, we do not experiment with our method on larger models\. In addition, our training and evaluation focus on the instruction\-following task\. Analysing the performance and generalizability of AttentionPO on other models and tasks is an important future direction\.

Moreover, the search space for the optimal choice of attention weights for TwDPO is large\. For example, while we simply take the mean across attention heads, future work could investigate selecting a specific head or group of heads for weight extraction\. However, we argue that such a choice of attention heads is likely model\-dependent and may not transfer across settings\.

## References

- Quantifying attention flow in transformers\.InProceedings of the 58th Annual Meeting of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3)\.
- J\. Achiam, S\. Adler, S\. Agarwal, L\. Ahmad, I\. Akkaya, F\. L\. Aleman, D\. Almeida, J\. Altenschmidt, S\. Altman, and S\. A\. et al\. \(2024\)GPT\-4o system card\.arXiv preprint arXiv:2410\.21276\.Cited by:[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1)\.
- M\. G\. Azar, M\. Rowland, B\. Piot, D\. Guo, D\. Calandriello, M\. Valko, and R\. Munos \(2024\)A general theoretical paradigm to understand learning from human preferences\.InProceedings of the 27th International Conference on Artificial Intelligence and Statistics,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§1](https://arxiv.org/html/2605.21883#S1.p2.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- R\. A\. Bradley and M\. E\. Terry \(1952\)Rank analysis of incomplete block designs: i\. the method of paired comparisons\.Biometrika39,pp\. 324–345\.External Links:[Document](https://dx.doi.org/https%3A//doi.org/10.2307/2334029)Cited by:[§2\.2](https://arxiv.org/html/2605.21883#S2.SS2.p1.5)\.
- K\. Clark, U\. Khandelwal, O\. Levy, and C\. D\. Manning \(2019\)What does bert look at?\.InProceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- G\. Cui, L\. Yuan, N\. Ding, G\. Yao, B\. He, W\. Zhu, Y\. Ni, G\. Xie, R\. Xie, Y\. Lin, Z\. Liu, and M\. Sun \(2024\)UltraFeedback: boosting language models with scaled ai feedback\.InProceedings of International Conference on Machine Learning,Cited by:[§F\.1](https://arxiv.org/html/2605.21883#A6.SS1.p1.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p2.2)\.
- N\. Ding, Y\. Chen, B\. Xu, Y\. Qin, S\. Hu, Z\. Liu, M\. Sun, and B\. Zhou \(2023\)Enhancing chat language models by scaling high\-quality instructional conversations\.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing,Cited by:[§F\.1](https://arxiv.org/html/2605.21883#A6.SS1.p1.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p1.1)\.
- Y\. Dubois, B\. Galambosi, P\. Liang, and T\. B\. Hashimoto \(2024\)Length\-controlled alpacaeval: a simple way to debias automatic evaluators\.InProceedings of Conference on Language Modeling,Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1)\.
- K\. Ethayarajh, W\. Xu, N\. Muennighoff, D\. Jurafsky, and D\. Kiela \(2024\)KTO: model alignment as prospect theoretic optimization\.InProceedings of International Conference on Machine Learning,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§1](https://arxiv.org/html/2605.21883#S1.p2.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle, A\. Letman, A\. Mathur, A\. Schelten, and A\. V\. et al\. \(2024\)The llama 3 herd of models\.arXiv preprint arXiv:2407\.21783\.Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p1.1)\.
- J\. Hong, N\. Lee, and J\. Thorne \(2024\)ORPO: monolithic preference optimization without reference model\.InProceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- C\. Huang and T\. Goyal \(2025\)DCRM: a heuristic to measure response pair quality in preference optimization\.InFindings of the Association for Computational Linguistics: EMNLP 2025,Cited by:[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1)\.
- A\. Hurst, A\. Lerer, A\. P\. Goucher, A\. Perelman, A\. Ramesh, A\. Clark, A\. Ostrow, A\. Welihinda, A\. Hayes, and A\. R\. et al\. \(2023\)GPT\-4 technical report\.arXiv preprint arXiv:2303\.08774\.Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1)\.
- Jain and Wallace \(2019\)Attention is not explanation\.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing \(EMNLP\-IJCNLP\),Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- A\. Q\. Jiang, A\. Sablayrolles, A\. Mensch, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, F\. Bressand, G\. Lengyel, G\. Lample, L\. Saulnier, L\. R\. Lavaud, M\. Lachaux, P\. Stock, T\. L\. Scao, T\. Lavril, T\. Wang, T\. Lacroix, and W\. E\. Sayed \(2023a\)Mistral 7b\.arXiv preprint arXiv:2310\.06825\.Cited by:[§F\.1](https://arxiv.org/html/2605.21883#A6.SS1.p1.1)\.
- D\. Jiang, X\. Ren, and B\. Y\. Lin \(2023b\)LLM\-blender: ensembling large language models with pairwise ranking and generative fusion\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),Cited by:[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p3.3)\.
- T\. Li, W\. Chiang, E\. Frick, L\. Dunlap, T\. Wu, B\. Zhu, J\. E\. Gonzalez, and I\. Stoica \(2025\)From crowdsourced data to high\-quality benchmarks: arena\-hard and benchbuilder pipeline\.InProceedings of International Conference on Machine Learning,Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1)\.
- T\. Li, W\. Chiang, E\. Frick, L\. Dunlap, B\. Zhu, J\. E\. Gonzalez, and I\. Stoica \(2024\)From live data to high\-quality benchmarks: the arena\-hard pipeline\.External Links:[Link](https://lmsys.org/blog/2024-04-19-arena-hard/)Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1)\.
- X\. Li, T\. Zhang, Y\. Dubois, R\. Taori, I\. Gulrajani, C\. Guestrin, P\. Liang, and T\. B\. Hashimoto \(2023\)AlpacaEval: an automatic evaluator of instruction\-following models\.GitHub\.Note:[https://github\.com/tatsu\-lab/alpaca\_eval](https://github.com/tatsu-lab/alpaca_eval)Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1)\.
- Z\. Lin, T\. Liang, J\. Xu, Q\. Lin, X\. Wang, R\. Luo, C\. Shi, S\. Li, Y\. Yang, and Z\. Tu \(2024\)Critical tokens matter: token\-level contrastive estimation enhances llm’s reasoning capability\.arXiv preprint arXiv:2411\.19943\.Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- A\. Liu, H\. Bai, Z\. Lu, Y\. Sun, X\. Kong, S\. Wang, J\. Shan, A\. M\. Jose, X\. Liu, L\. Wen, P\. S\. Yu, and M\. Cao \(2025\)TIS\-dpo: token\-level importance sampling for direct preference optimization with estimated weights\.InProceedings of International Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- I\. Loshchilov and F\. Hutter \(2019\)Decoupled weight decay regularization\.InProceedings of International Conference on Learning Representations,Cited by:[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px3.p1.3)\.
- R\. T\. McCoy, E\. Pavlick, and T\. Linzen \(2019\)Right for the wrong reasons: diagnosing syntactic heuristics in natural language inference\.InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- Y\. Meng, M\. Xia, and D\. Chen \(2024\)SimPO: simple preference optimization with a reference\-free reward\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[Appendix E](https://arxiv.org/html/2605.21883#A5.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p2.1),[§1](https://arxiv.org/html/2605.21883#S1.p5.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p1.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px1.p3.3),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1),[§4\.3](https://arxiv.org/html/2605.21883#S4.SS3.SSS0.Px2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- L\. Ouyang, J\. Wu, X\. Jiang, D\. Almeida, C\. L\. Wainwright, P\. Mishkin, C\. Zhang, S\. Agarwal, K\. Slama, A\. Ray, J\. Schulman, J\. Hilton, F\. Kelton, L\. Miller, M\. Simens, A\. Askell, P\. Welinder, P\. Christiano, J\. Leike, and R\. Lowe \(2022\)Training language models to follow instructions with human feedback\.arXiv preprint arXiv:2203\.02155\.Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§1](https://arxiv.org/html/2605.21883#S1.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- R\. Park, R\. Rafailov, S\. Ermon, and C\. Finn \(2024\)Disentangling length from quality in direct preference optimization\.InProceedings of Findings of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- R\. Rafailov, A\. Sharma, E\. Mitchell, S\. Ermon, C\. D\. Manning, and C\. Finn \(2023\)Direct preference optimization: your language model is secretly a reward model\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[Appendix B](https://arxiv.org/html/2605.21883#A2.p1.5),[§1](https://arxiv.org/html/2605.21883#S1.p1.1),[§2\.4\.1](https://arxiv.org/html/2605.21883#S2.SS4.SSS1.p1.1),[§2\.4](https://arxiv.org/html/2605.21883#S2.SS4.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov \(2017\)Proximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§1](https://arxiv.org/html/2605.21883#S1.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- S\. Serrano and N\. A\. Smith \(2019\)Is attention interpretable?\.InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- R\. Shao, B\. Li, G\. Liu, Y\. Chen, X\. Zhou, J\. Wang, X\. Cai, and P\. Li \(2025\)Earlier tokens contribute more: learning direct preference optimization from temporal decay perspective\.InProceedings of International Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- L\. Tunstall, E\. Beeching, N\. Lambert, N\. Rajani, K\. Rasul, Y\. Belkada, S\. Huang, L\. von Werra, C\. Fourrier, and e\. a\. Nathan Habib \(2023\)Zephyr: direct distillation of lm alignment\.\.arXiv preprint arXiv:2310\.16944\.Cited by:[§F\.1](https://arxiv.org/html/2605.21883#A6.SS1.p1.1)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. Polosukhin \(2017\)Attention is all you need\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- J\. Vig \(2019\)A multiscale visualization of attention in the transformer model\.InProceedings of the 57th Annual Meeting of the Association for Computational Linguistics: System Demonstrations,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- G\. Wiedemann, S\. Remus, A\. Chawla, and C\. Biemann \(2019\)Does bert make any sense? interpretable word sense disambiguation with contextualized embeddings\.InConference on Natural Language Processing,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px3.p1.1)\.
- G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. Lewis \(2024\)Efficient streaming language models with attention sinks\.InProceedings of International Conference on Machine Learning,Cited by:[§2\.6](https://arxiv.org/html/2605.21883#S2.SS6.SSS0.Px2.p1.5)\.
- H\. Xu, A\. Sharaf, Y\. Chen, W\. Tan, L\. Shen, B\. V\. Durme, K\. Murray, and Y\. J\. Kim \(2024\)Contrastive preference optimization: pushing the boundaries of llm performance in machine translation\.InProceedings of International Conference on Machine Learning,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- K\. Yang, Z\. Liu, Q\. Xie, J\. Huang, E\. Min, and S\. Ananiadou \(2025\)Selective preference optimization via token\-level reward function estimation\.InProceedings of Empirical Methods in Natural Language Processing,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3),[§4\.3](https://arxiv.org/html/2605.21883#S4.SS3.SSS0.Px2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- N\. Yang, H\. Lin, Y\. Liu, B\. Tian, G\. Liu, and H\. Zhang \(2026\)Token\-importance guided direct preference optimization\.InProceedings of International Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§4\.2](https://arxiv.org/html/2605.21883#S4.SS2.p1.3),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- Z\. Yuan, H\. Yuan, C\. Tan, W\. Wang, S\. Huang, and F\. Huang \(2023\)RRHF: rank responses to align language models with human feedback without tears\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- Y\. Zeng, G\. Liu, W\. Ma, N\. Yang, H\. Zhang, and J\. Wang \(2024\)Token\-level direct preference optimization\.InProceedings of International Conference on Machine Learning,Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2605.21883#S1.p3.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px2.p1.1)\.
- Y\. Zhao, R\. Joshi, T\. Liu, M\. Khalman, M\. Saleh, and P\. J\. Liu \(2023\)SLiC\-hf: sequence likelihood calibration with human feedback\.arXiv preprint arXiv:2305\.10425\.Cited by:[Appendix A](https://arxiv.org/html/2605.21883#A1.SS0.SSS0.Px1.p1.2),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px4.p1.1),[§5](https://arxiv.org/html/2605.21883#S5.SS0.SSS0.Px1.p1.1)\.
- L\. Zheng, W\. Chiang, Y\. Sheng, S\. Zhuang, Z\. Wu, Y\. Zhuang, Z\. Lin, Z\. Li, D\. Li, E\. P\. Xing, H\. Zhang, J\. E\. Gonzalez, and I\. Stoica \(2023\)Judging llm\-as\-a\-judge with mt\-bench and chatbot arena\.InAdvances in Neural Information Processing Systems, Datasets and Benchmarks Track,Cited by:[§1](https://arxiv.org/html/2605.21883#S1.p5.1),[§3](https://arxiv.org/html/2605.21883#S3.SS0.SSS0.Px5.p1.1)\.

## Appendix AExtended Related Work

##### Preference Optimization\.

The alignment with human preferences and values is critical for LLM\-generated output\. Early workOuyanget al\.\([2022](https://arxiv.org/html/2605.21883#bib.bib24)\)trains LLM using policy\-gradient RL\-algorithm such as PPOSchulmanet al\.\([2017](https://arxiv.org/html/2605.21883#bib.bib23)\)\. However, these methods require training a separate reward model on annotated preference pairs to provide RL rewards\. Preference Optimization \(PO\) methods avoid the need for an extra reward model and train the main policy directly on the preference pairs to learn the reward landscape\. On a high level, PO methods use a contrastive learning objective between a preferred responseywy\_\{w\}and a dispreferred responseyly\_\{l\}\. For example, RRHFYuanet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib20)\)and SLiC\-HFZhaoet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib22)\)use ranking losses\. DPORafailovet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib9)\)drives its training objective from policy gradient, which normalizes the model’s probability with the reference model’s probability\. IPOAzaret al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib16)\)uses a bounded objective in contrast to DPO\. KTOEthayarajhet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib15)\)avoids the pairwise assumption and uses a pointwise formulation\. CPOXuet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib19)\)and ORPOHonget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib18)\)add an SFT term in addition to the contrastive term\. R\-DPOParket al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib17)\)adds an extra term to the original DPO objective to avoid the length exploitation\. SimPOMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)modifies DPO by removing dependencies on the reference model, incorporating length normalization, and adding a reward margin term\. However, all the methods above assign equal weights to every response token, which cannot account for the fine\-grained importance of each token to the overall response quality\.

##### Token\-level PO\.

Various attempts are made to enable fine\-grained, token\-level credit assignment during PO training\. TDPOZenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib12)\)is grounded on token\-level policy gradient methods\. It adds two sequence KL terms to the original DPO objective\. TIS\-DPOLiuet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib10)\)and cDPOLinet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib51)\)estimate the token importance via contrasting the probability to produce the token from a positive policy with the one from a negative policy, given the same prefix\. The positive and negative policies are induced either from prompts or training\. SePOYanget al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib11)\)takes a similar approach, but it trains an oracle model with DPO on a subset of the original training set and contrasts the probability from the oracle against the one from the reference model\. TI\-DPOYanget al\.\([2026](https://arxiv.org/html/2605.21883#bib.bib8)\)estimates the token weights by combining a pre\-defined Gaussian prior distribution with the gradient norm of the last token’s maximum logit with respect to each token\. Certain methods use heuristics to assign weights\. For example, D2POShaoet al\.\([2025](https://arxiv.org/html/2605.21883#bib.bib13)\)gives earlier tokens more weight than later ones\. However, many of these methods require training extra models to estimate weights\. The rest compute weights using heuristic functions, which are purely based on token positions rather than the tokens themselves\.

##### Attention Weights\.

Early studies on transformersVaswaniet al\.\([2017](https://arxiv.org/html/2605.21883#bib.bib33)\)indicate that specific attention heads capture linguistic properties like coreference resolution, framing attention weights as transparent windows into model reasoningClarket al\.\([2019](https://arxiv.org/html/2605.21883#bib.bib34)\)\. However, subsequent research challenges this notion by demonstrating that attention distributions can be manipulated without altering model predictions, implying that attention correlates with but does not cause model outputsJain and Wallace \([2019](https://arxiv.org/html/2605.21883#bib.bib35)\); Serrano and Smith \([2019](https://arxiv.org/html/2605.21883#bib.bib36)\)\. Counter\-counter\-arguments suggest attention remains a useful, albeit incomplete, diagnostic signal when contextualizedWiedemannet al\.\([2019](https://arxiv.org/html/2605.21883#bib.bib40)\)\. While individual attention heads specializeVig \([2019](https://arxiv.org/html/2605.21883#bib.bib38)\), they operate redundantly\. Furthermore, attention patterns are frequently associated with superficial heuristics rather than deep semantic understandingMcCoyet al\.\([2019](https://arxiv.org/html/2605.21883#bib.bib39)\)\. When aggregated across multiple layers with residual connections, the aggregated attentions show a higher correlation with gold token importance scores, compared with raw attentionsAbnar and Zuidema \([2020](https://arxiv.org/html/2605.21883#bib.bib41)\)\. These findings suggest that attention could potentially provide a complementary view of information flow\. In summary, whether attention is interpretable is still under debate\. As such, in this work, we do not explicitly assume the interpretability of attentions\.

## Appendix BTheoretical Analysis and Error Bounds

Because the sequence\-level expectation𝔼y∼π\\mathbb\{E\}\_\{y\\sim\\pi\}in the TwDPO objective is coupled with token\-weighted log\-probabilities, deriving an exact analytical closed\-form solution for the optimal policy is intractable without invoking complex token\-level dynamic programming\. In this section, we establish a formal theoretical guarantee for TwDPO by framing it as a bounded perturbation around the standard Direct Preference Optimization \(DPO\) objective\(Rafailovet al\.,[2023](https://arxiv.org/html/2605.21883#bib.bib9)\)\. We first bound the divergence between the true, intractable TwDPO optimal policyπopt\\pi\_\{\\text\{opt\}\}and the standard DPO optimal policyπDPO\\pi\_\{\\text\{DPO\}\}\. We then bound the distance between our empirical autoregressive heuristic policyπ~∗\\tilde\{\\pi\}^\{\*\}andπDPO\\pi\_\{\\text\{DPO\}\}, concluding with a unified Total Variation \(TV\) distance bound between the empirical heuristic and the true optimum\.

### B\.1Objective Formulation and Perturbation Setup

Let𝒟\\mathcal\{D\}denote the dataset of promptsxx, and lety=\(y1,y2,…,y\|y\|\)y=\(y^\{1\},y^\{2\},\\dots,y^\{\|y\|\}\)represent a generated sequence of length\|y\|\|y\|\. For a given promptxx, policyπ\\pi, and reference policyπref\\pi\_\{\\text\{ref\}\}, the TwDPO objectiveJTwDPO​\(π\)J\_\{\\text\{TwDPO\}\}\(\\pi\)modifies the standard sequence\-level KL penalty by weighting individual token steps using attention weightsata^\{t\}:

JTwDPO​\(π\)=𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)\[r​\(x,y\)−β​\|y\|​∑t=1\|y\|at​log⁡π​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)\],\\begin\{split\}&J\_\{\\text\{TwDPO\}\}\(\\pi\)=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\\\ &~~~~~~~~~~\\left\[r\(x,y\)\-\\beta\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\frac\{\\pi\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\\right\],\\end\{split\}\(16\)whereβ\>0\\beta\>0is the global scaling parameter\. By definition, the token attention weights are normalized such that∑t=1\|y\|at=1\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}=1\. Let us define the effective token weights aswt=\|y\|​atw^\{t\}=\|y\|a^\{t\}, satisfying1\|y\|​∑t=1\|y\|wt=1\\frac\{1\}\{\|y\|\}\\sum\_\{t=1\}^\{\|y\|\}w^\{t\}=1\.

We parameterize the deviation of these attention weights from a uniform distribution as a perturbation vectorϵt\\epsilon^\{t\}:

wt=1\+ϵt,where∑t=1\|y\|ϵt=0\.w^\{t\}=1\+\\epsilon^\{t\},\\quad\\text\{where\}\\quad\\sum\_\{t=1\}^\{\|y\|\}\\epsilon^\{t\}=0\.\(17\)
###### Assumption 1\(Bounded Attention Deviation\)\.

There exists a constantδ≥0\\delta\\geq 0such that the maximum token\-level attention perturbation is strictly bounded for all sequences:

maxt⁡\|ϵt\|≤δ\.\\max\_\{t\}\|\\epsilon^\{t\}\|\\leq\\delta\.\(18\)

###### Assumption 2\(Bounded Trust Region Log\-Ratios\)\.

The log\-probability ratio between any valid policyπ\\piunder consideration and the reference policyπref\\pi\_\{\\text\{ref\}\}is bounded at the token level by a finite constantC\>0C\>0:

maxt⁡\|log⁡π​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)\|≤C\.\\max\_\{t\}\\left\|\\log\\frac\{\\pi\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\\right\|\\leq C\.\(19\)

Using this formulation, we decompose the TwDPO objective into the standard DPO objectiveJDPO​\(π\)J\_\{\\text\{DPO\}\}\(\\pi\)and a descriptive perturbation term:

JTwDPO​\(π\)=𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)\[r​\(x,y\)−β​∑t=1\|y\|\(1\+ϵt\)​log⁡π​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)\]=JDPO​\(π\)−β​𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)​\[Rϵ​\(π;x,y\)\],\\begin\{split\}J\_\{\\text\{TwDPO\}\}\(\\pi\)&=\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\\\ &\\left\[r\(x,y\)\-\\beta\\sum\_\{t=1\}^\{\|y\|\}\(1\+\\epsilon^\{t\}\)\\log\\frac\{\\pi\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\\right\]\\\\ &=J\_\{\\text\{DPO\}\}\(\\pi\)\-\\beta\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\\left\[R\_\{\\epsilon\}\(\\pi;x,y\)\\right\],\\end\{split\}where the point\-wise sequence perturbationRϵ​\(π;x,y\)R\_\{\\epsilon\}\(\\pi;x,y\)is defined as:

Rϵ​\(π;x,y\)=∑t=1\|y\|ϵt​log⁡π​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)\.R\_\{\\epsilon\}\(\\pi;x,y\)=\\sum\_\{t=1\}^\{\|y\|\}\\epsilon^\{t\}\\log\\frac\{\\pi\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\.\(20\)

### B\.2Bounding the Objective Discrepancy

We now bound the magnitude of the expected perturbation term\. Under Assumptions[1](https://arxiv.org/html/2605.21883#Thmassumption1)and[2](https://arxiv.org/html/2605.21883#Thmassumption2), the sequence\-level perturbation can be upper\-bounded via the triangle inequality:

\|Rϵ​\(π;x,y\)\|≤∑t=1\|y\|\|ϵt\|⋅\|log⁡π​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)\|≤∑t=1\|y\|δ​C=\|y\|​δ​C\.\\begin\{split\}\|R\_\{\\epsilon\}\(\\pi;x,y\)\|&\\leq\\sum\_\{t=1\}^\{\|y\|\}\|\\epsilon^\{t\}\|\\cdot\\left\|\\log\\frac\{\\pi\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\\right\|\\\\ &\\leq\\sum\_\{t=1\}^\{\|y\|\}\\delta C=\|y\|\\delta C\.\\end\{split\}\(21\)
Taking the expectation over the prompt and sequence distributions, the discrepancy between the two objectives is strictly bounded by:

\|JTwDPO​\(π\)−JDPO​\(π\)\|≤β​δ​C​𝔼x∼𝒟​𝔼y∼π\(⋅\|x\)​\[\|y\|\]\.\\left\|J\_\{\\text\{TwDPO\}\}\(\\pi\)\-J\_\{\\text\{DPO\}\}\(\\pi\)\\right\|\\leq\\beta\\delta C\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{y\\sim\\pi\(\\cdot\|x\)\}\[\|y\|\]\.\(22\)For ease of notation, letB​\(π\)=β​δ​C​𝔼π​\[\|y\|\]B\(\\pi\)=\\beta\\delta C\\mathbb\{E\}\_\{\\pi\}\[\|y\|\]denote the policy\-dependent bound constraint\.

### B\.3Distance Bound for True Optima

Letπopt=arg⁡maxπ⁡JTwDPO​\(π\)\\pi\_\{\\text\{opt\}\}=\\arg\\max\_\{\\pi\}J\_\{\\text\{TwDPO\}\}\(\\pi\)be the true, intractable optimal policy under the token\-weighted framework, and letπDPO=arg⁡maxπ⁡JDPO​\(π\)\\pi\_\{\\text\{DPO\}\}=\\arg\\max\_\{\\pi\}J\_\{\\text\{DPO\}\}\(\\pi\)be the analytical optimal policy for the standard DPO objective\.

From the exact mathematical properties of KL\-constrained reinforcement learning objectives, the suboptimality gap of any arbitrary policyπ\\pievaluated underJDPOJ\_\{\\text\{DPO\}\}is exactly equal to its generalized KL divergence from the optimal DPO policy:

JDPO​\(πDPO\)−JDPO​\(π\)=β𝔼x∼𝒟\[DKL\(π\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]\.\\begin\{split\}&J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\-J\_\{\\text\{DPO\}\}\(\\pi\)\\\\ &=\\beta\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\left\[D\_\{\\text\{KL\}\}\(\\pi\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\)\\right\]\.\\end\{split\}\(23\)
Evaluating the true optimal policyπopt\\pi\_\{\\text\{opt\}\}using the identity in Equation[23](https://arxiv.org/html/2605.21883#A2.E23)yields:

JDPO​\(πDPO\)−JDPO​\(πopt\)=β𝔼x∼𝒟\[DKL\(πopt\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]\.\\begin\{split\}&J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\-J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{opt\}\}\)\\\\ &=\\beta\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\left\[D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{opt\}\}\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\)\\right\]\.\\end\{split\}\(24\)
By the definition of functional optimality,πopt\\pi\_\{\\text\{opt\}\}maximizesJTwDPOJ\_\{\\text\{TwDPO\}\}, meaningJTwDPO​\(πopt\)≥JTwDPO​\(πDPO\)J\_\{\\text\{TwDPO\}\}\(\\pi\_\{\\text\{opt\}\}\)\\geq J\_\{\\text\{TwDPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\. We construct a double inequality chain by introducing the objective boundariesB​\(πopt\)B\(\\pi\_\{\\text\{opt\}\}\)andB​\(πDPO\)B\(\\pi\_\{\\text\{DPO\}\}\):

JDPO​\(πopt\)\+B​\(πopt\)≥JTwDPO​\(πopt\)≥JTwDPO​\(πDPO\)≥JDPO​\(πDPO\)−B​\(πDPO\)\.\\begin\{split\}J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{opt\}\}\)\+B\(\\pi\_\{\\text\{opt\}\}\)&\\geq J\_\{\\text\{TwDPO\}\}\(\\pi\_\{\\text\{opt\}\}\)\\\\ &\\geq J\_\{\\text\{TwDPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\\\\ &\\geq J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\-B\(\\pi\_\{\\text\{DPO\}\}\)\.\\end\{split\}\(25\)
Rearranging the outermost components of this inequality gives an upper bound on the DPO objective gap:

JDPO​\(πDPO\)−JDPO​\(πopt\)≤B​\(πDPO\)\+B​\(πopt\)\.J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{DPO\}\}\)\-J\_\{\\text\{DPO\}\}\(\\pi\_\{\\text\{opt\}\}\)\\leq B\(\\pi\_\{\\text\{DPO\}\}\)\+B\(\\pi\_\{\\text\{opt\}\}\)\.\(26\)
Substituting this directly back into the identity in Equation[24](https://arxiv.org/html/2605.21883#A2.E24), we obtain:

β𝔼x∼𝒟\[DKL\(πopt\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]≤β​δ​C​\(𝔼πDPO​\[\|y\|\]\+𝔼πopt​\[\|y\|\]\)\.\\begin\{split\}\\beta\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\left\[D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{opt\}\}\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\)\\right\]\\\\ \\leq\\beta\\delta C\\left\(\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\pi\_\{\\text\{opt\}\}\}\[\|y\|\]\\right\)\.\\end\{split\}\(27\)
Dividing both sides by the regularization parameterβ\\betaestablishes our first localized distance guarantee\.

###### Lemma 1\(Optima Distance Bound\)\.

The expected KL divergence between the true optimal token\-weighted policyπopt\\pi\_\{\\text\{opt\}\}and the standard DPO policyπDPO\\pi\_\{\\text\{DPO\}\}is linearly bounded by the maximum attention perturbationδ\\delta:

𝔼x∼𝒟\[DKL\(πopt\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]≤δ​C​\(𝔼πDPO​\[\|y\|\]\+𝔼πopt​\[\|y\|\]\)\.\\begin\{split\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\left\[D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{opt\}\}\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\)\\right\]\\\\ \\leq\\delta C\\left\(\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\pi\_\{\\text\{opt\}\}\}\[\|y\|\]\\right\)\.\\end\{split\}\(28\)

### B\.4Bounding the Autoregressive Heuristic Policy

We now address our empirical heuristic policy, denoted asπ~∗\\tilde\{\\pi\}^\{\*\}\. Whileπ~∗\\tilde\{\\pi\}^\{\*\}is derived via a localized parameterization mapping, it is generated sequentially in an autoregressive fashion, ensuring thatπ~∗​\(y∣x\)=∏t=1\|y\|π~∗​\(yt∣x,y<t\)\\tilde\{\\pi\}^\{\*\}\(y\\mid x\)=\\prod\_\{t=1\}^\{\|y\|\}\\tilde\{\\pi\}^\{\*\}\(y^\{t\}\\mid x,y^\{<t\}\)constitutes a mathematically proper, normalized probability distribution over sequence space \(∑yπ~∗​\(y∣x\)=1\\sum\_\{y\}\\tilde\{\\pi\}^\{\*\}\(y\\mid x\)=1\)\.

By definition, our empirical token\-weighted optimal parameterization satisfies the equality:

∑t=1\|y\|wt​log⁡π~∗​\(yt∣x,y<t\)πref​\(yt∣x,y<t\)=1β​r​\(x,y\)−log⁡Z​\(x\)\.\\sum\_\{t=1\}^\{\|y\|\}w^\{t\}\\log\\frac\{\\tilde\{\\pi\}^\{\*\}\(y^\{t\}\\mid x,y^\{<t\}\)\}\{\\pi\_\{\\text\{ref\}\}\(y^\{t\}\\mid x,y^\{<t\}\)\}=\\frac\{1\}\{\\beta\}r\(x,y\)\-\\log Z\(x\)\.\(29\)Expandingwt=1\+ϵtw^\{t\}=1\+\\epsilon^\{t\}yields:

log⁡π~∗​\(y∣x\)πref​\(y∣x\)\+Rϵ​\(π~∗;x,y\)=1β​r​\(x,y\)−log⁡Z​\(x\)\.\\log\\frac\{\\tilde\{\\pi\}^\{\*\}\(y\\mid x\)\}\{\\pi\_\{\\text\{ref\}\}\(y\\mid x\)\}\+R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)=\\frac\{1\}\{\\beta\}r\(x,y\)\-\\log Z\(x\)\.\(30\)
Recall that the standard sequence\-level DPO optimal policy satisfies:

log⁡πDPO​\(y∣x\)πref​\(y∣x\)=1β​r​\(x,y\)−log⁡ZDPO​\(x\)\.\\log\\frac\{\\pi\_\{\\text\{DPO\}\}\(y\\mid x\)\}\{\\pi\_\{\\text\{ref\}\}\(y\\mid x\)\}=\\frac\{1\}\{\\beta\}r\(x,y\)\-\\log Z\_\{\\text\{DPO\}\}\(x\)\.\(31\)Subtracting Equation[31](https://arxiv.org/html/2605.21883#A2.E31)from Equation[30](https://arxiv.org/html/2605.21883#A2.E30)isolates the log\-ratio between our heuristic policy and the standard DPO policy:

log⁡π~∗​\(y∣x\)πDPO​\(y∣x\)=log⁡ZDPO​\(x\)−log⁡Z​\(x\)−Rϵ​\(π~∗;x,y\)\.\\log\\frac\{\\tilde\{\\pi\}^\{\*\}\(y\\mid x\)\}\{\\pi\_\{\\text\{DPO\}\}\(y\\mid x\)\}=\\log Z\_\{\\text\{DPO\}\}\(x\)\-\\log Z\(x\)\-R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\.\(32\)
LetΔ​Z​\(x\)=log⁡ZDPO​\(x\)−log⁡Z​\(x\)\\Delta Z\(x\)=\\log Z\_\{\\text\{DPO\}\}\(x\)\-\\log Z\(x\)\. Taking the expectation of Equation[32](https://arxiv.org/html/2605.21883#A2.E32)underπ~∗\\tilde\{\\pi\}^\{\*\}yields the forward KL divergence:

DKL​\(π~∗∥πDPO\)=Δ​Z​\(x\)−𝔼y∼π~∗​\[Rϵ​\(π~∗;x,y\)\]\.D\_\{\\text\{KL\}\}\(\\tilde\{\\pi\}^\{\*\}\\parallel\\pi\_\{\\text\{DPO\}\}\)=\\Delta Z\(x\)\-\\mathbb\{E\}\_\{y\\sim\\tilde\{\\pi\}^\{\*\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\.\(33\)Similarly, taking the expectation of Equation[32](https://arxiv.org/html/2605.21883#A2.E32)under the standard DPO policy yields the reverse KL divergence:

−DKL​\(πDPO∥π~∗\)=Δ​Z​\(x\)−𝔼y∼πDPO​\[Rϵ​\(π~∗;x,y\)\]\.\-D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{DPO\}\}\\parallel\\tilde\{\\pi\}^\{\*\}\)=\\Delta Z\(x\)\-\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{DPO\}\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\.\(34\)
Subtracting Equation[34](https://arxiv.org/html/2605.21883#A2.E34)from Equation[33](https://arxiv.org/html/2605.21883#A2.E33)cancels the intractable partition scaling termΔ​Z​\(x\)\\Delta Z\(x\):

DKL​\(π~∗∥πDPO\)\+DKL​\(πDPO∥π~∗\)=𝔼y∼πDPO​\[Rϵ​\(π~∗;x,y\)\]−𝔼y∼π~∗​\[Rϵ​\(π~∗;x,y\)\]\.\\begin\{split\}&D\_\{\\text\{KL\}\}\(\\tilde\{\\pi\}^\{\*\}\\parallel\\pi\_\{\\text\{DPO\}\}\)\+D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{DPO\}\}\\parallel\\tilde\{\\pi\}^\{\*\}\)\\\\ &=\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{DPO\}\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\-\\mathbb\{E\}\_\{y\\sim\\tilde\{\\pi\}^\{\*\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\.\\end\{split\}\(35\)
Since the reverse KL divergence is strictly non\-negative \(DKL​\(πDPO∥π~∗\)≥0D\_\{\\text\{KL\}\}\(\\pi\_\{\\text\{DPO\}\}\\parallel\\tilde\{\\pi\}^\{\*\}\)\\geq 0\), we drop it to establish an upper bound, and apply the triangle inequality across expectations:

DKL​\(π~∗∥πDPO\)≤\|𝔼y∼πDPO​\[Rϵ​\(π~∗;x,y\)\]\|\+\|𝔼y∼π~∗​\[Rϵ​\(π~∗;x,y\)\]\|≤δ​C​𝔼y∼πDPO​\[\|y\|\]\+δ​C​𝔼y∼π~∗​\[\|y\|\]\.\\begin\{split\}D\_\{\\text\{KL\}\}\(\\tilde\{\\pi\}^\{\*\}\\parallel\\pi\_\{\\text\{DPO\}\}\)&\\leq\\left\|\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{DPO\}\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\\right\|\+\\\\ &~~~~~~~~~~~~\\left\|\\mathbb\{E\}\_\{y\\sim\\tilde\{\\pi\}^\{\*\}\}\\left\[R\_\{\\epsilon\}\(\\tilde\{\\pi\}^\{\*\};x,y\)\\right\]\\right\|\\\\ &\\leq\\delta C\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\delta C\\mathbb\{E\}\_\{y\\sim\\tilde\{\\pi\}^\{\*\}\}\[\|y\|\]\.\\end\{split\}
Taking the expectation across the entire prompt distribution𝒟\\mathcal\{D\}establishes our second localization lemma\.

###### Lemma 2\(Heuristic Distance Bound\)\.

The expected KL divergence between the empirical autoregressive heuristic policyπ~∗\\tilde\{\\pi\}^\{\*\}and the closed\-form standard DPO policyπDPO\\pi\_\{\\text\{DPO\}\}is strictly bounded by:

𝔼x∼𝒟\[DKL\(π~∗\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]≤δ​C​\(𝔼πDPO​\[\|y\|\]\+𝔼π~∗​\[\|y\|\]\)\.\\begin\{split\}\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\left\[D\_\{\\text\{KL\}\}\(\\tilde\{\\pi\}^\{\*\}\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\)\\right\]\\\\ \\leq\\delta C\\left\(\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\tilde\{\\pi\}^\{\*\}\}\[\|y\|\]\\right\)\.\\end\{split\}\(36\)

### B\.5Unified Suboptimality Guarantee via Total Variation Distance

We now unify Lemma[1](https://arxiv.org/html/2605.21883#Thmlemma1)and Lemma[2](https://arxiv.org/html/2605.21883#Thmlemma2)to bound the final distance between our empirical heuristic policyπ~∗\\tilde\{\\pi\}^\{\*\}and the true intractable optimal policyπopt\\pi\_\{\\text\{opt\}\}\. We map our KL bounds to Total Variation \(TV\) distance using Pinsker’s inequality, which states that for any distributionsPPandQQ,DTV​\(P,Q\)≤12​DKL​\(P∥Q\)D\_\{\\text\{TV\}\}\(P,Q\)\\leq\\sqrt\{\\frac\{1\}\{2\}D\_\{\\text\{KL\}\}\(P\\parallel Q\)\}\.

Applying Pinsker’s inequality to Lemma[1](https://arxiv.org/html/2605.21883#Thmlemma1)and Lemma[2](https://arxiv.org/html/2605.21883#Thmlemma2)gives:

𝔼x∼𝒟​\[DTV​\(πopt,πDPO\)\]≤𝔼x∼𝒟​\[12DKL\(πopt\(⋅∣x\)∥πDPO\(⋅∣x\)\)\]≤𝔼x∼𝒟​\[δ​C2​\(𝔼y∼πDPO​\[\|y\|\]\+𝔼y∼πopt​\[\|y\|\]\)\]≤δ​C2​\(𝔼πDPO​\[\|y\|\]\+𝔼πopt​\[\|y\|\]\),\\begin\{split\}&\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\[D\_\{\\text\{TV\}\}\(\\pi\_\{\\text\{opt\}\},\\pi\_\{\\text\{DPO\}\}\)\]\\\\ &\\leq\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\biggl\[\\sqrt\{\\frac\{1\}\{2\}D\_\{\\text\{KL\}\}\\left\(\\pi\_\{\\text\{opt\}\}\(\\cdot\\mid x\)\\parallel\\pi\_\{\\text\{DPO\}\}\(\\cdot\\mid x\)\\right\)\}\\biggr\]\\\\ &\\leq\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\\biggl\[\\sqrt\{\\frac\{\\delta C\}\{2\}\\left\(\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{y\\sim\\pi\_\{\\text\{opt\}\}\}\[\|y\|\]\\right\)\}\\biggr\]\\\\ &\\leq\\sqrt\{\\frac\{\\delta C\}\{2\}\\left\(\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\pi\_\{\\text\{opt\}\}\}\[\|y\|\]\\right\)\},\\end\{split\}\(37\)where the first inequality comes from Pinsker’s inequality, the second comes from equation[28](https://arxiv.org/html/2605.21883#A2.E28)with𝒟\\mathcal\{D\}containing only a single queryxx, the third comes from Jensen’s inequality\.

Similarly, applying Pinsker’s inequality, equation[36](https://arxiv.org/html/2605.21883#A2.E36)with𝒟\\mathcal\{D\}containing only a single queryxx, and Jensen’s inequality, we get

𝔼x∼𝒟​\[DTV​\(π~∗,πDPO\)\]≤δ​C2​\(𝔼πDPO​\[\|y\|\]\+𝔼π~∗​\[\|y\|\]\)\.\\begin\{split\}&\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\[D\_\{\\text\{TV\}\}\(\\tilde\{\\pi\}^\{\*\},\\pi\_\{\\text\{DPO\}\}\)\]\\leq\\sqrt\{\\frac\{\\delta C\}\{2\}\\left\(\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\tilde\{\\pi\}^\{\*\}\}\[\|y\|\]\\right\)\}\.\\end\{split\}\(38\)
Finally, by utilizing the triangle inequality property of Total Variation distances, we bound the direct distance between the empirical heuristic and the true functional optimum:

DTV​\(πopt,π~∗\)≤DTV​\(πopt,πDPO\)\+DTV​\(π~∗,πDPO\)\.D\_\{\\text\{TV\}\}\(\\pi\_\{\\text\{opt\}\},\\tilde\{\\pi\}^\{\*\}\)\\leq D\_\{\\text\{TV\}\}\(\\pi\_\{\\text\{opt\}\},\\pi\_\{\\text\{DPO\}\}\)\+D\_\{\\text\{TV\}\}\(\\tilde\{\\pi\}^\{\*\},\\pi\_\{\\text\{DPO\}\}\)\.\(39\)
Taking expectations over𝒟\\mathcal\{D\}and substituting Equations[37](https://arxiv.org/html/2605.21883#A2.E37)and[38](https://arxiv.org/html/2605.21883#A2.E38)leads directly to our main result\.

###### Theorem 1\(TwDPO Suboptimality Consistancy\)\.

Under Assumptions[1](https://arxiv.org/html/2605.21883#Thmassumption1)and[2](https://arxiv.org/html/2605.21883#Thmassumption2), the expected Total Variation distance between the empirical autoregressive heuristic policyπ~∗\\tilde\{\\pi\}^\{\*\}and the true, intractable token\-weighted optimal policyπopt\\pi\_\{\\text\{opt\}\}is strictly bounded by𝒪​\(δ\)\\mathcal\{O\}\(\\sqrt\{\\delta\}\):

𝔼x∼𝒟\[DTV\(πopt\(⋅∣x\),π~∗\(⋅∣x\)\)\]≤δ​C2\(𝔼πDPO​\[\|y\|\]\+𝔼πopt​\[\|y\|\]\+𝔼πDPO​\[\|y\|\]\+𝔼π~∗​\[\|y\|\]\)\.\\begin\{split\}&\\mathbb\{E\}\_\{x\\sim\\mathcal\{D\}\}\[D\_\{\\text\{TV\}\}\(\\pi\_\{\\text\{opt\}\}\(\\cdot\\mid x\),\\tilde\{\\pi\}^\{\*\}\(\\cdot\\mid x\)\)\]\\leq\\sqrt\{\\frac\{\\delta C\}\{2\}\}\\\\ &~~~~\\left\(\\sqrt\{\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\pi\_\{\\text\{opt\}\}\}\[\|y\|\]\}\+\\sqrt\{\\mathbb\{E\}\_\{\\pi\_\{\\text\{DPO\}\}\}\[\|y\|\]\+\\mathbb\{E\}\_\{\\tilde\{\\pi\}^\{\*\}\}\[\|y\|\]\}\\right\)\.\\end\{split\}\(40\)

This theorem provides a rigorous theoretical safety net for TwDPO\. It demonstrates that while a clean closed\-form sequence\-level derivation is structurally impossible due to token\-weight expectation coupling, the empirical objective remains mathematically consistent\. As long as the attention variance deviation from uniform is controlled \(δ→0\\delta\\to 0\), the empirical policy behaves predictably and stays bounded within a localized neighborhood of the true mathematical optimum\.

## Appendix CGradient Analysis

The gradient of objective[5](https://arxiv.org/html/2605.21883#S2.E5)is:

∇θℒ=−β𝔼\(x,yw,yl\)∼𝒟\[σ\(rθ′\(x,yl\)−rθ′\(x,yw\)\)\(\|yw\|∑t=1\|yw\|awt∇θlogπθ\(ywt\|x,yw<t\)−\|yl\|∑t=1\|yl\|alt∇θlogπθ\(ylt\|x,yl<t\)\)\],\\begin\{split\}\\nabla\_\{\\theta\}\\mathcal\{L\}&=\-\\beta\\mathbb\{E\}\_\{\(x,y\_\{w\},y\_\{l\}\)\\sim\\mathcal\{D\}\}\\biggl\[\\sigma\\biggl\(r^\{\\prime\}\_\{\\theta\}\(x,y\_\{l\}\)\-r^\{\\prime\}\_\{\\theta\}\(x,y\_\{w\}\)\\biggr\)\\\\ &\\biggl\(\|y\_\{w\}\|\\sum\_\{t=1\}^\{\|y\_\{w\}\|\}a\_\{w\}^\{t\}\\nabla\_\{\\theta\}log\\pi\_\{\\theta\}\(y\_\{w\}^\{t\}\|x,y\_\{w\}^\{<t\}\)\-\\\\ &~~~\|y\_\{l\}\|\\sum\_\{t=1\}^\{\|y\_\{l\}\|\}a\_\{l\}^\{t\}\\nabla\_\{\\theta\}log\\pi\_\{\\theta\}\(y\_\{l\}^\{t\}\|x,y\_\{l\}^\{<t\}\)\\biggr\)\\biggr\],\\end\{split\}\(41\)whererθ′​\(x,y\)=β​\|y\|​∑t=1\|y\|at​log⁡πθ​\(yt\|x,y<t\)πr​e​f​\(yt\|x,y<t\)r^\{\\prime\}\_\{\\theta\}\(x,y\)=\\beta\|y\|\\sum\_\{t=1\}^\{\|y\|\}a^\{t\}\\log\\frac\{\\pi\_\{\\theta\}\(y^\{t\}\|x,y^\{<t\}\)\}\{\\pi\_\{ref\}\(y^\{t\}\|x,y^\{<t\}\)\}\.

Similar to DPO, theσ​\(rθ′​\(x,yl\)−rθ′​\(x,yw\)\)\\sigma\(r^\{\\prime\}\_\{\\theta\}\(x,y\_\{l\}\)\-r^\{\\prime\}\_\{\\theta\}\(x,y\_\{w\}\)\)term scales up our combined gradient when the reward margin betweenywy\_\{w\}andyly\_\{l\}is smaller or negative\. The main difference from DPO is that we weigh the gradient of each token at positionttbyata^\{t\}, so that higher\-weighted tokens get more gradient updates and the model learns more on these tokens\.

## Appendix DToken Matching

We note that certain attention scores may not be matched to any tokens, and certain tokens may not be assigned an attention weight\. This is because when we extract the attention weights in §[2\.5](https://arxiv.org/html/2605.21883#S2.SS5), the responseyyis tokenized with the pairwise judge prompt as the surrounding context\. During training,yyis tokenized with a different surrounding context, and consequently, the first and last several tokens do not match the original tokens extracted from the pairwise judge prompt\. In case no perfect match is found, we use the "edit\_distance" Python library to match the attention weight to the response tokens during training\. For tokens without a match, we assign a weight of 0\. We observe that in most cases, more than 95% of the tokens are matched to attention weights\.

## Appendix EHyperparameters

Table 6:Training Hyperparameters##### Training\.

We tune our training hyperparameters in the rangeβ∈\{2e\-2, 1e\-2, 5e\-3, 1e\-3\}\\beta\\in\\\{\\text\{2e\-2, 1e\-2, 5e\-3, 1e\-3\}\\\}and show the final training hyperparameters for AttentionPOfor both LLaMA\-3\-8B\-Base\-SFT and LLaMA\-3\-8B\-Instruct in Figure[6](https://arxiv.org/html/2605.21883#A5.T6)\. We train for around 2K steps and validate every 500 steps to pick the best checkpoint for evaluation\. For the baselines, we take the existing checkpoints fromMenget al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib21)\)and use their hyperparameters\.

Table 7:Hyperparameters for AlpacaEvalTable 8:Hyperparameters for MT\-Bench\. We use the default temperature, which varies between different task subsets\.Table 9:Hyperparameters for ArenaHard
##### Evaluation\.

The default hyperparameters are used whenever they are available\. See the details in Figure[7](https://arxiv.org/html/2605.21883#A5.T7),[8](https://arxiv.org/html/2605.21883#A5.T8), and[9](https://arxiv.org/html/2605.21883#A5.T9)\.

## Appendix FAdditional Results

### F\.1Results on other models

We additionally experiment with zephyr\-7b\-sft\-fullTunstallet al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib47)\), which is a model supervised\-finetuned from Mistral\-7B\-v0\.1Jianget al\.\([2023a](https://arxiv.org/html/2605.21883#bib.bib49)\)on UltraChat\-200KDinget al\.\([2023](https://arxiv.org/html/2605.21883#bib.bib48)\)\. We denote this model asMistral\-7B\-Base\-SFT\. We use the same UltraFeedback binarized datasetCuiet al\.\([2024](https://arxiv.org/html/2605.21883#bib.bib45)\)\(HuggingFaceH4/ultrafeedback\_binarized\) and the same hyperparameters to train it\. The results are shown in Table[10](https://arxiv.org/html/2605.21883#A6.T10)\.

Table 10:Results on Mistral\-7B\-Base\-SFT\.Compared withπr​e​f\\pi\_\{ref\}, AttentionPO significantly enhances performances across different benchmarks\. We observe a 14% increase on AlpacaEval, a 0\.82 increase on MT\-Bench, and a 12% increase on ArenaHard\. In addition, AttentionPO outperforms DPO in terms of AlpacaEval and ArenaHard, but DPO has a higher MT\-Bench score\.

### F\.2Using attentions from other models

We investigate whether attention weights from one model can be used to train another model\. In particular, we use the attentions from LLaMA\-3\-8B\-Base\-SFT to train LLaMA\-3\-8B\-Instruct \(w/ attn from sft\) and vice versa \(w/ attn from inst\)\. The results are shown in Table[11](https://arxiv.org/html/2605.21883#A6.T11)\.

Table 11:Results of using attentions from other models\.Training the SFT model with attentions from the stronger Instruct model \(w/ attn from inst\) leads to a further increase in performance on AlpacaEval and ArenaHard on top of the original AttentionPO results\. In contrast, training the instruction model with attentions from the SFT model \(w/ attn from sft\) gives worse performance than the original AttentionPO, but the results are nonetheless better thanπr​e​f\\pi\_\{ref\}by a significant margin across benchmarks\. This suggests that attentions from a stronger model gives stronger performance\. AttentionPO can still deliver performance gains even with attentions from a model weaker thanπr​e​f\\pi\_\{ref\}, but to a lesser extent\.

## Appendix GExamples

We show an example of attention weights for LLaMA\-3\-8B\-Base\-SFT in Figure[4](https://arxiv.org/html/2605.21883#A7.F4)and LLaMA\-3\-8B\-Instruct in Figure[5](https://arxiv.org/html/2605.21883#A7.F5)\.

![Refer to caption](https://arxiv.org/html/2605.21883v1/x4.png)Figure 4:An example for LLaMA\-3\-8B\-Base\-SFT\.![Refer to caption](https://arxiv.org/html/2605.21883v1/x5.png)Figure 5:An example for LLaMA\-3\-8B\-Instruct\.
## Appendix HCompute

We run all experiments on a single node with 40 CPU cores, 256 GB of memory, and 2 Nvidia A100 80 GB GPUs\.

## Appendix IPrompt

![Refer to caption](https://arxiv.org/html/2605.21883v1/x6.png)Figure 6:Pairwise Judge Prompt![Refer to caption](https://arxiv.org/html/2605.21883v1/x7.png)Figure 7:Verbalized Self\-judge System Prompt![Refer to caption](https://arxiv.org/html/2605.21883v1/x8.png)Figure 8:Verbalized Self\-judge User Prompt