Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
Authors: Maksim Zhdanov
Max Welling
Jan-Willem van de Meent
Venue: ICML 2025 poster (2025)
Keywords: hierarchical transformer
ball tree
physical simulations
tree-based methods
window attention
AI4Science
Primary area: deep_learning
TL;DR: We present a hierarchical transformer that uses ball tree partitioning to process physical data on irregular grids in linear time while capturing both local and global interactions.

Abstract: Large-scale physical systems defined on irregular grids pose significant scalability challenges for deep learning methods, especially in the presence of long-range interactions and multi-scale coupling. Traditional approaches that compute all pairwise interactions, such as attention, become computationally prohibitive as they scale quadratically with the number of nodes. We present Erwin, a hierarchical transformer inspired by methods from computational many-body physics, which combines the efficiency of tree-based algorithms with the expressivity of attention mechanisms. Erwin employs ball tree partitioning to organize computation, which enables linear-time attention by processing nodes in parallel within local neighborhoods of fixed size. Through progressive coarsening and refinement of the ball tree structure, complemented by a novel cross-ball interaction mechanism, it captures both fine-grained local details and global features. We demonstrate Erwin's effectiveness across multiple domains, including cosmology, molecular dynamics, and particle fluid dynamics, where it consistently outperforms baseline methods both in accuracy and computational efficiency.

Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems
Maksim Zhdanov
Max Welling Jan- Wan-Willem van de Meent
Abstract
MHSA
MHSA
Large-scale physical systems defined irregular grids pose significant scalability challenges for deep learning methods, especially in the presence of long-range interactions and multi-scale cou- pling. Trrddiiooall approaches that compute all pairwise interactions, such as attention, become computationally prohibitive as they scale quadrati- cally with the number of nodes. We present Erwin. hierarchical transformer inspired by methods from computational many-body physics, which combines the efficiency of tree-based algorithms with the expressivity of attention mechanisms. Erwin employs ball tree partitioning to organize computation, which enables linear-time attention by processing nodes in parallel within local neigh- borhoods of fixed size. Through progressive coarsening and refinement of the ball tree struc- ture, complemented by novel cross-ball inter- action mechanism, it captures both fine-grained local details and global features. We demonstrate Erwin's effectiveness across multiple domains, including cosmology, molecular dynamics, PDE solving, and partiele fluid dynamics, where it con- sistently outperforms baseline methods both in accuracy and computational efficiency.
layer
layer 1
layer
runtime vs. size
runtime
receptive field
62
62
2.6x
(Ous)
Viewins
32
Strpet mode Bowis 10ad) MUNS
24
Broonn
MPNN
1k
10k
20k
Figure I. Top: Ball attention over a molecular graph. Multi- head self-attention (MHSA) is computed in parallel at fixed hi- erarchy levels (bold circles). the following layers, progressively coarsened to learn global features, while the partition size is fixed. Bottom: Computational advantages of our model.
tree
Such large-scale systems pose a significant challenge to ex- isting methods that were developed and validated at smaller scales. For example, in computational chemistry, models are typically trained on molecules with tens of atoms (Kovács et al., 2023), while molecular dynamics simulations often exceed thousands of atoms. This scale disparity might result in prohibitive runtimes that render models inapplicable in high-throughput scenarios such as protein design (Watson et al., 2023) screening (Fu et al., 2022).
1. Introduction
Scientific deep learning is tackling increasingly computa- tionally intensive tasks, following the trajectory of com- puter vision and natural language processing. Applications range from molecular dynamics (MD) (Arts et al., 2023) and computational particle mechanics (Alkin et al., 2024b) to weather forecasting (Bodnar et al., 2024), where simulations often involve data deñned on irregular grids with thousands millions of nodes, depending on the required resolution and complexity of the system.
A key challenge in scaling to larger system sizes is that computational methods that work well at small scales break down at larger scales. For small systems, all pairwise inter- actions can be computed explicitly, allowing deep learning models focus on properties like equivariance (Cohen & Welling, 2016). However, this brute-force approach be- comes intractable as the system size grows. At larger scales, approximations are required to efficiently capture both long- range effects from slowly decaying potentials and multi- scale coupling (Majumdar et al., 2020). As a result, models validated only on small systems often lack the architectural components necessary for efficient scaling.
AMLab, University of Amsterdam CCuspAI. Correspondence to: Maksim Zhdanov <m.zhdanov uva.nl>.
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Erwin Transformer
This problem has been extensively studied in computational many-body physics (Hockney & Eastwood, 2021), where the need for evaluating long-range potentials for large-scale particle systems led to the development of uub-quadratic tree-based algorithms (Barnes & Hut, 1986; Carrier et al., 1988). These methods are based on the intuition that distant particles can be approximated through their mean-field ef- feet rather than individual interactions (Pfalzner & Gibbon, 1996). The computation is then structured using hierarchical trees to efficiently organize operations at multiple scales. While highly popular for numerical simulations, these tree- based methods have seen limited adoption in deep learning due to poor synergy with GPU architectures.
eral approaches (Liu al., 2023; Sun et al., 2022) transform point clouds into sequences, most notably Pranoorra v3 (PTv3) (Wu et al., 2024b), which projects points into vox- els and orders them using space-filling curves (e.g., Hilbert curve). While scalable, these curves introduce artificial discontinuitics that can break local spatial relationships.
Particularly relevant to our work are hierarchical attention methods. In the context of ID sequences, approaches like the H-transformer (Zhu & Soricut, 2021) and Fast Multipole Attention (Kang et al., 2023) approximate self-attention through multi-level decomposition: tokens interact at full resolution locally while distant interactions are computed using learned or fixed groupings at progressively coarser scales. For point clouds, OctFormer (Wang. 2023) converts spatial data into a sequence by traversing an octree, ensuring spatially adjacent points consecutive in memory. While conceptually similar to our approach, OctFormer relies on computationally expensive octree convolutions, whereas our utilization of ball trees leads to significant efficiency gains.
Transformers (Vaswani et al., 2017), on the other hand, employ the highly optimized attention mechanism, which comes with the quadratic cost of computing all-to-all inter- actions. In this work, we combine the efficiency of hierarchi- cal tree methods with the expressivity of attention to create a scalable architecture for processing large-scale particle systems. Our approach leverages ball trees to organize com- putation at multiple scales, enabling both local accuracy and global feature capture while maintaining linear complexity in the number of nodes.
Rather than using a hierarchical decomposition, another line of work proposes cluster attention (Janny et al., 2023; Alkin et al., 2024a; Wu et al., 2024a). These methods first group points into clusters and aggregate their features at the cluster centroids through message passing or cross-attention. After computing attention between the centroids, the updated features are then distributed back to the original points. While these approaches yield the quadratic cost only in the number of clusters, they introduce an information bottleneck at the clustering step that may sacrifice fine- ane-grained details and fail to capture features multiple scales limitation our hierarchical approach aims to overcome.
The main contributions of the work are the following:
We introduce ball tree partitioning for efficient point cloud processing, enabling linear-time self-attention through localized computation within balls at different hierarchical levels.
We present Erwin, a hierarchical transformer that pro- cesses data through progressive coarsening and re- finement of ball tree structures, effectively capturing both fine-gra ane-grained local interactions and global features while maintaining computational efficiency.
Beyond attention, several alternatives have been developed for processing large-scale systems with irregular geome- tries. Message-passing neural networks (Gilmer et al., 2017) typically address scalability through multi-level graph rep- resentations (Lam et al., 2023; Cao et al., 2023; Valencia et al., 2025). Additionally, recent works have explored sub-quadratic convolution-based architectures like Hyena (Moskalev et al., 2025) and state-space models (Zhang et al., 2025), which offer promising alternatives for efficient pro- cessing of geometric data without the computational over- head of attention mechanisms.
We validate Erwin's performance across multiple large- scale physical domains:
Capturing long-range interactions (cosmology) - Computational efficiency (molecular dynamics) Expressivity on large-scale phenomena (PDE benchmarks, turbulent fluid dynamics)
achieving state-of-the-ard performance in both compu- tational efficiency and prediction accuracy.
2. Related Works: sub-quadratic attention
3. Background
One way to avoid the quadratic cost of self-attention is to linearize attention by performing it on non-overlapping patches. For data on regular grids, like images, the Swin Transformer (Liu et al., 2021) achieves this by limiting at- tention to local windows with cross-window connections enabled by shifting the windows. However, for irregular data such as point clouds or non-umiform meshes, one first needs to induce a structure that will allow for patching. Sev-
Our work revolves around attention, which we aim to lin- earize by imposing structure onto point clouds using ball trees. We formally introduce both concepts in this section.
3.1. Attention
The standard self-attention mechanism is based on the scaled dot-product attention (Vaswani et al., 2017). Given a set X
2

---

Erwin Transformer
of N input feature vectors of dimension C, self-attention is computed
and its radius is determined by the furthest point it contains:
TB mar - call2 p B1 U B2}
Q K. V XW. XWk. XW
QK" VC + B
(L)
where [B] denotes the number of points contained in B.
Att(Q, Q,K, K, V) = softmax
To construct the ball tree, we recursively split the data points into two sets starting from P. In each recursive step, we find the dimension of the largest spread (i.e., the max- min value) and split at its median (Pedregosa al., 2012), constructing covering balls per Def.3.2. For details, Appendix Alg.1'.
where W WR. W E Rxx' are learnable weights and RN is the bias term.
Multi-heade self-attention (MHSA) improves expressivity by computing attention H times with different weights and concatenating the output before the final projection:
Tree Completion want to work with perfect binary trees, i.e., trees where all internal nodes have exactly two children and all leaf nodes appear at the same depth. To achieve this, pad the leaf level of a ball tree with virtual nodes, yielding the total ceil cei((log2nn)..
To enable efficient implementation, we
( = I.. Yu) w
Y, = Att XXW, XW, R. XW:)
where [....., , denotes concatenation along the feature di- mension, and W y. W W E CCx(("/H) and wo E RCx"" are learnable weights.
number of nodes 2 where m
3.2.1. 1. BALL TREE PROPERTIES
The operator explicitly computes interactions between all elements in the input set without any locality constraints. This yields the quadratic computational cost w.r.t. the input set size (N2)) Despite being heavily optimized (Dao, 2024), this remains bottleneck for large-scale applications.
In the context of our method, there are several properties of ball trees that enable efficient hierarchical partitioning:
Proposition 3.3 (Ball Tree Properties). The ball tree T constructed as described satisfes the following properties:
3.2. Ball tree
The tree a perfect binary tree.
ball tree is hierarchical data structure that recursively par- titions points into nested sets of equal size, where each is represented by ball that covers all the points in the set. As- sume operate the d-dim. Euclidean space (I24) 11 l12) where we have a point cloud (set) P (PI,... Pn. Rd.
2. At each level i, each ball contains exactly 2' leaf nodes.
3. Balls at each level cover the point set
U B= BEL,
Vi (O,..., m.
Definition 3.1 (Ball). A ball is a region bounded by hy- persphere in Red Each ball is represented by the coordinates of its center e E _Rd and radius r E IR.::
Proposition 3.4 (Contiguous Storage). For ball tree
T = {Lo, LI....,mmi point cloud P {Pi,... ........ Pu there exists bijective mapping т [1....,m} - (1,..., n. such that points belonging to the same ball B E Li have contiguous indices under x.
B B(c,r) r) {z IRd ||z c|22 sr..
We will omit the parameters (c,r) for brevity from now Definition 3.2 (Ball Tree). A ball tree T point set P is hierarchical sequence of partitions (Le, Li...., Lm where each level I, consists of disjoint balls that cover P. At the leaf level = 0, the nodes are the original points:
As a corollary, the hicrarchical structure at each level can be represented by nested intervals of contiguous indices:
Example. Let P {pi, ... Ds}, then ball tree T (To, In L2, La} is stored after the permutation as
Lo ((p,) P)
L3
For each subsequent level i > 0, each ball B E L, is formed by merging two balls at the previous level B1, B2 E L1-1:
L2
T
L B1 U B2) B1, Bg E L-1)
(4)
Lo === (P))
Pa Ps Pe Pa P. P/ Ps PA
such that its center is computed as the center of mass:
Note that since we split along coordinate axes, the resulting structure depends on the orientation of the input data and thus breaks rotation invariance. We will rely this property in Sec- tion 4.1 to implement brlls-ball connections.
B1|c, + B2lc2
CB
[B| Ba
3

---

Erwin Transformer
The contiguous storage property, combined with the fixed size of balls at each level, enables efficient implementation through tensor operations. Specifically, accessing any ball B E L, simply requires selecting contiguous sequence of 2' indices. For instance, in the example above, for i = 2, we select a:d and e:h to access the balls. Since the balls are equal size, we can simply reshape Lo to access any level. This representation makes particularly efficient to implement our framework's ' core operations ball attention and coarsening/refinementt which we will introduce next.
We then compute self-attention independently on each ball?:
X/3 BAtt(XB) == Att(Xg X WR, W.) (8)
where weights are shared between balls and the output X'B maintains row correspondence with XB.
Computational cost As attention is computed indepen- dently for each ball B E LR. the computational cost is reduced from quadratic to linear. Precisely, for ball atten- tion, the complexity is O(BI|2 TH) ). i.e., quadratic in the ball size and linear in the number of balls:
Another important property of ball trees is that while they cover the whole point set, they are not required to partition the entire space. Coupled with completeness, means that at each tree level, the nodes are essentially associated with the same scale. This contrasts with other structures such as octrees that cover the entire space and whose nodes at the same level can be associated with regions of different sizes:
Attention (n
Ball Attention (n)
Ball tree
Oct-tree
Figure 4. For highlighted points, standard attention computes inter- actions with all other points the point cloud, while hall attention only considers points within their balls.
Positional encoding We introduce positional information to the attention layer in two ways. First, augment the fea- tures of leaf nodes with their relative positions with respect to the ball's center of mass (relative position embedding):
Figure Ball tree vs. octree construction. Colors highlight the difference in scales for nodes including the same number of points.
RPE:
XB = X + (PB cBB
-B) W pos
(9)
4. Erwin Transformer
where Pn contains positions of leaf nodes, Cn is the center of mass, and Wpos the layer to incorporate geometric structure within each ball.
Following the notation from the Background Section 3.2, we consider a point cloud p = {PI...., pa] C Rd. Addi- tionally, each point is now endowed with feature vector yielding feature set X {x,,..., x,, RC
pos is a learnable projection. This allows
Second, we introduce a distance-based attention bias:
On top of the point cloud, we build a ball tree T = {Lo.... .... }. We initialize Lieat == Lo to denote the cur- rent finest level of the tree. As each leaf node contains a
BB = - |C - CB" B'BB" B" E leavesg (10)
with a learnable parameter o E R (Wessels et al., 2024).
single point, it inherits its feature vector:
The term decays rapidly as the distance between two nodes increases, which enforces locality and helps to mitigate potential artifacts from the tree building, particularly in cases where distant points are grouped together.
Xteff (xB =; B {p:} TAcas)
(5)
4.1. Ball tree attention
Ball attention For each ball attention operator, we specify a level k of the ball tree where each ball B E Le contains 24 leaf nodes. The choice of & presents a trade-off: larger balls capture longer-range dependencies, while smaller balls are more resource-effiieett. For each ball B E Lg, we collect the leaf nodes within B:
Cross-ball connection To increase the receptive field of our attention operator, we implement cross-bar ball connections inspired by the shifted window approach in Swin Trans- former (Liu et al., 2021). There, patches are displaced diagonally by half their size to obtain two different image partitioning configurations. This operation can be equiva- lently interpreted as keeping the patches fixed while sliding the image itself.
leaves {B' E Lient BB
(6)
along with their features from Xteat:
FFor any set of vectors X, we abuse notation by treating X (7) matrix with vectors as its rows.
XB B = {XB' E Xieat B' E leavves)
4

---

Erwin Transformer
point cloud P
ball tree, level Lu
coarsening
Emhedding
Erwin Transformer
Point Cloud
Ball Tree
@ @
MPNN
Encoder
Decoder
Bottleneck
coarsened tree
ball tree, level Li+11
refinement
ErwinLayer
XS
rotate
ball tree
ErwinBlock
Coursesing
1.Norm
Attention
1.Norm
SwiG
original balls
leaves
"rotated" balls
Figure Overview Erwin. Left: A sequence two attention layers with intermediate tree coarsening. In every layer, attention computed on partitions of size 16, which correspond to progressively higher levels of hierarchy. Center (top): Coarsening and refinement of a ball tree. Center (bottom): Building a tree on top a rotated configuration for cross- ball interaction. Right: Architecture of Erwin.
Following this interpretation, we rotate the point cloud and construct the second ball tree Tros = [T.0 ..... 1. J. which induces permutation not of Ieaf nodes (see Fig. 3, center). We can then compute ball attention on the rotated config- uration by first permuting the features according to not applying attention, and then permuting back:
coarsening ball tree x rearrange [x, rel. rel.pas), " (n 2 di - n. 12 d1 W. pos = reduce (poa, "In 2 - n d', "mean"]
The inverse operation, refinement, allocates information from coarse representation back to finer scales. More precisely, for ball B E Lx, its features are distributed back to the nodes at level L contained within B as:
X noo (BAtt (Frot (XB)))
(11)
{x B' E Lx-1) [XB, PB ca) W, (13)
By alternating between the original and rotated configura- tions in consecutive layers, ensure the interaction be- tween leaf nodes in otherwise separated balls.
where P3 contains positions of all nodes at level k - within ball B with center of mass CB, and W, Cx'' is learnable projection. After refinement, Lieaf and Xleaf are updated accordingly. In pseudocode:
Tree coarsening/reffimmennt For larger systems, are interested in coarser representations to capture features at larger scales. The coarsening operation allows to hier- archically aggregate information by pooling leaf nodes to the centers of containing balls at levels higher (see Fig. 3, top, = 1). Suppose the Icaf level k. For every ball B E Li+1, we concatenate features of all interior leaf nodes along with their relative positions with respect to cg and project them higher-dimenssional representation:
refining ball tree x [zearraage (x, *n (2 d) + (a 2f) ar), rel.pos] e W.
4.2. Model architecture
We are now ready to describe the details of the main model to which we refer as Erwin (see Fig. 3) hierarchical transformer operating on ball trees.
XB
X CB'
CB]
W
(12)
Embedding At the embedding phase, we first construct ball tree on top of the input point cloud and pad the leaf layer to complete the tree, as described in Section 3.2. To capture local geometric features, we employ a small-scale MPNN, which is conceptually similar to PointTransformer's embedding module using sparse convolution. When input connectivity is not provided (e.g., mesh), utilize the ball tree structure for fast nearest neighbor search.
Cleavve
where ④ denotes leaf-wise concatenation, and W. E C2' (C++d) is a learnable projection that increases the fea- ture dimension to maintain expressivity. After coarsening, balls at level /++ become the new leaf nodes, Lieat := Lx+1, with features Xieaf = {xx B E T.+1). To highlight the simplicity of our method, we provide the pseudocode":
"We pay homage to Swin Transformer our model is hased on rotating windows instead sliding, hence Rwin - Erwin.
We use e nops (Rogozhnikov, 2022) primitives.
5

---

Erwin Transformer
Runtime scaling
Receptive field size
Sculing with training sot size
1.2
Roatino tims, log poabo)
Oyeieiny (8 9999
Target mode Erwin (Ours)
equivarrian
non-quisstionn
10-
1.1
SEGNNN
Featn-5
Rel The bnzidng
(breaed pum))
MPNN
1.0
Neguil
SEGNN (bams 2)
Ernin-h (Ours)
Test: MSB
NegulP frem =2)
MPNN
0.0
1124
0.7
Size single data point (log scale)
Figure 5. Left: Computational cost of Erwin. We split the total runtime into building a ball tree and running model. The input is batch of 16 point clouds, each of size We fit power law which indicates close to linear scaling. Right: Receptive field 800. A node is in the receptive field changing its features affects the target node's output, MPNN consists of 6 layers, each node connected to 16 nearest neighbours.
64
612
2, 2:048
8, 192192
Trining act size
Figure 6. Test meen--suared error (MSE) on the predicted veloci- ties as a function of training set size for the cosmology task, 5 runs per point. Point transformers indicate favourable scaling surpass- ing graph-based models with sufficiently many training samples.
of MPNN Erwin,
Computational cost To experimentally evaluate Erwin's scaling, we learn the power-law' form Runtime C' by first applying the logarithm transform to both sides and then using the least squares method to evaluate B. The result is an upproximately linear scaling with 3=1.054 with =0909 see Fig. 5. Ball tree construction accounts for <5% of the overall time (see Tables 7,8), proving the efficiency of our method for linearizing attention for point clouds.
ErwinBlock The core building block of Erwin follows a standard pre-norm transformer structure: LayerNorm fol- lowed by ball attention with a residual connection, and a SwiGLU feed-forward network (Shazeer, 2020). For the ball attention, the size 2 of partitions a hyperparameter. To ensure cross-ball interaction, alternate between the original and rotated ball tree configurations, using even number of blocks per ErwinL in our experiments.
Receptive field One of the theoretical properties of our model is that with sufficiently many layers, its receptive field is global. To verify this claim experimentally, for an arbitrary target node, run the forward pass of Erwin and MPNN and compute gradients of the node output with respect to all input nodes' features. If the gradient is non- zero, the node considered to be in the receptive field of the target node. The visualization is provided in Fig. 5, right, where we compare the receptive field of our model with that of MPNN. As expected, the MPNN has a limited receptive field, as it cannot exceed N hops, where N is the number of paassig-p layers. Conversely, Erwin implicitly computes all-to-all interactions, enabling it to capture long-range interactions in the data.
Overall architecture Following a UNet structure (Ron- neberger et al., 2015; Wu et al., 2024b), Erwin processes features at multiple scales through encoder and decoder paths (Fig. 3, right)) The encoder progressively coarsens the ball tree while increasing feature dimensionality to maintain expressivity. The coarsening factor is hyperparameter that takes values that are powers of 2. At the decoder stage, the representation is refined back to the original resolution, with skip connections from corresponding encoder levels enabling multi-scale feature integration.
5. Experiments
The code is available at https://qithut com/ mrxxxxdn/erwii We summarize datasets in Table and provide implementation details in Appendix B.
5.1. Cosmological simulations
To demonstrate our model's ability capture long-range interactions, we use the cosmology benchmark (Balla et al., 2024), which consists of large-scale point clouds represent- ing potential galaxy distributions.
Table 1. Summary of benchmark datasets. i indicates varying size, for which we report the average number of nodes.
GEOMETRY
BENCHMARKS
#DIM
#NODES
POINT CLOUD
ELASTICITY
2D
972
COSMOLOGY
3D
5,000
The dataset is derived from N-body simulations that evolve dark matter particles from the early universe to the present time. After the simulation, gravitationally bound structures (halos) are identified, from which the 5,000 heaviest ones are selected as potential galaxy locations. The halos form local clusters through gravity while maintaining
Dataset
MD
3D+TIME
890
STRUCTURED
PLASTICITY
2D+TIME
3,131
AIRFOIL
2D
11,271
PIPE
2D
16,641
UNSTRUCTURED SHAPE-NE CAR
3D
32,186
EAGLE!
2D+TIME
3,388
'We only use data for 2 1024 to exclude overhead costs.
6

---

Erwin Transformer
Performance vs. Runtine
Table 2. Ablation on the cosmology task. Increasing window size
0.73
improves performance at the cost of slower runtime (Erwin-S).
BALL SIZE
256
128
64
32
Tost NLL
0.72
IBN
TEST Loss
0.595
0.603
0.612
0.620
0.71
RUNTIME, MS
229.6
165.2
135.3
126.0
0.70
Ervin (Ours)
Table 3. Ablation study on architectural choices: using MPNN in embedding, RPE, and cross-ball connection via rotating trees.
Point aonfoooux MPNN PaintNat+
0.69
TEST Loss
0.75
1.00
1:25
1.50
1.75
2.00
2:25
2.50
MODEL
Speedup over MPNN (med/um) times
MD
SHAPENET-CAL
Figure 7. Test negative log-likelingg log-likeldhond (NI..) of the predicted ac- celeration distribution for the molecular dynamics task (averaged over 3 runs). The baseline MPNN is taken from (Fu et al., 2022). The size of the markers reflects the number of parameters.
w/o
0.738
30.39
+ MPNN
0.720
30.49
RPE
0.715
30.02
+ ROTATING TREE
0.712
15.85
range correlations that originated from interactions in the early universe, reflecting the initial conditions.
ate on the original representation but improve the runtime by employing our hardware-efficint model. Therefore, the question we ask is how much we can accelerate simulation w.r.t. MPNN without compromising the performance.
Task The input is a point cloud X E R5000033 where each row corresponds to a galaxy and each column to x, y, coordinate respectively. The task is a regression problem to predict the velocity of every galaxy Y E IR 1000003, We vary the size of the training dataset from 64 to 8,192, while the validation and test datasets have a fixed size of 512. The models are trained using the mean squared error between predicted and ground truth velocities.
Dataset The dataset consists of single-chain coarse- grained polymers (Webb et al., 2020; Fu et al., 2022) sim- ulated using MD. Each system includes types of coarse- grained beads interacting through bond, angle, dihedral, and non-bonded potentials. The training set consists of poly- mers with repeated bead patterns while the test set polymers are constructed by randomly sampling bead sequences, thus introducing a challenging distribution shift. The training set contains 100 short trajectories (50k T), while the test set contains 40 trajectories that are 100 times longer.
The results are shown in Fig. 6. We compare against multiple equivariant (NequIP (Batzner et al., 2021), SEGNN (Brandstetter et al., 2022)) and non-equivariant (MPNN (Gilmer et al., 2017), PTv3 (Wu et al., 2024b)) baselines. In the small data regime, graph-based equivariant models are preferable. However, as the training set size in- creases, their performance plateaus. We note that this is also the case for non-equivariant MPNNs, suggesting the issue might arise from failing to capture medium to large-scale interactions, where increased local expressivity has mini- mal impact. Conversely, transformer-basee models scale favorably with training set size and eventually surpass graph- based models, highlighting their ability to capture both small and surge-scale interactions. Our model demonstrates par- ticularly strong performance and significantly outperforms other baselines for larger training sets.
Results
Task We follow the experimental setup from Fu al. (2022). The model takes as input a polymer chain of N coarse- graeee-rriedd beads. Each bead has a specific weight and is associated with the history (x (x1-16 ..., x/- malized) velocities from 16 previous timesteps at intervals of A = 57. The model predicts the mean 14 E N and variance of E I of (normalized) acceleration for each bead, assuming normal distribution. We train us- ing the negative log- likelihood loss between predicted and ground truth accelerations computed from the ground truth trajectories.
of (nor-
Results The results are given in Fig. 7. As baselines, we use MPNN (Gilmer et al., 2017) as well as two hardware- efficient architectures: PointNet++ (Qi et al., 2017) and PTv3 (Wu al., 2024b). Notably, model choice has mini- mal impact on performance, potentially due to the absence of long-range interactions as the CG beads do not carry any charge. Furthermore, it is sufficient to only learn local bonded interactions. There is, however, a considerable im- provement in runtime for Erwin (1.7-2.5 times depending on the size), which is only matched by a smaller MPNN or PointNet++, both having significantly higher test loss.
5.2. Molecular dynamics
Molecular dynamics (MD) is essential for understanding physical and biological systems at the atomic level but re- mains computationally expensive even with neural network potentials due to all-atom force calculations and femtosec- ond timesteps required to maintain stability and accuracy. Fu et al. (2022) suggested accelerating MD simulation through coarse- grasseegganeddd dynamics with an MPNN. In this experiment, we take a different approach and instead oper-
7

---

Erwin Transformer
Table RMSE PDE benchmarks from Li et al. (2023h). Transformer-
Table 5. Test MSE for ShapeNet-Car pressure prediction. Baseline results are taken from Bleeker et al. (2025).
based models are taken as baselines from Luo et al. (2025).
RMSE
MODEL
ELASTICITY
PLASTICITY
AIRFOIL
PIPE
MODEL
MSE
LNO (2024)
0.69
0.29
0.53
0.31
POINTNET (2017)
43.36
GALERKIN (2021)
2,40
1.20
1.18
0.98
GINO (2023B)
35.24
HT-NET (2022)
/
3.33
0.65
0.59
UPT (2024A)
31.66
OFORMER (2023C)
1.83
0.17
1.83
1.68
TRANSOLVER (2024A)
19.88
GNOT (2023)
0.86
3.36
0.76
0.47
GP-UPT (2025)
17.02
FACTFORMER (2023D)
3.12
0.71
0.60
PTv3-S (2024B)
19.09 + 0.67
ONO (2024)
1.18
0.48
0.61
0.52
PTv3-M (2024B)
17.42 + 0.38
TRANSOLVER++ (2025)
0.52
0.11
0.48
0.27
ERWIN-S (OURS)
15.85 0.19
ERWIN (OURS)
0.34
0.10
2.57
0.61
ERWIN-M (OURS)
15.43 0.45
5.3. PDE benchmarks and airflow pressure
ShapeNet-Car See Table 5 for results. Both Erwin and PTv3 achieve significantly lower test MSE compared to other models We note that the best performing configura- tion of Erwin and PTv3 did not include any coarsening, thus operating directly on the original point cloud. This indicates that the task favors the ability of model to capture fine ge- ometric details. In comparison, other approaches introduce information loss through compression UPT (Alkin et al., 2024a) and Transolver (Wu et al., 2024a) involve pooling to the latent space, while GINO (Li et al., 2023b) interpolates the geometry onto regular grids and back.
Deep learning models have emerged as surrogate solvers of partial differential equations (PDEs), learning to approxi- mate solutions from data (Li et al., 2021; Wu et al., 2024a). While their advantages over traditional numerical methods remain unclear (McGreivy Hakim, 2024), the task it- self serves surrogate for large-scale applications like weather forecasting (Price al., 2025) and fluid dynamics (Bleeker al., 2025), where conventional solvers become computationally prohibitive. Furthermore, in this task, we are interested in the model's ability to scale to large domains while capturing complex patterns in underlying physics.
5.4. Turbulent fluid dynamics
Dataset We benchmark on multiple datasets taken from Li et al. (2023a). Each dataset is defined either on point cloud (Elasticity) structured mesh (Plasticity, Airfoil, Pipe). Additionally, we evaluate our model on airflow pressure modeling (Umetani Bickel, 2018; Alkin et al., 2024a). It consists of 889 car models, each car represented by 3,586 surface points in 3D space. Airflow was simulated around each car for 10s (Re 5 10") and averaged over the last 4 to obtain pressure values at each point.
In the last experiment, we demonstrate the expressivity of our model by simulating turbulent fluid dynamics. The prob- lem notoriously challenging due to multiple factors: the inherently nonlinear behavior of fluids, the multiscale and chaotic nature of turbulence, and the presence of long-range dependencies. Moreover, the geometry of the simulation do- main and the presence of objects introduce complex bound- ary conditions, thus adding another layer of complexity.
Dataset We use EAGLE (Janny et al., 2023), large-scale benchmark of unsteady fluid dynamics. Each simulation includes flow source (drone) that moves in 2D environ- ments with different boundary geometries, producing air- flow. The time evolution of velocity and pressure fields is recorded along with dynamically adapting meshes. The dataset contains 600 different geometries of 3 types, with ap- proximately 1.1 million 2D meshes averaging 3,388 nodes each. The total dataset includes 1,184 simulations with 990 time steps per simulation. The dataset split with 80% for training and 10% each for validation and testing.
Task For the PDE benchmarks, we follow the pipeline from Wu et al. (2024a) and minimize the relative L2 error; see Appendix B for details. For ShapeNet- the task is estimate the surface pressure Y R3388x1 given surface points X E p3388x3 We train by optimizing the mean squared error between predicted and ground truth pressure.
PDE benchmarks The results are given Table 4, where compare against other transformer-based methods. Er- win achieves oftte--f-thearr- performance on 2 out of 4 tasks. Interestingly, it dramatically underperforms on the Airfoil task, which indicates failure mode. We speculate that this is related to a specific structure of the data the density of the mesh decreases dramatically moving away from the center of mass. This means that points across different balls have considerably varying density, which poses a challenge.
Task We follow the original experimental setup of the benchmark. The input is the velocity V E N and pres-
'The baseline results for the ShapeNet-Car task, except for PTv3, were taken from Bleeker et al. (2025).
8

---

Erwin Transformer
Ground truth
Table 6. RMSF velocity V and pressure P fields across different prediction horizons (mean + std over 5 runs). Inference runtime and memory use are computed for a batch avg. 3,500 nodes.
Prediction
HORIZON
+1
+50
TIME (MS)
MEM.
FIELD/ UNIT
(GB)
MGN
0.081
0.43
0.592
2.25
Normalized RMSE
GAT
0.170
61.6
0.855
163
41
0.5
0.4
DRN
0.251
1.45
0.537
2.46
42
0.2
EAGLE
0.053
0.349
0.2
ERWIN
0.044
0.31
0.281
1.15
11
0.2
(OURS)
1:0.001
4:0.01
=0.001
4:006
0.0
35
95
Figure 8. The norm of the velocity field at different steps of the rollout trajectories.
improvements, likely due to its effectiveness in learning local interactions. At the same time, that did not generalize to ShapeNet-Car, where the most decisive factor was includ- ing cross- bass-ball interactions, highlighting the importance of capturing fine-grained details in this task.
sure P E R- x2 fields evaluated at every node of the mesh at time step f, along with the type of the node. The task is to predict the state of the system at the next time step + The training done by predicting trajectory of states of length 5 and optimizing the loss
6. Conclusion
We present Erwin, a hierarchical transformer that uses ball tree partitioning to process large-scale physical systems with linear complexity. Erwin achieves state-of-the-art perfor- mance in cosmology, turbulent fluid dynamics, and, partially, standard PDE benchmarks, demonstrating its effective- ness across diverse physical domains. The efficiency of Erwin makes it suitable candidate for tasks that require modeling large particle systems, such as computational chemistry (Fu et al., 2024) or diffusion-basee molecular dynamics (Jing et al., 2024).
C= -=E Σ
MSS((V++,, Ve+s) + MSE (EBE(P+i..)) Pe+t)).
where a = 0.1 is the parameter that balances the importance of the pressure field over the velocity field.
Results For comparison, we include the baselines from the original benchmark: MeshGraphNet (MGN; Pfaff al., 2021), GAT (Velickovic et al., 2018), DilResNet (DRN: Stachenfeld et al., 2021) and EAGLE (Janny et al., 2023)? The first two baselines are based on message-passing, while DilResNet operates on regular grids, hence employing in- terpolation for non-uniform meshes. EAGLE uses message- passing to pool the mesh to a coarser representation with a fixed number of clusters, on which attention then com- puted. The quantitative results are given in Table and unrolling trajectories shown in Fig. 8, 9. Erwin demon- strates strong results on the benchmark and outperforms every baseline, performing especially well at predicting pressure. In terms of inference time and memory consump- tion, Erwin achieves substantial gains over EAGLE, being 3 times faster and using times less memory.
Limitations and Future Work Because Erwin relies on perfect binary trees, we need to pad the input set with virtual nodes, which induces computational overhead for ball atten- tion computed over non-coarsened trees (first ErwinBlock). This issue can be circumvented by employing learnable pooling to the next level of the ball tree, which is always full, ensuring the remaining tree is perfect. Whether we can perform such pooling without sacrificing expressivity is a question that we leave future research.
Erwin relies cross-ball interaction and coarsening to capture long-range interactions. Both mechanisms have in- herent limitations: the former requires multiple steps for signals to propagate between balls, while the latter sacri- fices fine- ime-grained detail. promising approach to address these issues adapting sparse attention methods such as Native Sparse Attention (Yuan et al., 2025). This frame- work aligns naturally with Erwin's ball tree structure and would enable learning distant interactions while preserving full resolution. Finally, Erwin is neither permutation nor rotation equivariant, although rotation equivariance can be incorporated without compromising scalability, such as via Geometric Algebra Transformers (Brehmer et al., 2023) or Fast Fuclidean Attention (Frank et al., 2024).
5.5. Ablation study
We also conducted an ablation study to examine the effect of increasing ball sizes on the model's performance in the cosmology experiment; see Table 2. Given the presence of long-range interactions in the data, larger window sizes (and thus receptive fields) improve model performance, albeit at the cost of increased computational runtime. Our architec- tural ablation study the MD task (Table 3) reveals that using MPNN at the embedding step produces substantial
We additionally trained UPT (Alkin et al., 2024a), but were not able to obtain competitive results in our initial experiments.

---

Erwin Transformer
Acknowledgements
Bodnar, C.,, Bruinsma, W. P., Lucie, A., Stanley, M., Brand-
stetter, J., Garvan, P., Riechert, M., Weyn, J., Dong, H.,
We are grateful to Evgenii Egorov and Ana Lučié for their feedback and inspiration. MZ acknowledges support from Microsoft Research AI4Science. JWvdM acknowledges support from the European Union Horizon Framework Pro- gramme (Grant agreement ID: 101 120237).
Vaughan, A., Gupta, J. K., Tambiratnam, K., Archibald, A., Heider, E... Welling, M.. Tumer, R. E.,, and Perdikaris, P. Aurora: A foundation model of the atmosphere. arXiv preprint arXiv:2405.1306. 2024.
Brandstetter, J., Hesselink, R., van der Pol. E., Bekkers, E. J., and Welling, M. Geometric and physical quantities improve E(3) equivariant message passing. In Interna- tional Conference on Learning Representations (ICLR), 2022.
Impact statement
The broader implications of our work are primarily in moderate- large-scale scientific applications, such as molecular dynamics or computational fluid dynamies. We believe that efficient and expressive architectures like Er- win could become foundation for resource- resource-intensive deep learning frameworks and therefore help in better understand- ing the physical systems governing our world.
Brehmer, J., de Haan, P., Behrends, S., and Cohen, T. S. Geometric algebra transformer. In Conference on Neural Information Processing Systems (NeurIPS), 2023.
Cao, S. Choose transformer: Fourier galerkin. In Conference on Neural Information Processing Systems (NeurIPS), 2021.
References
Alkin, B., Fürst, A., Schmid, S., Gruber, L.,, Holzleitner, M., and Brandstetter, J. Universal physics transformers: framework for efficiently scaling neural operators. In Conference on Neural Information Processing Systems (NeurIPS), 2024a.
Cao, Y., Chai, M., Li, M., and Jiang, C. Efficient learning of mesh-based physical simulation with bi-stride multi-scale graph neural network. In International Conference on Machine Learning (ICML), 2023.
Carrier, J., Greengard, L., and Rokhlin, V. A fast adaptive multipole algorithm for partiele simulations. SIAM Jour- nal Scientific and Statistical Computing, 9())669-686, 1988. doi: 10. 1137/0909044.
Alkin, B., Kronlachner, T., Papa, S., Pirker, S., Lichteneg- ger, T., and Brandstetter, J. Neuraldem real-time sim- ulation of industrial particulate flows. arXiv preprint arXiv:2411.09678 2024b.
Cohen, T. and Welling, M. Group equivariant convolu- tional networks. In International Conference on Machine Learning (ICML), 2016.
Arts, M., Satorras, V., Huang, C.-W., Zuegner, D., Federici, M., Clementi, C., Noé, F., Pinsler, R., and Berg, R. Two for one: Diffusion models and force fields for coarse- grained molecular dynamics. Journal of chemical theory and computation, 19, 09 2023. doi: 10.1021/acs.jctc. 3c00702.
Dao, T. Flashattention-2: Faster attention with better paral- lelism and work partitioning. In International Conference Learning Representations (ICLR), 2024.
Balla, J., Mishra-Sharma. S., Lázsro- C., Jaakkola, T. S., and Smidt, T. E. A cosmic-scale benchmark for symmetrry-reseerinng data processing. arXiv preprint arXiv:2410.20516, 2024.
Frank, J. T., Chmiela, S., Müller, K., and Unke, O. T. Eu- clidean fast attention: Machine lcarning global atomic representations at linear cost. CoRR, abs/2412.08541, 2024.
Barnes, J. and Hut, P. A hierarchical O(N log N) force- calculation algorithm. Nature, 32466096):446-44,, 1986. doi: 10. 10.00//22444.0..
Fu, X., Xie, T., Rebello, N. J., Olsen, B. D., and Jaakkola, T. Simulate time-interrated coarse roarse-grained molecular dynamics with multi-scale graph networks. Trans. Mach. Learn. Res., 2023, 2022.
Batzner, S. L., Musaclian, A., Sun, L.. Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13, 2021.
Fu, X., Xie, T., Rosen, A. S., Jaakkola, T. S., and Smith, J. Mofdiff: Coarse Coarse-grained diffusion for -egal-organic framework design. In International Conference on Learn- ing Representations (ICLR), 2024.
Bleeker, M., Dorfer, M., Kronlachner, T., Sonnleitner, R., Alkin, B., and Brandstetter, J. Neuralefd: Deep learning on high-fidelity automotive aerodynamics simulations. CORR, abs/2502.09692. 2025.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural message passing for quantum chem- istry. In International Conference on Machine Learning (ICML), 2017.
10

---

Erwin Transformer
Hao, Z., Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng, Z.,, Song, J., and Zhu, GNOT: A general neural operator transformer for operator learning. In Interna- tional Conference on Machine Learning (ICML), 2023.
Li, Z.., Shu, D., and Farimani, A. B. Scalable transformer for PDE surrogate modeling. In Conference on Neural Information Processing Systems (NeurIPS), 2023d.
Liu, X., Xu, B., and Zhang, L. Ht-net: Hierarchical trans- former based operator learning model for multiscale pdes. 2022.
Hockney, R. and Eastwood, J. Computer Simulation Using Particles. CRC Press, 2021. ISBN 9781439822050, URL https:/ /bookss..... google.oooooooooooooooooooo..o htt.l/ee./bbooooiieeo IIOFKmnCQuIC.
CoRR, abs/2210.1088
Liu, Z., Lin, Y., Cao. Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. In International Con- ference on Computer Vision (ICCV), 2021.
Janny, S., Béneteau, A., Nadri, M., Digne, J., Thome, N., and Wolf, C. EAGLE: large-scale learning of turbulent fluid dynamics with mesh transformers. In International Conference on Learning Representations (ICLR), 2023.
Liu, Z., Yang, X., Tang, H., Yang, S., and Han, S. Flat- former: Flattened window attention for efficient point cloud transformer. Conference on Computer Vision and Pattern Recognition( 2023.
Jing, B., Stärk, H., Jaakkola, T. S., and Berger, B. Genera- tive modeling of molecular dynamics trajectories. arXiv preprint 1XX.i::2499.1880 2024.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations (ICLR)) 2019.
Kang, Y., Tran, G., and Sterck, H. D. Fast multipole atten- tion: A divide-and-conquer attention mechanism for long sequences, arXiv preprint arXiv:2310.11960, 2023.
Luo, H., Wu, H., Zhou, H., Xing. L., Di, Y., Wang, J., and Long, M. Transolver++: An accurate neural solver for pdes on million-scal scale geometries. CoRR, abs/2502.02414, 2025.
Kovács, D. P., Moore, J. H., Browning, N. J., Batatia, I., Horton, J. T., Kapil, V., Witt, W. C., Magdäu, I.-B., Cole, D. J., and Csányi, G. Mace-off23: Transferable machine learning force fields for organic molecules, 2023.
Majumdar, S., Sun, J., Golding, B., Joe, P., Dudhia, J., Caumont, O., Gouda, K. C., Steinle, P., Vincendon, B., Wang. J., and Yussouf, N. Multiscale forecasting of high- impact weather: Current status and future challenges. Bulletin of the American Meteorological Society, 102: 1-65, 10 2020. doi: 10.1 10.11755BAMSS--2000000
Lam, R., Sonzaez- A., Willson, M., Wirnsberger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Eaton- Rosen, Z., Hu, W., Merose, A., Hoyer. S., Holland, G., Vinyals, O., Stott, J., Pritzel, A., Mohamed, S., and Battaglia, P. Learning skillful medium-range global weather forecasting. Science, 382(6677; 2023.
McGreivy, N. and Hakim, A. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid- related partial differential equations. Nature Machine Intelligence, 6(10), 2024.
Li, Z., Kovachki, N. B., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A. M., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. In International Conference on Learning Represen- tations (ICLR), 2021.
Moskalev, A., Prakash, M., Xu, J., Cui, T., Liao, R., and Mansi, T. Geometric hyena networks for large-scale equivariant learning. In International Conference on Machine Learning (ICML), 2025.
Li, Z., Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. JMLR, 24:388:1 1-38:22,,...6,,,..111111111 2023a.
Pedregosa, F,, Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., VanderPlas, J., Passos, A., Cour- napeau, D.,, Brucher, M., Perrot, M., and Duchesnay, Scikit-learn: Machine learning in python. arXiv preprint 1:Xiv:200.04900 2012.
Li, Z., Kovachki, N. B., Choy, C. B., B., Kossaifi, J., Otta, S. P., Nabian, M. A., Stadler, M., Hundt, C., Azizzade- nesheli, K., and Anandkumar, A. Geometry-informedd neural operator for large-scale 3d pdes. In Conference on Newral Information Processing Systems (NeurIPS), 2023b.
Pfaff, T. Fortunato, M., Sanchez-Gonzalez, A., and Battaglia, P. W. Learning mesh-based simulation with graph networks. In International Conference Learning Representations (ICLR), 2021.
Li, Z., Meidani, K., and Farimani, A. B. Transformer for partial differential equations' operator learning. Trans. Mach. Learn. Res., 2023c.
Pfalzner, S. and Gibbon, P. Many-Bos Body Tree Methods in Physics. Cambridge University Press, 1996.

---

Erwin Transformer
Price, I., Sanchez-Gonzale,, A., Alet, F., Andersson, T. R., El-Kadi,, A., Masters, D., Ewalds, T., Stott, J., Mohamed, S., Battaglia, P. W., Lam, R. R., and Willson, M. Proba- bilistic weather forecasting with machine learning. Na- ture, (337(8044:844-99) 2025.
Wang, T. and Wang, C. Latent neural operator for solving forward and inverse PDE problems. In Conference on Neural Information Processing Systems (NeurIPS), 2024.
Watson, L.,, Juergens, D., Bennett, N. R., Trippe, B. L.,
Yim, J., Eisenach, H. E., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., Wicky, B. I. M., Hanikel, N., Pellock, S. J., Courbet, A., Sheffler, W., Wang, J., Venkatesh, P., Sappington, I., Torres, S. V., Lauko, A., Bortoli, V. D., Mathieu, E., Ovchinnikov, S., Barzilay, R., Jaakkola, T. DiMaio, F., Back, M., and Baker, D. De novo design of protein structure and function with rfdiffusion. Nature, 620: 1100, 2023.
Qi. C. Yi, L.,, Su, H., and Guibas, L. Pointnet++: Deep hierarchical feature learning on point sets a metric space. In Conference Neural Information Processing Systems (NeurIPS), 2017.
Rogozhnikov, A. Einops: Clear and reliable tensor ma- nipulations with einstein-like notation. In International Conference on Learning Representations (ICLR), 2022.
Webb, M., Jackson, N., Gil, P., and de Pablo, J. Tar- geted sequence design within the coarse-grained polymer genome. Science Advances, 6:eabc6216, 10 2020. doi: 10.1126/sciadv.abb6211
Ronneberger; O., Fischer, P., and Brox, T. U-net: Convolu- tional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Compuler-Assisted Intervention (MICCAI), 2015.
Wessels, D. R., Knigge, D. M., Papa, S., Valperga, R., Vadgama, S. P., Gavves, E., and Bekkers, E. Ground- ing continuous representations in geometry: Equivariant neural fields. arXiv preprint arXhv:2406.05753 2024.
Shazeer. N. GLU variants improve transformer. arXiv preprint ar arXiv:2002.05202 2020.
Stachenfeld, K. L., Fielding. D. B., Kochkov, D., Cranmer, M. D., Pfaff, T., Godwin, J., Cui, C., Ho, S., Battaglia, P. W., and Sanchez-Gonzales A. Learned coarse mod- els for efficient turbulence simulation. arXiv preprint arXiv:21 2021.
Wu, H., Luo, H., Wang, H., Wang, J., and Long, M. Tran- solver: A fast transformer solver for pdes on general geometries. In International Conference on Machine Learning (ICML), 2024a.
Sun, P., Tan, M., Wang, W., Liu, C., Xia, F., Leng, Z., and Anguelov, D. Swformer: Sparse window transformer for 3d object detection in point clouds. In European Conference on Computer Vision (ECCV), 2022.
Wu, X., Jiang, L., Wang, P., Liu, Z., Liu, X., Qiao, Y., Ouyang, W., He, T., and Zhao, H. Point transformer V3: simpler, faster, stronger. In Conference on Computer Vision and Puttern Recognition( (VVP)) 2024b.
Xiao, Z., Hao, Z., Lin, B., Deng, Z., and Su, H. Improved op- erator learning by orthogonal attention. In International Conference on Machine Learning (ICML), 2024.
Umetani, N. and Bickel, B. Learning firee-dimensiona, flow for interactive aerodynamic design. ACM Trans. Graph., 37(4):89, 2018. URL /
3201325.
Yuan, J., Gao, H., Dai, D., Luo, J., Zhao, L., Zhang, Z., Xie, Z., Wei, Y. X., Wang, L., Xiao, Z., Wang, Y., Ruan, C., Zhang, M., Liang, W., and Zeng, W. Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025.
Valencia, M. L., Pfaff, T., and Thuerey, N. Learning distri- butions of complex fluid simulations with diffusion graph networks. In International Conference Learning Rep- resentations (ICLR), 2025.
Zhang, T., Yuan, H., Qi, L.., Zhang, J., Zhou, Q., Ji, S., Yan, S., and Li, X. Point cloud mamba: Point cloud learning via state space model. In Association for the Advancement of Artificial Intelligence, AAAI, 2025.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. In Conference on Neural Information Processing Systems (NeurIPS), pp. 5998-6008, 2017.
Zhu, Z. and Soricut, R. H-transformer-ld: Fast one- dimensional hierarchical attention for sequences. In Conference on Neural Information Processing Systems (NeurIPS), pp. 3801-3815, Association for Computa- tional Linguistics, 2021.
Velickovic, P., Cucurull, G., Casanova, A., Romero, A., Liô, P., and Bengio, Y. Graph attention networks. In International Conference Learning Representations (ICLR), 2018
Wang, P.-.. Octformer: Octree-based transformers for 3D point clouds. ACM Transactions on Graphics (SIG- GRAPH), 42(4), 2023.
12

---

Erwin Transformer
Algorithm 1 BULLEE
ball tree. In our experiments, consistently maintain the size of MLPP, and ML.,, to be small (H I 32) such that embedding accounts for less than 5% of total runtime.
input Array data points D in IR output Ball tree node B
if then
B. Experimental details
Create leaf node B containing single point in D return B
In this section, we provide comprehensive experimental details regarding hardware specifications, hyperparameter choices, and optimization procedures.
end if'
Find dimension of greatest spread
< argmax (maxxee (maxzeD 24 mingeD x4)
B.1. Hardware and Software
All experiments were conducted on a single NVIDIA RTX A6000 GPU with 48GB memory and 16 AMD EPYCTM 7543 CPUs. Erwin and all baselines except those for cos- mology were implemented in PyTorch 2.6. For the cosmol- ogy benchmark, SEGNN, NequlP, and MPNN baselines were implemented in JAX as provided by the original bench- mark repository. Training times for Erwin varied by task: cosmology (5-10 minutes depending on training set size), molecular dynamics (2-4 hours depending on model size), PDE benchmarks (8 hours for Elasticity, 48 hours for oth- ers), ShapeNet-Car (2 hours), and EAGLE (48 hours).
Find the median point along 8 < (adian(as x ( D}
# Points left of median along
L + + fØ D I $ ps}
# Points right of median along
R + (x E D xz > p8}
Recursively construct children
B.child,+ +- BUL.DBALITREE(L)
B.childz < BUILDBALLITEE(()]
return
B.2. Training Details
All models were trained using the AdamW optimizer (Loshchilov & Hutter, 2019) with weight decay 10-5, The learning rate was tuned in the range 10 to 10 to mini- mize loss on the respective validation sets with cosine decay to 10-7 Gradient clipping by norm with value 1.0 was applied across experiments. Early stopping was used only for ShapeNet-Car and molecular dynamics tasks, while all other models were trained until convergence. In every experiment, normalize inputs to the model. Hyperpa- rameter optimization was performed using grid search with single trials.
A. Implementation details
The algorithm used for construct- ing ball trees (Pedregosa al., 2012) can be found in Alg. 1. Note that this implementation is not rotationally equivariant as it relies on choosing the dimension of the greatest spread, which in turn depends on the original orientation. Examples of ball trees built in our experiments are shown in Fig. 9.
Ball construction
MPNN in the embedding Erwin employs a small-scal. MPNN in the embedding. More precisely, given a graph
G (V, E) with nodes v; V and edges eij E, compute multiple layers of message-passing as proposed in (Gimeer et al., 2017):
B.3. Dataset Splits and Evaluation
Dataset splits followed the original benchmarks:
Training set varied from 64 to 8192 exam- ples, with validation and test sets of 512 examples each
Cosmology:
miy -MLP, (h;, hy. Pj),
message
m; - myj
aggregate (14)
Molecular Dynamics: 100 short trajectories for training,
JEN())
40 long trajectories for testing
h. = MLP,, (L,
update
1000 training / 200 test examples (except Plasticity: 900/80)
PDE Benchmarks:
where h, E " the neighborhood of ug The motivation for using an MPNN is to incorporate local neighborhood information into the model. Theoretically, attention should be able to capture this information as well; however, this might require sub- stantially increasing the feature dimension and the number of attention heads, which would be prohibitively expen- sive for a large number of nodes the original level of
H is feature vector of 0/ and (i) denotes
ShapeNet-Car: 700 training / 189 test examples
EAGLE: 1184 trajectories with 80%/10%//10 split
For statistical significance, we ran each experiment 5 times and report mean and standard deviation for cosmology, molecular dynamics, ShapeNet-Car, and EAGLE. PDE ex- periments were run once due to computational constraints.
13

---

Erwin Transformer
ball size 512
ball size 256
ball size 128
Figure Examples of ball trees built on top of data. Partitions at different levels trees are shown. Top: A polypeptide from molecular dynamics task. Center: A domain from the EAGLE dataset. Bottom: A car surface from the ShapeNet-Car dataset.
B.4. Evaluation Metrics
B.6. Computational Efficiency
For RMSE computation, we the relative L2 error: RMSE : [x) wll/Ily/l. For molecular dynamies, randomly sample 16 history points from each trajectory and predict one future point. Acceleration is predicted by the neural network based on history and compared against ground truth computed using forward differences. Model performance evaluated using negative ligllikll loss between predicted and ground truth accelerations.
Table shows the inference time breakdown for Erwin on NVIDIA RTX A6000 with batch size 16. Ball tree con- struction (Table 8) consistently accounts for less than 5% of total runtime, demonstrating the efficiency of our optimized implementation compared to standard libraries.
Table 7. Erwin runtime with torch. compile
NODES PER BATCH
RUNTIME (MS)
16 x
2048
4096
8192
16384
B.5. Baseline Implementations
FWD
17.3
31.6
79.7
189
All baseline results were taken from official implementa- tions or reported values:
FWD BWD
26.4
45.4
232
Cosmology: Official JAX implementation (Balla et al., 2024), code from https://gitthnn com/ smsharma/eegnn-jan
Table 8, Ball tree construction on 16 AMD EPYCTM 7543 CPUs.
NODES PER BATCH,
RUNTIME (MS)
x
2048
4096
8192
16384
ShapeNet-Car: Results from Bleeker et al. (2025), code from https://githutt www/mmljkummmmmnamm
16.3
21.2
24.1
44.0
OURS SPEED-UP
0.73
1.54
3.26
PDE Benchmarks: Results from Luo al. (2025),
22.3x
13.8x
7.4x
6.3x
https://github.com//
EAGLE: Original implementation (Janny al., 2023), https:h/tttthhhiee com/caoc////eeeooo hatmmeagle-daaassee
B.7. Further details per experiment
EaglelleshTransfomeer
Cosmological simulations We follow the experimental setup of the benchmark. The training was done for 5,000 epochs with batch size 16 for point transformers and batch size for message- message-passing- tion of SEGNN, NequlP, and MPNN was done in JAX and
Molecular Dynamies: MPNN and PointNet++ imple- mented by us, other models taken from official codebases.
PTv3: https: httpt//gttttttt//e www.babccco///nnccee..e PointTransformerV
models. The implementa-
14

---

Erwin Transformer
taken from the original benchmark repository (Balla al., 2024). We maintained the hyperparameters of the baselines used in the benchmark. For Erwin and PTv3, the hyperpa- rameters are provided in Table 9. In Erwin's 's embedding, we conditioned messages on Bessel hasis functions rather than the relative position, which significantly improved overall performance.
Molecular dynamics All models were trained with batch size 32 for 50,000 training iterations with an initial learning rate of 5- 10 We fine-tuned the hyperparameters of every model on the validation dataset (reported in Table 11).
PDE benchmarks The baseline results and experimental setups are taken from Luo et al. (2025). We adjusted batch size, ball size, and the number of attention heads per block for the best performance. The hidden dimensionality Er- win was adjusted such that the overall number of parameters is around 106, which is comparable with other baselines. Constrained by the parameter size, the same configuration worked the best; see Table 12 for details.
Airflow pressure modeling We take the results of base- line models from Bleeker et al. (2025). Both Erwin and PTv3 were trained with batch size 32 for 1,000 epochs, and their hyperparameters are given in Table 10. We tuned the number of blocks, the number of message-passingg steps, hidden dimensionality, and ball size per block for the best performance. When experimenting with pooling, we found that not involving any coarsening significantly improves model performance; hence, we used stride in each block.
Turbulent fluid dynamics Baseline results are taken from (Janny al., 2023), except for runtime and peak memory us- age, which we measured ourselves. Erwin batch size 12 for 4,000 epochs. We tuned hidden dimen- sionality and ball size per block for the lowest validation loss.
trained with
15

---

Erwin Transformer
Table 9. Model architectures for the cosmological simulations task. varying sizes Erwin, values are given (S/M).
Table 11. Model architectures for the molecular dynamics task. For models of varying sizes, the values are given as (S/M/L).
Model
Parameter
Value
Model
Parameter
Value
PTv3
Grid size
0.01
MPNN
Hidden dim.
48/64/128
Enc. depths
(2, 2, 2)
MP steps
6
Enc. channels
(32, 64, 128, 256)
MLP layers
Enc. heads
(2, 4, 8, 16)
Enc. patch size
Message agg-n
mean
Dec. depths
(2, 2,
PointNet++
Hidden dim.
64/128/196
Dec. channels
(64, 64, 128)
MLP layers
Dec. heads
(2, 8)
PTv3
Grid size
0.025
Dec. patch size
64
Enc. depths
(2, 2, 6, 2)
Pooling
2,
Enc. channels
16-192/24-384/64-1022
Erwin
MPNN dim.
Enc. heads
(2, 8, 16, 32)
Channels
32-512/64-1024
Enc. patch size
128
Window size
Dec. depths
(2, 2,
Enc. heads
(2, 4, 8, 16)
Dec. channels
11-66/48-112/44-12
Enc. depths
(2, 6, 2)
Dec. heads
(4,4, 8, 16)
Dec. heads
42, 8)
Dec. patch size
128
Dec. depths
(2, 2,)
Erwin
MPNN dim.
16/16/32
Pooling
(2, 1)
Channels
(16-256/22-512/4441022)
Window size
128
Enc. heads
(2, 4, 8, 16, 32)
Enc. depths
(2, 2, 6,
heads
(4, 4, 16)
Dec. depths
(2, 2,
Pooling
(2, 22, 1)
Table Model architectures for airflow pressure task.
Model
Parameter
Value
PTv3
Grid size
0.01
Enc. depths
(2. 2,
Enc. channels
24-384
Enc. heads
(2, 4, 16, 32)
Enc. patch size
256
Dec. depths
(2, 2, 2, 2)
Table 12. Model architectures for the PDE benchmarks.
Dec. channels
48-192
Dec. heads
4, 16)
Model
Parameter
Value
Dec. patch
Erwin
MPNN dim.
64
Erwin
MPNN dim.
8
Channels
64
Channels
96
Window size
256
Window size
256
Enc. heads
(8, 16)
Enc. heads
(8, 8)
Enc. depths
(6, 2)
Enc. depths
(6, 6)
Dec. heads
(8.)
Dec. heads
(8, 8)
Dec. depths
(2.)
Dec. depths
(6)
Pooling
1)
Pooling
(1,
steps
16

---

Erwin Transformer
Figure 10. The norm of the velocity field at different steps of the rollout trajectories, predicted by Erwin.
Ground truth
10
Prediction
5
Normalized RMSE
0.4
0.2
0.0
t 5
== 30
= 55
f = 80
== 105
Ground truth
10
Prediction
3
Normalized RMSE
0.4
0.2
0.0
5
30
55
f 80
105
17

---

Erwin Transformer
Figure The norm the pressure field at different steps of rollout trajectories, predicted by Erwin.
Ground truth
60
40
Prediction
20
Normalized RMSE
0.4
0.2
0.0
5
t= 30
55
80
t= 105
Ground truth
60
40
Prediction
20
0
Normalized RMSE
0.4
0.2
0.0
t
== 30
f == 105
18


Reviews:
[{"review_id": "ou8Aaj0GBi", "reviewer": "Reviewer_dSeu", "summary": "The paper introduces a transformer-based approach for processing point cloud data. It utilizes ball-tree-based structures, enabling the attention-based framework to operate in linear time rather than quadratic time. This also allows the model to capture interactions at both fine and coarse scales.\n\nThe proposed method is evaluated on three problems: capturing long-range interactions (cosmology); computational efficiency (molecular dynamics); model expressivity on large-scale multi-scale phenomena (turbulent fluid dynamics), demonstrating strong performance in both computational efficiency and prediction accuracy.\n\n\nupdate after rebuttal\nThank the authors for the responses. I decide to keep my original score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging feedback. We address their questions below.\n\n## Questions For Authors\n> The paper mentions \"The code is available at anonymized link\". But there seems to be no link attached at the \"anonymized link\".\n\nWe apologize for the inconvenience. An anonymized version of the repo can be found here: https://github.com/erwin-transformer/erwin.\n\n> During the simulation process, the positions of particles will change. Does the proposed method re-build the tree after each step?\n\nRebuilding the tree depends entirely on the properties of the modelled system, particularly, how quickly the arrangement of particles changes:\n- The mesh is *static*: the tree should only be computed once during preprocessing.\n- The mesh changes *slowly*: the tree should be rebuilt each N step.\n- The mesh is dynamic and changes *quickly*: ideally, the tree should be rebuilt at every step of a simulation.\n\nDuring the code implementation, we assumed that we would need to build the tree as fast as possible. Hence, we implemented ball tree construction in C++ & OpenMP and according to our benchmarks, it consistently takes less than 10% of the forward pass. For details, see https://github.com/erwin-transformer/erwin?tab=readme-ov-file#benchmark.\n\n## Additional experiments\n\nWe also conducted additional experiments to further highlight the strength of our method. \n### OctFormer on the molecular dynamics task\n\nWe will include OctFormer results for the molecular dynamics task in the final manuscript. We have evaluated OctFormer while maintaining similar parameter count, depth, and partitioning size (128) as Erwin for each model size.\n|  Model || model size ||\n|--------------------------------|------------|--------------|-------------|\n|                               | small (4M) | medium (19M) | large (43M) |\n| OctFormer                      | 0.725      | 0.715        | 0.712       |\n| Erwin                          | **0.712**      | **0.693**        | **0.691**   |\n\n### Ablation study on ShapeNet-Car dataset\nWe conducted an ablation study to gather additional information concerning the role of the MPNN embedding and Ball Attention in Erwin's performance. Here, Erwin has a small size (1.5M parameters) and a fixed width (we do not use any upsampling/downsampling to achieve the best performance)\n\n| | Test MSE |\n|-------------|--------|\n| w/o         | 30.39  |\n| + MPNN      | 30.49  |\n| + RPE       | 30.02  |\n| + Rotating trees | **16.58** |\n| Transolver (3.9M) | 19.88 |\n\nNotice that in this experiment, including the MPNN leads to deteriorating performance, while rotating trees essentially allow Erwin to achieve state-of-the-art performance. The reason is that the ball attention is able to process large-scale data (3.6K points) at full fidelity, and rotating trees allow for information to propagate from one ball to another, essentially making the receptive field encompass the full car.\n\n## Conclusion\n\nWe are grateful for the reviewer's feedback. Should you have any questions or concerns, we are always ready to discuss them further.", "claims_and_evidence": "The claims made in the submission is supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "The experimental designs make sense.", "relation_to_broader_scientific_literature": "The paper discussed the existing works of sub-quadratic attention. Compared to those works, the paper uses ball-tree structures to capture the hierarchical structures of point cloud data.", "essential_references_not_discussed": "no.", "other_strengths_and_weaknesses": "The proposed ball-tree-based attention seems novel to me.\nBesides the algorithm, the paper also discusses the implementation details of the ball-tree attention, such as, the representation of the tree which makes the tensor rearranging on the nodes more efficient.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "- The paper mentions \"The code is available at anonymized link\". But there seems to be no link attached at the \"anonymized link\".\n- During the simulation process, the positions of particles will change. Does the proposed method re-build the tree after each step?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QkpSvM2prh", "reviewer": "Reviewer_eeiZ", "summary": "The paper introduces Erwin, a hierarchical transformer model designed for large-scale physical systems. The model employs a ball tree partitioning strategy to group nodes, enabling linear-time attention by processing local neighborhoods in parallel. This hierarchical approach allows progressive coarsening and refinement, capturing both fine-grained local details and global interactions. The paper demonstrates Erwin’s effectiveness across multiple domains, including cosmology, molecular dynamics, and particle fluid dynamics.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the actionable feedback. We address their concerns and questions in the following.\n\n## Questions For Authors\n### Q1\n\nThat is an insightful question. In general, ball trees are geometry-adaptive but do not guarantee that balls will cover points in a perfectly natural way and certain irregularities can indeed appear in practice. To address it, we introduce the distance-based attention bias (Eq. 10), which penalizes the strength of interaction between distant tokens. It provides a soft mechanism for creating more “natural” partitions. In our opinion, such partitioning irregularities are a byproduct of *any* attempt to impose a rigid structure onto irregular data, such as converting a point cloud into a sequence (e.g. space-filling curves in PTv3 produce discontinuities).\n\nRegarding tree depth, we clarify that it is fixed, as the tree is always constructed until leaf size 2 and is not a variable design parameter. For the splitting strategy, we employ Algorithm 1, selected primarily for its speed (also used by scikit-learn).\n\n### Q2\n\nFor both these datasets ball tree partitioning aligns very well with KD-tree partitioning and KD-tree would be a viable alternative. However, one advantage of ball trees is that they offer a rotation invariant partitioning which might be critical for certain physical applications.\n\n### Q3\n\nWe are thankful to the reviewer for bringing our attention to it and apologize for overlooking it in the original submission. We will include details in the final version. Tree construction creates leaf nodes with either 1 point (incomplete) or 2 points (complete). For single-point leaves, we duplicate the point as padding, creating virtual nodes that are copies of original points. This simplifies implementation, allowing pooling to work without edge cases and minimal computational overhead. Speaking of learnable pooling, this is one of the directions for future work as it would make Erwin applicable to gigantic-scale industrial applications and we aim to explore it.\n\n### Q4\n\nBall size is chosen according to computational constraints, i.e. available GPU memory. In our experiments, it was a hyperparameter which we tuned for the best performance, choosing between 128 and 256. Generally, it is a trade-off as 256 provides better coverage but 128 allows for using larger batch size. \n\n### Q5\n\nWe specify the number of  message-passing steps in Tables 6 and 7 for ShapeNet ($1$ step) and MD ($2$ steps). We will also add the detail for EAGLE and cosmology, where the number of steps was $3$.\n\n### Q6\n\nFor 2D cases, we specify degrees of rotation, and for 3D, we specify Euler angles. In all experiments we used fixed value of 45 which we empirically found to generate sufficiently different partitions. The difference between partitionings influences the performance significantly, as it is the mechanism that allows points within different original balls to interact. Furthermore, we avoided using small angles, but there was no difference for sufficiently large angles.\n\n### Q7\n\nThe reviewer refers to the cosmology experiment where the performance of MP models plateaus with increasing number training points. This is consistent with results of the original benchmark paper (Balla et al., Fig. 3). In the task, each data point is a point cloud with 5000 points. Furthermore, even for small training sizes, there is enough signal to capture local interactions for message-passing models. We note, however, that there are also long-range dependencies in data (Balla et al., Fig. 1), which MP models are not able to capture. Furthermore, regardless of how much data is given to them, they are not able to get that part of the signal due to short-range receptive field, which is different for Erwin or PointTransformer v3 that both have large receptive fields.\n\n### Q8\n\nWe did control explicitly in *all* the tasks that both Erwin and state-of-the-art PT v3 have the same depth and similar number of parameters as those models are comparable in their architecture.\n- cosmology: we use the same hyperparameters for each model as in the original paper (Balla et al.), as those are the best performing models according to the hyperparameter tuning from the authors (see Appendix C, Balla et al.).\n- EAGLE: we take results from the original paper where authors do hyperparameter sweep and thus report best-performing configurations. Parameter-wise, Erwin has 18M params and the closest competitor, EAGLE, has 10M params.\n- ShapeNet, all models have similar parameter count - between 4M and 5M parameters.\n\n## Conclusion\nUnfortunately, we could only use 5k characters for the rebuttal and do not have space to address weaknesses one-by-one despite having prepared answers. For additional experiments that address W1 and W4, we kindly refer to another rebuttal: https://openreview.net/forum?id=MrphqqwnKv&noteId=BfX9JC9Pxg. We would be happy to engage in further discussion if the reviewer would like us to clarify further.", "claims_and_evidence": "Most claims are supported.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes.\n\nWhile the paper provides training information for Erwin, it lacks precise training details of compared methods.", "relation_to_broader_scientific_literature": "Erwin builds on transformer-based approaches, but introduces hierarchical ball tree partitioning to improve efficiency.", "essential_references_not_discussed": "The paper includes most of the relevant prior work and provides thorough comparisons. However, since one of Erwin’s key contributions is hierarchical modeling, it might be beneficial to discuss and compare with existing hierarchical models in physical systems, such as BSMS-GNN [1] and HCMT [2]. These models incorporate multi-scale representations in physics simulations and could provide additional context for understanding Erwin’s approach.\n\n[1] Efficient Learning of Mesh-Based Physical Simulation with BSMS-GNN.\n\n[2] Learning Flexible Body Collision Dynamics with Hierarchical Contact Mesh Transformer.", "other_strengths_and_weaknesses": "**Strengths**\n- The paper presents a creative combination of ball tree partitioning and transformer-based attention, leading to an efficient and expressive hierarchical model for large-scale physical systems. \n- By leveraging ball tree partitioning, the model reduces computation cost, which is important for large-scale physics systems.\n\n**Weakness**\n- The paper discusses OctFormer, another tree-based hierarchical attention model, but does not include a direct comparison. Since both methods aim to improve efficiency through hierarchical partitioning, an empirical evaluation would better highlight Erwin’s advantages.\n- While the paper compares Erwin to UPT and EAGLE, which use two-level hierarchical pooling, it does not consider more flexible multi-level approaches like BSMS-GNN[1] and HCMT[2], which incorporate multi-scale message passing.\n- Erwin performs worse than some baselines when trained on small datasets. This suggests that the model may rely more heavily on large-scale data to learn effective representations, whereas certain baselines generalize better with limited training samples.\n- Table 2 shows that most improvements come from MPNN in the encoder, while Erwin Blocks (ball attention) contribute less. This raises the question of whether hierarchical attention is essential, or if similar gains could be achieved with a strong graph-based encoder and standard attention. \n- The paper directly adopts baseline results from EAGLE for turbulent flow experiments, but it is unclear whether Erwin’s model depth is aligned with the baselines. Additionally, the appendix lacks details on Erwin’s architecture for this experiment, making it difficult to assess the fairness of the comparison.\n- Erwin’s MPNN and the compared MPNN baseline rely on k-nearest neighbors, while a more general and physically meaningful approach is ball neighborhoods, as used in MGN paper.\n- Some important methodological and experimental details are missing (see questions), which significantly affect my recommendation.", "other_comments_or_suggestions": "Where is the anonymous link?", "questions_for_authors": "1. In Figure 9, the ball tree partitions appear to vary significantly across different datasets. Did you analyze how the choice of tree depth or splitting strategy affects performance? For example, in molecular dynamics, does the partitioning align well with physical structures like polymer chains, or are there cases where the tree creates unnatural groupings?\n2. The partitions in ShapeNet and EAGLE seem mostly aligned with coordinate axes. Given this, how does ball tree partitioning compare to KD-tree partitioning for these datasets? Would KD-trees provide similar efficiency, and what advantages does the ball tree offer in these cases?\n3. Erwin pads the ball tree to form a perfect binary tree, but the paper does not specify how virtual nodes are positioned or initialized. Are they assigned based on nearby real nodes, or are they purely placeholders? Could learnable pooling or adaptive padding reduce computational overhead?\n4. Table 1 ablates ball size for cosmology dataset, but how was ball size chosen for each datasets?\n5. How many layers does the MPNN in the embedding stage have?\n6. The paper introduces cross-ball connections by rotating the point cloud before constructing a second ball tree. What specific rotation matrix is used? Is it a fixed transformation, randomly sampled, or learned? How does the choice of rotation affect performance?\n7. The results suggest that MPNN performance plateaus as training data increases, but my experience, especially MGN models, is that they usually continue improving. Can you give me a suitable explanation?\n8. Did you ensure that all baseline models have comparable depth and capacity with Erwin?\n\nA clear response to these questions would help clarify my key concerns and could positively influence my evaluation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nPFR0sVTTC", "reviewer": "Reviewer_qhYr", "summary": "The authors propose a tree-based hierarchical transformer model to capture large-scale and small-scale interactions in exhaustively large datasets. The motivation is that traditional transformers are not scalable and efficient, so the authors propose a hierarchical model to allow course granting to capture large scale interactions that scales ~linearly with the data size. They test the proposed model on three empirical data sets.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback. We shall address their questions below. When discussing the questions regarding the cosmology dataset, we will refer to the original dataset paper by Balla et al.\n\n## Questions\n\n> \"The input is a point cloud $X \\in \\mathbb{R}^{5000 \\times 3}$\" -> I assume 3, the the 3-spacial dimensions. Can the authors elaborate on what time the positions and velocity vectors are measured? \n\nYes, that is correct - 3 stands for spatial dimensions. Regarding the measurement time, according to A.1. Details of simulation of Balla et al., it is the last step of a Quijote simulation (*present time*): \"Each simulation models the evolution of the large-scale structure of the Universe by following the dynamics of 5123 cold dark matter particles in a cubic comoving volume of side ∼ 1 Gigaparsec from redshift z = 127 to z = 0 (present time).\" The simulations are defined in Villaescusa-Navarro et al.\n\n> \"The total dataset includes 1184 simulations with 990 time steps per simulation. The dataset is split with 80% for training and 10% each for validation and testing.\" -> Can you describe how the training, validation, and test are selected?\n\nSince we use the benchmark, we use the predefined validation and test partitions as defined in Balla et al. (See https://zenodo.org/records/11479419 and https://github.com/smsharma/eqnn-jax/blob/main/benchmarks/galaxies/dataset.py for details). Training samples are taken randomly from the total number of 11,200.\n\n> I assume the proposed solution has application mainly to problems where data are in d-dimensional spatial domain. Are there any other applications that the author can think of where the unit balls are not defined in the spatial domain?\n\n1) One application that comes to mind is the relativistic case where particles live in spacetime, and one would need to adjust the algorithm to account for the pseudo-Euclidean structure.\n2) Related to spacetime, if one axis represents time, unit balls will not make much sense. In this scenario, however, one can simply build a ball tree based on the spatial coordinates and omit time during the construction. After all, a ball tree can be seen just as a means of organizing computation to linearize attention.\n\n## Other comments\n\n> \"halos form local clusters through gravity while maintaining long-range correlations originated from interactions in the early universe before cosmic expansion\" -> the universe has been expanding since the big bang. Also, the long range correlations are vanishing (unless a non-standard cosmology is assumed).\n\nRegarding the universe expansion, we thank the reviewer for their careful reading and for identifying this textual error. Indeed, what we meant to say was that the structures were initially in causal contact. We will remove \"before cosmic expansion\" in the revised manuscript. Regarding long-range correlations, although vanishing, they are still present according to the benchmark description (see Fig. 1 and \"Introduction, Information across scales\" in Balla et al.).\n\n## Experimental Designs Or Analyses\n\nWe appreciate the reviewer's critical assessment of the benchmarks used in our evaluation. We note that both cosmology and EAGLE datasets are benchmarks with established tasks/preprocessing/partitioning/losses. That is the reason we predict velocity from positions - it is the original task of the cosmology benchmark (see Balla et al., Section 3.3, Node-level prediction). For the same reason, we stick to the losses as proposed in the EAGLE benchmark. Overall, we used the original settings defined in the benchmarks for consistent comparison with other baselines validated on each benchmark.\n\n## Conclusion\n\nWe thank the reviewer for the valuable feedback and are hopeful that the reviewer will continue to support the paper. If the are any questions, we would be happy to answer them.\n\n## References\n\n- Balla, J., Mishra-Sharma, S., Cuesta-Lázaro, C., Jaakkola, T., & Smidt, T.E. (2024). A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processing. Learning on Graphs Conference, 2024.\n\n- Villaescusa-Navarro, F., Hahn, C., Massara, E., Banerjee, A., Delgado, A.M., Ramanah, D.K., Charnock, T., Giusarma, E., Li, Y., Allys, E., Brochard, A., Chiang, C., He, S., Pisani, A., Obuljen, A., Feng, Y., Castorina, E., Contardo, G., Kreisch, C.D., Nicola, A., Scoccimarro, R., Verde, L., Viel, M., Ho, S., Mallat, S., Wandelt, B.D., & Spergel, D.N. (2019). The Quijote Simulations. The Astrophysical Journal Supplement Series, 250.", "claims_and_evidence": "The claims around computational efficiency, linear scalability, and performance are evaluated empirically and backed by different experiments.", "methods_and_evaluation_criteria": "The methods and evaluation criteria, including the benchmarks, are suitable for the problem and demonstrate the superiority of the proposed solution.", "theoretical_claims": "N/A. No theoretical claims are made in this paper.", "experimental_designs_or_analyses": "The cosmology experiment has some logical gaps. This paper takes the location of galaxies (halos) and wants to predict their velocity. Location and velocity, even though correlated, are two independent random variables. Their correlation is due to the initial condition, therefore the model is not generalizable to another initial condition or universe with different properties. I am not sure what the motivation is behind setting up the problem in this way.\n\nMy intuition is that writing the loss function as the MSE of log Y for the turbulent fluid dynamics problem is more appropriate. The fractional error is important, not the actual value. The actual value has a unit and two quantities with different units (strictly speaking) should not be added together. However, fractional error does not have a unit and the fractional error of two quantities can be added.", "relation_to_broader_scientific_literature": "The proposed solution si appropriate for the target problems. The authors consider a set of scientific problems in which the traditional transformer-based models are not scalable and limited. Therefore, they propose a hierarchical solution that is suitable for the class of problems studied in this work and can be readily applied to those problems.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The main target audience is people who use AI for science applications, and I feel this ICML is an appropriate venue for this kind of publication.", "other_comments_or_suggestions": "- \"halos form local clusters through gravity while maintaining long-range correlations originated from interactions in the early universe before cosmic expansion\" -> the universe has been expanding since the big bang. Also, the long range correlations are vanishing (unless a non-standard cosmology is assumed).", "questions_for_authors": "- \"The input is a point cloud $X \\in \\mathcal{R}^{5000\\times3}$\" -> I assume 3, the the 3-spacial dimensions. Can the authors elaborate on what time the positions and velocity vectors are measured? \n\n- \"The total dataset includes 1184 simulations with 990 time steps per simulation. The dataset is split with 80% for training and 10% each for validation and testing.\" -> Can you describe how the training, validation, and test are selected?\n\n- I assume the proposed solution has application mainly to problems where data are in d-dimensional spatial domain. Are there any other applications that the author can think of where the unit balls are not defined in the spatial domain?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NWvJp193jo", "reviewer": "Reviewer_CTuS", "summary": "**Review updates after rebuttal**\nThe authors correspond to many of my commets well. As my initial score is already the positive side and I find it fair, hence I will keep my score unchanged.\n**End Review updates after rebuttal**\n\nThe paper introduces Erwin, a hierarchical transformer that combines ball tree partitioning with attention mechanisms to efficiently process large-scale physical systems defined on irregular grids. Erwin achieves linear-time attention by organizing computation hierarchically, enabling it to capture both fine-grained local details and global features. The method is validated across multiple domains (cosmology, molecular dynamics, and turbulent fluid dynamics), demonstrating state-of-the-art performance in both accuracy and computational efficiency.\n\nHowever, the paper has several areas that need improvement (details see in later sections):\n\n- No comprehensive reports for baselines;\n\n- Incomplete dataset descriptions;\n\n- Lack of literature review for multiscale GNNs for Irregular Grids;\n\n- Rushed presentation;\n\nOverall, I give weak accept at current phase; upon resolving the above bullet points, I will improve my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their encouraging feedback and hope to address their concerns with the rebuttal.\n\n## Experimental Designs Or Analyses\n> Not all baselines are applied to all datasets. For example, PointTransformer v3 is not included in the fluid dynamics experiments, and EAGLE is not included in the cosmology experiments. Suggestion: Include a comprehensive table listing all models and baselines for each dataset, with explanations for missing entries (e.g., due to incompatibility or heavy modifications required to the source).\n\nIn our experiments we did not aim at evaluating every baseline on every dataset. For each benchmark, except for MD (as it is not an established benchmark), we report results for baselines that have been previously evaluated on the benchmark. That being said, we do compare extensively against PointTransformer v3 when possible as it is the state-of-the-art model from computer vision and is very close to Erwin in spirit as it also suggest a way to linearize attention. Unfortunately, as the reviewer pointed out, we did not evaluated it on the fluid dynamics dataset as PTv3 is only implemented for 3D, while the benchmark is 2D. Doing so would require heavy modifications from us, particularly in serialization (https://github.com/Pointcept/PointTransformerV3/blob/main/serialization/z_order.py), which is the cornerstone of the approach. If there are specific experiments the reviewer would like us to run for the rebuttal, we are happy to do it.\n\n**Action taken**: In the final version, we will further explain our logic in baseline selection in the experimental section for each benchmark stressing out the comparison against PTv3. We also conducted additional experiments on ShapeNet-Car which now includes comparison to Transolver - state-of-the-art model for PDE modelling.\n\n## Supplementary Material\n> Only two datasets (cosmology, molecular dynamics) are introduced in detail. The fluid dynamics and airflow pressure datasets are mentioned briefly. Suggestion: Provide a coherent and complete introduction to all datasets, organized in a structured layout. Ensure all components (e.g., dataset statistics, preprocessing steps) are included.\n\n**Action taken**: We appreciate reviewer's suggestion and will use increased space in the final version to describe datasets in detail, particularly statistics and preprocessing steps. For the ease of comprehension, for each dataset, we will compile a table which describes\n- dataset statistics\n- train/validation/test split\n- input type (mesh/point cloud), feature type\n- target type\n- preprocessing (normalization, connectivity type, etc.)\n\nIf there are any other descriptors that reviewers would like to be included, we would be happy to do so.\n\n## Essential References Not Discussed\n**Action taken**: We appreciate the reviewer's suggestions and will use the increased page count to expand the related works section and include multiscale frameworks.\n\n## Conclusion\nWe are thankful to the reviewer for the valuable suggestions that we believe will improve the presentation of our framework. If the are any questions, we would be happy to answer them.", "claims_and_evidence": "Claim: Erwin achieves linear-time attention and outperforms baseline methods in accuracy and computational efficiency.\n\nEvidence:\n\nFigure 5~7: Shows linear scaling of runtime, lower MSE, and faster runtime compared to baselines\nTable 3, 4: lower RMSE and MSE on two datasets.\n\nClaim: Erwin captures long-range interactions and multi-scale phenomena effectively.\n\nEvidence:\n\nFigure 6~8", "methods_and_evaluation_criteria": "Methods:\n\n- Ball tree partitioning to organize hierarchy.\n\n- UNet-like structure, or Pointnet++, SWIN transformer, etc.\n\nEvaluation Criteria:\n\nAccuracy: Measured by MSE (cosmology, airflow pressure), NLL (molecular dynamics), and RMSE (fluid dynamics).\n\nEfficiency: Measured by runtime and memory usage.\n\nScalability: Demonstrated through linear scaling of runtime with the number of nodes.", "theoretical_claims": "NA; no concern: The method builds on mature techniques (ball trees, attention mechanisms) and provides solid incremental improvements.", "experimental_designs_or_analyses": "The paper includes experiments across multiple domains (cosmology, molecular dynamics, fluid dynamics, airflow pressure), demonstrating the generalizability of Erwin.\n\nThe results show consistent improvements in both accuracy and computational efficiency.\n\nWeaknesses:\n\nNot all baselines are applied to all datasets. For example, PointTransformer v3 is not included in the fluid dynamics experiments, and EAGLE is not included in the cosmology experiments.\n\nSuggestion: Include a comprehensive table listing all models and baselines for each dataset, with explanations for missing entries (e.g., due to incompatibility or heavy modifications required to the source).", "relation_to_broader_scientific_literature": "The paper builds on hierarchical attention methods and tree-based algorithms from computational physics. It also draws inspiration from vision transformers (e.g., SwinTransformer) and point cloud transformers (e.g., PointTransformer v3).\nThe authors cover regular grids, particles, transformers, but miss some related literature for multiscale structures on irregular mesh, see below section for suggested refs to add.", "essential_references_not_discussed": "Multiscale learning on irregular mesh\n- Multi-scale rotation-equivariant graph neural networks for unsteady Eulerian fluid dynamics\n- Efficient Learning of Mesh-Based Physical Simulation with Bi-Stride Multi-Scale Graph Neural Network\n- Learning Distributions of Complex Fluid Simulations with Diffusion Graph Networks", "other_strengths_and_weaknesses": "Weaknesses:\n\n- Incomplete sections: Dataset introduction, experimental designs, baseline comparisons, and supplementary material need refinement.\n\n- Rushed presentation: The paper feels slightly rushed, with some sections (e.g., dataset descriptions) lacking detail. Many tables in appendix lack text accompanying them.\n\n- Suggestion: Revise the paper to address these issues, ensuring completeness and coherence.", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "NA"}}]

Author rebuttal:

Meta-review: This paper presents Erwin, a tree-based hierarchical transformer tailored for large-scale physical systems defined on irregular grids. By leveraging ball tree partitioning and cross-ball interactions, Erwin achieves linear-time attention, capturing both local and global interactions effectively. The model is evaluated across diverse scientific domains, including cosmology, molecular dynamics, and turbulent fluid simulations, and demonstrates improvements in both accuracy and computational efficiency.

The authors spent quite some time on the rebuttal dispersing some/most(?) of the initial criticism which lead to increased scores and the general sentiment is in favor of the paper (averaging around weak accept) and so am I. I recommend weak acceptance; a bit of rework could have improved the submission further and I sincerely hope that the authors are doing so in their camera-ready or the revision.
Decision: Accept (poster)