Weighted Resolvent-Equivariant Attention

1. Integration point: apply the method to the row-stochastic attention matrix P produced by one transformer attention head, or to the normalized adjacency/message-passing matrix in a graph neural network. Choose a known index involution \sigma, such as reversing sequence positions, swapping two mirrored graph halves, or pairing two sensor types. Store positive pair weights \kappa(i); if no prior weights are known, initialize all \kappa(i)=1 and optionally learn log \kappa(i) with a small L2 penalty. Construct the sparse matrix G=\Delta^2\mathcal{R}, whose only nonzero entry in row i is G[i,\sigma(i)]=\kappa(i). 2. Pseudocode: compute logits L=QK^T/\sqrt{d}; compute P=softmax(L,dim=-1); compute GP by permuting the columns of P according to \sigma and multiplying row i by \kappa(i); compute PG by permuting the rows according to \sigma and multiplying the appropriate columns by \kappa; add lambda_comm * ||GP-PG||_F^2/(||G||_F^2||P||_F^2+1e-8) to the task loss; then compute Y=PV. For a graph layer, replace P with the normalized message matrix. 3. The paper supplies the exact commutator criterion and its consequence for all powers and resolvents. Estimate the empirical penalty directly from the current finite-batch P; do not explicitly invert I-(1-gamma)P during training. For validation, compute R_gamma by solving (I-(1-gamma)P)X=I for several gamma values and measure ||GR_gamma-R_gamma G||_F. 4. Cheap experiment: train a 4-layer, 8-head transformer on a synthetic sequence task whose target is invariant under reversal but whose two mirrored halves have unequal noise scales. Compare baseline attention, reversal augmentation, and the weighted-commutator model at matched parameters and FLOPs. Track validation accuracy, equivariance error ||Gf(x)-f(Gx)||, commutator error, and loss variance across seeds. Success means higher accuracy or lower sample complexity at equal compute, at least a 10x reduction in equivariance error, and no increase in optimization instability. Test a second small graph benchmark with a known automorphism and measure whether the benefit remains when \kappa is learned.