# Coarse-to-fine active-support transport attention

- ID: 2163
- Canonical URL: https://synthcore.org/idea/2163/coarse-to-fine-active-support-transport-attention
- API JSON: https://synthcore.org/api/idea/2163.json
- API Markdown: https://synthcore.org/api/idea/2163.md
- Verification status: mechanism_failed
- Source: [arXiv:2608.12060](https://arxiv.org/abs/2608.12060)
- Category: memory
- Solves: memory, speedup, accuracy
- ML areas: attention, memory, inference-speedup
- Math tags: optimal-transport, optimization, linear-algebra, numerical-analysis
- Ratings: usefulness 7/10; difficulty 7/10; novelty 6/10

## Idea description

Replace dense cross-attention weights with a balanced transport plan whose nonzero query-key edges are maintained by a multiscale active-set procedure. Solve the coarse token-group problem first, lift its support to the fine token grid, add only edges indicated by local cost or marginal residuals, and warm-start the fine problem from the lifted plan. This should provide a principled sparse attention pattern rather than fixing a global top-k pattern before seeing the transport solution.

## Mathematical statement

The paper formulates balanced transport as the nonnegative linear program $\min c^{\top}x$ subject to $\bar A x=\bar b$ and $x\ge 0$, equivalently $\min\langle C,X\rangle$ subject to $X\mathbf 1_n=\boldsymbol\mu$ and $X^{\top}\mathbf 1_m=\boldsymbol\nu$. Here $X_{ij}$ is the attention mass from query/source token $u_i$ to key/target token $v_j$, $C_{ij}$ is a learned attention cost, and $\boldsymbol\mu,\boldsymbol\nu$ are prescribed row and column masses. Let $\mathcal N$ be the current active edge set, $A_{\mathcal N}$ the equality-constraint matrix restricted to those edges, and $Z=\operatorname{diag}(z)$ the diagonal matrix of positive primal-dual slack variables on active edges. The primal-dual Newton direction is obtained from the reduced Schur system $A_{\mathcal N}Z^2A_{\mathcal N}^{\top}\Delta\lambda=r$, with $r$ the current residual right-hand side. The paper diagonally scales this system as $N=\begin{pmatrix}I&Q\\Q^{\top}&I\end{pmatrix}$, where $Q=U^{-1/2}DV^{-1/2}$, $U$ and $V$ are diagonal row and column sums of $z_t^2$, and obtains $\kappa(N)=(1+\sigma_{\max}(Q))/(1-\sigma_{\max}(Q))$ together with an explicit upper bound involving the largest diagonal sum and smallest active slack. The adaptation uses these equations to solve only the active attention edges and to expand supports whose residuals are too large.

## Key formulas

- $$\min_{X\ge 0}\ \langle C,X\rangle\quad\text{s.t.}\quad X\mathbf 1_n=\boldsymbol\mu,\qquad X^{\top}\mathbf 1_m=\boldsymbol\nu.$$
- $$A_{\mathcal N}Z^{2}A_{\mathcal N}^{\top}\Delta\lambda=r,$$
- $$N=\begin{pmatrix}U^{-1/2}&0\\0&V^{-1/2}\end{pmatrix}\begin{pmatrix}U&D\\D^{\top}&V\end{pmatrix}\begin{pmatrix}U^{-1/2}&0\\0&V^{-1/2}\end{pmatrix}=\begin{pmatrix}I&Q\\Q^{\top}&I\end{pmatrix},\quad Q=U^{-1/2}DV^{-1/2},$$
- $$\kappa(N)=\frac{1+\sigma_{\max}(Q)}{1-\sigma_{\max}(Q)},\qquad \kappa(N)\le 2(m+n-1)^2\frac{\max\{U_{\max},V_{\max}\}}{\min\{z_t^2\mid t\in\widehat{\mathcal N}\}}.$$

## Implementation notes

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.

## Verification

- Status: mechanism_failed
- Mechanism evidence: yes
- Mechanism confirmed: no
- Verdict: Built a readable NumPy/SciPy MVP for balanced transport, coarse-support lifting, active-edge expansion, Schur conditioning diagnostics, and a fixed-top-k comparison. Global slack scaling preserved kappa to 1.35e-13 and the identity kappa=(1+sigma)/(1-sigma) had zero measured error, but the hypothesized coarse-support transition at separation >0.2 was falsified: the lifted support was already exact at separation 0.1. The idea matched dense transport exactly on the representative case, but used 50% of all edges and did not improve the top-k objective, so there is no demonstrated practical win.

### Mechanism check

- Verdict: Built a readable NumPy/SciPy MVP for balanced transport, coarse-support lifting, active-edge expansion, Schur conditioning diagnostics, and a fixed-top-k comparison. Global slack scaling preserved kappa to 1.35e-13 and the identity kappa=(1+sigma)/(1-sigma) had zero measured error, but the hypothesized coarse-support transition at separation >0.2 was falsified: the lifted support was already exact at separation 0.1. The idea matched dense transport exactly on the representative case, but used 50% of all edges and did not improve the top-k objective, so there is no demonstrated practical win.
- Confidence: 9/10
- Limitations: Only a small synthetic balanced-transport experiment was tested; no Transformer, CIFAR/WikiText task, GPU implementation, Newton primal-dual solver, wall-clock attention benchmark, memory measurement, nonuniform marginals, masking, or accuracy/perplexity evaluation was included. The active-set expansion uses a simple cheapest-edge proposal rather than full reduced-cost and residual-based Newton updates.

## Artifacts

- [report.md](https://synthcore.org/code/867/report.md)
- [results.json](https://synthcore.org/code/867/results.json)
- [transport_active_support.py](https://synthcore.org/code/867/transport_active_support.py)
- [Download all files as ZIP](https://synthcore.org/download/867)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.