Token-Count-Aware Attention Projection

(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.