DimPO: Dimensionality Reduction for Attention using Preference Optimization

Vojtěch Lanz; Yufei Cui; Prasanna Parthasarathi

cs.LGSep 26, 2026

DimPO: Dimensionality Reduction for Attention using Preference Optimization

Authors: Vojtěch Lanz, Yufei Cui, Prasanna Parthasarathi

Organizations: Charles University, Faculty of Mathematics and Physics Institute of Formal and Applied Linguistics · Noah’s Ark Lab, Huawei Technologies Ltd.

Abstract

A linear projection can reduce the dimension of query and key vectors without updating the pretrained model, but it remains unclear which training objective best preserves model behavior. We ask whether preferences over keys and attention mass on the highest-weighted keys provide a better signal than matching the full attention distribution with KL divergence, especially in long-context settings. We introduce DimPO, which combines listwise preference optimization with a lightweight top-k cross-entropy term for head-fidelity. DimPO is trained offline from the attention patterns of a frozen language model, with one map per layer, shared by the query and the keys and trained separately from the other layers. Across LLaMA3.2-3B, LLaMA3.1-8B, Qwen2.5-7B, and Qwen3-4B Instruct models, pairwise preference objectives outperform the triplet baseline and retain 98% of the original score on short-context tasks when projecting to half the dimension on the last 40% of the layers. With more projected layers or on long-context RULER, they degrade rapidly. In contrast, KL and DimPO, which use every key during training, retain about 95% of the original RULER 4k score on the 8B model when projecting up to 50% of the layers. KL-based projections remain closer to the original attention distribution and attention output, yet DimPO achieves better downstream performance. Beyond 50% of projected layers, DimPO increasingly outperforms KL on tasks including SQuAD, common-word extraction, frequent-word extraction, and variable tracking. These results suggest that under dimensionality reduction, preserving the ordering and concentration of task-relevant attention can matter more than reproducing the full attention distribution.

Figures & tables

Figure 1: Siamese architecture applying the same linear map $F_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d^{\prime}}$ to a query and its keys, mapping them to a shared lower-dimensional space.

	64		32		16		8		4		2		1
	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE
Rand	12.87	0.079	15.02	0.096	16.06	0.111	16.58	0.112	16.83	0.116	16.94	0.117	16.84	0.115
PCA	3.10	0.010	5.96	0.011	7.07	0.011	7.15	0.011	6.49	0.014	6.27	0.016	6.45	0.019
Triplet	3.29	0.009	3.77	0.009	4.07	0.009	4.25	0.010	4.32	0.009	4.33	0.010	4.43	0.010
CPO	2.06	0.007	1.64	0.005	2.59	0.007	3.44	0.009	3.95	0.010	4.30	0.010	4.53	0.010
SimPO	2.88	0.008	2.14	0.007	1.97	0.006	2.19	0.007	2.57	0.008	2.91	0.012	3.50	0.013

Table 1: Attention fidelity of different projection approaches on Llama3.1-8B-Instruct, measured by KL divergence of the attention weights and MSE of the attention outputs, averaged across layers.

Figure 2: Average score across ARC-Challenge, HellaSwag, MMLU, TruthfulQA, and WinoGrande against the percentage of projected layers on Llama3.1-8B-Instruct, with

d^{\prime}{=}64

Figure 3: RULER 4k score against the percentage of projected layers for different projection objectives on Llama3.1-8B-Instruct (left) and Qwen3-4B-Instruct (right).

Method	NIAH	SQuAD	HotpotQA	VT	CWE	FWE	RULER Avg
Base ( $l{=}0$ )	99.9	81.0	63.2	99.8	99.6	92.3	95.0
DimPO ( $k{=}0$ , $\lambda{=}0$ )	27.1	60.1	37.2	28.8	63.1	56.7	35.6
DimPO ( $k{=}64$ , $\lambda{=}1$ )	51.3	60.3	38.4	39.6	61.2	71.5	52.4
KL	39.1	43.6	32.0	26.7	39.0	51.5	38.9

Table 2: RULER 4k subtasks on Llama3.1-8B-Instruct at

l{=}20

(

62.5\%

of the layers), with

d^{\prime}{=}64

Method	NIAH	SQuAD	HotpotQA	VT	CWE	FWE	RULER Avg
Base ( $l{=}0$ )	100.0	76.7	62.8	100.0	99.3	81.2	93.8
DimPO ( $k{=}0$ , $\lambda{=}0$ )	75.1	47.3	36.4	20.2	61.3	57.4	63.4
DimPO ( $k{=}64$ , $\lambda{=}1$ )	73.0	46.2	36.4	69.6	56.5	64.5	65.9
KL	62.1	32.7	28.0	11.2	31.3	53.1	50.2

Table 3: RULER 4k subtasks on Qwen3-4B-Instruct at

l{=}24

(

66.7\%

of the layers), with

d^{\prime}{=}64

	LLaMA 3B		LLaMA 8B		Qwen 4B
Layers	4k	8k	4k	8k	4k	8k
0%	92.5	87.7	95.0	93.9	93.8	93.2
25%	87.9	82.0	94.2	91.7	93.4	89.2
40%	88.2	80.1	93.8	90.6	88.6	81.6
50%	76.3	64.9	90.5	87.1	78.6	62.9

Table 4: RULER 4k and 8k scores for Llama3.2-3B-Instruct, Llama3.1-8B-Instruct, and Qwen3-4B-Instruct when applying DimPO (

k{=}64

\lambda{=}1

) with

d^{\prime}{=}64

25\%

38\%

, and

50\%

of the attention layers.

	0% ( $l{=}0$ )		25% ( $l{=}8$ )		38% ( $l{=}12$ )		50% ( $l{=}16$ )
Method	4k	8k	4k	8k	4k	8k	4k	8k
Base	95.0	93.9	94.2	91.7	93.8	90.6	90.5	87.1
KIVI 2-bit	94.5	92.6	91.6	86.5	89.3	84.9	83.4	76.2
KIVI 4-bit	95.0	93.9	94.1	91.7	93.8	90.9	90.6	86.7
SnapKV 1024	95.0	93.9	94.1	91.7	93.7	90.6	90.5	87.0
SnapKV 2048	95.0	93.9	94.1	91.7	93.7	90.6	90.5	87.0

Table 5: RULER 4k and 8k average for combining DimPO with SnapKV or KIVI on Llama3.1-8B-Instruct with

d^{\prime}{=}64

. Columns show the percentage of layers projected, with

l{=}0,8,12,16

. The

l{=}0

column corresponds to using SnapKV or KIVI without DimPO.

	LLaMA 8B		Qwen 4B
Objective	RULER 4k	KL	RULER 4k	KL
DimPO ( $k{=}0$ , $\lambda{=}0$ )	35.6	0.74	63.4	0.73
KL	38.9	0.26	50.2	0.27
DimPO ( $k{=}0$ , $\lambda{=}0$ ) + KL	35.8	0.27	44.4	0.27
DimPO ( $k{=}32$ , $\lambda{=}1$ )	48.5	0.68	62.9	0.65
DimPO ( $k{=}64$ , $\lambda{=}1$ )	52.4	0.57	65.9	0.53

Table 6: Ablation of DimPO on RULER 4k and attention KL at

d^{\prime}{=}64

. Results are reported at

l{=}20

for Llama3.1-8B-Instruct and

l{=}24

for Qwen3-4B-Instruct.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: Effect of key-pair selection on pairwise losses. (a) All key pairs from a sampled subset. The DimPO curve is the listwise term, DimPO ( $k{=}0$ , $\lambda{=}0$ ), and it keeps improving as more keys are included, whereas pairwise methods plateau or degrade. (b) Training with multiple non-overlapping pairs per query does not yield improvements, suggesting that additional pairs introduce noise rather than a meaningful signal. (c) Increasing the diversity between chosen and rejected keys consistently improves performance, indicating that pairwise losses benefit most from highly diverse key pairs.

Method	$\beta$	$\gamma$	Learning rate	Batch size	Time/layer
DimPO	1.0	0.0001	0.0001	1	$\sim$ 5 min
SimPO	1.0	1.0	0.001	32	$\sim$ 10 s
Triplet	–	–	0.0001	32	$\sim$ 10 s
ORPO	0.1	–	0.001	32	$\sim$ 10 s
CPO	1.0	0.1	0.0001	32	$\sim$ 10 s

Appendix

Table 7: Final hyperparameter settings for all preference optimization methods.

Hyperparameter	Tested values
$\beta$	0.0001, 0.001, 0.01, 0.1, 1.0, 2, 2.5, 5.0
$\gamma$	0, 0.00001, 0.0001, 0.001, 0.01, 0.1, 1.0
Learning rate	1e-5, 1e-4, 1e-3, 1e-2, 0.1
Batch size	1, 2, 4, 8, 16, 32, 64

Appendix

Table 8: Hyperparameter values explored during tuning for all methods, where applicable.

	64		32		16		8		4		2		1
	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE
Rand	12.56	0.084	14.50	0.104	15.78	0.119	16.52	0.129	16.60	0.128	16.82	0.132	16.70	0.126
PCA	2.65	0.011	5.71	0.013	7.28	0.013	7.57	0.013	6.66	0.016	6.44	0.019	6.49	0.021
Triplet	3.19	0.014	3.63	0.014	3.94	0.014	4.12	0.014	4.13	0.014	4.19	0.015	4.28	0.014
CPO	2.20	0.009	1.69	0.007	2.19	0.009	3.36	0.014	3.86	0.014	4.17	0.014	4.39	0.014
SimPO	2.90	0.010	2.12	0.009	1.95	0.008	2.20	0.009	2.58	0.011	2.92	0.014	3.69	0.019

Appendix

Table 9: Attention fidelity of different projection approaches on Llama3.2-3B-Instruct, measured by attention-weight KL divergence and attention-output MSE, averaged across layers.

	64		32		16		8		4		2		1
	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE
Rand	14.07	1.269	14.85	1.577	15.49	1.837	15.75	2.045	15.89	2.131	16.01	2.173	15.88	1.929
PCA	4.42	0.217	6.51	0.267	8.01	0.360	9.16	0.503	10.02	0.554	10.34	0.533	10.36	0.534
Triplet	2.50	0.149	2.91	0.169	3.20	0.174	3.38	0.177	3.55	0.179	3.63	0.183	3.69	0.180
CPO	2.52	0.216	2.14	0.187	2.55	0.190	3.04	0.161	3.38	0.184	3.63	0.184	3.80	0.193
SimPO	1.44	0.115	1.74	0.156	2.21	0.146	2.85	0.177	3.56	0.224	3.95	0.202	4.65	0.371

Appendix

Table 10: Attention fidelity of different projection approaches on Qwen3-4B-Instruct, measured by attention-weight KL divergence and attention-output MSE, averaged across layers.

	64		32		16		8		4		2		1
	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE	KL	MSE
Rand	14.65	1.181	15.58	1.459	15.82	1.594	16.02	1.789	16.08	1.911	16.09	1.739	15.96	1.705
PCA	3.65	0.826	7.49	0.891	10.39	0.936	11.32	0.939	11.40	0.985	11.43	1.138	11.41	1.205
Triplet	2.51	0.104	2.95	0.085	3.22	0.088	3.44	0.088	3.52	0.090	3.64	0.093	3.69	0.094
CPO	2.53	0.120	2.24	0.087	2.72	0.088	3.17	0.090	3.55	0.092	3.73	0.095	3.87	0.101
SimPO	1.58	0.094	1.86	0.083	2.38	0.096	3.21	0.129	3.95	0.146	4.45	0.178	4.69	0.210

Appendix

Table 11: Attention fidelity of different projection approaches on Qwen2.5-7B-Instruct, measured by attention-weight KL divergence and attention-output MSE, averaged across layers.

$d^{\prime}$ / $l$	0	2	4	8	12	16	20	24	28	30	32
64	66.2/0	66.0/1.6	65.8/3.1	65.5/6.2	65.0/9.4	63.6/12.5	60.2/15.6	57.6/18.8	52.0/21.9	48.0/23.4	43.9/25
16	66.2/0	65.1/2.7	65.2/5.5	64.3/10.9	62.4/16.4	53.7/21.9	43.7/27.3	37.9/32.8	34.7/38.3	34.7/41	34.7/43.8
4	66.2/0	64.9/3	64.5/6.1	62.4/12.1	55.2/18.2	42.2/24.2	37.0/30.3	34.9/36.3	34.9/42.4	34.8/45.4	35.1/48.4
1	66.2/0	64.8/3.1	64.2/6.2	60.3/12.4	52.3/18.6	39.7/24.8	35.2/31	34.5/37.2	34.8/43.4	35.4/46.5	35.8/49.6

Appendix

Table 12: Average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA and WinoGrande on Llama3.1-8B-Instruct. Each cell is the score and the percentage of the KV cache removed, for DimPO (

k{=}64

\lambda{=}1

Figure 5: Harness task scores for different projection methods on Llama3.1-8B-Instruct against the percentage of the KV cache removed.

$d^{\prime}$ / $l$	0	2	4	6	9	12	15	21	27	32	36
64	65.7/0	65.7/1.4	65.6/2.8	65.7/4.2	65.5/6.2	65.1/8.3	63.7/10.4	60.9/14.6	57.3/18.8	56.3/22.2	55.0/25
16	65.7/0	64.7/2.4	64.4/4.9	63.7/7.3	62.5/10.9	56.4/14.6	48.5/18.2	41.5/25.5	38.2/32.8	35.5/38.9	34.1/43.8
4	65.7/0	64.2/2.7	63.1/5.4	62.0/8.1	59.7/12.1	49.5/16.1	43.0/20.2	38.7/28.3	36.5/36.3	34.8/43.1	35.2/48.4
1	65.7/0	64.2/2.8	63.1/5.5	61.4/8.3	59.0/12.4	47.1/16.5	41.8/20.7	38.1/28.9	35.5/37.2	34.8/44.1	34.9/49.6

Appendix

Table 13: Average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA and WinoGrande on Qwen3-4B-Instruct. Each cell is the score and the percentage of the KV cache removed, for DimPO (

k{=}64

\lambda{=}1

Figure 6: Harness task scores for different projection methods on Qwen3-4B-Instruct against the percentage of the KV cache removed.

$d^{\prime}$ / $l$	0	2	4	7	10	12	14	16	18	21	24	26	28
64	58.8/0	58.8/1.8	58.7/3.6	58.5/6.2	58.0/8.9	57.9/10.7	56.2/12.5	54.2/14.3	52.1/16.1	47.4/18.8	45.3/21.4	43.3/23.2	40.6/25
16	58.8/0	58.1/3.1	57.8/6.2	57.1/10.9	53.8/15.6	52.1/18.8	44.4/21.9	41.2/25	37.6/28.1	35.7/32.8	34.5/37.5	35.2/40.6	34.8/43.8
4	58.8/0	57.9/3.5	55.6/6.9	52.0/12.1	43.7/17.3	39.9/20.8	36.9/24.2	35.8/27.7	34.9/31.1	35.0/36.3	34.5/41.5	34.5/45	35.4/48.4
1	58.8/0	57.6/3.5	54.5/7.1	50.2/12.4	40.9/17.7	36.9/21.3	35.9/24.8	34.8/28.3	34.8/31.9	34.5/37.2	33.9/42.5	34.2/46.1	35.0/49.6

Appendix

Table 14: Average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA and WinoGrande on Llama3.2-3B-Instruct. Each cell is the score and the percentage of the KV cache removed, for DimPO (

k{=}64

\lambda{=}1

Figure 7: Harness task scores for different projection methods on Llama3.2-3B-Instruct against the percentage of the KV cache removed.

$d^{\prime}$ / $l$	0	2	4	7	10	12	14	16	18	21	24	26	28
64	68.7/0	68.4/1.8	68.7/3.6	68.2/6.2	65.2/8.9	64.4/10.7	63.4/12.5	62.6/14.3	61.5/16.1	60.8/18.8	60.2/21.4	57.9/23.2	56.7/25
16	68.7/0	68.0/3.1	67.6/6.2	65.7/10.9	53.1/15.6	50.3/18.8	47.3/21.9	44.6/25	43.0/28.1	40.8/32.8	39.8/37.5	37.8/40.6	34.9/43.8
4	68.7/0	67.9/3.5	66.4/6.9	64.2/12.1	49.9/17.3	46.7/20.8	43.5/24.2	41.3/27.7	40.3/31.1	37.1/36.3	37.2/41.5	35.0/45	34.4/48.4
1	68.7/0	67.7/3.5	66.2/7.1	64.0/12.4	49.5/17.7	46.5/21.3	42.8/24.8	40.4/28.3	39.6/31.9	36.1/37.2	36.5/42.5	34.6/46.1	33.9/49.6

Appendix

Table 15: Average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA and WinoGrande on Qwen2.5-7B-Instruct. Each cell is the score and the percentage of the KV cache removed, for DimPO (

k{=}64

\lambda{=}1

Figure 8: Harness task scores for different projection methods on Qwen2.5-7B-Instruct against the percentage of the KV cache removed.

Figure 9: Average of ARC-Challenge, HellaSwag, MMLU, TruthfulQA, and WinoGrande against the percentage of the KV cache removed.

Figure 10: RULER 4k score against the percentage of projected layers for different projection objectives on Llama3.2-3B-Instruct.

Subtask	Base	DimPO ( $k{=}0$ , $\lambda{=}0$ )	DimPO ( $k{=}64$ , $\lambda{=}1$ )	KL
NIAH single 1	100.0	55.8	87.6	70.4
NIAH single 2	99.8	24.4	83.4	76.4
NIAH single 3	100.0	0.4	8.2	0.2
NIAH multi-key 1	100.0	37.8	67.8	64.8
NIAH multi-key 2	99.8	44.8	66.6	38.4
NIAH multi-key 3	100.0	0.2	0.8	0.0

Appendix

Table 16: RULER 4k subtask scores on Llama3.1-8B-Instruct after projecting

l{=}20

(

62.5\%

) attention layers to

d^{\prime}{=}64

Subtask	Base	DimPO ( $k{=}0$ , $\lambda{=}0$ )	DimPO ( $k{=}64$ , $\lambda{=}1$ )	KL
NIAH single 1	100.0	75.4	79.8	70.6
NIAH single 2	100.0	80.4	59.0	75.2
NIAH single 3	100.0	0.6	3.4	0.0
NIAH multi-key 1	100.0	63.8	66.8	62.2
NIAH multi-key 2	100.0	36.8	37.0	45.8
NIAH multi-key 3	100.0	0.0	0.0	0.0

Appendix

Table 17: RULER 4k subtask scores on Llama3.2-3B-Instruct after projecting

l{=}16

(

57.1\%

) attention layers to

d^{\prime}{=}64

Subtask	Base	DimPO ( $k{=}0$ , $\lambda{=}0$ )	DimPO ( $k{=}64$ , $\lambda{=}1$ )	KL
NIAH single 1	100.0	98.4	99.6	98.6
NIAH single 2	100.0	98.8	97.4	99.0
NIAH single 3	99.8	29.0	23.6	28.0
NIAH multi-key 1	100.0	86.4	85.2	68.8
NIAH multi-key 2	100.0	86.4	74.6	50.4
NIAH multi-key 3	100.0	25.6	24.6	1.8

Appendix

Table 18: RULER 4k subtask scores on Qwen3-4B-Instruct after projecting

l{=}24

(

66.7\%

) attention layers to

d^{\prime}{=}64

Figure 11: Inference latency with an

8

k input and

2048

generated tokens. Dashed lines show base model and SnapKV alone, while marked lines add the projection to the last

l

layers. The x-axis shows the percentage of layers receiving the projection.

	LongBench	RULER
		4K	8K	16K	32K	65K
Llama3.1-8B-Instruct	37.83	95.05	93.94	93.39	87.74	84.75
MagicPIG	35.84	92.63	92.35	91.64	86.71	83.67
MagicPIG, $\ell{=}4$	34.96	91.60	90.27	87.90	83.74	81.42
MagicPIG, $\ell{=}8$	33.38	89.95	87.12	83.83	81.83	72.08
MagicPIG, $\ell{=}12$	32.58	87.59	83.41	79.56	75.29	63.66

Appendix

Table 19: Llama3.1-8B-Instruct, MagicPIG (

K{=}8

L{=}75

), and MagicPIG with DimPO (

k{=}0

\lambda{=}0

) at

d^{\prime}{=}64

on the last

\ell

layers.

Transformer Attention

Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.

See all 320 papers

Token-weighted Direct Preference Optimization with Attention
May 21, 2026Chengyu Huang, Zhuohang Li, Sheng-Yen Chou +1Pairwise Preference Learning LLM Alignment
Cornell University · Vanderbilt University
RACE: RoPE-Aligned Attention Compression Based on Query-Key Activation Energies
Oct 20, 2026Irene Tenison, Abdulrahman Alabdulkareem, Daniel Karl I. Weidele +2KV-Cache Compression Transformer Attention
MIT CSAIL · MIT-IBM Computing Research Lab
Hierarchical Alignment: Block-Level LoRA-DPO for Preference Tuning in Large Language Models
Oct 14, 2025Yishen Chen, Qi Dong, Yukun ZhangLLM Alignment Preference Alignment
Fudan University Shanghai, China · The Chinese University of Hong Kong Hong Kong, China

May 21, 2026J/K move · Enter open · S save

Token-weighted Direct Preference Optimization with Attention

Chengyu Huang, Zhuohang Li, Sheng-Yen Chou, Claire Cardie

Cornell University · Vanderbilt University

Pairwise Preference Learning LLM Alignment Direct Preference Optimization Preference Optimization Preference-Based RL

Open paper AskPDF arXiv

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 AttentionPO content-aware, adjusting weights based on response content, and efficient, 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.

Abstract

Transformer Attention

Explore similar work

Token-weighted Direct Preference Optimization with Attention

RACE: RoPE-Aligned Attention Compression Based on Query-Key Activation Energies

Hierarchical Alignment: Block-Level LoRA-DPO for Preference Tuning in Large Language Models