Structure-Adaptive Tree Field Integrators

Millend Roy; Soham Samal; Ivan Zelich; Krzysztof Marcin Choromanski

cs.LGSep 27, 2026

Structure-Adaptive Tree Field Integrators

Authors: Millend Roy, Soham Samal, Ivan Zelich, Krzysztof Marcin Choromanski

Organizations: Columbia University New York, NY 10027, USA

Abstract

We present a new class of near-linear algorithms for efficiently integrating general tensor fields defined on trees with distance dependent kernels, the Structure-Adaptive Tree Field Integrators (STAD-TFIs). STAD-TFIs exploit the tree's underlying structure through decompositions built around path backbones and single vertex separators, and use two-dimensional fast Fourier transforms to compute interactions jointly. By exploiting this structural information, STAD-TFIs achieve more computationally efficient integration than their regular efficient tree field integrators (TFI) counterparts. We provide a detailed theoretical analysis of our proposed approach and complement it with an exhaustive empirical evaluation, ranging from speed tests on synthetic trees, through accelerated Sinkhorn-based relaxations of the Optimal Transport algorithms on real meshes, to Topological Attention Transformers for vision tasks. To the best of our knowledge, we provide some of the first results showing that efficient to compute and accurate relaxations of the geodesic Sinkhorn-based solutions of the Optimal Transport problem can be derived by applying fast TFI methods.

Figures & tables

Figure 2: Integration time as a function of number of vertices for kernel

f(d)=1/(1+d)

. Diameter-1D versus Diameter-2D isolates the convolution strategy; Diameter-2D versus Adaptive-2D measures the effect of separator selection. See Table 5 in Appendix C.3 for detailed numbers.

Tree Approximation	Rel. OT Error	Plan TV	Kernel Error
Minimum Spanning Tree (MST)	$0.028\pm 0.027$	$0.329\pm 0.155$	$0.680\pm 0.135$
Random Shortest-Path Tree (Random SPT)	$0.033\pm 0.047$	$0.183\pm 0.082$	$0.677\pm 0.050$
Diameter/Backbone Spanning Tree	$0.021\pm 0.020$	$\mathbf{0.124\pm 0.063}$	$0.599\pm 0.155$
Alon et al. (1995) (AKPW) Tree	$\mathbf{0.018\pm 0.025}$	$0.236\pm 0.101$	$\mathbf{0.573\pm 0.100}$
Abraham and Neiman (2012) Tree	$0.026\pm 0.048$	$0.211\pm 0.100$	$0.677\pm 0.072$

Table 1: Tree approximations for geodesic Sinkhorn on Thingi10K. Results are reported as mean

\pm

standard deviation across the evaluated meshes. Lower is better for all metrics.

Figure 3: Runtime comparison as a function of the number of vertices. Speedup Numbers are in Table 6 in Appendix D.4 . StAd-TFI (SpclK-Diameter) is

\approx 3.35\times

its FTFI counterpart.

	CIFAR-100		Tiny-ImageNet
Configuration	$N=196$	$N=1{,}024$	$N=196$
Linear, unmasked	$49.94\pm 0.26$	$51.89$	$36.77\pm 0.49$
Softmax, unmasked	$51.71\pm 0.26$	$53.19\pm 0.49$	$37.96\pm 0.78$
Linear, masked (random)	$\mathbf{55.27\pm 0.68}$	$\mathbf{58.02\pm 0.08}$	$\mathbf{40.14\pm 0.33}$
Linear, masked (serpentine)	$\mathbf{54.14\pm 1.37}$	$\mathbf{57.84\pm 0.89}$	$\mathbf{41.35\pm 0.99}$

Table 2: Top-1 accuracy (%, mean

\pm

standard deviation over seeds) for ViT-S, trained from scratch under a single fixed recipe. Seed counts and protocol are given in Appendices E.2 and E.7 .

Method (tree)	$1{,}024$	$4{,}096$	$16{,}384$	$65{,}536$	$262{,}144$
Dense ( $\Theta(N^{2})$ )	0.15	2.24	36.07	—	—
FTFI (random)	1.10	4.93	25.20	99.99	424.97
StAd-TFI (serpentine)	0.39	1.76	8.02	37.82	165.67

Table 3: Runtime (seconds) of the masked-attention workload as a function of

N

(median of at least five runs; protocol in Appendix E.6 ). FTFI uses its most favorable configuration (random tree).

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: The schematic grid presents a matrix where each entry $G_{q}[s]$ sums over source depths $r$ and anchors $t$ .

Family	Construction and vertex count	Tested parameters
Path	One chain, $N$ vertices.	$N\in\mathcal{N}$ .
Balanced binary	Complete binary tree; $N=2^{d+1}-1$ .	$d\in\{4,5,6,7,8,9\}$ .
Complete 3-ary	Complete ternary tree; $N=(3^{d+1}-1)/2$ .	$d\in\{3,4,5,6\}$ .
Caterpillar	Three leaves per spine vertex, $N=4\ell$ .	$\ell\in\mathcal{L}$ .
Caterpillar + hanging paths	One $\ell$ -edge path at every spine vertex, $N=\ell+\ell^{2}$ .	$\ell\in\{8,12,16,24,32\}$ .
Comb	Two leaves per spine vertex, $N=3\ell$ .	$\ell\in\mathcal{L}$ .

Appendix

Table 4: Deterministic tree constructions and complete parameter sweeps. All leaves and hanging paths are distinct.

Family	$N$	BF	FTFI	D1	D2	A2
Path	800	250.93	59.31	1.19	1.15	2.09
Balanced binary	1023	399.17	64.10	49.46	23.80	56.63
Complete 3-ary	1093	458.63	56.36	63.56	29.72	61.63
Caterpillar	1600	1,004.16	84.25	35.77	35.43	38.02
Caterpillar + paths	1056	435.47	64.77	62.78	6.68	16.67
Comb	1200	577.43	67.53	24.20	24.11	25.74

Appendix

Table 5: Largest-size integration times (in milliseconds). BF denotes brute force; D1, D2, and A2 denote Diameter-1D, Diameter-2D, and Adaptive-2D. Random-family entries are means across three tree seeds. Bold indicates the lowest recorded mean in each row.

Figure 5: Additional runtime sweeps for complete 3-ary trees, caterpillars with three branches, combs with two teeth, brooms, and double brooms.

Figure 6: Additional runtime sweeps for uniform spiders, short-arm spiders, and stars.

Figure 7: Runtime sweeps for random trees, random weighted MST trees, and preferential attachment trees. Curves show means of per-seed timing minima across three tree seeds; bands show approximate

95\%

confidence intervals.

Figure 8: The 28 Thingi10K meshes used in the geodesic Sinkhorn experiments.

StAd-TFI variant ( $A$ )	Baseline ( $B$ )	Meshes	Single $Kx$	Per iteration	Sinkhorn	With setup
SpclK-Diameter	FTFI ( SpclK-Centroid )	28	$3.31\times$	$3.38\times$	$3.38\times$	$3.35\times$
SpclK-Adaptive	FTFI ( SpclK-Centroid )	28	$3.00\times$	$3.07\times$	$3.07\times$	$2.98\times$
Adaptive-2D-FFT	FTFI ( 1D-FFT )	28	$5.99\times$	$6.32\times$	$6.32\times$	$5.84\times$
SpclK-Diameter	FTFI ( SpclK-LinearExactDP )	28	$1.24\times$	$1.28\times$	$1.28\times$	$1.20\times$
SpclK-Adaptive	FTFI ( SpclK-LinearExactDP )	28	$1.12\times$	$1.17\times$	$1.17\times$	$1.07\times$

Appendix

Table 6: Geometric-mean speedups

t_{B}/t_{A}

over 28 meshes. Values greater than one indicate that

A

is faster.

Metric	Mean $\pm$ std.	Median	Maximum
Kernel error $E_{K,Q}$	$0.3158\pm 0.1409$	$0.3201$	$0.6909$
Relative OT error $E_{\mathrm{rel}}(Q)$	$0.0295\pm 0.0262$	$0.0274$	$0.1014$
Plan variation $\operatorname{TV}_{Q}$	$0.0295\pm 0.0258$	$0.0238$	$0.1140$

Appendix

Table 7: Quantization and transport-cost errors on the MSTs of 28 Thingi10K meshes at

B=64

. Kernel error and plan TV compare the quantized and original tree problems; relative OT error uses the original-graph transport-cost reference. Each mesh has equal weight. We report the arithmetic mean

\pm

sample standard deviation, median, and maximum. All values are dimensionless.

Figure 9: CIFAR-100 top-1 accuracy versus training epoch. (a) Masked softmax and masked linear attention at $N{=}196$ (curves are means over three seeds). (b, c) Masked linear attention across spanning trees at $N{=}196$ and $N{=}1{,}024$ ; multi-seed configurations are drawn as seed means, the remaining configurations as single seeds. In all panels, the dashed and dotted horizontal lines mark the best accuracy of the unmasked softmax and unmasked linear baselines, respectively. Every masked tree exceeds both baselines (Appendix E.7 ).

Figure 10: The five spanning-tree families of the patch grid evaluated in the tree sweep, from bushy (random) to path-like (serpentine, Hilbert). Every spanning tree of the unit-weight grid is a minimum spanning tree, so all are admissible under the topological-mask construction.

tree	family	avg. edge stretch	top-1 (%)	$\Delta$ vs. unmasked
unmasked	—	—	49.79	—
fishbone	caterpillar	4.50	55.95	$+6.16$
comb	caterpillar	7.50	54.86	$+5.07$
random	Kruskal	5.35	54.80	$+5.01$
serpentine	path	7.50	52.82	$+3.03$
Hilbert	path	8.81	51.74	$+1.95$

Appendix

Table 8: Accuracy as a function of tree choice (CIFAR-100, ViT-S, masked linear attention,

N{=}196

, single seed), together with the average edge stretch on the

14\times 14

grid. Accuracy decreases with stretch (Pearson

-0.86

tree	avg. edge stretch ( $32^{2}$ )	top-1 (%)	$\Delta$ vs. softmax
random	7.21	58.02	$+4.83$
serpentine	16.50	57.84	$+4.65$
fishbone	9.00	56.38	$+3.19$
rec. bisection	6.99	n.c.	—

Appendix

Table 9: Accuracy as a function of tree choice at

N{=}1{,}024

(patch

7

32\times 32

grid; CIFAR-100, ViT-S, masked linear attention); random and serpentine are means over two seeds (Section E.7.5 ), fishbone single seed; average edge stretch on the

32\times 32

grid. References: unmasked linear

51.89

, plain softmax

53.19\pm 0.49

(three seeds). “n.c.”: not converged at this scale (Section E.7.5 ).

$N$	$16{,}384$	$65{,}536$	$147{,}456$	$262{,}144$
StAd-TFI , serpentine (path, $H{=}0$ )	8.89	41.80	96.71	196.49
StAd-TFI , fishbone ( $H{\approx}\sqrt{N}/2$ )	13.41	103.19	492.82	813.26
StAd-TFI , comb ( $H{\approx}\sqrt{N}$ )	30.70	264.71	786.96	2593.84
StAd-TFI , random	27.63	133.01	397.78	1618.01
FTFI , serpentine	240.31	1099.08	$\dagger$	$\dagger$
FTFI , fishbone	34.69	225.81	$\dagger$	$\dagger$

Appendix

Table 10: Single-pass runtime (seconds) of the masking primitive per tree at

C{=}4{,}160

on a

32

-core CPU; all methods within a row are timed in the same process. The upper block reports StAd-TFI ; the middle block reports FTFI on the same tree (

\dagger

: not completed within a

15

-minute budget); the final row is the incumbent FTFI -random baseline. The path scales near-linearly for StAd-TFI , deeper trees super-linearly, and the relative efficiency shifts toward StAd-TFI as trees become more backbone-dominated (cf. Table 11 ).

Tree	$H$	$T_{\textsc{FTFI}{}}$	$T_{\mathrm{StAd}}$	$T_{\textsc{FTFI}{}}/T_{\mathrm{StAd}}$
Random	616	100.0	121.4	$0.82\times$
Comb	255	379.3	231.1	$1.6\times$
Fishbone	128	217.1	91.4	$2.4\times$
Serpentine	0	964.1	37.8	$\mathbf{25.5\times}$

Appendix

Table 11: Runtime crossover at

N{=}65{,}536

C{=}4{,}160

(seconds; protocol as in Table 3 ).

H

is the off-backbone depth of the diameter backbone; the relative efficiency shifts toward StAd-TFI on backbone-dominated trees.

path construction	stretch ( $14^{2}$ )	stretch ( $32^{2}$ )
row-snake (serpentine)	7.50	16.50
column-snake	7.50	16.50
Hilbert	8.81	19.62
diagonal-snake	9.50	21.50
spiral	16.64	40.56

Appendix

Table 12: Hamiltonian-path constructions ranked by average edge stretch; all attain identical

H{=}0

runtime. The row-snake (serpentine) ordering attains the lowest stretch.

Explore similar work

Stochastic Optimization of Tree Tensor Networks
Sep 1, 2026Marius Willner, Maximilian Scharf, André Uschmajew +2Tensor Networks Neural Network Optimization
Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany · Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany · Tensor AI Solutions GmbH, 89284 Pfaffenhofen an der Roth, Germany +1
Riemannian Optimization on Tree Tensor Networks with Application in Machine Learning
Jul 29, 2025Marius Willner, Marco Trenti, Dirk LebiedzTensor Networks Riemannian Optimization
Chair of Mathematical Data Science, University of Augsburg, Augsburg, Germany · Tensor AI Solutions GmbH, Pfaffenhofen a.d. Roth, Germany · Institute for Numerical Mathematics, University of Ulm, Ulm, Germany
Towards Scalable Persistence-Based Topological Optimization
May 9, 2026Abderrahim Bendahi, Alexandre Duplessis, Arnaud FickingerTopology Optimization Persistent Homology
ENS Ulm, PSL Paris, France · UC Berkeley Berkeley, CA, USA

Sep 1, 2026, math.OCJ/K move · Enter open · S save

Stochastic Optimization of Tree Tensor Networks

Marius Willner, Maximilian Scharf, André Uschmajew, Timo Felser+1

Institute of Mathematics, University of Augsburg, 86159 Augsburg, Germany · Centre for Advanced Analytics and Predictive Sciences, University of Augsburg, 86159 Augsburg, Germany · Tensor AI Solutions GmbH, 89284 Pfaffenhofen an der Roth, Germany+1

Tensor Networks Neural Network Optimization Riemannian Optimization Stochastic Optimization Batch Many-Body Physics+1

Open paper AskPDF arXiv

Tensor networks, originally developed for quantum many-body physics, are promising models for machine learning. We derive stochastic Riemannian optimizers for tree tensor networks (TTNs) on both their parameter and quotient manifolds, including adaptive and learning-rate-free schemes suitable for minibatch training. Using a hybrid CNN-TTN architecture, we evaluate the methods on Fashion-MNIST, CIFAR10, and Imagenette. The proposed optimizers achieve predictive performance comparable to unconstrained optimization while enabling numerically stable downstream compression.

Abstract

Explore similar work

Stochastic Optimization of Tree Tensor Networks

Riemannian Optimization on Tree Tensor Networks with Application in Machine Learning

Towards Scalable Persistence-Based Topological Optimization