Simultaneous-Label Sparse Attention

Integrate this into the attention-logit and aggregation stage of a Transformer or graph-attention layer. Let X\in\mathbb{R}^{n imes D} be token features, let H provide neighbor lists N_H(u), and add a routing head R=\mathrm{MLP}(X)\in\mathbb{R}^{n imes d}. During an initial warm-up, use p_{u,i}=\sigma(R_{u,i}/ au) and calculate the soft shared-label score q_{uv}=1-\prod_i(1-p_{u,i}p_{v,i}) only for (u,v)\in E(H). Compute logits s_{uv}=(Q_uK_v^ op/\sqrt{D})+\log(q_{uv}+\epsilon) on those edges, apply a segmented softmax over v\in N_H(u), and aggregate V_v. Anneal au from 1.0 to 0.1, then convert p to top-k binary memberships z_{u,i} and use the exact mask A_{uv} from the simultaneous representation at inference. To avoid testing every pair, maintain inverted label buckets B_i=\{u:z_{u,i}=1\}; intersect each bucket with the base-neighbor list, deduplicate candidate edges, and run attention only there. Add \lambda\sum_{u,i}p_{u,i} to discourage excessive labels and a load-balancing term that penalizes unequal bucket sizes. Pseudocode: compute R; compute p; for each u, candidates = union_i(B_i intersect N_H(u)); for each candidate v, calculate q or A; masked-softmax attention; backpropagate through q during training; periodically rebuild B_i. The mathematical quantities computed exactly are the intersection mask and shared-label predicate; the empirical choices are H, d, label sparsity k, temperature schedule, and regularization weights. First test on a 125M-parameter Transformer on WikiText-103 or C4, using a local-window H with occasional global landmark edges. Compare dense attention and fixed local attention at equal sequence length and hidden size. Measure peak attention memory, attention FLOPs, tokens/sec, validation perplexity, average candidate degree, and the fraction of disconnected tokens. A successful result should show substantially lower memory and higher throughput than dense attention with no more than 1-2% perplexity degradation, or better perplexity than fixed local attention at the same candidate-edge budget.