Coarse-to-fine active-support transport attention

Integrate this at the cross-attention weight computation, before multiplying attention weights by values. For a query block of size $m$ and key block of size $n$, form $q_i=W_Qh_i$, $k_j=W_Kh_j$, and cost $C_{ij}=-q_i^{\top}k_j/\sqrt d$; use uniform or mask-adjusted marginals $\mu_i$ and $\nu_j$. Build a two-level hierarchy by grouping consecutive tokens, image patches, or graph nodes into coarse groups. Pseudocode: (1) compute coarse group costs; (2) solve the coarse balanced transport problem; (3) lift every positive coarse edge to its fine Cartesian product and use those fine edges as the initial active set $\mathcal N$; (4) initialize fine primal masses and dual/slack variables from the lifted coarse plan; (5) repeatedly compute row and column feasibility residuals and complementarity residuals, form $r$, assemble $A_{\mathcal N}Z^2A_{\mathcal N}^{\top}$, diagonally scale it with $U_{ii}=\sum_{t\in i}z_t^2$ and $V_{jj}=\sum_{t\in j}z_t^2$, solve for $\Delta\lambda$, recover edge updates, and damp the step to retain nonnegative masses; (6) add candidate edges with low $C_{ij}$ or large marginal residual and remove edges whose mass remains below a threshold for several iterations; (7) stop when marginal residuals are below tolerance and use sparse $X$ as attention weights. Estimate $\sigma_{\max}(Q)$ with 2--5 power iterations and log it as a conditioning diagnostic. The Schur system and scaling come from the paper; grouping, edge proposals, damping, and pruning are empirical choices. Start with a 2-layer, 8-head Transformer on CIFAR-10 patch classification or a small language model on WikiText-2, comparing dense softmax attention, fixed top-k attention, and this method at equal hidden size. Measure peak attention memory, wall-clock time, marginal residual, validation accuracy or perplexity, and retained-edge fraction. A successful result is 4--8x fewer active edges and at least 2x lower attention memory with no more than 1% relative accuracy loss, while warm-started fine levels should require fewer Newton iterations than independently solved sparse problems.