Support-Scenario Attention Pruning

Modify the attention mask after computing scaled dot-product logits QK^T/sqrt(d), before softmax. Partition a sequence into overlapping windows, for example windows [i,i+w) and [i+s,i+s+w), and treat each query row or query block as a context c. Define p_c(j)=softmax(logits_c)_j and an initial support F(c)={j:p_c(j)>epsilon}; use epsilon between 1e-3 and 1e-2. Construct restriction maps by deleting queries outside an overlap and keys outside the corresponding local key set. Apply a greedy support-cleaning loop: remove the lowest-weight edge (c,j) if its removal leaves every context nonempty, every overlap child with at least one parent continuation, and no query with zero admissible keys. Next, build a graph over overlapping windows; mark a link bijective when the surviving restricted supports have equal cardinality and the restriction is one-to-one and onto. Keep the largest strongly connected component, or add penalties instead of hard removal during the first training phase. Pseudocode: logits -> dense softmax -> threshold -> validity-preserving edge deletions -> sparse softmax over the mask -> FlashAttention-style block dispatch. The mathematical checks are exact on the finite mask; epsilon selection, component tradeoffs, and whether to retain one or several components are empirical. First test a 6-layer, 8-head Transformer on ListOps or WikiText-2 with sequence length 512, comparing dense attention and local top-k attention at equal quality targets. Measure attention FLOPs, peak memory, retained edges, perplexity, and whether the structured mask reaches 2x fewer edges without the validation loss increase seen from independent thresholding.