Erwin: A Tree-based Hierarchical Transformer for Large-scale Physical Systems

We thank the reviewer for their encouraging feedback and hope to address their concerns with the rebuttal.

Experimental Designs Or Analyses

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).

In 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.

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.

Supplementary Material

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.

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

• dataset statistics
• train/validation/test split
• input type (mesh/point cloud), feature type
• target type
• preprocessing (normalization, connectivity type, etc.)

If there are any other descriptors that reviewers would like to be included, we would be happy to do so.

Essential References Not Discussed

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.

Conclusion

We 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.

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.

Questions

"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?

Yes, 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.

"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?

Since 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.

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?

1. 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.
2. 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.

Other comments

"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).

Regarding 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.).

Experimental Designs Or Analyses

We 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.

Conclusion

We 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.

References

• 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.

• 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.

We thank the reviewer for the actionable feedback. We address their concerns and questions in the following.

Questions For Authors

Q1

That 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).

Regarding 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).

Q2

For 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.

Q3

We 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.

Q4

Ball 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.

Q5

We 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$.

Q6

For 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.

Q7

The 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.

Q8

We 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.

• 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.).
• 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.
• ShapeNet, all models have similar parameter count - between 4M and 5M parameters.

Conclusion

Unfortunately, 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.

We thank the reviewer for their encouraging feedback. We address their questions below.

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".

We apologize for the inconvenience. An anonymized version of the repo can be found here: https://github.com/erwin-transformer/erwin.

During the simulation process, the positions of particles will change. Does the proposed method re-build the tree after each step?

Rebuilding the tree depends entirely on the properties of the modelled system, particularly, how quickly the arrangement of particles changes:

• The mesh is static: the tree should only be computed once during preprocessing.
• The mesh changes slowly: the tree should be rebuilt each N step.
• The mesh is dynamic and changes quickly: ideally, the tree should be rebuilt at every step of a simulation.

During 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.

Additional experiments

We also conducted additional experiments to further highlight the strength of our method.

OctFormer on the molecular dynamics task

We 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.

Ablation study on ShapeNet-Car dataset

We 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)

Notice 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.

Conclusion

We are grateful for the reviewer's feedback. Should you have any questions or concerns, we are always ready to discuss them further.