Critical power-law sparse attention

1. Integration point: replace the dense \(QK^\top\) attention matrix in one or more Transformer blocks by a precomputed sparse edge list. Assume tokens have an order, such as text positions, image patches in raster order, or time steps. Always include a local window of radius \(r\), then add long-range edges using the power-law rule. For large sequences, do not enumerate all Bernoulli trials: draw \(m\) distances from the discrete distribution \(q(d)=d^{-(1+\sigma)}/\sum_{u=1}^{\lfloor n/2\rfloor}u^{-(1+\sigma)}\), choose a random direction, and deduplicate edges; this is a fixed-degree approximation to the paper's mask. 2. Pseudocode: inputs \(X\in\mathbb{R}^{n\times d}\), local radius \(r\), long-range budget \(m\), exponent \(\sigma\). Build \(N(i)\leftarrow\{i-r,\ldots,i+r\}\); repeat \(m\) times: sample \(d\sim q\), sample sign \(s\in\{-1,+1\}\), set \(j=(i+sd)\bmod n\), and add \(j\) to \(N(i)\). Compute only logits \(Q_iK_j^\top/\sqrt{d_k}\) for \(j\in N(i)\), apply the masked softmax and weighted value sum shown above. 3. Compute from the paper: the distance law, finite expected-degree scaling, and recommended transition range \(\sigma\approx0.8\text{--}0.85\); estimate empirically: validation loss, gradient propagation across positions, graph diameter, and the spectrum of the normalized mask Laplacian or adjacency. 4. First experiment: train a 6-layer small Transformer on WikiText-103 or enwik8 with sequence length 1024, comparing dense attention, local-window attention, and this mask at matched heads and hidden size. Sweep \(\sigma\in\{0.5,0.7,0.8,0.85,1.0,1.2\}\) and long-range budgets giving 32, 64, and 128 total neighbors. Measure training loss versus attended edges/FLOPs, validation perplexity, throughput, and attention-gradient influence between tokens separated by at least 256 positions. Success is lower perplexity than local attention at the same edge budget, with at least 2x lower attention memory than dense attention and no degradation relative to dense attention at fixed training compute.