# Token-Count-Aware Attention Projection

- ID: 2262
- Canonical URL: https://synthcore.org/idea/2262/token-count-aware-attention-projection
- API JSON: https://synthcore.org/api/idea/2262.json
- API Markdown: https://synthcore.org/api/idea/2262.md
- Verification status: unverified
- Source: [arXiv:2608.13782](https://arxiv.org/abs/2608.13782)
- Category: memory
- Solves: scalability, memory, speedup
- ML areas: attention, transformer, kv-cache, inference-speedup, memory
- Math tags: linear-algebra, metric-geometry, probability, approximation-theory
- Ratings: usefulness 6/10; difficulty 6/10; novelty 6/10

## Idea description

Use the number of tokens participating in an attention operation to adapt the projected query/key dimension according to the sharp Johnson–Lindenstrauss bound. Short attention windows receive a dimension limited by n-1, while longer windows use a logarithmic dimension, avoiding a fixed large head dimension when the current attention set contains too few points to require it.

## Mathematical statement

For an n-point set X in R^d, the paper gives r = Theta(min{d,n-1,log(2+epsilon^2 n)/epsilon^2}) for preserving every pairwise Euclidean distance with distortion at most 1+epsilon. Here d is the source dimension, n is the number of tokens in the local attention window, epsilon in (0,1/2) is the distortion tolerance, and r is the projected dimension. Apply this to normalized key vectors k_1,...,k_n and queries, treating Euclidean geometry as a proxy for attention similarity. Let P_Q and P_K have entries distributed as N(0,1/r), and use projected logits (P_Q q_i)^T(P_K k_j)/sqrt(r). The theorem directly controls distances rather than dot products or softmax probabilities, so the implementation must empirically measure logit and attention-distribution preservation. Using one shared projection for Q and K preserves a common coordinate system and is the simplest initial variant.

## Key formulas

- $$r=\Theta\left(\min\left\{d,n-1,\frac{\log(2+\varepsilon^{2}n)}{\varepsilon^{2}}\right\}\right).$$
- $$\lambda\lVert x-y\rVert_2\leq\lVert Px-Py\rVert_2\leq(1+\varepsilon)\lambda\lVert x-y\rVert_2.$$
- $$\operatorname{Attn}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt d}\right)V\quad\longrightarrow\quad\operatorname{softmax}\!\left(\frac{(QP_Q^{\mathsf T})(KP_K^{\mathsf T})^{\mathsf T}}{\sqrt r}\right)V.$$
- $$P_{ij}\sim\mathcal N(0,1/r),\qquad r=\left\lceil C\min\left\{d,n-1,\frac{\log(2+\varepsilon^2 n)}{\varepsilon^2}\right\}\right\rceil.$$

## Implementation notes

(1) Integration point: modify each attention head's Q/K output dimension based on the token count n in its local window. Keep V at its standard dimension in the first experiment, so only score formation and Q/K activation storage are reduced. For efficient GPU execution, bucket sequence lengths and compile one kernel for each resulting r. (2) Pseudocode: for each batch or window, set r = min(d_h, ceil(C * min(n-1, log(2+epsilon^2*n)/epsilon^2))); compute Q = X W_Q and K = X W_K; create a fixed projection P with entries N(0,1/r), or use a learned projection initialized this way; compute q = Q P.T and k = K P.T; optionally center q and k across tokens; form logits = q @ k.T / sqrt(r); apply the usual mask and softmax; output A @ V. Start with the same P for Q and K, then compare independent P_Q and P_K. (3) The paper determines the dimension schedule and pairwise-distance target. Estimate practical validity using the correlation between original and projected logits, KL divergence between original and projected attention distributions, relative output error, and retrieval accuracy. Tune C and epsilon on a calibration set rather than assuming the hidden theorem constant is one. (4) First experiment: train or evaluate a four-layer Transformer on WikiText-2 and a synthetic copy/retrieval task at sequence lengths 32, 128, 512, and 2048. Compare fixed head dimensions 64 and 32 against the adaptive schedule for epsilon in {0.1,0.2,0.3} and C in {1,2,4}. Measure attention FLOPs, Q/K activation memory, tokens per second, perplexity, attention KL divergence, and retrieval accuracy. The desired signal is lower score-computation cost and memory at long sequences with no more than 1% perplexity degradation. Failure is systematic retrieval loss or high softmax KL despite acceptable pairwise-distance distortion, which would show that the Euclidean JL proxy is insufficient for attention logits.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.