Dimensionality reduction method preserves attention to improve language models
DimPO: Dimensionality Reduction for Attention using Preference Optimization
Machine Learning
Summary
Large language models use attention to focus on important parts of a text, but this attention can be very complex and slow. The authors propose DimPO, a way to simplify the attention process by reducing the size of the information it uses. Instead of trying to copy every detail, DimPO focuses on keeping the order of important parts and where the attention is strongest. Tests show that DimPO keeps model accuracy well while making the process faster, especially in long-text scenarios.
What this means in practice
- •For machine learning engineers: Reduce computation in large language models by applying DimPO to maintain attention quality while lowering dimension in inference tasks.
- •For natural language processing developers: Improve long sequence understanding in language models by using DimPO to preserve task-relevant attention more effectively than conventional methods.
Authors
Vojtěch Lanz, Yufei Cui, Prasanna Parthasarathi
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.