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

Authors: krzysztof  choromanski
arijit  sehanobish
somnath  basu
roy  chowdhury
han  lin
avinava  dubey
tamas  sarlos
snigdha  chaturvedi
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:29:07

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.

Main paper: 
 introduction: 
 Matrix-vector multiplication remains a key computational block of virtually all modern machine learning (ML) algorithms. For this reason, decades of research have been dedicated towards making this fundamental operation more efficient. One approach to achieve this goal is through efficient hardware design, e.g., using modern GPU and TPU accelerators [Abadi et al., 2016 , Yu et al., 2022 , 2020] . The alternative method involves developing algorithms for efficient matrix-vector multiplication by leveraging either (1) sparse matrices [Wang, 2021 , Beniamini et al., 2020] , or (2) structured dense matrices [Thomas et al., 2018 , Chandrasekaran et al., 2018] . These algorithms can be applied in modern neural network systems, where weights are pruned to encourage sparsity [Blalock et al., 2020] or they can be parameterized with structured matrices [Sindhwani et al., 2015] .In this work, we aim to accelerate multiplications with a large class of matrices, that we refer to as f -distance matrices, which play an important role in several ML algorithms. Consider a matrix M G f = [f (dist(i, j))] i,j=1,...,N ∈ R N ×N , where dist(i, j) stands for the shortest-path distance between the i-th and j-th vertex of an undirected graph G = (V, E, W). Here V = {1, ..., N } stands for the set of vertices (nodes), E denotes the set of edges, W : E → R + maps them to their positive weights, and f : R → R. We call M G f a f -distance matrix in G. Note that if f (x) def = x, then M G f is the Shortest Path Kernel matrix.The product M G f x (where x ∈ R N ) represents a scalar field on V obtained by discretely integrating the field defined by x. In this integration, a new field value at a vertex v is calculated by averaging the old field values at all vertices u, weighted according to the function f (dist (v, u) ). This integration can be extended to general tensor fields by replacing vector x ∈ R N with a tensor X ∈ R N ×d1×d2×... :M G f X[i] = j∈V(G)f (dist(i, j))X [j] (1)We refer to the above procedure as the f -integration of a field X on G. We will use the terms graph field integration (GFI) and multiplication with f -distance matrices interchangeably throughout the paper. When the graph, G, is a tree, we call this procedure (Eq. 1) tree field integration. Next, we highlight several applications that rely on multiplications with f -distance matrices, M G f . 1. Interpolation on manifolds: This task involves predicting unseen values on a manifold from a set of known values. For example, predicting the velocities of all points on a flag with known velocities for a few points [Pfaff et al., 2021] . For a discretized manifold, the interpolated values can be obtained using a weighted average using graph field integration (Eq. 1). 2. Optimal Transport (OT): A popular method used to solve the entropic OT problem [Peyré and Cuturi, 2019] is the Sinkhorn algorithm [Eckstein and Nutz, 2022] . Sinkhorn relies on multiplications with cost matrices, which are special cases of f -distance matrices for metric spaces induced by shortest-path distances in graphs. This can be efficiently solved using graph field integration. 3. Topological Transformers (TTs): Topological Transformers [Choromanski et al., 2022] are extensions of traditional Transformers [Vaswani et al., 2017] for graph inputs. TTs modify the 1-D relative positional encoding (RPE) using "mask matrices", which are f -distance matrices. We show how these matrices can be efficiently integrated into the attention mechanism (Sec. 4.4).In the above applications, apart from the graph field integration step, the bottleneck lies in the process of explicitly materializing the f -distance matrix. Naively performing the integration in Eq 1 consists of two steps: (a) computing the f -distance matrix, M G f , which requires O(N 3 ) time in the worst case (which we call preprocessing), and (b) performing the multiplication takes O(N 2 ) time. This is prohibitively expensive while using large graphs.In this paper, we introduce a new class of fast polylog-linear algorithms for graph field integration that uses low displacement rank (LDR) matrices [Thomas et al., 2018 , Chandrasekaran et al., 2018] . To summarize, our primary contributions are given below:1. We provide the first exact polylog-linear multiplication algorithms called Fast Tree-Field Integrators (FTFIs), for general weighted trees and a rich class of maps f , including rational, trigonometric, exponential and exponentiated quadratic functions (Sec. 3.2). 2. We show how Fast Tree-Field Integrators can be applied to support fast computations on general graphs by approximating graph metrics with tree metrics (Sec. 4). 3. We show that FTFIs are 5.7-10x faster than baseline graph field integration methods for large-scale graphs (Sec. 4.1 and 4.2). 4. We showcase the efficacy of FTFIs in several applications including graph classification (Sec. 4.2), interpolation on meshes (Sec. 4.2), and Topological Vision Transformers (TVTs) (Sec. 4.4) . For TVTs, we propose new relative position encoding (RPE) masking mechanisms by introducing only three extra learnable parameters, which leads to 1.0-1.5% accuracy gains. We provide an exhaustive evaluation on Vision Performers (25 models on multiple datasets). Some of our best models use exponentiated quadratic functions f , which has not been applied in this context before.For completeness, we also propose approximate FTFI extensions via Non-Uniform FFT (NU-FFT) [Kircheis et al., 2023] and random Fourier features (RFFs) [Rahimi and Recht, 2007] (Sec. A.2).Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We give detailed explanations of our contributions in the introduction (page 2).Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Efficient graph field integration (Eq. 1) has been studied by prior works for different classes of matrices. For example, Al-Mohy and Higham [2011] considered exponentiated adjacency matrixvector multiplication, Spielman and Teng [2012] targeted symmetric diagonally dominant matrices (e.g., Laplacian), Arrigo et al. [2018] analyzed matrices that are power series of random walk kernels.In contrast to these approaches, Saad and Schultz [1986] proposed general iterative methods for solving certain linear systems using Arnoldi's iterations. However, These iterative methods can suffer from convergence issues. Williams [2007] showed that it is possible to pre-process any boolean matrix to achieve sub-quadratic matrix-vector multiplication.The general problem of computing the action of a matrix on a vector, where the matrix is the graph kernel, in sub-quadratic time is intractable, except for a few special cases [Al-Mohy and Higham, 2011, Choromanski et al., 2023] . In this work, we embed the graph G under consideration in a tree (replacing the graph metric by the underlying tree metric). Then, we leverage the tree structure to approximate the action of the kernel on a given vector by providing exact integration on a tree.Previous works [Bartal et al., 2022 , 2019 , Abraham et al., 2008 , Bartal, 1998 ] have used the theory of tree metrics (TMs) in several applications in mathematics and computer science. TMs are widely used to embed a complex metric space (e.g., a Riemannian manifold) into a more tractable one, while approximately preserving (all or most of the) pairwise distances. They find applications in distributed & online algorithms [Khan et al., 2008 , Bubeck et al., 2018] , biology [Mossel, 2007] , vision, robotics [Athitsos and Sclaroff, 2003] , and ML (e.g., metric spaces' regression [Gottlieb et al., 2011] ).Tree metrics for fast matrix multiplication: Applying tree metrics (TM) to compute approximate M G f is a natural approach to scale up matrix multiplications. If a TM approximates the metric space well, then the derived embeddings should have low distortion. However, in the worst-case scenario, this is not true for deterministic tree embeddings. A natural alternative is to sample trees from probabilistic distributions, which are shown to provide logarithmic distortion in expectation [Fakcharoenphol et al., 2004b , Bartal et al., 2022] . This can be further improved to constant distortion for certain classes of metrics, e.g., celebrated snowflake metics [Leeb, 2016] . For graph metrics defined by shortest-path distances, there exist spanning trees providing constant average distortion (over all pairs of nodes). These spanning trees can be constructed as near minimum weight spanning trees [Bartal et al., 2016] . Unfortunately, explicit application of any tree metric still requires O(N 2 ) time (impractical for large N ) to: (1) compute all shortest-path distances via the breadth-first-search algorithm (BFS), even if sub-quadratic methods were used to construct a tree (e.g. minimum spanning tree), (2) store the matrix, and (3) perform matrix-vector multiplications. We provide more details about work related to graph field integration in Appendix B.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The limitations are clearly explained in Appendix F Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
fast tree-field integrators (ftfi): 
 In this section, we present our approach for performing efficient field integration on a tree, which we call fast tree field integrator. We begin by introducing the concept of integrator trees (ITs), which is a specialized decomposition of a tree using the theory of balanced separators (Sec 3.1). Subsequently, we leverage these integrator trees to execute efficient integration on a tree via a divide-and-conquer algorithm (Sec 3.2).To support fast integration for various tensor fields X ∈ R N ×d1×...×ds defined on a given input tree T , we first design a special data structure that we refer to as an IntegratorTree (IT). An object of this type is constructed only once per T , regardless of the number of tensor fields used. An IT is a rooted binary tree. To avoid confusion, we will refer to its vertices as nodes, reserving term vertices for those of T . Each node of IT corresponds to the induced sub-tree ST of T . For every non-leaf node corresponding to some ST , a pivot point p along with two sub-trees: ST left and ST right are constructed. The following needs to be satisfied:• |ST x | ≥ |ST | 4 for x ∈ {left, right}, • ST x ∩ ST y = {p} (| • | denotes the number of vertices).The next lemma shows that every tree K with |K| ≥ 6 has the above decomposition and it can be efficiently found.Lemma 3.1 (Pivoting). If K is a tree with |K| ≥ 6, then K admits a decomposition (K left , K right , p) given above and it can be constructed in linear time.The algorithmic proof is provided in Appendix A.1 and uses standard tools from the theory of balanced separators.The left child of the non-leaf node for ST corresponds to ST left and the right child to ST right . In addition to these two pointers, a non-leaf node also contains eight extra fields, partitioned into two groups, one corresponding to its left child and one to its right children. The fields corresponding to the left child are as follows:• Left-ids: an array of the ids (in T ) of those vertices that are in ST left , mapping the ids of vertices in ST left to the original ids in T (each sub-tree uses consecutive numbers from 0 as ids locally). In principle, the leaf nodes of IT correspond to sub-trees with less than t = 6 vertices each. In practice, we choose higher t, for more efficient integration (see: discussion in Sec. 4.1).From what we have said so far, it is clear that an IT can be constructed by applying breadth first search (BFS) and the linear algorithmic procedure for constructing the decomposition from Lemma 3.1. Note that every vertex of the input tree appears in the logarithmic number of nodes in the IT since the size of the sub-tree is at most 3 4 × the size of its parent in IT. We conclude that IT for the given input tree T can be computed in O(N log(N )) time, where N stands for the number of vertices |T | of T .We are ready to explain how ITs allow us to efficiently integrate any given tensor field X ∈ R N ×d1×...×ds defined on T for a wide class of function f : R → R. We will apply a divide-andconquer strategy.We start in the root node of IT. If that node is a leaf then the f -transformed distance matrix is stored and can be directly used for matrix-tensor multiplication. If this node is not a leaf, then it encodes the decomposition (T left , T right , p). Take some v ∈ V(T left ). Note that the value M G f X[v] of the new field in v after f -integration is given as follows for W = V(T right )\{p}:j∈V(T left ) f (dist(v, j))X[j] Finner(v) + j∈W f (dist(v, j))X[j] Fcross(v).(2)To compute the new values of the field for nodes v ∈ T left , one needs to:1. Compute the contribution to it from T left (F inner (v)-terms). This can be done simply by applying Eq. 2 recursively for T left , which means traversing to the left child of the root. 2. Add the so-called cross-terms contributions coming from the vertices of W (F cross (v)-terms).The key observation is that the latter (cross-term) contributions can be retrieved simply by computing CX ′ , where: ..×ds is defined as follows:X ′ [j] def = bj -1 z=0 X[right-ids[right-s[j][z]]].(3)Given the structure of IT, tensor X ′ can be computed in linear time. Note that the following holds:EQUATIONwhere τ (v) = left-id-d [v] . Analogous analysis can be derived for v ∈ T right , with matrix C ⊤ replacing C. Thus the overall time complexity of the cross-terms computations is determined by the algorithm for matrix-tensor multiplications with matrices C and C ⊤ . Next, we demonstrate the importance of cordial functions in our FTFI framework. Lemma 3.3 (f -integration with cordial functions). If f is d-cordial then f -integration for the general weighted tree of N vertices can be conducted in time O(N log d+1 (N )).Proof. Denote by T (N ) time complexity for running FTFI on the N -vertex tree. We have the following recursive formula for T , where 1 4 ≤ c ≤ 3 4 :EQUATIONThis is implied by the fact that: (1) the size of each sub-tree is at most 3 4 × the size of its parent, (2) the computation across left and right children is dominated by multiplications with matrices C and C ⊤ . The solution of this recursion leads to the statement.Next, we show some practical implications of Lemma 3.3, where tree weights are completely arbitrary. Additional results are given in Sec. A.2.3.We claim that every rational f is (2 + ϵ)-cordial for any ϵ > 0. We will use Lemma 1 from [Cabello, 2022] stating that: given any set of b rational functions R j (x) = )) (by applying FFT). For a given vector v ∈ R b , it thus suffices to define: R j (x) = v j f (x + y j ) and that lemma can be applied to efficiently compute Mv. We conclude that for any ϵ > 0, f -integration can be conducted in O(N log 3+ϵ (N )) time for N -vertex weighted trees and any rational f : R → R (see also : Sec. 4.3, Sec. 4.2, Sec. 4.4) . x+c . In the polynomial case, 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, M is high-rank, but the so-called low displacement rank operator ∆D 1 ,D 2 : X → D1M -MD2 for diagonal D1, D2 can be applied to make it a low-rank outer-product matrix. The multiplication with M can be efficiently performed using the theory of LDR matrices [Thomas et al., 2018] .Polynomial functions: The above result on rational functions clearly applies also to polynomial f , but here we can do better. We show that f is 0-cordial. Assume that f (x) = B t=0 a t x t . We have: M = B t=0 t l=0 a t t l M l,t-l , where matrix M u,v ∈ R a×b is defined as an outer-product of two vectors:(x u 1 , ..., x u a ) ∈ R a and (y v 1 , ..., y v b ) ∈ R b .Thus each M u,v supports linear matrixvector multiplication (via associativity property). The proof is completed, since B is a constant. We conclude that f -integration can be conducted in O(N log(N )) time for N -vertex weighted trees and any polynomial f : R → R (see: Fig. 2 Exponential functions: Take f (x) = exp(λx). Then M is an outer-product of two vectors: (exp(λx i )) a i=1 ∈ R a and (exp(λy j )) b j=1 ∈ R b . The remaining analysis and conclusion is thus the same as for the polynomial case (see also: Sec. 4.4).Function: f (x) = exp (λx) x+c : (c is a constant) We claim that f is 2-cordial. In that setting, matrix M satisfies:M(i, j) = exp(λxi) exp(λyj ) (xi+ c 2 )+(yj + c2 ) and thus is a Cauchy-like LDR, supporting fast O(N log 2 (N )) matrix-vector multiplication [Victor Y. Pan, 2000] . We conclude that f -integration can be conducted in O(N log 3 (N )) time for N -vertex weighted trees and f (x) = exp (λx) x+c (see: Fig. 2 ).Functions f (x) = exp(ux 2 + vx + w) and trees with positive rational weights: Now matrix M can be re-written as M = exp(w)D 1 VD 2 , where D 1 ∈ R a×a and D 2 ∈ R b×b are diagonal, with diagonal entries given by sequences {exp(ux 2 i +vx i )} a i=1 and {exp(uy 2 j +vy j )} b j=1 respectively, and furthermore V is the generalized Vandermonde matrix (GVM) (using arbitrary nonnegative integers as exponents). It is defined as: V(i, j) = r sj i , where r i = exp( 2uxi q ) and s j = y j q ∈ N. As in the previous case, the embedding trick can be applied, but we will use it only for columns. That effectively leads to the completion of the set of exponents {s j } to the set of consecutive integers starting from 0 and a regular Vandermonde matrix, that supports O(N log 2 (N )) matrix-vector multiplication, replacing GVM. The benefit of this embedding, as compared to the previous one, is that even though it still increases the number of columns by a multiplicative factor of p, the number of rows does not change. Therefore, for p ≫ log(N ), substantial computational speedups are achieved (see: Sec. 4.4).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: We provide the code as well as details to run our experiments in Appendix D.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experiments: 
 In this section, we outline the experimental setup and report the performance of FTFI across various settings. For all the experiments, we only consider minimum spanning tree (MST) as an approximation of our graph. Specifically, we design experiments to answer these research questions: (Q1) How efficient are FTFIs for tree field integration? (Q2) How does the approximation quality of FTFI compare to other integration algorithms? (Q3) How can we further improve the approximation quality in FTFI? (Q4) How can we use FTFI in real-world large-scale settings? The main goal of this experiment is to evaluate the speedups obtained by FTFI as compared to brute-force tree field integrator (BTFI) i.e. the explicit calculation of Eq 1 on a tree. We consider two classes of graphs: (a) synthetic, obtained from a path-graph by adding random edges and (b) mesh graphs from Thingi10K [Zhou and Jacobson, 2016] dataset. For BTFI, we compute the MST and then integrate a random scalar field X on the vertices of the MST. Since BTFI & FTFI are numerically equivalent, we report the pre-processing time and integration as a function of vertex count (N ) in Fig. 3 . We observe that FTFI achieves up to 13x speedups for 20K-vertex meshes and 5.7x+ for synthetic graphs with over 10K vertices compared to BTFI.We evaluate the approximation quality achieved by FTFI across a wide range of graph-based tasks.Interpolation on meshes. We compare the efficiency of FTFI with baselines on the normal vector prediction task. Every node of the considered mesh G with a vertex-set V, is associated with a location x i ∈ R 3 and a vertex normal F i ∈ R 3 . For each mesh, we randomly select a subset V ′ ⊆ V with |V ′ | = 0.8|V| and mask out their vertex normals (set as zero vectors). The interpolation task involves predicting the vertex normals of each masked node i ∈ V ′ as:F i = j∈V\V ′ K f (i, j)F j , where K f (w, v) = f (dist(w, v)), with dist(w, v) being the shortest path distance between node w and v, and f is a rational function f (x) = 1/(1 + λx 2 ). We perform a grid search to set hyperparameter λ for each mesh and report the result with the highest cosine similarity between predicted and ground truth vertex normals, averaged over all the nodes. We run tests on 40 meshes of the 3D-printed objects with a wide range of sizes from the Thingi10K dataset (details in Appendix D.3). We compare FTFI with BTFI, low-distortion tree-based algorithms such as Bartal Trees [Bartal, 1996] and FRT trees [Fakcharoenphol et al., 2004a] alongside the state-of-the-art method for graph-field integration, the Separator Factorization (SF) algorithm [Choromanski et al., 2023] . We also compare against the baseline BGFI which entails explicitly materializing the kernel matrix of G and then performing matrix tensor multiplication with a tensor field F defined by the F i 's.Preprocessing involves building specific tree structures (FRT, Bartal) , calculating the kernel matrices (BGFI, BTFI), or creating specialized data structures (SF, FTFI) for efficient later use. The first two plots in Fig. 4 shows the pre-processing time and cosine similarity for various algorithms applied to meshes of different sizes. FTFI is the fastest in terms of pre-processing time and achieves competitive performance in terms of cosine similarity (between predicted and actual vertex normals) when compared with the SF algorithm while being numerically equivalent to BTFI. FTFI is a few orders of magnitude faster than BTFI and the tree-based methods while maintaining accuracy.Graph classification. Graph kernels have been widely used for graph classification tasks in previous works [Kriege et al., 2020 , Nikolentzos et al., 2021] . We compare the classification results obtained using the approximate kernel from FTFI with those from the exact SP kernel. In this setting, we use the Shortest Path (SP) kernel, f (dist(i, j)). We perform experiments on a wide range of bioinformatics and social networks datasets like D&D, MUTAG, REDDIT, IMDB, among others. We follow [de Lara and Pineau, 2018] and construct the graph feature for both kernels by using the smallest k eigenvalues (k is a hyperparameter). This feature set is then used for classification, using a random forest classifier. We observe that FTFI achieves significant speed improvements while achieving similar accuracy compared to its brute-force counterpart, BGFI (see Fig. 5 ). We provide more details about the experimental setup and baselines Appendix D.4. We also report additional experiments on meshes and point clouds in Appendix D.1.We propose to further improve the approximation quality of FTFI by learning a f -distance matrix on metrics derived from the MST. As an application, we choose general graph metrics, where our goal is to learn the shortest-path distance d v,w between a given pair of nodes (v, w) in a graph. Given a f -distance matrix and tree-derived metric d v,w the objective is to learn a mapping to minimizeEQUATION)Rather than using a fixed f , we parameterize and train it. We consider rational function f : f a0,...,at b0,...,bs (x) =EQUATIONwhere a 0 , ..., a t , b 0 , ..., b s ∈ R are trainable parameters. Final evaluation. To evaluate the quality of the approximation, we compute the relative Frobeniusnorm error: ϵ = ∥M T f -M G id ∥F ∥M G id ∥F, where ∥ • ∥ F stands for the Frobenius norm, T is a tree for a given graph G and id is an identity function (see: our notation from Sec. 1). It quantifies how closely the distance matrix of G is approximated by the f -distance matrix of T. Computing ϵ is expensive and our training does not rely on it. Our empirical results show that the relative error, ϵ, can be substantially improved by using the light-weight MSE training loss (defined in Eq. 6).We report the evaluation error for these experiments in Fig. 6 (with additional results in Fig. 8 in the Appendix). We observe that a rational function with quadratic numerator and denominator provides strong performance across different graphs. We notice that increasing the training set to > 100 data points does not have a substantial impact on the final error. Estimating the coefficients of f provides approximation improvements across all graphs in as few as 40 training steps.These above results show that tree-based estimators are expressive enough to emulate integration on arbitrary graphs. This expressive power can be further enhanced by pairing them with "nonlinear" functions f . Thus, they explain why the presented techniques are relevant for general graphs.For large-scale applications of FTFI, we select Topological Vision Transformers (TopViT), [Choromanski et al., 2022] , and leverage it for efficient incorporation of masking within ViTs. We provide detailed description of masked Transformers in Appendix C.We propose an extension to TopViT that seamlessly integrates FTFI. In this extension, we model the mask matrix as an f -distance matrix (with learnable f ) defined on the minimum spanning tree (MST) obtained from the 2D grid graph image encoding, where vertices correspond to different patches. We parameterize f as f t g def = g( t i=0 a t x t ). We use the linear attention mechanism introduced in Performers [Choromanski et al., 2021] , where the attention kernel is written as: K(q, k) = ϕ(q) ⊤ ϕ(k) for a deterministic ϕ : R d QK → R, applied element-wise. We experiment with different values of hyperparameters g, t, ϕ and cross-heads parameter sharing strategies as shown in Table 1 (synced indicates that RPE-parameters are shared across different attention heads).We run experiments on ImageNet and Places365 datasets using ViT-B/16 (see Table 1 ). For all the kernels, our variants beat the baselines. For ϕ(x) = x 4 , the best variant applies an exponentiated quadratic function, for which we apply Vandermonde matrices (see: discussion in Sec. 3.2.1). Our Larger Transformer models: We scale our experiments to run on the larger ViT-L architectures and evaluate on ImageNet. In this setting, we use RPE mechanism with g = exp and t = 1 (that provided strong performance in previous experiments) and asynced strategy. We observe that FTFI provides 7% accuracy improvement (see: Fig. 7 ).Further results on Video Transformer (ViViT) [Arnab et al., 2021] are provided in Appendix D.6. We also provide additional experiments including Gromov-Wasserstein distance computation [Vayer et al., 2018] (see Sec. D.2), along with code pointers (Appendix D). 
conclusion: 
 We provided a new class of algorithms for fast and exact integration of tensor fields defined on weighted trees, relying on the theory of structured (in particular low displacement rank) matrices.We showed how those algorithms can be applied for accurate integration on general graphs, in particular via their minimum weight spanning trees. We presented several applications of the presented methods, from graph classification and interpolation on meshes, through graph metric approximation to Topological Vision Transformers. Our methods provide significant (5-13x) speedups while maintaining the quality of their exact counterparts. 
author contributions: 
 KC conceived the idea behind FTFI, proved the theoretical results, implemented FTFI algorithm and ran the vision experiments in this paper. AS integrated the FTFI algorithm in the GW style algorithms and ran some graph and point cloud classification tasks. SBRC ran graph classification experiments as well as experiments on the CUBES dataset. HL ran the experiments on the meshes. AD helped develop methods, and along with TS and SC acted as senior advisors for the project. All authors contributed to the writing of the manuscript.Denote:g = Mv for a given v = (v 1 , ..., v b ) ⊤ ∈ R b . Define: g(x) = R f (x -z)P (z)dz, where P is given as: P (z) = b j=1 v j δ(z -z j ), and furthermore: (1) δ is a delta-Dirac function, (2) z j = -y j . Our goal is to efficiently evaluate function g in points: {x 1 , ..., x a }.Assume that the inverse FT of g exists and denote it by η : C → C. Note that g is the FT of η and can be written as: g(x) = R η(ω) exp(-2πiωx)dω. Since g is also a convolution of f and P , η is a product of the inverse FTs: ρ and R respectively. Therefore, we can write: g(x) = R ρ(ω)R(ω) exp(-2πiωx)dω, where R(ω) = b j=1 v j exp(2πiωz j ). Now, function g can be evaluated for {x 1 , . . . , x a } as follows: (1) a quadrature method is applied to obtain points: ω 1 , ..., ω r (and corresponding weights) for the approximate computation of the integral defining g, (2) the NU-FFT is applied to compute R(ω) simultaneously in those points in polylog-linear time, (3) given pre-computed (ρ(ω i )R(ω i )) r i=1 (and the quadrature weights), NU-FFT is applied again to compute quadrature-based approximation of g.The f -integration process applying this method runs in polylog-linear time since the computation of g = Mv takes polylog-linear time. A prominent application is f given as: f (x) = sin(x)x , with ρ being a renormalized indicator of belonging to interval [-0.5, 0.5]. In this setting, the integral defining g is thus limited to [-0.5, 0.5]. Interestingly, for f (x) = sin(x)x we can also apply methods from Sec. 3.2.1 (see: our discussion below on the trigonometric case).Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: All details are presented in Sec 4 and Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: All experiments in the paper except for the ones using large Transformer models have been run multiple times using various random seeds and we report the relevant statistics. The experiments using Transformers are too expensive to run multiple times as the experiments are run on a huge dataset like ImageNet.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We report the compute resources used in Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: All authors have reviewed the NeurIPS code of ethics and the research conform to the code.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: The broader impacts of our work is detailed in Appendix E.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper is theoretical in nature and we are not releasing any new models or data.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We have properly cited all the papers that introduced various algorithms and data that are used in this work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We release the code for the main algorithm. The usage is detailed in the anonymous github repo. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Products of exponentials and polynomials: Note that a Hadamard (element-wise) product of two outer-product matrices is itself an outer-product matrix. Using the analysis from the polynomial and exponential cases, we conclude that M is a sum of a constant number of terms, each being an outer-product matrix. Thus the same conclusion follows.The case of the trigonometric f : If f (x) = cos(x) then it can be re-written as:f (x) = exp(ix)+exp(-ix) 2. Observe that the cordiality property is preserved under linear combination of the finite number of cordial functions. We can thus conclude that analogous results as the above for f (x) = exp(λx) can be derived for f (x) = cos(x). That is also the case for f (x) = sin(x) that can be re-written as: f (x) = exp(ix)-exp (-ix) 2i. In both cases, we extend the domain from R to C, but this does not affect the analysis. So far we have not put any restrictions on the tree weights. If we restrict all weights to be the same (without loss of generality, equal to one), then the problem becomes easier. In this case for any function f , matrices C and C ⊤ are Hankel [Brent, 2010] (constant on each anti-diagonal and belonging to LDR class). Then, matrix-vector multiplication can be done in O((a + b) log(a + b)). The analysis from the proof of Lemma 3.3 for d = 1 can be repeated. We conclude that f -integration can be conducted in O(N log 2 (N )) time for N -vertex unweighted trees and any f : R → R. This was already proven in [Choromanski et al., 2022] .Trees with positive rational weights: Assume that tree weights take values of the form: { e q : e ∈ {1, ..., p}} for some p, q ∈ N + . Then, matrices C and C ⊤ do not need to be Hankel, but can be embedded into Hankel matrices with rows/columns corresponding to distances l q from the pivot, where l = {0, ..., mp} and mp q is the largest distance between a vertex and the pivot. Tensor X can also be padded into a larger one with extra rows/columns (corresponding to unrealized distances) set to zero. If p is constant, the asymptotic time complexity remains the same as in the previous case, but the algorithm might not be practical since the number of rows and columns grows by a multiplicative factor of p. For certain non-cordial f , the algorithm can be modified for potential gains.In this section we provide additional related works. One of the methods to tackle this problem is via iterative methods [Koutis et al., 2012] like Arnoldi iteration [Arnoldi, 1951] , Conjugate Gradient [Shewchuk, 1994] and the celebrated Spielman-Teng algorithm [Spielman and Teng, 2012] for symmetric diagonally dominant (SDD) matrices. There are a number of extensions and variations of the above methods [Blelloch et al., 2011 , Boman et al., 2008 , Christiano et al., 2010 , Koutis and Miller, 2007 , Spielman and Teng, 2008 , Daitch and Spielman, 2008 , Koutis and Miller, 2008] .They mainly take into account the structure of the matrix (SDD) [Koutis et al., 2010 [Koutis et al., , 2011a [Koutis et al., , 2012]] , embedding of a graph into low stretch spanning trees [Elkin et al., 2005] , graph sparsification [Spielman and Teng, 2010] and the choice of a good pre-conditioner [Maggs et al., 2003 , Koutis et al., 2011b] . We want to emphasize that the research on low stretch trees for general graphs is orthogonal to the main topic of this work. In our manuscript, we show in particular how to conduct efficient integration on arbitrary trees. Thus our work can be naturally combined with those algorithms to leverage all the above low stretch tree constructions for a better approximation of the graph's metric.The other class of method comes from the celebrated work of [Al-Mohy and Higham, 2011] and there are a number of extensions of this work [Kloster and Gleich, 2023 , Al-Mohy and Higham, 2010 , Moore, 2011 , Moler and Van Loan, 2003 , Auckenthaler et al., 2010] .Another class of methods is via sampling, where one samples a subset of a large matrix, which is then used to approximate the matrix-vector multiplication (i.e. Monte Carlo sampling) methods [Drineas et al., 2006 , Drineas and Kannan, 2001 , Acebron, 2019 , Acebron et al., 2019 , Benzi et al., 2017 , Martinsson, 2019] .We note that none of these methods are directly applicable in our cases as our f -matrix is neither Hermitian or SDD. The randomized algorithms are harder to use in the setting of training of a neural network. Moreover our method is exact on trees, where all the above methods are approximations.Algorithm 1 General Efficient Low-Rank Masked Attention from Choromanski et al. [2022] Input: Query/key matrices:Q, K ∈ R L×d QK , value matrix V ∈ R L×d , mask M ∈ R L×L , procedure FastMult M : R L → R L calculating Mx (or its approximation) for the input x ∈ R L , kernel feature map: ϕ : R d QK → R m . vec(•) denotes vectorization.Output: Masked low-rank attention embeddings using ϕ.1. Compute matrices V 1 ∈ R L×(md) , V 2 ∈ R L×m with rows defined as: V 1 i: = vec(ϕ(k ⊤ i )v i ), V 2 i: = ϕ(k ⊤ i ) ⊤ , where k i /v i stands for the ith row of K/V. 2. D1 := [FastMult M (V 1 :1 ), ..., FastMult M (V 1 :md )] ∈ R L×md , D2 := [FastMult M (V 2 :1 ), ..., FastMult M (V 2 :m )] ∈ R L×m for V 1/2 :i denoting ith column of V 1/2 . 3.Output the embedding r i of the ith tokens as:r i = ϕ(q ⊤ i ) ⊤ devec( D1 i: ) ϕ(q ⊤ i ) ⊤ ( D2 i: ) ⊤, where q i is the ith row of Q and devec(•) devectorizes its input back to R m×d .We now recall the formulation of general masked transformers.Let us denote by L the number of input tokens. The attention used in a regular Transformer linearly projects their representations into three learnable matrices Q, K ∈ R L×d QK , V ∈ R L×d called queries, keys and values respectively. Definition C.1 (general masked attention). General masked attention is given by the following equation, where M ∈ R L×L is the mask matrix, and A ∈ R L×L is the so-called masked attention matrix (MAM): which is defined as:EQUATIONwhere ⊙ denotes the element-wise (Hadamard) matrix product, K : R d × R d → R is some kernel function and K(Q, K) is a kernel matrix defined as: K(Q, K) i,j = K(q ⊤ i , k ⊤ j ) for the ith row q i of Q and the jth row k j of K respectively. We call A ′ = K(Q, K) the unmasked attention matrix (UAM). Note that when K is the softmax function, we recover the well-known attention mechanism in Transformers.Here 1 L is the all-ones vector of length L, and diag(•) is a diagonal matrix with the input vector as the diagonal. The time complexity of computing ( 9) is O(L 2 d). If the kernel K admits (at least in expectation) a dot-product decomposition, i.e. K(x, y) = E[ϕ(x) ⊤ ϕ(y)] for some mapping: ϕ : R d QK → R m (and some m > 0). ϕ(u) is called a (random) feature map (RFM) for u ∈ R d . For Q ′ , K ′ ∈ R L×m with rows given as ϕ(q ⊤ i ) ⊤ and ϕ(k ⊤ i ) ⊤ respectively, RFM-based kernel linearization leads directly to the efficient unmasked attention mechanism of the form: The central question in [Choromanski et al., 2022] was how to incorporate the masking in the linear attention as above. Note that in this case A ′ is never materialized. Building on the work of [Luo et al., 2021] , the authors [Choromanski et al., 2022] propose a general algorithm that efficiently implements masked linear attention.Att K (Q, K, V) = D -1 (Q ′ ((K ′ ) ⊤ V)), D = diag(Q ′ ((K ′ ) ⊤ 1 L )).In this work, we use different mappings ϕ (see Table 1 ). Our key contribution in this work is to propose a novel mask matrix M and the implementation of a fast matrix multiplication by M. The above result then allows us to construct novel classes of Topological Transformers.In this section, we provide additional details regarding the experimental setup and present additional results from our experiments. Our code is available at https://github.com/brcsomnath/ FastTreeIntegrator. Specifically, we provide there the code for: (1) our algorithm leveraging IntegratorTree data structure (depicted in Fig 1 ), ( 2) adaptation to the Gromov-Wasserstein-type computation, (3) graph classification and (4) experiments on interpolation on meshes.We present additional results for the training loss, relative Frobenius Norm Error (ϵ), for more samples from the Thingi10K dataset (to complement the results in Fig. 6 ). In Fig. 9 , we observe that in most cases having rational functions with higher polynomial degrees results in lower training loss.We also perform similar experiments for graph classification on the CUBES dataset Hanocka et al. [2019] . Specifically, we investigate how the polynomial degree affects the graph classification performance in Fig. 9 (left). We observe that increasing the polynomial degree improves the classification accuracy up to a certain degree. For the same dataset, we also compute the training loss for different polynomial degrees in Fig. 9 (right). Similarly, we observe that higher-degree rational functions achieve lower training loss for fitting the polynomial coefficients.Moreover, we benchmark FTFI on ModelNet10 [Wu et al., 2015] , a dataset for 3D Point Cloud (PC) classification. For each PC, we create an ϵ-neighborhood-graph and use FTFI for graph classification The Shortest Path kernel achieves an accuracy of 39.6%, whereas our FTFI with the degree-2 f(x)=x GRF(1) GRF( 2) GRF( 3) GRF( 4 Error GRF(1) GRF( 2) GRF( 3) GRF(4)Figure 9 : Left: Variation in FTFI performance with different f -distance functions on the CUBES dataset. We use general rational functions (GRF) of varying polynomial degrees. GRF(i) indicates a rational function of the i-th degree. We observe a general trend of accuracy increase with function complexity up to a certain degree. The coefficients of the GRF were learnt using a few graph instances.Right: We show the training loss curves for estimating the coefficients of the rational function, f , for samples in the CUBES dataset. We report the training loss for rational functions with varying polynomial degrees. We observe that the training loss is lower when we use rational functions with high-degree polynomials.polynomial improves the accuracy to 44.2% (10% relative improvement over the baseline), similar to the observation in 9. Wasserstein distance has found many uses in ML, particularly due to it's principled approach to compare probability distributions. Gromov Wasserstein Mémoli [2011] discrepancy is an extension of Wasserstein distance to graph structured data, with a lot of downstream applications like graph clustering and classification. Inspired by the work of [Choromanski et al., 2023] , we follow the exact same procedure in the integration of FTFI in the conditional gradient algorithm. The FTFI can be injected seamlessly in place of the Fast Matrix Multiplication (FMM) algorithms in Algorithm 2 and Algorithm 3 (see [Choromanski et al., 2023] ).Our method GW-FTFI run consistently 2-6x faster than the baseline methods using the Shortest Path kernel, with no drop in accuracy in computing the associated costs (Figure 10 ). The plots shown are obtained by averaging over 10 seeds and random trees of various sizes. For the baseline experiments, we use the implementation from the POT library [Flamary et al., 2021] .In this section, we present implementation details for the mesh interpolation experiments in Section 4.2. All experiments were run on a computer with an i9-12900k CPU and 64GB memory.In the vertex normal prediction task in Section 4.2, we choose 40 meshes for 3D-printed objects with a wide range of size from the Thingi10K [Zhou and Jacobson, 2016] dataset with the File IDs: [60246, 85580, 40179, 964933, 1624039, 91657, 79183, 82407, 40172, 65414, 90431, 74449, 73464, 230349, 40171, 61193, 77938, 375276, 39463, 110793, 368622, 37326, 42435, 1514901, 65282, 116878, 550964, 409624, 101902, 73410, 87602, 255172, 98480, 57140, 285606, 96123, 203289, 87601, 409629, 37384, 57084] For both our FTFI and the baseline BFFI methods, we do a grid-search over the hyper-parameter λ for each mesh and report the pre-processing time associated with the hyper-parameter(s) that give(s) us the best cosine similarity.We conduct graph classification experiments on a wide range of benchmark datasets. We report the dataset statistics for the graph classification datasets in Table 2 . More details about the datasets are available in Morris et al. [2020] . To evaluate the performance of the different kernels, we employ the framework proposed by [Errica et al., 2020] . In particular, 10-fold cross-validation is used to obtain an estimate of the generalization performance of our method and the baseline method. We repeat this cross validation experiment 5 times to get a robust estimation and report the standard deviation for each setup.To obtain graph features, we follow the approach presented in [de Lara and Pineau, 2018] . In this setting, we obtain the k-smallest eigenvalues from the approximated kernel from FTFI and forward these features to a random forest classifier for classification. For BGFI, we perform the same process obtaining the k-smallest eigenvalues from the exact shortest kernel. FTFI achieves similar performance to the BGFI while being significantly faster. We tune the hyperparameter k independently for each method.In Table 4 , we report the results for a wide range of baselines and compare FTFI. We observe that FTFI achieves competitive performance among various strong kernel-based classification baseline approaches. Note that FTFI results are not directly comparable with other approaches, as FTFI constructs an intra-graph kernel while other methods use inter-graph kernels. Despite the aforementioned considerations, we contend that positioning our results within the broader framework of alternative methodologies demonstrates that FTFI remains a compelling approach, owing to its speed and comparable classification accuracy.In this subsection, we provide additional training details for our image classification tasks. Table 5 and table 6 present the architectural as well as the training details.We train the ViT models starting from their pretrained checkpoint (pretrained on ImageNet-21k). We replace the dense attention in ViT by the Performer attention (see Equation 10). We use Algorithm 1 to efficiently incorporate the mask matrix M in the attention mechanism.We have already provided comparison with SOTA efficient-attention methods: low-rank attention Transformers in Sec 4.4, quality-wise. On standard ImageNet benchmark, our best Transformer with FTFI provide 78.15% accuracy, as compared to 76.37% of the best low-rank -attention variant (obtained by testing three different linear variants). That gives 1.78% accuracy improvement with only 3 extra trainable parameters per head (36 extra trainable parameters per layer). We have also run the experiments with cosFormer. It achieved 76.3% accuracy (consistent with what is reported in the literature, see [8] ), lower than both: our method and the best tested low-rank attention variant. The RF-Gate-Gaussian achieved 76.35% accuracy, which is is still lower than both: FTFI and the best tested low-rank attention variant.We have also conducted tests on another challenging dataset: Places365. In the paper, we report 1.71% accuracy improvement over low-rank attention Transformer (56.51% accuracy vs 54.8% accuracy). For the rebuttal, we also run the experiment with cosFormer which achieved 55.4% accuracy (consistent with what is reported in the literature, see: [8] ). This is still 0.93% behind our method. The RF-Gate-Gaussian achieved accuracy 55.1%, lower than this of cosFormer.I-naturalist is yet another challenging dataset, with 10K classes, diverse image quality and significant class imbalance. Transformer with FTFI provides 1% accuracy improvement over its regular lowrank attention counterpart and the cosFormer. Furthermore, FTFI achieved 0.8% improvement over RF-Gate-Gaussian. The convergence of the FTFI variant is 20-23% faster than this of its regular low-rank attention counterpart, the cosFormer and RF-Gate-Gaussian.ViViT ( [Arnab et al., 2021] ) is a novel architecture that adapts the Vision Transformer (ViT) for video processing. It efficiently handles the spatiotemporal dimensions of video data by factorizing the input and applying attention mechanisms across both space and time. This allows ViViT to capture complex motion patterns and long-range dependencies in videos.Applying FTFI with a topological masking mechanism to the ViViT architecture (factorized Transformer model variant, trained from scratch, as described in Arnab et al. [2021] ) results in a 0.8% absolute improvement on the Kinetics dataset ( [Kay et al., 2017] ). The experimental setup follows Arnab et al. [2021] . To the best of our knowledge, this is the first application of Topological Transformers to video data. 51.7 (± 3.7) 34.5 (± 3.0) 63.9 (± 3.0) 34.9 (± 1.7) 23.9 (± 1.2) 57.0 (± 1.2) ML 69.9 (± 4.8) 47.7 (± 3.2) 89.4 (± 2.1) 35.4 (± 2.0) OOM 75.6 (± 1.6) CORE-WL-VH 73.5 (± 6.1) 51.7 (± 4.1) 73.0 (± 4.5) 51.1 (± 1.6) 40.2 (± 1.8) 84.5 (± 2.0) CORE-SP 68.5 (± 3.9) 51.0 (± 3.5) 91.0 (± 1.8) OOT OOM OOT FTFI 65.1 (± 1.6) 46.4 (± 1.9) 83.7 (± 1.3) 43.8 (± 2.0) 31.8 (± 0.3) 63.7 (± 0.3) BGFI 65.1 (± 2.0) 47.6 (± 2.0) 84.3 (± 3.5) 44.0 (± 1.9) 37.6 (± 0.3) 75.5 (± 0. 81.4 (± 8.9) 16.7 (± 1.8) OOT 54.4 (± 9.8) OOM 69.5 (± 5.1) SP 82.4 (± 5.5) 37.3 (± 8.7) 72.5 (± 2.0) 60.2 (± 9.4) 77.9 (± 4.5) 74.9 (± 3.6) WL-VH 86.7 (± 7.3) 50.7 (± 7.3) 85.2 (± 2.2) 64.9 (± 6.4) 78.7 (± 2.3) 76.2 (± 3.5) WL-SP 81.4 (± 8.7) 27.3 (± 7.4) 60.8 (± 2.4) 54.5 (± 9.8) 76.0 (± 3.5) 72.1 (± 3.1) WL-PM 88.3 (± 7.1) 57.5 (± 6.8) 85.6 (± 1.7) 65.1 (± 7.5) OOM 75.9 (± 3.8) WL-OA 87.2 (± 5.4) 58.0 (± 5.0) 86.3 (± 1.6) 65.7 (± 9.6) 77.6 (± 3.0) 76.2 (± 3.9) NH 88.3 (± 6.3) 54.5 (± 3.6) 84.7 (± 1.9) 63.4 (± 9.2) 74.6 (± 3.5) 75.0 (± 4.2) NSPDK 85.6 (± 8.9) 42.2 (± 8.0) 74.3 (± 2.1) 59.1 (± 7.3) 78.9 (± 4.7) 72.5 (± 2.9) ODD-STh 80.4 (± 8.8) 32.3 (± 4.8) 75.2 (± 2.0) 59.4 (± 9.8) 76.4 (± 4.5) 70.9 (± 4.1) PM 85.1 (± 5.8) 43.2 (± 5.3) 73.5 (± 1.9) 60.2 (± 8.2) 77.9 (± 3.7) 70.9 (± 4.4) GH 82.5 (± 5.8) We do believe that the potential impact of this work is significant, as providing both: (a) theoretical advancements in structural graph theory as well as (b) practical applications in (1) designing computationally efficient Transformers leveraging topological inductive priors, (2) graph classification and (3) interpolation on manifolds. The core problem of fast multiplication with f -distance matrices plays an important role in various fields: physical sciences, chemistry, and network sciences. Our main contributions are algorithmic, with no clear negative side effects. While used in the context of Transformers, they should be though applied cautiously due to the nontrivial carbon emission footprint associated with training large Transformer models.Currently, FTFI can be applied on general graphs via certain classes of trees defined on these graphs (e.g. spanning trees), with low-distortion trees being more preferable. It would be interesting to see whether the main concepts used in the FTFI algorithm (such as the theory of balanced separators) can be directly incorporated into efficient and exact algorithms operating on general graphs (or general sparse graphs that appear in most machine learning applications). Determining general conditions on the classes of graphs and functions f under consideration that are sufficient for exact sub-quadratic time integration is yet another important problem for future work.See[Greengard and Lee, 2004] for an excellent introduction. 


Decision: All reviewers argue for acceptance, albeit with partial low scores. The improvements the paper has made are noted as well as the interesting theoretical contributions. The paper meets the bar for acceptance.

Reviews: review: paper_summary: paper_summary: This paper proposes faster matrix multiplication algorithms for a class of structured matrices, namely, f-distance matrices, where the $(i, j)^{th}$ entry of the matrix is $f(\text{dist}(i, j))$, for nodes $i, j$ in a graph. The authors propose an algorithm, referred to as Fast Tree-Field Integrators (FTFI), for the case of tree graphs (and in the experiments, for general graphs, the graph is approximated by its spanning tree).

The FTFI algorithm works as follows: first, an "integrator tree" is constructed (which is created only once for the f-distance matrix, and reused whenever this matrix is multiplied by another matrix). The integrator tree is a binary tree where each node is a subtree of the original tree. To obtain the two children of a node, they pick some vertex in the subtree (of the original tree) that this node corresponds to - this node will be called the "pivot point", and the subtree will be split into a "left" and "right" subtree at this pivot point.

This integrator tree is used to recursively compute the product of the $f$-distance matrix (denoted $M^G_f$) with another arbitrary matrix/tensor $X$. To compute the row of $M^G_f X$ corresponding to the vertex $v$, suppose $v$ is in the left subtree of the current node of the integrator tree. Then, we have to consider the contribution of rows of $X$ from vertices $j$ that are also in the left subtree (which is dealt with recursively) and vertices $j$ that are in the right subtree. In the latter case, the entry $f(dist(v, j))$ of the $f$-distance matrix can be written as $f(dist(v, p) + dist(p, j))$ where $p$ is the pivot point, and this property is leveraged to achieve a running time faster than naive matrix multiplication in several cases, such as when $f$ is a rational function or an exponential function.

The authors compare FTFI to BTFI (brute-force matrix multiplication in the tree setting), and other algorithms including BGFI (brute-force matrix multiplication without approximating the graph by a tree). The tasks studied include interpolation on meshes and graph classification. For interpolation on meshes, FTFI is the fastest in terms of preprocessing. It performs similarly to the SF algorithm (a state of the art algorithm for this problem), and performs worse than BGFI, which is a more brute force, but more accurate, algorithm. In the graph classification setting, FTFI gets similar accuracy as BGFI while being significantly faster.

The authors finally do experiments with FTFI on topological vision transformers. In a large scale setting, FTFI obtains a 7% improvement on ImageNet compared to the standard ViT performer (which obtains around 70% accuracy).
strength_weakness: strength_weakness: - There are some questions I have about the experimental results, mentioned below in the questions section. If these questions are addressed, I would be willing to raise the score.
questions: questions: - Why does BTFI require more time for processing than BGFI in Figure 4? Intuitively, they would require similar effort, and BGFI should perhaps be slower (with computing the shortest paths potentially having a longer running time for more general graphs than for trees).
- Also, given that BGFI does significantly better than BTFI/FTFI in terms of cosine similarity, at the cost of a roughly 5x slowdown compared to FTFI according to Figure 4, is it better to use BGFI compared to FTFI? What are the considerations here?
- Do the plots on the right-hand side of Figure 4 contradict the plots on the left-hand side of Figure 4? The second plot from the left in Figure 4 suggests that BGFI generally achieves greater cosine similarity than FTFI for various numbers of points (though it is slower) while the third plot from the left shows that FTFI is both faster and achieves higher cosine similarity than BGFI. How do you interpret the difference?
- In Figure 5, on PTC-MR, why is BGFI slightly faster than FTFI? Is this a specific distribution of graphs?
- For the experiments in Figure 6, the goal is to show that tree-based estimators can emulate integration on arbitrary graphs, using a rational function with quadratic numerator and denominator. I am not convinced that this experiment shows that tree-based metrics can emulate general graph metrics, since the MSE loss seems to plateau around 1. This seems a bit large given that the distribution of weights is taken from (0, 1). What is the distribution of graphs?

========================================

These questions are addressed by the rebuttal, and I am increasing my score from 6 to 7.
limitations: limitations: Yes.
paper_summary: paper_summary: The authors propose an algorithm for exact polylog-linear multiplication for general weighted trees and cordial functions f, which leads to a fast algorithm for distance-matrix tensor multiplication as used in transformers and graph kernels. The core of the algorithm is a binary tree structure called integration tree, which allows to perform the integration using an efficient divide and conquer scheme. The authors show that, for matrices defined by d-cordial functions, the integration can then be done by O(N log^{d+1} (N)). The authors show that their algorithm improves runtime against naive baselines in graph classification, vertex normal prediction, interpolation on meshes and topological vision transformers.

The work in general is far away from my research area. Thus, my confidence is low and I feel not qualified to judge the significance of  theoretical contributions in this work.
strength_weakness: strength_weakness: I agree with the authors that the tackled integration problem is relevant in many machine learning / deep learning models. It is unclear to me though how often the tree assumption is applicable. The practical experiments on topological vision transformers and graph classification, showing usefulness in practice. However, the experiments seem to be a bit inconclusive due to lacking trade-off analysis

What this paper lacks is an experiment that conclusively shows that the presented algorithm can be integrated into a state-of-the-art method (such as general masked transformers on multiple tasks) and that it has a positive impact on the quality/efficiency trade-off.

It would be also helpful if source code would be released in order to replicate/verify results.
questions: questions: How does the quality vs. efficiency trade-off behave in the topological vision transformers experiment? It is only shown that the quality slightly improves but not if the method has an impact on transformer efficiency.
limitations: limitations: The authors briefly discuss limitations in the appendix.
paper_summary: paper_summary: The paper tackles the problem of integrating tensor fields defined on graphs. The paper suggests Fast Tree-Field Integrators (FTFI) that integrate tensors on weighted trees with reduced time complexity, which is based on their data structure called IntegralTrees (IT). The idea is original and new, and the algorithm is carefully designed. All the properties of Fast Tree-Field Integrators are theoretically supported and experimentally verified.
strength_weakness: strength_weakness: The paper seems to consider the machine learning community as potential readers, but the underlying concepts are fairly unknown to the machine learning community. The reviewer also does not do research in integrating tensor fields on graphs, and it was a bit difficult to follow the concepts and structures in Fast Tree-Field Integrators (FTFI). I think this presentation problem can be greatly improved by providing examples or figures. For example, Section 3.2 would have been more easily grasped if the authors provided pictorial examples as Figure 1.

========================================

This weakness about the presentation of FTFI is addressed by the rebuttal, so I am increasing my score from 6 to 7.
questions: questions: In the experiments, the authors have only used minimum spanning tree (MST) as the tree approximation of a graph. On one side, this makes sense since the paper focuses on the fast computation of matrix-vector multiplication on a graph, not on choosing an appropriate tree. However, it seems to me that one reason of the relatively unsatisfactory accuracy (cosine similarity) of Fast Tree-Field Integrators (FTFI) is due to the distortion of the geometry of a graph by MST. As the authors have presented related work in embedding graphs to trees, I was wondering if we can get some improvements in accuracy when we use more advanced methods of embedding graphs to trees.

========================================

The question is addressed by the rebuttal.
limitations: limitations: The authos have addressed the limitations that Fast Tree-Field Integrators (FTFI) can work well provided that a graph is well approximated by a tree. I do not find any societal impact of this work.
paper_summary: paper_summary: In "Fast Tree-Field Integrators: From Low Displacement Rank to Topological Transformers", the authors propose a novel method for integrating scalar fields, and more generally tensor fields, on trees. Given an input tree and potentially multiple tensor fields, FTFI constructs an auxiliary binary "IntegratorTree" with nodes corresponding to sub-trees of the input tree. The efficiency of FTFIs is then demonstrated on a variety of experiments and tasks on synthetic and real-world datasets.
strength_weakness: strength_weakness: 1. Figure 5 claims that FTFI achieves similar accuracy as its brute force counterpart BGFI. However, BGFI seems to outperform FTFI on almost all datasets and often with accuracy gains of >1%.
2. Many graphs in practice are not trees. The paper does not discuss any theoretical guarantees or qualitative patterns on how well FTFI can be used as approximate integrators for almost tree-like and general graphs.
3. In the conclusion, speedups of 5--13x while maintaining quality are claimed. Both of these claims don't seem to be supported by all experiments and should be framed more cautiously.
questions: questions: In the introduction and related work, you talk about general graphs as an input (line 84), however, in section 3.1 you explicitly talk about the input tree. How do you turn the general graph into a tree? What information or what accuracy is lost in that process? Are there theoretical considerations that quantify this?

Thank you in advance for your answers!
limitations: limitations: The authors discuss limitations in the appendix.

Comments: title: Addressing comments of Reviewer H9MT
Addressing comments of Reviewer o1jE
Addressing comments of Reviewer Lag9
Addressing comments of Reviewer SX39
Additional responses for Reviewer Lag9
Additional responses for Reviewer H9MT
General comment
comment: I would like to thank the authors for addressing weakness and questions I proposed. I think all of them are addressed, so I increased the score from 6 to 7.
Dear Reviewer Lag9,

We would like to once more sincerely apologize for taking your time. **We believe we have addressed all Reviewer's comments**. In particular, we have conducted a comprehensive set of additional experiments, **showing conclusively** that Transformers with FTFI are SOTA efficient Transformers methods for image data. We also included new experiments with other modality. To be more specific, we have run additional experiments with Transformers on new image dataset (I-naturalist) as well as on a new modality, namely: videos. In those experiments we also tested many other efficient Transformer-architectures. These additional experiments confirm all our previous findings. In particular, FTFI applied in a topological masking mechanism in the ViViT architecture ([1], factorized Transformer model for videos, trained from scratch) leads to the +0.8% absolute improvement as compared on the Kinetics dataset ([2]). To the best of our knowledge, this is the first application of the Topological Transformers for video data.

[1] ViViT: A Video Vision Transformer, Arnab et al., ICCV 2021.

[2] The kinetics human action video dataset, Kay et al., 2017.  

For image modality, we have conducted a comprehensive comparison, not only with one linear-attention Transformer-model, but with **three different classes** of efficient Transformer-architectures, namely: (1) linear-attention methods based on low-rank decomposition (we tested several different instantiations with different attention-kernel functions) ([3]), (2) cosFormers ([4]) and (3) RF-Gate-Gaussian models ([5]). The FTFI-Transformer outperforms all of them quality-wise, while being at least as fast as those methods.

[3] Rethinking Attention with Performers, Choromanski et al., ICLR 2021.

[4] cosFormer: Rethinking Softmax in Attention, Qin et al., ICLR 2022.

[5] Gated Linear Attention Transformers with Hardware-Efficient Training, Yang et al., ICML 2024.


We thus would like to sincerely ask the Reviewer to comment on the rebuttal and update the score accordingly. If the Reviewer has any additional questions please let us know and we will be happy to answer them. **In particular, if there are any other specific experiments regarding Transformers experiments on image data the Reviewer would like to see, please let us know which specific experiments the Reviewer would like us to conduct and we will be more than happy to do it**. The discussion period is almost over thus we want to make sure that all additional comments the Reviewer will have (if any) can be promptly addressed by us.

Once more, thank you very much for your comments and feedback !

Yours sincerely,

The Authors
Dear Reviewer o1jE,

We would like to once more sincerely apologize to you for taking your time. We believe that we have addressed all Reviewer's questions, in particular: (1) provided additional figure explaining main algorithmic concepts (see: attached pdf) to improve presentation, (2) discussed results with more advanced trees (than minimum spanning trees), **which are actually already in the paper** (the Reviewer might have missed them so we clarified this point in the rebuttal). In addition, we have included in particular new experiments with Transformers on video input (as well as experiments with additional image datasets and speed tests). The experiments with this new data modality confirm all our findings that were already included in the paper.

Since the end of the discussion period is approaching, we would like to sincerely ask the Reviewer to comment on the rebuttal and update the score accordingly. If the Reviewer has any additional questions, please let us know so that we can address them before the end of the discussion period.

Once more, thank you very much for your reviews and feedback !

Yours sincerely,

The Authors
Thank you for your reply, which clarifies my questions - I will increase my score. In the final version, it would indeed be useful to have a version of the second plot in Figure 4 where the points do not overlap. The explanation of Figure 6 is helpful, since it gives a sense of what error can be typically achieved using tree-based metrics.
Thank you very much for your in-depth response and your proposed clarifications! I especially find the part on the almost-tree-likeness very interesting. I have no further questions.

I will raise my score as I believe this to be a very good paper. However, I am interested to learn of the opinion of the other reviewers and the AC, particularly from those with a stronger background in this topic.
Dear Reviewer H9MT,

We would like to once more sincerely apologize for taking your time. As we mentioned before, we believe we have addressed all Reviewer’s comments. We do hope that the Reviewer can update the score correspondingly. If the Reviewer has any additional questions, please let us know and we will be happy to address them. Thank you very much !

Yours sincerely,

The Authors
Dear Reviewer o1jE,

We would like to once more sincerely apologize for taking your time. As we mentioned before, we believe we have addressed all Reviewer’s comments. We do hope that the Reviewer can update the score correspondingly. If the Reviewer has any additional questions, please let us know and we will be happy to address them. Thank you very much !

Yours sincerely,

The Authors
Dear Reviewer Lag9,

We would like to once more sincerely apologize for taking your time. As we mentioned before, we believe we have addressed all Reviewer’s comments. We do hope that the Reviewer can update the score correspondingly. If the Reviewer has any additional questions, please let us know and we will be happy to address them. Thank you very much !

Yours sincerely,

The Authors
Dear Reviewer SX39,

We would like to once more sincerely apologize for taking your time. As we mentioned before, we believe we have addressed all Reviewer’s comments. We do hope that the Reviewer can update the score correspondingly. If the Reviewer has any additional questions, please let us know and we will be happy to address them. Thank you very much !

Yours sincerely,

The Authors
Dear Reviewer H9MT,

We would like to once more sincerely thank you for all the comments and very useful feedback. We think that we have addressed in depth all Reviewer's questions. Please let us know. If the Reviewer has any additional questions, we would be more than happy to answer them.

Yours sincerely,

The Authors
Dear Reviewer o1jE,

We would like to once more sincerely thank you for all the comments and very useful feedback. We think that we have addressed in depth all Reviewer's questions. Please let us know. If the Reviewer has any additional questions, we would be more than happy to answer them.

Yours sincerely,

The Authors
Dear Reviewer Lag9,

We would like to once more sincerely thank you for all the comments and very useful feedback. We think that we have addressed in depth all Reviewer's questions. Please let us know. If the Reviewer has any additional questions, we would be more than happy to answer them.

Yours sincerely,

The Authors
Dear Reviewer SX39,

We would like to once more sincerely thank you for all the comments and very useful feedback. We think that we have addressed in depth all Reviewer's questions. Please let us know. If the Reviewer has any additional questions, we would be more than happy to answer them.

Yours sincerely,

The Authors
**“It is unclear to me though how often the tree assumption is applicable”:**

Thank you for the comment. The quality of the general graph metric approximation with tree-metrics is a research area on its own, with voluminous corresponding literature. We provided a summary in Sec. 2 as well as in Appendix B. Notably, tree-metrics is a well-established tool, used in several applications, e.g. in distributed & online algorithms and biology (see for example: [1,2,3]). Theoretical analysis of tree-metrics is beyond the scope of this paper. In the paper we show that those tools can be effectively applied in machine learning. Rather than cherry-picking particular applications, where tree-metrics might be relevant, we provide extensive experimental evaluations for: graph classification, mesh modeling and Transformer training, where those methods were (to the best of our knowledge) never used before, showing the generality of the approach.

Furthermore, in Sec. 4.3 we provide a clear evidence that machine learning provides yet another path for effectively applying tree-metrics. We show that by applying relatively simple learnable nonlinear functions f (ratios of quadratic functions; only **six** trainable parameters in total) on the shortest path distances in the minimum spanning tree, one can substantially improve the quality of the approximation of the general graph metrics. That results in **9x+** more accurate approximation than this given by classic SOTA algorithmic and non-learnable approaches, such as Bartal trees ([4]) or Fakcharoenphol trees ([5]), providing log(N)lo log(N) and log(N) distortion ratio respectively. Importantly, as we show in this paper, integration with those nonlinear functions f can be done efficiently, in polylog-linear time. Thus we showed that machine learning methods can be effectively applied to improve the accuracy of the approximation with tree-metrics. To the best of our knowledge, this is the first such result.

[1] Efficient distributed approximation algorithms via probabilistic tree embeddings, Khan et al.,
     PODC 2008.

[2] K-server via multiscale entropic regularization, Bubeck et al., STOC 2018

[3] Distorted metrics on trees and phylogenetic forests, Mossel et al., IEEE ACM Trans. Comput.
     Biol. Bioinform, 2007. 

[4] On approximating arbitrary metrics by tree metrics, Bartal, Y., STOC 1998

[5] A tight bound on approximating arbitrary metrics by tree metrics, Fakcharoenphol et al., . J. Comput. Syst. Sci.

**Source code:**

Thank you very much for the comment. Following Reviewer’s suggestion, we provide the pointer to the anonymous repository with the code:

https://anonymous.4open.science/r/FastTreeFieldIntegrator/README.md.

That repository contains in particular: **detailed implementation** of the FTFI algorithm (key contribution of the paper) with several unit tests, as well as the code for graph classification and vertex-normal prediction tasks (we will release any remaining code upon the acceptance of the paper).
**"How do you turn the general graph into a tree ?". Information lost and theoretical analysis**:

Thank you very much for the comment. In practice, we construct minimum spanning trees (MSTs) for those graphs. In principle, trees providing low distortion of the general graph metric can be used (the best algorithmic constructions provide log(N)-distortion ratio, where N is the number of nodes of the graph). Low distortion trees is a subject of an extensive research, as we show in in Sec. 2 and Appendix B. However with learnable functions f applied for the shortest-path lengths in the trees, the approximation guarantees can be substantially improved, as we demonstrate in Sec. 4.3. We provide additional theoretical discussion (also regarding near minimum spanning trees) in other responses to Reviewer H9MT, in particular below.

**Theoretical guarantees for almost-tree like graphs:**

This is a second part of the response to Reviewer H9MT from the paragraph: **FTFI to approximate almost-tree like graphs and general graphs: theoretical analysis** in the rebuttal for Reviewer H9MT below.

Following Reviewer’s comment, we would like to give examples of the almost-tree like graphs, where our method can be applied to produce integration of arbitrary precision in polylog-linear time, as long as all edge-weights are taken from the bounded interval. We also provide detailed sketch of the proof below and will add the full proof in the camera-ready version of the paper.

The class of graphs we consider are called bounded connected treewidth graphs, e.g. the graphs that admit a treewidth decomposition with all bags being connected graphs of bounded size. We refer the Reviewer to [1] for an introduction to the theory of treewidth.
To see that these are natural candidates for almost tree-like graphs, note that in the treewidth decomposition of trees every bag is a graph corresponding to an edge (thus connected and of two nodes). It turns out that a treewidth decomposition of graphs with bounded connected treewidth, where all bags are connected and bounded can be found in linear time.
Furthermore, in such graphs a balanced separators of constant size can be found in linear time (it can be found among bags of the treewidth decomposition). Now note that the key trick of the FTFI algorithm is a divide-and-conquer strategy, where a “pivot node” is found that splits the tree into two substantial parts. For bounded connected treewidth graphs, this node is simply replaced by the balanced separator. In the regular FTFI algorithm, vertices are then partitioned based on the distance from the pivot node. For bounded connected treewidth graphs, we instead apply two levels of partitioning. For every vertex, we first register its shortest path distance d to the balanced separator (defined as the distance to the closest vertex of the separator) and then we compute its so-called signature vector (d1,...,ds), where d+di is the shortest distance to the ith vertex of the separator and s is the number of all vertices of the separator. Note that since the separator is of constant size and connected and furthermore the weights are taken from the bounded set, (d1,...,ds) is taken from the bounded s-dimensional cube C. We then construct an epsilon-net N to cover that cube and assign for each vertex of the graph the node of the epsilon-net that is closest to its signature vector. Note that the number of nodes of the epsilon-net is constant. To effectively conduct the divide-and-conquer strategy in the setting with three sets of vertices: A, S and B, where S is s separator that separates A and B, we proceed as follows. We choose a pair of nodes (n1, n2) in the epsilon-net N. We then consider vertices from A for which n1 was assigned and vertices from B for which n2 was assigned (first level of partitioning). We then partition those vertices based on the distance from the separator S (second level of partitioning). Then the corresponding cross-terms can be computed in the completely analogous way as from the analysis in Sec. 3.2. We repeat this procedure for all pairs of nodes from the epsilon-net N. That completes the sketch of the proof since the size of the epsilon-net N is constant, and thus the number of all the pairs of nodes of the epsilon-net N is also constant.

[1] Parameterized Algorithms, Pilipczuk et al., Springer 2015.
We would like to sincerely thank all the Reviewers for their very valuable feedback and comments. We addressed all of them in the rebuttals below and additional official responses. We also want to emphasize that we provide the pointer to the anonymous repository with the code:

https://anonymous.4open.science/r/FastTreeFieldIntegrator/README.md.

That repository contains in particular: **detailed implementation** of the FTFI algorithm (key contribution of the paper) with several unit tests, as well as the code for graph classification and vertex-normal prediction tasks (we will release any remaining code upon the acceptance of the paper).

Finally, we have also run additional experiments with Transformers on new image dataset (I-naturalist) as well as on a  **new modality**, namely: videos. In those experiments we also test many other efficient Transformer-architectures. These additional experiments **confirm all our previous findings**. In particular, FTFI applied in a topological masking mechanism in the ViViT architecture ([1], factorized Transformer model for videos,  trained from scratch) leads to the **+0.8% absolute** improvement as compared on the Kinetics dataset ([2]). **To the best of our knowledge, this is the first application of the Topological Transformers for video data**. 

[1] ViViT: A Video Vision Transformer, Arnab et al., ICCV 2021.

[2] The kinetics human action video dataset, Kay et al., 2017.

We also conducted additional speed tests for Transformers to describe speed-quality trade-offs, showing that our methods are as fast as previous efficient-attention Transformers, while providing better quality.