△ Mechanism confirmed, baseline not beaten 2026

Assignment Tree Attention

Implementation & benchmark of arXiv:2608.22308 — Spanning trees in the Assignment Problem: two theorems and two conjectures

Usefulness6/10

Difficulty7/10

Novelty7/10

Memory Memory Speedup Stability Attention Transformer Memory Graph nn Combinatorics Graph theory Optimal transport Geometry

Source paper: Spanning trees in the Assignment Problem: two theorems and two conjectures arXiv:2608.22308 ⓘ · analyzed Aug 29, 2026

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

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.

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}).$$

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.

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

Mechanism confirmed, baseline not beaten

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: No measurable effect

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

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.

Agent confidence: 8/10

Baseline: Dense attention: retrieval MSE 0.584946; 240 possible cross-token edges.
Idea: Tree sparse attention K=16: retrieval MSE 0.465771; exactly 30 edges; RMSE scaling slope -0.51284 versus predicted -0.5; frequency correlation alpha=0/1: 0.67094/0.78340; maximum row transport marginal error 0.02148 and key error 0.00417.

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.

How to run: python3 assignment_tree_attention.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.22308) Sep 2, 2026

Ahead of the baseline (−0.7%) but not statistically significant (p=0.7096, wins 3 / 8 paired seeds).

Benchmark: Sequence modeling
Model: transformer_tiny
Paired seeds: 8
Baseline mean: 0.958
Idea mean: 0.9509
Effect (Δ): -0.0072 (−0.7%; negative = idea better)
Wins: 3 / 8 paired seeds
p-value: 0.7096 (permutation test, 20 000 shuffles)
Smallest detectable effect: n/a
Mechanism: Confirmed ✓
Practical effect: No effect
Baseline tuning: swept over 0 configs

Limitations:

Only the sequence track was tested, with 400 training samples, 200 test samples, two epochs, K=4, and a cached batch-mean support for tractability. Wall-clock speed, activation memory, refresh interval R, and K/R sweeps were not measured; the implementation uses a practical greedy matching approximation rather than exact Hungarian assignments.

How to run: python3 bench_assignment_tree.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 31, 2026 latest Failed ✗

Agent confidence: 9/10

Built Assignment Tree Attention for the structurally matched sequence forecast track with dense and sparse systems sharing data, optimizer, architecture family, and training budget. Mean test MSE was 0.950891 for the idea versus 0.958047 for baseline; paired delta was -0.007156 but permutation p=0.70955, so there was no significant effect. The trained-model signature confirmed 63 observed tree edges versus 63 predicted for a 32-token support.

Baseline: 8-seed mean MSE 0.9580474123, std 0.0756922766; best lr=0.001, epochs=2; sweep means: lr=0.001 -> 0.9283310622, lr=0.003 -> 1.0033762753, lr=0.006 -> 1.0917343497.
Idea: 8-seed mean MSE 0.9508912265, std 0.0570249099; best lr=0.001, K=4, tau=0.7, alpha=0.5, epochs=2; paired delta=-0.0071561858, 3/8 wins, p=0.70955; 63 tree edges observed and predicted.

Benchmark result