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

> RACE compresses pretrained RoPE attention natively by retaining complete rotary pairs, cutting inference FLOPs by 47.8% and boosting throughput 1.79× with no runtime reconstruction.

- **Source:** [arXiv](https://arxiv.org/abs/2610.28385)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/6vnSyt
- **Whiteboard:** https://picx.dev/p/6vnSyt/image

## Summary

## Summary (Overview)

- **RACE (RoPE-Aligned Attention Compression based on activation Energy)** is a post-training method that converts pretrained RoPE-based attention layers into natively reduced-dimensional form, reducing KV-cache storage, attention computation, and parameters simultaneously.
- The method is **reconstruction-free**: queries, keys, and values operate entirely in reduced dimensions during inference, requiring no runtime reconstruction of discarded coordinates and no specialized kernels.
- RACE selects complete rotary pairs using a **joint query-key activation energy** criterion, $E_m = \|Q_m\|_F^2 + \|K_m\|_F^2$, which minimizes a separable, calibration-dependent upper bound on post-RoPE query-key product approximation error.
- A key theoretical contribution is the proof that **coordinate-selection reductions commute with pretrained RoPE if and only if complete rotary pairs are retained**, characterizing the admissible family of native reductions.
- Empirically, at 50% KV-cache retention and 256K context, RACE reduces total inference FLOPs by 47.8%, yielding 1.79× higher throughput, 43.2% lower time-to-first-token (TTFT), and 28.1% lower allocated GPU memory compared to the uncompressed model.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Autoregressive language-model inference is increasingly constrained by the **key-value (KV) cache**, especially at long contexts. Each attention layer stores key and value vectors for every preceding token, so cache memory grows linearly with context length and with the dimensionality of cached representations. Reducing key and value head dimensions offers the possibility of decreasing both KV-cache memory and native attention computation.

However, existing head-dimension compression methods (e.g., EigenAttention, PALU) project keys and values into compact representations but **reconstruct full-dimensional keys during inference** to effectively apply rotary positional embeddings (RoPE). This reduces cache storage but introduces decompression and additional arithmetic into the attention path during inference.

### Core Theoretical Problem

The fundamental question addressed is: *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 $\widetilde{Q}\widetilde{K}^\top$, head-dimension reduction is viewed 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—the reduction must be **realizable before RoPE**, imposing a structural constraint that the reduction must commute with the pretrained rotary transformation.

### Problem Formulation

For 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$, the goal is to find reduced-dimensional representations:

$$
Q' = Q P_Q, \qquad K' = K P_K, \qquad V' = V P_V \tag{1}
$$

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 layer stores only $K'$ and $V'$ in the KV cache and performs attention directly in dimensions $r_k$ and $r_v$.

---

## Methodology

### 3.1 Constrained AMM Formulation

The Q/K head-dimension reduction is formulated as an AMM problem. For a column-orthonormal reduction $\bar{\boldsymbol{P}} \in \mathbb{R}^{d_k \times r_k}$:

$$
(\widetilde{Q}P)(\widetilde{K}P)^\top = \widetilde{Q}PP^\top\widetilde{K}^\top \approx \widetilde{Q}\widetilde{K}^\top \tag{2}
$$

This formulation motivates two design principles:
1. **Shared reduction**: Queries and keys should share a common reduction $P_{QK}$, since their interaction is evaluated along the same feature dimension.
2. **Joint selection**: The selection criterion should reflect both matrices whose product is being approximated, using statistics of both query and key activations.

### 3.2 RoPE-Compatible Query–Key Reduction

Let $R_t = \text{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 \times 2}$ acts on one rotary pair. A reduction $P_{QK} \in \mathbb{R}^{d_k \times r_k}$ is **RoPE-compatible** if:

$$
R_t P_{QK} = P_{QK} R_t^{(\mathcal{S})} \qquad \text{for all } t \tag{3}
$$

**Theorem 1 (RoPE-compatible coordinate selection):** Suppose for every rotary pair $m$, there exists a position $t$ such that $\sin \phi_{m,t} \neq 0$. A coordinate-selection matrix $P_{QK} \in \mathbb{R}^{d_k \times r_k}$ satisfies Eq. 3 for all $t$ if and only if it selects exactly the coordinates belonging to a union of complete rotary pairs indexed by $\mathcal{S} \subseteq \{1, \ldots, M\}$, with the selected coordinates ordered consistently with $R_t^{(\mathcal{S})}$.

### 3.3 Calibration-Guided Selection

For each RoPE pair $m$, RACE assigns the **joint query-key energy**:

$$
E_m = \|Q_m\|_F^2 + \|K_m\|_F^2 \tag{4}
$$

For target dimension $r_k$, the $r_k/2$ pairs with the largest energies are retained. 

**Proposition 1 (Pair-energy approximation bound):** For any retained pair set $\mathcal{S}$:

$$
\left\| \widetilde{Q}\widetilde{K}^\top - \sum_{m \in \mathcal{S}} \widetilde{Q}_m \widetilde{K}_m^\top \right\|_F \leq \sum_{m \notin \mathcal{S}} \|\widetilde{Q}_m\|_F \|\widetilde{K}_m\|_F \leq \frac{1}{2} \sum_{m \notin \mathcal{S}} E_m \tag{5}
$$

Selecting pairs with the largest $E_m$ minimizes this separable upper bound.

### Value Compression

Values are not subject to RoPE, so RACE follows standard low-rank activation projection: retain the top-$r_v$ right singular vectors $P_V \in \mathbb{R}^{d_v \times r_v}$ from calibration activations.

### 3.4 Folding into Native Attention

The reductions are folded directly into pretrained projections:

$$
W_Q^{\text{fold}} = W_Q P_{QK}, \qquad W_K^{\text{fold}} = W_K P_{QK} \tag{6}
$$

$$
W_V^{\text{fold}} = W_V P_V, \qquad W_O^{\text{fold}} = P_V^\top W_O
$$

This converts the original attention layer into a native architecture with query/key dimension $r_k$ and value dimension $r_v$, where RoPE, caching, and attention operate entirely in reduced dimensions.

### Recovery

For aggressive compression, a short one-time post-fold recovery stage using KL distillation from the uncompressed model is recommended (1M FineWeb-Edu tokens), updating only the reduced attention parameters.

---

## Empirical Validation / Results

### Experimental Setup

- **Models**: Qwen2.5-7B, Llama-3.2-3B-Instruct, Granite-4.1-3B
- **Calibration**: 32,768 FineWeb-Edu tokens, FP32 statistics, BF16 inference
- **Baselines**: PALU-GLRD, EigenAttention, Thin Keys, RAP (concurrent work)
- **Metrics**: WikiText perplexity, 15 zero-shot tasks, LongBench (16 English tasks)

### Quality Results

At 75% KV-cache retention, RACE preserves:
- **98.5%** of uncompressed zero-shot accuracy on Qwen
- **96.5%** on Granite
- **97.5%** on Llama

Across compression levels, RACE closely matches PALU on zero-shot accuracy and LongBench without runtime reconstruction, and consistently outperforms RAP.

### Systems Benefits (Llama-3.2-3B-Instruct, 256K context, 50% retention)

| Metric | Improvement |
|--------|-------------|
| Prefill throughput | 1.76× speedup |
| Decode throughput | 1.90× speedup |
| Overall throughput | 1.79× speedup |
| Time-to-first-token | 43.22% lower |
| Allocated GPU memory | 28.06% lower |
| Total inference FLOPs | 47.75% lower |

### Compatibility with Other Compression Methods

RACE can be combined with token-level compression (H2O, SnapKV) and quantization (KIVI). At 75% RACE retention + 50% token/quantization compression, combined retention is 37.5%. RACE is especially complementary to quantization (KIVI), showing smaller degradation.

### Ablation: Selection Criteria

| Criterion | Backprop-free | Task-loss-free | Explicit $O(S^2)$ scoring | Time (sec) | Throughput (tok./sec) | Memory (GB) |
|-----------|:---:|:---:|:---:|:---:|:---:|:---:|
| **Energy (Ours)** | ✓ | ✓ | ✗ | 2.717 | 12061.6 | 6.142 |
| Magnitude | ✓ | ✓ | ✗ | 2.925 | 11202.4 | 6.142 |
| Attn. Product | ✓ | ✓ | ✓ | 11.616 | 2821 | 10.592 |
| Grad. Fisher | ✗ | ✗ | ✗ | 8.290 | 3952.7 | 9.218 |

RACE achieves the lowest perplexity among selectors while having the lowest calibration cost. The normalized query-key product error is 0.1051 for joint Q/K selection vs. 0.1159 (Q-only) and 0.1120 (K-only).

### Effect of Head Width

Wider-head models (4 × 256) retain perplexity more effectively under aggressive compression than standard (8 × 128) or narrow (16 × 64) heads at matched model size, as wider heads expose more RoPE pairs for selection.

---

## Theoretical and Practical Implications

### Theoretical Contributions

1. **Constrained-AMM formulation**: Frames native attention compression as preservation of the post-RoPE query-key product subject to a pre-RoPE realizability constraint.
2. **Structural characterization**: Proves that RoPE-compatible coordinate selection requires retaining complete rotary pairs, exactly characterizing the admissible family.
3. **Optimal selection criterion**: Shows that joint Q/K activation energy minimizes a separable, calibration-dependent upper bound on approximation error.

### Practical Implications

- **Reconstruction-free execution**: RACE eliminates the need for runtime reconstruction, reducing decoding FLOPs by up to 2× compared to reconstruction-based methods.
- **No specialized kernels**: Works with standard PyTorch attention operations (SDPA, eager, flash attention backends).
- **Forward-only calibration**: Requires only forward passes, avoiding backpropagation and task-loss dependence.
- **Complementary to existing methods**: Can be combined with token eviction, compaction, and quantization for further compression.
- **Architectural guidance**: Suggests that wider heads during pretraining may yield more favorable quality-efficiency trade-offs when post-training compression is anticipated.

---

## Conclusion

RACE addresses the fundamental question of when and how pretrained RoPE attention can be converted into a natively lower-dimensional form. Key takeaways:

1. **Theoretical foundation**: Coordinate-selection reductions are RoPE-compatible if and only if complete rotary pairs are retained, and joint Q/K activation energy minimizes a provable upper bound on approximation error.

2. **Practical effectiveness**: RACE outperforms evaluated reconstruction-free baselines and remains competitive with reconstruction-based methods while providing substantial systems benefits (47.8% FLOP reduction, 1.79× throughput improvement at 50% retention and 256K context).

3. **Flexibility**: The method is compatible with token-level compression and quantization approaches, and its forward-only calibration makes it practical for large models.

**Future directions** include:
- Exploring heterogeneous compression levels across layers and heads
- Investigating input-dependent or adaptive compression strategies within the RoPE-compatible framework
- Studying the interaction between head width and compressibility during pretraining
- Extending the theoretical framework to other positional encoding schemes

---

_Markdown view of https://picx.dev/p/6vnSyt, served by PicX — AI-generated visual whiteboard summaries of research papers._