Rethinking the role of frames for SE(3)-invariant crystal structure modeling

Authors: Yusei Ito, Tatsunori Taniai, Ryo Igarashi, Yoshitaka Ushiku, Kanta Ono
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Materials Science, Invariant Networks, Transformer, Physics-Informed ML

Abstract:
Crystal structure modeling with graph neural networks is essential for various applications in materials informatics, and capturing SE(3)-invariant geometric features is a fundamental requirement for these networks. A straightforward approach is to model with orientation-standardized structures through structure-aligned coordinate systems, or  “frames.” However, unlike molecules, determining frames for crystal structures is challenging due to their infinite and highly symmetric nature. In particular, existing methods rely on a statically fixed frame for each structure, determined solely by its structural information, regardless of the task under consideration. Here, we rethink the role of frames, *questioning whether such simplistic alignment with the structure is sufficient*, and propose the concept of *dynamic frames*. While accommodating the infinite and symmetric nature of crystals, these frames provide each atom with a dynamic view of its local environment, focusing on actively interacting atoms. We demonstrate this concept by utilizing the attention mechanism in a recent transformer-based crystal encoder, resulting in a new architecture called **CrystalFramer**. Extensive experiments show that CrystalFramer outperforms conventional frames and existing crystal encoders in various crystal property prediction tasks.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1729543439864

thread_type:
official_review

<thread_object>
creator: reviewer_dkni

<content>
summary:
This work explores an alternative, more powerful version of frame averaging networks and applies it to crystal property prediction tasks. Instead of using a fixed lattice frame or one calculated from the PCA of the moment tensor of the atoms, the authors propose to learn "dynamic frames" for each individual atom. They propose two methods to do this: Either by calculating the PCA from weighted atoms or constructing the frames by greedily choosing the highest weighted atoms, where the weights in both cases are equal to the attention weights in their transformer architecture. They show that this results in significant performance gains and that the model exhibits intuitive, sensible behaviour, where the earlier layers attend to closer atoms, and later layers attend to further away atoms.

strengths:
Crystal property prediction is an important topic, and (stochastic) frame averaging is a promising direction for the task. Using the novel idea of dynamic frames, the proposed method improves stochastic frame averaging for crystal property prediction by a large margin. Therefore, the contribution seems significant and novel. It is believable that the proposed "dynamic frames" make the model more expressive, although I find it hard to get an intuitive understanding of what the angles between the learned frames and direction vectors could learn or why they could represent useful features. The experimental results show strong improvements on two datasets and against the base model (Crystalformer) and other state-of-the-art models.

weaknesses:
The biggest weakness of this paper is, in my opinion, the writing. Many sentences are grammatically incorrect or use strangely constructed, overly long sentences. This makes it much harder to read the paper than it should be. A non-exhaustive list of examples:

Line 121: "... are linear projects of current state x, σi is a tail-length variable
of Gaussian distance-decay attention adaptively derived from xi, ψij(n) is geometric position
embedding that encodes interatomic". This misses several articles, and I assume the authors mean "projections" not "projects"

Line 198: "... and local view of entire crystal ..." misses the article "the"

Line 234: "This perturbation scheme is considered to implement the stochastic FA" I am not sure what this sentence means

I strongly recommend that the authors use a grammar tool like Grammarly or an LLM to help streamline their writing.

questions:
In the runtime comparison, Table 4: How can PotNet take only 43s/epoch to train but 313ms for inference, while CrystalFramer takes almost double the time to train with 74s but is nearly ten times faster during inference? Can you explain this large discrepancy?

In line 298: The model uses the cosine of the angles. This would map pi/2 and -pi/2 to the same value 0. Have you considered using a concatenation of both sine and cosine of the angles as features? It might make the model more expressive.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730611059105

thread_type:
official_review

<thread_object>
creator: reviewer_sjyw

<content>
summary:
The paper proposes a new method for modeling crystal based on dynamic frames that are based on the local neighborhood of each atom in the crystal structure. The crystal frames are then processed by a transformer based architecture called CrystalFramer. The paper starts by motivating the utility of geometric deep learning methods that can capture relevant physical inductive biases, such as invariance. Previous work captured such geometric information using frame averaging, which is based on taking canonical frames of the entire structure. By contrast, the paper proposes dynamic frames which are based on the local coordinate system of a given atom. Section 2 then outlines relevant preliminary details with relevant mathematical definitions, including the description of a crystal structure, how transformers have been applied to crystal structure modeling and prior work on frame averaging for molecular and materials modeling.  

Section 3 introduces and defines the concept of dynamic frames, which is the main part of the proposed new method. This includes the frame definitions in Section 3.1 and the CrystalFramer architecture in Section 3.2 that is based on Gaussian basis functions. This is followed by a description of related work on invariant features, frames and equivariant features in Section 4. 

Section 5 outlines the main experiments in the paper, including materials property prediction and an efficiency analysis . The results generally show that the proposed dynamic frames improve Crystalformer modeling performance across all datasets while achieving best performance in most cases. The efficiency analysis shows that CrystalFramer incurs an extra compute cost. Section 6 provides a visual analysis of the dynamic frames, as well as discussion on limitations along with suggestions for using dynamic frames for equivariant prediction and molecular modeling.

strengths:
* The idea of dynamic frames proposed in the paper is a novel and effective way to include relevant inductive biases for crystal structure modeling.
* The paper is well written with relevant details described and convincing experiments followed by detailed analysis.
* The paper candidly discusses advantages and limitations of the proposed method.

weaknesses:
* The paper could be further with additional discussion of related work for machine learning potentials [1 egraff]. The current focus of the study is mainly on properly modeling, so this would probably fit as potential future work. There be advantages and disadvantages of using CrystalFramer in such a context. Similarly, the paper can also discuss potential application to modeling of crystals with surfaces, such as catalyst [2 - OCP].
* The dynamic frames provide a new dimension for analysis for how local neighborhoods interact in geometric deep learning. The paper could be strengthened by discussing the frames across different materials and/or different stages of the training process.

questions:
* Could you talk more about scalability of your method? This would be especially relevant when modeling larger supercell that have more atoms.
* Did you notice an evolution of the dynamic frames as training progressed? Are there consistencies across some types of materials?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730684169289

thread_type:
official_review

<thread_object>
creator: reviewer_bzky

<content>
summary:
This paper proposes CrystalFramer and dynamic frames to capture SE(3)-invariant geometric features for crystal property prediction. Unlike using static frame in previous work, this work use dynamics frames to give each atom its own dynamics view of the structure, focusing only on those atoms actively interacting with it, meanwhile also accommodating the infinite and symmetric nature of crystal structures. The experiment shows superior performance of the proposed method in several prediction tasks.

strengths:
- This work proposes dynamic frames for crystal features modeling, which can dynamically accounting for the atoms actively engaged in learned interactions in each interatomic message-passing layer.
- The proposed method achieves superior performance on several commonly used benchmarks in crystal property prediction.

weaknesses:
- Previous relevant frame-based method ComFormer also provides equivariant version. I am wondering if this method can also be easily extended to SE(3)-equivariant? There’re are many other equivariant properties for the materials such as force or high-order tensors. I would like to see the performance of equivariant version of this method if the extension is straightforward. Even for the invariant properties, it’s interesting to know whether equivariant network can help invariant properties prediction.

questions:
- Please refer to the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730701735825

thread_type:
official_review

<thread_object>
creator: reviewer_akn2

<content>
summary:
The authors change the fixed, static frames for SE(3)-invariant modeling in crystal structures into dynamic frames that adapt to each atom's local environment through the network's attention mechanism, then integrate into a new architecture **CrystalFramer**. Experimental results on large-scale datasets (JARVIS, MP, OQMD) demonstrate that **CrystalFramer** with dynamic frames outperforms models using static frames in several material property prediction tasks.

strengths:
1. The proposed method is straightforward and presented with commendable clarity.
2. The shift from static to dynamic frames enables atom-specific adjustments, significantly enhancing SE(3) invariance handling in highly symmetric crystal structures. Furthermore, the extensive evaluations on multiple datasets (JARVIS, MP, OQMD) highlight the effectiveness of dynamic frames, with the "max frames" variation demonstrating particular superiority.

weaknesses:
### Major
How do the "max frame" construction and its usage within the networks differ from those proposed in [1], [2], [3], and [4]? In particular, [1] and [4] use similar atomic local frames to address equivariance in molecular structures. A comparative discussion of these methods would be valuable to clarify whether the proposed approach is a direct adaptation from molecular modeling to materials or if there are novel insights beyond those methods.

### Minor
Using frame-based techniques will hinder the model's continuity, potentially affecting its generalization capabilities, as noted in [5]. Does the proposed model address this limitation? It will be valuable to discuss the continuity of the proposed dynamic frame approach and empirically analyze if this affects generalization in the experiments.



[1] "SE (3) equivariant graph neural networks with complete local frames." Du, et al.

[2] "Smooth, exact rotational symmetrization for deep learning on point clouds." Pozdnyakov, Sergey, and Michele Ceriotti. 

[3] "Equivariance via Minimal Frame Averaging for More Symmetries and Efficiency." Lin, et al. 

[4] "A new perspective on building efficient and expressive 3D equivariant graph neural networks." Du, et al. 

[5] "Equivariant frames and the impossibility of continuous canonicalization." Dym, Nadav, Hannah Lawrence, and Jonathan W. Siegel.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731487132048

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
The authors are preparing point-by-point responses to each reviewer

comment:
Dear Reviewers and Meta-Reviewers, 

We sincerely appreciate the time and effort you have put into reviewing our work. We are pleased to receive the positive feedback, including "excellent" ratings for both presentation and contribution. We are confident that the remaining concerns raised by Reviewers AkN2 and dkNi will be addressed in our forthcoming responses, which we are currently preparing. 

We will provide point-by-point responses to each Reviewer within a few days.
Thank you again for your valuable feedback and commitment.

Sincerely,  
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732247332736

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Request for Reactions from Reviewers

comment:
Dear All Reviewers and Area Chairs,

We have responded to all the Reviewers' comments in their respective threads and updated the manuscript accordingly. As the end of the rebuttal period is approaching, we sincerely ask all Reviewers for their reactions, such as adjusting scores or posting follow-up questions. Below, we summarize the **main concerns and our responses** to assist Reviewers and Area Chairs.

---
### **Main Concerns and Our Responses**
- **Novelty Compared to Existing Local Frames for Molecules** (Reviewer AkN2):
  - Our work introduces **a clear conceptual difference** by proposing *dynamic frames* that account for models' internal states, whereas existing *static frames* do not. (See `Appendix B` for detailed discussions.)
  - Additional experiments validate the **superiority of (dynamic) max frames over a static counterpart using static local frames** (see `Tables 1 and 2` and  `Sec 5.1`).
- **Clarity and Writing** (Reviewer dkNi):
  - Three reviewers rated the presentation as "excellent" (4) and "good" (3) with positive comments, such as “presented with commendable clarity” (Reviewer AkN2) and “paper is well written” (Reviewer SjyW)
  - We have **revised texts throughout to further enhance readability** (see `texts highlighted in green`).
- **Application to Equivariant Tasks** (Reviewer bzKy & Reviewer SjyW):
  - This was raised as **a constructive suggestion** rather than a critical weakness by two positive reviewers.
  - However, **addressing both invariant and equivariant tasks in a single study is challenging**, given the complexities of equivariant tasks and crystal-structure modeling as well as the recent demand for extensive evaluations.
  - Prior works on crystals (even the mentioned eComFormer) also focus solely on invariant tasks.
  - We have incorporated **references to equivariant-task benchmarks (Open Catalyst and EGraFFBench) for future research**, as suggested by Reviewer SjyW. (See `"Equivariant prediction" in Sec 6`)

---
### **Summary of Updates**
**Highlight colors** in the revised manuscript are intended as follows
- `Red`: Major changes introducing **new information**, such as additional references, discussions, and results.
- `Blue`: Relatively minor changes to enhance **clarity and understanding**, including improved interpretations and corrections of inaccuracies.
- `Green`: Minor changes made to improve **overall readability** (e.g., breaking up long sentences), addressing the comments from Reviewer dkNi.

**◆ Major Changes** (red highlights)
- New comparative discussions on **existing local frames for molecules** (Du et al., 2022; 2023;
Pozdnyakov & Ceriotti, 2023) in **`Appendix B`** and `Sec 4 (last paragraph)`. [Reviewer AkN2]
- New **benchmark results for a static local frame method** in `Tables 1 and 2` and `Sec 5.1`.  [Reviewer AkN2]
- New experimental analysis on **model's continuity** in **`Appendix H`** and `"Visual analysis" in Sec 6`.  [Reviewer AkN2]
- New visual analysis on **frame evolution during training** in **`Appendix F.3`** and `"Visual analysis" in Sec 6`.  [Reviewer SjyW]
- New dicussions on **scalability for large structures** in **`Appendix G`** with reference from `Sec 5.2`.  [Reviewer SjyW]
- New references to **equivariant task benchmarks (Open Catalyst and EGraFFBench)** as important potential applications in `"Equivariant prediction" in Sec 6`.  [Reviewer bzKy & Reviewer SjyW]
- New reference to **the Minimal Frame Averaging framework (Lin et al., 2024)** in `"Frames" in Sec 4`.  [Reviewer AkN2]

**◆ Minor and Other Changes** (blue and green highlights)
- **Enhanced readability** throughout the paper (green parts) [Reviewer dkNi]
- Added **intuitive interpretations of the frame-based edge features (Eq 7)** in `Sec 3.2`.  [Reviewer dkNi]
- Clarified **the reason for slower inference times of existing methods** at `the end of Sec 5.2`. [Reviewer dkNi]
- **Shrinked `Tables 1 and 2`** and moved the full results to `Appendix E` (`Tables A2 and A3`)
- *Clarified the contents of Appendix F* in `"Visual analysis" in Sec 6`. [Reviewer SjyW]
- *Clarified the relation between the perturbation sheme and stochastic FA* in `"Max frames" in Sec 3.2`. [Reviewer dkNi]
- *Corrected the description of lattice frames* in `"Sec 2.3`.
- *Corrected the architectural description (i.e., final "linear" → "feed forward")* in `"Overall architecture" in Sec 3.2` and `Figure 2`.

---
Given these updates, we feel the depth and quality of our paper have been greatly enhanced. We believe that these revisions, together with our responses in the review threads, adequately address all concerns raised by the Reviewers.

We sincerely appreciate your time and thoughtful feedback.

Best Regards,  
The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors presented a new architecture CrystalFramer for frame-based SE(3)-invariant crystal structure modeling, which introduces new dynamic local frames instead of existing static frames to better capture symmetry and equivariance in crystal materials with experiments showing better crystal property prediction performances. 

The proposed dynamic frame as well as corresponding attention mechanisms and invariance neural network for message passing is an interesting methodological contribution for crystal materials modeling, with strong empirical results.