Budgeted Random Tree Attention
Implementation & benchmark of arXiv:2608.18967 — Exact random covers of metric trees: balanced rounding, duality, and sharp thresholds
Source paper: Exact random covers of metric trees: balanced rounding, duality, and sharp thresholds arXiv:2608.18967 ⓘ · analyzed Sep 1, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace dense attention on tree-structured inputs with stochastic attention neighborhoods formed by metric balls of sampled radii. Use the paper's exact trimming rule to ensure that every sampled cover remains valid while its total radius budget is bounded, then average predictions over several independent covers during training. This creates sparse, globally covering attention masks with an explicit locality-versus-coverage control.
Formulas
Mathematical statement
Let T be a finite metric tree with total edge length L, let d_T be its tree metric, and let B_T(v,s)={u in T:d_T(u,v)<=s} be the radius-s ball centered at v. A finite radius tuple is 0-good when its balls can be recentered to cover T and its retained radii satisfy the budget condition sum_i s_i<=L. The paper's trimming rule sorts radii s_1>=s_2>=...>=s_m; if their sum exceeds L, it chooses the smallest k satisfying 2s_1+sum_{i=2}^k s_i>=L, keeps (s_1,...,s_k), and rec enters the retained balls to obtain a cover. The proof also uses the arbitrary-budget dual criterion: for a finite positive Borel measure beta on [0,R], beta dominates the expected radius measure of a random 0-good cover if and only if sigma(T)<=integral_[0,R] max_{v in T} sigma(B_T(v,s)) d beta(s) for every finite positive measure sigma on T; finite atomic sigma suffice. In the neural adaptation, beta is the learned distribution of attention radii, sigma is an empirical nonnegative mass over tokens, and the right-hand side measures how much token mass can be covered by the best local ball at each radius.
Implementation notes
Integrate this into a tree, AST, or graph Transformer before the attention softmax. Inputs are a rooted tree, node features x, total tree length L, and a maximum radius R. First compute pairwise tree distances or an LCA-based distance routine. For each training example, sample m candidate radii s_i from beta_theta, sort them descending, and apply the exact paper rule: if sum(s_i)<=L retain all radii; otherwise find the smallest k with 2*s[0]+sum(s[1:k])>=L and retain s[:k]. For each retained radius, choose a center using a greedy farthest-uncovered-node procedure: repeatedly select the uncovered node maximizing distance from the currently covered set, then mark all nodes within radius s_i as covered. If the greedy procedure fails to cover the tree, increase the number of candidate balls or reject the sample; this failure rate is an explicit diagnostic. Construct a sparse attention mask by allowing token q to attend only to tokens in the union of retained balls, optionally adding the tree root as a global token. Run standard scaled dot-product attention on this mask and average over two to four independently sampled covers during training; use one deterministic seed or the union of two covers at inference. Estimate the dual coverage certificate empirically using atomic masses sigma_j on nodes: compute C_beta(s)=max_v sum_{j:d_T(v,j)<=s} sigma_j, approximate the integral by Monte Carlo samples, and penalize max(0, sigma(T)-mean(C_beta(s))). The paper's budget and trimming steps are exact; center placement, beta_theta parameterization, and the dual penalty are engineering approximations. Start with a 2-layer Transformer on 1k-node synthetic trees and molecular or AST node classification, comparing dense attention, fixed-radius local attention, and random untrimmed balls at equal attention FLOPs. Success means full-node coverage near 100%, lower peak attention memory, and equal-or-better validation accuracy at the same number of attended edges; additionally test whether the trimmed model has fewer catastrophic examples where distant dependencies are unreachable.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.