Table of Contents
Fetching ...

Fast Tree-Field Integrators: From Low Displacement Rank to Topological Transformers

Krzysztof Choromanski, Arijit Sehanobish, Somnath Basu Roy Chowdhury, Han Lin, Avinava Dubey, Tamas Sarlos, Snigdha Chaturvedi

TL;DR

This work introduces Fast Tree-Field Integrators (FTFI), a class of exact polylog-linear algorithms for integrating tensor fields on weighted trees by exploiting low displacement rank (LDR) structure. By constructing IntegratorTrees (ITs) and leveraging cross-term decompositions, FTFI achieves exact $f$-distance matrix multiplications with strong complexity guarantees for a broad family of functions $f$ (including rational, polynomial, exponential, and exponentiated-quadratic), enabling efficient approximations of graph metrics via tree metrics and accelerating applications from mesh interpolation to Graph Classification and Topological Vision Transformers (TopViT). Empirically, FTFI delivers 5.7–13x speedups over brute-force methods on large graphs, and when integrated into TT masking mechanisms, yields notable accuracy gains with only a few extra trainable parameters per layer. The work further demonstrates learnable $f$-distance matrices to improve approximation quality and provides extensive experiments on vision and graph tasks, along with a thorough theoretical treatment of the method. Overall, FTFI offers a practical, theoretically grounded path to scalable graph- and transformer-based computation on large-scale structured data.

Abstract

We present a new class of fast polylog-linear algorithms based on the theory of structured matrices (in particular low displacement rank) for integrating tensor fields defined on weighted trees. Several applications of the resulting fast tree-field integrators (FTFIs) are presented, including (a) approximation of graph metrics with tree metrics, (b) graph classification, (c) modeling on meshes, and finally (d) Topological Transformers (TTs) (Choromanski et al., 2022) for images. For Topological Transformers, we propose new relative position encoding (RPE) masking mechanisms with as few as three extra learnable parameters per Transformer layer, leading to 1.0-1.5%+ accuracy gains. Importantly, most of FTFIs are exact methods, thus numerically equivalent to their brute-force counterparts. When applied to graphs with thousands of nodes, those exact algorithms provide 5.7-13x speedups. We also provide an extensive theoretical analysis of our methods.

Fast Tree-Field Integrators: From Low Displacement Rank to Topological Transformers

TL;DR

This work introduces Fast Tree-Field Integrators (FTFI), a class of exact polylog-linear algorithms for integrating tensor fields on weighted trees by exploiting low displacement rank (LDR) structure. By constructing IntegratorTrees (ITs) and leveraging cross-term decompositions, FTFI achieves exact -distance matrix multiplications with strong complexity guarantees for a broad family of functions (including rational, polynomial, exponential, and exponentiated-quadratic), enabling efficient approximations of graph metrics via tree metrics and accelerating applications from mesh interpolation to Graph Classification and Topological Vision Transformers (TopViT). Empirically, FTFI delivers 5.7–13x speedups over brute-force methods on large graphs, and when integrated into TT masking mechanisms, yields notable accuracy gains with only a few extra trainable parameters per layer. The work further demonstrates learnable -distance matrices to improve approximation quality and provides extensive experiments on vision and graph tasks, along with a thorough theoretical treatment of the method. Overall, FTFI offers a practical, theoretically grounded path to scalable graph- and transformer-based computation on large-scale structured data.

Abstract

We present a new class of fast polylog-linear algorithms based on the theory of structured matrices (in particular low displacement rank) for integrating tensor fields defined on weighted trees. Several applications of the resulting fast tree-field integrators (FTFIs) are presented, including (a) approximation of graph metrics with tree metrics, (b) graph classification, (c) modeling on meshes, and finally (d) Topological Transformers (TTs) (Choromanski et al., 2022) for images. For Topological Transformers, we propose new relative position encoding (RPE) masking mechanisms with as few as three extra learnable parameters per Transformer layer, leading to 1.0-1.5%+ accuracy gains. Importantly, most of FTFIs are exact methods, thus numerically equivalent to their brute-force counterparts. When applied to graphs with thousands of nodes, those exact algorithms provide 5.7-13x speedups. We also provide an extensive theoretical analysis of our methods.
Paper Structure (44 sections, 3 theorems, 10 equations, 10 figures, 6 tables, 1 algorithm)

This paper contains 44 sections, 3 theorems, 10 equations, 10 figures, 6 tables, 1 algorithm.

Key Result

Lemma 3.1

If $\mathcal{K}$ is a tree with $|\mathcal{K}| \geq 6$, then $\mathcal{K}$ admits a decomposition ($\mathcal{K}_{\mathrm{left}}, \mathcal{K}_{\mathrm{right}}, p)$ given above and it can be constructed in linear time.

Figures (10)

  • Figure 1: Pictorial representation of the IntegratorTree (see: Sec \ref{['sec:it']}) data structure for the nine-vertex input tree $\mathcal{T}$ on the left. Numbers in blue next to the input tree denote the weights of its edges. Leaves of the IntegratorTree object represent $f$-transformed (element-wise) distance matrices: $\mathbf{D}_{0},\mathbf{D}_{1},\mathbf{D}_{2},\mathbf{D}_{3}$ for sub-trees induced by vertex-sets: $\{1,2,4\},\{1,3,0\},\{5,7,8\}$ and {5,6,0} respectively. Different levels correspond to different distances from the pivot point.
  • Figure 2: Pictorial representations of the main concepts behind efficient matrix-vector multiplications $\mathbf{Mv}$ with $\mathbf{M} \in \mathbb{R}^{5 \times 4}$, for the polynomial $f$ and $f(x)=\frac{\exp(\lambda x)}{x+c}$. In the polynomial case, $\mathbf{M}$ is re-written as a sum of low-rank outer-product matrices corresponding to terms of different degrees (e.g., constant, linear, quadratic, etc.). Matrix associativity property is applied for efficient calculations (dotted-border blocks indicating the order of computations). In the second case, $\mathbf{M}$ is high-rank, but the so-called low displacement rank operator$\Delta_{D_{1},D_{2}}:\mathbf{X} \rightarrow \mathbf{D}_{1}\mathbf{M}-\mathbf{M}\mathbf{D}_{2}$ for diagonal $\mathbf{D}_{1},\mathbf{D}_{2}$ can be applied to make it a low-rank outer-product matrix. The multiplication with $\mathbf{M}$ can be efficiently performed using the theory of LDR matrices ldr-1.
  • Figure 3: Runtime comparison of FTFI with BTFI as a function of the number of vertices, $N$. Left: Synthetic graphs. Right: Mesh-graphs from https://ten-thousand-models.appspot.com/. The speed is not necessarily monotonic in $N$ as it depends on the distribution of lengths of the shortest paths. For each graph, 10 experiments were run (std. shown via dotted lines).
  • Figure 4: Speed (pre-processing time) and accuracy (cosine similarity) comparison of the FTFI and other baselines for vertex normal prediction on meshes. Cosine similarity of BFFI and FTFI almost overlaps. The last two figures are qualitative examples showcasing the tradeoff between cosine similarity and preprocessing time for meshes of sizes 3K and 5K nodes respectively.
  • Figure 5: Trade-off plot comparing graph classification accuracy and feature processing time for the classifiers using FTFI and BGFI. FTFI achieves similar accuracy as BGFI while significantly reducing fp time across most datasets. We report the reduction in FTFI's processing time ($\pm$x%) compared to BGFI using a dotted line.
  • ...and 5 more figures

Theorems & Definitions (7)

  • Lemma 3.1: Pivoting
  • Definition 3.2: cordial functions
  • Lemma 3.3: $f$-integration with cordial functions
  • proof
  • proof
  • Lemma A.1
  • Definition C.1: general masked attention