Shifted-Mask Defect Regularization

Integrate the method at the binary sparse-mask or router-assignment level rather than in dense attention logits. For a self-attention layer with n tokens, obtain dense scores S, select the k largest neighbors per row, and symmetrize the resulting mask M in {0,1}^{n x n}; alternatively apply the same procedure to a token-expert coactivation graph. Compute d_i=sum_j M_ij, m=sum_{i,j}M_ij/2, sort the degrees, and calculate c_r=sum_i 1[d_i>=r] for r=1,...,max_i d_i. Add lambda_shift times D/(n^2+epsilon) to the task loss, where D=sum_r c_r^2-sum_i d_i^2-2m. Because top-k selection and threshold counts are discrete, use a straight-through estimator: forward with the hard mask, backward through M_ij=sigmoid((S_ij-tau)/T), annealing T. During training, recompute the defect every few steps to reduce overhead. The paper supplies the exact defect identity and exchange-distance guarantee; no eigenvalue estimation is needed in the graph case. For post-training structured sparsity, order vertices by degree and repeatedly apply failed elementary shifts, replacing an eligible edge containing j by the corresponding edge containing i when absent while preserving m edges. First test a 2-4 layer Transformer on WikiText-2 or CIFAR-10 with fixed average attention degree, comparing ordinary top-k sparse attention against defect-regularized masks. Record validation perplexity or accuracy at equal FLOPs, defect value, number of unique sparse-index patterns, routing imbalance, and mask/index construction time. A positive result is lower indexing cost or better load balance at equal accuracy, or faster and less oscillatory loss descent at equal sparsity.