Local-Tail Adaptive Attention

(1) Integrate this into the inference and training attention kernel, immediately after computing query-key logits and before the softmax. Start with causal self-attention. For each query position i, define a candidate ordering of keys: local keys first, then remaining keys sorted by descending a_ij, where a_ij=|logit_ij| times ||v_j||; detach a_ij from gradients to avoid introducing a difficult discrete optimization problem. Select the smallest prefix S_i whose estimated tail T_i=sum_{j outside S_i} a_ij is at most epsilon, while enforcing a minimum window w_min and a maximum budget k_max. Compute exact softmax attention only over S_i. (2) Pseudocode: compute Q,K,V; logits=QK^T/sqrt(d_k); magnitudes=abs(logits)*rowwise_norm(V); order local keys and then descending magnitudes; for each i, tail=sum(magnitudes[i]); retained=[]; scan ordered keys, append j and subtract magnitudes[i,j] until tail<=epsilon; output_i=softmax(logits[i,retained]) @ V[retained]. During training, use a straight-through or soft mask for the first few epochs, then hard selection. (3) The paper supplies the local-tail construction and the power-law prediction, but not the neural constants: estimate epsilon from a target output-drift calibration, and estimate an empirical decay exponent p by fitting log a_ij versus log distance. Also record maximum and percentile tail masses, not merely their mean. (4) First experiment: compare dense attention, fixed-window attention, top-k attention, and this adaptive rule on a 2-layer Transformer language model trained on WikiText-2 or a synthetic long-range-copy task. Sweep epsilon and k_max, measuring validation perplexity, attention FLOPs, peak memory, and the norm of the difference between dense and sparse layer outputs. Success means at least 2x lower attention FLOPs or memory at matched perplexity, with output error increasing monotonically with measured T_i and without unstable loss spikes.