RACE: RoPE-Aligned Attention Compression Based on Query-Key Activation Energies

Introduction

Autoregressive language-model inference is increasingly constrained by the key-value (KV) cache, especially at long contexts. Each attention layer stores a key and value vector for every preceding token, so cache memory grows linearly with context length and with the dimensionality of the cached representations. Reducing the key and value head dimensions offers the possibility of decreasing both KV-cache memory and native attention computation. Realizing these savings is more difficult than merely storing a low-dimensional encoding. Several head-dimension compression methods project keys and values into compact representations but reconstruct full-dimensional keys during inference to effectively apply rotary positional embeddings (RoPE) [10, 3]. Such methods reduce cache storage but introduce decompression and additional arithmetic into the attention path during inference. To obtain the full benefit of head-dimension compression, the attention layer itself must operate directly in the reduced dimensions.

This raises a more fundamental question: when can a pretrained RoPE-based attention layer be converted into a natively lower-dimensional attention layer without reconstructing the discarded coordinates? Since attention logits depend on the post-RoPE product $QK^{\mathsf{T}}$, we view head-dimension reduction as an approximate matrix multiplication (AMM) problem along the shared feature dimension. However, preserving this interaction in the post-RoPE space is not sufficient for native inference. It must be realizable before RoPE. This imposes an additional structural constraint that the reduction must commute with the pretrained rotary transformation. We characterize this condition exactly for coordinate-selection reductions.

These observations define an admissible family of native Q/K reductions: queries and keys can share a reduced feature space, and the reduction must respect the pretrained RoPE structure. The remaining question is which admissible dimensions to retain. We introduce RACE, which ranks complete rotary pairs using joint activation energy for each RoPE pair $m$, $E_m = \lVert Q_m\rVert_F^2 + \lVert K_m\rVert_F^2$, and retains the highest-energy pairs. This criterion minimizes a separable, calibration-dependent upper bound on post-RoPE query-key product approximation error. The resulting Q/K reduction, together with a calibrated value subspace, is folded into the attention projections, enabling RoPE, caching, and attention to operate entirely in the reduced dimensions without runtime reconstruction or specialized kernels. Summarizing our contributions:

• Constrained-AMM formulation. We formulate native attention compression as preservation of the post-RoPE query-key product subject to pre-RoPE realizability constraint.

• Structural Characterization. We prove that a coordinate-selection reduction can commute with the pretrained RoPE if and only if it retains complete rotary pairs, characterizing the admissible family of coordinate reductions.

• Forward-only pair selection within the family. RACE ranks RoPE pairs using joint query-key activation energy, minimizing a separable calibration-dependent upper bound.

• Empirical validation. Across multiple model families and compression levels, RACE remains competitive with reconstruction-based methods while reducing KV-cache memory, attention cost, parameters, TTFT, and runtime overhead without specialized kernels.

Related Work

Compression along other KV-cache axes. KV-cache memory can be reduced along the token, precision, or head-dimension axes as shown in Figure 8. Token-level methods reduce the number of cached positions ($S$ in Eq. (7)), through eviction, retention, or compaction ([21 Xiao]彩票论坛 et al., 2024; Li et al., 2024; Cai et al., 2024; Fu et al., 2025; Shi et al., 2025). Quantization methods instead reduce the number of bits used for each cached element, ($b$ in Eq. (7)), [9, 6, 15]. Our method preserves the retained token positions and reduces the number of elements in every key and value vector. It is therefore complementary to both token reduction and quantization; we evaluate representative combinations in § 4.3.

Head-dimension KV-cache compression. Head-dimension methods reduce the width of the cached representations ($d_k$ and $d_v$ in Eq. (7)) for keys and values respectively. EigenAttention [10] derives low-rank subspaces from calibrated activations, whereas PALU [3] factorizes attention projection weights. The dense low-rank compression subspaces in these and similar methods are not generally invariant under the pretrained RoPE transformation and so these methods recover full-dimensional keys before applying RoPE. This can be compute-intense and so, methods like PALU [3] mitigate this by using a custom kernel that does not materialize the full-dimensional reconstructed key tensor. However the inference FLOPs remains dominated by the dense reconstruction basis.

Thin Keys [20] avoids key reconstruction by reducing the shared Q/K dimension, but retains all dimensions of Values. A concurrent preprint, RAP [18], performs RoPE-aligned structured parameter pruning using empirical-Fisher scores computed from gradients of $W_K$ and arrives at a similar RoPE-aligned native architecture from a different starting point to RACE. While RAP approaches the problem as structured parameter pruning, RACE begins from the problem of constrained-AMM as shown in § 3.1. The methods can therefore produce similar reduced-dimensional architectures, but differ in their compression objective and calibration signal. Additionally, our calibration is forward-only and does not depend on a task loss or the batch-sensitivity of Fisher estimates [13]. Figure 1 compares our solution to previous techniques in terms of accuracy and FLOPs across a range of compression levels.

Methods including MatryoshkaKV [8], STAR-KV [1], Think [19], and OjaKV [22] additionally study heterogeneous ranks or input-dependent compression. These methods primarily address where or when to allocate compression across layers, heads, or inputs. Our work instead addresses how to construct a fixed, RoPE-compatible reduced representation that operates natively without runtime selection, adaptation, or full-dimensional reconstruction.

Method

Problem Formulation. Consider an attention head with pre-RoPE activations $Q, K \in\mathbb{R}^{S \times d_k}$ and $V \in\mathbb{R}^{S \times d_v}$. Given target dimensions $r_k < d_k$ and $r_v < d_v$, we seek reduced-dimensional representations

with $P_Q, P_K \in\mathbb{R}^{d_k \times r_k}$ and $P_V \in\mathbb{R}^{d_v \times r_v}$, such that the resulting layer stores only $K^\prime$ and $V^\prime$ in the KV cache and performs attention directly in dimensions $r_k$ and $r_v$, without reconstructing the discarded coordinates at inference.

For native reduced-dimensional execution, the Q/K reductions must ultimately be applied to the pre-RoPE activations so that they can be folded into the query and key projection weights. However, the quantity that determines attention logits is the post-RoPE query-key product. We therefore first characterize the low-dimensional Q/K representation that preserves this interaction, and then determine when it can be realized equivalently before RoPE. Finally we introduce RACE, a criterion that satisfies these conditions and effectively convert a pretrained attention to a natively reduced-dimensional attention.

Native Attention Compression as Constrained AMM

Let $\widetilde{Q} = \operatorname{RoPE}(Q)$, $\widetilde{K} = \operatorname{RoPE}(K)$. Since attention logits depend on the post-RoPE product $\widetilde{Q}\widetilde{K}^{\mathsf{T}}$, we formulate Q/K head-dimension reduction as an approximate matrix multiplication (AMM) problem. We are seeking a column-orthonormal reduction $P \in\mathbb{R}^{d_k \times r_k}$ such that,

This is consistent with classical matrix-sketching and AMM results, where a common lower-dimensional map is applied along the shared dimension of the two operands [16, 4]. Detailed explanation is given in A.5.

This view motivates two design principles. First, queries and keys should be represented in a shared low-dimensional subspace, since their interaction is evaluated along the same feature dimension. We therefore seek a common reduction $P_{QK}$ for both operands. Second, when this subspace is chosen from calibration data, its selection criterion should reflect both matrices whose product is being approximated. Because the target interaction depends jointly on $\widetilde{Q}$ and $\widetilde{K}$, we construct $P_{QK}$ from statistics of both query and key activations rather than either operand in isolation as in the concurrent work, RAP [18] or ThinKeys [20]. This is validated empirically in Figure 12.

These observations specify how the Q/K interaction should be reduced, but not yet whether the reduction can be executed natively before RoPE to get the desired systems benefits. Applying $P_{QK}$ only after RoPE would still require full-dimensional query and key projections and full-dimensional rotary encoding, reducing the KV cache width without fully reducing the attention computation as in reconstruction-based methods. Native execution therefore requires the shared reduction to be realizable before RoPE. Section 3.1 establishes how the query-key interaction should be reduced, but not whether the shared reduction can be realized before RoPE. Section 3.2 characterizes the additional structural condition under which this is possible.

RoPE-Compatible Query–Key Reduction

Section 3.1 motivates representing queries and keys through a shared low-dimensional map. To apply this reduction before RoPE, it must preserve the corresponding blocks of the pretrained rotary transformation as detailed in A.3.

Let $R_t = \operatorname{diag}(R_{1,t}, \ldots, R_{M,t}) \in\mathbb{R}^{d_k \times d_k}$ be the pretrained RoPE operator at position $t$, where each $R_{m,t} \in\mathbb{R}^{2 \times2}$ acts on one rotary pair. For a retained pair set $\mathcal{S} \subseteq\{1, \ldots, M\}$, define $R_t^{(\mathcal{S})} = \operatorname{diag}(R_{m,t}: m \in\mathcal{S}) \in\mathbb{R}^{r_k \times r_k}$. Thus, $R_t^{(\mathcal{S})}$ is the restriction of the original pretrained operator to the retained pairs. We say that a reduction $P_{QK} \in\mathbb{R}^{d_k \times r_k}$ is RoPE-compatible if

This condition ensures that applying the reduction before RoPE is equivalent to applying full-dimensional RoPE and then restricting to the retained coordinates. Consequently, query/key projection, rotary encoding, and attention can all be performed directly in dimension $r_k$. A simple admissible family of $P_{QK}$ is obtained by retaining complete RoPE pairs. The following theorem shows that this family exactly characterizes RoPE-compatible coordinate-selection reductions.

Algorithm 1 RACE: Post-Training RoPE-Aligned Attention Compression

Require: Pretrained model $\mathcal{M}$, calibration set $\mathcal{D}_{\mathrm{cal}}$, target dimensions $r_k, r_v$

1. for each attention layer do

2. Collect post-RoPE $Q$, $K$, $V$ activations on $\mathcal{D}_{\mathrm{cal}}$

3. For RoPE pairs, compute $E_m = \lVert Q_m\rVert_F^2 + \lVert K_m\rVert_F^2$

4. Construct $P_{QK}$ that selects the top-$r_k/2$ pairs based on $E_m$

5. Fold in $P_{QK}$. $W_Q^{\mathrm{fold}} = W_QP_{QK}$; $W_K^{\mathrm{fold}} = W_KP_{QK}$

6. Compute the top-$r_v$ right singular vectors of $V$ to get $P_V$

7. Fold in $P_V$. $W_V^{\mathrm{fold}} = W_VP_V$; $W_O^{\mathrm{fold}} = P_V^{\mathsf{T}}W_O$

8. end for

9. Optional: recover performance of $\mathcal{M}'$ from $\mathcal{M}$

10. return Native attention-compressed model $\mathcal{M}'$

The proof is provided in A.6. The remaining question is which admissible pairs should be retained. Section 3.3 introduces a calibration criterion for ranking the admissible pairs.

Calibration-Guided Selection Within the RoPE-Compatible Family

Section 3.2 characterizes a family of structurally valid native reductions by restricting selection to complete RoPE pairs. What remains is no longer a compatibility question, but a ranking question over this admissible pairs to determine what pairs should be retained. Because the approximation target in Section 3.1 depends on both queries and keys, we rank candidate pairs using a joint calibration statistic over the two operands.

Let $\mathcal{P}_{\mathrm{rope}} = \{(i_m,j_m)\}_{m=1}^{M}$ denote the RoPE pairs of an attention head, and let $Q_m = Q_{:,\{i_m,j_m\}}, K_m = K_{:,\{i_m,j_m\}}$. We assign pair $m$ the joint query-key energy

For target dimension $r_k$, we retain the $r_k/2$ pairs with the largest energies, $\mathcal{S} = \operatorname{TopK}_{m\in[M]}(E_m, r_k/2)$ and construct $P_{QK} = [e_i]_{i\in I_{QK}}$ where $I_{QK} = \bigcup_{m\in\mathcal{S}}\{i_m,j_m\}$.

The following result provides a calibration-dependent justification for this selection rule.

Value compression. Values are not subject to RoPE, so we follow the standard low-rank activation projection used in prior head-dimension compression methods [10]. Given calibration activations $V \in\mathbb{R}^{D_{\mathrm{cal}}\times d_v}$, we retain the top-$r_v$ right singular vectors $P_V \in\mathbb{R}^{d_v\times r_v}$ and compute $V' = VP_V$.

Folding into Native Reduced-Dimensional Attention

Having obtained the RoPE-compatible query-key selector $P_{QK}$ and the value subspace $P_{V}$, we fold both transformations directly into the pretrained attention projections. For queries and keys,

$Q'$ and $K'$ are $r_{k}$-dimensional queries and keys to which the restricted RoPE operator from § 3.2 is applied. Hence, neither the discarded query/key coordinates nor a full-dimensional RoPE representation is materialized during inference.

For values, we similarly fold the calibrated subspace into the value and output projections - $W_{V}^{\mathrm{fold}} = W_{V}P_{V}$, $W_{O}^{\mathrm{fold}} = P_{V}^{T}W_{O}$ and $V' = XW_{V}^{\mathrm{fold}}$, $Y' = (A'V')W_{O}^{\mathrm{fold}}$ where $A'$ is computed entirely in the reduced dimension $r_{k}$ using the retained RoPE pairs.

This folding converts the original pretrained attention layer into a native architecture with query/key dimension $r_{k}$ and value dimension $r_{v}$. The resulting model stores only reduced-dimensional keys and values and performs RoPE, query-key comparison, value aggregation, and output projection without reconstructing the discarded dimensions. Consequently, the reduction decreases KV-cache storage, attention arithmetic, and attention projection parameters simultaneously.

Recovery. For aggressive compression, we recommend a short one-time post-fold recovery stage using KL distillation from the uncompressed model (Figure 10). We update the reduced attention parameters $W_{Q}^{\mathrm{fold}}$, $W_{K}^{\mathrm{fold}}$, $W_{V}^{\mathrm{fold}}$, $W_{O}^{\mathrm{fold}}$ while keeping the rest of the model frozen. Recovery is performed directly on the natively compressed checkpoint using 1M FineWeb-Edu tokens. Its effect is modest at upto 75% KV retention but substantial at 50%, where it restores much of the performance. Therefore, while moderate compression requires little or no adaptation; for 50% retention, we recommend a short one-time recovery stage for RACE and RAP.

Experiments and Results

Implementation Details. We implement RACE in PyTorch and Hugging Face Transformers. Calibration uses 32,768 FineWeb-Edu tokens, with statistics and SVDs computed in FP32 during calibration and inference is performed in BF16. We respect each architecture’s physical rotary-pair layout and retain the original frequency indices and RoPE scaling rules; selected pairs are not re-indexed as consecutive frequencies. Compressed checkpoints resize the attention projections and cache dimensions while maintaining the $\frac{1}{\sqrt{d_k}}$ scaling, and run using standard layers and PyTorch attention operations. Thus, RACE requires no specialized CUDA/Triton kernel. While all the results in § 4 were reported with SDPA attention backend, Table 2 shows systems results with eager and flash attention backends. Although different kernels optimize for wall-clock time, the number of FLOPs remain the same. Although optional, all the results of RACE and the concurrent work, RAP involve recovery with KL-distillation over approximately 1M FineWeb-Edu tokens as detailed in A.7.

	Total Inference TFLOPs	TTFT (ms)	Prefill T.put (tok/sec)	Decode T.put (tok/sec)	Total T.put (tok/sec)	GPU Memory (GB)
Eager	3.973	41.051	13937	45.186	224.44	5.884
SDPA	3.973	28.854	17755.36	56.864	282.904	5.882
Flash Attention 3	3.973	34.170	15054	46.23	230.127	5.882

Quality at Matched KV-Cache Budgets

We evaluate Qwen2.5-7B, Llama-3.2-3B-Instruct, and Granite-4.1-3B against RoPE-compatible implementations of PALU-GLRD ([3]), EigenAttention([10]), Thin Keys([20]) and concurrent work, RAP ([18]). The full results are shown in Table 3 in the Appendix. Figure 3 shows the same on Qwen2.5-7B and Llama-3.2-3B-Instruct for RACE, the concurrent work and structured pruning baseline: RAP, and the reconstruction-based baseline: PALU. We finetune reduced attention weights of both RAP and RACE for 1M tokens of FineWeb-Edu using KL distillation from the respective uncompressed model for fair comparison. We report WikiText perplexity, mean accuracy across 15 zero-shot tasks (Table 3), and LongBench score average over 16 English tasks (Table 4). At 75% KV-cache retention, RACE preserves 98.5%, 96.5%, and 97.5% of the uncompressed zero-shot accuracy on Qwen, Granite, and Llama, respectively. Across compression levels, on Llama and Qwen models (Figure 3), RACE closely matches PALU on zero-shot accuracy and LongBench without runtime reconstruction. RACE consistently outperforms RAP. RACE also achieves the lowest perplexity at 50% and 75% retention.

Method	KV Cache	QK Dim.	V Dim.	Wiki (ppl)	Completion lambada	Reasoning hella	Reasoning piqa	Reasoning wino	Knowledge QA arc_e	Knowledge QA sciq	Knowledge QA obqa	Knowledge QA arc_c	Understanding boolq	Understanding race	Truth fulqa	Specialized Knowledge college_CS	Specialized Knowledge college_Phy	Specialized Knowledge Jaw	Specialized Knowledge HS-psych	Avg
Base	100%	128	128	10.976	0.706	0.795	0.784	0.707	0.759	0.960	0.460	0.495	0.845	0.418	0.524	0.700	0.480	0.529	0.897	0.671
	87.50%	112	112	11.438	0.696	0.791	0.771	0.703	0.772	0.954	0.456	0.504	0.839	0.421	0.457	0.710	0.500	0.511	0.881	0.664
	87.50%	96	128	11.080	0.693	0.797	0.780	0.711	0.782	0.956	0.462	0.514	0.844	0.414	0.516	0.700	0.490	0.512	0.883	0.670
	75%	96	96	12.290	0.687	0.782	0.756	0.703	0.777	0.953	0.460	0.508	0.838	0.414	0.411	0.660	0.451	0.486	0.868	0.650
	75%	64	128	11.325	0.695	0.793	0.774	0.706	0.790	0.964	0.470	0.525	0.853	0.410	0.532	0.660	0.471	0.521	0.862	0.668
PALU	50%	64	64	17.240	0.639	0.746	0.662	0.683	0.703	0.941	0.392	0.452	0.789	0.393	0.430	0.600	0.402	0.408	0.828	0.605
	87.50%	112	112	10.178	0.702	0.797	0.783	0.723	0.790	0.961	0.458	0.538	0.838	0.405	0.500	0.630	0.471	0.507	0.899	0.667
	87.50%	96	128	9.602	0.713	0.803	0.792	0.729	0.757	0.951	0.460	0.506	0.843	0.415	0.561	0.690	0.490	0.533	0.892	0.676
	75%	96	96	11.043	0.666	0.793	0.768	0.713	0.757	0.953	0.454	0.497	0.811	0.398	0.500	0.590	0.461	0.480	0.862	0.647
	75%	64	128	10.080	0.715	0.804	0.783	0.702	0.722	0.949	0.462	0.495	0.829	0.398	0.545	0.650	0.431	0.497	0.884	0.658
Eigen Attention	50%	64	64	21.050	0.376	0.747	0.647	0.619	0.617	0.897	0.410	0.398	0.636	0.369	0.429	0.330	0.275	0.374	0.659	0.519
Thin Keys Full Values	87.50%	96	128	30.547	0.294	0.733	0.634	0.547	0.669	0.847	0.330	0.416	0.650	0.321	0.390	0.210	0.314	0.275	0.428	0.470
	75%	64	128	50.975	0.284	0.737	0.611	0.544	0.687	0.834	0.376	0.436	0.611	0.298	0.419	0.250	0.314	0.250	0.363	0.467
RAP	87.50%	112	112	10.743	0.691	0.794	0.782	0.708	0.739	0.961	0.450	0.466	0.843	0.393	0.517	0.660	0.441	0.500	0.884	0.655
	87.50%	96	128	10.594	0.705	0.794	0.784	0.707	0.753	0.963	0.432	0.482	0.854	0.407	0.528	0.680	0.402	0.507	0.888
Granite 4.0 8B	Base	100%	64	64	11.578	0.662	0.794	0.746	0.695	0.790	0.960	0.434	0.504	0.751	0.366	0.476	0.570	0.461	0.414	0.813	0.629
Granite 4.0 8B	PALU	87.50%	48	64	11.775	0.671	0.796	0.736	0.693	0.780	0.958	0.426	0.502	0.736	0.380	0.466	0.530	0.490	0.419	0.806	0.626
Granite 4.0 8B	75%	48	48	13.615	0.626	0.787	0.738	0.679	0.736	0.949	0.436	0.500	0.626	0.385	0.471	0.520	0.412	0.380	0.772	0.601
Granite 4.0 8B	50%	32	64	11.918	0.675	0.799	0.731	0.693	0.776	0.957	0.426	0.500	0.748	0.376	0.468	0.540	0.490	0.422	0.818	0.628
Granite 4.0 8B	50%	32	32	17.483	0.539	0.760	0.646	0.628	0.652	0.914	0.412	0.445	0.708	0.362	0.456	0.440	0.333	0.332	0.642	0.551
Eigen Attention	87.50%	48	64	11.775	0.671	0.796	0.736	0.693	0.779	0.958	0.426	0.502	0.736	0.379	0.466	0.530	0.490	0.419	0.806	0.626
	75%	48	48	13.481	0.577	0.785	0.723	0.657	0.757	0.952	0.416	0.513	0.729	0.376	0.440	0.560	0.421	0.360	0.761	0.601
	50%	32	32	23.361	0.162	0.744	0.545	0.539	0.556	0.765	0.379	0.357	0.388	0.301	0.429	0.330	0.275	0.157	0.419	0.465
Thin Keys Full Values	87.50%	48	64	26.024	0.221	0.719	0.586	0.517	0.643	0.813	0.356	0.357	0.583	0.294	0.395	0.230	0.235	0.257	0.215	0.428
RAP	87.50%	48	64	11.928	0.642	0.791	0.732	0.678	0.784	0.942	0.430	0.478	0.777	0.382	0.469	0.530	0.490	0.396	0.782	0.620
RACE (Ours)	87.50%	48	64	12.092	0.573	0.782	0.712	0.662	0.755	0.948	0.424	0.453	0.759	0.373	0.467	0.340	0.373	0.357	0.672	0.577

	KV Cache Retained	Narrative QA	Qasper	Multifield QA	Hotpot QA	Wikimqa	Musique	gov_report	qmsum	multi_news	trec	triviaqa	samsum	passage count	passage retrieval en	lcc	repobench	Average
Llama3.2 3B Instruct.	Uncompressed 100%	21.25	39.89	50.06	47.31	36.75	19.15	33.91	23.5	26.64	71	88.41	42.62	8.5	75	52.01	52.5	43.03
	87.5%	24.81	33.42	49.61	47.52	36.52	17.9	33.95	23	26.98	73	90.47	42.91	4	63.5	51.26	53.9	42.05
RACE	75%	22.2	30.79	45.3	40.44	35.29	15.96	32.88	22.73	25.77	71	87.78	41.16	3.5	58.5	44.34	44.45	38.88
	50%	11.45	15.08	25.46	22.86	20.01	15.52	21.44	20.5	20.97	59.5	78.45	35.32	1.18	5.75	34.83	35.91	26.51
RAP	87.5%	23.49	32.55	46.14	41.67	36.7	14.63	30.91	22.6	26.31	68.5	87.58	40.96	4.5	45	49.3	52.24	38.94
	75%	19.49	27.26	41.91	30.35	31.14	13.96	25.55	21.63	25.39	64	84.46	39.81	4	23	47.55	49.75	34.33
	50%	7.79	6.42	19	6.16	6.81	8.77	9.2	7.94	5.48	28	40.2	11.99	4	5	34.8	32.67	14.64
PALU	87.5%	21.9	31.34	49.22	45.26	36.56	17.96	31.82	23.32	25.58	67	89.59	30.3	2.79	74.5	46.9	50.69	40.30
	75%	21.91	30.94	46.85	40.73	32.06	17.78	30.95	22.63	24.6	66	88.7	28.8	2.21	56.5	43.62	38.81	37.07
	50%	20.47	30.26	37.77	33.88	29.15	14.18	24.79	22.24	21.8	48	79.63	25.5	1	9	25.74	34.71	28.63
Qwen2.5 7B	Uncompressed 100%	17.56	14.86	27.7	29.5	15.51	11.51	30.96	23.07	26.75	72.5	89.31	45.83	3.5	24.5	63.07	51.58	34.23
	87.5%	17.98	20.87	41.23	41.75	31.3	18.85	28	22.99	26.7	72.5	90.25	44.98	2.5	19	53.35	44.22	36.03
RACE	75%	17.11	18.33	40.39	41.08	30.36	18.76	25.78	22.27	26.25	70.5	87.13	44.07	3	8.5	40.51	46.8	33.80
	50%	4.9	14.27	22.66	19.21	7.23	3.84	18.37	21.1	20.89	69	84.26	39.01	2.8	8.5	37.1	38.07	25.70
RAP	87.5%	13.1	21.46	27.31	10.36	18.43	14.52	26.22	21.15	25.7	72.5	89.56	44.53	4	8	44.44	50.62	30.74
	75%	5.78	21.38	26.8	13.88	11.9	5.26	23.31	23.32	23.75	68.5	86.99	44.21	3	4.5	44.79	42.49	28.12
	50%	5.16	9.4	15.36	7.57	8.48	4.09	15.67	21.49	18.78	52	73.1	33.92	1.17	2.8	27.8	42.13	21.18
PALU	87.5%	7.94	12.25	29.8	36.21	26.83	18.46	29.05	22.63	26.35	72.5	89.49	34.25	2	19	49.72	40.75	32.33
	75%	8.15	16.48	28.76	36.94	28.71	18.8	27.25	22.58	25.16	70.5	89.93	34.45	4	10	41.39	33.24	31.02
	50%	4.15	8.95	23.03	10.8	9.37	5.83	21.21	21.25	22.3	69.5	83.81	33.82	1.5	8.5	31.17	31.17	24.15

Measured Systems Benefits

Reducing KV-cache storage does not automatically reduce inference cost. As discussed in § 1, this requires native reduced attention. Such native reduced attention methods like RACE eliminates reconstruction at the architectural level, exposing standard reduced-dimensional attention operations that can be supported by general-purpose backends. Figure 4 shows that these measured gains become larger with prompt length as attention accounts for a greater fraction of inference cost. At 256K tokens, 50% retention achieves $1.76\times$, $1.90\times$, and $1.79\times$ prefill, decoding, and overall throughput speedups, respectively. It also reduces TTFT by 43.22%, allocated GPU memory by 28.06%, and total inference FLOPs by 47.75% relative to the uncompressed model. The end-to-end memory reduction is smaller than the 50% KV-cache reduction because total allocated memory also includes model parameters, activations, and other components unaffected by the proposed method. Figure 11 shows the memory reduction with compression on total model parameters, attention parameters and KV cache. While parameter memory is unaffected by sequence length, KV Cache memory grows with it.

Compatibility with Token-Dimension Compression and Quantization

RACE reduces the KV-cache feature dimension or head dimension and can therefore be combined with methods that compress its token dimension or numerical precision. We apply 75% RACE retention before H2O and SnapKV token compression and KIVI quantization, each with a 50% cache footprint, yielding 37.5% combined retention. Figure 5 shows that applying RACE to an already compressed model causes greater degradation than applying it to the uncompressed model, indicating that the effects of the two compression stages are not independent. Nevertheless, the combinations provide an additional 25% reduction in cache size and feature-dependent attention FLOPs as well as attention parameters. The degradation is smaller with KIVI, particularly in perplexity, suggesting that RACE is especially complementary to quantization. Overall, RACE can augment existing KV-cache methods with further memory, bandwidth, and, for token compression, decoding-FLOP reductions.

Effect of Calibration Data

Our selection criterion explained in § 3.3 is derived from a calibration-dependent upper bound on the query-key approximation error, so the retained subspace is expected to reflect the calibration distribution. Figure 6 evaluates this dependence by varying the calibration dataset and reporting performance relative to FineWeb-Edu calibration. Note that for WikiText, the perplexity change is sign-inverted so that positive values consistently denote improvement. Using the same dataset as the downstream task generally gives the strongest results - upto +5.4% on the evaluated datasets with respect to calibration on FineWeb-Edu, and calibration on a related task family may also improve performance as observed in evaluation on arc-c with calibration on SciQ. However, the differences are modest (< 6%) for most tasks. Therefore, *FineWeb-Edu remains a strong generic calibration source, especially when in-distribution task-specific samples are unavailable or the downstream task family is unknown.

Ablation: Selection Criteria

The shared Q/K reduction and RoPE-compatibility constraint define a family of native coordinate reductions; the remaining options are on how to rank candidate rotary pairs. We first compare RACE’s joint Q/K activation energy with otherwise identical Q-only and K-only selectors. As shown in Figure 7(b), K-only selection degrades perplexity more strongly, while joint Q/K selection performs best overall. A similar result was observed for Q-only (Table 6). This is consistent with preservation of the query-key interaction: the average normalized error $\lVert QK^{\mathsf{T}}-Q_{\mathcal{S}}K_{\mathcal{S}}^{\mathsf{T}}\rVert_{\mathsf{F}}/\lVert QK^{\mathsf{T}}\rVert_{\mathsf{F}}$ is 0.1159, 0.1120, and 0.1051 for Q-only, K-only, and joint Q/K selection, respectively.

KV Cache Retained	Selection Criteria	wikitext	piqa	wino	arc_c	sciq	arc_easy	boolq	race	mmlu_cs	mmlu_phy	Avg
87.5%	Joint QK Energy (RACE)	13.88	0.74	0.67	0.46	0.93	0.70	0.77	0.42	0.48	0.35	0.61
87.5%	Magnitude	14.58	0.74	0.64	0.44	0.89	0.68	0.71	0.41	0.37	0.25	0.57
87.5%	Random	14.95	0.71	0.63	0.45	0.89	0.67	0.65	0.39	0.35	0.25	0.55
87.5%	Gradient Fisher	14.58	0.76	0.66	0.44	0.87	0.69	0.74	0.41	0.38	0.25	0.58
87.5%	Attention Score	14.13	0.74	0.66	0.44	0.92	0.69	0.75	0.40	0.40	0.27	0.59
87.5%	Energy Q only	14.02	0.75	0.66	0.45	0.93	0.69	0.75	0.42	0.40	0.32	0.60
87.5%	Energy K only	14.03	0.76	0.67	0.45	0.93	0.67	0.75	0.43	0.45	0.30	0.60
75%	Joint QK Energy (RACE)	14.95	0.76	0.66	0.43	0.93	0.70	0.76	0.42	0.48	0.33	0.61
75%	Magnitude	15.64	0.73	0.61	0.43	0.88	0.65	0.67	0.39	0.34	0.24	0.55
75%	Random	17.74	0.72	0.56	0.39	0.84	0.62	0.60	0.36	0.25	0.21	0.50
75%	Gradient Fisher	14.76	0.75	0.60	0.42	0.86	0.67	0.73	0.39	0.40	0.27	0.57
75%	Attention Score	14.87	0.74	0.64	0.42	0.91	0.67	0.73	0.40	0.42	0.32	0.58
75%	Energy Q only	16.78	0.74	0.66	0.42	0.91	0.68	0.75	0.41	0.43	0.33	0.59
75%	Energy K only	16.17	0.75	0.65	0.44	0.92	0.68	0.74	0.42	0.43	0.33	0.60
50%	Joint QK Energy (RACE)	16.02	0.69	0.59	0.38	0.89	0.62	0.69	0.37	0.34	0.33	0.54
50%	Magnitude	18.73	0.71	0.54	0.32	0.77	0.49	0.50	0.31	0.23	0.22	0.45
50%	Random	26.09	0.62	0.50	0.30	0.65	0.44	0.47	0.27	0.26	0.22	0.41
50%	Gradient Fisher	17.80	0.72	0.55	0.35	0.80	0.53	0.44	0.33	0.30	0.27	0.48
50%	Attention Score	17.86	0.70	0.59	0.33	0.73	0.49	0.50	0.35	0.33	0.31	0.48
50%	Energy Q only	21.01	0.72	0.57	0.35	0.86	0.55	0.66	0.37	0.30	0.33	0.52
50%	Energy K only	20.32	0.71	0.57	0.35	0.87	0.54	0.66	0.38	0.30	0.32	0.52

We further compare RACE with random, magnitude, attention-product, and gradient-Fisher scoring. Random selection is data-oblivious; energy(RACE), magnitude and attention-product use joint Q/K forward statistics; and gradient-Fisher uses joint Q/K gradients from backward-pass on the calibration data. RACE achieves the lowest perplexity across the evaluated retention ratios among these selectors. Table 7(a) further compares calibration cost of methods that uses RoPE-compatible shared joint Q/K statistics for selection. Attention-product scoring explicitly incurs $O(S^{2})$ computation and gradient-Fisher requires backpropagation and a task loss. Accordingly, these methods incur substantially higher calibration time and peak memory. Joint QK energy selection therefore offers a favorable trade-off between performance and calibration cost.

Towards Models with Fewer and Wider Heads

If post-training attention compression is anticipated at deployment time, larger head dimensions may be a useful architectural choice during pretraining. We study three matched $\sim245$ M-parameter GPT-NeoX-style models pretrained from scratch on FineWeb-Edu for 2.5B tokens: wide ($4 \times256$), standard ($8 \times128$), and narrow ($16 \times64$). All three reach comparable uncompressed WikiText vali-dation perplexity. Figure 7(c) shows that wider-head models retain perplexity more effectively under aggressive RACE compression. At each retention ratio, the models have approximately matched aggregate retained attention dimension, KV-cache memory, attention FLOPs, and compressed attention parameter count; the primary difference is how this capacity is partitioned across heads. Wider heads expose more RoPE pairs from which RACE can select, providing greater flexibility in structured dimensional reduction. Thus, when post-training attention compression is a foreseeable deployment requirement, allocating attention width across fewer, wider heads during pretraining may yield a more favorable quality-efficiency trade-off.

Conclusion

In this work, we study when and how pretrained RoPE attention can be converted into a natively lower-dimensional form that reduces KV-cache storage, attention computation, and parameters without runtime reconstruction. Viewing Q/K compression as constrained approximate matrix multiplication motivates a shared reduction derived jointly from queries and keys. This reduction has to be RoPE-compatible to enable reconstruction-free compression. We further show that, for coordinate-selection reductions family, RoPE compatibility holds if and only if complete rotary pairs are retained. RACE selects within this admissible family using joint Q/K activation energy, a forward-only criterion that minimizes a separable calibration-dependent upper bound on post-RoPE product error. The resulting reductions are folded directly into the projection weights, enabling attention to execute entirely in the compressed dimensions without specialized kernels. Across multiple model families and compression levels, RACE outperforms evaluated reconstruction-free baselines and remains competitive with reconstruction-based methods while reducing KV-cache memory, attention cost, and latency.

AI Use Statement

We used generative AI tools to assist with editing portions of the manuscript, retrieving related works, generating and debugging experimental and plotting code, and formatting figures and tables. We also used AI tools to discuss and refine aspects of theoretical presentation and interpretation of experimental results. All AI-assisted technical suggestions, mathematical arguments, code, and interpretations were independently reviewed and verified by the authors against the underlying derivations, implementations, and experimental outputs. No AI generated datasets were used in this work. We take responsibility for the final content of this work, including all text, claims, code, and artifacts produced with AI assistance.

Ethics Statement

The proposed method is intended to reduce the memory and computational cost of long-context inference, which may improve the accessibility and efficiency of deploying language models. As with other efficiency techniques, however, reducing inference cost can also lower the barrier to large-scale model deployment; the downstream societal effects therefore depend on how the resulting models are used. Our work does not modify model objectives, safety mechanisms, or training data, and we do not make claims about changes to model safety, fairness, or robustness beyond the evaluated performance metrics

Reproducibility Statement

We provide sufficient implementation and experimental details to reproduce the results in this work. A.7 specifies the RACE calibration and folding procedure, RoPE-pair handling, value compression, optional KL-based recovery, model and benchmark configurations, long-context evaluation settings, hardware/software environment, precision, number of tokens, kernels, and analytical FLOP derivations. Detailed per-task results, selector ablations, recovery settings, and systems measurements are also provided in the A.8.

Appendix

Axes of KV-Cache Compression

For a model with $L$ layers, batch size $B$, context length $S$, $H_{kv}$ key–value heads, key dimension $d_k$, value dimension $d_v$, and $b$ bytes per element, the KV-cache memory is

This expression exposes three independent factors that can be targeted by compression.

Token-level methods reduce $S$ by evicting, retaining, or consolidating cached positions. Head-dimension methods reduce $d_k$ and $d_v$ for each retained position. Quantization reduces $b$, the storage precision of each element. As illustrated in Figure 8, these approaches operate on different factors of the cache size and can therefore be applied independently or jointly. In particular, combining token and head-dimension compression changes the cache size to

where $S' \le S$, $r_k \le d_k$, and $r_v \le d_v$.

Attention and KV caching.

For hidden states $X \in\mathbb{R}^{S \times d_{\mathrm{model}}}$ and attention weights $W_Q, W_K \in\mathbb{R}^{d_{\mathrm{model}} \times d_k}$ and $W_V \in\mathbb{R}^{d_{\mathrm{model}} \times d_v}$, an attention layer computes $Q = XW_Q$, $K = XW_K$, $V = XW_V$. For each head, let $Q_h, K_h \in\mathbb{R}^{S \times d_k}$ and $V_h \in\mathbb{R}^{S \times d_v}$. Its output is $O_h = A_hV_h = \operatorname{softmax}\left(\frac{Q_hK_h^\top}{\sqrt{d_k}}\right)V_h$.

During autoregressive decoding, previously computed keys and values are stored in the KV cache. For $L$ layers, batch size $B$, context length $S$, $H_{kv}$ KV heads, $H_q$ Q heads, and $b$ bytes per element, its memory is

Thus, reducing $d_k$ and $d_v$ decreases cache memory without removing cached tokens. It also reduces the query-key and attention-value products, whose context-dependent decoding costs scale as $O(BSH_qd_k)$ and $O(BSH_qd_v)$, respectively. Their prefill costs scale as $O(BS^2H_qd_k)$ and $O(BS^2H_qd_v)$, respectively. Complete FLOP expressions are provided in A.4.

RoPE-Compatible Reduction

Rotary positional embeddings. RoPE partitions the rotary query and key dimensions into coordinate pairs and applies a position-dependent two-dimensional rotation to each pair [12]. For the pair associated with frequency $\theta_i$ at position $t$,

Although implementations may store paired coordinates consecutively or across two halves of the head dimension, the two coordinates remain coupled by the rotation. A coordinate-selection method

can therefore preserve the pretrained RoPE operation without reconstruction only if it retains or removes complete rotary pairs. Architectures with partial RoPE additionally contain non-rotary coordinates, which need not satisfy this pairwise constraint.

Logical pairs and physical coordinate layouts. Let $d_{\mathrm{rot}} \le d_k$ denote the number of query and key coordinates to which RoPE is applied, and let $M = d_{\mathrm{rot}/2}$. We represent the logical rotary pairs as

where $(i_m,j_m)$ identifies the two physical coordinates rotated using the $m$-th rotary frequency. The mapping from logical pairs to physical coordinates depends on the model implementation. With one-based indexing, common layouts include

The former is used by GPT-J-style rotary implementations, whereas the latter is used by LLaMA-style implementations. Our selection operates on these logical pairs rather than on neighboring memory locations, so the same procedure applies to either layout after constructing the appropriate pair map.

Restriction of the original RoPE operator. Let $S \subseteq[M]$ be the retained set of rotary pairs and let $P_S \in\{0,1\}^{d_k \times2|S|}$ contain the standard basis vectors corresponding to both coordinates of every pair in $S$. The reduced RoPE operator is formed by retaining the original rotation block associated with each selected pair:

Provided that the columns of $P_S$ and the blocks of $R_t^{(S)}$ use the same pair ordering,

Consequently, for row-wise query activations,

and analogously for keys. Thus, selecting coordinates before RoPE and applying the restricted operator produces exactly the selected coordinates of the original RoPE output. This is an algebraic restriction of the pretrained operator, not an approximation introduced by the reduced-dimensional implementation.

Preserving pretrained rotary frequencies. Reducing the query–key dimension must not cause the rotary frequencies to be recomputed from the new dimension $r_k$. Standard RoPE implementations derive frequencies from their original rotary dimension; reinitializing RoPE with $r_k$ would generally assign different angles to the retained coordinates and would therefore alter the pretrained positional encoding. Instead, for every retained pair $m \in S$, we copy its original frequency $\theta_m$:

The selected pairs may be stored in score order or original-coordinate order, provided that the queries, keys, and frequencies use the same ordering. In our implementation, we retain the original pair order after selection to make the folded checkpoints deterministic.

For models using scaled RoPE, the same principle applies: the model’s original position-dependent frequency transformation and any associated attention scaling are preserved, and only the entries corresponding to selected pairs are gathered. This avoids silently replacing the pretrained positional encoding with a new reduced-dimensional frequency schedule.

Partial RoPE. Some architectures apply RoPE to only $d_{\mathrm{rot}} < d_k$ coordinates. Let $\mathcal{I}_{\mathrm{NoPE}}$ denote the complementary non-rotary coordinates. The complete-pair constraint is imposed only on $\mathcal{P}_{\mathrm{rope}}$; non-rotary coordinates may be selected individually. If $\mathcal{J} \subseteq\mathcal{I}_{\mathrm{NoPE}}$ is the retained non-rotary set, the complete selector is

and the corresponding reduced positional operator is

The identity block leaves the retained non-rotary coordinates unchanged. Hence, the same commutation relation holds for architectures with full or partial RoPE.

Grouped-query attention. Under grouped-query attention, several query heads attend to a shared key head. For each KV head $g$, let $\mathcal{G}_g$ denote its associated query heads. We compute the Q/K pair statistics jointly over the key activations of head $g$ and the query activations from all heads in $\mathcal{G}_g$. A single pair set $\mathcal{S}_g$ is then applied to the shared key head and to every query head in its group. This ensures that each query and the key cache against which it is evaluated use identical retained coordinates and rotary frequencies. Pair sets may differ across KV heads, but they cannot differ among query heads that interact with the same cached key representation.

Implementation consequence. After selection, the columns indexed by $P_{QK}$ are folded into $W_Q$ and $W_K$, and the RoPE module is restricted to the corresponding original frequencies. The resulting layer produces $r_k$-dimensional queries and keys directly. Neither the discarded coordinates nor a full-dimensional RoPE output is materialized during prefill or decoding.

Inference FLOPs Analysis

We analyze the attention-related floating-point operations (FLOPs) incurred during prefill and autoregressive decoding. We count a multiply-add as two FLOPs and omit element-wise operations such as softmax, normalization, and RoPE, which are lower-order relative to the matrix multiplications considered here.

Let $B$ denote the batch size, $S$ the context length, $d$ the model hidden dimension, $H_q$ the number of query heads, and $H_{kv}$ the number of key-value heads. The original query/key and value head dimensions are $d_k$ and $d_v$, respectively. After compression, their corresponding dimensions are $r_k$ and $r_v$.

For prefill, projections are applied to all $S$ input tokens and attention is computed over an $S \times S$ attention matrix. For decoding, we report the FLOPs for generating one new token given a cache of $S$ previous tokens.

Full Attention For standard attention, the query, key, and value projections during prefill cost

The query-key product and attention-value product cost

and the output projection costs

Thus,

During decoding, projections are required only for the newly generated token:

The query attends to all $S$ cached keys and values, giving

while the output projection costs

Therefore,

For long contexts, the cache-length-dependent term dominates:

EigenAttention with RoPE EigenAttention projects attention representations K and V into learned low-dimensional subspaces. In its RoPE-compatible formulation, compressed keys must be decompressed into the original dimension, which is RoPE-compatible, before query-key attention.

During prefill, the projection costs

Transforming the compressed keys to the original key space costs

Since queries and reconstructed keys interact in the original key dimension, the query-key product is the same as full attention. The value representation remains compressed to dimension $r_v$, giving

Thus, the prefill cost is

During decoding, all $S$ cached compressed keys must be reconstructed before RoPE, giving

The subsequent query-key, attention-value operations, and output cost

Therefore, the cache-length-dependent cost of generating one token is

This additional $O\left(BSH_{kv}r_kd_k\right)$ term arises from reconstructing cached keys into the RoPE-compatible space.

PALU with G-LRD and RoPE PALU (Chang et al., [3]) applies group-head low-rank decomposition (G-LRD) to the key and value projection matrices. Let $s$ denote the number of KV heads in each group and $G = \frac{H_{kv}}{s}$ the number of groups. For each key group $g$, PALU factorizes the concatenated projection of $s$ heads as

where

The intermediate representation produced by $A_K^{(g)}$ is stored in the KV cache. For RoPE-based attention, the cached key latent is reconstructed to the original $sd_k$-dimensional group representation before applying RoPE and computing attention. Assuming uniform group ranks $r_k^{(g)} = r_k^G$ and $r_v^{(g)} = r_v^G$, the Q/K/V projection cost during prefill is

For one decoding step,

For each group, reconstructing $S$ cached key latents from dimension $r_k^G$ to $sd_k$ costs $2BSr_k^Gsd_k$. Across all $G$ groups,

The same cache-wide reconstruction is required for each decoding step:

Because reconstructed keys have the original head dimension $d_k$, the query-key products and value products cost

PALU folds the value-side reconstruction into the output projection, so attention-value aggregation can be performed in the grouped latent value space. Let the total latent value width across all groups be $d_V^{\mathrm{latent}} = Gr_v^G$. The precise arithmetic depends on the mapping between query heads and value groups; we therefore retain the value-side latent width explicitly.

The fused output projection maps the concatenated latent value representation back to the model dimension, giving

The prefill cost is

Ignoring the $S$-independent projection and output terms, the dominant cache-length-dependent decoding cost is

Ours Our method reduces the architectural query/key and value head dimensions to $r_k$ and $r_v$, respectively. All compression transformations are folded into the query, key, value, and output projection matrices. Consequently, neither key nor value reconstruction is required during inference. The native-reduced dimensional attention FLOPs are as follows

During prefill,

During decoding,

The cache-length-dependent decoding cost is therefore

For $d_k = d_v$ and $r_k = r_v = d_k/2$, this term is exactly half that of full attention:

RAP follows similar compressed attention as Ours and therefore the prefill and decoding FLOPs would be alike. Similarly, Think Keys Full Values also compresses attention while keeping Values in full dimension $r_v = d_v$.

Derivation of Table 1

Method	Reconstruction free	Custom-Kernel free	Fwd.-only Calibration	Shared Q/K Map	Joint Q/K Signal	Native RoPE	KV Cache Retained	Attn. Params Retained	Prefill FLOPs	Decode FLOPs
Attention ([14])	-	-	-	-	-	-	100%	100%	$\times1$	$\times1$
EigenAttention ([10])	✗	✓	✓	✗	✓	✗	50%	69.8%	$\times0.75$	$\times11.42$
PALU ([3])	✗	✗	✓	✗	✗	✗	50%	58.3%	$\times0.75$	$\times21.92$
Thin Keys ([20])	✓	✓	✓	✗	✗	✗	75%	50%	$\times0.75$	$\times0.75$
RAP ([18])	✓	✓	✗	✓	✗	✓	50%	50%	$\times0.5$	$\times0.5$
RACE (Ours)	✓	✓	✓	✓	✓	✓	50%	50%	$\times0.5$	$\times0.5$

Table 1 compares methods under the representative grouped-query attention configuration $d = 1536$, $H_q = 6$, $H_{kv} = 2$, and $d_k = d_v = 128$. For RACE, RAP, and EigenAttention, we use $r_k = r_v = 64$, while Thin Keys uses $r_k = 64$ and $r_v = 128$. For PALU G-LRD, the two KV heads form a single group ($G = 1$, $s = 2$) with group ranks $r_k^G = r_v^G = 128$; these ranks are defined per group rather than per head. We count a multiply–add as two FLOPs and omit elementwise operations. For prefill, we report the coefficient of the dominant $O(S^2)$ attention computation, while for single-token decoding, we report all operations that scale linearly with the cached sequence length $S$. Common factors such as $2BS^2$ for prefill and $2BS$ for decoding therefore cancel in the normalized ratios. This convention captures the distinction between native dimensional reduction, where attention is performed directly in the compressed representation, and reconstruction-based methods, which incur additional cache-length-dependent work during decoding.

RACE, RAP, and Thin Keys. Both methods execute attention natively with $r_k = r_v = 64$, without reconstructing the discarded dimensions. Their cache retention is

The same dimensional reduction applies to the Q, K, V, and output projections, giving $P_{\mathrm{native}}/P_{\mathrm{full}} = 0.5$. Since both the query–key and attention–value products are evaluated directly in dimensions $r_k$ and $r_v$, the normalized dominant prefill and decoding costs are also

For the Thin Keys baseline, $r_v = 128$ as they do not compress values. This gives $\rho_{\mathrm{cache}} = 0.75$ and $\rho_{\mathrm{prefill}} = \rho_{\mathrm{decode}} = 0.75$.

EigenAttention. EigenAttention stores keys and values in reduced dimensions $r_k = r_v = 64$, giving 50% KV-cache retention. Following the parameter-counting convention in Table 1, we count the compressed projections together with the learned key and value bases:

For the representative configuration used for parameter counting, $d = 1536$, $H_q = H_{kv} = 6$, $d_k = d_v = 128$, and $r_k = r_v = 64$, this gives

During prefill, key reconstruction scales as $O(S)$, while the query–key and attention–value products scale as $O(S^2)$. Hence, asymptotically the dominant prefill cost is

For decoding, however, every cached compressed key must be mapped back to the original key dimension before query-key attention. The cache-length-dependent cost $F_{\mathrm{Eigen,LC}}^{\mathrm{dec}} \simeq2BS[H_{kv}r_kd_k + H_q(d_k + r_v)]$. Normalizing by full attention, $F_{\mathrm{Full,LC}}^{\mathrm{dec}} = 2BSH_q(d_k + d_v)$, gives

PALU. PALU G-LRD jointly compresses groups of KV heads. For the parameter comparison in Table 1, we use $H_{kv} = 2$ heads divided into $G = 1$ groups of $s = 2$ heads, with group ranks $r_k^G = r_v^G = 128$. The resulting cache stores

scalars per token, corresponding to 50% of the full KV cache.

Following the parameter-counting convention in Table 1, we count the query projection, both low-rank factor matrices for the grouped key and value projections, and the fused value-output map:

which gives

PALU reconstructs cached keys into the original key dimension before query-key attention, while value aggregation is performed directly in the grouped latent value space. Consequently, the dominant cache-length-dependent decoding cost is

Normalizing by full attention gives

Fore prefill,

Approximate Matrix Multiplication View of Attention Compression

This section formalizes the connection between head-dimension compression and approximate matrix multiplication (AMM), and clarifies how our RoPE-constrained coordinate selection relates to existing sampling-based AMM results.

AMM formulation. Consider a single attention head with post-RoPE query and key matrices $Q, K \in\mathbb{R}^{S \times d_k}$. The pre-softmax attention interaction is $L = QK^{\top}$.

Following the standard AMM formulation, let

so that $A^{\top}B = QK^{\top}$. Given a dimensionality-reducing operator $\Pi\in\mathbb{R}^{r_k \times d_k}$, AMM replaces this product by

Thus, dimensionality reduction along the shared row dimension of $A$ and $B$ is exactly dimensionality reduction along the head-feature dimension of $Q$ and $K$.

Stable-rank AMM. For conforming matrices $A$ and $B$, Cohen et al. [4] define a $(k, \epsilon)$ spectral-norm AMM property of the form

They show that suitable oblivious embeddings can satisfy this property using

rows with probability at least $1-\delta$.

The stable rank of a matrix $A$ is

Setting

makes the Frobenius-to-spectral terms in Eq. equation 41 controlled by the corresponding spectral norms. Consequently, the required sketch dimension can scale with the effective spectral dimension of the operands rather than their ambient dimension. Applied to attention with $A = Q^{\top}$ and $B = K^{\top}$, this provides theoretical motivation for reducing $d_k$ when the query and key activations have low stable rank.

From row sampling in AMM to feature selection in attention. AMM sampling methods act on the shared row dimension of the two operands. Under the identification $A=\widetilde{Q}^{\mathsf{T}}$; $B=\widetilde{K}^{\mathsf{T}}$ the rows of $A$ and $B$ correspond to the feature dimensions of $\widetilde{Q}$ and $\widetilde{K}$. Let

select the coordinate set $\mathcal{I}$, and take $\Pi=P_{\mathcal{I}}^{\mathsf{T}}$. The AMM approximation then becomes

Thus, row selection in AMM is exactly feature selection in attention, with the same selection matrix applied to queries and keys.

Equation 43 motivates two design choices. First, queries and keys should use an identical reduction $P_{\mathcal{I}}$ or in the context of RACE, $P_{QK}$ rather than independently constructed maps $P_Q$ and $P_K$.

Separate maps would produce $\left(\widetilde{Q}P_Q\right)\left(\widetilde{K}P_K\right)^{\mathsf{T}}=\widetilde{Q}P_QP_K^{\mathsf{T}}\widetilde{K}^{\mathsf{T}}$, which is not the shared-sketch AMM approximation $\widetilde{Q}P_{QK}P_{QK}^{\mathsf{T}}\widetilde{K}^{\mathsf{T}}$.

Second, because the approximation target is the joint product $\widetilde{Q}\widetilde{K}^{\mathsf{T}}$, a data-dependent choice of $P_{QK}$ should account for both operands. Selecting dimensions using only query or only key statistics can discard features that have low importance for one operand but materially affect their interaction.

These observations determine the form of the reduction - a common, jointly informed Q/K selector, but not its precise selection rule. General AMM results show that appropriately chosen subsets of the shared dimension can preserve matrix products. RACE further restricts this family to selectors that are compatible with RoPE and derives an activation-energy criterion for choosing among them.

Relation to our selector. RACE’s compression operator is a more restricted member of this general family. For each attention head, we use an unweighted binary selection matrix $P_{QK}=[e_i]_{i\in\mathcal{I}_{QK}}$, where $\mathcal{I}_{QK}$ is chosen from calibration activations. Unlike the general data-dependent sampling operators covered by AMM theory, our selector must additionally satisfy the structural constraints induced by RoPE. Specifically, if $\mathcal{P}_{\mathrm{rope}}=\{(i_m,j_m)\}_{m=1}^{M}$ denotes the rotary coordinate pairs, then

The selected set is therefore a union of complete two-dimensional rotary blocks.

Energy-based selection and scope of the guarantee. The AMM results above establish that low-dimensional and sampling-based representations capable of preserving matrix products exist under appropriate conditions. They do not imply that every coordinate subset, or specifically our RoPE-constrained subset, satisfies the same AMM bound.

Our method chooses among the admissible RoPE pairs using calibration energy. For pair $(i_m,j_m)$, we define

and retain the $r_k/2$ highest-energy pairs. This criterion favors coordinates that carry substantial query/key activation energy while remaining compatible with RoPE.

We therefore use stable-rank AMM and sampling-based AMM as theoretical motivation for reducing the shared query–key dimension, rather than claiming a direct AMM guarantee for the proposed energy-based selector. The exact property guaranteed by our construction is structural: selecting complete rotary pairs permits RoPE and attention to be executed natively in the reduced space.

Proofs

Proof of Theorem 1. [RoPE-compatible coordinate selection] Proof. Suppose $P_{QK}$ selects exactly the coordinates belonging to the complete rotary pairs indexed by $S \subseteq\{1,\ldots,M\}$, with its columns ordered consistently with the blocks of $R_t^{(S)}$. Since

acts independently on each rotary pair, applying $R_t$ to the selected coordinates applies precisely the blocks $R_{m,t}$ for $m \in S$. Therefore,

for every $t$. Hence $P_{QK}$ is RoPE-compatible.

Conversely, suppose

The right-hand side lies in $\operatorname{col}(P_{QK})$, and therefore

Thus, the coordinate subspace selected by $P_{QK}$ must be invariant under every $R_t$. Consider any rotary pair $(i_m,j_m)$. Suppose, for contradiction, that $P_{QK}$ selects $i_m$ but not $j_m$. Then $e_{i_m}\in\operatorname{col}(P_{QK})$, whereas $e_{j_m}\notin\operatorname{col}(P_{QK})$. For the $m$-th rotary block,

so, up to the coordinate convention for the pair,

By assumption, pair $m$ undergoes a nontrivial rotation for at least one position, so there exists $t$ such that $\sin\phi_{m,t}\ne0$. For this $t$, $R_te_{i_m}$ has a nonzero component along $e_{j_m}$ and therefore

contradicting the invariance established above.

The same argument applies if $j_m$ is selected but $i_m$ is not. Hence, for every rotary pair, either both coordinates are selected or neither is selected. Therefore the coordinates selected by $P_{QK}$ form a union of complete rotary pairs, indexed by some $S \subseteq\{1,\ldots,M\}$.

Proof of Proposition 1. [Error bound for RoPE-pair truncation]

Proof. Partition the post-RoPE query and key matrices into RoPE-aligned coordinate pairs,

where each

Because matrix multiplication along the feature dimension decomposes additively across these blocks,

For a retained pair set $S\subseteq[M]$, the corresponding truncated product is

Hence the truncation error would be,

Applying the triangle inequality for the Frobenius norm gives

Since for any conforming matrices $A$ and $B$, $\left\lVert AB^{\mathsf{T}}\right\rVert_{F}\leq\left\lVert A\right\rVert_{F}\left\lVert B\right\rVert_{F}$, the above equation can be written as

Finally, by the arithmetic–geometric mean inequality, we obtain

Combining the above inequalities yields

For a fixed budget $\lvert S\rvert=k$, minimizing the right-hand side is equivalent to maximizing

Therefore, the subset consisting of the $k$ largest-energy RoPE pairs minimizes this upper bound among all subsets of $k$ pairs. □

Evaluation in the pre- or post-RoPE space. Because each RoPE pair is transformed by an orthogonal $2\times2$ rotation,

Hence,

The pair-energy criterion can therefore be evaluated equivalently in the pre- or post-RoPE space. The implementation accumulates the joint Q/K second moment and ranks each pair by the sum of its two diagonal entries, which is proportional to this $E_m$. In pour implementations we use the post-RoPE values for the calibration.

Connection to attention-output perturbation. Let

The Jacobian of softmax at a probability vector $p$ is $\operatorname{diag}(p)-pp^{\mathsf{T}}$, whose spectral norm is at most $1/2$. Applying this result row-wise gives

Because $Z-Z_{S}\in\mathbb{R}^{S_{\mathrm{cal}}\times S_{\mathrm{cal}}}$,

For fixed values $V$,

where the final inequality follows from Proposition 1.

More generally, if the downstream calibration loss is locally $L_{\mathrm{down}}$-Lipschitz with respect to the attention output, then

This is a conditional, calibration-dependent bound and does not imply that energy selection globally minimizes language-model loss.

Implementation Details

Software and model execution. We implement RACE in PyTorch and Hugging Face Transformers. Our implementation supports GPT-NeoX, Llama, Qwen2, Pythia, and Granite-style attention modules, including both multi-head and grouped-query attention. Although all models reported in the main paper use full-RoPE or RoPE applied to all dimensions, our implementation supports partial RoPE as well by applying pairwise selection only within the rotary-encoded subspace while leaving any non-rotary dimensions unconstrained by the RoPE-pair requirement. Pre-trained models are loaded and evaluated in BF16 on GPU, while calibration statistics and projection construction use higher-precision accumulation during calibration. The folded compressed model uses standard attention operations and does not require a custom CUDA or Triton kernel. Unless otherwise stated, we use PyTorch scaled dot-product attention (SDPA) for the compressed runtime. The implementation also supports eager attention and FlashAttention-3 as shown below.

Furthermore, PALU rely on method-specific execution paths for efficient inference, including custom fused kernels. We were unable to reliably reproduce these optimized kernels in our hardware/software environment and therefore avoid presenting potentially misleading wall-clock comparisons. Instead, we compare algorithmic attention FLOPs across methods and measure RACE end-to-end using standard PyTorch attention backends. This separates implementation-specific kernel efficiency from the architectural distinction that RACE performs attention directly in the reduced dimensions without runtime reconstruction. The original rotary-frequency indices and the model’s RoPE scaling rule are preserved; selected frequencies are not re-indexed as consecutive frequencies. For grouped-query attention (GQA), all query heads associated with a given KV head use the same selected key subspace. For partial RoPE models like Pythia models, the RoPE pairs are considered as pairs and the non-RoPE dimensions are paired with adjacent dimensions as in interleaved RoPE as shown in Figure 9.

Calibration in RACE. For each attention layer, we collect query and key activations after the pre-trained RoPE transformation. Our implementation also support calibration with pre-RoPE queries and keys and we also show that they are equivalent in A.6. We preserve the physical rotary-pair layout of the original architecture rather than assuming that adjacent tensor coordinates necessarily form a rotary pair. For each rotary pair $m$, we accumulate the joint activation-energy score as in Eq. (4) and retain the $r_{k}/2$ highest-scoring pairs. Selection is performed independently for each layer and attention head (or KV head under grouped-query attention), subject to retaining both coordinates of every selected RoPE pair.

Note that models like Llama, Qwen, and Granite have full RoPE where all dimensions have RoPE and in such cases the sequential selection shown in 9 Left will not be present. Such models also have sequential RoPE $[x_0, x_{d/2}], [x_1, x_{d/2+1}], \ldots$. Our implementations supports both sequential and interleaved $[x_0, x_1], [x_2, x_3], \ldots$ RoPE models. The values are not subject to RoPE. We employ $P_V$, the right singular vectors of $V$ for the compression. This do not follow selection as shown in Figure 9 Right.

Optional post-compression recovery. RACE is a post-training compression method and can be evaluated immediately after constructing and folding the reduced Q/K and V subspaces into the pretrained attention projections. At aggressive compression ratios, however, this discrete change in attention dimensionality can perturb the pretrained representation enough that a short optimization stage is highly beneficial. Although optional, all results in the main paper involve this recovery. Following folding, we use the original uncompressed checkpoint as a frozen teacher and optimize the compressed model using KL distillation on FineWeb data. The implementation uses a temperature of $T = 2$ and minimizes the divergence between the teacher and compressed-model output distributions. We train for 250 optimizer steps with sequences of length 512, a micro-batch size of 2, and gradient accumulation over 4 micro-batches, corresponding to approximately 1.0M training tokens.

Figure 10 separates the effect of this recovery from the compression procedure itself by reporting both the directly folded models (“w/o FT”) and their recovered counterparts. At 75% and 87.5% KV-cache retention, the directly compressed models are already relatively close to their recovered versions, indicating that moderate dimensionality reduction introduces limited disruption. In contrast, at 50% retention, both RACE and RAP experience substantially larger degradation without recovery, particularly in perplexity. The optimization stage is not necessary to define or execute the compressed model. RACE produces a valid reconstruction-free reduced-dimensional attention layer immediately after folding. Rather, recovery provides an optional mechanism for adapting the pretrained weights when the dimensionality change becomes sufficiently aggressive. Moreover, the same recovery procedure also substantially improves RAP, suggesting that much of the recoverable degradation at aggressive compression reflects adaptation to the compressed representation rather than a method-specific training advantage.

Systems Details All systems experiments were conducted on a single NVIDIA H100 80GB HBM3 GPU. The machine used NVIDIA driver 590.48.01 with CUDA 13.1, and all reported memory, latency, throughput, and FLOP-related measurements were collected on this hardware configuration. Unless otherwise noted, experiments were run on a single GPU without model parallelism using SDPA kernels, so comparisons between the base model and compressed variants reflect changes induced by the compression method rather than differences in hardware allocation.

Results

Language Modeling and Zero-Shot Performance The main quality experiments evaluate Qwen2.5-7B, Llama-3.2-3B-Instruct, and Granite-4.1-3B. We compare against PALU-GLRD, EigenAttention, Thin Keys, and RAP at matched or reported KV-cache budgets. The full results are presented in Table 3 while only RAP, a reconstruction-free method like RACE, and PALU-GLRD, a reconstruction-based method, are included in Figure 3. Across all three model families, RACE consistently outperforms the reconstruction-free baselines RAP and Thin Keys, while remaining competitive with the reconstruction-based EigenAttention and PALU. The difference is most pronounced under aggressive compression: At 50% KV-cache retention, RACE achieves the lowest WikiText perplexity across all three model families while remaining competitive with PALU in average zero-shot accuracy, among the evaluated compression methods on each of Qwen2.5, Granite, and Llama. At more moderate compression, RACE remains close to the reconstruction-based methods; at 75% retention it preserves 98.5%, 96.5%, and 97.5% of the corresponding uncompressed average zero-shot accuracy on Qwen, Granite, and Llama, respectively. The per-task results also show that the effect of compression is not uniform across benchmarks: some tasks remain close to the uncompressed baseline even at lower retention, whereas others are more sensitive to reduced attention dimensionality. Table 4 presents the per-task results of Qwen2.5-7B and Llama-3.2-3B-Instruct on 16 English tasks in LongBench. RACE outperforms RAP at all evaluated retention levels and exceeds PALU at 87.5% and 75% retention on both model families. At 50% retention, RACE remains competitive with PALU, outperforming it on Qwen while trailing it modestly on Llama. Together, these results show that the quality-compression behavior is consistent across model families rather than being specific to a single architecture, while retaining the reconstruction-free execution of RACE.

Additional Systems Results

The systems results in Figure 11 show that RACE reduces both the fixed attention-parameter footprint and the context-dependent KV-cache footprint. As expected, the reduction in attention parameter memory aligns with the compression. However, since the folded projection weights are smaller in comparison to the result of the model parameters and memory consumption elements, the reduction in memory is less prominent in the total parameter memory.

Additionally, parameter memory is independent of context length. KV-cache memory scales with sequence length and therefore yields increasingly large absolute savings at longer contexts. At 256K tokens, the gap between the base model and the compressed variants becomes substantial, illustrating that the main systems benefit of RACE grows with context length while preserving native reduced-dimensional execution.

Compatibility with Token compression and Quantization

Table 5 reports the full per-task results when RACE is combined with token-compression and KV-quantization methods. Across H2O, SnapKV, and KIVI, applying RACE on top of an already compressed KV cache provides additional multiplicative memory reduction, but also introduces further accuracy and perplexity degradation. This effect is generally more pronounced for token-compression methods than for quantization.

Tokens /Precision Retained	RACE KV Cache Retained	Total KV Cache Retained	wikitext	hellaswag	piqa	arc_challenge	boolq	mmlu_college_cs	Average
512 (50%)	100%	50%	12.66	0.66	0.85	0.48	0.82	0.72	0.706
	75%	37.5%	15.8	0.65	0.83	0.45	0.81	0.63	0.674
	50%	25%	17.44	0.62	0.79	0.44	0.74	0.43	0.604
256 (25%)	100%	25%	13.62	0.66	0.85	0.48	0.79	0.66	0.688
	75%	18.75%	16.12	0.65	0.83	0.45	0.78	0.63	0.668
	50%	12.5%	21.81	0.62	0.79	0.42	0.71	0.4	0.588
512 (50%)	100%	50%	12.55	0.68	0.86	0.48	0.82	0.71	0.71
	75%	37.5%	15.61	0.65	0.83	0.46	0.81	0.63	0.676
	50%	25%	18.63	0.62	0.79	0.45	0.76	0.44	0.612
256 (25%)	100%	25%	12.76	0.68	0.85	0.48	0.81	0.7	0.704
	75%	18.75%	15.66	0.65	0.83	0.45	0.79	0.63	0.67
	50%	12.5%	18.44	0.62	0.79	0.46	0.74	0.44	0.61
2 bits (25%)	100%	25%	17.04	0.69	0.86	0.48	0.78	0.7	0.702
	75%	18.75%	18.17	0.66	0.84	0.48	0.76	0.6	0.668
	50%	12.5%	21.81	0.64	0.8	0.44	0.71	0.42	0.602
4 bits (50%)	100%	50%	15.7	0.69	0.86	0.48	0.78	0.7	0.702
	75%	37.5%	16.58	0.66	0.84	0.48	0.78	0.63	0.678
	50%	25%	17.14	0.64	0.8	0.44	0.74	0.42	0.608

For H2O and SnapKV, reducing the head dimension with RACE consistently lowers the total KV-cache footprint from 50% to 37.5% and 25%, or from 25% to 18.75% and 12.5%, depending on the token-retention ratio. The corresponding degradation is gradual across most zero-shot tasks, while WikiText perplexity increases more noticeably at aggressive combined compression. For example, with H2O at 50% token retention, adding RACE at 75% and 50% head retention reduces total KV-cache usage to 37.5% and 25%, respectively, while average zero-shot accuracy decreases from 0.706 to 0.674 and 0.604. SnapKV follows a similar trend, with average accuracy decreasing from 0.710 to 0.676 and 0.612 under the same settings.

RACE is comparatively more complementary with KIVI. With 4-bit quantization, adding RACE at 75% and 50% head retention reduces the total KV cache from 50% to 37.5% and 25%, while average accuracy decreases from 0.702 to 0.678 and 0.608. Similar behavior is observed with 2-bit quantization. Overall, the results suggest that RACE composes cleanly with both token and precision compression, but the cost of applying additional head-dimension compression is larger once information has already been removed along the token dimension. In contrast, quantization leaves more redundancy for RACE to exploit, resulting in a milder degradation at comparable total cache budgets.

Ablations We first examine whether the retained RoPE pairs should be selected using information from queries and keys jointly or from either projection alone. Across all retention ratios, the joint Q/K energy criterion gives the strongest overall results. At 87.5% retention, Q-only and K-only selection remain close to the joint criterion, but the gap becomes clearer as compression increases. At 75% retention, joint Q/K energy achieves 0.61 average accuracy with a WikiText perplexity of Table 6: Ablation of selection criteria for RoPE-pair retention. We compare RACE’s joint Q/K activation-energy criterion against magnitude, random selection, gradient Fisher, attention-score selection, and Q-only or K-only energy. Results are reported on WikiText perplexity and zero-shot accuracy across nine tasks at 87.5%, 75%, and 50% KV-cache retention. Joint Q/K energy is consistently strongest overall, with the advantage becoming more pronounced at higher compression.

14.95, compared with 0.59/16.78 for Q-only and 0.60/16.17 for K-only. At 50% retention, the joint criterion retains 0.54 average accuracy and 16.02 perplexity, while both Q-only and K-only fall to 0.52 average accuracy with substantially higher perplexity.

This trend is also reflected directly in the approximation of the uncompressed attention interaction. Figure 12 shows the layer-wise error norm of the compressed $QK^{T}$ product relative to the uncompressed model. Joint Q/K selection gives the lowest average error, 0.105, compared with 0.112 for K-only and 0.116 for Q-only, and the reduction is observed across most layers rather than being concentrated in a small subset. Together, the downstream results and the layer-wise approximation error support deriving the shared selector from both Q and K rather than from either side in isolation.

We next compare different scoring criteria that all operate on joint Q/K information. Joint activation energy consistently outperforms magnitude, random selection, gradient-Fisher scoring, and attention-based scoring, with the differences becoming larger at stronger compression. At 50% retention, RACE achieves 0.54 average accuracy, compared with 0.48 for gradient Fisher and attention score, 0.45 for magnitude, and 0.41 for random selection. The same trend is reflected in WikiText perplexity, where joint Q/K energy also achieves the lowest value. Notably, the gradient- and attention-based alternatives do not recover the performance of the simple forward-pass energy criterion despite requiring more complex statistics. Overall, these ablations indicate that both components of the selector are important: the ranking should be derived jointly from Q and K, and joint activation energy provides an effective criterion for identifying the RoPE pairs that best preserve the original attention interaction.