Key Stage Attention Sharing in Transformers

Updated 6 July 2026

Key Stage Attention Sharing is a design principle where attention computations are reused or aligned across designated stages to reduce redundancy and computational costs.
It involves sharing various elements such as attention weights, key-value caches, and projection matrices, optimizing performance in tasks from language modeling to image restoration.
Empirical studies demonstrate that stage-constrained sharing can significantly improve efficiency and throughput while maintaining high model accuracy.

“Key Stage Attention Sharing” (Editor’s term) denotes a family of mechanisms in which attention-related computation is not treated as independent at every layer, block, or timestep, but is instead reused, aligned, factorized, or constrained across a designated stage. In the literature, the shared object varies by architecture: some methods share the attention matrix itself across adjacent layers, others share key–value caches, projection matrices, static keys, or stage-level semantic dictionaries, while still others treat attention sharing as supervision transfer or as the alignment of foreground maps in interactive settings. The unifying premise is that attention exhibits structured redundancy or reusable structure across depth, stages, or interaction phases, and that exploiting this structure can reduce computation, memory, or supervision difficulty without proportionate loss of function (Liao et al., 2024, Zhussip et al., 6 Aug 2025, Ren et al., 2024, Zhao et al., 2023, Chen et al., 16 Dec 2025).

1. Conceptual scope and taxonomy

A “stage” is defined differently across subfields. In decoder-centric LLM work it is usually a contiguous span of layers; in hierarchical vision transformers it is a block of layers with fixed token resolution; in image restoration it can be a restoration stage with a shared semantic graph; in diffusion editing it can refer to denoising timesteps whose attention behavior is modulated jointly. The shared entity is correspondingly heterogeneous: attention weights $A$ , cached $K/V$ activations, projection matrices $W_Q,W_K,W_V,W_O$ , static key parameters, or explicit foreground maps (Liao et al., 2024, Chen et al., 2021, Brandon et al., 2024, Zhussip et al., 6 Aug 2025, Hu et al., 2024).

Setting	Shared object	Representative mechanism
Transformer depth	Attention weights, $K/V$ , or Q/K/V/O projections	SA, San, CLA, MASA, LiSA
Vision/restoration stage	Attention maps or key-semantic dictionary	PSViT, SemanIR, NASNet
Supervision/interaction	Weights, values, or foreground maps	KS-DETR, joint guided search, Attentional Push

This diversity matters because “attention sharing” is not a single operator. In some papers, sharing is a hard equality constraint such as $A^{(l)}=A^\star$ over a layer span. In others, it is a structured approximation such as $\hat W_\ell=\sum_{s=1}^S c_{\ell s}D_s$ , or a shared base with low-rank residuals. A plausible implication is that the term is best understood as a design principle rather than a standardized module name (Zhussip et al., 6 Aug 2025, O'Neill et al., 16 Jan 2026).

2. Reusing attention weights across depth

The most direct form of key-stage attention sharing reuses the attention matrix across adjacent or nearby layers. In Shared Attention for LLMs, a contiguous stage $S=\{l_1,\dots,l_m\}$ computes a single attention matrix in the first layer,

$A^\star=\mathrm{softmax}\!\left(\frac{Q^{(l_1)}K^{(l_1)\top}}{\sqrt{d_k}}\right),$

and all later layers in that stage use $A^{(l)}=A^\star$ while retaining their own values $V^{(l)}$ , so that $K/V$ 0 (Liao et al., 2024). The same paper argues that middle and late layers exhibit highly similar attention maps, whereas early and final layers are more distinct; this is the empirical basis for stage selection.

In PSViT, attention sharing is local and stage-constrained: if layer $K/V$ 1 is a sharing layer, it reuses the preceding layer’s attention map $K/V$ 2 and computes

$K/V$ 3

without recomputing $K/V$ 4, $K/V$ 5, or the softmax. Because attention maps change size when token pooling changes resolution, sharing is only allowed within a stage of fixed token count. The architecture search in PSViT treats both token counts and sharing locations as learnable hyper-parameters, and the reported scheme achieves up to 6.6% accuracy improvement in ImageNet classification compared with DeiT (Chen et al., 2021).

A closely related decoder-side formulation appears in “Sharing Attention Weights for Fast Transformer”. There, for a block of $K/V$ 6 decoder layers, the base layer $K/V$ 7 computes $K/V$ 8, and upper layers reuse it as

$K/V$ 9

For encoder–decoder attention the reuse is stronger: $W_Q,W_K,W_V,W_O$ 0 The sharing policy is learned from layer-wise JS divergence, and the method yields an average 1.3X speed-up with almost no decrease in BLEU on top of a cached baseline; with AAN it reaches 1.8X speed-up and is 16 times faster than the baseline with no use of the attention cache (Xiao et al., 2019).

LiSA generalizes direct weight reuse by addressing two observed failure modes: direct sharing without head rearrangement is ineffective, and shallow layers are vulnerable to small deviations in attention weights. LiSA therefore augments cross-layer sharing with tiny feed-forward networks for head alignment and low-rank matrices for layer-wise attention differences. Evaluated on 13 benchmarks, it reduces redundant attention calculations within 53–84% of total layers, achieves a 6X compression of Q and K, and reports maximum throughput improvements of 19.5% for LLaMA3-8B and 32.3% for LLaMA2-7B (Mu et al., 2024).

A second line of work shares not the attention matrix itself but the structures from which attention is formed. Cross-Layer Attention (CLA) extends Multi-Query Attention by sharing key and value heads between adjacent layers. In CLA2, each pair of layers shares one KV cache; in CLA3, groups of three layers share one KV cache. The result is another 2x reduction in KV cache size relative to unmodified MQA while maintaining nearly the same accuracy in models trained from scratch at 1B and 3B scale (Brandon et al., 2024).

MASA moves the sharing locus from activations to parameters. For each projection type, it stacks all layer matrices into a dictionary-learning problem,

$W_Q,W_K,W_V,W_O$ 1

so that the $W_Q,W_K,W_V,W_O$ 2-th layer is approximated by

$W_Q,W_K,W_V,W_O$ 3

The atoms $W_Q,W_K,W_V,W_O$ 4 are shared across all layers, while the coefficients $W_Q,W_K,W_V,W_O$ 5 remain layer-specific. With $W_Q,W_K,W_V,W_O$ 6, the reduction per projection is approximately $W_Q,W_K,W_V,W_O$ 7, and the method reports on-par or better performance than GQA, low-rank baselines, and rigid layer-sharing baselines. The paper further states that separate dictionaries for $W_Q,W_K,W_V,W_O$ 8, $W_Q,W_K,W_V,W_O$ 9, $K/V$ 0, and $K/V$ 1 outperform a common dictionary across all projections (Zhussip et al., 6 Aug 2025).

LRKV defines a continuum between complete KV sharing and fully independent heads by using a shared full-rank KV projection plus low-rank head-specific residuals: $K/V$ 2 Its KV cache per layer becomes

$K/V$ 3

which is lower than standard $K/V$ 4 when $K/V$ 5. At the 2.5B scale, LRKV is reported to outperform standard attention while using roughly half the KV cache, and to reach equivalent model quality with up to 20–25% less training compute in cumulative FLOPs (O'Neill et al., 16 Jan 2026).

Static Key Attention in Vision simplifies the key stage itself by replacing the dynamic key $K/V$ 6 with a learned input-independent parameter per head. In SKA, keys are static within a layer and shared across all tokens and inputs; in CSKA, grouped convolution provides the static key mechanism. The paper reports that static-key modules can match or exceed standard self-attention on several vision tasks and work best as an intermediate stage in hierarchical hybrid architectures, rather than as a uniform replacement of all attention blocks (Hu et al., 2024).

In hierarchical vision models, key-stage attention sharing often exploits the fact that resolution, semantics, and computational cost are stage-dependent. PSViT formalizes a stage as a sequence of layers with fixed token number and channel dimension, and because the attention map size is constant within such a stage, adjacent-layer attention reuse is feasible there but not across token-pooling boundaries (Chen et al., 2021).

SemanIR makes this stage-wise definition explicit by constructing a sparse key-semantic dictionary once per transformer stage. For a patch set $K/V$ 7, it computes self-similarity $K/V$ 8, keeps only the top- $K/V$ 9 neighbors per patch to form $A^{(l)}=A^\star$ 0, and then reuses $A^{(l)}=A^\star$ 1 across all subsequent transformer blocks within the same stage. Attention is restricted to the semantically related components stored in this dictionary, and the paper states that attention calculation thereby achieves linear computational complexity within each window while delivering state-of-the-art results across 6 image restoration tasks (Ren et al., 2024).

NASNet provides an earlier stage-by-stage restoration formulation outside the transformer setting. Its abstract states that the stage-by-stage strategy guides the later stage by incorporating the useful information in previous stage, and that stage-level information is concatenated and fused dynamically by a Neuron Attention module composed of Depthwise Conv and Pointwise Conv. The method is presented as a unified network for rain streak, rain drop, and rain mist deraining, with reported gains on six public large-scale datasets for three rain model tasks (Qin et al., 2019).

KS-DETR uses a different stage-specific sharing mechanism in the last encoder layer. Its triple-attention module contains a plain student attention and two GT-mask-assisted teacher attentions: one produces higher-quality values and shares them with the student, the other produces higher-quality attention weights and shares them with the student. The second and third attentions are removed during inference, so the extra sharing is a training-time device for attention learning rather than a test-time computational shortcut. The reported COCO results show consistent AP gains across several DETR variants (Zhao et al., 2023).

5. Empirical regularities, gains, and failure modes

Across model families, the empirical case for key-stage attention sharing rests on nonuniform redundancy. In PSViT, adjacent-layer normalized correlations of attention maps are reported as 0.89, 0.91, and 0.91 for representative layer pairs. In Shared Attention for LLMs, middle and late layers of pretrained models show pairwise cosine similarities around 0.8–0.9, while early and final layers are outliers. LiSA sharpens this observation by showing that direct head matching across layers is little better than random matching, whereas most-similar head matching substantially lowers JS divergence; this is why head alignment is necessary for high-fidelity cross-layer reuse (Chen et al., 2021, Liao et al., 2024, Mu et al., 2024).

These regularities do not imply that “more sharing is always better.” Several papers document the opposite. Shared Attention reports that late-stage sharing such as layers 27–30 causes small or negligible losses, whereas mid-stage sharing such as 15–18 can severely degrade GSM8K and MMLU. LiSA likewise finds that shallow layers are especially sensitive to small attention perturbations. Static Key Attention reports that fully replacing all attention blocks with static-key variants is not always optimal on large-scale ImageNet models, and that using SKA/CSKA as an intermediate stage works better (Liao et al., 2024, Mu et al., 2024, Hu et al., 2024).

Another recurring pattern is that the granularity of sharing matters. MASA reports that per-projection dictionaries for $A^{(l)}=A^\star$ 2, $A^{(l)}=A^\star$ 3, $A^{(l)}=A^\star$ 4, and $A^{(l)}=A^\star$ 5 are preferable to a common dictionary across all projections. CLA reports that uniform pairwise sharing is better than more exotic non-uniform layer groupings. LRKV reports that moderate low-rank specialization preserves nearly all functional head diversity relative to standard attention, whereas more aggressive KV-sharing mechanisms rely on compensatory query specialization (Zhussip et al., 6 Aug 2025, Brandon et al., 2024, O'Neill et al., 16 Jan 2026).

Diffusion editing introduces a stage-aware variant in which the shared mechanism itself is scheduled across denoising time. SynPS diagnoses two failure modes of existing attention sharing: positional collapse when RoPE dominates, and semantic collapse when RoPE is removed. It defines an editing measurement

$A^{(l)}=A^\star$ 6

and uses it to modulate the RoPE scaling factor $A^{(l)}=A^\star$ 7 across denoising steps. On PIE-Bench Change Pose and a curated 200-prompt Non-Rigid Editing benchmark, this dynamic stage-aware sharing improves faithfulness over FreeFlux and CharaConsist, which the paper attributes to avoiding both over-editing and under-editing (Chen et al., 16 Dec 2025).

6. Broader interpretations beyond standard transformers

Outside mainstream transformer efficiency work, the same general idea appears as guided sharing of attention-related structure. In “Attentional Push,” shared attention is treated as a social process in which scene actors redirect the viewer’s attention. The method combines saliency $A^{(l)}=A^\star$ 8 with Attentional Push maps $A^{(l)}=A^\star$ 9 through

$\hat W_\ell=\sum_{s=1}^S c_{\ell s}D_s$ 0

thereby modeling how gaze direction, body pose, center bias, and dynamic cues jointly shape fixation prediction (Gorji et al., 2016).

In the joint guided search task for human–robot interaction, shared attention is operationalized through the alignment of binary foreground maps. A deictic action is parameterized as $\hat W_\ell=\sum_{s=1}^S c_{\ell s}D_s$ 1, candidate foregrounds are ranked against the gesture center, and prediction quality is measured by normalized mean squared error between the predicted and reported foreground maps. Here, key-stage attention sharing is not a network-layer mechanism but an interaction cycle in which gesture production, gesture interpretation, foreground prediction, and fixation alignment form successive stages of attention coordination (DePalma et al., 2016).

A broader, plausible interpretation extends the term to modular multi-task systems in which attention routes state-dependent computation among shared sub-networks. Attentive Multi-Task Deep Reinforcement Learning uses a task- and state-conditioned attention network to weight multiple sub-networks, aiming to maximize positive transfer and avoid negative transfer across tasks. The paper does not use the phrase “key stage attention sharing,” but it does instantiate attention-mediated sharing at a critical routing stage between shared perception and task-specific behavior (Bram et al., 2019).

Taken together, these works show that key-stage attention sharing is best understood as a recurrent architectural motif: identify a stage where attention-related structure is redundant, stable, or socially meaningful; expose that structure as an explicit shared object; and let later computation either reuse it directly or adapt it with a lightweight correction. The differences across papers concern what is shared, how rigidly it is shared, and which stages are considered safe for sharing, rather than whether stage-wise sharing exists at all.