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

# Faster Query-Key Learning Sharpens Attention in Self-Attention Models
Source: [https://arxiv.org/html/2608.06776](https://arxiv.org/html/2608.06776)
###### 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 on 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\.

Machine Learning, ICML

## 1Introduction

Transformer models\(Vaswani et al\.,[2017](https://arxiv.org/html/2608.06776#bib.bib32)\)are now standard across language\(Brown et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib7)\), vision\(Dosovitskiy et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib12)\), and speech\(Latif et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib18)\)\. Their success is driven by self\-attention, which produces contextual representations and is typically trained using next\-token prediction\. Empirical studies have shown that trained models with similar predictive performance can exhibit a wide range of attention patterns\(Jain & Wallace,[2019](https://arxiv.org/html/2608.06776#bib.bib16); Wiegreffe & Pinter,[2019](https://arxiv.org/html/2608.06776#bib.bib33); Serrano & Smith,[2019](https://arxiv.org/html/2608.06776#bib.bib29)\)\. Prior work has also studied formal settings in which attention weights may or may not align with attention\-based explanations\(Pandey et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib25)\)\. However, the role of architectural and parameterization choices in the emergence of such attention patterns remains unclear\.

Existing theoretical works have established that transformers are highly expressive, including universality and Turing completeness under suitable parameter settings\(Yun et al\.,[2019](https://arxiv.org/html/2608.06776#bib.bib34); Bhattamishra et al\.,[2020a](https://arxiv.org/html/2608.06776#bib.bib4),[b](https://arxiv.org/html/2608.06776#bib.bib5); Dehghani et al\.,[2018](https://arxiv.org/html/2608.06776#bib.bib8); Pérez et al\.,[2021](https://arxiv.org/html/2608.06776#bib.bib26)\)\. These results, however, rely on idealized parameter constructions and therefore do not explain how attention structure arises under standard gradient\-based optimization\. As a result, expressivity alone provides limited insight into how attention mechanisms and prediction layers co\-evolve during training\.

In this work, we analyze the training dynamics of a single\-layer transformer trained for next\-token prediction\(Radford et al\.,[2018](https://arxiv.org/html/2608.06776#bib.bib27)\)\. The model decomposes naturally into two interacting linear circuits\(Elhage et al\.,[2021](https://arxiv.org/html/2608.06776#bib.bib14); Olsson et al\.,[2022](https://arxiv.org/html/2608.06776#bib.bib24); Bietti et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib6)\):

- •thequery\-key circuit\(𝑾K​𝑾Q⊤\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\), which determines the attention pattern, and
- •theoutput\-value circuit\(𝑾O​𝑾V\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\), which maps attended representations to predictions\.

![Refer to caption](https://arxiv.org/html/2608.06776v1/x1.png)\(a\)Single Layer Self\-Attention Model Illustration
![Refer to caption](https://arxiv.org/html/2608.06776v1/x2.png)\(b\)Synthetic Data Illustration

Figure 1:Overall Setting: \(a\) A sequence𝐗\{\\mathbf\{X\}\}with contextual tokens and query token\{𝒙1,𝒙2,…,𝒙T−1,𝒙T\}\\\{\{\\bm\{x\}\}\_\{1\},\{\\bm\{x\}\}\_\{2\},\\ldots,\{\\bm\{x\}\}\_\{T\-1\},\{\\bm\{x\}\}\_\{T\}\\\}is fed into1−1\-layer transformer to predict the next tokeny=𝒙T\+1y=\{\\bm\{x\}\}\_\{T\+1\}\(b\) Synthetic data illustration, For each sequence class there exists a set of tokens specific to that class, called distinct tokens and a set of tokens common to all sequence classes called common tokens\. For example ‘the’ ‘sat’ ‘ate’ and ‘on’ are the common tokens, and ‘cat’ ‘mat’ are tokens distinct to class 1We show that the relative optimization speeds of these two circuits play a central role in shaping attention structure\. Varying the learning rates of the query\-key and output\-value parameters leads to qualitatively different attention patterns, even when predictive performance remains similar\. In particular, faster learning of the query\-key circuit leads to sharper attention on task\-relevant tokens\.

A key focus of this paper is the role of parameterization\. We compare factorized and collapsed parameterizations of both the query\-key and output\-value circuits and show that factorization induces a state\-dependent rescaling of gradient updates in the corresponding collapsed parameter space\. As a result, factorized and collapsed models can follow different optimization trajectories, even when they achieve comparable losses\.

We then characterize the training dynamics under collapsed parameterizations in a controlled synthetic setting\. Under orthogonality and symmetry assumptions, we derive closed\-form population gradient\-flow dynamics for both circuits and show that their effective parameters evolve at different rates\. In particular, the output\-value scalar grows on the order oflog⁡\(t\)\\log\(t\), while the query\-key scalar grows on the order oflog2⁡\(t\)\\log^\{2\}\(t\), with constants controlled by the learning\-rate ratioηQ​K/ηO​V\\eta\_\{QK\}/\\eta\_\{OV\}\. Because the query\-key circuit affects attention through a saturating nonlinearity, this difference in growth provides a mechanistic explanation for how relative learning speeds shape attention sharpening during training\.

We support the theory with experiments on synthetic data and real\-world datasets with token\-level relevance annotations, including HateXplain\(Mathew et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib23)\), subject\-verb agreement\(Linzen et al\.,[2016](https://arxiv.org/html/2608.06776#bib.bib21)\), and SQuAD\(Rajpurkar et al\.,[2016](https://arxiv.org/html/2608.06776#bib.bib28)\)\. Across settings, we observe that changes in relative optimization speed and parameterization alter attention structure in ways predicted by the theory, while predictive performance remains largely unchanged\.

## 2Related Works

![Refer to caption](https://arxiv.org/html/2608.06776v1/x3.png)\(a\)FAFO
![Refer to caption](https://arxiv.org/html/2608.06776v1/x4.png)\(b\)FACO
![Refer to caption](https://arxiv.org/html/2608.06776v1/x5.png)\(c\)CAFO
![Refer to caption](https://arxiv.org/html/2608.06776v1/x6.png)\(d\)CACO
![Refer to caption](https://arxiv.org/html/2608.06776v1/x7.png)\(e\)FAFO, QK learning rate increased10×10\\times
![Refer to caption](https://arxiv.org/html/2608.06776v1/x8.png)\(f\)FAFO, QK learning rate increasing linearly1−20×1\-20\\times
![Refer to caption](https://arxiv.org/html/2608.06776v1/x9.png)\(g\)CACO, QK learning rate increased5×5\\times
![Refer to caption](https://arxiv.org/html/2608.06776v1/x10.png)\(h\)CACO, QK learning rate increasing linearly1−10×1\-10\\times

Figure 2:Distinct Token Attention\-Prediction heat map of test data for different parameterizations under SGD with same learning rates \(Top\-Row\) and Faster Learning for Query\-Key \(Bottom\-Row\)\. All models are trained until the negative log\-likelihood is close to zero \(≤0\.001\\leq 0\.001\)\. Here, we train the model for55random seeds and report the average, refer to appendix for standard deviation in the heatmaps\.Many recent works have studied the optimization dynamics of attention and self\-attention\-based models, particularly in simplified or single\-layer settings\(Tian et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib30); Deora et al\.,[2024a](https://arxiv.org/html/2608.06776#bib.bib9); Lu et al\.,[2021](https://arxiv.org/html/2608.06776#bib.bib22); Vashisht & Ramaswamy,[2023](https://arxiv.org/html/2608.06776#bib.bib31)\)\.Tian et al\. \([2023](https://arxiv.org/html/2608.06776#bib.bib30)\)analyze attention dynamics under the assumption that output–value parameters learn faster than query–key parameters\.Deora et al\. \([2024a](https://arxiv.org/html/2608.06776#bib.bib9)\)study the optimization and generalization properties of multi\-head self\-attention models under realizability assumptions on the data\. Across these analyses, optimization dynamics are typically examined under fixed or implicit assumptions about the relative learning behavior of attention components\. In contrast, we vary the relative learning rates of the query–key and output–value circuits and show that faster learning of query–key parameters sharpens attention\. While prior work often focuses on single\-sample or batch\-size\-one regimes for analytical tractability, we analyze full\-batch training dynamics that more closely reflect modern practice\. Moreover, we extend the analysis beyond the attention mechanism to the final prediction layer, providing a unified view of how interacting circuits shape the training trajectory\.

A related line of work adopts a mechanistic perspective on transformer models\. The transformer circuits framework ofElhage et al\. \([2021](https://arxiv.org/html/2608.06776#bib.bib14)\)provides a systematic decomposition of attention layers into interacting components, which has motivated subsequent analyses of how structure emerges in trained models, including clustering behavior in attention patterns\(Geshkovski et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib15); Karagodin et al\.,[2024](https://arxiv.org/html/2608.06776#bib.bib17)\), memory\-like mechanisms\(Bietti et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib6)\), and the emergence of topic structure\(Li et al\.,[2023](https://arxiv.org/html/2608.06776#bib.bib19)\)\. These works offer valuable insight into the functional roles of attention heads and learned representations after training\. However, they primarily characterize attention structure post hoc and do not analyze how optimization dynamics shape the relative evolution of interacting circuits during training\. Our work instead focuses on how the training trajectories of the query–key and output–value circuits jointly determine attention sharpening under gradient\-based optimization\.

## 3Problem Setting

We study a single\-layer self\-attention model trained for next\-token prediction\. Let𝐗=\[𝒙1,…,𝒙T\]⊤∈ℝT×d\{\\mathbf\{X\}\}=\[\{\\bm\{x\}\}\_\{1\},\\ldots,\{\\bm\{x\}\}\_\{T\}\]^\{\\top\}\\in\{\\mathbb\{R\}\}^\{T\\times d\}denote an input sequence, where𝒙T\{\\bm\{x\}\}\_\{T\}is the query token\. The model parameters are𝑾Q,𝑾K,𝑾V∈ℝd×d\{\\bm\{W\}\}\_\{Q\},\{\\bm\{W\}\}\_\{K\},\{\\bm\{W\}\}\_\{V\}\\in\{\\mathbb\{R\}\}^\{d\\times d\}and𝑾O∈ℝM×d\{\\bm\{W\}\}\_\{O\}\\in\{\\mathbb\{R\}\}^\{M\\times d\}, whereMMis the vocabulary size\.

Given a sequence𝐗\{\\mathbf\{X\}\}, the model output is

𝒇​\(𝐗\)=𝑾O​𝑾V​𝐗⊤​𝕊​\(𝐗​𝑾K​𝑾Q⊤​𝒙T\),\{\\bm\{f\}\}\(\{\\mathbf\{X\}\}\)=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\{\\mathbf\{X\}\}^\{\\top\}\{\\mathbb\{S\}\}\\\!\\left\(\{\\mathbf\{X\}\}\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{x\}\}\_\{T\}\\right\),\(1\)where𝕊​\(⋅\)\{\\mathbb\{S\}\}\(\\cdot\)denotes the softmax operator\. Figure[1\(a\)](https://arxiv.org/html/2608.06776#S1.F1.sf1)illustrates this computation for a single\-layer self\-attention model\.

The model decomposes naturally into two linear circuits\. The query\-key circuit𝑾K​𝑾Q⊤\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}determines the attention distribution over tokens, while the output\-value circuit𝑾O​𝑾V\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}maps the attended representation to prediction logits\. We study four parameterizations that differ in whether these circuits are factorized or collapsed:

- •FAFO:factorized attention and factorized output\-value\.
- •CAFO:collapsed attention \(𝑾Q​K=𝑾K​𝑾Q⊤\{\\bm\{W\}\}\_\{QK\}=\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\) and factorized output\-value\.
- •FACO:factorized attention and collapsed output\-value \(𝑾O​V=𝑾O​𝑾V\{\\bm\{W\}\}\_\{OV\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\)\.
- •CACO:collapsed attention and collapsed output\-value\.

Let𝑬=\[𝐞1,…,𝐞M\]⊤∈ℝM×d\{\\bm\{E\}\}=\[\{\\mathbf\{e\}\}\_\{1\},\\ldots,\{\\mathbf\{e\}\}\_\{M\}\]^\{\\top\}\\in\{\\mathbb\{R\}\}^\{M\\times d\}denote the embedding matrix\. Given a sequence𝐗\{\\mathbf\{X\}\}with next\-token labely∈\[M\]y\\in\[M\], the training objective is to minimize the negative log\-likelihood

ℒ=−log⁡𝕊y​\(𝒇​\(𝐗\)\)\.\\mathcal\{L\}=\-\\log\{\\mathbb\{S\}\}\_\{y\}\(\{\\bm\{f\}\}\(\{\\mathbf\{X\}\}\)\)\.\(2\)
To analyze attention behavior in a controlled setting, we adopt a synthetic data construction based onTian et al\. \([2023](https://arxiv.org/html/2608.06776#bib.bib30)\)\. Each sequence class is associated with a fixed query and next\-token pair\(q,n\)\(q,n\), where the next tokennnuniquely identifies the class\. Context tokens preceding the query are sampled from a mixture of class\-specific*distinct*tokens and class\-agnostic*common*tokens, as illustrated in Figure[1\(b\)](https://arxiv.org/html/2608.06776#S1.F1.sf2)\.

For example, consider three sequence classes:

- •Class 1: “cat sat on the mat”, “cat ate on the mat”,
- •Class 2: “dog sat on the rug”, “dog ate on the rug”,
- •Class 3: “bird sat on the branch”, “bird ate on the branch”\.

Here, tokens such as “sat”, “ate”, and “on” appear across all classes and are common, while tokens such as “cat”, “mat”, “dog”, and “rug” are class\-specific and distinct\.

Formally, defineΩ​\(l\)=\{n:ℙ​\(l∣n\)\>0\}\\Omega\(l\)=\\\{n:\\mathbb\{P\}\(l\\mid n\)\>0\\\}\. Tokens with\|Ω​\(l\)\|=1\|\\Omega\(l\)\|=1are distinct, while tokens with\|Ω​\(l\)\|\>1\|\\Omega\(l\)\|\>1are common\. Context tokens are sampled according to

ℙ​\(l∣n\)=\{γ/C,\|Ω​\(l\)\|=1,\(1−γ\)/D,\|Ω​\(l\)\|\>1,0,otherwise,\\mathbb\{P\}\(l\\mid n\)=\\begin\{cases\}\\gamma/C,&\|\\Omega\(l\)\|=1,\\\\ \(1\-\\gamma\)/D,&\|\\Omega\(l\)\|\>1,\\\\ 0,&\\text\{otherwise\},\\end\{cases\}whereCCandDDdenote the number of distinct and common tokens, andγ∈\(0,1\)\\gamma\\in\(0,1\)controls the mixture weight\. This construction separates informative and non\-informative tokens and enables precise analysis of how attention and prediction evolve under different parameterizations and optimization speeds\.

## 4Self\-Attention Analysis under Different Parameterizations

We identify an empirical phenomenon that motivates our theory\. Different self\-attention parameterizations exhibit qualitatively different attention patterns even when trained under identical optimization settings and achieving comparable predictive performance\.

We consider a synthetic next\-token prediction task with four sequence classes\. Each sequence containsmmclass\-specific tokens, which are informative for prediction,nncommon tokens have no class information, and a query token\.

To assess alignment between attention and prediction, we introduce the Distinct Token Attention–Prediction \(DTAP\) heatmap\. DTAP is a two\-dimensional histogram in which the x\-axis represents the fraction of attention assigned to class\-specific tokens and the y\-axis represents the predicted probability of the correct class\. Each bin reports the fraction of examples that fall within the corresponding attention and confidence ranges\. An ideal model concentrates most mass in the top\-right corner, indicating both strong attention on informative tokens and high prediction confidence\.

We compare four parameterizations: factorized attention–factorized output \(FAFO\), collapsed attention–collapsed output \(CACO\), collapsed attention–factorized output \(CAFO\), and factorized attention–collapsed output \(FACO\)\. Figure[2](https://arxiv.org/html/2608.06776#S2.F2)\(a–d\) shows that all parameterizations achieve high prediction confidence, yet their attention distributions differ substantially\.

- •Parameterizations with collapsed output–value structure \(CACO and FACO\) concentrate mass near the top\-right region of the DTAP heatmap\.
- •The factorized output–value parameterization \(FAFO\) places substantial mass in the top\-left region, corresponding to correct predictions made with weak attention on informative tokens\.
- •Increasing the learning rate of the query–key parameters shifts mass toward the top\-right region across all parameterizations\.

As shown in Figure[2](https://arxiv.org/html/2608.06776#S2.F2)\(e–h\), accelerating query–key learning causes the attention patterns of FAFO and CACO to resemble those of FACO\. In particular, attention mass shifts toward distinct tokens without degrading prediction accuracy\. This behavior is observed under both fixed and gradually increasing query–key learning rates, indicating that relative optimization speed governs attention alignment beyond parameterization alone\. These observations motivate the optimization analysis in Section[6](https://arxiv.org/html/2608.06776#S6)\.

The dataset for this task was generated as follows\. We consider44sequence classes, with44query and corresponding next tokens\. For each query token, we define a set of1010distinct tokens\. We also have a set of common tokens with cardinality1010\. Note that the position \(or index\) of the distinct and common tokens can be arbitrary\. We sample multiple such triples \(context tokens, query token, next token\), train a single\-layer self\-attention model on a subset of this dataset, and evaluate it on the remaining data\. The vocabulary size is5050tokens and sequence length is6464\. We train the model with200200batches with batch size3232using stochastic gradient descent\.

## 5Isolating Output–Value Dynamics by Fixing Attention

![Refer to caption](https://arxiv.org/html/2608.06776v1/x11.png)\(a\)Factorized Output\-Value
![Refer to caption](https://arxiv.org/html/2608.06776v1/x12.png)\(b\)Collapsed Output\-Value
![Refer to caption](https://arxiv.org/html/2608.06776v1/x13.png)\(c\)Factorized Output\-Value
![Refer to caption](https://arxiv.org/html/2608.06776v1/x14.png)\(d\)Collapsed Output\-Value

Figure 3:\(a\) and \(b\) Evolution of the training loss for fixed values ofα\\alphaunder gradient descent\. \(c\) and \(d\) show improvement incentive curves, measuring the change in loss resulting from a small increaseδ=0\.005\\delta=0\.005inα\\alpha\.The experiments in Section[4](https://arxiv.org/html/2608.06776#S4)show that both parameterization and relative learning rates influence attention allocation\. To isolate the source of this effect, we analyze the learning dynamics of the output–value circuit while holding attention over distinct tokens fixed\.

We consider a query–key circuit that assigns a total attention massα∈\[0,1\]\\alpha\\in\[0,1\]to the distinct tokens, distributed uniformly among them, with the remaining mass assigned uniformly to the common tokens\. Although fixing attention in this manner is not a realistic training regime, it serves as a diagnostic intervention that isolates the optimization behavior of the output–value circuit\.

Under fixedα\\alpha, training the output–value parameters reduces to a multi\-class classification problem with an effective input representation

ℒ​\(α\)=−log⁡\(softmaxy​\(𝑾O​𝑾V​𝒙~\)\),\\mathcal\{L\}\(\\alpha\)=\-\\log\\Big\(\\mathrm\{softmax\}\_\{y\}\\big\(\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\\widetilde\{\{\\bm\{x\}\}\}\\big\)\\Big\),where

𝒙~=αm​∑j∈Sdistinct𝒙j\+1−αn​∑j∈Scommon𝒙j\.\\widetilde\{\{\\bm\{x\}\}\}=\\frac\{\\alpha\}\{m\}\\sum\_\{j\\in S\_\{\\text\{distinct\}\}\}\{\\bm\{x\}\}\_\{j\}\+\\frac\{1\-\\alpha\}\{n\}\\sum\_\{j\\in S\_\{\\text\{common\}\}\}\{\\bm\{x\}\}\_\{j\}\.
Asα\\alphaincreases, the effective input𝒙~\\widetilde\{\{\\bm\{x\}\}\}becomes more informative, simplifying the optimization problem faced by the output–value circuit\. We define the fixed\-attention loss curve as the training loss evaluated under gradient descent while holding the attention massα\\alphafixed for distinct tokens\. To quantify the incentive for attention sharpening, we measure the change in loss induced by a small increase in attention mass,δ=0\.005\\delta=0\.005, evaluated at a fixed training time\.

Under the factorized output\-value parameterization, the loss reaches near\-zero atα=0\.2\\alpha=0\.2\(blue curve in Figure[3](https://arxiv.org/html/2608.06776#S5.F3)a\)\. The improvement incentive curve in Figure[3](https://arxiv.org/html/2608.06776#S5.F3)c shows a positive incentive for attention sharpening for up to100100epochs, allowingα\\alphato increase to0\.40\.4\. However, forα≥0\.4\\alpha\\geq 0\.4\(orange curve\), the incentive rapidly diminishes\. As a consequence, attention tends to stabilize at values such asα=0\.4\\alpha=0\.4\.

In contrast, under the collapsed output\-value parameterization, the loss does not decrease to zero as quickly forα≥0\.2\\alpha\\geq 0\.2andα≤0\.8\\alpha\\leq 0\.8\(blue, orange, green, and red curves in Figure[3](https://arxiv.org/html/2608.06776#S5.F3)b\)\. As a result, the incentive for increasingα\\alpharemains high over a wider range, allowing attention to continue sharpening\.

Taken together, the experiments in Sections[4](https://arxiv.org/html/2608.06776#S4)and[5](https://arxiv.org/html/2608.06776#S5)reveal an empirical pattern\. Architectural parameterization shapes the relative learning dynamics of the query–key and output–value circuits, which in turn determines whether attention continues to sharpen or stabilizes early\. These effects cannot be explained by predictive performance alone and persist across learning\-rate interventions\.

In the next section, we develop a theoretical framework that explains the observations by analyzing how different parameterizations induce distinct optimization trajectories under gradient flow\.

## 6Optimization Trajectories under Different Parameterization

Theoretical analyses of self\-attention often simplify optimization dynamics\(Deora et al\.,[2024b](https://arxiv.org/html/2608.06776#bib.bib10); Li et al\.,[2024](https://arxiv.org/html/2608.06776#bib.bib20); Bao et al\.,[2024](https://arxiv.org/html/2608.06776#bib.bib3)\)by working with collapsed parameter matrices, rather than explicitly modeling the factorized query–key and output–value parameters\(Elhage et al\.,[2021](https://arxiv.org/html/2608.06776#bib.bib14); Olsson et al\.,[2022](https://arxiv.org/html/2608.06776#bib.bib24)\)\. This simplification enables tractable analysis, but it is not a priori clear whether it faithfully reflects the training dynamics of factorized models\. Empirically, as shown in Section[4](https://arxiv.org/html/2608.06776#S4), factorized and collapsed parameterizations can achieve comparable predictive performance while exhibiting qualitatively different attention behaviors\.

In a single\-layer self\-attention model, each circuit consists of a composition of linear transformations with no intervening nonlinearity\. Thus, the optimization dynamics of the query\-key and output\-value circuits fall within the class of homogeneous linear models for which conservation laws under gradient flow hold\(Arora et al\.,[2018](https://arxiv.org/html/2608.06776#bib.bib2); Du et al\.,[2018](https://arxiv.org/html/2608.06776#bib.bib13)\)\. Building on results from deep linear network theory, we show that factorized training induces a well\-defined dynamics on the collapsed parameters, differing from collapsed training only through a state\-dependent preconditioning operator\.

Let𝒛=𝑾O​V​𝑿⊤​𝒂\{\\bm\{z\}\}=\{\\bm\{W\}\}\_\{OV\}\{\\bm\{X\}\}^\{\\top\}\\bm\{\{\\bm\{a\}\}\}and𝒛~=𝑾O​𝑾V​𝑿⊤​𝒂\\widetilde\{\{\\bm\{z\}\}\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\{\\bm\{X\}\}^\{\\top\}\\bm\{\{\\bm\{a\}\}\}\. We define the loss functionsℒ1​\(𝑾O​V\)=−log⁡𝕊y​\(𝒛\),ℒ2​\(𝑾O,𝑾V\)=−log⁡𝕊y​\(𝒛~\),\\mathcal\{L\}^\{1\}\(\{\\bm\{W\}\}\_\{OV\}\)=\-\\log\\,\{\\mathbb\{S\}\}\_\{y\}\(\{\\bm\{z\}\}\),~~\\mathcal\{L\}^\{2\}\(\{\\bm\{W\}\}\_\{O\},\{\\bm\{W\}\}\_\{V\}\)=\-\\log\\,\{\\mathbb\{S\}\}\_\{y\}\(\\widetilde\{\{\\bm\{z\}\}\}\),where the attention weights are given by𝒂=𝕊​\(𝑿​𝑾K​𝑾Q⊤​𝒙T\)\\bm\{\{\\bm\{a\}\}\}=\{\\mathbb\{S\}\}\\\!\\left\(\{\\bm\{X\}\}\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{x\}\}\_\{T\}\\right\)\.

![Refer to caption](https://arxiv.org/html/2608.06776v1/x15.png)\(a\)Log\-Loss Curves
![Refer to caption](https://arxiv.org/html/2608.06776v1/x16.png)\(b\)μO​V​\(t\)\\mu\_\{OV\}\(t\)evolution
![Refer to caption](https://arxiv.org/html/2608.06776v1/x17.png)\(c\)μQ​K​\(t\)\\mu\_\{QK\}\(t\)evolution
![Refer to caption](https://arxiv.org/html/2608.06776v1/x18.png)\(d\)α​\(t\)\\alpha\(t\)evolution

Figure 4:Population gradient\-flow simulations for the orthogonal synthetic setting of Theorem[1](https://arxiv.org/html/2608.06776#Thmtheorem1)withm=5m=5,n=50n=50,b=50b=50,M=5M=5, varying the ratior=ηQ​K/ηO​Vr=\\eta\_\{QK\}/\\eta\_\{OV\}\. \(a\) Training loss vs\. normalized time, \(b\)μO​V​\(t\)\\mu\_\{OV\}\(t\)at matched stopping loss, \(d\)μQ​K​\(t\)\\mu\_\{QK\}\(t\)at matched stopping loss, and attention massα​\(t\)\\alpha\(t\)evaluated at matched stopping losses\.We now relate the gradient flow dynamics of the factorized model to those of the collapsed parameterization\. Let𝑾O​\(t\)\{\\bm\{W\}\}\_\{O\}\(t\)and𝑾V​\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)denote the gradient flow trajectories underℒ2\\mathcal\{L\}^\{2\}, and define the effective matrix𝑾~O​V​\(t\)=𝑾O​\(t\)​𝑾V​\(t\)\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\(t\)=\{\\bm\{W\}\}\_\{O\}\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)\.

The proposition[1](https://arxiv.org/html/2608.06776#Thmproposition1)is a direct application of the chain rule\.

###### Proposition 1\.

LetC=A​BC=AB, whereAA,BB, andCCare matrices of shapeM×dM\\times d,d×dd\\times dandM×dM\\times d\.

Then the gradients satisfy

∇𝑾Oℒ2\|\(A,B\)\\displaystyle\\nabla\_\{\{\\bm\{W\}\}\_\{O\}\}\\mathcal\{L\}^\{2\}\\big\|\_\{\(A,B\)\}=∇𝑾O​Vℒ1\|C​B⊤,\\displaystyle=\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}\}\\mathcal\{L\}^\{1\}\\big\|\_\{C\}\\,B^\{\\top\},\(3\)∇𝑾Vℒ2\|\(A,B\)\\displaystyle\\nabla\_\{\{\\bm\{W\}\}\_\{V\}\}\\mathcal\{L\}^\{2\}\\big\|\_\{\(A,B\)\}=A⊤​∇𝑾O​Vℒ1\|C\.\\displaystyle=A^\{\\top\}\\,\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}\}\\mathcal\{L\}^\{1\}\\big\|\_\{C\}\.

The gradient identities in Proposition[1](https://arxiv.org/html/2608.06776#Thmproposition1)do not by themselves imply a closed evolution for the effective matrix𝑾~O​V=𝑾O​𝑾V\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}, since the dynamics of𝑾O\{\\bm\{W\}\}\_\{O\}and𝑾V\{\\bm\{W\}\}\_\{V\}may depend on their relative scaling\. We therefore next characterize a structural condition under which factorized gradient flow admits a well\-defined induced dynamics on the effective parameters\.

###### Lemma 1\.

Let𝐖O​\(t\)\{\\bm\{W\}\}\_\{O\}\(t\)and𝐖V​\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)denote the gradient flow trajectories with lossℒ2\\mathcal\{L\}^\{2\}\. Let𝐖O​\(0\)⊤​𝐖O​\(0\)\{\\bm\{W\}\}\_\{O\}\(0\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\(0\)=𝐖V​\(0\)​𝐖V​\(0\)⊤=\{\\bm\{W\}\}\_\{V\}\(0\)\{\\bm\{W\}\}\_\{V\}\(0\)^\{\\top\}, Then𝐖O​\(t\)⊤​𝐖O​\(t\)\{\\bm\{W\}\}\_\{O\}\(t\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\(t\)=𝐖V​\(t\)​𝐖V​\(t\)⊤=\{\\bm\{W\}\}\_\{V\}\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)^\{\\top\}

Lemma[1](https://arxiv.org/html/2608.06776#Thmlemma1)ensures that the evolution of the effective matrix does not depend on arbitrary rescalings of the factorized parameters\. Under this condition, the factorized gradient flow induces a closed evolution for the effective matrix that depends only on the collapsed loss\.

###### Lemma 2\.

Let𝐖O​\(t\)\{\\bm\{W\}\}\_\{O\}\(t\)and𝐖V​\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)follow gradient flow onℒ2\\mathcal\{L\}^\{2\}, and define the effective matrix𝐖~O​V​\(t\)=𝐖O​\(t\)​𝐖V​\(t\)\\widetilde\{\\bm\{W\}\}\_\{OV\}\(t\)=\{\\bm\{W\}\}\_\{O\}\(t\)\{\\bm\{W\}\}\_\{V\}\(t\)\. Assume𝐖O​\(0\)⊤​𝐖O​\(0\)=𝐖V​\(0\)​𝐖V​\(0\)⊤\{\\bm\{W\}\}\_\{O\}\(0\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\(0\)=\{\\bm\{W\}\}\_\{V\}\(0\)\{\\bm\{W\}\}\_\{V\}\(0\)^\{\\top\}\. Then𝐖~O​V​\(t\)\\widetilde\{\\bm\{W\}\}\_\{OV\}\(t\)evolves according to the preconditioned gradient flow forℒ1\\mathcal\{L\}^\{1\}

dd​t​vec⁡\(𝑾~O​V​\(t\)\)=−Ω​\(𝑾~O​V​\(t\)\)​vec⁡\(∇ℒ1​\(𝑾~O​V​\(t\)\)\),\\frac\{d\}\{dt\}\\operatorname\{vec\}\(\\widetilde\{\\bm\{W\}\}\_\{OV\}\(t\)\)=\-\\Omega\(\\widetilde\{\\bm\{W\}\}\_\{OV\}\(t\)\)\\,\\operatorname\{vec\}\\\!\\left\(\\nabla\\mathcal\{L\}^\{1\}\(\\widetilde\{\\bm\{W\}\}\_\{OV\}\(t\)\)\\right\),whereΩ​\(𝐂\)=\(𝐂⊤​𝐂\)1/2⊕\(𝐂​𝐂⊤\)1/2,\\Omega\(\{\\bm\{C\}\}\)=\(\{\\bm\{C\}\}^\{\\top\}\{\\bm\{C\}\}\)^\{1/2\}\\oplus\(\{\\bm\{C\}\}\{\\bm\{C\}\}^\{\\top\}\)^\{1/2\},

Ω​\(𝑪\)\\Omega\(\{\\bm\{C\}\}\)a mapping fromℝM×d→ℝM​d×M​d\{\\mathbb\{R\}\}^\{M\\times d\}\\rightarrow\{\\mathbb\{R\}\}^\{Md\\times Md\}is a degree 1\-homogeneous map\.⊕\\oplusdenotes the Kronecker sum\.

Lemma[2](https://arxiv.org/html/2608.06776#Thmlemma2)shows that, under the balancedness condition, the evolution of the effective matrix𝑾~O​V​\(t\)\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\(t\)induced by factorized training is fully determined by the collapsed lossℒ1\\mathcal\{L\}^\{1\}, up to the application of a state\-dependent preconditioning operator\. In contrast to collapsed training, which follows standard gradient flow onℒ1\\mathcal\{L\}^\{1\}, factorized training follows a preconditioned gradient flow in the same parameter space\.Ω\\Omegabeing a homogeneous map implies that using factorized parameterization is closely approximated by collapsed factorization with a learning rate that increases over time\.

The same argument applies to the query\-key factorization𝑾~Q​K=𝑾K​𝑾Q⊤\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}=\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\. Its gradient flow admits an analogous preconditioned representation of the collapsed query\-key dynamics\. We provide the full derivation in Appendix[D](https://arxiv.org/html/2608.06776#A4)\. Together, these results justify analyzing optimization trajectories in the collapsed parameter space for both circuits, provided the induced preconditioning is taken into account\.

### 6\.1Synthetic Orthogonal Data Setting

This section isolates the mechanism by which relative learning rates between the query\-key and output\-value circuits shape attention during training\. We introduce a simplified synthetic setting that enables a transparent, closed\-form characterization of how attention sharpening emerges under gradient flow\. To enable such a characterization, we make the following assumptions on the data distribution\.

###### Assumption 1\.

For each classy∈\[M\]y\\in\[M\], the class\-specific token distribution is supported onbbdistinct tokens,Dy=1b​∑τ=1bδ​\(𝐬yτ\),D\_\{y\}=\\frac\{1\}\{b\}\\sum\_\{\\tau=1\}^\{b\}\\delta\(\{\\mathbf\{s\}\}\_\{y\}^\{\\tau\}\),where for eachτ∈\[b\]\\tau\\in\[b\], the vectors𝐬1τ,𝐬2τ,…,𝐬Mτ∈ℝd\{\\mathbf\{s\}\}\_\{1\}^\{\\tau\},\{\\mathbf\{s\}\}\_\{2\}^\{\\tau\},\\ldots,\{\\mathbf\{s\}\}\_\{M\}^\{\\tau\}\\in\{\\mathbb\{R\}\}^\{d\}are mutually orthogonal\.

###### Assumption 2\.

The background token distributionD0D\_\{0\}has zero mean\.

###### Assumption 3\.

The support ofD0D\_\{0\}is orthogonal to all class\-specific token vectors\{𝐬yτ\}y∈\[M\],τ∈\[b\]\\\{\{\\mathbf\{s\}\}\_\{y\}^\{\\tau\}\\\}\_\{y\\in\[M\],\\,\\tau\\in\[b\]\}\.

Together, these assumptions eliminate cross\-token interactions, reducing training to a low\-dimensional system in which the coupling between attention learning and output learning can be analyzed explicitly\.

Each training sequence consists ofmmclass\-specific token positions \(sampled with replacement from thebbtokens of the corresponding class\),nnbackground token positions, and a final query token, giving total lengthT=m\+n\+1T=m\+n\+1\. The attention vector is𝜶=softmax​\(𝑿​𝑾Q​K​𝒙T\)\\bm\{\\alpha\}=\\mathrm\{softmax\}\(\{\\bm\{X\}\}\{\\bm\{W\}\}\_\{QK\}\{\\bm\{x\}\}\_\{T\}\), and we denote byα​\(t\):=∑j∈𝒟αj\\alpha\(t\):=\\sum\\limits\_\{j\\in\\mathcal\{D\}\}\\alpha\_\{j\}the total attention mass on class\-specific token positions \(relevant token for a class\)\.

The following theorem shows that, under Assumptions[1](https://arxiv.org/html/2608.06776#Thmassumption1)\-[3](https://arxiv.org/html/2608.06776#Thmassumption3), training dynamics collapse to two coupled scalar trajectories governing output\-value learning and attention sharpening\.

###### Theorem 1\.

Under the orthogonal full\-data setting and initialization𝐖O​V​\(0\)=𝐖Q​K​\(0\)=𝟎\{\\bm\{W\}\}\_\{OV\}\(0\)=\{\\bm\{W\}\}\_\{QK\}\(0\)=\\bm\{0\}, the output\-value and query\-key parameters evolve under population gradient flow with learning ratesηO​V\\eta\_\{OV\}andηQ​K\\eta\_\{QK\}as

𝑾O​V\(k,:\)​\(t\)=μO​V​\(t\)​\[∑τ=1bskτ−1M​∑y=1M∑τ=1bsyτ\],\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)=\\mu\_\{OV\}\(t\)\\left\[\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\],
𝑾Q​K\(k,:\)​\(t\)=μQ​K​\(t\)​∑y=1M\(∑τ=1bsyτ\)​\(∑τ′=1bsyτ′,k\),\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\(t\)=\\mu\_\{QK\}\(t\)\\sum\_\{y=1\}^\{M\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)\\left\(\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\},k\}\\right\),
where the scalar coefficients satisfy

d​μO​V​\(t\)d​t=ηO​V​α​\(t\)b​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\),\\frac\{d\\mu\_\{OV\}\(t\)\}\{dt\}=\\frac\{\\eta\_\{OV\}\\,\\alpha\(t\)\}\{b\\big\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\big\)\},d​μQ​K​\(t\)d​t=ηQ​K​\(M−1\)​\(α​\(t\)−α​\(t\)2\)​μO​V​\(t\)M​b2​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\),\\frac\{d\\mu\_\{QK\}\(t\)\}\{dt\}=\\frac\{\\eta\_\{QK\}\\,\(M\-1\)\\big\(\\alpha\(t\)\-\\alpha\(t\)^\{2\}\\big\)\\mu\_\{OV\}\(t\)\}\{Mb^\{2\}\\big\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\big\)\},with

α​\(t\)=m​exp⁡\(μQ​K​\(t\)\)m​exp⁡\(μQ​K​\(t\)\)\+n\.\\alpha\(t\)=\\frac\{m\\,\\exp\(\\mu\_\{QK\}\(t\)\)\}\{m\\,\\exp\(\\mu\_\{QK\}\(t\)\)\+n\}\.

In particular, both parameter matrices evolve along fixed data\-dependent directions, with all learning dynamics captured by the scalar coefficientsμO​V​\(t\)\\mu\_\{OV\}\(t\)andμQ​K​\(t\)\\mu\_\{QK\}\(t\)\.

Table 1:Train and test performance together with attention\-based interpretability metrics on real\-world datasets\.↑\\uparrowindicates higher is better, while↓\\downarrowindicates lower is better\. All results are reported as average over55runs\. We report95%95\\%confidence intervals over five random seeds in Table[4](https://arxiv.org/html/2608.06776#A1.T4)\.###### Lemma 3\(Relative growth rates \(pre\-saturation regime\)\)\.

Under the population gradient\-flow dynamics of Theorem[1](https://arxiv.org/html/2608.06776#Thmtheorem1), assume that the attention mass satisfiesα​\(t\)≤1−δ\\alpha\(t\)\\leq 1\-\\deltafor someδ\>0\\delta\>0over a time interval\[0,T\]\[0,T\]\. Then for allt∈\[0,T\]t\\in\[0,T\],

μO​V​\(t\)=Θ​\(ln⁡\(1\+t\)\),μQ​K​\(t\)=Θ​\(ηQ​KηO​V​\(ln⁡\(1\+t\)\)2\)\.\\mu\_\{OV\}\(t\)=\\Theta\\\!\\big\(\\ln\(1\+t\)\\big\),~\\mu\_\{QK\}\(t\)=\\Theta\\\!\\left\(\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\(\\ln\(1\+t\)\)^\{2\}\\right\)\.

Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3)formalizes the core mechanism: the output\-value scale grows logarithmically in time, while the growth of the query\-key scale that governs attention concentration is controlled by the learning\-rate ratior=ηQ​K/ηO​Vr=\\eta\_\{QK\}/\\eta\_\{OV\}through a log\-squared dependence\. When output learning is slowed relative to attention learning \(ηQ​K\>ηO​V\\eta\_\{QK\}\>\\eta\_\{OV\}\), optimization compensates by increasing the magnitude of the query\-key parameters, thereby concentrating attention mass on class\-relevant tokens\.

Figure[4](https://arxiv.org/html/2608.06776#S6.F4)a shows that the training loss converges to zero for all values of the ratior=ηQ​K/ηO​Vr=\\eta\_\{QK\}/\\eta\_\{OV\}, indicating that final predictive accuracy alone does not distinguish between learning regimes\. Figure[4](https://arxiv.org/html/2608.06776#S6.F4)b shows the evolution ofμO​V\\mu\_\{OV\}evaluated at matched stopping losses varying values ofrr\. Figures[4](https://arxiv.org/html/2608.06776#S6.F4)c and d similarly show the evolution ofμQ​K\\mu\_\{QK\}and the attention massα\\alphaat matched stopping loss thresholds\.

To understand the differing attention behaviors, we examine the intermediate dynamics in Fig\.[4](https://arxiv.org/html/2608.06776#S6.F4)b–c\. Figure[4](https://arxiv.org/html/2608.06776#S6.F4)b shows the evolution ofμO​V​\(t\)\\mu\_\{OV\}\(t\), which serves as a proxy for logit growth and directly governs loss decay, while Fig\.[4](https://arxiv.org/html/2608.06776#S6.F4)c shows the evolution ofμQ​K​\(t\)\\mu\_\{QK\}\(t\), which acts as a proxy for growth of attention on relevant tokens\.

For all values ofrr,μO​V​\(t\)\\mu\_\{OV\}\(t\)increases sufficiently to drive the loss to zero, explaining the similar convergence behavior observed in Fig\.[4](https://arxiv.org/html/2608.06776#S6.F4)a\. In contrast,μQ​K​\(t\)\\mu\_\{QK\}\(t\)depends strongly onrr: for large ratios \(e\.g\.,r=100r=100\),μQ​K​\(t\)\\mu\_\{QK\}\(t\)grows rapidly, leading to attention mass on relevant tokens close to11, whereas for small ratios \(e\.g\.,r=0\.01r=0\.01\),μQ​K​\(t\)\\mu\_\{QK\}\(t\)remains small, resulting in attention close to uniform despite loss becoming zero\. This decoupling between logit growth and attention sharpening shows that selective attention emerges from the relative optimization dynamics of theQ​KQKandO​VOVrather than from improved prediction\.

## 7Experiments

The goal of our experiments is to test qualitative predictions of the theory developed in Section[6](https://arxiv.org/html/2608.06776#S6)\. We do not aim to optimize task performance or explore architectural design choices\. Instead, we evaluate whether changing the relative learning rates of the query\-key and output\-value circuits produces the attention behavior predicted by the theory\.

The theory makes two core predictions\. First, increasing the ratior=ηQ​KηO​Vr=\\dfrac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}should lead to sharper attention\. Second, this attention sharpening should occur while predictive performance remains comparable\. All experiments in this section are designed as controlled interventions on this learning rate ratio\.

We compare a baseline transformer to a variant trained in a higher query\-key learning\-rate regime\. In this variant, we increase the learning rate of the query\-key parameters while keeping all other settings fixed\. This directly increases the ratiorr\. This intervention corresponds directly to the mechanism analyzed in Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3)\.

We evaluate three datasets with token\-level relevance annotations: SQuAD, HateXplain, and subject\-verb agreement \(SVA\)\. For SQuAD, we use a 6\-layer GPT\-style model\. For HateXplain and SVA, we use 1\-layer models to match the theoretical setting\. Additional multi\-layer results are provided in the Appendix[D\.5](https://arxiv.org/html/2608.06776#A4.SS5)\. Here, predictive performance refers to the standard task metric for each dataset: F1 score for SQuAD and classification accuracy for HateXplain and SVA\.

All models are trained with AdamW\. We first tune a single base learning rate using validation performance\. This rate is shared across all parameters in the baseline model\. In the increasedrrregime, we multiply the query\-key learning rate by a fixed factor while keeping the output\-value learning rate unchanged\. We emphasize that adaptive optimization alone does not induce this behavior, as adaptive optimizers treat all parameter groups equally\. We tune the query–key multiplier using validation performance to ensure stable training and comparable predictive accuracy across settings\. Importantly, this tuning is not used to optimize attention metrics\. For concreteness, we use base learning rates of5×10−55\\times 10^\{\-5\},5×10−55\\times 10^\{\-5\}, and1×10−51\\times 10^\{\-5\}for SQuAD, HateXplain, and SVA respectively, and increase the query\-key learning rate by fixed multipliers of20×20\\times,30×30\\times, and100×100\\timeswhile keeping the output\-value learning rate unchanged\. Hyperparameter details are in Appendix[A\.2](https://arxiv.org/html/2608.06776#A1.SS2)\.

We also report attention\-based metrics that directly test the theoretical predictions\. We measure mean relevant token attention \(MRTA\) and use attention rollout\(Abnar & Zuidema,[2020](https://arxiv.org/html/2608.06776#bib.bib1)\)to aggregate the scores across layers\.

We report sufficiency and comprehensiveness scores as faithfulness metrics\(DeYoung et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib11)\)\. Sufficiency measures if important tokens are enough to retain the original prediction \(lower is better\), while comprehensiveness measures the change in probability of the predicted class after removing important tokens \(higher is better\)\.

Finally, we use DTAP\-based summary metrics\. Attention Confidence \(AC\) measures the percentage of instances in which a large fraction of attention mass is assigned to task\-relevant tokens\. Attention Confidence & Model Confidence \(ACMC\) measures the percentage of confidently correct instances that also exhibit high attention concentration\. These metrics are designed to capture the increase in attention massα​\(t\)\\alpha\(t\)predicted by the theory\. We discuss these metrics in detail in Appendix[D\.6](https://arxiv.org/html/2608.06776#A4.SS6)\.

Table[1](https://arxiv.org/html/2608.06776#S6.T1)summarizes the results across all datasets\. Predictive performance remains stable across tasks\. Increasingrrdoes not degrade test performance, and in some cases yields small improvements\. This confirms that attention sharpening is not driven by improved prediction alone\.

At the same time, attention becomes substantially sharper\. Both AC and ACMC increase by large margins across all datasets\. Sufficiency decreases and comprehensiveness increases in all settings, indicating more faithful attention\. These effects are consistent across datasets and persist in deeper models\. They match the qualitative behavior predicted by the optimization dynamics in Section[6](https://arxiv.org/html/2608.06776#S6)\.

Overall, the results support the theoretical mechanism proposed in this work\. When output\-value learning is slower, optimization increases the query\-key scale\. This concentrates attention through the softmax nonlinearity\. The model compensates for slower output learning by attending more strongly to informative tokens\. Importantly, this effect does not require larger output weights or improved predictive performance\. It arises from relative optimization speeds alone\.

## 8Conclusion

In this work, we study how optimization dynamics shape attention structure in self\-attention models\. By analyzing the interaction between the query\-key and output\-value circuits, we showed that attention sharpening is controlled by the relative learning speeds of these components rather than by predictive performance alone\. In particular, faster learning of the query\-key circuit leads to sharper attention, even when performance is comparable\. Our theoretical analysis provides a mechanistic explanation for this behavior, which is supported by experiments on synthetic data and real\-world benchmarks\. These results highlight the importance of parameterization and learning\-rate choices\. These factors are often treated as implementation details but play a central role in shaping attention behavior\. Our findings suggest that attention interpretability can be systematically influenced through optimization dynamics, without requiring architectural changes\.

## 9Limitations

The theoretical analysis in the paper focuses on a single\-layer attention\-only transformer\. Although our empirical results suggest that the core findings extend to standard multi\-layer transformers, additional components such as layer normalization and feedforward networks may introduce dynamics that require further theoretical study\. Our conclusions are also tied to the simplifying assumptions used in the theory and experiments\. Sharper attention or improved alignment with annotation\-based metrics does not necessarily imply a fully trustworthy explanation of the model’s reasoning, especially in high\-stakes settings\. The learning\-rate intervention should be viewed as one concrete probe of relative QK\-OV optimization dynamics rather than as the only possible causal mechanism\. Other circuit\-specific interventions, such as initialization imbalance, weight decay, or gradient normalization, are left for future work\. We do not claim that sharper attention directly improves downstream task performance\. Our results show that attention structure can be controlled, with improvements on attention\-based interpretability proxies\.

## Impact Statement

This work studies how optimization dynamics influence attention patterns in self\-attention models\. A potential positive impact is improved understanding of internal model behavior for interpretability and downstream performance\. However, sharper attention or improved alignment with annotation\-based metrics does not necessarily imply faithful explanation, and should not be used on its own in high\-stakes settings\.

## Acknowledgements

We thank the Department of Computer Science and Engineering at IIT Madras for providing travel support through the Kris Gopalakrishnan Endowment Fund\.

## References

- Abnar & Zuidema \(2020\)Abnar, S\. and Zuidema, W\.Quantifying attention flow in transformers\.In Jurafsky, D\., Chai, J\., Schluter, N\., and Tetreault, J\. \(eds\.\),*Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics*, pp\. 4190–4197, Online, July 2020\. Association for Computational Linguistics\.doi:10\.18653/v1/2020\.acl\-main\.385\.URL[https://aclanthology\.org/2020\.acl\-main\.385/](https://aclanthology.org/2020.acl-main.385/)\.
- Arora et al\. \(2018\)Arora, S\., Cohen, N\., and Hazan, E\.On the optimization of deep networks: Implicit acceleration by overparameterization\.*ArXiv*, abs/1802\.06509, 2018\.URL[https://api\.semanticscholar\.org/CorpusID:3331630](https://api.semanticscholar.org/CorpusID:3331630)\.
- Bao et al\. \(2024\)Bao, H\., Hataya, R\., and Karakida, R\.Self\-attention networks localize when qk\-eigenspectrum concentrates\.In*Proceedings of the 41st International Conference on Machine Learning*, ICML’24\. JMLR\.org, 2024\.
- Bhattamishra et al\. \(2020a\)Bhattamishra, S\., Ahuja, K\., and Goyal, N\.On the Ability and Limitations of Transformers to Recognize Formal Languages\.In Webber, B\., Cohn, T\., He, Y\., and Liu, Y\. \(eds\.\),*Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing \(EMNLP\)*, pp\. 7096–7116, Online, November 2020a\. Association for Computational Linguistics\.doi:10\.18653/v1/2020\.emnlp\-main\.576\.URL[https://aclanthology\.org/2020\.emnlp\-main\.576/](https://aclanthology.org/2020.emnlp-main.576/)\.
- Bhattamishra et al\. \(2020b\)Bhattamishra, S\., Patel, A\., and Goyal, N\.On the computational power of transformers and its implications in sequence modeling\.In Fernández, R\. and Linzen, T\. \(eds\.\),*Proceedings of the 24th Conference on Computational Natural Language Learning*, pp\. 455–475, Online, November 2020b\. Association for Computational Linguistics\.doi:10\.18653/v1/2020\.conll\-1\.37\.URL[https://aclanthology\.org/2020\.conll\-1\.37/](https://aclanthology.org/2020.conll-1.37/)\.
- Bietti et al\. \(2023\)Bietti, A\., Cabannes, V\., Bouchacourt, D\., Jégou, H\., and Bottou, L\.Birth of a transformer: a memory viewpoint\.In*Proceedings of the 37th International Conference on Neural Information Processing Systems*, NIPS ’23, Red Hook, NY, USA, 2023\. Curran Associates Inc\.
- Brown et al\. \(2020\)Brown, T\. B\., Mann, B\., Ryder, N\., Subbiah, M\., Kaplan, J\., Dhariwal, P\., Neelakantan, A\., Shyam, P\., Sastry, G\., Askell, A\., Agarwal, S\., Herbert\-Voss, A\., Krueger, G\., Henighan, T\., Child, R\., Ramesh, A\., Ziegler, D\. M\., Wu, J\., Winter, C\., Hesse, C\., Chen, M\., Sigler, E\., Litwin, M\., Gray, S\., Chess, B\., Clark, J\., Berner, C\., McCandlish, S\., Radford, A\., Sutskever, I\., and Amodei, D\.Language Models are Few\-Shot Learners, July 2020\.URL[http://arxiv\.org/abs/2005\.14165](http://arxiv.org/abs/2005.14165)\.arXiv:2005\.14165 \[cs\]\.
- Dehghani et al\. \(2018\)Dehghani, M\., Gouws, S\., Vinyals, O\., Uszkoreit, J\., and Kaiser, L\.Universal transformers\.*ArXiv*, abs/1807\.03819, 2018\.URL[https://api\.semanticscholar\.org/CorpusID:49667762](https://api.semanticscholar.org/CorpusID:49667762)\.
- Deora et al\. \(2024a\)Deora, P\., Ghaderi, R\., Taheri, H\., and Thrampoulidis, C\.On the optimization and generalization of multi\-head attention\.*Transactions on Machine Learning Research*, 2024a\.ISSN 2835\-8856\.URL[https://openreview\.net/forum?id=wTGjn7JvYK](https://openreview.net/forum?id=wTGjn7JvYK)\.
- Deora et al\. \(2024b\)Deora, P\., Ghaderi, R\., Taheri, H\., and Thrampoulidis, C\.On the optimization and generalization of multi\-head attention, 2024b\.URL[https://arxiv\.org/abs/2310\.12680](https://arxiv.org/abs/2310.12680)\.
- DeYoung et al\. \(2020\)DeYoung, J\., Jain, S\., Rajani, N\. F\., Lehman, E\., Xiong, C\., Socher, R\., and Wallace, B\. C\.ERASER: A benchmark to evaluate rationalized NLP models\.In Jurafsky, D\., Chai, J\., Schluter, N\., and Tetreault, J\. \(eds\.\),*Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics*, pp\. 4443–4458, Online, July 2020\. Association for Computational Linguistics\.doi:10\.18653/v1/2020\.acl\-main\.408\.URL[https://aclanthology\.org/2020\.acl\-main\.408/](https://aclanthology.org/2020.acl-main.408/)\.
- Dosovitskiy et al\. \(2020\)Dosovitskiy, A\., Beyer, L\., Kolesnikov, A\., Weissenborn, D\., Zhai, X\., Unterthiner, T\., Dehghani, M\., Minderer, M\., Heigold, G\., Gelly, S\., Uszkoreit, J\., and Houlsby, N\.An image is worth 16x16 words: Transformers for image recognition at scale\.*ArXiv*, abs/2010\.11929, 2020\.URL[https://api\.semanticscholar\.org/CorpusID:225039882](https://api.semanticscholar.org/CorpusID:225039882)\.
- Du et al\. \(2018\)Du, S\. S\., Hu, W\., and Lee, J\. D\.Algorithmic regularization in learning deep homogeneous models: layers are automatically balanced\.In*Proceedings of the 32nd International Conference on Neural Information Processing Systems*, NIPS’18, pp\. 382–393, Red Hook, NY, USA, 2018\. Curran Associates Inc\.
- Elhage et al\. \(2021\)Elhage, N\., Nanda, N\., Olsson, C\., Henighan, T\., Joseph, N\., Mann, B\., Askell, A\., Bai, Y\., Chen, A\., Conerly, T\., DasSarma, N\., Drain, D\., Ganguli, D\., Hatfield\-Dodds, Z\., Hernandez, D\., Jones, A\., Kernion, J\., Lovitt, L\., Ndousse, K\., Amodei, D\., Brown, T\., Clark, J\., Kaplan, J\., McCandlish, S\., and Olah, C\.A mathematical framework for transformer circuits\.*Transformer Circuits Thread*, 2021\.https://transformer\-circuits\.pub/2021/framework/index\.html\.
- Geshkovski et al\. \(2023\)Geshkovski, B\., Letrouit, C\., Polyanskiy, Y\., and Rigollet, P\.The emergence of clusters in self\-attention dynamics\.In*Proceedings of the 37th International Conference on Neural Information Processing Systems*, NIPS ’23, Red Hook, NY, USA, 2023\. Curran Associates Inc\.
- Jain & Wallace \(2019\)Jain, S\. and Wallace, B\. C\.Attention is not explanation\.In*North American Chapter of the Association for Computational Linguistics*, 2019\.URL[https://api\.semanticscholar\.org/CorpusID:67855860](https://api.semanticscholar.org/CorpusID:67855860)\.
- Karagodin et al\. \(2024\)Karagodin, N\., Polyanskiy, Y\., and Rigollet, P\.Clustering in causal attention masking\.In*Proceedings of the 38th International Conference on Neural Information Processing Systems*, NIPS ’24, Red Hook, NY, USA, 2024\. Curran Associates Inc\.ISBN 9798331314385\.
- Latif et al\. \(2023\)Latif, S\., Zaidi, A\., Cuayáhuitl, H\., Shamshad, F\., Shoukat, M\., and Qadir, J\.Transformers in speech processing: A survey\.*ArXiv*, abs/2303\.11607, 2023\.URL[https://api\.semanticscholar\.org/CorpusID:257636830](https://api.semanticscholar.org/CorpusID:257636830)\.
- Li et al\. \(2023\)Li, Y\., Li, Y\., and Risteski, A\.How do transformers learn topic structure: towards a mechanistic understanding\.In*Proceedings of the 40th International Conference on Machine Learning*, ICML’23\. JMLR\.org, 2023\.
- Li et al\. \(2024\)Li, Y\., Huang, Y\., Ildiz, M\. E\., Rawat, A\. S\., and Oymak, S\.Mechanics of next token prediction with self\-attention\.*ArXiv*, abs/2403\.08081, 2024\.URL[https://api\.semanticscholar\.org/CorpusID:268379753](https://api.semanticscholar.org/CorpusID:268379753)\.
- Linzen et al\. \(2016\)Linzen, T\., Dupoux, E\., and Goldberg, Y\.Assessing the ability of LSTMs to learn syntax\-sensitive dependencies\.*Transactions of the Association for Computational Linguistics*, 4:521–535, 2016\.doi:10\.1162/tacl˙a˙00115\.URL[https://aclanthology\.org/Q16\-1037/](https://aclanthology.org/Q16-1037/)\.
- Lu et al\. \(2021\)Lu, H\., Mao, Y\., and Nayak, A\.On the dynamics of training attention models\.In*International Conference on Learning Representations*, 2021\.URL[https://openreview\.net/forum?id=1OCTOShAmqB](https://openreview.net/forum?id=1OCTOShAmqB)\.
- Mathew et al\. \(2020\)Mathew, B\., Saha, P\., Yimam, S\. M\., Biemann, C\., Goyal, P\., and Mukherjee, A\.Hatexplain: A benchmark dataset for explainable hate speech detection\.In*AAAI Conference on Artificial Intelligence*, 2020\.URL[https://api\.semanticscholar\.org/CorpusID:229332119](https://api.semanticscholar.org/CorpusID:229332119)\.
- Olsson et al\. \(2022\)Olsson, C\., Elhage, N\., Nanda, N\., Joseph, N\., DasSarma, N\., Henighan, T\., Mann, B\., Askell, A\., Bai, Y\., Chen, A\., Conerly, T\., Drain, D\., Ganguli, D\., Hatfield\-Dodds, Z\., Hernandez, D\., Johnston, S\., Jones, A\., Kernion, J\., Lovitt, L\., Ndousse, K\., Amodei, D\., Brown, T\., Clark, J\., Kaplan, J\., McCandlish, S\., and Olah, C\.In\-context learning and induction heads\.*Transformer Circuits Thread*, 2022\.https://transformer\-circuits\.pub/2022/in\-context\-learning\-and\-induction\-heads/index\.html\.
- Pandey et al\. \(2023\)Pandey, L\. N\., Vashisht, R\., and Ramaswamy, H\. G\.On the interpretability of attention networks\.In Khan, E\. and Gonen, M\. \(eds\.\),*Proceedings of The 14th Asian Conference on Machine Learning*, volume 189 of*Proceedings of Machine Learning Research*, pp\. 832–847\. PMLR, 12–14 Dec 2023\.URL[https://proceedings\.mlr\.press/v189/pandey23a\.html](https://proceedings.mlr.press/v189/pandey23a.html)\.
- Pérez et al\. \(2021\)Pérez, J\., Barceló, P\., and Marinkovic, J\.Attention is turing complete\.*J\. Mach\. Learn\. Res\.*, 22\(1\), January 2021\.ISSN 1532\-4435\.
- Radford et al\. \(2018\)Radford, A\., Narasimhan, K\., Salimans, T\., and Sutskever, I\.Improving language understanding by generative pre\-training\.2018\.
- Rajpurkar et al\. \(2016\)Rajpurkar, P\., Zhang, J\., Lopyrev, K\., and Liang, P\.SQuAD: 100,000\+ questions for machine comprehension of text\.In Su, J\., Duh, K\., and Carreras, X\. \(eds\.\),*Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing*, pp\. 2383–2392, Austin, Texas, November 2016\. Association for Computational Linguistics\.doi:10\.18653/v1/D16\-1264\.URL[https://aclanthology\.org/D16\-1264/](https://aclanthology.org/D16-1264/)\.
- Serrano & Smith \(2019\)Serrano, S\. and Smith, N\. A\.Is attention interpretable?In Korhonen, A\., Traum, D\., and Màrquez, L\. \(eds\.\),*Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics*, pp\. 2931–2951, Florence, Italy, July 2019\. Association for Computational Linguistics\.doi:10\.18653/v1/P19\-1282\.URL[https://aclanthology\.org/P19\-1282/](https://aclanthology.org/P19-1282/)\.
- Tian et al\. \(2023\)Tian, Y\., Wang, Y\., Chen, B\., and Du, S\.Scan and snap: understanding training dynamics and token composition in 1\-layer transformer\.In*Proceedings of the 37th International Conference on Neural Information Processing Systems*, NIPS ’23, Red Hook, NY, USA, 2023\. Curran Associates Inc\.
- Vashisht & Ramaswamy \(2023\)Vashisht, R\. and Ramaswamy, H\. G\.On the learning dynamics of attention networks\.In*European Conference on Artificial Intelligence*, 2023\.URL[https://api\.semanticscholar\.org/CorpusID:260155204](https://api.semanticscholar.org/CorpusID:260155204)\.
- Vaswani et al\. \(2017\)Vaswani, A\., Shazeer, N\., Parmar, N\., Uszkoreit, J\., Jones, L\., Gomez, A\. N\., Kaiser, L\. u\., and Polosukhin, I\.Attention is all you need\.In Guyon, I\., Luxburg, U\. V\., Bengio, S\., Wallach, H\., Fergus, R\., Vishwanathan, S\., and Garnett, R\. \(eds\.\),*Advances in Neural Information Processing Systems*, volume 30\. Curran Associates, Inc\., 2017\.URL[https://proceedings\.neurips\.cc/paper\_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa\-Paper\.pdf](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)\.
- Wiegreffe & Pinter \(2019\)Wiegreffe, S\. and Pinter, Y\.Attention is not not explanation\.In Inui, K\., Jiang, J\., Ng, V\., and Wan, X\. \(eds\.\),*Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing \(EMNLP\-IJCNLP\)*, pp\. 11–20, Hong Kong, China, November 2019\. Association for Computational Linguistics\.doi:10\.18653/v1/D19\-1002\.URL[https://aclanthology\.org/D19\-1002/](https://aclanthology.org/D19-1002/)\.
- Yun et al\. \(2019\)Yun, C\., Bhojanapalli, S\., Rawat, A\. S\., Reddi, S\. J\., and Kumar, S\.Are transformers universal approximators of sequence\-to\-sequence functions?*ArXiv*, abs/1912\.10077, 2019\.URL[https://api\.semanticscholar\.org/CorpusID:209444410](https://api.semanticscholar.org/CorpusID:209444410)\.

## Appendix ASupplementary Material

### A\.1Code Reproducibility

### A\.2Dataset Details

#### A\.2\.1Synthetic Data Section[4](https://arxiv.org/html/2608.06776#S4)

The dataset for the task in Section[4](https://arxiv.org/html/2608.06776#S4)is generated as follows\. We consider44sequence classes, with44query and corresponding next tokens\. For each query token, we define a set of1010distinct tokens\. We also have a set of common tokens with cardinality1010\. Note that the position \(or index\) of the distinct and common tokens can be arbitrary\. We sample multiple such triples \(context tokens, query token, next token\), train a single\-layer self\-attention model on a subset of this dataset, and evaluate it on the remaining data\. The vocabulary size is5050tokens and sequence length is6464\. We train the model with200200batches with batch size3232using stochastic gradient descent\.

#### A\.2\.2Real\-World Dataset Section[7](https://arxiv.org/html/2608.06776#S7)

Table[2](https://arxiv.org/html/2608.06776#A1.T2)lists the three datasets used in the experiments\.

Table 2:Dataset statistics: number of instances and rationales for each split\.Table[3](https://arxiv.org/html/2608.06776#A1.T3)gives model architecture details for each dataset\.

Table 3:Model architecture configurations used in experiments\.HateXplain Data: We describe the validation and test results for the HateXplain dataset here\. For the preprocessing of text data, we follow\(Mathew et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib23)\)and clean the data using the same method\. For hyperparameter tuning, we first consider the same\-learning\-rate setting and then use the same initialization for the faster\-learning setting, except for the query\-key parameters, for which we use a learning rate of30×30\\timesthe base rate of5​e−55\\mathrm\{e\}\{\-5\}\.

Hyperparameter Tuning: We choose the base learning rate by hyperparameter tuning on validation data for learning rates in range\[0\.01,0\.000001\]\[0\.01,0\.000001\]\. Once we choose the base learning rate, we increase the learning rate of query\-key parameters by5×5\\times,10×10\\times,20×20\\times,30×30\\times, and100×100\\times\. In the main paper, we observe that increasing the learning rate of the query\-key parameters improves interpretability in terms of heatmaps for all the cases\. We report the best performing hyperparameters in the paper\. For this data, we also experiment with2−2\-layer and4−4\-layer models\.

Subject Verb Agreement Data: We further test our hypothesis on the Subject\-Verb Agreement \(SVA\) dataset \(binary\-classification\)\(Linzen et al\.,[2016](https://arxiv.org/html/2608.06776#bib.bib21)\), which evaluates the model’s ability to capture long\-range dependencies between subjects and verbs\. The dataset contains sentences designed to check grammatical number agreement \(e\.g\., whether a singular subject like “the cat” pairs with a singular verb like “runs,” or a plural subject like “the cats” pairs with “run”\)\. This task provides a controlled setting to examine whether faster learning for query\-key parameters, which improved interpretability for HateXplain, generalizes to core syntactic tasks\. We use the same single\-layer attention\-only model\. We train the model with SGD and batch size 256, and consider two settings: a same learning\-rate setting and a faster\-query\-key\-learning setting\.The validation and test results for SVA are described here\. For the preprocessing of text data, we follow\(Linzen et al\.,[2016](https://arxiv.org/html/2608.06776#bib.bib21)\)and clean the data using the same method\. For hyperparameter tuning, we first consider the same\-learning\-rate setting and then use the same weights for the faster\-learning setting, except for the query\-key parameters, for which we use a learning rate of100×100\\timesthe base rate of1​e−51\\mathrm\{e\}\{\-5\}\.

Hyperparameter Tuning: We choose the base learning rate by hyperparameter tuning on validation data for learning rates in range\[0\.01,0\.000001\]\[0\.01,0\.000001\]\. Once we choose the base learning rate, we increase the learning rate of query\-key parameters by5×5\\times,10×10\\times,20×20\\times,30×30\\times, and100×100\\times\. In the main paper, we observe that increasing the learning rate of the query\-key parameters improves interpretability in terms of heatmaps for all the cases\. We report the best performing hyperparameters in the paper\.

SQuAD QA Data: We also run large\-scale experiments on the subset of SQuAD question\-answering dataset\. Here, we finetune a 6\-layer GPT model to generate answers based on given contexts and questions\. We use the answer span information only for evaluation of attention\-based interpretability\. We train the model under two settings: \(1\) same learning rate for all parameters and \(2\) faster learning for query\-key parameters\(20×\(20\\timesthe base rate of5​e−55\\mathrm\{e\}\{\-5\}\)\. We also show the results on proxies for attention\-based interpretability after aggregating attention scores using the rollout method\(Abnar & Zuidema,[2020](https://arxiv.org/html/2608.06776#bib.bib1)\)\.

Hyperparameter Tuning: We choose the base learning rate by hyperparameter tuning on validation data for learning rates in range\[0\.01,0\.00001\]\[0\.01,0\.00001\]\. Once we choose the base learning rate, we increase the learning rate of query\-key parameters by5×5\\times,10×10\\times,20×20\\times,30×30\\times, and100×100\\times\. In the main paper, we observe that increasing the learning rate of the query\-key parameters improves interpretability proxies for interpretability in terms of heatmaps for all the cases\. We report the best performing hyperparameters in the paper\.

All DTAP heatmaps averaged over 5 runs are provided in Appendix[D\.5](https://arxiv.org/html/2608.06776#A4.SS5)\.

Standard Deviation Table for Table[1](https://arxiv.org/html/2608.06776#S6.T1)

Table 4:We report 95% confidence intervals, shown as±\\pmvalues, for each metric in Table[1](https://arxiv.org/html/2608.06776#S6.T1)\.

## Appendix BTheorem[1](https://arxiv.org/html/2608.06776#Thmtheorem1)Proof

Theorem[1](https://arxiv.org/html/2608.06776#Thmtheorem1)\[Parameter Trajectories\] Under the orthogonal full\-data setting, consider sequences of lengthTTwithmmdistinct tokens andnncommon tokens\. Suppose that𝑾O​V​\(0\)=0\{\\bm\{W\}\}\_\{OV\}\(0\)=0and𝑾Q​K​\(0\)=0\{\\bm\{W\}\}\_\{QK\}\(0\)=0, and letηO​V\\eta\_\{OV\}andηQ​K\\eta\_\{QK\}be the learning rates of𝑾O​V\{\\bm\{W\}\}\_\{OV\}and𝑾Q​K\{\\bm\{W\}\}\_\{QK\}, respectively\. Then the parameters evolve under population gradient flow as follows:

𝑾O​V\(k,:\)​\(t\)=μO​V​\(t\)​\[∑τ=1bskτ−1M​∑y=1M∑τ=1bsyτ\],\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)=\\mu\_\{OV\}\(t\)\\Bigg\[\\sum\\limits\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\dfrac\{1\}\{M\}\\sum\\limits\_\{y=1\}^\{M\}\\sum\\limits\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\],
𝑾Q​K\(k,:\)​\(t\)=μQ​K​\(t\)​\[∑y=1M∑τ=1bsyτ​∑τ′=1bsyτ′,k\],\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\(t\)=\\mu\_\{QK\}\(t\)\\Bigg\[\\sum\\limits\_\{y=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\},k\}\\Bigg\],
where the scalarsμO​V​\(t\)\\mu\_\{OV\}\(t\)andμQ​K​\(t\)\\mu\_\{QK\}\(t\)vary as follows:

d​μO​V​\(t\)d​t=ηO​V​α​\(t\)b∗\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\),\\dfrac\{d\\mu\_\{OV\}\(t\)\}\{dt\}=\\dfrac\{\\eta\_\{OV\}\\alpha\(t\)\}\{b\*\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\)\},
d​μQ​K​\(t\)d​t=ηQ​K​\(M−1\)​\(α​\(t\)−α2​\(t\)\)​μO​V​\(t\)b2​M​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\),\\dfrac\{d\\mu\_\{QK\}\(t\)\}\{dt\}=\\dfrac\{\\eta\_\{QK\}\(M\-1\)\(\\alpha\(t\)\-\\alpha^\{2\}\(t\)\)\\mu\_\{OV\}\(t\)\}\{b^\{2\}M\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\)\},
α​\(t\)=m×exp⁡\(μQ​K​\(t\)\)m∗\(exp⁡\(μQ​K​\(t\)\)\)\+n\.\\alpha\(t\)=\\dfrac\{m\\times\\exp\(\\mu\_\{QK\}\(t\)\)\}\{m\*\(\\exp\(\\mu\_\{QK\}\(t\)\)\)\+n\}\.
###### Proof\.

We consider the following form of the network output:

𝒛=𝑾O​V​𝒙~,𝒙~=𝑿⊤​𝜶,𝜶=softmax​\(𝑿​𝑾Q​K​𝒙T\),\{\\bm\{z\}\}=\{\\bm\{W\}\}\_\{OV\}\\,\\widetilde\{\{\\bm\{x\}\}\},\\qquad\\widetilde\{\{\\bm\{x\}\}\}=\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\},\\qquad\\bm\{\\alpha\}=\\mathrm\{softmax\}\(\{\\bm\{X\}\}\{\\bm\{W\}\}\_\{QK\}\{\\bm\{x\}\}\_\{T\}\),where𝑾O​V∈ℝM×d\{\\bm\{W\}\}\_\{OV\}\\in\\mathbb\{R\}^\{M\\times d\},𝑾Q​K∈ℝd×d\{\\bm\{W\}\}\_\{QK\}\\in\\mathbb\{R\}^\{d\\times d\},𝑿∈ℝT×d\{\\bm\{X\}\}\\in\\mathbb\{R\}^\{T\\times d\}stacks the sequence tokens, and𝒙T∈ℝd\{\\bm\{x\}\}\_\{T\}\\in\\mathbb\{R\}^\{d\}is the query token\.

Gradient of cross\-entropy with softmax\.

Let𝒛=𝑾O​V​𝑿⊤​𝜶\{\\bm\{z\}\}=\{\\bm\{W\}\}\_\{OV\}\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\}and𝒑=softmax​\(𝒛\)\{\\bm\{p\}\}=\\mathrm\{softmax\}\(\{\\bm\{z\}\}\)\. For the cross\-entropy lossℒ=−log⁡𝒑y\\mathcal\{L\}=\-\\log\{\\bm\{p\}\}\_\{y\},

−∇𝑾O​V\(k,:\)ℒ=\(𝟏​\[y=k\]−𝒑k\)​𝑿⊤​𝜶\.\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=\\big\(\\mathbf\{1\}\[y=k\]\-\{\\bm\{p\}\}\_\{k\}\\big\)\\,\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\}\.
Under population gradient flow \(averaging over the data distribution\),

−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=𝔼​\[\(𝟏​\[y=k\]−𝒑k\)​∑j=1T𝜶j​𝒙j⊤\]\.\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\big\(\\mathbf\{1\}\[y=k\]\-\{\\bm\{p\}\}\_\{k\}\\big\)\\,\\sum\_\{j=1\}^\{T\}\\bm\{\\alpha\}\_\{j\}\{\\bm\{x\}\}\_\{j\}^\{\\top\}\\right\]\.\(4\)
Note that𝜶\\bm\{\\alpha\}depends on the sampled sequence\. Using linearity of expectation,

𝔼​\[∑j=1T𝜶j​𝒙j\]=∑j=1T𝔼​\[𝜶j​𝒙j\]\.\\mathbb\{E\}\\\!\\left\[\\sum\_\{j=1\}^\{T\}\\bm\{\\alpha\}\_\{j\}\{\\bm\{x\}\}\_\{j\}\\right\]=\\sum\_\{j=1\}^\{T\}\\mathbb\{E\}\[\\bm\{\\alpha\}\_\{j\}\{\\bm\{x\}\}\_\{j\}\]\.
Let𝒟\\mathcal\{D\}denote themmdistinct\-token positions and𝒞\\mathcal\{C\}denote thenncommon\-token positions\. Define the scalar

α​\(t\):=∑j∈𝒟𝜶j,\\alpha\(t\):=\\sum\_\{j\\in\\mathcal\{D\}\}\\bm\{\\alpha\}\_\{j\},i\.e\., the total attention mass on distinct\-token positions\. By symmetry, all distinct positions have the same expected weight, and all common positions have the same expected weight:

𝔼​\[𝜶j\]=\{α​\(t\)/m,j∈𝒟,\(1−α​\(t\)\)/n,j∈𝒞\.\\mathbb\{E\}\[\\bm\{\\alpha\}\_\{j\}\]=\\begin\{cases\}\\alpha\(t\)/m,&j\\in\\mathcal\{D\},\\\\ \(1\-\\alpha\(t\)\)/n,&j\\in\\mathcal\{C\}\.\\end\{cases\}
Under the full\-data symmetry assumptions for distinct\-token sampling,

𝔼\[∑j∈𝒟𝒙j\|y\]=mb∑τ=1bsyτ\.\\mathbb\{E\}\\\!\\left\[\\sum\_\{j\\in\\mathcal\{D\}\}\{\\bm\{x\}\}\_\{j\}\\,\\middle\|\\,y\\right\]=\\frac\{m\}\{b\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\.Moreover, since the common\-token distribution has zero mean,𝔼​\[ν\]=0\\mathbb\{E\}\[\\nu\]=0, we have𝔼​\[∑j∈𝒞𝜶j​𝒙j\]=0\\mathbb\{E\}\\\!\\left\[\\sum\_\{j\\in\\mathcal\{C\}\}\\bm\{\\alpha\}\_\{j\}\{\\bm\{x\}\}\_\{j\}\\right\]=0\.

Therefore,

𝔼​\[𝑿⊤​𝜶∣y\]\\displaystyle\\mathbb\{E\}\[\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\}\\mid y\]=𝔼\[∑j∈𝒟𝜶j𝒙j\|y\]=α​\(t\)m𝔼\[∑j∈𝒟𝒙j\|y\]=α​\(t\)b∑τ=1bsyτ\.\\displaystyle=\\mathbb\{E\}\\\!\\left\[\\sum\_\{j\\in\\mathcal\{D\}\}\\bm\{\\alpha\}\_\{j\}\{\\bm\{x\}\}\_\{j\}\\,\\middle\|\\,y\\right\]=\\frac\{\\alpha\(t\)\}\{m\}\\,\\mathbb\{E\}\\\!\\left\[\\sum\_\{j\\in\\mathcal\{D\}\}\{\\bm\{x\}\}\_\{j\}\\,\\middle\|\\,y\\right\]=\\frac\{\\alpha\(t\)\}\{b\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\.\(5\)
Substituting \([5](https://arxiv.org/html/2608.06776#A2.E5)\) into \([4](https://arxiv.org/html/2608.06776#A2.E4)\) yields

−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=1M​∑y=1M\(𝟏​\[y=k\]−𝒑k\)​α​\(t\)b​∑τ=1bsyτ\.\\displaystyle=\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\\Big\(\\mathbf\{1\}\[y=k\]\-\{\\bm\{p\}\}\_\{k\}\\Big\)\\,\\frac\{\\alpha\(t\)\}\{b\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\.\(6\)
Now consider the ansatz

𝑾O​V\(k,:\)​\(t\)=μO​V​\(t\)​\[∑τ=1bskτ−1M​∑y′=1M∑τ=1bsy′τ\]\.\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)=\\mu\_\{OV\}\(t\)\\left\[\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{y^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y^\{\\prime\}\}^\{\\tau\}\\right\]\.Let

𝒔~k:=∑τ=1bskτ−1M​∑y′=1M∑τ=1bsy′τ\.\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}:=\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{y^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y^\{\\prime\}\}^\{\\tau\}\.
Using orthogonality of the vectors\{syτ\}y,τ\\\{s\_\{y\}^\{\\tau\}\\\}\_\{y,\\tau\},

𝒔~k⊤​∑τ=1bsyτ\\displaystyle\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}^\{\\top\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}=\(∑τ=1bskτ\)⊤​\(∑τ=1bsyτ\)−1M​\(∑y′=1M∑τ=1bsy′τ\)⊤​\(∑τ=1bsyτ\)\\displaystyle=\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\\right\)^\{\\top\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)\-\\frac\{1\}\{M\}\\left\(\\sum\_\{y^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y^\{\\prime\}\}^\{\\tau\}\\right\)^\{\\top\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)=b​1​\[k=y\]−bM=b​\(𝟏​\[k=y\]−1M\)\.\\displaystyle=b\\,\\mathbf\{1\}\[k=y\]\-\\frac\{b\}\{M\}=b\\\!\\left\(\\mathbf\{1\}\[k=y\]\-\\frac\{1\}\{M\}\\right\)\.\(7\)
Hence, for a sequence with labelyy,

𝑾O​V​𝔼​\[𝑿⊤​𝜶∣y\]\\displaystyle\{\\bm\{W\}\}\_\{OV\}\\,\\mathbb\{E\}\[\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\}\\mid y\]=μO​V​\(t\)​α​\(t\)b​\[𝒔~1⊤​∑τ=1bsyτ⋮𝒔~M⊤​∑τ=1bsyτ\]=μO​V​\(t\)​α​\(t\)​\(ey−1M​𝟏\)\.\\displaystyle=\\mu\_\{OV\}\(t\)\\frac\{\\alpha\(t\)\}\{b\}\\begin\{bmatrix\}\\widetilde\{\{\\bm\{s\}\}\}\_\{1\}^\{\\top\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\\\ \\vdots\\\\ \\widetilde\{\{\\bm\{s\}\}\}\_\{M\}^\{\\top\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\end\{bmatrix\}=\\mu\_\{OV\}\(t\)\\alpha\(t\)\\left\(e\_\{y\}\-\\frac\{1\}\{M\}\\mathbf\{1\}\\right\)\.\(8\)
Using shift\-invariance of softmax, this implies

𝒑=softmax​\(μO​V​\(t\)​α​\(t\)​ey\),\{\\bm\{p\}\}=\\mathrm\{softmax\}\\\!\\left\(\\mu\_\{OV\}\(t\)\\alpha\(t\)\\,e\_\{y\}\\right\),and hence

𝒑k=y=exp⁡\(α​\(t\)​μO​V​\(t\)\)exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1,𝒑k≠y=1exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\.\\displaystyle\{\\bm\{p\}\}\_\{k=y\}=\\frac\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\},\\qquad\{\\bm\{p\}\}\_\{k\\neq y\}=\\frac\{1\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\.\(9\)
Using \([9](https://arxiv.org/html/2608.06776#A2.E9)\) into \([6](https://arxiv.org/html/2608.06776#A2.E6)\), we have,

−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=1M​∑y=1M\(𝟏​\[y=k\]−𝒑k\)​α​\(t\)b​∑τ=1bsyτ\.\\displaystyle=\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\\Big\(\\mathbf\{1\}\[y=k\]\-\{\\bm\{p\}\}\_\{k\}\\Big\)\\,\\frac\{\\alpha\(t\)\}\{b\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\.\(10\)
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b∗M​\[\[1−𝒑k=y\]​∑τ=1bskτ\+∑y=1,y≠kM−𝒑k≠y​∑τ=1bsyτ\]\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\*M\}\\Bigg\[\\Big\[1\-\{\\bm\{p\}\}\_\{k=y\}\\Big\]\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\+\\sum\_\{y=1,y\\neq k\}^\{M\}\-\{\\bm\{p\}\}\_\{k\\neq y\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\(11\)
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b∗M\[\[1−exp⁡\(α​\(t\)​μO​V​\(t\)\)exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\]∑τ=1bskτ\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\*M\}\\Bigg\[\\Big\[1\-\\frac\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\Big\]\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}−∑y=1,y≠kM1exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1∑τ=1bsyτ\]\.\\displaystyle\-\\sum\_\{y=1,y\\neq k\}^\{M\}\\frac\{1\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\.
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b∗M\[\[M−1exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\]∑τ=1bskτ\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\*M\}\\Bigg\[\\Big\[\\frac\{M\-1\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\Big\]\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}−∑y=1,y≠kM1exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1∑τ=1bsyτ\]\.\\displaystyle\-\\sum\_\{y=1,y\\neq k\}^\{M\}\\frac\{1\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\.
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b∗M​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​\[\[M−1\]​∑τ=1bskτ−∑y=1,y≠kM∑τ=1bsyτ\]\.\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\*M\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\)\}\\Bigg\[\\Big\[M\-1\\Big\]\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\sum\_\{y=1,y\\neq k\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\.
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b∗M​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​\[M​∑τ=1bskτ−∑y=1M∑τ=1bsyτ\]\.\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\*M\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\)\}\\Bigg\[M\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\sum\_\{y=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\.
−∇𝑾O​V\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=α​\(t\)b​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​\[∑τ=1bskτ−1M​∑y=1M∑τ=1bsyτ\]\.\\displaystyle=\\frac\{\\alpha\(t\)\}\{b\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\)\}\\Bigg\[\\sum\_\{\\tau=1\}^\{b\}s\_\{k\}^\{\\tau\}\-\\dfrac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\Bigg\]\.
−∇𝑾O​V\(k,:\)ℒ=α​\(t\)b​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​𝒔~k\.\\displaystyle\\boxed\{\-\\nabla\_\{\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\}\\mathcal\{L\}=\\frac\{\\alpha\(t\)\}\{b\\left\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\right\)\}\\,\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}\}\.\(12\)
Therefore, under gradient flow,

d​𝑾O​V\(k,:\)​\(t\)d​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)\}\{dt\}=ηO​V​α​\(t\)b​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​𝒔~k\.\\displaystyle=\\eta\_\{OV\}\\,\\frac\{\\alpha\(t\)\}\{b\\left\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\right\)\}\\,\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}\.\(13\)
On the other hand, differentiating the ansatz𝑾O​V\(k,:\)​\(t\)=μO​V​\(t\)​𝒔~k\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)=\\mu\_\{OV\}\(t\)\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}with respect to time yields

d​𝑾O​V\(k,:\)​\(t\)d​t=d​μO​V​\(t\)d​t​𝒔~k\.\\frac\{d\{\\bm\{W\}\}\_\{OV\}^\{\(k,:\)\}\(t\)\}\{dt\}=\\frac\{d\\mu\_\{OV\}\(t\)\}\{dt\}\\,\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}\.
Since𝒔~k≠0\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}\\neq 0and the same direction𝒔~k\\widetilde\{\{\\bm\{s\}\}\}\_\{k\}appears in \([13](https://arxiv.org/html/2608.06776#A2.E13)\), we obtain the scalar differential equation

d​μO​V​\(t\)d​t=ηO​V​α​\(t\)b​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)\.\\frac\{d\\mu\_\{OV\}\(t\)\}\{dt\}=\\frac\{\\eta\_\{OV\}\\,\\alpha\(t\)\}\{b\\left\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\right\)\}\.
Key\-Query Parameter Trajectory Proof

Recallx~=X⊤​𝜶\\widetilde\{x\}=X^\{\\top\}\\bm\{\\alpha\}and𝜶=softmax​\(𝑿​𝑾Q​K​xT\)\\bm\{\\alpha\}=\\mathrm\{softmax\}\(\{\\bm\{X\}\}\{\\bm\{W\}\}\_\{QK\}x\_\{T\}\)\. Let𝒛=𝑾O​V​𝒙~\{\\bm\{z\}\}=\{\\bm\{W\}\}\_\{OV\}\\widetilde\{\{\\bm\{x\}\}\}and𝒑=softmax​\(𝒛\)\{\\bm\{p\}\}=\\mathrm\{softmax\}\(\{\\bm\{z\}\}\)\.

−∇𝑾Q​K\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\}\\mathcal\{L\}=1M​∑y=1M𝒙T​∑j=1T𝜶j​𝒙j,k​\(𝒙j−𝒙~\)⊤​\[𝑾O​V\(y,:\)−∑k′=1M𝒑k′​𝑾O​V\(k′,:\)\]\.\\displaystyle=\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\{\\bm\{x\}\}\_\{T\}\\sum\_\{j=1\}^\{T\}\\bm\{\\alpha\}\_\{j\}\\,\{\\bm\{x\}\}\_\{j,k\}\\,\(\{\\bm\{x\}\}\_\{j\}\-\\widetilde\{\{\\bm\{x\}\}\}\)^\{\\top\}\\left\[\{\\bm\{W\}\}\_\{OV\}^\{\(y,:\)\}\-\\sum\_\{k^\{\\prime\}=1\}^\{M\}\{\\bm\{p\}\}\_\{k^\{\\prime\}\}\\,\{\\bm\{W\}\}\_\{OV\}^\{\(k^\{\\prime\},:\)\}\\right\]\.\(14\)
Using \([8](https://arxiv.org/html/2608.06776#A2.E8)\), the logits satisfy

𝑾O​V​\(𝑿⊤​𝜶\)=μO​V​\(t\)​α​\(t\)​\[𝒆y−1M​𝟏\],\{\\bm\{W\}\}\_\{OV\}\(\{\\bm\{X\}\}^\{\\top\}\\bm\{\\alpha\}\)=\\mu\_\{OV\}\(t\)\\alpha\(t\)\\left\[\{\\bm\{e\}\}\_\{y\}\-\\frac\{1\}\{M\}\\mathbf\{1\}\\right\],and therefore

𝑾O​V\(y,:\)−∑k′=1M𝒑k′​𝑾O​V\(k′,:\)\\displaystyle\{\\bm\{W\}\}\_\{OV\}^\{\(y,:\)\}\-\\sum\_\{k^\{\\prime\}=1\}^\{M\}\{\\bm\{p\}\}\_\{k^\{\\prime\}\}\\,\{\\bm\{W\}\}\_\{OV\}^\{\(k^\{\\prime\},:\)\}=M​μO​V​\(t\)exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1​\[∑τ=1bsyτ−1M​∑k′=1M∑τ′=1bsk′τ′\]\.\\displaystyle=\\frac\{M\\mu\_\{OV\}\(t\)\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\left\[\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{k^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{k^\{\\prime\}\}^\{\\tau^\{\\prime\}\}\\right\]\.\(15\)
We now analyze the remaining factor in \([14](https://arxiv.org/html/2608.06776#A2.E14)\) under population expectation\. Let𝒟\\mathcal\{D\}denote the multiset of distinct\-token positions in the sequence \(\|𝒟\|=m\|\\mathcal\{D\}\|=m, counting multiplicities\), and𝒞\\mathcal\{C\}denote the multiset of common\-token positions \(\|𝒞\|=n\|\\mathcal\{C\}\|=n\)\.

Under the trajectory ansatz forWQ​KW\_\{QK\},

WQ​K\(k,:\)​\(t\)=μQ​K​\(t\)​∑y=1M\(∑τ=1bsyτ\)​\(∑τ′=1bsyτ′,k\),W\_\{QK\}^\{\(k,:\)\}\(t\)=\\mu\_\{QK\}\(t\)\\sum\_\{y=1\}^\{M\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)\\left\(\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\},k\}\\right\),the attention scores𝒖j=𝒙j⊤​𝑾Q​K​𝒙T\{\\bm\{u\}\}\_\{j\}=\{\\bm\{x\}\}\_\{j\}^\{\\top\}\{\\bm\{W\}\}\_\{QK\}\{\\bm\{x\}\}\_\{T\}take only two values: for distinct\-token positionsj∈𝒟j\\in\\mathcal\{D\},𝒖j=μQ​K​\(t\)\{\\bm\{u\}\}\_\{j\}=\\mu\_\{QK\}\(t\), while for common\-token positionsj∈𝒞j\\in\\mathcal\{C\},𝒖j=0\{\\bm\{u\}\}\_\{j\}=0\. Consequently, the softmax attention vector𝜶\\bm\{\\alpha\}is uniform within each group:

𝜶j=\{α​\(t\)/m,j∈𝒟,\(1−α​\(t\)\)/n,j∈𝒞,α​\(t\)=∑j∈𝒟𝜶j\.\\bm\{\\alpha\}\_\{j\}=\\begin\{cases\}\\alpha\(t\)/m,&j\\in\\mathcal\{D\},\\\\ \(1\-\\alpha\(t\)\)/n,&j\\in\\mathcal\{C\},\\end\{cases\}\\qquad\\alpha\(t\)=\\sum\_\{j\\in\\mathcal\{D\}\}\\bm\{\\alpha\}\_\{j\}\.
Writing𝒙j=syτ\{\\bm\{x\}\}\_\{j\}=s\_\{y\}^\{\\tau\}for distinct tokens and𝒙j=ντ\{\\bm\{x\}\}\_\{j\}=\\nu^\{\\tau\}for common tokens, the population expectation of the first factor in \([14](https://arxiv.org/html/2608.06776#A2.E14)\) becomes

1M​∑y=1M𝒙T​∑j=1Tαj​𝒙j,k​\(𝒙j−𝒙~\)⊤\\displaystyle\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\{\\bm\{x\}\}\_\{T\}\\sum\_\{j=1\}^\{T\}\\alpha\_\{j\}\\,\{\\bm\{x\}\}\_\{j,k\}\\,\(\{\\bm\{x\}\}\_\{j\}\-\\widetilde\{\{\\bm\{x\}\}\}\)^\{\\top\}=1M∑y=1M\(1b∑r=1bsyr\)\{α​\(t\)m∑τ=1bsyτ,k\(syτ−α​\(t\)m∑τ′=1bsyτ′−1−α​\(t\)n∑τ′′=1nντ′′\)⊤\\displaystyle=\\frac\{1\}\{M\}\\sum\_\{y=1\}^\{M\}\\left\(\\frac\{1\}\{b\}\\sum\_\{r=1\}^\{b\}s\_\{y\}^\{r\}\\right\)\\Bigg\\\{\\frac\{\\alpha\(t\)\}\{m\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau,k\}\\Bigg\(s\_\{y\}^\{\\tau\}\-\\frac\{\\alpha\(t\)\}\{m\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\}\}\-\\frac\{1\-\\alpha\(t\)\}\{n\}\\sum\_\{\\tau^\{\\prime\\prime\}=1\}^\{n\}\\nu^\{\\tau^\{\\prime\\prime\}\}\\Bigg\)^\{\\\!\\top\}\+1−α​\(t\)n∑τ=1nντ,k\(ντ−α​\(t\)m∑τ′=1bsyτ′−1−α​\(t\)n∑τ′′=1nντ′′\)⊤\}\.\\displaystyle\\qquad\\qquad\\qquad\\qquad\+\\frac\{1\-\\alpha\(t\)\}\{n\}\\sum\_\{\\tau=1\}^\{n\}\\nu^\{\\tau,k\}\\Bigg\(\\nu^\{\\tau\}\-\\frac\{\\alpha\(t\)\}\{m\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\}\}\-\\frac\{1\-\\alpha\(t\)\}\{n\}\\sum\_\{\\tau^\{\\prime\\prime\}=1\}^\{n\}\\nu^\{\\tau^\{\\prime\\prime\}\}\\Bigg\)^\{\\\!\\top\}\\Bigg\\\}\.\(16\)
Substituting \([15](https://arxiv.org/html/2608.06776#A2.E15)\) and \([16](https://arxiv.org/html/2608.06776#A2.E16)\) into \([14](https://arxiv.org/html/2608.06776#A2.E14)\) yields

−∇𝑾Q​K\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\}\\mathcal\{L\}=μO​V​\(t\)exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1​\[∑y=1M\(1b​∑r=1bsyr\)​\{⋯\}\]​\[∑τ=1bsyτ−1M​∑k′=1M∑τ′=1bsk′τ′\],\\displaystyle=\\frac\{\\mu\_\{OV\}\(t\)\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\Bigg\[\\sum\_\{y=1\}^\{M\}\\left\(\\frac\{1\}\{b\}\\sum\_\{r=1\}^\{b\}s\_\{y\}^\{r\}\\right\)\\\{\\cdots\\\}\\Bigg\]\\left\[\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{k^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{k^\{\\prime\}\}^\{\\tau^\{\\prime\}\}\\right\],\(17\)where\{⋯\}\\\{\\cdots\\\}denotes the bracketed expression in \([16](https://arxiv.org/html/2608.06776#A2.E16)\)\.

Under the orthogonality assumptions and the zero\-mean condition𝔼​\[ντ\]=0\\mathbb\{E\}\[\\nu^\{\\tau\}\]=0, all cross terms involving common tokens vanish under the population expectation\. The expression therefore reduces to the distinct\-token component:

−∇𝑾Q​K\(k,:\)ℒ\\displaystyle\-\\nabla\_\{\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\}\\mathcal\{L\}=μO​V​\(t\)exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1​∑y=1M\(1b​∑r=1bsyr\)​C~y​Vy,\\displaystyle=\\frac\{\\mu\_\{OV\}\(t\)\}\{\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\}\\sum\_\{y=1\}^\{M\}\\left\(\\frac\{1\}\{b\}\\sum\_\{r=1\}^\{b\}s\_\{y\}^\{r\}\\right\)\\widetilde\{C\}\_\{y\}\\,V\_\{y\},\(18\)where

C~y\\displaystyle\\widetilde\{C\}\_\{y\}=α​\(t\)m​∑τ=1bsyτ,k​\(syτ−α​\(t\)m​∑τ′=1bsyτ′\)⊤,\\displaystyle=\\frac\{\\alpha\(t\)\}\{m\}\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau,k\}\\left\(s\_\{y\}^\{\\tau\}\-\\frac\{\\alpha\(t\)\}\{m\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\}\}\\right\)^\{\\\!\\top\},\(19\)Vy\\displaystyle V\_\{y\}=\[∑τ=1bsyτ−1M​∑k′=1M∑τ′=1bsk′τ′\]\.\\displaystyle=\\left\[\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\-\\frac\{1\}\{M\}\\sum\_\{k^\{\\prime\}=1\}^\{M\}\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{k^\{\\prime\}\}^\{\\tau^\{\\prime\}\}\\right\]\.\(20\)
Using orthogonality of\{syτ\}\\\{s\_\{y\}^\{\\tau\}\\\}, the product\(1b​∑r=1bsyr\)​C~y​Vy\\left\(\\frac\{1\}\{b\}\\sum\_\{r=1\}^\{b\}s\_\{y\}^\{r\}\\right\)\\widetilde\{C\}\_\{y\}V\_\{y\}simplifies to a rank\-one term, yielding

−∇𝑾Q​K\(k,:\)ℒ=μO​V​\(t\)​\(M−1\)​α​\(t\)​\(1−α​\(t\)\)M​b2​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​∑y=1M\(∑τ=1bsyτ\)​\(∑τ′=1bsyτ′,k\)\.\\displaystyle\\boxed\{\-\\nabla\_\{\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\}\\mathcal\{L\}=\\frac\{\\mu\_\{OV\}\(t\)\(M\-1\)\\alpha\(t\)\\big\(1\-\\alpha\(t\)\\big\)\}\{Mb^\{2\}\\big\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\big\)\}\\sum\_\{y=1\}^\{M\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)\\left\(\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\},k\}\\right\)\}\.\(21\)
In continuous time,

d​𝑾Q​K\(k,:\)d​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{QK\}^\{\(k,:\)\}\}\{dt\}=ηQ​K​μO​V​\(t\)​\(M−1\)​α​\(t\)​\(1−α​\(t\)\)M​b2​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)​∑y=1M\(∑τ=1bsyτ\)​\(∑τ′=1bsyτ′,k\),\\displaystyle=\\eta\_\{QK\}\\,\\frac\{\\mu\_\{OV\}\(t\)\(M\-1\)\\alpha\(t\)\\big\(1\-\\alpha\(t\)\\big\)\}\{Mb^\{2\}\\big\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\big\)\}\\sum\_\{y=1\}^\{M\}\\left\(\\sum\_\{\\tau=1\}^\{b\}s\_\{y\}^\{\\tau\}\\right\)\\left\(\\sum\_\{\\tau^\{\\prime\}=1\}^\{b\}s\_\{y\}^\{\\tau^\{\\prime\},k\}\\right\),\(22\)and hence

d​μQ​K​\(t\)d​t\\displaystyle\\frac\{d\\mu\_\{QK\}\(t\)\}\{dt\}=ηQ​K​μO​V​\(t\)​\(M−1\)​α​\(t\)​\(1−α​\(t\)\)M​b2​\(exp⁡\(α​\(t\)​μO​V​\(t\)\)\+M−1\)\.\\displaystyle=\\eta\_\{QK\}\\,\\frac\{\\mu\_\{OV\}\(t\)\(M\-1\)\\alpha\(t\)\\big\(1\-\\alpha\(t\)\\big\)\}\{Mb^\{2\}\\big\(\\exp\(\\alpha\(t\)\\mu\_\{OV\}\(t\)\)\+M\-1\\big\)\}\.\(23\)
Finally, since distinct\-token positions have scoreμQ​K​\(t\)\\mu\_\{QK\}\(t\)and common\-token positions have score0, the total attention mass on distinct\-token positions is

α​\(t\)=m​exp⁡\(μQ​K​\(t\)\)m​exp⁡\(μQ​K​\(t\)\)\+n\.\\alpha\(t\)=\\frac\{m\\exp\(\\mu\_\{QK\}\(t\)\)\}\{m\\exp\(\\mu\_\{QK\}\(t\)\)\+n\}\.∎

## Appendix CLemma[3](https://arxiv.org/html/2608.06776#Thmlemma3)Proof

Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3)\[Part I:μO​V​\(t\)\\mu\_\{OV\}\(t\)Bounds\] LetμO​V​\(t\)\\mu\_\{OV\}\(t\)evolve as

μ˙O​V​\(t\)=ηO​V​α​\(t\)b​\(eα​\(t\)​μO​V​\(t\)\+M−1\),μO​V​\(0\)=0,\\dot\{\\mu\}\_\{OV\}\(t\)=\\frac\{\\eta\_\{OV\}\\,\\alpha\(t\)\}\{b\\big\(e^\{\\alpha\(t\)\\mu\_\{OV\}\(t\)\}\+M\-1\\big\)\},\\qquad\\mu\_\{OV\}\(0\)=0,whereα​\(t\)=m​eμQ​K​\(t\)m​eμQ​K​\(t\)\+n\\alpha\(t\)=\\frac\{me^\{\\mu\_\{QK\}\(t\)\}\}\{me^\{\\mu\_\{QK\}\(t\)\}\+n\}andα0:=α​\(0\)=mm\+n\\alpha\_\{0\}:=\\alpha\(0\)=\\frac\{m\}\{m\+n\}\. Then for allt≥0t\\geq 0,

log⁡\(1\+ηO​V​α0M​b​t\)≤μO​V​\(t\)≤1α0​log⁡\(1\+ηO​V​α0b​t\)\.\\displaystyle\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\\Big\)\\;\\leq\\;\\mu\_\{OV\}\(t\)\\;\\leq\\;\\frac\{1\}\{\\alpha\_\{0\}\}\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\\Big\)\.\(24\)
###### Proof\.

Fromμ˙Q​K​\(t\)≥0\\dot\{\\mu\}\_\{QK\}\(t\)\\geq 0andμQ​K=0\\mu\_\{QK\}=0, we haveμQ​K​\(t\)≥0\\mu\_\{QK\}\(t\)\\geq 0and non\-decreasing; sinceα​\(\)\\alpha\(\)is increasing,α​\(t\)\\alpha\(t\)is non\-decreasing andα0≤α​\(t\)≤1\\alpha\_\{0\}\\leq\\alpha\(t\)\\leq 1

For anyx≥0x\\geq 0,ex≤ex\+M−1≤M​exe^\{x\}\\leq e^\{x\}\+M\-1\\leq Me^\{x\}, hence

1M​e−x≤1ex\+M−1≤e−x\.\\frac\{1\}\{M\}e^\{\-x\}\\leq\\frac\{1\}\{e^\{x\}\+M\-1\}\\leq e^\{\-x\}\.Apply this withx=α​\(t\)​μO​V​\(t\)x=\\alpha\(t\)\\mu\_\{OV\}\(t\)in the ODE to obtain

ηO​V​α​\(t\)M​b​e−α​\(t\)​μO​V​\(t\)≤μ˙O​V​\(t\)≤ηO​V​α​\(t\)b​e−α​\(t\)​μO​V​\(t\)\.\\displaystyle\\frac\{\\eta\_\{OV\}\\alpha\(t\)\}\{Mb\}e^\{\-\\alpha\(t\)\\mu\_\{OV\}\(t\)\}\\;\\leq\\;\\dot\{\\mu\}\_\{OV\}\(t\)\\;\\leq\\;\\frac\{\\eta\_\{OV\}\\alpha\(t\)\}\{b\}e^\{\-\\alpha\(t\)\\mu\_\{OV\}\(t\)\}\.\(25\)
Usingα​\(t\)≤1\\alpha\(t\)\\leq 1ande−α​\(t\)​μ≤e−α0​μe^\{\-\\alpha\(t\)\\mu\}\\leq e^\{\-\\alpha\_\{0\}\\mu\}\(sinceα​\(t\)≥α0\\alpha\(t\)\\geq\\alpha\_\{0\}andμ≥0\\mu\\geq 0\), the right inequality of \([25](https://arxiv.org/html/2608.06776#A3.E25)\) implies

μ˙O​V​\(t\)≤ηO​Vb​e−α0​μO​V​\(t\)\.\\dot\{\\mu\}\_\{OV\}\(t\)\\leq\\frac\{\\eta\_\{OV\}\}\{b\}\\,e^\{\-\\alpha\_\{0\}\\mu\_\{OV\}\(t\)\}\.Multiplying byα0​eα0​μO​V​\(t\)\\alpha\_\{0\}e^\{\\alpha\_\{0\}\\mu\_\{OV\}\(t\)\}gives

α0​eα0​μO​V​\(t\)​μ˙O​V​\(t\)≤ηO​V​α0b,\\alpha\_\{0\}e^\{\\alpha\_\{0\}\\mu\_\{OV\}\(t\)\}\\dot\{\\mu\}\_\{OV\}\(t\)\\leq\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\},i\.e\.,

dd​t​\(eα0​μO​V​\(t\)\)≤ηO​V​α0b\.\\frac\{d\}\{dt\}\\Big\(e^\{\\alpha\_\{0\}\\mu\_\{OV\}\(t\)\}\\Big\)\\leq\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\.Integrating from0tottand usingeα0​μO​V​\(0\)=1e^\{\\alpha\_\{0\}\\mu\_\{OV\}\(0\)\}=1yields

eα0​μO​V​\(t\)≤1\+ηO​V​α0b​t\.e^\{\\alpha\_\{0\}\\mu\_\{OV\}\(t\)\}\\leq 1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\.Applyinglog⁡\(⋅\)\\log\(\\cdot\)and dividing byα0\\alpha\_\{0\}gives

μO​V​\(t\)≤1α0​log⁡\(1\+ηO​V​α0b​t\)\.\\mu\_\{OV\}\(t\)\\leq\\frac\{1\}\{\\alpha\_\{0\}\}\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\\Big\)\.
Usingα​\(t\)≥α0\\alpha\(t\)\\geq\\alpha\_\{0\}ande−α​\(t\)​μ≥e−μe^\{\-\\alpha\(t\)\\mu\}\\geq e^\{\-\\mu\}\(sinceα​\(t\)≤1\\alpha\(t\)\\leq 1andμ≥0\\mu\\geq 0\), the left inequality of \([25](https://arxiv.org/html/2608.06776#A3.E25)\) implies

μ˙O​V​\(t\)≥ηO​V​α0M​b​e−μO​V​\(t\)\.\\dot\{\\mu\}\_\{OV\}\(t\)\\geq\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,e^\{\-\\mu\_\{OV\}\(t\)\}\.Multiplying byeμO​V​\(t\)e^\{\\mu\_\{OV\}\(t\)\}gives

eμO​V​\(t\)​μ˙O​V​\(t\)≥ηO​V​α0M​b,e^\{\\mu\_\{OV\}\(t\)\}\\dot\{\\mu\}\_\{OV\}\(t\)\\geq\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\},i\.e\.,

dd​t​\(eμO​V​\(t\)\)≥ηO​V​α0M​b\.\\frac\{d\}\{dt\}\\Big\(e^\{\\mu\_\{OV\}\(t\)\}\\Big\)\\geq\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\.Integrating from0tottand usingeμO​V​\(0\)=1e^\{\\mu\_\{OV\}\(0\)\}=1yields

eμO​V​\(t\)≥1\+ηO​V​α0M​b​t\.e^\{\\mu\_\{OV\}\(t\)\}\\geq 1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\.Applyinglog⁡\(⋅\)\\log\(\\cdot\)gives

μO​V​\(t\)≥log⁡\(1\+ηO​V​α0M​b​t\)\.\\mu\_\{OV\}\(t\)\\geq\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\\Big\)\.Combining the two bounds proves \([24](https://arxiv.org/html/2608.06776#A3.E24)\)\. ∎

Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3)\[Part II:μQ​K​\(t\)\\mu\_\{QK\}\(t\)Bounds\] LetμO​V​\(t\),μQ​K​\(t\)\\mu\_\{OV\}\(t\),\\mu\_\{QK\}\(t\)satisfy the coupled ODEs

μ˙O​V​\(t\)=ηO​V​α​\(t\)b​\(eα​\(t\)​μO​V​\(t\)\+M−1\),μ˙Q​K​\(t\)=ηQ​K​\(M−1\)​α​\(t\)​\(1−α​\(t\)\)​μO​V​\(t\)M​b2​\(eα​\(t\)​μO​V​\(t\)\+M−1\),\\dot\{\\mu\}\_\{OV\}\(t\)=\\frac\{\\eta\_\{OV\}\\,\\alpha\(t\)\}\{b\\big\(e^\{\\alpha\(t\)\\mu\_\{OV\}\(t\)\}\+M\-1\\big\)\},\\qquad\\dot\{\\mu\}\_\{QK\}\(t\)=\\eta\_\{QK\}\\frac\{\(M\-1\)\\alpha\(t\)\(1\-\\alpha\(t\)\)\\mu\_\{OV\}\(t\)\}\{Mb^\{2\}\\big\(e^\{\\alpha\(t\)\\mu\_\{OV\}\(t\)\}\+M\-1\\big\)\},withμO​V​\(0\)=μQ​K​\(0\)=0\\mu\_\{OV\}\(0\)=\\mu\_\{QK\}\(0\)=0andα​\(t\)=m​eμQ​K​\(t\)m​eμQ​K​\(t\)\+n\\alpha\(t\)=\\frac\{me^\{\\mu\_\{QK\}\(t\)\}\}\{me^\{\\mu\_\{QK\}\(t\)\}\+n\}\. Letα0=mm\+n\\alpha\_\{0\}=\\frac\{m\}\{m\+n\}\.

For allt≥0t\\geq 0,

μQ​K​\(t\)≤ηQ​KηO​V​\(M−1\)​\(1−α0\)2​M​b​α02​log2⁡\(1\+ηO​V​α0b​t\)\.\\displaystyle\\mu\_\{QK\}\(t\)\\;\\leq\\;\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\(1\-\\alpha\_\{0\}\)\}\{2Mb\\,\\alpha\_\{0\}^\{2\}\}\\,\\log^\{2\}\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\\Big\)\.\(26\)
Fixδ∈\(0,1\)\\delta\\in\(0,1\)and assume

α​\(t\)≤1−δfor all​t∈\[0,T\]\.\\displaystyle\\alpha\(t\)\\leq 1\-\\delta\\qquad\\text\{for all \}t\\in\[0,T\]\.\(27\)Then for allt∈\[0,T\]t\\in\[0,T\],

μQ​K​\(t\)≥ηQ​KηO​V​\(M−1\)​δ2​M​b​log2⁡\(1\+ηO​V​α0M​b​t\)\.\\displaystyle\\mu\_\{QK\}\(t\)\\;\\geq\\;\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\\delta\}\{2Mb\}\\,\\log^\{2\}\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\\Big\)\.\(28\)
###### Proof\.

Divide theμQ​K\\mu\_\{QK\}\-ODE by theμO​V\\mu\_\{OV\}\-ODE to cancel the common denominator:

d​μQ​Kd​μO​V=μ˙Q​Kμ˙O​V=ηQ​KηO​V​M−1M​b​\(1−α​\(t\)\)​μO​V​\(t\)\.\\displaystyle\\frac\{d\\mu\_\{QK\}\}\{d\\mu\_\{OV\}\}=\\frac\{\\dot\{\\mu\}\_\{QK\}\}\{\\dot\{\\mu\}\_\{OV\}\}=\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{M\-1\}\{Mb\}\\,\(1\-\\alpha\(t\)\)\\,\\mu\_\{OV\}\(t\)\.\(29\)
Sinceα​\(t\)\\alpha\(t\)is nondecreasing andα​\(0\)=α0\\alpha\(0\)=\\alpha\_\{0\}, we haveα​\(t\)≥α0\\alpha\(t\)\\geq\\alpha\_\{0\}and hence1−α​\(t\)≤1−α01\-\\alpha\(t\)\\leq 1\-\\alpha\_\{0\}\. Using this in \([29](https://arxiv.org/html/2608.06776#A3.E29)\) yields

d​μQ​Kd​μO​V≤ηQ​KηO​V​M−1M​b​\(1−α0\)​μO​V\.\\frac\{d\\mu\_\{QK\}\}\{d\\mu\_\{OV\}\}\\leq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{M\-1\}\{Mb\}\\,\(1\-\\alpha\_\{0\}\)\\,\\mu\_\{OV\}\.Integrating overμO​V∈\[0,μO​V​\(t\)\]\\mu\_\{OV\}\\in\[0,\\mu\_\{OV\}\(t\)\]gives

μQ​K​\(t\)≤ηQ​KηO​V​M−1M​b​\(1−α0\)​∫0μO​V​\(t\)u​𝑑u=ηQ​KηO​V​\(M−1\)​\(1−α0\)2​M​b​μO​V​\(t\)2\.\\mu\_\{QK\}\(t\)\\leq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{M\-1\}\{Mb\}\\,\(1\-\\alpha\_\{0\}\)\\int\_\{0\}^\{\\mu\_\{OV\}\(t\)\}u\\,du=\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\(1\-\\alpha\_\{0\}\)\}\{2Mb\}\\,\\mu\_\{OV\}\(t\)^\{2\}\.Applying the upper bound from Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3),

μO​V​\(t\)≤1α0​log⁡\(1\+ηO​V​α0b​t\),\\mu\_\{OV\}\(t\)\\leq\\frac\{1\}\{\\alpha\_\{0\}\}\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\\Big\),we obtain

μQ​K​\(t\)≤ηQ​KηO​V​\(M−1\)​\(1−α0\)2​M​b​α02​log2⁡\(1\+ηO​V​α0b​t\),\\mu\_\{QK\}\(t\)\\leq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\(1\-\\alpha\_\{0\}\)\}\{2Mb\\,\\alpha\_\{0\}^\{2\}\}\\,\\log^\{2\}\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{b\}\\,t\\Big\),which is exactly \([26](https://arxiv.org/html/2608.06776#A3.E26)\)\.

Assume \([27](https://arxiv.org/html/2608.06776#A3.E27)\)\. Then1−α​\(t\)≥δ1\-\\alpha\(t\)\\geq\\deltafor allt∈\[0,T\]t\\in\[0,T\]\. Using this in \([29](https://arxiv.org/html/2608.06776#A3.E29)\) yields

d​μQ​Kd​μO​V≥ηQ​KηO​V​M−1M​b​δ​μO​V\.\\frac\{d\\mu\_\{QK\}\}\{d\\mu\_\{OV\}\}\\geq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{M\-1\}\{Mb\}\\,\\delta\\,\\mu\_\{OV\}\.Integrating overμO​V∈\[0,μO​V​\(t\)\]\\mu\_\{OV\}\\in\[0,\\mu\_\{OV\}\(t\)\]gives

μQ​K​\(t\)≥ηQ​KηO​V​M−1M​b​δ​∫0μO​V​\(t\)u​𝑑u=ηQ​KηO​V​\(M−1\)​δ2​M​b​μO​V​\(t\)2\.\\mu\_\{QK\}\(t\)\\geq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{M\-1\}\{Mb\}\\,\\delta\\int\_\{0\}^\{\\mu\_\{OV\}\(t\)\}u\\,du=\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\\delta\}\{2Mb\}\\,\\mu\_\{OV\}\(t\)^\{2\}\.Applying the lower bound ofμO​V\\mu\_\{OV\}from Lemma[3](https://arxiv.org/html/2608.06776#Thmlemma3),

μO​V​\(t\)≥log⁡\(1\+ηO​V​α0M​b​t\),\\mu\_\{OV\}\(t\)\\geq\\log\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\\Big\),we obtain

μQ​K​\(t\)≥ηQ​KηO​V​\(M−1\)​δ2​M​b​log2⁡\(1\+ηO​V​α0M​b​t\),\\mu\_\{QK\}\(t\)\\geq\\frac\{\\eta\_\{QK\}\}\{\\eta\_\{OV\}\}\\,\\frac\{\(M\-1\)\\delta\}\{2Mb\}\\,\\log^\{2\}\\\!\\Big\(1\+\\frac\{\\eta\_\{OV\}\\alpha\_\{0\}\}\{Mb\}\\,t\\Big\),which is exactly \([28](https://arxiv.org/html/2608.06776#A3.E28)\)\. ∎

## Appendix DRelation Between Collapsed and Factorized Parameterization

In this section, we provide proofs for Proposition[1](https://arxiv.org/html/2608.06776#Thmproposition1), Lemma[1](https://arxiv.org/html/2608.06776#Thmlemma1), and Lemma[2](https://arxiv.org/html/2608.06776#Thmlemma2)\. we also state the corresponding proposition and lemmas for QK circuit in proposition[2](https://arxiv.org/html/2608.06776#Thmproposition2),[4](https://arxiv.org/html/2608.06776#Thmlemma4)and[5](https://arxiv.org/html/2608.06776#Thmlemma5)\.

### D\.1Proposition[1](https://arxiv.org/html/2608.06776#Thmproposition1)Proof

Proposition[1](https://arxiv.org/html/2608.06776#Thmproposition1)Consider the network output

y^=𝑾O​𝑾V​x~,𝑾O∈ℝM×d,𝑾V∈ℝd×d,x~∈ℝd,\\widehat\{y\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\\widetilde\{x\},\\qquad\{\\bm\{W\}\}\_\{O\}\\in\\mathbb\{R\}^\{M\\times d\},\\ \{\\bm\{W\}\}\_\{V\}\\in\\mathbb\{R\}^\{d\\times d\},\\ \\widetilde\{x\}\\in\\mathbb\{R\}^\{d\},and define the collapsed output\-value parameter𝑾O​V:=𝑾O​𝑾V\{\\bm\{W\}\}\_\{OV\}:=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\. Lety~=𝕊​\(y^\)\\widetilde\{y\}=\{\\mathbb\{S\}\}\(\\widehat\{y\}\)and letℒ\\mathcal\{L\}be the cross\-entropy loss

ℒ=−log⁡\(exp⁡\(y^c\)∑kexp⁡\(y^k\)\)\.\\mathcal\{L\}=\-\\log\\left\(\\frac\{\\exp\(\\widehat\{y\}\_\{c\}\)\}\{\\sum\_\{k\}\\exp\(\\widehat\{y\}\_\{k\}\)\}\\right\)\.Then,

∂ℒ∂𝑾O​V\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{OV\}\}=\(y~−y\)​x~⊤\.\\displaystyle=\(\\widetilde\{y\}\-y\)\\,\\widetilde\{x\}^\{\\top\}\.\(30\)By chain rule,

∂ℒ∂𝑾O\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{O\}\}=\(y~−y\)​\(𝑾V​x~\)⊤,\\displaystyle=\(\\widetilde\{y\}\-y\)\\,\(\{\\bm\{W\}\}\_\{V\}\\widetilde\{x\}\)^\{\\top\},\(31\)∂ℒ∂𝑾V\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{V\}\}=𝑾O⊤​\(y~−y\)​x~⊤\.\\displaystyle=\{\\bm\{W\}\}\_\{O\}^\{\\top\}\(\\widetilde\{y\}\-y\)\\,\\widetilde\{x\}^\{\\top\}\.\(32\)Under gradient flow, the dynamics are

d​𝑾O​Vd​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{OV\}\}\{dt\}=\(y~−y\)​x~⊤,\\displaystyle=\(\\widetilde\{y\}\-y\)\\,\\widetilde\{x\}^\{\\top\},\(33\)d​𝑾Od​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{O\}\}\{dt\}=\(y~−y\)​x~⊤​𝑾V⊤,\\displaystyle=\(\\widetilde\{y\}\-y\)\\,\\widetilde\{x\}^\{\\top\}\{\\bm\{W\}\}\_\{V\}^\{\\top\},\(34\)d​𝑾Vd​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{V\}\}\{dt\}=𝑾O⊤​\(y~−y\)​x~⊤\.\\displaystyle=\{\\bm\{W\}\}\_\{O\}^\{\\top\}\(\\widetilde\{y\}\-y\)\\,\\widetilde\{x\}^\{\\top\}\.\(35\)
###### Proof\.

Equation \([30](https://arxiv.org/html/2608.06776#A4.E30)\) is the standard softmax cross\-entropy gradient\. Equations \([31](https://arxiv.org/html/2608.06776#A4.E31)\) \([32](https://arxiv.org/html/2608.06776#A4.E32)\) follow by chain rule, and \([33](https://arxiv.org/html/2608.06776#A4.E33)\)\-\([35](https://arxiv.org/html/2608.06776#A4.E35)\) are the corresponding gradient\-flow updates\. ∎

### D\.2Lemma[1](https://arxiv.org/html/2608.06776#Thmlemma1)Proof

Lemma[1](https://arxiv.org/html/2608.06776#Thmlemma1)\[OV Parameters Balancedness\] Under gradient flow and balanced initialization,

𝑾V​\(t\)​𝑾V⊤​\(t\)=𝑾O⊤​\(t\)​𝑾O​\(t\)∀t\.\{\\bm\{W\}\}\_\{V\}\(t\)\{\\bm\{W\}\}\_\{V\}^\{\\top\}\(t\)=\{\\bm\{W\}\}\_\{O\}^\{\\top\}\(t\)\{\\bm\{W\}\}\_\{O\}\(t\)\\qquad\\forall t\.
###### Proof\.

Using \([34](https://arxiv.org/html/2608.06776#A4.E34)\)–\([35](https://arxiv.org/html/2608.06776#A4.E35)\),

dd​t​\(𝑾V​𝑾V⊤\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{V\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}\\big\)=𝑾V​\(d​𝑾Vd​t\)⊤\+\(d​𝑾Vd​t\)​𝑾V⊤\\displaystyle=\{\\bm\{W\}\}\_\{V\}\\left\(\\frac\{d\{\\bm\{W\}\}\_\{V\}\}\{dt\}\\right\)^\{\\top\}\+\\left\(\\frac\{d\{\\bm\{W\}\}\_\{V\}\}\{dt\}\\right\)\{\\bm\{W\}\}\_\{V\}^\{\\top\}=𝑾V​x~​\(y~−y\)⊤​𝑾O\+𝑾O⊤​\(y~−y\)​x~⊤​𝑾V⊤\.\\displaystyle=\{\\bm\{W\}\}\_\{V\}\\widetilde\{x\}\(\\widetilde\{y\}\-y\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\+\{\\bm\{W\}\}\_\{O\}^\{\\top\}\(\\widetilde\{y\}\-y\)\\widetilde\{x\}^\{\\top\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}\.\(36\)Similarly,

dd​t​\(𝑾O⊤​𝑾O\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{O\}^\{\\top\}\{\\bm\{W\}\}\_\{O\}\\big\)=\(d​𝑾Od​t\)⊤​𝑾O\+𝑾O⊤​\(d​𝑾Od​t\)\\displaystyle=\\left\(\\frac\{d\{\\bm\{W\}\}\_\{O\}\}\{dt\}\\right\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\+\{\\bm\{W\}\}\_\{O\}^\{\\top\}\\left\(\\frac\{d\{\\bm\{W\}\}\_\{O\}\}\{dt\}\\right\)=𝑾V​x~​\(y~−y\)⊤​𝑾O\+𝑾O⊤​\(y~−y\)​x~⊤​𝑾V⊤\.\\displaystyle=\{\\bm\{W\}\}\_\{V\}\\widetilde\{x\}\(\\widetilde\{y\}\-y\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\+\{\\bm\{W\}\}\_\{O\}^\{\\top\}\(\\widetilde\{y\}\-y\)\\widetilde\{x\}^\{\\top\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}\.\(37\)Comparing \([36](https://arxiv.org/html/2608.06776#A4.E36)\) and \([37](https://arxiv.org/html/2608.06776#A4.E37)\) gives

dd​t​\(𝑾V​𝑾V⊤−𝑾O⊤​𝑾O\)=0\.\\frac\{d\}\{dt\}\\left\(\{\\bm\{W\}\}\_\{V\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}\-\{\\bm\{W\}\}\_\{O\}^\{\\top\}\{\\bm\{W\}\}\_\{O\}\\right\)=0\.Hence the difference is constant in time, and under \(approximately\) balanced initialization it is0, so𝑾V​𝑾V⊤=𝑾O⊤​𝑾O\{\\bm\{W\}\}\_\{V\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}=\{\\bm\{W\}\}\_\{O\}^\{\\top\}\{\\bm\{W\}\}\_\{O\}for alltt\. ∎

### D\.3Lemma[2](https://arxiv.org/html/2608.06776#Thmlemma2)Proof

Lemma[2](https://arxiv.org/html/2608.06776#Thmlemma2)\[Relation Between Factorized and Collapsed OV parameterization\] Let𝑾~O​V=𝑾O​𝑾V\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}and𝑾O​V\{\\bm\{W\}\}\_\{OV\}be the collapsed parameter\. Then

vec⁡\(∂𝑾~O​V∂t\)\\displaystyle\\operatorname\{vec\}\\\!\\left\(\\frac\{\\partial\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\}\{\\partial t\}\\right\)=Ω​\(𝑾O​V\)​vec⁡\(∂𝑾O​V∂t\),\\displaystyle=\\Omega\(\{\\bm\{W\}\}\_\{OV\}\)\\,\\operatorname\{vec\}\\\!\\left\(\\frac\{\\partial\{\\bm\{W\}\}\_\{OV\}\}\{\\partial t\}\\right\),\(38\)where

Ω​\(C\):=\(C⊤​C\)1/2⊗IM\+Id⊗\(C​C⊤\)1/2,\\Omega\(C\):=\(C^\{\\top\}C\)^\{1/2\}\\otimes I\_\{M\}\+I\_\{d\}\\otimes\(CC^\{\\top\}\)^\{1/2\},andΩ​\(α​C\)=α​Ω​\(C\)\\Omega\(\\alpha C\)=\\alpha\\,\\Omega\(C\)\(degree\-1 homogeneous\)\.

###### Proof\.

From Lemma[1](https://arxiv.org/html/2608.06776#Thmlemma1),𝑾V​𝑾V⊤=𝑾O⊤​𝑾O\{\\bm\{W\}\}\_\{V\}\{\\bm\{W\}\}\_\{V\}^\{\\top\}=\{\\bm\{W\}\}\_\{O\}^\{\\top\}\{\\bm\{W\}\}\_\{O\}\. For fixedtt, write𝑾O=𝑼O​ΣO​𝑽O⊤\{\\bm\{W\}\}\_\{O\}=\{\\bm\{U\}\}\_\{O\}\\Sigma\_\{O\}\{\\bm\{V\}\}\_\{O\}^\{\\top\}and𝑾V=𝑼V​ΣV​𝑽V⊤\{\\bm\{W\}\}\_\{V\}=\{\\bm\{U\}\}\_\{V\}\\Sigma\_\{V\}\{\\bm\{V\}\}\_\{V\}^\{\\top\}\. The Gram matrices share eigenvalues

ΣO⊤​ΣO=ΣV​ΣV⊤=d​i​a​g​\(ρ1​Id1,…,ρm​Idm\),\\Sigma\_\{O\}^\{\\top\}\\Sigma\_\{O\}=\\Sigma\_\{V\}\\Sigma\_\{V\}^\{\\top\}=diag\(\\rho\_\{1\}I\_\{d\_\{1\}\},\\ldots,\\rho\_\{m\}I\_\{d\_\{m\}\}\),henceΣO=ΣV=d​i​a​g​\(ρ1​Id1,…,ρm​Idm\)\\Sigma\_\{O\}=\\Sigma\_\{V\}=diag\(\\sqrt\{\\rho\_\{1\}\}I\_\{d\_\{1\}\},\\ldots,\\sqrt\{\\rho\_\{m\}\}I\_\{d\_\{m\}\}\)\. Therefore,

𝑾O​𝑾O⊤\\displaystyle\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{O\}^\{\\top\}=𝑼O​d​i​a​g​\(ρ1​Id1,…,ρm​Idm\)​𝑼O⊤,\\displaystyle=\{\\bm\{U\}\}\_\{O\}diag\(\\rho\_\{1\}I\_\{d\_\{1\}\},\\ldots,\\rho\_\{m\}I\_\{d\_\{m\}\}\)\{\\bm\{U\}\}\_\{O\}^\{\\top\},\(39\)𝑾V⊤​𝑾V\\displaystyle\{\\bm\{W\}\}\_\{V\}^\{\\top\}\{\\bm\{W\}\}\_\{V\}=𝑽V​d​i​a​g​\(ρ1​Id1,…,ρm​Idm\)​𝑽V⊤\.\\displaystyle=\{\\bm\{V\}\}\_\{V\}diag\(\\rho\_\{1\}I\_\{d\_\{1\}\},\\ldots,\\rho\_\{m\}I\_\{d\_\{m\}\}\)\{\\bm\{V\}\}\_\{V\}^\{\\top\}\.\(40\)
Moreover, since𝑾O​V=𝑾O​𝑾V\{\\bm\{W\}\}\_\{OV\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\},

𝑾O​V​𝑾O​V⊤\\displaystyle\{\\bm\{W\}\}\_\{OV\}\{\\bm\{W\}\}\_\{OV\}^\{\\top\}=𝑼O​d​i​a​g​\(ρ12​Id1,…,ρm2​Idm\)​𝑼O⊤,\\displaystyle=\{\\bm\{U\}\}\_\{O\}diag\(\\rho\_\{1\}^\{2\}I\_\{d\_\{1\}\},\\ldots,\\rho\_\{m\}^\{2\}I\_\{d\_\{m\}\}\)\{\\bm\{U\}\}\_\{O\}^\{\\top\},\(41\)𝑾O​V⊤​𝑾O​V\\displaystyle\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\{\\bm\{W\}\}\_\{OV\}=𝑽V​d​i​a​g​\(ρ12​Id1,…,ρm2​Idm\)​𝑽V⊤\.\\displaystyle=\{\\bm\{V\}\}\_\{V\}diag\(\\rho\_\{1\}^\{2\}I\_\{d\_\{1\}\},\\ldots,\\rho\_\{m\}^\{2\}I\_\{d\_\{m\}\}\)\{\\bm\{V\}\}\_\{V\}^\{\\top\}\.\(42\)Thus,

𝑾O​𝑾O⊤=\(𝑾O​V​𝑾O​V⊤\)1/2,𝑾V⊤​𝑾V=\(𝑾O​V⊤​𝑾O​V\)1/2\.\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{O\}^\{\\top\}=\(\{\\bm\{W\}\}\_\{OV\}\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\)^\{1/2\},\\qquad\{\\bm\{W\}\}\_\{V\}^\{\\top\}\{\\bm\{W\}\}\_\{V\}=\(\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\{\\bm\{W\}\}\_\{OV\}\)^\{1/2\}\.
Using the product rule,

d​𝑾~O​Vd​t=d​𝑾Od​t​𝑾V\+𝑾O​d​𝑾Vd​t\.\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\}\{dt\}=\\frac\{d\{\\bm\{W\}\}\_\{O\}\}\{dt\}\{\\bm\{W\}\}\_\{V\}\+\{\\bm\{W\}\}\_\{O\}\\frac\{d\{\\bm\{W\}\}\_\{V\}\}\{dt\}\.Substituting \([34](https://arxiv.org/html/2608.06776#A4.E34)\)\-\([35](https://arxiv.org/html/2608.06776#A4.E35)\) and the identities above yields

d​𝑾~O​Vd​t\\displaystyle\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\}\{dt\}=\(y−y~\)​𝒙~⊤​\(𝑾O​V⊤​𝑾O​V\)1/2\+\(𝑾O​V​𝑾O​V⊤\)1/2​\(y−y~\)​𝒙~⊤\.\\displaystyle=\(y\-\\widetilde\{y\}\)\\widetilde\{\{\\bm\{x\}\}\}^\{\\top\}\(\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\{\\bm\{W\}\}\_\{OV\}\)^\{1/2\}\+\(\{\\bm\{W\}\}\_\{OV\}\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\)^\{1/2\}\(y\-\\widetilde\{y\}\)\\widetilde\{\{\\bm\{x\}\}\}^\{\\top\}\.\(43\)vec⁡\(d​𝑾~O​Vd​t\)\\displaystyle\\operatorname\{vec\}\\\!\\left\(\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\}\{dt\}\\right\)=\(\(𝑾O​V⊤​𝑾O​V\)1/2⊗IM\+Id⊗\(𝑾O​V​𝑾O​V⊤\)1/2\)​vec⁡\(\(y−y~\)​𝒙~⊤\)\.\\displaystyle=\\Big\(\(\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\{\\bm\{W\}\}\_\{OV\}\)^\{1/2\}\\otimes I\_\{M\}\+I\_\{d\}\\otimes\(\{\\bm\{W\}\}\_\{OV\}\{\\bm\{W\}\}\_\{OV\}^\{\\top\}\)^\{1/2\}\\Big\)\\operatorname\{vec\}\\\!\\big\(\(y\-\\widetilde\{y\}\)\\widetilde\{\{\\bm\{x\}\}\}^\{\\top\}\\big\)\.\(44\)Using \([30](https://arxiv.org/html/2608.06776#A4.E30)\) and \([33](https://arxiv.org/html/2608.06776#A4.E33)\) ,

vec⁡\(d​𝑾~O​Vd​t\)=Ω​\(𝑾O​V\)​vec⁡\(d​𝑾O​Vd​t\),\\operatorname\{vec\}\\\!\\left\(\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{OV\}\}\{dt\}\\right\)=\\Omega\(\{\\bm\{W\}\}\_\{OV\}\)\\,\\operatorname\{vec\}\\\!\\left\(\\frac\{d\{\\bm\{W\}\}\_\{OV\}\}\{dt\}\\right\),which is \([38](https://arxiv.org/html/2608.06776#A4.E38)\)\. ∎

### D\.4Proposition and Lemma for QK circuit

###### Proposition 2\.

Considery^=𝐖O​𝐖V​x~\\widehat\{y\}=\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\\widetilde\{x\}withx~=X⊤​A\\widetilde\{x\}=X^\{\\top\}AandA=𝕊​\(𝐗​𝐖K​𝐖Q⊤​𝐱T\)A=\{\\mathbb\{S\}\}\(\{\\mathbf\{X\}\}\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{x\}\}\_\{T\}\), and define𝐖Q​K=𝐖K​𝐖Q⊤\{\\bm\{W\}\}\_\{QK\}=\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\.

LetZ=D​i​a​g​\(A\)−A​A⊤Z=Diag\(A\)\-AA^\{\\top\}\. Then the gradient w\.r\.t\.𝐖Q​K\{\\bm\{W\}\}\_\{QK\}is

∂ℒ∂𝑾Q​K\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{QK\}\}=𝒙T\(y~−y\)⊤𝑾O𝑾V𝐗⊤Z𝐗=:𝑮\.\\displaystyle=\{\\bm\{x\}\}\_\{T\}\(\\widetilde\{y\}\-y\)^\{\\top\}\{\\bm\{W\}\}\_\{O\}\{\\bm\{W\}\}\_\{V\}\{\\mathbf\{X\}\}^\{\\top\}Z\{\\mathbf\{X\}\}=:\\bm\{G\}\.\(45\)By chain rule,

∂ℒ∂𝑾Q\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{Q\}\}=𝑮⊤​𝑾K,\\displaystyle=\\bm\{G\}^\{\\top\}\{\\bm\{W\}\}\_\{K\},\(46\)∂ℒ∂𝑾K\\displaystyle\\frac\{\\partial\\mathcal\{L\}\}\{\\partial\{\\bm\{W\}\}\_\{K\}\}=𝑮​𝑾Q\.\\displaystyle=\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\.\(47\)
Under gradient flow, the dynamics are

d​𝑾Q​Kd​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{QK\}\}\{dt\}=𝑮,\\displaystyle=\\bm\{G\},\(48\)d​𝑾Qd​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{Q\}\}\{dt\}=𝑮⊤​𝑾K,\\displaystyle=\\bm\{G\}^\{\\top\}\{\\bm\{W\}\}\_\{K\},\(49\)d​𝑾Kd​t\\displaystyle\\frac\{d\{\\bm\{W\}\}\_\{K\}\}\{dt\}=𝑮​𝑾Q\.\\displaystyle=\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\.\(50\)

###### Proof\.

Equation \([45](https://arxiv.org/html/2608.06776#A4.E45)\) is the standard gradient with respect to the𝑾Q​K\{\\bm\{W\}\}\_\{QK\}parameters\. Equations \([46](https://arxiv.org/html/2608.06776#A4.E46)\) and \([47](https://arxiv.org/html/2608.06776#A4.E47)\) follow by chain rule, and \([49](https://arxiv.org/html/2608.06776#A4.E49)\) and \([50](https://arxiv.org/html/2608.06776#A4.E50)\) are the corresponding gradient\-flow updates\. ∎

###### Lemma 4\(QK Parameter Balancedness\)\.

Let𝐖Q​K=𝐖K​𝐖Q⊤\{\\bm\{W\}\}\_\{QK\}=\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}and define𝐆\\bm\{G\}as in Equation \([45](https://arxiv.org/html/2608.06776#A4.E45)\)\. Under gradient flow and balanced initialization,

𝑾Q⊤​\(t\)​𝑾Q​\(t\)=𝑾K⊤​\(t\)​𝑾K​\(t\)∀t\.\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\(t\)\{\\bm\{W\}\}\_\{Q\}\(t\)=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\(t\)\{\\bm\{W\}\}\_\{K\}\(t\)\\qquad\\forall t\.

###### Proof\.

dd​t​\(𝑾Q⊤​𝑾Q\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{W\}\}\_\{Q\}\\big\)=𝑾Q⊤​\(d​𝑾Qd​t\)\+\(d​𝑾Qd​t\)⊤​𝑾Q\\displaystyle=\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\\left\(\\frac\{d\{\\bm\{W\}\}\_\{Q\}\}\{dt\}\\right\)\+\\left\(\\frac\{d\{\\bm\{W\}\}\_\{Q\}\}\{dt\}\\right\)^\{\\top\}\{\\bm\{W\}\}\_\{Q\}=𝑾Q⊤​𝑮⊤​𝑾K\+𝑾K⊤​G​𝑾Q\.\\displaystyle=\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\\bm\{G\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}\+\{\\bm\{W\}\}\_\{K\}^\{\\top\}G\{\\bm\{W\}\}\_\{Q\}\.\(51\)
Similarly,

dd​t​\(𝑾K⊤​𝑾K\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{K\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}\\big\)=𝑾K⊤​\(d​𝑾Kd​t\)\+\(d​𝑾Qd​t\)⊤​𝑾Q\\displaystyle=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\\left\(\\frac\{d\{\\bm\{W\}\}\_\{K\}\}\{dt\}\\right\)\+\\left\(\\frac\{d\{\\bm\{W\}\}\_\{Q\}\}\{dt\}\\right\)^\{\\top\}\{\\bm\{W\}\}\_\{Q\}=𝑾K⊤​𝑮​𝑾Q\+𝑾Q⊤​G⊤​𝑾K\.\\displaystyle=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\+\{\\bm\{W\}\}\_\{Q\}^\{\\top\}G^\{\\top\}\{\\bm\{W\}\}\_\{K\}\.\(52\)
dd​t​\(𝑾Q⊤​𝑾Q\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{W\}\}\_\{Q\}\\big\)=𝑾K⊤​𝑮​𝑾Q\+𝑾Q⊤​𝑮⊤​𝑾K,\\displaystyle=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\+\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\\bm\{G\}^\{\\top\}\{\\bm\{W\}\}\_\{K\},\(53\)dd​t​\(𝑾K⊤​𝑾K\)\\displaystyle\\frac\{d\}\{dt\}\\big\(\{\\bm\{W\}\}\_\{K\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}\\big\)=𝑾Q⊤​𝑮⊤​𝑾K\+𝑾K⊤​𝑮​𝑾Q\.\\displaystyle=\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\\bm\{G\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}\+\{\\bm\{W\}\}\_\{K\}^\{\\top\}\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\.\(54\)Comparing \([53](https://arxiv.org/html/2608.06776#A4.E53)\) and \([54](https://arxiv.org/html/2608.06776#A4.E54)\) yields

dd​t​\(𝑾Q⊤​𝑾Q−𝑾K⊤​𝑾K\)=0\.\\frac\{d\}\{dt\}\\left\(\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{W\}\}\_\{Q\}\-\{\\bm\{W\}\}\_\{K\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}\\right\)=0\.Thus the difference is constant in time and equals0under balanced initialization, so𝑾Q⊤​𝑾Q=𝑾K⊤​𝑾K\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{W\}\}\_\{Q\}=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}for alltt\. ∎

###### Lemma 5\(QK preconditioning identity\)\.

Let𝐖~Q​K=𝐖K​𝐖Q⊤\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}=\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}and𝐖Q​K\{\\bm\{W\}\}\_\{QK\}be the collapsed parameter\. Then

vec⁡\(∂𝑾~Q​K∂t\)\\displaystyle\\operatorname\{vec\}\\\!\\left\(\\frac\{\\partial\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}\}\{\\partial t\}\\right\)=Ω​\(𝑾Q​K\)​vec⁡\(∂𝑾Q​K∂t\),\\displaystyle=\\Omega\(\{\\bm\{W\}\}\_\{QK\}\)\\,\\operatorname\{vec\}\\\!\\left\(\\frac\{\\partial\{\\bm\{W\}\}\_\{QK\}\}\{\\partial t\}\\right\),\(55\)where

Ω​\(C\):=\(C⊤​C\)1/2⊗Id\+Id⊗\(C​C⊤\)1/2,\\Omega\(C\):=\(C^\{\\top\}C\)^\{1/2\}\\otimes I\_\{d\}\+I\_\{d\}\\otimes\(CC^\{\\top\}\)^\{1/2\},andΩ​\(α​C\)=α​Ω​\(C\)\\Omega\(\\alpha C\)=\\alpha\\,\\Omega\(C\)\.

###### Proof\.

Using Lemma[4](https://arxiv.org/html/2608.06776#Thmlemma4),𝑾Q⊤​𝑾Q=𝑾K⊤​𝑾K\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\{\\bm\{W\}\}\_\{Q\}=\{\\bm\{W\}\}\_\{K\}^\{\\top\}\{\\bm\{W\}\}\_\{K\}, SVD argument of the proof of Lemma[2](https://arxiv.org/html/2608.06776#Thmlemma2), we obtain

𝑾Q​𝑾Q⊤=\(𝑾Q​K⊤​𝑾Q​K\)1/2,𝑾K​𝑾K⊤=\(𝑾Q​K​𝑾Q​K⊤\)1/2\.\{\\bm\{W\}\}\_\{Q\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}=\(\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\{\\bm\{W\}\}\_\{QK\}\)^\{1/2\},\\qquad\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{K\}^\{\\top\}=\(\{\\bm\{W\}\}\_\{QK\}\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\)^\{1/2\}\.
Using the product rule,

d​𝑾~Q​Kd​t=d​𝑾Kd​t​𝑾Q⊤\+𝑾K​\(d​𝑾Qd​t\)⊤=𝑮​𝑾Q​𝑾Q⊤\+𝑾K​𝑾K⊤​𝑮\.\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}\}\{dt\}=\\frac\{d\{\\bm\{W\}\}\_\{K\}\}\{dt\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\+\{\\bm\{W\}\}\_\{K\}\\left\(\\frac\{d\{\\bm\{W\}\}\_\{Q\}\}\{dt\}\\right\)^\{\\top\}=\\bm\{G\}\{\\bm\{W\}\}\_\{Q\}\{\\bm\{W\}\}\_\{Q\}^\{\\top\}\+\{\\bm\{W\}\}\_\{K\}\{\\bm\{W\}\}\_\{K\}^\{\\top\}\\bm\{G\}\.Thus,

d​𝑾~Q​Kd​t\\displaystyle\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}\}\{dt\}=𝑮​\(𝑾Q​K⊤​𝑾Q​K\)1/2\+\(𝑾Q​K​𝑾Q​K⊤\)1/2​𝑮\.\\displaystyle=\\bm\{G\}\(\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\{\\bm\{W\}\}\_\{QK\}\)^\{1/2\}\+\(\{\\bm\{W\}\}\_\{QK\}\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\)^\{1/2\}\\bm\{G\}\.\(56\)Vectorizing \([56](https://arxiv.org/html/2608.06776#A4.E56)\) and usingv​e​c​\(A​B\)=\(B⊤⊗I\)​v​e​c​\(A\)vec\(AB\)=\(B^\{\\top\}\\otimes I\)vec\(A\),

v​e​c​\(d​𝑾~Q​Kd​t\)\\displaystyle vec\\\!\\left\(\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}\}\{dt\}\\right\)=\(\(𝑾Q​K⊤​𝑾Q​K\)1/2⊗Id\+Id⊗\(𝑾Q​K​𝑾Q​K⊤\)1/2\)​v​e​c​\(𝑮\)\.\\displaystyle=\\Big\(\(\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\{\\bm\{W\}\}\_\{QK\}\)^\{1/2\}\\otimes I\_\{d\}\+I\_\{d\}\\otimes\(\{\\bm\{W\}\}\_\{QK\}\{\\bm\{W\}\}\_\{QK\}^\{\\top\}\)^\{1/2\}\\Big\)vec\(\\bm\{G\}\)\.Under gradient flow, \([45](https://arxiv.org/html/2608.06776#A4.E45)\) implies

v​e​c​\(d​𝑾~Q​Kd​t\)=Ω​\(𝑾Q​K\)​v​e​c​\(d​𝑾Q​Kd​t\),vec\\\!\\left\(\\frac\{d\\widetilde\{\{\\bm\{W\}\}\}\_\{QK\}\}\{dt\}\\right\)=\\Omega\(\{\\bm\{W\}\}\_\{QK\}\)\\,vec\\\!\\left\(\\frac\{d\{\\bm\{W\}\}\_\{QK\}\}\{dt\}\\right\),which is \([55](https://arxiv.org/html/2608.06776#A4.E55)\)\. ∎

### D\.5Additional Results DTAP Heatmaps and Multilayer Setting Results

Table 5:Train and test performance together with attention\-based interpretability metrics on real\-world datasets\.↑\\uparrowindicates higher is better, while↓\\downarrowindicates lower is better\. All results are reported as mean over 5 random seeds\. We report 95% confidence intervals over five random seeds in Table[6](https://arxiv.org/html/2608.06776#A4.T6)\.Table 6:In this table we report 95% CI as±\\pmvalues for each metric for Table[5](https://arxiv.org/html/2608.06776#A4.T5)\.![Refer to caption](https://arxiv.org/html/2608.06776v1/x19.png)\(a\)Baseline Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x20.png)\(b\)Baseline STD Dev
![Refer to caption](https://arxiv.org/html/2608.06776v1/x21.png)\(c\)Faster QK Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x22.png)\(d\)Faster QK STD Dev

Figure 5:66layer model Distinct Token Attention\-Prediction heat map of SQuAD QA Data \(a\)\(b\) Baseline Model Average and Standard Deviation \(c\)\(d\) FasterQK Model Average and Standard Deviation \(5 runs\)![Refer to caption](https://arxiv.org/html/2608.06776v1/x23.png)\(a\)Baseline Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x24.png)\(b\)Baseline STD Dev
![Refer to caption](https://arxiv.org/html/2608.06776v1/x25.png)\(c\)Faster QK Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x26.png)\(d\)Faster QK STD Dev

Figure 6:11layer model Distinct Token Attention\-Prediction heat map of SVA data \(a\)\(b\) Baseline Model Average and Standard Deviation \(c\)\(d\) Faster QK Model Average and Standard Deviation \(5 runs\)![Refer to caption](https://arxiv.org/html/2608.06776v1/x27.png)\(a\)Baseline Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x28.png)\(b\)Baseline STD Dev
![Refer to caption](https://arxiv.org/html/2608.06776v1/x29.png)\(c\)Faster QK Average
![Refer to caption](https://arxiv.org/html/2608.06776v1/x30.png)\(d\)Faster QK STD Dev

Figure 7:11layer model Distinct Token Attention\-Prediction heat map of HateXplain data \(a\)\(b\) Baseline Model Average and Standard Deviation \(c\)\(d\) Faster QK Model Average and Standard Deviation \(5 runs\)
### D\.6Detailed Metric Discussion

We adopt the definitions of sufficiency and comprehensiveness from\(DeYoung et al\.,[2020](https://arxiv.org/html/2608.06776#bib.bib11)\)\.

Sufficiency: Sufficiency captures whether the important tokens are sufficient to retain the original prediction\.

Sufficiency=f​\(x\)−f​\(r:k%\)\\textrm\{Sufficiency\}=f\(x\)\-f\(r\_\{:k\\%\}\),

wherer:k%r:k\\%denotes the top\-k%k\\%most important tokens ranked by attention score\.

Comprehensiveness: Comprehensiveness measures the change in the predicted\-class probability after removing important tokens\.

Comprehensiveness=f​\(x\)−f​\(x/r:k%\)\\textrm\{Comprehensiveness\}=f\(x\)\-f\(x/r\_\{:k\\%\}\)
wherer:k%r\_\{:k\\%\}refers to top\-k% most important tokens chosen based on attention scores\.

Mean Relevant Token Attention \(MRTA\): Using the learned attention vector, we compute the average attention mass assigned to relevant tokens for an instance\. This metric can be computed exactly using DTAP heatmaps, if the full attention distribution is known\.

Attention Confidence \(AC\): We compute fraction of instances where attention confidence is high\. This refers to right\-half column sum in DTAP heatmap\. For example: Fig[5](https://arxiv.org/html/2608.06776#A4.F5)a for baseline model we will have AC value11\.13\+1\.47\+0\.01\+0\.04=12\.6511\.13\+1\.47\+0\.01\+0\.04=12\.65\(Table[1](https://arxiv.org/html/2608.06776#S6.T1)column55row11\)\.

Attention Confidence Model Confidence \(ACMC\): We compute fraction of instances where attention confidence is high as well as prediction probability is high\. This refers to right\-half columns and top\-half rows sum in DTAP heatmap\. For example: Fig[5](https://arxiv.org/html/2608.06776#A4.F5)a for baseline model we will have ACMC value11\.13\+1\.47=12\.611\.13\+1\.47=12\.6\(Table[1](https://arxiv.org/html/2608.06776#S6.T1)column55row11\)\.