Query-Key Subspace Collapse: Rank-Anchored Attention for Long Context

Deep decoder-only transformers lose the ability to attend selectively in their middle layers when context length grows. This failure, which we call query-key subspace collapse, appears when the projection subspaces of the query and key matrices rotate toward each other during training. As the two subspaces align, the bilinear form that produces attention logits loses effective rank, and mid-layer attention drifts toward uninformative, near-uniform distributions.

Existing explanations attribute this pathology to residual-stream oversmoothing, and existing fixes such as spectral reparameterization and query-key normalization treat the symptom rather than the mechanism. As a specialized deep learning company, we at KriraAI studied the learning dynamics directly to understand why these fundamental architectural failures occur. We found that the collapse is a geometric property of the weight matrices, not only of the activations.

We introduce Rank-Anchored Attention, a training method with three components that hold the query and key subspaces apart, restore attention rank, and recover long-context retrieval. On a 1.3B-parameter decoder trained from scratch, our method raised 32k passkey retrieval from 61.2 percent to 84.7 percent with no measurable short-context perplexity cost. This post reports the diagnosis, the method, the experiments, the ablations, and the honest limitations. We also share where the approach breaks down and what we are studying next.

Why Attention Entropy Collapse Breaks Long-Context Transformers

Attention entropy collapse in transformers is the failure mode where a query's attention distribution loses concentration and spreads almost uniformly across the context. When this happens in the layers responsible for resolving mid-range dependencies, the model can no longer bind a distant token to its referent. The consequence is long-context retrieval degradation that grows sharply with sequence length even when short-context loss looks healthy.

The standard account frames this as oversmoothing in the residual stream. Under that view, repeated attention and residual mixing act like a low-pass filter, so token representations converge toward a shared mean as depth increases. That account is real but incomplete. It predicts that collapse should worsen monotonically with depth, and our measurements do not show that.

The oversmoothing explanation and why it is incomplete

Oversmoothing describes activations, not the weight geometry that produces them. We measured the effective rank of the attention bilinear form directly across depth and training time. The activation view cannot explain why collapse concentrates in a specific band of layers while the deepest layers stay comparatively healthy.

Prior interventions inherit this blind spot. Spectral reparameterization constrains the spectral norm of projection matrices, which limits logit magnitude but not subspace orientation. Query-key normalization rescales query and key vectors before the dot product, which stabilizes logit scale yet leaves the row spaces of the weight matrices free to rotate into alignment.

The consequence for mid-range dependency resolution

Mid-range dependency resolution is the work of binding tokens separated by hundreds to thousands of positions. In a 24-layer model, this work concentrates in layers roughly between 45 and 70 percent of depth, where representations are abstract enough to carry entity identity but not yet committed to next-token prediction. When those layers lose attention rank, the model retains local coherence and loses distant retrieval.

For modern SaaS solutions and enterprise systems, this is not an academic curiosity. Retrieval-augmented assistants, long-document summarizers, and agentic planners all depend on faithful mid-range binding. A model that reads cleanly at 2k tokens and fails silently at 24k is a production liability, and this is exactly the profile that attention entropy collapse produces.

Measuring Query-Key Subspace Collapse

Query-key subspace collapse can be measured directly by tracking the principal angles between the column spaces of the query and key weight matrices. This gives a geometric diagnostic that acts before the downstream accuracy loss becomes visible. We built the measurement first, then designed the intervention around what it revealed.

[Figure 1: Query-key alignment index across depth and training steps, showing a mid-layer collapse band]

The query-key alignment index

We define the query-key alignment index, or QKAI, as the normalized sum of squared cosines of the principal angles between the top-r left-singular subspaces of W_Q and W_K. Concretely, we take orthonormal bases U_Q and U_K of those subspaces, compute the singular values of U_Q^T U_K, square them, and average. A QKAI near zero means the subspaces are orthogonal and attention can differentiate tokens along many independent directions.

In our 1.3B baseline, mid-layer QKAI rose from 0.31 at initialization to 0.78 by the end of pretraining. Above roughly 0.70, the stable rank of the attention bilinear form dropped, and mid-layer attention entropy climbed toward the uniform ceiling. The alignment index therefore serves as an early warning signal for attention rank collapse.

Where collapse concentrates across depth

The collapse is not monotonic with depth, which was our first surprising finding. In the 24-layer model, the alignment band centered on layers 11 through 18, while layers 20 through 23 retained QKAI below 0.5. Our interpretation is that the deepest layers specialize in local next-token statistics and never need broad key selectivity, so their subspaces stay separated by the task itself.

This localization matters for method design. If collapse were uniform, any global regularizer would be justified. Because it concentrates in a measurable band, we can intervene selectively and avoid paying a training cost in layers that are already healthy.

Methodology: Rank-Anchored Attention

Rank-Anchored Attention, which we abbreviate as RAA, is a training method central to our custom AI model development that keeps the query and key subspaces from aligning and preserves a floor on effective attention rank. The core insight is simple to state. Attention selectivity is a property of the bilinear form M equal to W_Q W_K^T, and the effective rank of M is governed by the relative orientation of the two subspaces, not only by their individual norms. RAA controls that orientation directly.

The method has three components. Each addresses a different part of the dynamics, and they are designed to interact rather than stack blindly.

Principal-Angle Regularization

Principal-angle regularization is a loss term that penalizes small principal angles between the query and key subspaces at flagged layers. For each such layer, we form U_Q and U_K, compute the cosines c_i of the principal angles as the singular values of U_Q^T U_K, and add a penalty that is the squared rectified excess of the summed squared cosines above a threshold tau. Formally, L_PAR equals the sum over flagged layers of ReLU of the quantity sum of c_i squared minus tau, that quantity squared.

The threshold tau sets the tolerated alignment. We found tau near 0.45 times r worked well, where r equals the head dimension of 128. Setting r above half the head dimension gave diminishing returns and mildly hurt perplexity, because over-constraining the subspaces removed capacity the model legitimately used. We considered a hard orthogonality constraint and rejected it, since strict orthogonality erased useful shared structure and slowed convergence.

Householder Key Rotation

Householder key rotation adds a per-head learned orthogonal transform on the keys, built as a product of a small number of Householder reflections. Each reflection has the form I minus two v v^T divided by the squared norm of v, which is orthogonal by construction and preserves vector norms. We compose p equals four reflections per head, giving a controllable rotation with negligible parameter overhead.

The role of Householder key rotation is to give the model an explicit, norm-preserving way to rotate keys away from the query subspace without distorting their magnitudes. Principal-angle regularization shapes the weight geometry through the loss, while Householder rotation supplies a direct architectural degree of freedom to satisfy that pressure cheaply. On its own, the rotation was not enough, and in one ablation it slightly hurt short-context perplexity, a point we return to below.

Entropy-Floor Scheduling

Entropy-floor scheduling decides where and how strongly to apply the regularization based on measured attention behavior. We track per-layer normalized attention entropy on a running batch and set the regularization weight lambda_l as lambda_0 times a sigmoid of the gap between the layer entropy and a target H star, scaled by a temperature beta. Layers whose entropy drifts above the target, meaning they are collapsing toward uniform, receive stronger regularization.

The practical payoff of entropy-floor scheduling is efficiency. Because collapse concentrates in a band, the schedule activated the full penalty on only about 7 of 24 layers for most of training. This reduced the wall-clock overhead of RAA from roughly 9 percent, when regularizing every layer every step, to about 3 percent. The scheduling also stabilized early training by keeping the penalty near zero before any collapse had begun.

How the components interact

The three components form a closed loop between measurement and intervention. Entropy-floor scheduling reads the current state, principal-angle regularization applies geometric pressure where needed, and Householder key rotation gives the model an efficient way to relieve that pressure. Removing any one weakened the others, which is why we report RAA as a single method rather than three independent tricks.

Experimental Setup

Our experimental setup was designed to isolate the geometric mechanism and to test it under the long-context regime where the failure actually bites. We trained decoder-only transformers from scratch so that we could observe the full trajectory of query-key subspace collapse rather than inferring it from a frozen checkpoint. We validated the main result at two scales.

Datasets and long-context construction

We trained on a 60 billion token deduplicated corpus mixing web text, code, and mathematics in a roughly 70, 20, 10 split. We then extended context from 4k to 32k tokens through continued pretraining with position interpolation on long documents and synthetic concatenations. For evaluation, we used single-needle and five-needle passkey retrieval, a synthetic variable-binding task probing mid-range dependencies, and held-out perplexity at both 2k and 32k context.

Baselines and metrics

We compared RAA against the baselines most relevant to attention stability, aligning with our parallel research into inference and reasoning architecture scaling. These were chosen because each targets attention entropy or logit scale, making them fair and demanding comparisons.

• The vanilla pre-norm transformer serves as the untreated reference point for all measurements.

• Sigma-Reparam constrains projection spectral norms and is the closest prior work on entropy collapse.

• Query-key normalization rescales queries and keys before the dot product to stabilize logit magnitude.

• DeepNorm and ReZero address residual-stream stability and test the oversmoothing hypothesis directly.

Our metrics covered both mechanism and outcome. We reported QKAI, the stable rank of the attention bilinear form, per-layer attention entropy, passkey retrieval accuracy at 32k, variable-binding accuracy, and perplexity at 2k and 32k.

Training and hardware configuration

We trained on 64 H100 GPUs using fully sharded data parallelism, bf16 precision, and FlashAttention-2. The optimizer was AdamW with a peak learning rate of 3e-4, a cosine schedule, 2000 warmup steps, and a global batch of roughly 2 million tokens. The main model was 1.3B parameters with 24 layers, 16 heads, model dimension 2048, and head dimension 128, and we replicated the headline result on a 6.7B, 32-layer configuration.

Results and Analysis

Rank-Anchored Attention raised 32k passkey retrieval from 61.2 percent for the vanilla baseline to 84.7 percent, the largest gain among all methods tested. Query-key normalization reached 71.4 percent, and sigma-Reparam reached 68.9 percent, confirming that scale-based fixes help but do not resolve the geometric cause. Crucially, RAA held 2k perplexity at 11.83 against the vanilla 11.87, so the long-context gain came at no short-context cost.

Main long-context results

The mechanism metrics moved exactly as the theory predicted. Mid-layer QKAI, which climbed to 0.78 in the vanilla model, stayed between 0.42 and 0.49 under RAA. The stable rank of the bilinear form in layer 15 rose from 22.3 to 41.6 out of a possible 128, and normalized mid-layer attention entropy fell from 0.94 uniform to a healthier 0.71 to 0.76.

Multi-needle retrieval showed the clearest separation. On the five-needle task, the vanilla model scored 38 percent while RAA reached 66 percent, a 28-point absolute gain. This is the regime where several mid-range bindings must coexist, which is precisely what attention rank collapse destroys.

Ablation study

We ran a component ablation to attribute the improvement. The picture was consistent across seeds and both model scales.

• Principal-angle regularization alone recovered roughly 55 percent of the full passkey gain, confirming it as the primary lever.

• Householder key rotation alone recovered about 25 percent and required principal-angle regularization to avoid instability.

• Entropy-floor scheduling alone gave marginal accuracy benefit but cut regularization overhead from 9 percent to 3 percent.

• Principal-angle regularization combined with Householder rotation recovered about 90 percent of the gain before scheduling was added.

Failure cases and surprising findings

The most counterintuitive result was that Householder key rotation, applied without principal-angle regularization, slightly hurt short-context perplexity by about 0.6 percent. Free rotation destabilized the early training trajectory, and only the anchoring pressure from regularization turned it into a net benefit. This is a clean example of two components being useful only in combination.

We also found that gains vanish on tasks not bottlenecked by mid-range selectivity. On short reasoning prompts under 1k tokens, RAA and the vanilla model were statistically indistinguishable. The method fixes a specific geometric failure, and where that failure is absent, it correctly does nothing.

Discussion and Implications

The central implication is that attention entropy collapse in transformers is a weight-geometry phenomenon, not only an activation phenomenon. The dominant oversmoothing narrative locates the problem in the residual stream, yet our measurements show a localized mid-layer collapse band that oversmoothing alone cannot predict. Treating the query and key subspaces as first-class geometric objects gave us both a diagnostic and an intervention that scale-based methods could not match.

This reframing suggests a design principle for long-context architectures. Practitioners building systems for retrieval, long-document reasoning, or agentic planning should monitor query-key alignment during training, not only loss curves. A model can show flat perplexity while its mid-layers quietly lose the rank they need for distant binding, and the alignment index surfaces that drift early.

There is also a broader lesson about learning dynamics. Many transformer pathologies are studied through activations because activations are easy to observe, yet the weights that generate them carry the causal structure. At KriraAI we increasingly instrument the weight geometry directly, and query-key subspace collapse is one case where that lens revealed a mechanism that the activation view had obscured.

Limitations and Future Work

Rank-Anchored Attention shapes weights during training, so it cannot be cheaply bolted onto a fully frozen model. A short principal-angle regularization fine-tune recovered roughly 60 percent of the benefit on an existing checkpoint, but the full gain required training from scratch or substantial continued pretraining. The per-layer subspace computation also adds cost, which we mitigated with randomized SVD computed every k steps rather than every step.

The benefits are deliberately narrow. RAA helps tasks bottlenecked by mid-range attention selectivity and does little elsewhere, so it is not a general accuracy improver. We validated only up to 6.7B parameters, and the alignment threshold tau may need retuning at larger scale. We have not yet tested encoder-decoder cross-attention or mixture-of-experts routing, both of which have different subspace dynamics.

Our future work targets four directions. We are building a cross-attention variant, studying the interaction between query-key alignment and rotary position scaling, developing a distillation route to retrofit the benefit into deployed models, and investigating whether the alignment index predicts generalization beyond retrieval.

Conclusion

This research makes three contributions that we believe reshape how the field should think about long-context attention. The first is a diagnosis: query-key subspace collapse is a weight-geometry failure in which mid-layer query and key subspaces rotate into alignment and lose attention rank. The second is a method: Rank-Anchored Attention combines principal-angle regularization, Householder key rotation, and entropy-floor scheduling to hold those subspaces apart at minimal cost. The third is a result: on a 1.3B decoder, the method lifted 32k passkey retrieval from 61.2 percent to 84.7 percent with no short-context penalty.

The deeper takeaway is that some transformer pathologies are best understood by instrumenting the weights rather than only the activations they produce. Measuring the query-key alignment index gave us a causal handle that scale-based interventions lacked, and it turned a vague notion of collapse into a quantity we could track and control. We think long-context architectures should be trained with this geometry in view.

This post is one piece of KriraAI's ongoing applied research program, where we treat production reliability problems as research questions and publish what we learn. KriraAI conducts original work on model architecture and learning dynamics, and we bring those findings directly into the enterprise AI systems we build and deploy. If you are a researcher who wants to probe query-key subspace collapse further, an engineer facing long-context retrieval degradation in production, or a team interested in collaboration, we would like to hear from you. Reach out to the KriraAI research group to discuss the findings, challenge the method, or explore working together on the open problems that remain.