Bridging Geometric States via Geometric Diffusion Bridge

Summary

The paper introduces the Geometric Diffusion Bridge (GDB), a novel framework designed to generate the evolution of geometric states in geometric (coordinate) systems. GDB uses a diffusion bridge connecting initial and target geometric states with equivariant transition kernels, preserving symmetry and joint state distributions. Furthermore, GDB can use a chain of equivariant diffusion bridges to leverage trajectory data for more accurate dynamic modeling.

Strengths

- The presentation of theorems in Section 3.1 is clear and straightforward, establishing a solid theoretical foundation for GDB. The authors effectively derive theorems and integrate them with point cloud states. - GDB demonstrates strong performance across various tasks, including QM9, Molecule3D, and OpenCatalyst IS2RS.

Weaknesses

I have no complaints regarding the technical and experimental sections, as they are well-written. However, I wonder existing works, such as [1] and [2], also use diffusion bridges over molecular data. What advantages does your approach have over theirs? [1] Diffusion-based Molecule Generation with Informative Prior Bridges. Lemeng Wu, et al. NeurIPS 2022. [2] DiSCO: Diffusion Schrödinger Bridge for Molecular Conformer Optimization. Danyeong Lee, et al. AAAI 2024.

Questions

See weaknesses.

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors note the need of exploring better implementation strategies for their framework to enhance performance.

Summary

This paper proposes a generative model for bridging initial and target geometric states using diffusion bridge. This work introduces an equivariant diffusion bridge based on equivariant transition kernels for symmetry constraints. The proposed method was validated on diverse settings including simple molecules and adsorbate-catalyst complex, outperforming previous MLFF baselines.

Strengths

- The motivation of using diffusion bridge to bridge initial and target geometrical states is reasonable. - Using diffusion bridge model for equilibrium state prediction and structure relaxation is novel to the best of my knowledge, and the paper shows that GDB significantly outperforms previous methods with diverse datasets. - Equivariant design of bridge process is based on solid theory. - The paper is well written except for some missing relevant works on diffusion bridge.

Weaknesses

- Related works on diffusion bridges or diffusion mixtures were not discussed. Diffusion bridges has been studied in [1,2,3,4] with applications to molecules, graphs, point clouds, and images, and more recent works have studied general framework for diffusion bridges [5, 6] which is worth discussing. While GDB has a contribution for using diffusion bridges in new tasks, discussing related works and clarifying the novel contributions is necessary in particular for strengthening the contribution of this work. - Contribution seems limited as using diffusion bridge as generative modeling was already studied [1,2,3,4], in particular deriving diffusion bridges using Doob's h-transform. Designing an equivariant diffusion process (not necessarily bridge) specifically in SE(3) group has been covered in [7,8, 9]. What is the difference of designing equivariant diffusion bridges compared to equivariant diffusion processes? [1] Peluchetti, Diffusion Bridge Mixture Transports, Schrodinger Bridge Problems and Generative Modeling, JMLR 2023 [2] Liu et al., Learning Diffusion Bridges on Constrained Domains, ICLR 2023 [3] Wu et al., Diffusion-based Molecule Generation with Informative Prior Bridges, NeurIPS 2022 [4] Jo et al., Graph Generation with Destination-Predicting Diffusion Mixture, arXiv 2023 [5] Albergo et al., Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, arXiv 2023 [6] Shi et al., Diffusion Schrodinger Bridge Matching, NeurIPS 2023 [7] Xu et al., GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation, ICLR 2022 [8] Xu et al., Geometric Latent Diffusion Models for 3D Molecule Generation, ICML 2023 [9] Yim et al., SE(3) diffusion model with application to protein backbone generation, ICML 2023

Questions

- What is the reason for using deterministic process (i.e., probability flow ODE) instead of the original stochastic process? Does ODE results in better performance? - Is GDB scalable to geometric states of high dimensions? While analysis on this may not be necessary, it could strengthen the work.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

While the paper discusses future direction for the proposed method, specific limitations of the work is not specified. One potential issue might the scalability of GDB as the model has transformer architecture, and other issue could be long inference time which is a typical problem of diffusion models.

Summary

This paper proposes a type of diffusion model that captures the evolution of geometric states. The model is characterized by a diffusion SDE that couples the initial state with the target state, in the middle of which trajectory guidance is enabled when such data present. The framework is designed to yield equivariant density similar to other geometric diffusion models. Experiments on equilibrium state prediction with or without trajectory data have been performed to verify the applicability of the proposed approach.

Strengths

1. The distinction between existing works has been elaborated in Table 1, which is clear. 2. The method is designed with an option to leverage additional trajectory data, which is quite interesting.

Weaknesses

1. The experimental setup and comparison with baselines on equilibrium state prediction is a bit troublesome which requires more clarification or additional comparisons. Please refer to Q1. 2. The presentation is a bit unclear. Please refer to Q2. 3. Additional baselines may be considered. The baselines selected in the paper are not closely connected to the proposed approach. See Q3. 4. Missing ablation studies. In the current shape it is unclear where the performance gain comes from. See Q4.

Questions

Q1. The evaluation protocol on QM9 and Molecule3D, especially to compare with direct prediction approaches, is not a common practice. A more convincing benchmark protocol would be to compare with methods such as GeoDiff [1] on molecule generation tasks since they are also generative models. Since the paper is positioned to tackle generative modeling, the experiments should also be designed to align with the goal. Q2. Could the authors provided detailed sampling algorithm this approach adopts? If the model uses sampling approach similar to other diffusion models, there should be related discussions on sampling steps/sampling time the method consumes. Q3. A more reasonable baseline would be to directly apply existing bridge models (e.g., [2]) to the current task by switching the backbone to the one this paper adopts. This would help the audience understand the unique contribution of this work since both bridge models and equivariant (geometric) diffusion models have been proposed in literature. Q4. Ablation studies such as investigating the importance of preserving equivariance of the density modeled should be included. This would help justify the necessity of the proposed components. [1] Xu et al. GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation. In ICLR'22. [2] Zhou et al. Denoising Diffusion Bridge Models. In ICLR'24.

Rating

3

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

There seem to be no discussions on limitations in the paper. It would be better to discuss potential limitations from perspectives such as scalability and sampling time.

Summary

In this paper, the authors introduce a Geometric Diffusion Bridge (GDB) framework, which aims to predict the evolution of geometric states in complex systems accurately, crucial for fields such as quantum chemistry and material modeling. Traditional methods face computational challenges, while deep learning approaches lack precision and generality. The authors use Doob’s h-transform to construct an equivariant diffusion bridge. By applying Doob’s h-transform, the authors adjust the SDE to ensure that the process starts from an initial geometric state and is conditioned to reach a target geometric state. This ensures that the transformed process respects the symmetry constraints of the geometric states, leading to more accurate and physically meaningful predictions.

Strengths

+ The framework utilizes an equivariant diffusion bridge derived from a modified Doob’s h-transform. This ensures that the diffusion process respects symmetry constraints, making the predictions more robust and reliable. + The paper provides a theoretical framework analysis about preserving symmetries and accurately modeling evolution dynamics. + Experimental evaluations show that GDB is better than state-of-the-art approaches in various real-world scenarios, including equilibrium state prediction and structure relaxation tasks. + The framework achieves significant error reduction compared to strong baseline models, particularly in challenging tasks such as structure relaxation in the Open Catalyst 2022 dataset

Weaknesses

- The framework, especially when leveraging trajectory data, might introduce significant computational overhead. The simulation-free matching objective is designed to be efficient, but the overall framework’s computational demands might still be high - Some mathematical notations and definitions in the paper could be made clearer. For instance, explicitly defining all variables and functions used in the modified Doob’s h-transform and constructing equivariant diffusion bridges would improve readability and understanding.

Questions

See above

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

No limitations are addressed in the paper by the authors