Sub-Gaussian Reachability Sparse Attention

Integrate this into the attention-logit and neighborhood-construction path of a graph transformer or a long-context token transformer. Inputs are queries \(q_i\), keys \(k_j\), values \(v_j\), an effective metric \(\widetilde d(i,j)\), propagation time \(t\), target tail budget \(\varepsilon\), and parameters \(\beta,C_3,C_4\). First compute \(r=[t(\log(C_3/\varepsilon)/C_4)^{\beta-1}]^{1/\beta}\). Construct a sparse mask containing only pairs with \(\widetilde d(i,j)\leq r\), using a radius-neighbor index for graph data or a banded/local index for sequences. Compute standard scaled dot-product logits only on this mask, add an optional learnable distance bias \(b_{ij}=-a(\widetilde d(i,j)^\beta/t)^{1/(\beta-1)}\), apply a row-wise softmax, and multiply by values. To emulate several propagation times, use heads with \(t_h\) geometrically spaced and corresponding radii, then concatenate heads. The paper's mathematics supplies the radius formula and the exponential tail shape. The coder must estimate or tune the constants: build the empirical one-step transition matrix \(P\) from detached attention probabilities on a calibration batch, estimate \(\widehat p_{>r}(t)\) from \(P^t\) or repeated propagation, fit \(\log \widehat p_{>r}(t)\) against \((r^\beta/t)^{1/(\beta-1)}\), and use conservative lower estimates of \(C_4\) and upper estimates of \(C_3\). For a first experiment, train a 6-layer graph transformer on ZINC or a 4-layer long-context transformer on synthetic sequences with local-plus-rare-long-range dependencies, comparing dense attention, fixed-window attention, and this adaptive mask at equal hidden size. Measure validation accuracy, attention approximation error versus dense attention, peak memory, and wall-clock time. Success means at least 2x lower attention memory or FLOPs at matched accuracy, with the measured discarded propagation mass below \(\varepsilon\) and no instability from disconnected masks.