DimPO: Dimensionality Reduction for Attention using Preference Optimization

2026-09-26

AI Summary

This study addresses how to optimally preserve model behavior during attention dimensionality reduction, particularly in long-context scenarios. To this end, it proposes DimPO, a method that offline trains projection mappings based on the attention patterns of frozen language models. The core innovation lies in revealing that preserving the ranking and concentration of task-relevant attentions is more critical than reproducing the full distribution. Accordingly, DimPO replaces traditional KL divergence matching with listwise preference optimization combined with a lightweight top-K cross-entropy objective. Experimental results demonstrate that when reducing dimensions by 50% per layer on LLaMA and Qwen architectures, the proposed approach retains approximately 95% of long-context performance, significantly outperforming KL-based baselines across downstream tasks.

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.

Problem

• Dimensionality Reduction
• Attention Mechanism
• Preference Optimization
• Long-context
• Training Objective

Innovation

• Dimensionality Reduction
• Preference Optimization
• Attention Mechanism
• Listwise Preference
• Long-context