Collision-Aware Subset Attention
Implementation & benchmark of arXiv:2606.31716 — Gaussian Belief Propagation for Tracking With Unresolved Measurements
Source paper: Gaussian Belief Propagation for Tracking With Unresolved Measurements arXiv:2606.31716 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace ordinary one-token-to-one-expert or one-token-to-one-attention routing with a local latent subset router: a pooled observation can be explained by a compatible subset of tokens. Pairwise compatibility scores assign probability to subsets, and each token receives the marginal probability that it belongs to the selected subset. This should help when tokens represent overlapping objects, occluded entities, or multiple features that should be processed jointly.
Formulas
Mathematical statement
The paper defines, for a candidate set of objects \({\cal A}\), the pairwise-coupling weight \(V^{{\cal A}}_{{\mathrm{u}},k}(\underline{\bm{x}}^{{\cal A}}_k)\triangleq\prod_{\{i,l\}\subset{{\cal A}}}Q_{{\mathrm{u}},k}(\bm{x}^{(i)}_k,\bm{x}^{(l)}_k)\), where \({\cal A}\) is a subset of entities, \(\{i,l\}\) ranges over unordered pairs in that subset, and \(Q_{{\mathrm{u}},k}\) is a pairwise coupling function. The graph construction interprets connected components as groups, so a subset with strong pairwise couplings represents entities likely to produce an unresolved joint observation. Adapt this into a differentiable local posterior: for token neighborhood \(N_j\), subset \(S\subseteq N_j\), token embeddings \(x_i\), pooled observation \(z_j\), pairwise compatibility \(Q(i,l)=\sigma(q(x_i,x_l))\), and subset decoder \(g(S)\), define \(w_j(S)=\exp(\ell_j(S))\prod_{\{i,l\}\subset S}Q(i,l)\), \(\pi_j(S)=w_j(S)/\sum_{\varnothing\neq T\subseteq N_j}w_j(T)\), and token responsibility \(r_{ij}=\sum_{S\ni i}\pi_j(S)\). Here \(\ell_j(S)\) is a learned joint-observation score, for example \(-\|z_j-g(\sum_{i\in S}a_i x_i)\|^2/(2\tau^2)+\sum_{i\in S}s_i\). The marginal \(r_{ij}\) is the belief-propagation-style quantity used to distribute the joint observation back to its constituent tokens.
Implementation notes
Integrate this into a transformer attention block immediately after the normal token projections and before the value aggregation, or into an MoE router before expert dispatch. Partition tokens into local neighborhoods \(N_j\) of size at most 8, using spatial windows, nearby sequence positions, or k-nearest embedding neighbors; this keeps subset enumeration feasible. For each neighborhood and pooled query observation \(z_j\), compute token scores \(s_i\), pairwise logits \(q_{il}\), compatibilities \(Q(i,l)=\operatorname{sigmoid}(q_{il})\), and a joint decoder \(g\). Enumerate every nonempty subset \(S\subseteq N_j\), form \(a_S=\sum_{i\in S}a_i x_i\), compute \(\ell_j(S)\), add \(\sum_{\{i,l\}\subset S}\log(Q(i,l)+\epsilon)\), and normalize the resulting log-weights with log-sum-exp to obtain \(\pi_j(S)\). Accumulate \(r_{ij}\) by summing probabilities over subsets containing token \(i\). Use \(r_{ij}\) either as attention weights for \(u_j\), as a soft expert-dispatch probability, or as an auxiliary reconstruction path. The mathematical quantities computed exactly are the subset weights, posterior normalization, and inclusion marginals; the functions \(q,g,a,s\) are learned by backpropagation. Start with a synthetic collision dataset: generate 8 token features, randomly merge 1--3 tokens into one observation, and train a small 4-layer transformer against the known merged-group labels and a reconstruction target. Compare standard softmax attention, top-1 MoE routing, and this router at equal FLOPs. Success is higher group-membership F1 and lower reconstruction error at the same neighborhood size, with stable gradients and no increase in global sequence complexity.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: Significantly worse ✗
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 7/10
Built an exact local subset posterior router over 63 nonempty subsets for six-token neighborhoods, including pairwise compatibility products, normalized subset probabilities, and inclusion marginals. The math check passed: posterior mass summed to 1, marginal reconstruction error was 0, and boosting one pair compatibility increased its joint inclusion probability from 0.100 to 0.451. On the fixed synthetic collision task, the subset router outperformed the independent baseline in token F1 (0.987 vs 0.950) and reconstruction MSE (0.0055 vs 0.109), so the claimed signal was observed, although this is not yet evidence of a general transformer or FLOP-normalized win.
- Agent confidence
- 7/10
- Baseline
- F1=0.9499, reconstruction MSE=0.10925
- Idea
- F1=0.9874, reconstruction MSE=0.005536
Limitations: Tested only one fixed seed on a small synthetic six-token collision dataset and 500 optimization steps. The comparison was equal-step rather than equal-FLOP; subset enumeration uses 63 subsets and therefore has substantially higher local compute than the independent router. No transformer/MoE integration, multiple-seed confidence interval, scaling study, or real-data evaluation was performed.
How to run: python3 collision_subset_attention.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2606.31716)
Significantly WORSE than the tuned baseline (+42.8%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.1665
- Idea mean
- 0.2377
- Effect (Δ)
- 0.0712 (+42.8%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±14.7%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
The benchmark used the standard built-in sequence forecast track, reduced to 1,200 training and 400 test examples with 12 epochs for runtime. It did not test vision, dynamics, longer neighborhoods, larger subset sizes, equalized parameter/FLOP counts, or the original synthetic collision-label task in this round; the exact subset router also has higher local computational and parameter overhead than ordinary multi-head attention.
How to run: python3 bench_collision_subset.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Implemented Collision-Aware Subset Attention as a local exact subset-posterior mechanism over four-token neighborhoods inside a matched sequence transformer, using the canonical benchmark training path. The mechanism prediction held quantitatively: boosting pair compatibility increased trained-model joint inclusion mass from 0.0397 to 0.2282. However, test MSE was significantly worse than the standard attention baseline (0.23773 vs 0.16653, paired delta +0.07121, permutation p=0.0081), so there is no demonstrated transfer win on the standard sequence benchmark.
- Baseline
- test MSE mean=0.166526, std=0.019344; best lr=0.001; per-seed=[0.189505,0.150417,0.141310,0.178840,0.151882,0.159856,0.200021,0.160381]
- Idea
- test MSE mean=0.237731, std=0.022670; best lr=0.001; per-seed=[0.242086,0.213451,0.242339,0.202913,0.272864,0.231953,0.267943,0.228303]
- Paired comparison
- delta_mean=+0.071205 (idea-baseline; lower is better), idea_wins=0/8, p_value=0.0081, verdict='idea worse (significant)'
Benchmark result
Significantly WORSE than the tuned baseline (+42.8%, p=0.0081, wins 0 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.1665
- Idea mean
- 0.2377
- Effect (Δ)
- 0.0712 (+42.8%; negative = idea better)
- Wins
- 0 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±14.7%
- Mechanism
- Confirmed ✓
- Practical effect
- Hurts
- Baseline tuning
- swept over 3 configs
The benchmark used the standard built-in sequence forecast track, reduced to 1,200 training and 400 test examples with 12 epochs for runtime. It did not test vision, dynamics, longer neighborhoods, larger subset sizes, equalized parameter/FLOP counts, or the original synthetic collision-label task in this round; the exact subset router also has higher local computational and parameter overhead than ordinary multi-head attention.
How to run: python3 bench_collision_subset.py
Artifacts
- bench_collision_subset.py 4.7 KB View Raw
- bench_report.json 3.8 KB View Raw
- collision_subset_attention.py 5.1 KB View Raw
- report.md 1.4 KB View
- report_bench_2026-09-03T223727.md 5.2 KB View
- results.json 0.5 KB View Raw