Faster Query-Key Learning Sharpens Attention in Self-Attention Models

2026-08-06

AI Summary

This work addresses the issue of diffuse attention distributions in self-attention mechanisms, which often undermine model interpretability. By analyzing the parametric structures of the query–key and output–value circuits and their differing learning rates, the study reveals that the relative learning rate critically governs attention sharpness. Through gradient flow analysis and closed-form dynamical derivations, the authors theoretically demonstrate this relationship and empirically validate it using a single-layer self-attention model. The experiments show that assigning a higher learning rate to the query–key circuit naturally induces sharper, more concentrated attention patterns. This approach significantly improves interpretability metrics while maintaining comparable predictive performance, offering a novel pathway to enhance attention focus without requiring additional regularization.

Abstract

A standard self-attention layer consists of two interacting circuits: the query-key circuit that governs attention allocation, and the output-value circuit that maps attended representations to predictions. Collapsed and factorized parameterizations of the query-key and output-value circuits lead to qualitatively different attention patterns. In particular, some parameterizations give sharper attention to task-relevant tokens, at a similar training loss. We analyze how the parameterizations of these circuits shape the parameter trajectories in single-layer self-attention models trained for next-token prediction. Through gradient-flow analysis, we show that factorization induces implicit rescaling of the two circuits' learning rates. We derive closed-form dynamics showing that output-value and query-key parameters move along a line, with relative speeds determined by their learning rates. Faster query-key learning relative to output-value learning thus produces sharper attention, as the model compensates for slower output-value learning by increasing attention mass on relevant tokens. Experiments show that differences in the relative learning rates of the two circuits govern attention concentration. This improves attention interpretability proxies while maintaining comparable predictive performance.

Problem

• self-attention
• query-key circuit
• attention sharpness
• parameterization
• learning rate

Innovation

• self-attention
• query-key learning
• attention sharpness
• learning rate factorization
• gradient flow analysis