# Assignment Tree Attention

- ID: 53
- Canonical URL: https://synthcore.org/idea/53/assignment-tree-attention
- API JSON: https://synthcore.org/api/idea/53.json
- API Markdown: https://synthcore.org/api/idea/53.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.22308](https://arxiv.org/abs/2608.22308)
- Category: memory
- Solves: memory, speedup, stability
- ML areas: attention, transformer, memory, graph-nn
- Math tags: combinatorics, graph-theory, optimal-transport, geometry
- Ratings: usefulness 6/10; difficulty 7/10; novelty 7/10

## Idea description

Construct a sparse attention support by solving multiple small perturbed assignment problems between query and key embeddings and taking the union of the selected optimal matchings. Use the resulting spanning tree as the only set of cross-token edges, with edge biases determined by empirical assignment frequency.

## Mathematical statement

For query points X={x_i}_{i=1}^{n+1} and key points Y={y_j}_{j=1}^{n}, define a complete bipartite graph with edge costs W_{ij}=W(x_i,y_j), commonly W_{ij}=\|x_i-y_j\|^p. For each modified instance U in a finite family \mathcal{J}, let M_\star(G_U) be a minimum-weight matching and define H_{\mathcal{J}}(G)=\bigcup_{U\in\mathcal{J}}M_\star(G_U). The paper proves that H_{\mathcal{J}} is a spanning tree in its assignment settings. Let \mathcal{C}((a_i b_j))=\{U\in\mathcal{J}:(a_i b_j)\in M_\star(G_U)\} be the perturbation instances selecting edge (a_i,b_j). The paper gives the transport matrix T_{ij}=\frac{1}{n(n+1)}|\mathcal{C}((a_i b_j))| on H and zero elsewhere. Here T transports \mu=\frac{1}{n+1}\sum_{i=1}^{n+1}\delta_{x_i} to 
u=\frac{1}{n}\sum_{j=1}^{n}\delta_{y_j} and is optimal for costs W_{ij}. In the neural adaptation, K sampled perturbed assignments estimate edge inclusion probabilities; these probabilities define sparse attention support and an additive log-frequency bias.

## Key formulas

- $$H_{\mathcal{J}}(G)=\bigcup_{U\in\mathcal{J}}M_{\star}(G_U),$$
- $$T_{ij}=\begin{cases}\dfrac{1}{n(n+1)}\left|\mathcal{C}((a_i\,b_j))\right|,&(a_i\,b_j)\in H,\\0,&(a_i\,b_j)\notin H,\end{cases}$$
- $$\mu(z)=\frac{1}{n+1}\sum_{i=1}^{n+1}\delta(z-x_i),\qquad \nu(z)=\frac{1}{n}\sum_{j=1}^{n}\delta(z-y_j),$$
- $$\widehat T_{ij}=\frac{1}{K}\sum_{k=1}^{K}\mathbf{1}\{(i,j)\in M_k\},\qquad \ell_{ij}=\frac{q_i^\top k_j}{\sqrt d}+\alpha\log(\varepsilon+\widehat T_{ij}).$$

## Implementation notes

Integrate this at the cross-attention support-construction stage, before multiplying values. For each local sequence block, obtain query vectors Q_i and key vectors K_j from one transformer layer, detach them from backpropagation, and define costs W_{ij}=\|Q_i-K_j\|_2^2 or W_{ij}=1-\langle Q_i,K_j\rangle/(\|Q_i\|\|K_j\|). Use n+1 queries and n keys so the rectangular assignment matches the paper's dimensions. Generate K perturbed instances by adding independent Gumbel noise of scale     au to W and masking or modifying one query/key differently in each instance. Solve every instance with a Hungarian or auction algorithm, producing matchings M_k. Form the union of all selected edges and, if numerical ties create cycles, retain a maximum-weight spanning tree using edge score s_{ij}=K^{-1}\sum_k\mathbf{1}[(i,j)\in M_k]. Set non-tree logits to minus infinity and use the paper-inspired sparse logits \ell_{ij}=q_i^    op k_j/\sqrt d+\alpha\log(\varepsilon+\widehat T_{ij}). Apply row-wise softmax over tree neighbors and aggregate values normally. Cache the tree for R training steps and recompute it periodically. The mathematical construction supplies the tree and transport interpretation; edge frequencies are estimated empirically from K assignments. Start with a 4-layer, 128-dimensional transformer on WikiText-2 or CIFAR-10, comparing against dense local attention. Sweep K in {4,8,16} and R in {1,16,64}. Record retained edges, peak activation memory, wall-clock time, validation loss, and gradient variance. A successful result should show O(n) rather than O(n^2) attention edges, at least 1.5x lower attention memory or faster training at matched validation loss, and no significant loss spikes when the support is refreshed.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: Built assignment_tree_attention.py with Gumbel-perturbed rectangular Hungarian assignments, empirical edge frequencies, maximum-frequency bipartite spanning-tree projection, log-frequency attention bias, and a dense-versus-sparse retrieval comparison. The mechanism predictions manifested: frequency RMSE scaled as K^-0.513 versus predicted K^-0.5, union sizes closely matched the occupancy prediction for K=4..512, and alpha=1 increased frequency/probability correlation from 0.671 to 0.783. The sparse tree improved toy retrieval MSE (0.466 versus 0.585) while retaining exactly 30 tree edges versus 240 dense edges, although the rectangular transport claim was not exact because one query is unmatched per assignment.

### Mechanism check

- Verdict: Built assignment_tree_attention.py with Gumbel-perturbed rectangular Hungarian assignments, empirical edge frequencies, maximum-frequency bipartite spanning-tree projection, log-frequency attention bias, and a dense-versus-sparse retrieval comparison. The mechanism predictions manifested: frequency RMSE scaled as K^-0.513 versus predicted K^-0.5, union sizes closely matched the occupancy prediction for K=4..512, and alpha=1 increased frequency/probability correlation from 0.671 to 0.783. The sparse tree improved toy retrieval MSE (0.466 versus 0.585) while retaining exactly 30 tree edges versus 240 dense edges, although the rectangular transport claim was not exact because one query is unmatched per assignment.
- Confidence: 8/10
- Limitations: Only a small synthetic retrieval task was tested; no transformer training, WikiText/CIFAR benchmark, wall-clock or GPU memory measurement, cache-refresh study, gradient-variance measurement, or full K/R sweep was performed. The spanning-tree projection is a practical maximum-frequency heuristic rather than a proof of the paper's union-tree property, and the supplied rectangular transport normalization does not yield exact uniform query marginals when Hungarian matching leaves one query unmatched.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.70955
- Paired wins: 3/8
- Benchmark verdict: no measurable effect

## Artifacts

- [assignment_tree_attention.py](https://synthcore.org/code/695/assignment_tree_attention.py)
- [bench_assignment_tree.py](https://synthcore.org/code/695/bench_assignment_tree.py)
- [bench_report.json](https://synthcore.org/code/695/bench_report.json)
- [report.md](https://synthcore.org/code/695/report.md)
- [report_bench_2026-08-31T222936.md](https://synthcore.org/code/695/report_bench_2026-08-31T222936.md)
- [results.json](https://synthcore.org/code/695/results.json)
- [Download all files as ZIP](https://synthcore.org/download/695)

## Disclaimer

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