Wide-tree invariant alignment layer
Implementation & benchmark of arXiv:2607.08538 — High-Dimensional Procrustes Matching via Tree Counts
Source paper: High-Dimensional Procrustes Matching via Tree Counts arXiv:2607.08538 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace a purely pairwise embedding similarity used for set alignment with a sum of rooted-tree contraction scores. Each tree feature aggregates products of several coordinate-level interactions and can preserve correspondence information under an unknown orthogonal transformation, allowing matching from moderate correlation rather than nearly identical embeddings. Use the resulting score matrix for Hungarian matching, contrastive loss, or a differentiable Sinkhorn assignment.
Formulas
Mathematical statement
The paper represents a rooted tree T by a monomial X^T obtained by assigning data-point indices to sample vertices, coordinate indices in [d] to coordinate vertices, and multiplying the corresponding matrix entries along all edges; aut(T) is the number of automorphisms of T. For a candidate pair of roots i and j, the weighted score is a sum over tree shapes, with each tree contribution multiplied by aut(T). The paper's concentration requirement is Var(score_ij)/(E[score_ii])^2=o(1/n) for false pairs j != i and o(1) for true pairs j=i, so the score gap must dominate both false-pair fluctuations and the number of competing pairs. The false-pair variance expands into four-tree joint moments, as in equation (2.5), which is the mechanism by which tree overlap patterns control noise. The algorithmic threshold is rho^2 > sqrt(alpha), where rho is the coordinate-wise correlation between matched Gaussian vectors and alpha approximately 0.338 is Otter's tree-counting constant. In the adapted layer, X and Y are two n-by-d embedding matrices, T is a selected rooted bipartite tree with D edges, F_T(X,i) is its normalized rooted contraction at sample i, and S_ij=sum_T w_T F_T(X,i)F_T(Y,j) is the alignment score. The practical approximation uses a finite tree family and randomized coordinate contractions; the mathematical hypothesis is that higher-order tree aggregation increases the true/false score ratio.
Implementation notes
Integrate the layer immediately before set matching: inputs are two batches X,Y in R^{n x d}, optionally after separate MLP or transformer encoders. First normalize each embedding dimension across the batch and optionally whiten each set with a cheap diagonal variance estimate; do not learn an arbitrary post-matching rotation, since the point is to obtain robustness to such rotations. Select 4-8 rooted bipartite trees with 2-5 edges, emphasizing wide trees with several branches. For each tree T, compute F_T(X,i) and F_T(Y,j) by dynamic programming over the tree: leaves return the corresponding embedding coordinate, coordinate-vertex messages contract over a randomly sampled subset of d coordinates, and sample-vertex messages aggregate over a sampled subset of n candidate neighbors. Use K independent random coordinate/index samples and average them to estimate each contraction, with normalization Z_T chosen so F_T has unit batch variance. Form S = sum_T w_T F_T(X)F_T(Y)^T, where w_T can initially be aut(T)/sqrt(Var_batch(F_T(X))Var_batch(F_T(Y))). Add temperature-scaled S to the ordinary cosine-similarity logits, or apply Sinkhorn to S and train with permutation cross-entropy. Pseudocode: compute tree features for X and Y; for each T, estimate F_T by recursive sampled contractions; accumulate S += w_T outer(F_T(X),F_T(Y)); solve P=Sinkhorn(S/tau); optimize -sum_i log P_{i,pi(i)}. The paper supplies the tree weighting and variance-gap rationale; estimate feature variances, effective correlation rho, and contraction Monte Carlo error empirically. Start with n=128,d=256 synthetic pairs generated as Y=P X Q + Gaussian noise, compare cosine matching, orthogonal-Procrustes-plus-Hungarian, and the tree layer at equal wall-clock budget. Success is higher exact-match accuracy at fixed rho, especially when Q is random and rho is between 0.5 and 0.8; also measure whether 4-8 trees improve accuracy without more than 2x the baseline matching cost.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
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 6/10
Built a compact rooted-tree contraction matcher using Gram-matrix path and wide-branch features with Hungarian assignment. The orthogonal-invariance check passed at relative error 1.9e-15. For n=128,d=256, tree matching exceeded cosine accuracy at rho=.55, .70, and .85, but absolute recovery stayed low and computation was slower, so this is promising evidence rather than a demonstrated practical win.
- Agent confidence
- 6/10
- Baseline
- Cosine Hungarian accuracy: 0.94%, 0.31%, 0.78% at rho=0.55, 0.70, 0.85; chance is 0.78%.
- Idea
- Tree-contraction Hungarian accuracy: 1.09%, 2.03%, 4.06% at rho=0.55, 0.70, 0.85; orthogonal-invariance relative max error 1.9e-15. Runtime was approximately 1.6–17 ms versus 0.9–3.5 ms for cosine in the reported run.
Limitations: This is a compact Gram-based approximation rather than the paper's full wide rooted bipartite tree family. Randomized contractions, learned tree weights, Sinkhorn training, Procrustes control, larger repetitions, and realistic embeddings were not tested; the absolute accuracy advantage was small.
How to run: python3 tree_alignment.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.08538)
Significantly WORSE than the tuned baseline (+20.0%, p=0.0163, wins 1 / 8 paired seeds).
- Benchmark
- Unordered_pointset_denoising
- Model
- mlp_med
- Paired seeds
- 8
- Baseline mean
- 0.035
- Idea mean
- 0.042
- Effect (Δ)
- 0.007 (+20.0%; negative = idea better)
- Wins
- 1 / 8 paired seeds
- p-value
- 0.0163 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 0 configs
Only the unordered point-set denoising custom track and mlp_med architecture were tested. The implementation uses fixed low-order centered moments rather than randomized recursive bipartite contractions, and does not test Hungarian or Sinkhorn matching, larger tree families, alternative architectures, or larger datasets. A local training loop was necessary because the generic trainer assumes scalar regression targets while this track has six-dimensional targets.
How to run: python3 stage2_bench.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 the wide-tree invariant alignment mechanism as a tree-moment consistency regularizer in an end-to-end MLP on the structurally matched unordered point-set denoising track. The orthogonal-invariance sanity check passed with relative error 2.94e-7, but the full benchmark found the idea significantly worse: MSE rose from 0.03499 to 0.04198, paired delta +0.00699, permutation p=0.0163. The trained-model mechanism signature was not confirmed because mean tree-feature correlation decreased from 0.9267 to 0.9220.
- Baseline
- Best baseline lr=0.006, weight_decay=0.0; full 8-seed test MSE 0.0349912 ± 0.0021176.
- Idea
- Best idea config lr=0.003, weight_decay=0.0, tree_lambda=0.05; full 8-seed test MSE 0.0419774 ± 0.0049313; paired delta +0.0069862, 1/8 wins, permutation p=0.0163.
Benchmark result
Significantly WORSE than the tuned baseline (+20.0%, p=0.0163, wins 1 / 8 paired seeds).
- Benchmark
- Unordered_pointset_denoising
- Model
- mlp_med
- Paired seeds
- 8
- Baseline mean
- 0.035
- Idea mean
- 0.042
- Effect (Δ)
- 0.007 (+20.0%; negative = idea better)
- Wins
- 1 / 8 paired seeds
- p-value
- 0.0163 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Hurts
- Baseline tuning
- swept over 0 configs
Only the unordered point-set denoising custom track and mlp_med architecture were tested. The implementation uses fixed low-order centered moments rather than randomized recursive bipartite contractions, and does not test Hungarian or Sinkhorn matching, larger tree families, alternative architectures, or larger datasets. A local training loop was necessary because the generic trainer assumes scalar regression targets while this track has six-dimensional targets.
How to run: python3 stage2_bench.py
Artifacts
- bench_report.json 5.6 KB View Raw
- report.md 1.3 KB View
- report_bench_2026-09-04T194523.md 3.5 KB View
- results.json 0.8 KB View Raw
- stage2_bench.py 5.1 KB View Raw
- tree_alignment.py 3.2 KB View Raw