Title: 3D Structure Prediction of Atomic Systems with Flow-based Direct Preference Optimization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Flow Matching

Direct Preference Optimization

Geometric Graph Neural Networks

Structure Prediction
Abstract: Predicting high-fidelity 3D structures of atomic systems is a fundamental yet challenging problem in scientific domains. While recent work demonstrates the advantage of generative models in this realm, the exploration of different probability paths are still insufficient, and hallucinations during sampling are persistently occurring. To address these pitfalls, we introduce FlowDPO, a novel framework that explores various probability paths with flow matching models and further suppresses hallucinations using Direct Preference Optimization (DPO) for structure generation. Our approach begins with a pre-trained flow matching model to generate multiple candidate structures for each training sample. These structures are then evaluated and ranked based on their distance to the ground truth, resulting in an automatic preference dataset. Using this dataset, we apply DPO to optimize the original model, improving its performance in generating structures closely aligned with the desired reference distribution. As confirmed by our theoretical analysis, such paradigm and objective function are compatible with arbitrary Gaussian paths, exhibiting favorable universality. Extensive experimental results on antibodies and crystals demonstrate substantial benefits of our FlowDPO, highlighting its potential to advance the field of 3D structure prediction with generative models.

3D Structure Prediction of Atomic Systems with Flow-Based Direct Preference Optimization

Abstract

Predicting high-fidelity 3D structures of atomic systems is a fundamental yet challenging problem in scientific domains. While recent work demonstrates the advantage of generative models in this realm, the exploration of different probability paths are still insufficient, and hallucinations during sampling are persistently occurring. To address these pitfalls, we introduce FlowDPO, a novel framework that explores various probability paths with flow matching models and further suppresses hallucinations using Direct Preference Optimization (DPO) for structure generation. Our approach begins with a pre-trained flow matching model to generate multiple candidate structures for each training sample. These structures are then evaluated and ranked based on their distance to the ground truth, resulting in an automatic preference dataset. Using this dataset, we apply DPO to optimize the original model, improving its performance in generating structures closely aligned with the desired reference distribution. As confirmed by our theoretical analysis, such paradigm and objective function are compatible with arbitrary Gaussian paths, exhibiting favorable universality. Extensive experimental results on antibodies and crystals demonstrate substantial benefits of our FlowDPO, highlighting its potential to advance the field of 3D structure prediction with generative models.

1 Introduction

Predicting 3D structures of atomic systems is indispensable in various scientific domains, ranging from pharmaceutical drug design [1, 17] to materials science [7]. Accurate 3D modeling is not only crucial for understanding the physical and chemical properties of substances at the atomic level [2, 18] but also for simulating and predicting their behavior in various environments [3, 26]. Nevertheless, it remains challenging due to the intricate nature of atomic interactions, the vastness of the conformational space, as well as limited resources of structure data.

Conventional methods typically employ physics-based algorithms to derive structures at local energy optimum [24, 31, 32]. Recent advancements leverage deep generative model to learn the distribution of stable structures from available data, showcasing remarkable success across various domains. For example, DiffAb [20] designs a diffusion-based method for antigen-specific antibody design, which is further available for antibody structure prediction, and DiffCSP [14] proposes a joint diffusion framework for crystal structure prediction. Despite these advancements, generative models for structure prediction are confronted with two primary challenges.

First, existing structure prediction methods predominantly utilize diffusion-based generative models. While effective, this focus narrows the scope of exploration into other probability paths that could

𝑐

!

⋯

≻

≻

𝑥

!

% 𝑥

!"

!

% 𝑥

!"

"

!

,

% 𝑥

!"

)

𝛿

(b) Preference Dataset Construction

𝜃

#$%

𝜃

#$%

𝜃

&'(

≻

≻

⋯

𝒟

#$%

𝒟

')*#

ℒ

+,-

(c) Direct Preference Optimization

Init.

ℒ

./

ℒ

.,

ℒ

-0

𝜃

#$%

𝜃

#$%

𝜃

#$%

(a) Training Flow Models

potentially offer substantial benefits. A notable example is the Optimal Transport (OT) path, which has recently been demonstrated to be particularly effective in the field of molecular generation [28]. Second, current training paradigm frequently leads to hallucinated distribution peaks [1]. Most generative models are trained through maximizing the likelihood or its lower bound on the groundtruth structures, which are easily haunted by hallucinations due to the lack of negative samples during training. In the field of natural language processing or computer vision, Direct Preference Optimization (DPO) [25, 30] is proposed to align the model with human preferences, which effectively reduces hallucinations. For 3D structure prediction, such preferences can be naturally extended to similarity with the reference structure ( e.g. RMSD). However, it remains unclear whether the DPO method is compatible with arbitrary probability paths.

To address the above pitfalls, we introduce FlowDPO, a novel framework that explores flexible selection of Gaussian paths and enhances the quality of generated structures by alignment with the reference distribution. Specifically, we approach the structure prediction task via flow matching models regarding various paths. Given a pre-trained flow matching model, we sample multiple structures for each entry in the training set, evaluate these candidates against known ground truths to compute similarity, and construct an automatic preference dataset. Notably, we theoretically derive the unified objective of DPO for arbitrary Gaussian paths, and leverage the preference to enhance the performance of the original generative model. Intuitively, such a paradigm not only augment data with self-distilled samples, but also endow the model with the ability to distinguish between high-fidelity and hallucinated samples.

In summary, our contributions are threefold:

We explore multiple accessible probability paths for the 3D structure prediction task, and to the best of our knowledge, we are the first to theoretically prove the compatibility of DPO with arbitrary Gaussian paths by deriving a unified objective.

Based on the theoretical results, we develop a novel framework to encourage better alignment of flow matching models with desired reference distribution in 3D structure prediction, which effectively suppresses the probability of hallucinations.

Our approach yields promising results on antibody and crystal structure prediction tasks, showcasing the versatility and efficacy of our FlowDPO.

2 Related Work

Structure Prediction for Atomic Systems. 3D Structure prediction, including predicting conformations from molecular topological graphs [34], determining unit cell structures from crystal compositions [14], or inferring structures based on protein sequences [1], is crucial in computational chemistry and material science. Traditionally, these predictions have relied on physics-inspired scoring functions [21, 11] or density functional theory (DFT)-based energy calculations [10] to define the search space, with subsequent application of search algorithms to identify optimal structures. Recently, deep generative methods, particularly diffusion models [12, 27], have proven to be highly effective in this field. These models have been successfully applied across multiple specific domains, including small molecules [34], crystals [14], antibodies [20], complexes [6], and general biomolecules [1]. The emergence of flow matching models [19], which generalize diffusion paths to more flexible probability flows, has further enhanced the generative capabilities for geometric graphs [28]. The goal of our work is to explore structure prediction from the perspective of flow matching, and align these models towards more accurate predictions.

Aligning Generative Models. In the domain of generative model alignment, recent work has focused on refining models to better meet human preferences. Direct Preference Optimization (DPO), introduced by [25], offers a significant advancement over traditional Reinforcement Learning from Human Feedback (RLHF, [23]) methods by directly optimizing a policy based on human preference data. This approach has proven effective in aligning large language models (LLMs) with user expectations. Extending this concept, [30] propose Diffusion-DPO, a novel method that adapts DPO for text-to-image diffusion models. By reformulating the preference optimization for diffusion model likelihoods, Diffusion-DPO achieves state-of-the-art performance in generating images that are not only visually appealing but also closely aligned with textual prompts. Recently, [36] introduces ABDPO, a DPO-based method tailored for antibody design. Unlike ABDPO, which concentrates on guiding diffusion models to generate antibody candidates with lower energy, our approach emphasizes aligning flow models for precise structure predictions.

3 FlowDPO

3.1 Flow Matching for Geometric Graphs

Flow Matching (FM, [19]) is a general paradigm for generative tasks by learning a vector field to connect the pre-defined prior distribution with the targeted data distribution. Let q denote the data distribution, x 0 is a data point acquired from p 0 = q , and x 1 is a random sample from the prior distribution p 1 . A time-dependent flow ψ t is then defined to shift samples from the prior distribution to the time-dependent distribution p t via the vector field v t , that is

The vector field can be further parameterized by a time-dependent model v θ ( x t , t ) , leading to the continuous normalizing flows (CNFs, [5]). To avoid numerical ODE simulations to train v θ , FM simplifies the training target by aligning the model with a pre-defined vector field u t to yield p t , i.e. ,

However, as p t is still unknown, we are still unable to sample x t and apply the above objective. To address this gap, [19] leverages the more accessible conditional vector field u t ( x t | x 0 ) and its corresponding probability path p t ( x t | x 0 ) , resulting in the following Conditional Flow Matching (CFM) objective, which is equivalent to L FM in terms of gradients and accessible for sampling:

Different vector fields lead to different probability paths. For the commonly-used Gaussian distribution defined as

the corresponding vector field [19] is calculated as

where µ ′ t , σ ′ t are derivatives of µ t , σ t w.r.t. t . We consider three lines of Gaussian paths in this paper, namely the Variance Exploding (VE), Variance Preserving (VP) and Optimal Transport (OT) paths, which are listed in Table 1.

Based on these paths, we are capable of designing proper flow models to maintain symmetries for specific structure prediction tasks. In this paper, we mainly focus on the two typical tasks on atomic systems: antibody structure prediction and crystal structure prediction. Note that symmtries are crucial in 3D atomic systems, and we provide more discussions in Appendix B.

Example 1: Antibody Structure Prediction. Antibodies are Y-shaped proteins generated by the immune system to identify and bind to specific antigens, with the structure depicted in Figure 2. Researchers mainly center on the variable domains of antibodies, which comprise a heavy chain and a light chain. Each chain includes three Complementarity-Determining Regions (CDRs) and four framework regions in an alternating sequence. The six CDRs are volatile and crucial in defining the binding specificity and affinity, while the framework regions remain conserved. Among them, CDR-H3, which is the third CDR on the heavy chain, is the most diverse region and the primary focus of antibody design. Therefore, it is a fundamental yet challenging problem to accurately predict the structure of the CDRs upon binding.

FR1

H1

FR2

H2

FR3

H3

FR4

Heavy Chain

FR1

L1

FR2

L2

FR3

L3

FR4

Light Chain

Antigen

Task Definition: Let A = { a 1 , a 2 , · · · , a N } denote the sequence of the targeted CDR region with the length of N , where a i ∈ { 0 , 1 } 20 is the one-hot type of the amino acid, and ⃗ X = { ⃗ x 1 , ⃗ x 2 , · · · , ⃗ x N } is the corresponding 3D structures with x i ∈ R 3 × 4 as the backbone coordinates including N,C α , C, and O . Similarly, the sequence and structure of the context ( i.e. framework regions and the antigen) are defined as A C , ⃗ X C . The goal is to predict the structure of the CDR region given the context:

Probability Paths and Training Objectives: DiffAb [20] has designed the VP path for the coordinates of the CDR region as

where α t is scheduled as α t = e -1 2 ∫ t 0 β ( s ) ds . After sampling ⃗ X 0 = ⃗ ϵ ∼ N (0 , I ) , we have ⃗ X t = α t ⃗ X 0 + √ 1 -α 2 t ⃗ ϵ . With proper reparameterization, the training objective is defined as

which only requires a model θ to predict the denoising term given the current state.

Moreover, it is also practicable to linearly connect the data point ⃗ X 0 and the noisy prior ⃗ ϵ via the OT path as ⃗ X t = (1 -t ) ⃗ X 0 + t⃗ ϵ . The vector field is then defined as

The training objective directly align the model with the simple vector field:

Example 2: Crystal Structure Prediction. Crystal Structure Prediction (CSP), a fundamental aspect of material science, requires to predict the stable 3D structure of a compound solely from its composition. Unlike molecules or proteins, which have a finite number of atoms, the uniqueness of crystals lies in their periodic repetition in infinite 3D space. The infinite crystal structure is typically simplified by its repeating unit, which is called a unit cell . The key point of CSP is the representation and generation of the unit cell.

Task Definition: A unit cell is usually characterized by a triplet M = ( A , L , F ) , where A = [ a 1 , a 2 , ..., a N ] ∈ R h × N represents the one-hot encoded atom types, L = [ l 1 , l 2 , l 3 ] ∈ R 3 × 3 denotes the lattice matrix with three basis vectors describing the crystal's periodicity, and F = [ x 1 , x 2 , ..., x N ] ∈ R 3 × N [0 , 1) contains the fractional coordinates of the atoms, specifying their positions relative to the lattice matrix. The goal of CSP is to predict the lattice matrix and the atomic coordinates based on the given crystal composition as

Probability Paths and Training Objectives: As the lattice matrix L also lies in the Euclidean space, we can design similar VP and OT paths as Eq. (7-10). Given ϵ L ∼ N (0 , I ) , with L t = α t L 0 + √ 1 -α 2 t ϵ L , the loss function of the VP path is defined as

Besides, with L t = (1 -t ) L 0 + t ϵ L , the training objective of the OT path is similarly defined as

𝒇

!

𝒇

"

𝒇

#

The fractional coordinates lie in the torus space of R 3 × N [0 , 1) to inherently reflect the periodicity of the crystal. Previous works [15, 14] project the VE path to this manifold, and the Gaussian distribution is changed into the Wrapped Normal (WN) distribution as p t ( F t | F 0 ) = N w ( F t ; F 0 , σ 2 t I ) ,

where N w ( x ; · , · ) = ∑ ∞ i = -∞ N w ( x + i ; · , · ) . An accessible way to learn this path is to match the score, i.e. the negative logarithmic gradient, of p t , and the loss function is defined as

where λ t = E -1 [ ∥∇ log N w (0 , σ 2 t ) ∥ 2 2 ] is the pre-computed weight. If σ 1 in Eq. (14) is sufficiently large, p 1 would finally approach the uniform distribution, which can be selected as the prior distribution. Apart from the VE path, it is also applicable to directly connect the data point and the prior sample via the shortest path on the manifold. Specfically, given F 0 ∼ p 0 , F 1 ∼ p 1 , where p 1 is defined as the uniform distribution, the shortest path s ( F 0 , F 1 ) can be determined by the logarithmic map from F 0 to F 1 as s ( F 0 , F 1 ) = log F 0 F 1 = w ( F 1 -F 0 +0 . 5) -0 . 5 . Alternatively, F 1 can also be considered as the destination of s ( F 0 , F 1 ) via the exponential map from F 0 as exp F 0 s ( F 0 , F 1 ) = w ( F 0 + s ( F 0 , F 1 )) . To eliminate the effect of the overall translation introduced by the prior, we further normalize F 1 as ˆ F 1 = exp F 0 ˆ s ( F 0 , F 1 ) = exp F 0 ( s ( F 0 , F 1 ) -¯ s ( F 0 , F 1 ) ) ,

where ¯ s averages the paths of all atoms. With the path of F t defined as F t = exp F 0 ( t ˆ s ( F 0 , F 1 ) ) , the training objective for the OT path is

Generalized Notations. Overall, the structure prediction tasks aims at generating the targeted structure x given some condition c , i.e. to learn p 0 ( x | c ) . And the flow matching objective minimizes the Mean Square Error (MSE) of the predicted and pre-defined vector fields with proper reparameterization or simplification, which can be generalized as

Hereinafter, we use these generalized notations for simplicity.

3.2 Preference Dataset Construction

Building on the flow paths introduced in § 3.1, we now delve into the details of constructing a preference dataset, which is pivotal for the application of DPO, as detailed in § 3.3.

Candidate Generation As shown in Algorithm 1, the construction of the preference dataset begins with the generation of multiple candidate structures for each sample in our reference dataset, D ref . Leveraging the pre-trained flow-based generative model θ ref, we generate M candidate structures { ˆ x ij } j M =1 for each sample ( x i , c i ) via p ( x | c i ; θ ref ) , ensuring that each generated structure is contextually relevant and adheres to the geometric constraints discussed previously.

As each candidate is generated, we compute the distance between ˆ x ij and the original structure x i using a predefined metric d ( · , · ) . If this distance is less than or equal to a threshold δ , the candidate is considered a close match and is added to D pos , a subset of promising candidates. This step is crucial for efficiently filtering the generated data to retain only the most relevant candidates for DPO.

Preference Pairs Construction Subsequently, we construct K preference pairs ( x w ik , x l ik ) for each sample i by Algorithm 2, where x w ik is preferred over x l ik . This preference is determined based on their proximity to the original structure x i . Apart from sampling pairs from generated structures, we also use a ratio r to select the ground truth as the preferred sample. Moreover, if all generated structures for a sample are far from the original, the original structure x i is always preferred. The other pairs are formed by selecting x w ik from the promising subset D pos and x l ik from the broader set D gen . This process ensures that the pairs reflect a clear preference based on the closeness to the original structure, facilitating effective training through DPO, which is explored in the next section.

3.3 Direct Preference Optimization

To align large language models with human preference, DPO [25] is proposed to replace the RLHF [23] training objective with directly maximizing the likelihood of the preference. [30] extends DPO to text-to-image generation task, adapting the DPO target to diffusion models. Given the preference pair ( x w , x l ) , DPO [25] designs the training objective as

where p opt , p ref are probabilities yielded by the fine-tuned model θ opt and the pre-trained flow model, and β is a hyperparameter to control the KL divergence of these two distributions.

It is nontrivial to efficiently acquire p ( x ) via iterative generative models. Inspired by [30], we uniformly discretize the time interval into T steps, where step i is located at t = i/T . By formulating the probability from the path x 0: T , Eq. (18) can be rewritten as

To avoid costly sampling through the entire path, Jensen's inequality [30] is applied to bound Eq. (19) as

where B = βT servers as a hyperparameter. As directly sampling x i -1 , x i from an arbitrary intermediate step i is still unfeasible, we can estimate them via the accessible Gaussian paths p in Table 1 as

where J ( x w i ; p, p θ ) denotes D KL ( p ( x w i -1 | x w i, 0 ) ∥ p θ ( x w i -1 | x w i ) ) and the same for J ( x l i ; p, p θ ) . As p and p θ are Gaussian distributions with the same noise scheduler, the KL divergence can be formulated as

According to DDIM [27], if a time-dependent Gaussian path follows the form x i ∼ N ( x i ; k i x 0 , σ i I ) , we can further design p ( x i -1 | x i, 0 ) = N ( x ; µ ( x i -1 | x i, 0 ) , σ 2 i -1 | i ) . Given σ 2 i -1 | i , the mean can be formulated as

Fortunately, all paths defined in Table 1 follows this form. And µ θ ( x i -1 | x i ) can be parameterized similarly as Eq. (24), with estimating x 0 via predicted vector field or denoising terms. Hence, we can approximate J ( x i ; p, p θ ) by MSE i ( x 0 , x 1 ; θ ) . With sufficiently large T , Eq. (22) can be changed into an applicable form as follows, which is our final training objective.

4 Experiments

We validate our method on two distinct domains: antibody structure prediction (§ 4.1) and crystal structure prediction (§ 4.2).

4.1 Antibody Structure Prediction

Dataset Following previous literature [20], we extract antibody structures from the SAbDab database [8] for training and utilize the manually curated test set from DiffAb [20], which contains 19 antibody-antigen complexes. We first derive all structures deposited before April 11th, 2024, and remove those with resolution above 4.0Å or non-protein targets, resulting in 12,428 antibodies. Subsequently, we use mmseqs2 [29] to cluster the antibodies based on 50% sequence identity for each CDR, and exclude those in the same clusters as the test set antibodies. The dataset is then split into training and validation sets at a 9:1 ratio based on the clusters.

Metrics We employ the following metrics for evaluation. RMSD C α measures the Root Mean Square Deviation of the generated alpha carbon coordinates with respect to the reference. RMSD bb is the RMSD calculated on the four backbone atoms including C, C α , N, O . To better profile the generated distribution, for each antibody, we generate 20 structures and use two strategies to aggregate the results across different antibodies. Strategy worst select the worst generated structure per antibody according to RMSD and then averagse across different antibodies, while strategy mean averages the RMSD of 20 candidates first, and then across antibodies. Strategy worst measures the furthest deviation of the generated distribution compared to the reference, while strategy mean is commonly adopted in previous works [20, 16]. Results aggregated with worst are denoted as C α -w and bb-w , while those with mean are denoted as C α and bb .

Results We evaluate VP path (DiffAb) [20] and OT path [19] with the proposed FlowDPO on CDR structure prediction. Results in Table 2 illustrate that either using VP path or OT path, further training with DPO consistently enhances performance across different CDRs. Notably, on the most challenging part ( i.e. CDR-H3), the DPO phase yields the most significant improvement. Metrics aggregated with strategy worst demonstrate noticeable gains, indicating effective supperssion of low-quality samples by the DPO phase, which we attribute to the objective of DPO in distinguishing the prefered samples. Such characteristics are favorable in practical applications where it requires blind selection of generated structures without prior knowledge of which structures might be more correct. We also depict the distributions of RMSD and examples of generated CDR-H3 structures in Figure 4. It shows that the blue curves, yielded by the original flow models, often exhibit a bimodal distribution. While the first peak at a lower RMSD indicating higher quality generations, the second peak at a higher RMSD suggests the models experience hallucination , confidently generating conformations that significantly deviate from the ground truth. DPO effectively suppresses this erroneous second peak, leading to an overall improvement in the quality of generated samples. On closer inspection, this correction also addresses physical invalidities, such as the twisted backbone seen in Figure 4.

Reference (PDB: 7CHE)

VP path

VP path + DPO

RMSD

C

⍺

= 2.09Å

RMSD

bb

= 2.33Å

RMSD

C

⍺

= 1.69Å

RMSD

bb

= 1.67Å

Reference (PDB: 7D6I)

RMSD

C

⍺

= 1.30Å

RMSD

bb

= 1.83Å

OT path

OT path + DPO

RMSD

C

⍺

= 1.14Å

RMSD

bb

= 1.37Å

0.4

Model

VP

VP + DPO

0.3

Density

0.2

0.1

0.0

1

2

3

4

RMSD(CA)

0.7

Model

0.6

10

OT + DPO

0.5

Density

0.4

E'0

0.2

0.1

0.0

0.5

1.0

1.5

2.0

2.5

3.0

RMSD(CA)

4.2 Crystal Structure Prediction

Dataset We conduct the crystal structure prediction task on three datasets in line with previous works [33, 14]. Perov-5 [4] includes 18,928 perovskite crystals, each characterized by similar structures but varying compositions, and exactly 5 atoms per unit cell. MP-20 [13] comprises 45,231 materials from the Materials Project, featuring a wide range of compositions and structures, with each material containing no more than 20 atoms per unit cell. These materials predominantly represent crystals that have been synthesized experimentally. MPTS-52 is an advanced version of MP-20, containing 40,476 structures with unit cells that include up to 52 atoms, presenting a more complex challenge. For Perov-5 and MP-20, we maintain the conventional 60-20-20 split for training, validation, and testing. For the MPTS-52 dataset, we use a chronological split, assigning 27,380 crystals for training, 5,000 for validation, and 8,096 for testing.

Metrics For inference, we generate one structure given each composition. The predicted sample is then matched with the ground truth via the StructureMatcher class in pymatgen [22] with thresholds stol=0.5, angle_tol=10, ltol=0.3 as applied in previous works [33, 14]. We use Match Rate (MR) as the proportion of matched structures among the testing set, and the RMSD is averaged over the matched pairs, and normalized by 3 √ V/N where V is the volume of the unit cell.

Results We compare the results with two generative baselines P-cG-SchNet [9] and CDVAE [33]. The results are shown in Table 3, where we explore three combinations of paths for jointly generating the lattice and the fractional coordinates: VP+VE, OT+OT, and OT+VE. Notably, the VP+VE path is previously developed by DiffCSP [14]. We find that the OT path is more effective for lattice generation, while the VE path provides more accurate predictions of atomic coordinates within the cell. Overall, the OT+VE combination generally delivers the best performance. Furthermore, DPO consistently enhances the performance of the model trained on each combination, demonstrating its capability to refine the predictions to a more precise alignment with experimental structures. We additionally visualize the RMSD distribution of predicted structures from different Gaussian paths. Results in Figure 5 reveal a similar pattern to Figure 4, demonstrating that DPO reduces the probability of low-quality generations.

5 Conclusion

In this work, we propose FlowDPO, a novel framework for 3D structure prediction that integrates flowbased generative models with Direct Preference Optimization. We achieve 3D structure prediction via flow matching models with various probability paths, and generalize the DPO training objective to arbitrary Gaussian paths. To refine the model via DPO, we generate multiple candidate structures and construct the preference dataset by aligning with ground truth. The results demonstrate substantial

improvements in prediction accuracy for both antibody and crystal structures, highlighting the effectiveness and versatility of FlowDPO in the field of 3D structure prediction.

Acknowledgments

This work is jointly supported by the National Science and Technology Major Project under Grant 2020AAA0107300, the National Natural Science Foundation of China (No. 61925601, No. 62376276, No. 62236011), and Beijing Nova Program (20230484278).

References

J. Abramson, J. Adler, J. Dunger, R. Evans, T. Green, A. Pritzel, O. Ronneberger, L. Willmore, A. J. Ballard, J. Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature , pages 1-3, 2024.

I. Batatia, D. P. Kovacs, G. Simm, C. Ortner, and G. Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural

Information Processing Systems , 35:11423-11436, 2022.

S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications , 13(1):2453, 2022.

I. E. Castelli, D. D. Landis, K. S. Thygesen, S. Dahl, I. Chorkendorff, T. F. Jaramillo, and K. W. Jacobsen. New cubic perovskites for one-and two-photon water splitting using the computational materials repository. Energy & Environmental Science , 5(10):9034-9043, 2012.

R. T. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

G. Corso, B. Jing, R. Barzilay, T. Jaakkola, et al. Diffdock: Diffusion steps, twists, and turns for molecular docking. In International Conference on Learning Representations (ICLR 2023) , 2023.

G. R. Desiraju. Cryptic crystallography. Nature materials , 1(2):77-79, 2002.

J. Dunbar, K. Krawczyk, J. Leem, T. Baker, A. Fuchs, G. Georges, J. Shi, and C. M. Deane. Sabdab: the structural antibody database. Nucleic acids research , 42(D1):D1140-D1146, 2014.

N. W. Gebauer, M. Gastegger, S. S. Hessmann, K.-R. Müller, and K. T. Schütt. Inverse design of 3d molecular structures with conditional generative neural networks. Nature communications , 13(1):1-11, 2022.

J. Hafner. Ab-initio simulations of materials using vasp: Density-functional theory and beyond. Journal of computational chemistry , 29(13):2044-2078, 2008.

T. A. Halgren, R. B. Murphy, R. A. Friesner, H. S. Beard, L. L. Frye, W. T. Pollard, and J. L. Banks. Glide: a new approach for rapid, accurate docking and scoring. 2. enrichment factors in database screening. Journal of medicinal chemistry , 47(7):1750-1759, 2004.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

A. Jain, S. P. Ong, G. Hautier, W. Chen, W. D. Richards, S. Dacek, S. Cholia, D. Gunter, D. Skinner, G. Ceder, et al. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL materials , 1(1):011002, 2013.

R. Jiao, W. Huang, P. Lin, J. Han, P. Chen, Y. Lu, and Y. Liu. Crystal structure prediction by joint equivariant diffusion. Advances in Neural Information Processing Systems , 36, 2024.

B. Jing, G. Corso, J. Chang, R. Barzilay, and T. Jaakkola. Torsional diffusion for molecular conformer generation. Advances in Neural Information Processing Systems , 35:24240-24253, 2022.

X. Kong, W. Huang, and Y. Liu. Conditional antibody design as 3d equivariant graph translation. In The Eleventh International Conference on Learning Representations , 2022.

J. Kreitz, M. J. Friedrich, A. Guru, B. Lash, M. Saito, R. K. Macrae, and F. Zhang. Programmable protein delivery with a bacterial contractile injection system. Nature , 616(7956):357-364, 2023.

Y.-L. Liao and T. Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations , 2022.

Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2022.

S. Luo, Y. Su, X. Peng, S. Wang, J. Peng, and J. Ma. Antigen-specific antibody design and optimization with diffusion-based generative models for protein structures. Advances in Neural Information Processing Systems , 35:9754-9767, 2022.

G. M. Morris, R. Huey, W. Lindstrom, M. F. Sanner, R. K. Belew, D. S. Goodsell, and A. J. Olson. Autodock4 and autodocktools4: Automated docking with selective receptor flexibility. Journal of computational chemistry , 30(16):2785-2791, 2009.

S. P. Ong, W. D. Richards, A. Jain, G. Hautier, M. Kocher, S. Cholia, D. Gunter, V. L. Chevrier, K. A. Persson, and G. Ceder. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science , 68:314-319, 2013.

L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems , 35:27730-27744, 2022.

C. J. Pickard and R. Needs. Ab initio random structure searching. Journal of Physics: Condensed Matter , 23(5):053201, 2011.

R. Rafailov, A. Sharma, E. Mitchell, C. D. Manning, S. Ermon, and C. Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems , 36, 2024.

M. Schreiner, O. Winther, and S. Olsson. Implicit transfer operator learning: Multiple timeresolution surrogates for molecular dynamics. arXiv preprint arXiv:2305.18046 , 2023.

J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations , 2020.

Y. Song, J. Gong, M. Xu, Z. Cao, Y. Lan, S. Ermon, H. Zhou, and W.-Y. Ma. Equivariant flow matching with hybrid probability transport for 3d molecule generation. Advances in Neural Information Processing Systems , 36, 2024.

M. Steinegger and J. Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology , 35(11):1026-1028, 2017.

B. Wallace, M. Dang, R. Rafailov, L. Zhou, A. Lou, S. Purushwalkam, S. Ermon, C. Xiong, S. Joty, and N. Naik. Diffusion model alignment using direct preference optimization. arXiv preprint arXiv:2311.12908 , 2023.

Y. Wang, J. Lv, L. Zhu, and Y. Ma. Crystal structure prediction via particle-swarm optimization. Physical Review B , 82(9):094116, 2010.

B. D. Weitzner, J. R. Jeliazkov, S. Lyskov, N. Marze, D. Kuroda, R. Frick, J. Adolf-Bryfogle, N. Biswas, R. L. Dunbrack, and J. J. Gray. Modeling and docking of antibody structures with rosetta. Nature protocols , 12(2):401-416, 2017.

T. Xie, X. Fu, O.-E. Ganea, R. Barzilay, and T. S. Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. In International Conference on Learning Representations , 2021.

M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations , 2021.

C. Zeni, R. Pinsler, D. Zügner, A. Fowler, M. Horton, X. Fu, S. Shysheya, J. Crabbé, L. Sun, J. Smith, et al. Mattergen: a generative model for inorganic materials design. arXiv preprint arXiv:2312.03687 , 2023.

X. Zhou, D. Xue, R. Chen, Z. Zheng, L. Wang, and Q. Gu. Antigen-specific antibody design via direct energy-based preference optimization. arXiv preprint arXiv:2403.16576 , 2024.

A From RLHF to DPO

Given the preference pair ( x w , x l ) with condition c , the Bradley-Terry (BT) model considers a latent reward model r ( x | c ) behind them and formulates the preference as

RLHF [23] optimizes the generative model by explicitly training a reward model r ϕ , and maximizing the reward with a KL regularization term to control the model by the original reference p ref as

We omit the condition c for simplicity. As Eq. (27) exists a close-form solution p ∗ θ ( x ) = p ref ( x ) e r ∗ ( x ) /β /Z , where Z is the normalization term, we can rewrite the optimal reward model as

After introducing Eq. (28) into Eq. (26) and directly maximizing the log likelihood, DPO [25] simplifies the training objective as

B Required Symmetries of Atomic Systems

The design of flow paths is constrained by the symmetry requirements of specific atomic systems, which are detailed as follows.

Antibody Structure Prediction The designed vector field should maintain equivariance to any rotation Q ∈ SO (3) and be invariant to any translation ⃗ t ∈ R 3 :

Crystal Structure Prediction Previous works [14, 35] consider the task defined in Eq. (11) as a joint generation task on L and F . For the generative process, the vector field of the lattice should be equivariant to an arbitrary rotation Q ∈ SO (3) , and that of the coordinates is required to ensure the periodic translation invariance for any translation vector t . Specfically, we have

where the operation w ( F ) = F -⌊ F ⌋ ∈ [0 , 1) 3 × N returns the coordinates back to the unit cell.

After maintaining the required symmetries, the proposed flow model is capable of generating equivalent structures under different E(3) transformations. An example of the OT-OT path for the crystal is shown in Figure 6.

C Implementation Details

C.1 Antibody Structure Prediction

We use the framework of DiffAb [20] to train the flow models. The original denoising network in DiffAb requires orientation matrices as input, yet the OT path of the SO (3) matrices is not naive to derive, which is out of our main scope. Therefore, we replace the denoising network with the multi-channel EGNN proposed in MEAN [16] to avoid this problem. All experiments can be run on one GeForce RTX 3090 GPU. Detailed hyperparameters for our FlowDPO are presented in Table 4.

Ground Truth

Generated Samples

C.2 Crystal Structure Prediction

We use the denoising network designed in DiffCSP [14] as the backbone model to train the flow models. To predict the vector field via the denoising output, for the OT path designed on lattice, we apply the reparameterization as

And for the OT path on the fractional coordinates, we directly use v F ,θ ( L t , F t , A ) = ϵ F ,θ ( L t , F t , A ) . We select RMSD defined by StructureMatcher class in pymatgen [22] with thresholds stol=0.5, angle_tol=10, ltol=0.3 as the distance metric to construct the preference dataset. Specially, the RMSD of the unmatched structure is set as + ∞ , and such candidates will never be selected as the preferred sample.

The detailed hyperparameters for the FlowDPO pipeline on each crystal dataset are provided in Table 5. Each experiment is run on one GeForce RTX 3090 GPU.

D Comparison with Regressive Methods

To further investigate the advantages of the generative paradigm, we employ the same backbone model (MEAN) for a direct regression task as an additional baseline. The results are presented in Table 6. Our findings indicate that generative models outperform the regression model in 4 of the 6 CDRs, particularly in the highly variable and functionally critical regions, CDR-H3 and CDR-L3.

Additionally, we report both the mean and minimum RMSD values across 20 generations for each generative model. The significantly lower minimum RMSD values demonstrate that generative models not only yield predictions closer to the observed reference structures but also possess the capability to generate multiple feasible conformations around the stable state.

E Limitations

As detailed in § 3.3, the derivation of the rationality of DPO for flow models primarily focuses on Gaussian paths. However, flow models have the potential to learn mappings from an arbitrary prior to the data distribution if the probability path is appropriately defined. Therefore, a more general derivation that does not rely on Gaussian assumptions could be explored in future research. Additionally, our current evaluation is based predominantly on computational metrics. Conducting wet-lab experiments would provide a more robust validation of the model's effectiveness in practical applications.

F Broader Impacts

The introduction of FlowDPO marks a pivotal advancement in scientific domains such as drug development, materials research, and molecular informatics. Recent developments, such as AlphaFold3, have demonstrated remarkable accuracy in predicting structures across various domains [1]. However, issues such as hallucinations, like erroneous structural order in inherently disordered regions, remain a challenge. It is intriguing to explore whether alignment strategies based on DPO can mitigate these hallucinations and enhance overall prediction accuracy.

G Codes

Our codes are available at https://github.com/jiaor17/FlowDPO .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: In abstract and introduction, we summarize our contribution as enabling DPO for flow-based structure prediction models, constructing preference dataset to align model predictions with reference, and achieving promising results for antibody and crystal prediction tasks. These claims are detailed and verified by the method (§3) and experiment (§4) sections.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations of this paper are discussed in Appendix E.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The derivation are provided in §3.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The used datasets and evaluation setups are provided in §4, and we provide more details in Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Our codes are provided in Appendix G.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The hyperparameters are provided in Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The dataset used in the experiments are large and the results are relatively stable. Rerunning the pipeline for multiple times is costly. Instead, we further compare the prediction results for different models from the perspective of distribution in Figure 4.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The compute resources are provided in Appendix C.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This paper definitely follows the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impacts of this paper are discussed in Appendix ?? .

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All datasets used in this paper have been properly cited.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to apply the Direct Preference Optimization (DPO) procedure for the 3D Structure Prediction of atomic structure predictions. The approach is applied to crystal and antibody structures, considering different kinds of Gaussian Path for flow matching models, showing a general improvement in the state of the art.\n\nSTRENGTHS:\n- The method seems to consistently outperform previous methods on the two different tasks\n- The implementation of the code approach can foster the application of DPO protocol to other problems in the chemical and molecular domains. \n- Experimentations with different kinds of Gaussian paths seem to be an interesting analysis that could help future works to build on top of them.\n\nI may have missed some important aspects since I am not an expert in this domain.\n\nWEAKNESSES:\n- As a non-expert, I have a hard time understanding some of the merits. In particular, the approach seems mainly to be a straightforward application of DPO to the specific case of atomic systems, and I cannot find a technical contribution specifically designed for the specific case. Given that I do not know how relevant the problems considered are and what their impact could be, I am unsure about the significance of the work. In case I am missing something, I am looking forward to authors' clarification.\n- I feel the general presentation could be improved. Section 3.3 does not offer particular insight and does not blend with the previous applicative part. The paper does not offer many visualizations of the results or quantitative comparisons. These could help the reader to better understand the impact of DPO in this domain. For example, recent works suggest that DPO may be prone to overfitting [A]. I would like to know what the impact of the procedure is in the atomic structure prediction and if they observed similar problems. Similarly, although the authors offer some possible future directions, they don't really discuss failure cases of the method, and the presented limitations seem a bit generic.\n\n[A]: A general theoretical paradigm to understand learning from human preferences, Azar et al., AISTATS 2024\n\nQUESTIONS:\nAdding to the points above:\n1) How efficient is the proposed approach w.r.t. the competitors [9, 14, 33]?\n2) From my intuition, symmetries in the atomic structures may be a relevant aspect to consider in constructing the reference dataset since it could lead to ambiguities but also provide a source of augmentation. I would like to know what is the authors' thoughts on this aspect.\n\nLIMITATIONS:\nThe paper phrases the limitations mainly as future works but does not discuss the failure modes of the proposed approach. I suggest the authors comment on recurrent confusion cases for the observed generations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a framework called FlowDPO subsuming different Gaussian probability paths for flow matching models predicting the 3D structure of atomic systems. Moreover, the authors propose a method for the automatic creation of a preference pair dataset used for finetuning a pre-trained flow matching model with direct preference optimization (DPO). As examples of the structure prediction task for atomic systems, the paper explores antibody and crystal structure prediction, each with specific symmetry requirements. The choice of probability path is evaluated empirically including a validation of DPO in terms of consistent performance boosts.\n\nSTRENGTHS:\n* The paper is very well structured:\n  - The presented formal framework abstracts the choice of different probability paths for flow matching models.\n  - The general task of structure prediction of atomic systems is well motivated and explained by the two instances in form of antibody and crystal structure prediction.\n  - The structure of the method section follows the training procedure: Flow Matching training, preference dataset construction, direct preference optimization.\n* The paper is original: The authors apply general flow matching and direct preference optimization to the new domain of atomic system structure prediction.\n* The contributions of DPO for flow matching and automatic preference dataset construction are general and therefore also applicable to other domains.\n* The experimental results validate the proposed framework:\n  - Exploration of different probability paths can be beneficial, as there is not always one performing the best.\n  - Direct preference optimization on the proposed automatically constructed datasets consistently improves performance.\n\nWEAKNESSES:\n* Missing motivation for a generative approach to 3D structure prediction of atomic systems.\n* Unclear/imprecise use of the term \"hallucination\". The paper claims the model can distinguish between high-fidelity and hallucinated samples (cf. 52 f) without defining what hallucinated samples really are. Recombination / generalization is usually desired for generative models.\n* Self-contradictory method:\n  - The framework is about flow matching models, i.e., a generative method.\n  - The preference definition is solely based on distances to the ground truth training examples. This way of creating a preference dataset does not account for the possibility of multiple possible solutions given the conditioning. Optimization w.r.t. the RMSD metric is not even significantly different from the flow matching objective, begging the question for why this improves performance at all?\n  - Originally, DPO is used in combination with crowd-sourced human preferences for human alignment. This type of metric is significantly different from the pixel-wise Gaussian noise and MSE regression during training of text-to-image models, for example.\n* Contradictory method and evaluation:\n  - The paper uses generative flow matching models.\n  - The evaluation is done in terms of distances to a single ground truth, which does not account for the possibility of multiple valid solutions. However, this might be a general problem of generative approaches for regression tasks.\n* The comparison with baselines is limited to generative methods.\n\nQUESTIONS:\n* Is there uncertainty in two given tasks, i.e., can there be multiple different valid solutions to predicting the 3D structure given the respective prior knowledge as conditioning?\n  - If no, why would you use a generative approach?\n  - If yes:\n    - Why do you construct preferences based on distances to a single ground truth without accounting for the possibility of multiple solutions given the conditioning?\n    - Why does the evaluation consider distance to a single ground truth?\n* What do you understand as hallucinated samples?\n* Are there non-generative (regression) baselines and if so, how is the performance compared to them?\n\nLIMITATIONS:\nThe authors have addressed the limitations and broader impacts without any potential negative societal impacts in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces FlowDPO, a framework that predicts 3D structures of atomic systems using diffusion flow matching models. To suppress hallucinations and improve sample quality, Direct Preference Optimization (DPO) is adopted to finetune a pretrained model using a preference dataset consisting of winning and losing pairs, which are selected based on their distance to the ground truth. Details on the flow matching model and DPO training paradigm has been discussed. Experiments have been conducted on antibody and crystal structure prediction tasks, demonstrating better accuracy than the two baselines (P-cG-SchNet and CDVAE).\n\nSTRENGTHS:\n- This work introduces the application of diffusion models with DPO into the field of atomic structure predictions. Unlike the original diffusion DPO paper [30], where samples have to be ranked by humans, this task benefits from known ground truths from a training dataset, allowing the winning and losing pairs to be automatically constructed by comparing the distances to the ground truth.\n- Experimental results demonstrate the advantage of the proposed method over other generative models, and ablation studies have revealed the contribution of the DPO fine-tuning.\n\nWEAKNESSES:\n- The motivation for symmetry preservation and transform invariance of the flow trajectory is unclear and seems not quite relevant to the main idea of using DPO to improve sample accuracy. The theoretical justification for the trajectory needing to be transform invariant is not well explained, apart from the empirical results.\n- In L57, the authors claim that they are \"the first to theoretically prove the compatibility of DPO with arbitrary Gaussian paths by deriving a unified objective.\" However, I can hardly find anything relevant to this claim, and the paper seems more focused on application rather than theory.\n- In general, the contribution is limited, considering that diffusion DPO is a known method in [30], which is theoretically universally applicable to all diffusion flow matching models. While this work touches on the new field of atomic system prediction, using a universally applicable method in this field does not seem very novel to me.\n- The writing of this paper has some issues, making the presentation a bit unclear:\n  - The overall structure is not focused. While the main idea of this work is the application of diffusion DPO rather than proposing a new theory, the paper spends a lot of space on preliminaries with very specific examples. The main body, especially the experiments section, is short and lacking in detail.\n  -  The subscript 'i' is a bit confusing as it sometimes denotes the index of a data sample and sometimes the index of a timestep.\n\nQUESTIONS:\nAs mentioned in the weaknesses, I do not understand the motivation for symmetry preservation and transform invariance, and the claim about the theoretical proof.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}