Science, Faculty of
Computer Science, Department of
DSpace
UBCV
Shenoy, Nikhil
Master of Science - MSc
University of British Columbia
Efficiently sampling from the landscape of molecular conformations is an important task in computational drug discovery. Simulation approaches like Molecular Dynamics (MD) require the use of an energy function that is fast, accurate, transferable and scalable. Traditional energy function methods like Force Fields are fast, but inaccurate while Quantum Mechanical (QM) Methods are accurate but slow. Recently, Machine Learning (ML) potentials trained on datasets labeled with QM methods have become popular as they address the accuracy and speed trade-off. However, generating QM datasets is a cost-intensive exercise, and design choices like conformational diversity (coverage of the conformational landscape) and structural diversity (coverage of the chemical space) during the generation process can introduce biases into the dataset. In the first part of the thesis, we explore the intricate relationship between dataset biases, specifically conformational and structural diversity, and ML potential generalization. We investigate these dynamics through two distinct experiments: a fixed budget one, where the dataset size remains constant, and a fixed molecular set one, which focuses on fixed structural diversity while varying conformational diversity. Our results reveal the critical need for balanced structural and conformational diversity in QM datasets for optimal generalization, which current datasets lack. We believe that these findings can inform future data generation and development of ML potentials that generalize beyond training data. An alternative approach is directly sampling conformations given the molecular graph using deep generative models like Diffusion. Existing competitive approaches either use expensive local structure methods or rely on general-purpose large architecture models without task-specific inductive biases. In the second part of the thesis, we challenge the status quo and develop a simple and scalable deep generative method, Equivariant Transformer Flow (ET-Flow) incorporating flow-matching, an equivariant transformer, a harmonic prior, and an approximate Optimal Transport alignment. We achieve state-of-the-art performance on several molecular conformer generation benchmarks with significantly fewer parameters and inference steps than existing methods, highlighting the importance of inductive biases and well-informed modelling choices.
Investigating ML Potentials and DeepGenerative Models for EfficientConformational SamplingbyNikhil ShenoyB.Tech., Indian Institute of Technology, New Delhi, 2020A THESIS SUBMITTED IN PARTIAL FULFILLMENT OFTHE REQUIREMENTS FOR THE DEGREE OFMASTER OF SCIENCEinThe Faculty of Graduate and Postdoctoral Studies(Computer Science)THE UNIVERSITY OF BRITISH COLUMBIA(Vancouver)July 2024© Nikhil Shenoy 2024The following individuals certify that they have read, and recommend tothe Faculty of Graduate and Postdoctoral Studies for acceptance, the thesisentitled:Investigating ML Potentials and Deep Generative Models for Effi-cient Conformational Samplingsubmitted by Nikhil Shenoy in partial fulfillment of the requirements forthe degree of Master of Science in Computer Science.Examining Committee:Jiarui Ding, Assistant Professor, Computer Science, University of BritishColumbiaSupervisorDominique Beaini, Adjunct Professor, Computer Science, University of Mon-trealSupervisory Committee MemberKhanh Dao Duc, Assistant Professor, Mathematics, University of BritishColumbiaSupervisory Committee MemberiiAbstractEfficiently sampling from the landscape of molecular conformations is animportant task in computational drug discovery. Simulation approaches likeMolecular Dynamics (MD) require the use of an energy function that is fast,accurate, transferable and scalable. Traditional energy function methodslike Force Fields are fast, but inaccurate while Quantum Mechanical (QM)Methods are accurate but slow. Recently, Machine Learning (ML) potentialstrained on datasets labeled with QM methods have become popular as theyaddress the accuracy and speed trade-off. However, generating QM datasetsis a cost-intensive exercise, and design choices like conformational diversity(coverage of the conformational landscape) and structural diversity (coverageof the chemical space) during the generation process can introduce biasesinto the dataset.In the first part of the thesis, we explore the intricate relationship betweendataset biases, specifically conformational and structural diversity, and MLpotential generalization. We investigate these dynamics through two distinctexperiments: a fixed budget one, where the dataset size remains constant,and a fixed molecular set one, which focuses on fixed structural diversitywhile varying conformational diversity. Our results reveal the critical need forbalanced structural and conformational diversity in QM datasets for optimalgeneralization, which current datasets lack. We believe that these findingscan inform future data generation and development of ML potentials thatgeneralize beyond training data.An alternative approach is directly sampling conformations given the molecu-lar graph using deep generative models like Diffusion. Existing competitiveapproaches either use expensive local structure methods or rely on general-purpose large architecture models without task-specific inductive biases. Inthe second part of the thesis, we challenge the status quo and develop asimple and scalable deep generative method, Equivariant Transformer FlowiiiAbstract(ET-Flow) incorporating flow-matching, an equivariant transformer, a har-monic prior, and an approximate Optimal Transport alignment. We achievestate-of-the-art performance on several molecular conformer generation bench-marks with significantly fewer parameters and inference steps than existingmethods, highlighting the importance of inductive biases and well-informedmodelling choices.ivLay SummaryUnderstanding different molecular configurations (called conformations) viacomputational techniques is crucial for improving drug discovery efforts.A common approach, Molecular Dynamics (MD) simulates molecular move-ment requiring the use of a fast and accurate energy function. Traditionalenergy functions face an accuracy-speed trade-off, so Machine Learning (ML)potentials trained on datasets generated with accurate methods are a popularsolution. Since generating datasets is costly and existing ones are biased, weinvestigate the impact of biases on ML potential performance, highlightingthe need for balanced datasets and providing guidance for future development.A modern approach is directly generating conformations to study molecularproperties. Since existing competitive methods rely on expensive techniques,we present a simple approach, ET-Flow, that achieves competitive results onvarious benchmarks. We hope our approach highlights the benefits of usingkey modelling choices like Euclidean inductive biases, flow matching, and ameaningful prior.vPrefaceThe chapter presented in chapter 4 resulted in the paper,Role of Structural and Conformational Diversity for Machine LearningPotentials. Preprint, 2023. [Arxiv]Nikhil Shenoy, Prudencio Tossou, and Dominique Beaini were involved insetting the problem, designing the experiments, evaluating the results andwriting the paper. Jiarui Ding, Emmanuel Noutahi, Hadrien Mary andCristian Gabellini were involved in improving the draft during the paper-writing process. This paper was accepted at the following workshops, (1)Molecular Machine Learning Conference at MIT, 2023 in Boston (2) AI forScience, NeurIPS 2023 in New Orleans and (3) AI for Drug Discovery andDevelopment, NeurIPS 2023 in New Orleans. This paper was co-presentedwith Prudencio Tossou, Emmanuel Noutahi, Hadrien Mary, Jiarui Ding andDominique Beaini.The second part of the thesis presented in chapter 5 was done in collaborationwith members at Valence Labs (Recursion), MILA and MIT. The conceptualfoundation and direction of the project were provided by Dominique Beaini.The development of the codebase, training setup, and evaluation scriptswas carried out by Nikhil Shenoy. In method development, the use of theTorchMD-Net architecture and modifications, and the chirality correction,were conceptualized and implemented by Nikhil Shenoy. The experimentswere executed equally by Nikhil Shenoy, Majdi Hassan, and Jungyoon Lee.Dominique Beaini, Hannes Stark, and Stephan Thaler provided invaluablefeedback on the methodologies and models at each phase of the research.The paper-writing process primarily involved Nikhil Shenoy, Majdi Hassan,and Jungyoon Lee, and was reviewed and improved upon by Jiarui Ding,Dominique Beaini, Hannes Stark, and Stephan Thaler.viTable of ContentsAbstract . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . iiiLay Summary . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . vPreface . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . viTable of Contents . . . . . . . . . . . . . . . . . . . . . . . . . . . . viiList of Tables . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . xList of Figures . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . xiList of Abbreviations . . . . . . . . . . . . . . . . . . . . . . . . . xvAcknowledgements . . . . . . . . . . . . . . . . . . . . . . . . . . . xvi1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11.1 Structure . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 52 Background . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 72.1 Potential Energy Surface (PES) . . . . . . . . . . . . . . . . 72.2 Molecular Dynamics (MD) . . . . . . . . . . . . . . . . . . . 82.2.1 Quantum Mechanical (QM) Methods . . . . . . . . . 92.2.2 Semi-Empirical (SE) Methods . . . . . . . . . . . . . 102.2.3 Force Fields . . . . . . . . . . . . . . . . . . . . . . . 102.2.4 Machine Learning Potentials . . . . . . . . . . . . . . 122.3 Geometric Graph Neural Networks . . . . . . . . . . . . . . . 132.3.1 Graph Neural Networks (GNNs) . . . . . . . . . . . . 132.3.2 Permutation Symmetry . . . . . . . . . . . . . . . . . 142.3.3 Euclidean Symmetries . . . . . . . . . . . . . . . . . . 152.3.4 Equivariance and Invariance . . . . . . . . . . . . . . 152.3.5 How do inductive biases help? . . . . . . . . . . . . . 16viiTable of Contents2.4 Deep Generative Models . . . . . . . . . . . . . . . . . . . . 172.4.1 Diffusion Models . . . . . . . . . . . . . . . . . . . . . 172.4.2 Flow-Matching . . . . . . . . . . . . . . . . . . . . . . 212.4.3 Potential Benefits of Flow Matching over Diffusion . . 233 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . 253.1 QM Datasets . . . . . . . . . . . . . . . . . . . . . . . . . . . 253.2 Geometric-GNNs . . . . . . . . . . . . . . . . . . . . . . . . . 263.3 Molecular Conformer Generation . . . . . . . . . . . . . . . . 264 Role of Data Biases in ML Potential Generalization . . . . 284.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . 284.2 Method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 294.2.1 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . 294.2.2 Generalization metrics . . . . . . . . . . . . . . . . . 314.3 Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 314.3.1 Experimental details . . . . . . . . . . . . . . . . . . . 314.3.2 Structural generalization . . . . . . . . . . . . . . . . 324.3.3 Conformational generalization . . . . . . . . . . . . . 344.4 Discussion and Future Directions . . . . . . . . . . . . . . . . 375 Equivariant Flow Matching for Conformer Generation . . 395.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . 395.2 Method . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 415.2.1 Flow Matching . . . . . . . . . . . . . . . . . . . . . . 415.2.2 Alignment . . . . . . . . . . . . . . . . . . . . . . . . 425.2.3 Stochastic Sampling . . . . . . . . . . . . . . . . . . . 435.2.4 Chirality Correction . . . . . . . . . . . . . . . . . . . 435.2.5 Architecture . . . . . . . . . . . . . . . . . . . . . . . 445.3 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . 475.3.1 Experimental Setup . . . . . . . . . . . . . . . . . . . 485.3.2 Ensemble RMSD . . . . . . . . . . . . . . . . . . . . . 485.3.3 Coverage Threshold Plots . . . . . . . . . . . . . . . . 495.3.4 Ensemble Properties . . . . . . . . . . . . . . . . . . . 505.3.5 Results on GEOM-XL . . . . . . . . . . . . . . . . . . 505.3.6 Inference Steps Ablation . . . . . . . . . . . . . . . . 515.4 Discussion and Future Directions . . . . . . . . . . . . . . . . 526 Conclusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 54viiiTable of ContentsBibliography . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 56AppendicesA Appendix . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 69A.1 Structural differences between GEOM-Drugs and GEOM-QM9Distribution . . . . . . . . . . . . . . . . . . . . . . . . . . . 69A.2 ET-Flow Implementation Details . . . . . . . . . . . . . . . . 69A.2.1 Input Featurization . . . . . . . . . . . . . . . . . . . 69A.2.2 Evaluation Metrics . . . . . . . . . . . . . . . . . . . 70A.2.3 Training Details and Hyperparameters . . . . . . . . . 71A.3 Training and Sampling Algorithm . . . . . . . . . . . . . . . 71A.4 Visualizations . . . . . . . . . . . . . . . . . . . . . . . . . . 72ixList of Tables3.1 List of available QM-Datasets and their data generation char-acteristics . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 265.1 Molecule conformer generation results on GEOM-DRUGS (δ =0.75Ã ). SS is ET-Flow with stochastic sampling. For bothET-Flow and ET-Flow-SS, we sample conformations over 50time steps. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 485.2 Molecule conformer generation results on GEOM-QM9 (δ =0.5Ã ). For both ET-Flow, we sample conformations over 50time steps. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 495.3 Median averaged errors of ensemble properties between sam-pled and generated conformers (E, âÎµ, Emin in kcal/mol, andµ in debye). . . . . . . . . . . . . . . . . . . . . . . . . . . . . 505.4 Generalization results on GEOM-XL. . . . . . . . . . . . . . . 515.5 Ablation over number of inference steps on GEOM-DRUGS(δ = 0.75Ã ). Performance of ET-Flow at 5 steps is competentacross all metrics while also retaining state-of-the-art perfor-mance on precision metrics when compared with previousmethods. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 52A.1 Atomic features included in ET-Flow. . . . . . . . . . . . . . 69A.2 Hyperparameters for ET-Flow . . . . . . . . . . . . . . . . . . 71xList of Figures1.1 Trade-off between structural Diversity (number of uniquemolecules) vs conformational Diversity (average number ofconformations per molecule) in existing QM datasets. Eachpoint in the plot represents a QM dataset. The top left circledenotes an area of high structural diversity but lower confor-mational diversity. The bottom right denotes an area of lowerstructural diversity but higher conformational diversity. . . . 21.2 Overview of our first study: We vary dataset biases (confor-mational diversity, structural diversity and dataset size) tosub-sample a QM dataset, train an ML potential and thenstudy how these biases affect the performance of ML potentialson structural and conformational generalization. . . . . . . . . 31.3 Generating low-energy conformations given information abouta molecular graph (atoms and bonds). . . . . . . . . . . . . . 42.1 Trade-off between speed and accuracy for QM methods, SEmethods and Force Fields. . . . . . . . . . . . . . . . . . . . . 112.2 A visual description of the inputs and outputs for training anML potential. A single data point is usually a conformationwith N atoms (atomic positions (RNÃ3) and atomic numbers(RNÃ1)) and the target (energy scalar (R) and force vector(RNÃ3)) values. Forces (a vector per atom) are predicteddirectly from the ML potential or derived via the negativegradient w.r.t to initial positions. . . . . . . . . . . . . . . . . 122.3 Symmetries for 3D Atomistic Systems from Duval et al. [1]. Apermutation operation of the ordering of atoms changes theordering of the atom features and the 3D positions, but theenergy of the molecule remains the same. A rotation operationonly changes the 3D positions, whereas the energy and theatom features remain the same. . . . . . . . . . . . . . . . . . 14xiList of Figures2.4 Diffusion models gradually add noise to data samples in Tdiffusion steps (going from image x0 to noise xT ). The noiseaddition process is modelled as conditional probability q(xt |xtâ1) which is a Gaussian centered aroundâ1â Î²txtâ1 (β isthe variance schedule). The neural network learns the denoisingprocess modelled by pÎ¸(xtâ1 | xt). Figure used from DDPM [2]. 172.5 Training and Sampling Algorithm from DDPM [2]. . . . . . . 212.6 Diffusion and OT + Flow Matching Trajectories (Lipman et al.[3]). The figure represents a simple case where points lie ona 2D plane, and the data distribution is a 2D Gaussian (bluecolour signifying the density). The lines represent the pathbetween data samples and a fixed noise sample. One can seethat Flow-Matching provides straighter and shorter paths asopposed to Diffusion, which needs backtracking in some caseswhere it overshoots. . . . . . . . . . . . . . . . . . . . . . . . 244.1 Experimental setup: Left. (1) ns-Fixed: Keeping the numberof molecules ns fixed at 12.5k, 25k, 50k and 100k, we increasethe conformer per molecules nc. Right. (2) N -fixed: Keepingthe total number of conformers N fixed at 50k, 200k, 800kand 3.2M, we increase the conformer per molecule nc whiledecreasing the number of molecules ns. . . . . . . . . . . . . . 304.2 N -fixed: Performance on IID-S and OOD-S as we increase theconformational diversity (nc) and reduce structural diversity(ns) while keeping the number of conformers (N) fixed . . . . 334.3 ns-fixed: Performance on IID-S and OOD-S as we increase theconformational diversity (nc) while keeping structural diversity(ns) fixed . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 334.4 nc-fixed: Performance on IID-S and OOD-S as we increasethe number of conformers (N) while keeping conformationaldiversity (nc) fixed . . . . . . . . . . . . . . . . . . . . . . . . 344.5 N -fixed: Distribution of MAE performance for conformersfrom both IID-C and OOD-C based on their RMSD from thetraining distribution. The four plots represent fixed N values(50k, 200k, 800k, and 3.2M) with varying nc and ns. . . . . . 354.6 ns-fixed: Distribution of MAE performance for conformersfrom both IID-C and OOD-C based on their RMSD from thetraining distribution. The four plots represent fixed ns values(12k, 25k, 50k, and 100k) with varying nc and N . . . . . . . . 36xiiList of Figures4.7 nc-fixed: Distribution of MAE performance for conformersfrom both IID-C and OOD-C based on their RMSD from thetraining distribution. The four plots represent fixed nc values(1, 2, 4, and 8) with varying ns and N . . . . . . . . . . . . . . 375.1 Overview of ET-Flow. The model takes as input the featuresrelated to the molecular structure, like atom and bond features,as well as a radius graph of 10Ã from the atom positions. Sam-ples are drawn from the harmonic prior and then rotationallyaligned with the samples from the data. A probability path isconstructed between pairs of x0 and x1, and xt is then sampledfrom this path at a random time t. . . . . . . . . . . . . . . . 405.2 Rotation alignment using the Kabsh algorithm [4] resultingin shorter and straighter paths between the base and targetstructures. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 425.3 Two molecules that are isomorphic (one-to-one correspondencebetween atoms and bonds) can have different orientations in3D determined by the order of the tetrahedral chiral carbonâsneighbours. . . . . . . . . . . . . . . . . . . . . . . . . . . . . 435.4 Overall Architecture of ET-Flow consisting of 3 components.ET-Flow takes as input the molecular graph G, time-step t,relative positions Râ and predicts the vector field vâÎ¸, (1) Embed-ding Layer for encoding the atomic positions, molecular graphinformation and time-step into invariant features xi, equivari-ant features vâi (for atom i). (2) Update layers (with attentionlayers) updates the invariant features xi and equivariant fea-tures vâi at each layer and (3) Output Layer for predictingthe vector field vâÎ¸ given invariant features xi and equivariantfeatures vâi. . . . . . . . . . . . . . . . . . . . . . . . . . . . . 445.5 Recall and Precision Coverage result on GEOM-DRUGS asa function of the threshold distance. ET-Flow outperformsTorsionDiff by a large margin, especially in a lower thresholdregion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 49A.1 Structural differences between GEOM-Drugs (IID-S set) andGEOM-QM9 (OOD-S set) are evident from the distinct separa-tion between the two clusters. Each point in the plot representsa molecule. . . . . . . . . . . . . . . . . . . . . . . . . . . . . 70A.2 Examples of conformers of ground truth and ET-Flow forGEOM-DRUGS. . . . . . . . . . . . . . . . . . . . . . . . . . 74xiiiList of FiguresA.3 Examples of conformers of ground truth and ET-Flow fordifferent number of sampling steps. . . . . . . . . . . . . . . . 75xivList of Abbreviations3D 3-DimensionalCFM Conditional Flow MatchingCNN Convolutional Neural NetworkCNF Continuous Normalizing FlowsFF Force FieldsFM Flow MatchingGNN Graph Neural NetworksMD Molecular DynamicsML Machine LearningMLIP Machine Learning Interatomic PotentialsNN Neural NetworkNNP Neural Network PotentialsODE Ordinary Differential EquationOT Optimal TransportPES Potential Energy SurfaceQM Quantum MechanicsSDE Stochastic Differential EquationSE Semi-EmpiricalxvAcknowledgementsThroughout my graduate studies, I have accumulated a deep sense of gratitudefor the exceptional support of many individuals.First and foremost, I would like to express my sincerest thanks to my supervi-sors, Dr. Jiarui Ding and Dr. Dominique Beaini, and my internship manager,Prudencio Tossou. Their guidance and mentorship have been instrumentalin shaping my research journey.I am immensely grateful to the collaborative research environment fostered bymy colleagues, Majdi Hassan, Jungyoon Lee, Cristian Gabellini, and StephanThaler. Their willingness to share their expertise, engage in discussions,and provide constructive feedback has significantly enriched my researchexperience and my ability to deliver impactful results.Furthermore, I would like to extend my appreciation to Emmanuel Noutahi,Hadrien Mary, and Hannes Stark for their invaluable feedback on my writing.Their insights and suggestions have undoubtedly helped me strengthen mycommunication skills.A heartfelt thank you goes to my friends, Mehar Bhatia, Vedant Bahel, MajdiHassan, Jungyoon Lee, Cristian Gabellini, Jeffrey Niu, Sarah Chen, MinukMa, Wilson Tu, Kaiyun Guo, Praveen Gupta and Andrew Cheng. Theirunwavering presence and encouragement have provided a constant source ofsupport throughout my academic journey.Lastly, my deepest gratitude goes to my parents and sister for their unwaveringlove and unwavering support. Their belief in my abilities has been a constantsource of strength and inspiration, and I am incredibly fortunate to havethem in my life.xviChapter 1IntroductionRichard Feynman famously said [5],all things are made of atoms, and that everything that living thingsdo can be understood in terms of the jigglings and wigglings ofatomsAll biological processes essential for maintaining life, such as protein function,cell signalling, and metabolism, can be better understood by investigatingmolecules and their interactions at the atomic level. Atoms, which constitutemolecules, are particles consisting of electrons and protons that interact withother atoms through covalent and non-covalent interactions. These atoms arein constant motion due to thermal energy, causing them to vibrate, rotate,and translate, which causes the molecule to exist in an array of differentconfigurations or conformations. Specifically, a conformation is defined asthe three-dimensional structural arrangement of atoms in a molecule.Each conformation possesses a unique potential energy depending on thespecific arrangement of atoms. This variation in energy across all possibleconformations forms an energy distribution, which is also referred to asthe potential energy surface (PES). Accurately estimating the energy andsampling conformations from this distribution is a critical challenge due toseveral factors. First, the PES is a complex landscape consisting of low-energyvalleys and high-energy peaks, which makes it difficult to sample evenly fromall possible areas since simulation-based methods (covered later in Section 2.2)often get trapped in the valleys. Second, the complexity of the PES increaseswith the number of atoms in a molecule. The high dimensionality makesit computationally expensive to determine the energy of a conformationaccurately and to explore the entire landscape fully. Finally, interestingphenomena like chemical reactions occur only when molecules are in rareconformations, making it difficult to sample such events.Improving the computational techniques for efficient sampling of the PES1Chapter 1. IntroductionFigure 1.1: Trade-off between structural Diversity (number of uniquemolecules) vs conformational Diversity (average number of conformations permolecule) in existing QM datasets. Each point in the plot represents a QMdataset. The top left circle denotes an area of high structural diversity butlower conformational diversity. The bottom right denotes an area of lowerstructural diversity but higher conformational diversity.holds immense potential for accelerating drug discovery efforts. For instance,sampling diverse and relevant protein conformations can help improve iden-tify candidate drugs during virtual screening1. Other applications in drugdiscovery include understanding the function of proteins via protein foldingdynamics [6], identifying hidden drug-binding sites [7, 8], and uncoveringtransition pathways between meta-stable states [9].A popular method for exploring the potential energy surface is a simulation-based method called Molecular Dynamics (MD). To run an MD simulationfor a specific molecule, one needs to use an energy function that can ac-curately determine the potential energy given a conformation. Traditionalmethods for computing energy involve an accuracy-speed trade-off: ForceFields (FF) are fast but inaccurate, while Quantum Mechanical (QM) meth-ods, which approximate the Schrödinger Equation2, are accurate but slow.Recently, Machine Learning (ML) potentials have been proposed to addressthe accuracy-speed trade-off by training neural networks on QM datasetscontaining conformations and energy labels computed with accurate QM1https://en.wikipedia.org/wiki/Virtual_screening2https://en.wikipedia.org/wiki/Schrodinger_equation2Chapter 1. IntroductionFigure 1.2: Overview of our first study: We vary dataset biases (conforma-tional diversity, structural diversity and dataset size) to sub-sample a QMdataset, train an ML potential and then study how these biases affect theperformance of ML potentials on structural and conformational generaliza-tion.methods. Therefore, a well-trained ML potential should be able to predictenergies with the accuracy of QM methods while being fast, eliminating theneed for expensive QM calculations.Generating large and high-quality QM datasets is a computationally expen-sive exercise done primarily on CPUs, involving expensive matrix calculations.For example, PubChemQC [10] generated electronic properties for 86 millionmolecules, but the entire process took 1.5 years. Several QM dataset gener-ation efforts have recently been carried out to train reliable ML potentials.Several design choices need to be specified before generating QM datasets,such as the number of molecules to generate, the number of conformationsto generate per molecule, the QM method to use, and so on. These choicescan induce specific biases in the data, such as conformational diversity (thenumber of unique conformations) and structural diversity (the number ofunique molecules) (Figure 1.1). The performance of ML potentials trained onQM datasets can be impacted by these biases when tested on conformationsfrom novel parts of the potential energy surface (PES), novel molecules, anddifferent covalent 3 or non-covalent4 environments.3Bonds between atoms arising from sharing an electron pair.4Weaker interactions than covalent bonds, includes interactions like attraction be-tween oppositely charged ions, hydrogen bonds, Van der Waals interaction, Hydrophobic3Chapter 1. IntroductionFigure 1.3: Generating low-energy conformations given information about amolecular graph (atoms and bonds).In the first part of the thesis (Chapter 4), we aim to study the effect of thesebiases on ML potentials in a principled manner. Specifically, we aim to:Inform QM data generation efforts by understanding the intricate rela-tionship between dataset size, structural diversity, and conformationaldiversity on the generalization of ML Potentials.Studying biological, chemical and physical properties of molecules involvesanalyzing low-energy conformations, as molecules typically exist in these statesexperimentally. Generating low-energy conformations given the moleculargraph (atoms and the bonds involved) is commonly referred to as molecularconformer generation5 (Figure 1.3). Simulation-based methods like MD can becomputationally inefficient, requiring longer sampling times to produce diverselow-energy conformations. A popular alternative is using deep generativemodels like Diffusion [2, 11]. Diffusion models, which have been widelysuccessful in generating high-quality images, videos, and audio, are trainedby repeatedly adding noise into data (or diffusion noise into data) and thenlearning to reverse the noise-addition process (or reverse diffusion). Duringinference, one can sample from a noise distribution and repeatedly de-noiseusing the diffusion model to generate a data sample.Recently, diffusion models have been applied to molecular conformation gen-eration under different formulations. One of the top-performing approaches,interactions, etc.5Note that the term "molecular conformer generation" here specifically refers to thesampling of low-energy conformations, although the name may suggest a more generalsampling of all possible conformations.41.1. StructureTorsional Diffusion [12], decomposes the problem by constructing the con-formationâs local structure (bond lengths and bond angles) using a softwarepackage, RDKit [13], whereas the remaining degrees of freedom from torsionangles are modelled using a diffusion model. This approach is mainly bottle-necked by the performance and computational efficiency of RDKit since anyerrors in the bond lengths and bond angles propagate into the final confor-mation. The current state-of-the-art, Molecular Conformer Fields (MCF)[14], utilizes the framework from Zhuang et al. [15], where a transformer [16]is used to diffuse over functions that map molecular graphs to conformers.Although conformations are structural arrangements of atoms in Euclideanspace, which transform according to Euclidean symmetries (translations,rotations, and reflections), MCF does not incorporate any of these inductivebiases in its architecture, necessitating the use of a large number of parame-ters. The above two approaches achieve competitive results either by usingan expensive local structure method, RDKit, or by inefficiently utilizing alarge transformer architecture to compensate for the lack of inductive biases.In the second part of the thesis (Chapter 5), we improve upon the existingapproaches for molecular conformer generation. We utilize flow matching[3, 17, 18], a novel scalable approach to train a class of deep generative modelscalled continuous normalizing flows (CNFs) with a simple least squares ob-jective. Additionally, we incorporate symmetries into our neural network byutilizing an Equivariant Transformer from TorchMD-NET [19]. We showcasethat it is possible to design state-of-the-art molecular conformer generationmodels using molecule-specific inductive biases and efficient generative mod-elling, alleviating the need for complexities introduced in previous approaches.To summarize, we aim to:Develop a simple, scalable and efficient generative method calledEquivariant Transformer-Flow (ET-Flow) utilizing equivariant trans-formers and flow matching to achieve state-of-the-art results on variousmolecular conformer generation benchmarks with significantly fewerparameters and inference steps than previous methods.1.1 StructureThe thesis is structured as follows: In Chapter 2, I go through the background,which covers the fundamental concepts used in this thesis. In Chapter 3, Icover the various works that have been done in the domain of ML potentials51.1. Structurefor molecular dynamics and deep generative models for conformer generation.In Chapter 4, I cover the first part of my thesis, where we uncover how databiases affect ML potential generalization. In Chapter 5, I cover the secondaspect of my thesis, introducing a new conformer generation model. Finally,in Chapter 6, I conclude the thesis with a discussion on future directions.6Chapter 2BackgroundIn this chapter, we review some of the important concepts used in thethesis. First, we understand and discuss the relationship between molecularconformations, energy and the Boltzmann distribution in Section 2.1. Wecover the basics of Molecular Dynamics (MD) and the different ways tocompute energy for MD in Section 2.2. We look at a class of neural networks,specifically Geometric Graph Neural Networks in Section 2.3 as they arecommonly used for modelling conformations. We review a class of deepgenerative models in Section 2.4 to understand how to directly sample low-energy conformations, specifically diffusion models and flow matching.2.1 Potential Energy Surface (PES)Molecules are not static entities, they are constantly moving due to inter-actions among atoms, such as chemical bond stretching, electric dipoles,electrostatic interactions, and Van-der-Waals interactions. A single snapshotor a specific 3D configuration of a molecule is referred to as a conformation.Every conformation is associated with a scalar energy value depending onthe atoms present in the molecule and how the atoms are positioned in theEuclidean space (Born-Oppenheimer approximation [20]). The distributionof energy over all possible conformations forms a potential energy surface.The probability distribution over conformations is given by the Boltzmanndistribution where the probability of a molecule existing in a particularconformation xi â X (X being the set of all possible conformations) is afunction of the energy Exi and the temperature T as,pi â exp(âExikT) or pi =1Zexp(âExikT), (2.1)where pi is the probability of the i-th conformation, Exi is the energy of thei-th conformation, k is the Boltzmann constant6, T is the temperature of6The value of the Boltzmann constant is set to 1.3806491023 Joules per Kelvin72.2. Molecular Dynamics (MD)the system and Z is the normalization constant computed over all possibleconformations given by,Z =âxjâXexp(âExjkT). (2.2)Importance of Equilibrium Conformations: Conformations that exist inthe low-energy parts of the potential energy surface are more stable thereforehave a higher probability of existing (Exi is lower, therefore higher pi) com-pared to other high-energy conformations and are referred to as equilibriumconformations. Improved ability to sample equilibrium conformations allowsus to study important biological, chemical and physical properties [21â23]like charge distribution, docking poses [24] and shape similarity [25].2.2 Molecular Dynamics (MD)Molecular Dynamics (MD) is a computational method that simulates themovement of molecules over time, effectively creating a movie of their be-haviour. Starting with an initial conformation, the energy of the currentstate is calculated at each time step, and the molecule moves in the directionof lower energy, following Newtonian equations that govern the motion ofobjects. This process reveals how each atom moves in response to the influ-ence of other atoms, generating a dynamic picture of molecular behaviourand offering valuable insights into molecular interactions.We refer to a point cloud of atoms as a system7 where we have informationof the atomic numbers and their positions in 3D Euclidean space. Each stepof the simulation usually comprises of,1. Start with an initial set of positions and velocities for each atom in thesystem2. Compute the energy of the conformation E. Forces on each atom i aredefined as the gradient w.r.t to its 3D positions ri as,Fâi = ââEârâi, (2.3)Here, the negative sign is added because forces point in the direction ofdecreasing the systemâs energy.7We use this term interchangeably with atomistic system or 3D atomistic system82.2. Molecular Dynamics (MD)3. Update the positions and velocity of each atom based on an integrationalgorithm like Verlet Integration 8.Calculating the energy E: An important aspect of molecular dynamics isthe choice of the energy function used for simulation. To ensure a reliablesimulation, the main requirements of the energy function are,1. Speed: We need to be able to run simulations over large timescales.Longer timescales allow us to sample and understand relevant phenom-ena in large atomistic systems.2. Accuracy: Ability to accurately estimate experimental or theoreticalenergy values. An accurate method should be able to model bothcovalent and non-covalent interactions.3. Scalability: Ability to scale to a large number of atoms, like in proteins.4. Transferability: Ability to model systems, geometries and conditionsbeyond its initial training or parameterization.In the next few sections, we go over different ways to calculate the energy ofa system.2.2.1 Quantum Mechanical (QM) MethodsThe most accurate way to compute the energy of a system is to solve theSchrödinger equation [26] for a given system. However, the analytical solutionis only known for 2-body systems like Hydrogen. For moderately-sized systems,there exist tractable approximate methods such as (1) Density FunctionalTheory (DFT) [27], (2) Hartree-Fock (HF) methods [28] and (3) CoupledCluster (CC) methods [29â31].Density functional theory (DFT) methods usually involve the use of func-tionals, which are functions of another function. In the case of DFT, thefunctionals are a function of the electron density, which itself is a measure ofthe probability of an electron residing around a point. Functionals vary interms of accuracy, cost, and applicability (what system they can be appliedto), and there are different classes of functionals. Examples of functionalsinclude the Generalized Gradient Approximation (GGA), Meta, and Hybridfunctionals, etc. Additionally, Basis sets9 (groups of functions for approxi-mating molecular orbitals) are used for efficient calculation of DFT equations8https://en.wikipedia.org/wiki/Verlet_integration9https://en.wikipedia.org/wiki/Basis_set_(chemistry)92.2. Molecular Dynamics (MD)on computers. Examples of basis sets include Gaussian Type Orbitals andSlater-Type Orbitals. To summarize, the functional and the basis set de-termine the DFT methodâs accuracy and computational efficiency. Variousfunctionals and basis sets have been developed in the domain of DFT, forinstance,B3LYPï¸¸ ︷︷ ︸Functional/ 6-31Gï¸¸ ︷︷ ︸Basis Set, (2.4)where, the DFT method consists of the functional B3LYP from Becke [32]and the basis set 6-31G from Ditchfield et al. [33]. The specific mathematicaldetails of the DFT functional and basis set are out of the scope of this thesis.Limitation: QM methods are usually computationally expensive and notfeasible for large systems like proteins. For instance, CCSD(T) (CoupledCluster with Single, Double and Perturbative Triplets) is the gold standardfor quantum mechanical calculations but scales O(N7) where N is the numberof atoms in the system. This implies that using existing QM methods forlarge systems like proteins is impractical.2.2.2 Semi-Empirical (SE) MethodsSemi-empirical methods further approximate the above-discussed Hartree-Fock or DFT methods, providing more speed at the cost of accuracy. Forinstance, GFN2-xTB [34], a semi-empirical method, is designed for fastcalculation of energies for systems with up to 1000 atoms and is on average 100to 1000 times more computationally efficient than DFT or QM methods [34â36]. The speed of these methods allows for sampling conformational ensemblesof larger molecular systems [23, 37]. However, these systems are less accuratethan QM methods as they do not explicitly model specific interactions likedispersion, electronic polarizability, van der Waals interactions, and hydrogenbonding. Examples of SE methods include Density Functional Tight BindingMethods (GFN-xTB, GFN2-xTB), Parametric Methods (PM6, PM7), etc.2.2.3 Force FieldsForce fields are the simplest and fastest way to calculate the energy of thesystem. Classical force fields are usually parameterized in a simple form over102.2. Molecular Dynamics (MD)Figure 2.1: Trade-off between speed and accuracy for QM methods, SEmethods and Force Fields.all bonded and non-bonded interactions, such as,E = Ebonded + Enon-bonded, (2.5)Ebonded =âbondsEbondi +âanglesEanglei +âdihedralsEdihedrali ,Enon-bonded =âpairsEelectrostatici +âpairsEvanderwaalsi .Force fields are the industry standard for biomolecular simulations as theyare fast and can model systems with reasonable accuracy. For large systemslike proteins, the force field is usually parameterized carefully to reproducethe experimental data. Popular force fields are the industry standard, suchas AMBER [38], CHARMM [39] and GROMACS [40].Trade-offs for Conventional Methods: The Figure 2.1 is a summary ofhow the accuracy and speed of the above-discussed methods compare againsteach other. On the one hand, QM methods are slow and accurate; force fieldsare fast but inaccurate. Recently, machine learning potentials (discussednext) have been proposed as a solution to address this trade-off.112.2. Molecular Dynamics (MD)Figure 2.2: A visual description of the inputs and outputs for training an MLpotential. A single data point is usually a conformation with N atoms (atomicpositions (RNÃ3) and atomic numbers (RNÃ1)) and the target (energy scalar(R) and force vector (RNÃ3)) values. Forces (a vector per atom) are predicteddirectly from the ML potential or derived via the negative gradient w.r.t toinitial positions.2.2.4 Machine Learning PotentialsMachine Learning Potentials10 aim to bridge the gap between the accuracy ofab initio QM methods and the efficiency of classical force fields. ML Potentialsare trained to learn the mapping from conformations to energy and forceswhere the energy and force labels are computed using expensive and accurateQM methods. Various large QM datasets have been recently generated fortraining ML potentials, which will be covered later in Section 3.1.To be more precise, the input for an ML potential is the atomic numbers andpositions in Euclidean space (3D coordinates per atom), whereas the targetlabel is energy (a scalar value per conformation) and forces (a 3D vector peratom). Figure 2.2 is a visual description of the inputs and outputs involved.Computing Forces: Two methods are usually considered for computingforces with ML potentials. One way is to predict the forces from the neuralnetwork directly. This way is usually computationally efficient w.r.t memoryand runtime [1, 41]. Another way is to first predict the energy using thenetwork and then calculate the forces using Equation (2.3) by differentiatingw.r.t the atomic positions. This method, while slower, allows for more stable10Machine Learning Potentials are also referred to as Neural Network Potentials (NNPs),Machine Learning Interatomic Potentials (MLIPs) and Machine Learning Force Fields(MLFFs) in literature.122.3. Geometric Graph Neural Networksdynamics [42] as it ensures energy conservation.2.3 Geometric Graph Neural NetworksGraphs are a data structure that can be used to model complex relationshipsbetween entities. A graph G usually consists of a set of nodes V, and thenodes are connected by a set of edges E . In the context of molecules, thenodes are atoms, and the edges are bonds between atoms. Neural networksthat incorporate the structure of a graph as an inductive bias are referredto as Graph Neural Networks (covered in Section 2.3.1). A special case ofGraph Neural Neural Networks called Geometric Graph Neural Networks arenetworks that incorporate the structure of geometric graphs like moleculessince molecular attributes like atomic positions, energies, and forces transformaccording to Euclidean symmetries (covered in Section 2.3.3).2.3.1 Graph Neural Networks (GNNs)Graph Neural Networks [43â46] are a class of neural networks that usuallyact on graph-structured data and involve a message passing operation [47, 48]to learn node-level or graph-level representations based on the task at hand.Consider a graph G = (V, E) where V â {1, ..., n} is the set of n vertices withinitial features h0v and E â {(u, v) | (u, v) â V 2 â© u Ì¸= v} is the set of edgeswith associated edge features euv. The key component of GNNs is message-passing operations, which act iteratively on node-level representations. Thenode-level representations are then passed through a readout function toget a graph-level representation. Consider L message passing layers, thelth message passing operation has the following 3 steps to update noderepresentations,1. Compute the message muv for each edge (u, v) with,mluv = MESSAGEl(hlu, hlv, euv) (2.6)2. Aggregate the message using all connected edges with,mlv = AGGREGATEluâNv(mluv) (2.7)3. Apply an update function to get the updated node representation with,hl+1v = UPDATEl(mlv, hlv) (2.8)132.3. Geometric Graph Neural NetworksFigure 2.3: Symmetries for 3D Atomistic Systems from Duval et al. [1]. Apermutation operation of the ordering of atoms changes the ordering of theatom features and the 3D positions, but the energy of the molecule remainsthe same. A rotation operation only changes the 3D positions, whereas theenergy and the atom features remain the same.After all the message passing operations are performed, to get a graph-levelrepresentation, a readout function is used to aggregate node-level representa-tions hLv as,z = READOUT({hLv | v â V }) (2.9)Usually, the operations MESSAGE(Â·) and UPDATE(Â·) are Multi-LayeredPerceptrons whereas AGGREGATE(Â·) and READOUT(Â·) are aggregationoperations like sum, mean and maximum.Intuition: The rationale behind using message passing operations for graph-structured data is to incorporate local structure and relationships into noderepresentations by incorporating features from local neighbourhoods suchthat two nodes close together in a graph are more similar than two distantnodes.2.3.2 Permutation SymmetryAn important property of GNNs is that they respect permutation symmetries,i.e. if the input node ordering is permuted, then the graph-level outputsremain the same (or invariant), and the node-level outputs are permuted142.3. Geometric Graph Neural Networksaccordingly. For instance, in the case of ML potentials, if we permute thenode ordering, the forces should be permuted accordingly, whereas energywill stay invariant. We visualize the effect of permutation on energy, atomfeatures and atomic positions in Figure 2.3.2.3.3 Euclidean SymmetriesWhile working with ML potentials, the input to the model are molecularconformations where each atom is treated as a point in the Euclidean spacehaving 3D coordinates. The edges are defined based on the molecular graphor based on the distance between atoms. The target for ML potentials aregeometric attributes like energy and forces. Ideally, if we apply a globalrotation to the molecule in Euclidean space, the energy should remain thesame, whereas the forces should transform according to the rotation. Tra-ditional GNNs (or NNs) do not respect these symmetries. Therefore, arecent line of work has been to develop geometric graph neural networks thatincorporate Euclidean symmetries like rotations and translations, similar tohow Convolutional Neural Networks (CNNs) encode translational symmetriesin images. We visualize the effect of rotation on energy, atom features andatomic positions in Figure 2.3.2.3.4 Equivariance and InvarianceFollowing [1, 49], we denote the action11 of a group12 G on a space X as g Â·x,where g â G and x â X. If we consider a function f : X â Y and assumethat G acts on both spaces X and Y , then,â¢ A function is said to be G-equivariant, if f(g · x) = g · f(x). Intu-itively, this means transforming the input and applying the functionis equivalent to applying the function first and then transforming theoutput.â¢ A function is said to be G-invariant if f(g · x) = f(x). Intuitively, thismeans that the output does not change if the input is transformed withelements of group G.Say we have a neural network for predicting forces (a vector) and energy (ascalar) from input positions. We desire:â¢ energy to be invariant w.r.t rotations, translations and reflections.11https://en.wikipedia.org/wiki/Group_action12https://en.wikipedia.org/wiki/Group_(mathematics)152.3. Geometric Graph Neural Networksâ¢ forces to be equivariant w.r.t to rotations, translations and reflections.â¢ atomic positions to be equivariant w.r.t to rotations, translations andreflections.This means that rotating our input positions and using the neural networkshould be equivalent to using the neural network and then rotating only theforces, keeping the energy fixed. Mathematically, if R is a rotation matrix, Xare input positions, f(Â·) is a rotation equivariant neural network that outputsboth force vectors F and energy scalars E given positions, thenf(X) = E,F (2.10)f(RÃX) = E,RÃ F (2.11)The class of neural networks incorporating these symmetries via equivarianceare called Equivariant Neural Networks. While designing Equivariant GNNsfor atomistic systems, the following groups are usually considered,1. O(3): Contains distance-preserving orthogonal rotations and reflections.2. SO(3): Referred to as the 3D rotation group. A subgroup of O(3)containing only rotations with determinant= 1 (no reflections).3. E(3): Referred to as the Euclidean Group. It contains rotations,reflections, and translations in the Euclidean space.4. SE(3): A subgroup of E(3) containing only rotations and translationsin the Euclidean Space.2.3.5 How do inductive biases help?Inductive biases of an algorithm are the initial set of assumptions that enableit to generalize to unseen data. In the context of neural networks, these biases,informed by the structure of data and prior knowledge, influence architecturalchoices. For instance, Convolutional Neural Networks (CNNs), which areusually used for images, incorporate translational equivariance through weightsharing, which intuitively means applying the same convolutional filters acrossthe entire image, recognizing that features are equally relevant regardless ofspatial location. Notably, well-engineered inductive biases in convolutionshave been shown to achieve state-of-the-art results on vision benchmarks [50],outperforming architectures without such biases, like Vision Transformers[51]. Additionally, it has also been shown that using inductive biases helps162.4. Deep Generative ModelsFigure 2.4: Diffusion models gradually add noise to data samples in Tdiffusion steps (going from image x0 to noise xT ). The noise addition processis modelled as conditional probability q(xt | xtâ1) which is a Gaussian centeredaroundâ1â Î²txtâ1 (β is the variance schedule). The neural network learnsthe denoising process modelled by pÎ¸(xtâ1 | xt). Figure used from DDPM [2].models generalize more easily on out-of-distribution tasks and are less proneto memorization [52].2.4 Deep Generative ModelsThe second part of this thesis covered in Chapter 5 focuses on utilizingflow-matching to sample low-energy conformations. In this section, we goover Diffusion [2, 11, 53], a popular deep generative model to sample highquality and diverse samples from an unknown complex distribution. We thencover Flow-Matching [3, 17, 18, 54], another deep generative method thattrains continuous normalizing flows (CNFs) and improves upon some of thelimitations of diffusion.2.4.1 Diffusion ModelsDiffusion models [2, 11] learn to approximate a stochastic differential equation(SDE) to transform samples from a simple noise density into samples from amore complex data density. It does so by learning to reverse a Markov chainthat progressively adds noise to ground truth samples. This is visualized inFigure 2.4. During inference, one can sample from noise and then denoisemultiple times to get samples from the data distribution.Preliminaries: Consider the case that our ground truth data samples x0come from some unknown data distribution q(x). At time step t = T ,our samples come from the noise distribution N (0, I). For the noising anddenoising process, we consider the following conditional distributions,â¢ Forward Diffusion Distribution (q(xt | xtâ1)): This controls how we go172.4. Deep Generative Modelsfrom sample xtâ1 to xt (less noise to more noise). This will be usedduring the training process to add noise.â¢ Reverse Diffusion Distribution (q(xtâ1 | xt)): This controls how we gofrom sample xt to xtâ1 (more noise to less noise). This conditionaldistribution is modelled by a neural network (parameters given by θ).Forward Diffusion Process: The forward diffusion process is a Markovchain where a small amount of noise is added to the sample in T diffusionsteps modelled by the forward diffusion conditional probability q(xt | xtâ1),which produces samples x1, x2, ...., xT . The forward diffusion conditionalprobability that performs the noising process is given as,q(xt | xtâ1) = N (xt;â1â Î²txtâ1, βtI) (2.12)q(x1:T |x0) =Tât=1q(xt | xtâ1) (2.13)where βt is the variance schedule controlling the amount of noise addedat each step (where β1 < β2 < ... < βT ). The forward diffusion process(Equation (2.12)) using the reparameterization trick [55] gives us,xt =â1â Î²txtâ1 +âÎ²tÏµt, ϵt â¼ N (0, I). (2.14)This noising process can be intuitively understood as, given xtâ1 we can getxt by reducing the amount of xtâ1 by multiplying withâ1â βt and addingsome noise from ϵ â¼ N (0, I). At time step T , the factor â1â βt goes to 0,leaving us only with the isotropic Gaussian distribution or the noise term.Additionally, we can extend the Equation (2.14) using αt = (1 â βt) andα¯t =âti=1 αi to get xt from x0 as follows,xt =âÎ±txtâ1 +â1â αtϵt (2.15)=âÎ±tÎ±tâ1xtâ2 +âÎ±tâ1â Î±tâ1Ïµtâ1 +â1â αtϵt (2.16)=âÎ±tÎ±tâ1xtâ2 +â1â Î±tÎ±tâ1Ïµtâ1 (2.17)= ... (2.18)=âÎ±¯tx0 +â1â α¯tϵ ϵ â¼ N (0, I) (2.19)Note: We go from Equation (2.16) to Equation (2.17) by using the followingtrick,e1 â¼ N (0, (Ï21)I) e2 â¼ N (0, (Ï22)I)e3 = e1 + e2 â¼ N (0, (Ï21 + Ï22)I)182.4. Deep Generative ModelsTo summarize, we now have a distribution to go from our data sample x0 toany time step xt using the conditional probability,q(xt | x0) = N (xt;âÎ±¯tx0, (1â α¯t)) (2.20)Reverse Diffusion Process: The reverse diffusion process modelled byprobability q(xtâ1 | xt) can be rewritten using Bayesâ rule as,q(xtâ1 | xt) = q(xt | xtâ1)q(xtâ1)q(xt)(2.21)Unfortunately, we cannot compute a closed-form solution for q(xtâ1 | xt) sincewe do not have access to marginal densities q(xtâ1) and q(xt). A trick that isused to make this reverse conditional probability tractable is by conditioningit on the ground truth data point x0 as q(xtâ1 | xt, x0),q(xtâ1 | xt, x0) = q(xt | xtâ1, x0)q(xtâ1 | x0)q(xt | x0) (2.22)=q(xt | xtâ1)q(xtâ1 | x0)q(xt | x0) (2.23)=N (xt;â1â Î²txtâ1, βtI)N (xtâ1;âÎ±¯tâ1x0, (1â Î±¯tâ1))N (xt;âÎ±¯tx0, (1â α¯t))(2.24)= ... (2.25)= N (xtâ1 | µË(xt, x0), Î²ËtI) (2.26)where if x0 = 1âÎ±¯t (xt ââ1â α¯tϵt) (from Equation (2.19)), the mean µËt canbe derived as,µËt =1âÎ±¯t(xt â 1â Î±tâ1â α¯tϵt) (2.27)Note: We have skipped a few steps (covered in detail in the DDPM paper[2]) to derive the µËt.Training Objective: The setup of a diffusion model is very similar to thatof a Variational Auto-Encoder [55] where the variational lower-bound is usedto minimize the negative log-likelihood. We can derive the variational lower192.4. Deep Generative Modelsbound LVLB in our setup as follows,LNLL = âEq(x0) log pθ(x0)= âEq(x0) log(â«pÎ¸(x0:T )dx1:T)= âEq(x0) log(â«q(x1:T | x0) pθ(x0:T )q(x1:T | x0)dx1:T)= âEq(x0) log(Eq(x1:T |x0)pθ(x0:T )q(x1:T | x0))â¤ âEq(x0:T ) logpθ(x0:T )q(x1:T | x0)â¤ Eq(x0:T ) logq(x1:T | x0)pθ(x0:T )= LVLBWith the tractable forms derived for the reverse (Equation (2.26)) and forwardprocess (Equation (2.12)), we can get a tractable form of the objective functionas follows (proof in DDPM paper [2]),L = Etâ¼[1,T ],x0,ϵt [â¥Ïµt â ϵθ(xt, t)â¥2], (2.28)= Etâ¼[1,T ],x0,ϵt [â¥Ïµt â ÏµÎ¸(âÎ±¯tx0 +â1â α¯tϵt, t)â¥2], (2.29)where x0 is the ground truth sample, ϵt is the noise at time t, xt is thedenoised sample at time t and ϵθ(xt, t) is the noise predicted by the model attime t. The loss term has an intuitive explanation at any given time step t,we go from our data sample x0 to xt given the noise ϵt, and our model learnsto predict the noise added ϵt given the noisy sample xt and time step t.Inference: For inference, we sample from a Gaussian xT â¼ N (0, I), andthen iteratively use the denoising model to update our sample until we getto sample x0 according to the following equation,xtâ1 = (1/âÎ±t)(xt â 1â Î±tâ1â αtϵθ(xt, t)) (2.30)This equation can be seen as removing noise from a noisy sample xt byestimating the amount of noise ϵθ(xt, t) added at time step t = t.The training and sampling algorithm is described in Figure 2.5.Challenges with Diffusion-Based Models: Despite being state-of-the-artacross a wide range of domains, diffusion models still suffer from specificlimitations such as202.4. Deep Generative ModelsFigure 2.5: Training and Sampling Algorithm from DDPM [2].â¢ Several assumptions are made while designing the probability paths,leading to long training and inference time. For instance, we assumethe noise, forward and reverse processes as Gaussian distributions.â¢ Long inference times. While inference speed can be improved, it requiresspecialized tricks [11, 56].â¢ At large T values, the noised samples xT are not guaranteed to reachthe noise distribution N (0, I) when using the forward diffusion process.2.4.2 Flow-MatchingContinuous normalizing flows (CNFs) [57] are a way to learn a mappingbetween any two distributions without imposing any assumptions on theprobability paths or the distributions. However, traditional CNFs have beenheld back due to difficulty in training and scaling to large datasets. Recently,Flow-Matching [3, 18, 54, 58] proposed a way to train CNFs in a scalableand efficient manner with an objective similar to that of Diffusion.Preliminaries: Consider the case where we have samples [x1, ..., xd] â Rdfrom an unknown data distribution q(x1). We define 3 objects,1. a time-dependent probability density path, pt : [0, 1]Ã Rd â R>02. a time-dependent velocity field, vt : [0, 1]Ã Rd â Rd3. a time-dependent flow, Ït : [0, 1] Ã Rd â Rd, which is defined via anOrdinary Differential Equation (ODE) as,dÏt(x)dt= vt(Ït(x)), Ï0(x) = x. (2.31)212.4. Deep Generative ModelsNote: Diffusion models usually refer to t = 0 as the data distributionand t = T as the noise distribution. For Flow-Matching, itâs theopposite as t = 0 refers to the noise distribution and at t = 1 we getthe data distribution.A CNF models the velocity field vt(x; θ) using a neural network with param-eters θ. A CNF can transform a simple distribution (say a simple Gaussian)p(x0) to a more complicated one q(x1) (say conformations) using the followingpush-forward equation,pt = [Ït]âp0, (2.32)In our case, we want the probability field to be equal to the noise distributionat time step t = 0 and to be equal to the data distribution pt(x) = q(x1) attime step t = 1.The flow Ït uses the push forward equation to modify the probability fieldfrom pt=0 to pt=t using the change of variables operation â as,pt = [Ït]âp0(x) = p0(Ïâ1t (x)) det[âÏâ1tâx(x)]. (2.33)Flow-Matching Objective: Consider the simple case where we start fromthe noise distribution p0(x) = N (0, I) and at time step t = 1, we are atthe unknown data distribution p1 = q(x1). Assuming we have access tothe probability field pt(x) and the velocity field ut(x), we can use a neuralnetwork to approximate the velocity field using vt(x; θ) and the flow-matchingobjective as,LFM = Et,pt(x)[â¥ut(x)â vt(x; Î¸)â¥]. (2.34)Here, we simply sample a time step from a uniform distribution t â¼ U(0, 1),and then using the time step, sample from the time-dependent probabilityfield x â¼ pt(x) and then try to match the neural network output vt(x; θ) tothe ground truth velocity field ut(x).However, the above objective is intractable since we do not know about ptand ut. Lipman et al. [3] simplified this objective by designing conditionalprobability paths given a particular data sample x1, a conditional vector fieldut(x | x1) and a conditional flow matching objective,LCFM = Et,q(x1),pt(x|x1)[â¥ut(x | x1)â vt(x; Î¸)â¥]. (2.35)where the conditional probability path could be defined as,pt(x | x1) = N (x | µt(x1), Ït(x1)I). (2.36)222.4. Deep Generative ModelsHere, the time-conditioned mean and variance are set to µ0(x1) = 0 andÏ0(x1) = 1 at t = 0 and µ1(x1) = x1 and Ï1(x1) = Ïmin at t = 1 to ensurethat the probability path converges to the data density at t = 1 and noisedensity at t = 0. Many vector fields can generate this probability path, butfor simplicity, Lipman et al. [3] considers a simple flow,Ït(x) = Ït(x1)x+ µt(x1) (2.37)One can see that at t = 0, the flow Ït=0(x) = Ï0(x1)x+ µ0 = x, satisfyingthe ODE base condition in Equation (2.31).In the case of a Gaussian probability path (Equation (2.36)) and the above-mentioned flow (Equation (2.37)), we can get a closed form equation of theground truth vector field ut(Â·) using the Equation (2.31) as,ut(Ït(x) | x1) = ddtÏt(x) =Ïâ²t(x1)Ït(x1)(xâ µt(x1)) + µâ²t(x1). (2.38)Note: Refer to Theorem 3 from Lipman et al. [3] for derivation.A summary of the CFM Objective: Given a time step t â¼ U(0, 1) anda data point x1 â¼ q(x1), we sample from the conditional probability fieldpt(x | x1), and match the conditional vector field ut(x | x1) with the neuralnetwork output v(x; θ).Similarity to Diffusion Objective: We can see that both the diffusionobjective (Equation (2.28)) and the conditional flow matching objective(Equation (2.28)) can be trained in a simulation free manner. In the caseof diffusion, we want to estimate the noise added at a particular time step,whereas in the case of flow matching, we want to estimate the flow.2.4.3 Potential Benefits of Flow Matching over DiffusionSome of the advantages of using flow matching over diffusion for generativemodelling are as follows:â¢ Flow Matching makes no assumptions on the type of probability pathsas long as specific boundary conditions are specified, i.e. at time stept = 0, our probability field should match the noisy samples, whereas,at t = 1, it needs to match the data samples.â¢ The framework of flow-matching guarantees that we are at the noisedistribution at t = 0 as opposed to Diffusion where samples are notguaranteed to be exactly from the noise distribution at time step t = T(see equation Equation (2.19)).232.4. Deep Generative Modelsâ¢ Lipman et al. [3] highlights in Figure 2.6 how flow-matching coupledwith Optimal Transport (OT) can lead to straighter trajectories betweennoise and data samples. Straighter and shorter paths are helpful since,during inference, samples need to travel a shorter path to reach a datasample, ultimately helping in inference time and convergence speed.Figure 2.6: Diffusion and OT + Flow Matching Trajectories (Lipman et al.[3]). The figure represents a simple case where points lie on a 2D plane, and thedata distribution is a 2D Gaussian (blue colour signifying the density). Thelines represent the path between data samples and a fixed noise sample. Onecan see that Flow-Matching provides straighter and shorter paths as opposedto Diffusion, which needs backtracking in some cases where it overshoots.24Chapter 3Related WorkIn this section, we cover the previous works in the domain of ML poten-tials, quantum mechanical datasets and molecular conformer generation. InSection 3.1, we discuss the various Quantum Mechanics (QM) datasets inthe public domain for training ML Potentials. In Section 3.2 followed byunderstanding the different geometric graph neural network architectures formodelling atomistic systems. Finally, in Section 3.3, we look at conventionalsystematic methods and diffusion-based approaches for molecular conformergeneration.3.1 QM DatasetsTraining reliable ML potentials requires high-quality datasets where energiesare calculated using methods that are more accurate than classical forcefields. Ideal datasets in this domain satisfy the following requirements [59],(1) cover a wide range of chemical elements, (2) cover both low-energy andhigh-energy conformations, (3) include forces and energy labels and (4) haveenergy calculations based off an accurate level of theory. Recent open-sourceQM datasets are summarised in Table 3.1.Density Functional Theory (DFT) and Coupled Cluster (CC) are the gold-standard methods for accurately calculating energy. ANI [60] and its exten-sions, ANI-1x and ANI-1ccx [61], offer 20M off-equilibrium conformationsfor 57K molecules with energy computed using the ÏB97x:6-31G(d) DFTmethod. Whereas, Spice [59] provides a collection of 1.1M conformers for adiverse set of 19k molecules with both energy and force labels calculated usingthe ÏB97M-D3(BJ)/def2-TZVPPD DFT method. Datasets that cover a largenumber of molecules are PubchemQC-B3LYP [10] and Molecule3D [62], whichconsists of 3.6M and 4M molecules with 1 conformer per molecule and energycalculated using B3LYP/6-31G* DFT method and PM6 SE method. Finally,OrbNet Denali [63] contributes 2.3 million equilibrium and off-equilibriumconformers for 200K molecules with ÏB97X-D3/def2-TZVP DFT method.253.2. Geometric-GNNsTable 3.1: List of available QM-Datasets and their data generation character-isticsQM DatasetTotalConformers(N)EnergyMethodsForceMethodsAtomTypesAtomic NumberMin and MaxDFTTheoryLevelOnlyLow EnergyConformationsANI-1 [61] 22,057,374 1 0 4 2/26 ÏB97x:6-31G(d) NoANI-1x [61] 4,956,005 8 2 4 2/63 ÏB97x:6-31G(d) NoANI-1ccx [61] 489,571 4 0 4 2/63 ÏB97x:6-31G(d) NoGEOM [23] 33,078,483 1 0 13 3/181 GFN2-xTB YesMolecule3D [62] 3,899,647 1 0 36 1/137 B3LYP/6-31G*, PM6 YesOrbNet D. [63] 2,338,889 2 0 17 2/74 ÏB97X-D3/def2-TZVP NoPub. PM6 [65] 189,890,155 1 0 70 1/215 PM6 YesPub. B3lyp [10] 85,915,773 1 0 70 1/215 B3LYP/6-31G* YesQmugs [64] 1,992,984 2 0 10 4/228 GFN2-xTB, ÏB97X-D/def2-SVP YesSpice [59] 1,110,165 1 1 15 2/110 ÏB97M-D3(BJ)/def2-TZVPPD NoSemi-Empirical (SE) methods are significantly faster than DFT but moreaccurate than classical force fields and offer a good trade-off between speedand accuracy. GEOM [23] utilizes SE-method GFN2-xTB to generate 37Menergy-minimized conformations for 450k molecules. Similarly, QMugs [64]provides 665k molecules with 3 conformers per molecule using SE-methodGFN2-xTB and DFT method ÏB97X-D/def2-SVP.3.2 Geometric-GNNsInductive biases play an important role in generalization and sample effi-ciency as discussed in Section 2.3.5. In the case of 3D atomistic modelling,since atomic coordinates are 3D coordinates, useful inductive biases includeE(3) symmetries like rotations, translations and inversions in 3D. Recently,various equivariant architectures [1] have been developed that act on boththe Cartesian [19, 66â69] and spherical basis [70â76]. In the case of molecu-lar conformer generation, the use of equivariant GNNs has been limited toGeoDiff [77], and Torsional Diffusion [76]. GeoDiff [77] utilizes EGNN [66], acartesian basis equivariant architecture, while Torsional Diffusion [12] usesTensor Field Networks [76].3.3 Molecular Conformer GenerationVarious ML-based approaches [12, 77â81] have been developed to improveupon the limitations of conventional methods. GraphDG [82] and CGCF [78]generate pairwise distances between atoms using variational auto-encoders(VAE) [55] and flow-based models [83] respectively, which are then subse-quently used to generate the atom coordinates using the Distance Geometry263.3. Molecular Conformer Generation(DG) technique [84]. Some of the major drawbacks when using these methodsinclude (i) error accumulated in the predicted distances can mislead the atomcoordinate searching with DG, (ii) the distance matrix is overparameterizedcompared to the number of degrees of freedom, and (iii) specific geometricconstraints like torsion angles, rings are not explicitly modelled.ConfGF [79] and DGSM [85] uses Denoising Score Matching (DSM) [53, 86] tolearn the gradient of the log-density w.r.t inter-atomic distance. Subsequently,the gradient of the log-density w.r.t coordinates is derived using the chain rule.While this approach eliminates the need for using DG, it suffers from poorout-of-distribution performance, as the input to the model during trainingare noised distances which can be invalid (negative or not following triangularinequality); however, while testing, it is fed valid distances calculated fromcoordinates.GeoMol [80], a highly systematic method, uses MPNNs to predict a setof minimal geometric quantities (bond lengths, bond angles and torsionangles) along with a non-adversarial optimal transport based loss functionto promote diverse conformations. GeoDiff uses an E-GNN [66] that takesas input atom coordinates to learn an SE(3) equivariant score. TorsionDiff[12] designs a diffusion model on the torsion angles while incorporating thelocal structure from RDKiT ETKDG [87]. The current state-of-the-art isMolecular Conformer Fields (MCF) [14], which trains a diffusion model overfunctions, specifically functions that map elements from the molecular graphto points in 3D space using the framework of Diffusion Probabilistic Fields[15].27Chapter 4Role of Data Biases in MLPotential GeneralizationThis chapter is structured as follows. We introduce and motivate the problemto be addressed in Section 4.1. In Section 4.2, we cover the methodology interms of the structure of experiments, data splits and the evaluation strategy.In Section 4.3, we cover the results. Finally, in Section 4.4, we summarizeour efforts and conclude the chapter with limitations and future work.4.1 IntroductionML Potentials have gained popularity in the field of atomistic modellingand simulations over the past decade [60, 88â94]. Their appeal lies in theirtrade-off between speed and accuracy, enabling expedited calculations whilemaintaining comparable levels of precision compared to QM methods. Theyare mainly enabled by the recent developments in ML modelling for physicalsystems and the creation of large QM datasets that are made publicly available.The first is exemplified by the variety of model architectures and descriptorsallowing ML Potentials to comprehend the inherent symmetries and biaseswithin atomistic systems and QM modelling [22, 41, 66, 76, 95â98]. Thelatter is underscored by the increasing number of efforts to generate andpublicly release QM datasets, despite the substantial costs associated withsuch endeavours [10, 59â62, 64, 65, 99â102].The landscape of ML potentials and their inherent biases, as well as theirrole in generalization, has received some attention in the recent literature[103], whereas data biases, such as the QM level of theory, the numberof labelled molecules and conformers, and the diversity in chemical andconformational aspects, have been comparatively under-explored. These data-specific factors significantly affect the accuracy and generalization capabilitiesof ML Potentials. Consequently, the primary focus of this chapter is to shed284.2. Methodlight on the implications of data biases and to provide valuable insights andguidelines for optimizing the trade-off between the cost of data generationand the value it brings to modelling and generalization efforts.The main contributions in this chapter are as follows:1. We design and conduct experiments to understand the intricate relation-ship between dataset size, structural diversity, conformational diversityand model generalization.2. Our analysis of generalization is multifaceted, allowing the readers tounderstand how the performance of ML Potentials changes within andoutside the training distribution of both conformers and structures.4.2 MethodLetâs consider a QM dataset with N data points (conformers), encompassingns unique molecular structures, with fixed nc conformers per molecule (i.e.N = ns Ã nc). Our investigation seeks to analyze how generalization evolveswhen altering the dataset size (N), the structural diversity (ns), and theconformational diversity (nc). To give a comprehensive picture of ML Po-tential generalization, we consider four facets of model performance. In thesubsequent sections, we will delve deeper into the methodological setup andelaborate on the chosen generalization metrics. Itâs important to mentionthat for the present chapter, our definition of diversity is primarily based onthe count of unique molecules or conformations within a dataset. However,we intend to expand upon this definition in the future to incorporate measuresof (dis)similarity as well.4.2.1 SetupFor our investigation, we run two pivotal experiments, each involving trainingML Potentials on simulated QM datasets characterized by distinct values ofN , ns, and nc. For a visual representation of these experiments, please referto Figure 4.1.N-fixed experiment: Herein, we replicate a scenario where there is a fixedbudget for data generation. Our objective is to investigate the interplaybetween structural and conformational diversity and its influence on MLPotential generalization. By simulating the generation of QM datasets witha constant number of conformers (N), we concurrently vary the valuesof ns and nc. Specifically, as ns decreases, we proportionally increase nc294.2. Method1 2 4 8 16nc12.5k25k50k100k200k400k800k1.6MNns-Fixed Settingns12.5k25k50k100k1 2 4 8 16 32 64nc3.1k12.5k50k200k800k3.2Mn sN-Fixed SettingN50k200k800k3.2MFigure 4.1: Experimental setup: Left. (1) ns-Fixed: Keeping the number ofmolecules ns fixed at 12.5k, 25k, 50k and 100k, we increase the conformer permolecules nc. Right. (2) N -fixed: Keeping the total number of conformers Nfixed at 50k, 200k, 800k and 3.2M, we increase the conformer per moleculenc while decreasing the number of molecules ns.by the same factor. To illustrate, for N = 200K, we generate datasetswith (ns = 200K,nc = 1), (ns = 100K,nc = 2), (ns = 50K,nc = 4),(ns = 25K,nc = 8), and (ns = 12.5K,nc = 16). This gradual transitionspans from a setup featuring low conformational diversity but high structuraldiversity (ns = 200K,nc = 1) to one characterized by high conformationaldiversity and low structural diversity (ns = 12.5K,nc = 16). By varyingN â (50K, 200K, 800K, 3.2M), our aim is to explore the intricate relationshipbetween this trade-off and the generated dataset size.ns-fixed experiment: This experiment emulates a recent trend in QMdata generation, wherein an emphasis is placed on increasing conformationaldiversity due to its perceived importance in ML Potential generalization.Here, our goal is to evaluate the intrinsic impact of conformational diversityon ML Potential generalization. To achieve this, we simulate the creationof QM datasets where ns remains fixed, with values set at 12.5K, 25K, 50K,and 100K, while we systematically increase the value of nc from 1 to 16. Thetotal number of conformers (N) is defacto increasing with nc.nc-fixed experiment: We do not conduct any additional experimentsfor this setting where nc is fixed while N and ns increases. To observethe isolated impact of structural diversity on ML Potential generalization,we leverage the results from the N -fixed and ns-fixed experiments. Forinstance, for nc = 2, we gather results from experiments where nc = 2 andN â [25K, 50K, 100K, 200K] from the previous experiments.304.3. Results4.2.2 Generalization metricsThe distinct aspects of ML Potential model performance can be categorizedalong two axes of generalization. The first axis focuses on the similaritybetween test samples and the training distribution, distinguishing betweensamples that are Independent and Identically Distributed (IID) and those thatare Out-of-Distribution (OOD). As data points can exhibit variations alongboth structural and conformational dimensions, the second axis pertainsto differentiating chemical characteristics, encompassing both structuraland conformational aspects. Consequently, these axes yield four specificgeneralization metrics for analysis: IID structural (IID-S), OOD structural(OOD-S), IID conformational (IID-C), and OOD conformational (OOD-C).To calculate the IID-S metric, the test set consists of molecules with similarphysicochemical properties to those in the training set. Conversely, for OOD-S, the test molecules are drawn from a chemical subspace that is distant fromthe training set. For IID-C and OOD-C metrics, the test sets are composedof novel conformers belonging to molecules encountered during training. Todetermine whether a conformer is IID-C or OOD-C, we simply compute itsminimum Root Mean Square Distance (RMSD) to the training conformersand consider where it falls on that RMSD spectrum. We avoid choosing anarbitrary threshold herein because the spaces of conformers and RMSD arecontinuous, and what IID or OOD is might depend a lot on the potentialenergy surface, number of atoms, etc.4.3 Results4.3.1 Experimental detailsDatasets: For our experiments, we use the GEOM dataset [23], a largecollection comprising 37 million conformers covering 450K molecules. Ithas two subsets: GEOM-QM9 made of 133K small molecules from theQM9 dataset [99], with up to 9 heavy atoms (C, N, O, F) and GEOM-Drugs consisting of 317K larger and drug-like molecules. We simulate allour QM data generation by sampling from GEOM-Drugs, and we considerGEOM-QM9 as structurally OOD from it. The structural differences betweenGEOM-Drugs and GEOM-QM9 are illustrated in Appendix A.1.Model Training: To train our ML Potentials, we use the Equivariant Trans-former, a component of the TorchMD-NET models [91]. Our model has ap-proximately 2 million parameters over num_layers=8 and hidden_channels=128.314.3. ResultsOther hyperparameters are left to their default values 13. We trained withthe L2 loss and the Adam optimizer with a cosine annealing scheduler forthe learning rate between 10â8 and 10â4.Model Evaluation: We evaluate the modelsâ performance using the meanabsolute error (MAE) on the potential energy. The IID-S metric is computedusing unseen molecules from GEOM-Drugs, and the OOD-S is computedusing molecules from GEOM-QM9, whose chemical space is very differentfrom that of drug-like molecules. IID-C and OOD-C metrics are computedusing molecules that have been seen during training according to criteriadescribed in Section 4.2.2.Our experiments are repeated three times using different random seeds,leading to varied data splits and model initializations. For each result, weinclude error bars to illustrate the standard deviation across these three splits.4.3.2 Structural generalizationFigure 4.2 presents the structural generalization metrics for the N -fixedexperiment, illustrating their dependence on nc and, implicitly, on ns, as thetwo variables are inversely related in this setup. Across different values of N ,we observe a gradual increase in IID-S MAE as nc increases and ns decreases.Although the rate of this increase is less pronounced for larger values of N ,there remains a notable two-fold increase in IID-S MAE when structuraldiversity decreases by a factor of four and N = 3.2M . Conversely, OOD-SMAE also shows an increase with rising values of nc, but these trends are lesspronounced across all N values. This phenomenon can partly be attributedto the inherently larger OOD-S MAEs when compared to IID-S MAEs. Infact, the best IID-S MAEs remain in the low single digits, whereas the bestOOD-S MAEs hover around 50kcal/mol.Collectively, these results underscore that within fixed budget constraints,the structural generalization capabilities of ML Potentials significantly dete-riorate when prioritizing conformational diversity over structural diversity.Consequently, one should exercise caution when sacrificing structural diversityin favour of conformational diversity.Figure 4.3 shows the structural generalization metrics for the ns-fixed experi-ment, demonstrating their dependency on nc and implicitly on N which areproportional in this setup. For lower values of ns (i.e., ns â [12K, 25K]), we13Implementation as provided in https://github.com/torchmd/torchmd-net324.3. Results1 2 4 8 16 32 64nc110100MAE (kcal/mol)IID Structural1 2 4 8 16 32 64ncOOD StructuralN50k200k800k3.2mFigure 4.2: N -fixed: Performance on IID-S and OOD-S as we increasethe conformational diversity (nc) and reduce structural diversity (ns) whilekeeping the number of conformers (N) fixed .1 2 4 8 16nc110100MAE (kcal/mol)IID Structural1 2 4 8 16ncOOD Structuralns12k25k50k100kFigure 4.3: ns-fixed: Performance on IID-S and OOD-S as we increase theconformational diversity (nc) while keeping structural diversity (ns) fixed .observe a gradual reduction in both IID-S and OOD-S MAEs as conforma-tional diversity increases. Although the decrease in MAEs is less pronouncedfor OOD generalization, it remains notably significant. On the other hand,in cases with higher values of ns (i.e., ns â [50K, 100K]), both IID-S andOOD-S MAEs decrease rapidly with a small increase in conformational diver-sity, but when it increases further, IID-S MAE plateaus and OOD-S MAEbegins to increase. These findings suggest that when structural diversityis low, enhancing conformational diversity can be beneficial. However, asstructural diversity increases, the advantages of additional conformationaldiversity diminish significantly.Figure 4.4 shows the structural generalization for the nc-fixed experiment,highlighting the proportional relationship with N and implicitly with ns.334.3. Results10k 100k 1M 10MN110100MAE (kcal/mol)IID Structural10k 100k 1M 10MNOOD Structuralnc124816Figure 4.4: nc-fixed: Performance on IID-S and OOD-S as we increase thenumber of conformers (N) while keeping conformational diversity (nc) fixed .A clear trend is observed where increasing the total number of conformersN helps with better IID-S and OOD-S generalization. Additionally, theimportance of structural diversity can be observed as experiments with lowernc or higher ns generalize better than the ones with higher nc or lower ns.Across both experiments, irrespective of the particular values of N , nc, andns, we consistently observe that IID-S MAEs remain significantly lowerthan OOD-S MAEs. This emphasizes the ML Potentialâs limited capacityto generalize beyond its training distribution. Therefore, it is imperativefor both experimenters and model users to clearly understand the modelâsstructural applicability domain.4.3.3 Conformational generalizationIn Figure 4.5, we delve into conformational generalization in the N -fixedexperiment, examining its dependence on nc (implicitly ns). A consistentpattern emerges across all N values: the MAE remains relatively stable whenthe RMSD to the training conformers is below 2 Ã . However, beyond thisthreshold, we observe an increase in MAE, followed by a return to near-initialvalues as RMSD increases. Specifically, the plots reveal a steep MAE increasewhen 3.5 Ã â¤ RMSD â¤ 5 Ã in scenarios with low conformational diversity(nc â¤ 4) but high structural diversity in the training set. Conversely, lesssteep increases occur when MAE registers between 2.5 Ã and 4 Ã forhigh conformational diversity (nc â¥ 32) in the training set. The flattestcurves are evident when nc â [8, 16], highlighting the need for a delicatetrade-off between structural and conformational diversity to achieve effectivegeneralization to unseen conformers of seen molecules.344.3. Results0 1 2 3 4 5 6RMSD from Training Distribution01020304050MAE (kcal/mol)N: 50knc1248160 1 2 3 4 5 6RMSD from Training Distribution01020304050MAE (kcal/mol)N: 200knc12481632640 1 2 3 4 5 6RMSD from Training Distribution01234567MAE (kcal/mol)N: 800knc481632640 1 2 3 4 5 6RMSD from Training Distribution0.51.01.52.02.53.03.54.0MAE (kcal/mol)N: 3.2mnc163264Figure 4.5: N -fixed: Distribution of MAE performance for conformers fromboth IID-C and OOD-C based on their RMSD from the training distribution.The four plots represent fixed N values (50k, 200k, 800k, and 3.2M) withvarying nc and ns.In Figure 4.6, we explore conformational generalization in experiments wherestructural diversity is fixed and conformational diversity varies. Across allns values, we observe consistent MAE values for all RMSD when nc â[8, 16]. However, in low conformational diversity settings (i.e., nc â¤ 4), MAEremains steady when RMSD â¤ 3 Ã , but as RMSD increases, so does MAEbefore gradually decreasing. The steepness of these MAE increases, and themaximum values reached are inversely related to conformational diversity.This reaffirms the conclusions drawn from the fixed budget experiments:the trade-off between conformational and structural diversity significantlyimpacts conformational generalization.In Figure 4.7, we observe that across all nc plots, increasing the total number354.3. Results0 1 2 3 4 5RMSD from Training Distribution020406080100120MAE (kcal/mol)ns: 12k0 1 2 3 4 5RMSD from Training Distribution01020304050MAE (kcal/mol)ns: 25k0 1 2 3 4 5RMSD from Training Distribution0102030MAE (kcal/mol)ns: 50k0 1 2 3 4 5RMSD from Training Distribution01020304050MAE (kcal/mol)ns: 100knc124816Figure 4.6: ns-fixed: Distribution of MAE performance for conformers fromboth IID-C and OOD-C based on their RMSD from the training distribution.The four plots represent fixed ns values (12k, 25k, 50k, and 100k) with varyingnc and N .of conformers N , helps improve conformational generalization across allRMSD values, however, the performance gains reduce every time we increasethe value of N . Additionally, across the different nc plots, we observe thatthe maximum MAE observed decreases as nc increases, suggesting that highnc is essential for conformational generalization.While our experiments indicate that the optimal number of conformers permolecule for effective generalization across conformers in both IID and OODfalls between 8 and 16, itâs important to note that this may vary based onother experimental factors, such as the network architecture and the chemicalspace of the training set. Therefore, experimenters should determine theoptimal level of conformational diversity tailored to their specific chemicalspace and ML potential modelling approach.364.4. Discussion and Future Directions0 1 2 3 4 5RMSD from Training Distribution020406080100120MAE (kcal/mol)nc: 1N1250025000500001000000 1 2 3 4 5RMSD from Training Distribution020406080100MAE (kcal/mol)nc: 2N25000500001000002000000 1 2 3 4 5RMSD from Training Distribution01020304050MAE (kcal/mol)nc: 4N500001000002000004000000 1 2 3 4 5RMSD from Training Distribution0510152025MAE (kcal/mol)nc: 8N100000200000400000800000Figure 4.7: nc-fixed: Distribution of MAE performance for conformers fromboth IID-C and OOD-C based on their RMSD from the training distribution.The four plots represent fixed nc values (1, 2, 4, and 8) with varying ns andN .4.4 Discussion and Future DirectionsIn the pursuit of developing ML Potentials for atomistic modelling, ourchapter delved into the intricate interplay between conformational and struc-tural diversity, data size and model generalization. Through comprehensiveexperiments, we unravelled the following key insights that hold significantimplications for the ML Potential community,â¢ In the N -fixed experiment (fixed dataset size), we found that optimalstructural generalization requires a balance between structural andconformational diversity. Increasing conformational diversity at theexpense of structural diversity led to a steep rise in MAEs.â¢ In the ns-fixed experiment (fixed structural diversity), we observedthat the benefits of increased conformational diversity are more pro-374.4. Discussion and Future Directionsnounced when structural diversity is limited. However, as structuraldiversity increases, the advantages of additional conformational diversitydiminish.â¢ Throughout both experiments, the modelâs generalization capabili-ties were limited to its training distribution, with substantially lowerin-distribution MAEs compared to out-of-distribution MAEs. Thishighlights the need to define and recognize the modelâs applicabilitydomain.While our study has rigorously explored the influence of data biases on MLIPgeneralization, it uses a specific architecture and dataset, so we acknowledgethe need to enhance the validity of our conclusions. Consequently, we intendto conduct a more extensive analysis encompassing various MLIP modellingbiases and incorporating diverse QM datasets. Our plans involve the uti-lization of alternative QM datasets, employing improved DFT theory levels,incorporating force labels, and leveraging state-of-the-art MLIP architectures,such as Equiformer [104] and MACE [71]. Additionally, we plan to consideralternate definitions of structural diversity, specifically, Shannon Entropy[105] and TM-score [106]. This broader experimentation will provide a com-prehensive understanding of the impact of data biases on MLIP generalization,contributing to the advancement of atomistic modelling in various scientificdomains.38Chapter 5Equivariant Flow Matching forConformer GenerationThis chapter is structured as follows. We introduce and motivate the problemto be addressed in Section 5.1. In Section 5.2, we cover the details involvedin developing ET-Flow, ranging from flow-matching, optimal transport align-ment, stochastic sampling and architectural modifications. The followingsection Section 5.3 covers the various experiments performed to highlightthe efficacy of the method and the several ablations involved. In Section 5.4,we summarize our efforts and conclude the chapter with the limitations andfuture works.5.1 IntroductionAn alternate approach to sampling low-energy molecular conformations fromthe molecular graph is by using deep generative models like Diffusion [2, 11].Prior methods typically fall into two categories: diffusing along the internalgeometry such as pairwise distances, bond angles, and torsion angles [12, 80]or diffusing the atomic coordinates in the Cartesian space [77]. TorsionDiff[12] was the first to outperform cheminformatics methods by diffusing onlyon torsion angles. However, it has significant limitations from its relianceon RDKiT ETKDG [13] for generating local structures, such as being non-differentiable and exhibiting poor quality for large and complex molecules.Unlike prior approaches, the current state-of-the-art MCF [81] proposes adomain-agnostic approach by learning to diffuse over functions [15]. However,it comes with drawbacks such as high computational demands due to a largenumber of parameters, limited sample efficiency from lack of inductive biaseslike Euclidean symmetries, and potential difficulties in scenarios with sparsedata â a common challenge in this field.We propose Equivariant Transformer Flow (ET-Flow), a simple yet powerful395.1. IntroductionFigure 5.1: Overview of ET-Flow. The model takes as input the featuresrelated to the molecular structure, like atom and bond features, as well as aradius graph of 10Ã from the atom positions. Samples are drawn from theharmonic prior and then rotationally aligned with the samples from the data.A probability path is constructed between pairs of x0 and x1, and xt is thensampled from this path at a random time t.flow-matching model designed to generate low-energy 3D structures of smallmolecules with minimal assumptions. We utilize flow matching [3, 17, 54],which enables the learning of arbitrary probability paths beyond diffusionpaths, enhancing both training and inference efficiency compared to conven-tional diffusion generative models. Departing from traditional equivariantarchitectures like EGNN [66], we adopt an Equivariant Transformer [19] tobetter capture geometric features. Additionally, our method integrates a Har-monic Prior [107, 108], leveraging the inductive bias that atoms connected bya bond should be in close proximity. We further optimize our flow matchingobjective by initially conducting rotational alignment on the harmonic prior,thereby constructing shorter probability paths between source and targetdistributions at a minimal computational cost.The main contributions in this study are as follows,1. We obtain state-of-the-art precision results for molecule conformerprediction and improve upon the state-of-the-art by a large margin onensemble property prediction.2. Counter to the prevailing status quo in the literature, we provide strongevidence that previous innovations that improved performance at the405.2. Methodcost of modelling complexity are not necessary. The significantly simplerbut well-engineered ET-Flow can provide better performance.3. ET-Flow uses orders of magnitude fewer sampling steps than GeoDiff[77] and significantly fewer parameters than MCF [14].5.2 MethodGiven a molecular graph, we design a scalable equivariant model that generatesenergy-minimized conformers. In this section, we lay out the framework toachieve this objective by detailing the generative process in flow matching,the rotation alignment between distributions, stochastic sampling, and finally,the architecture details.Preliminaries We define notations that we use throughout this chapter.Inputs are continuous atom positions x â RNÃ3 where N is the number ofatoms. We interchangeably use the notation vt(x) with v(t,x) for a vectorfield.5.2.1 Flow MatchingThe aim is to learn a time-dependent vector field vt(x) : RNÃ3 Ã [0, 1] âRNÃ3 associated with the transport map Xt : RNÃ3 Ã [0, 1] â RNÃ3 thatpushes forward samples from a base distribution Ï0, often an easy-to-sampledistribution, to samples from a more complex target distribution Ï1, thelow-energy conformations of a molecule. This can be defined as an ordinarydifferential equation (ODE),XËt(x) = vt(Xt(x)), Xt=0 = x0, (5.1)where x0 â¼ Ï0. We can construct the vt via a time-differentiable interpolationbetween samples from Ï0 and Ï1 that gives rise to a probability path Ït thatwe can easily sample from [3, 18, 54, 58]. The general interpolation betweensamples x0 â¼ Ï0 and x1 â¼ Ï1 can be defined as:It(x0,x1) = αtx1 + βtx0. (5.2)Given this interpolant, we can define the probability path Ït(x) = N (x|It(x0,x1), Ï2t I),and the vector field can be computed as vt(x) = âtÏt(x) which has the fol-lowing form,vt(x) = Î±Ëtx1 + Î²Ëtx0 + ÏËtz z â¼ N (0, I). (5.3)415.2. MethodFigure 5.2: Rotation alignment using the Kabsh algorithm [4] resulting inshorter and straighter paths between the base and target structures.In our work, we use linear interpolation where αt = t, βt = 1 â t, andÏt = Ïât(1â t) (where Ï is a hyperparameter set to 0.1), resulting in thevector field,vt(x) = x1 â x0 + 1â 2t2ât(1â t)z. (5.4)Now, we can define the objective function for learning a vector field vθ(x)that generates a probability path Ït between a base density Ï0 and the targetdensity Ï1 as,L = Etâ¼U(0,1),xâ¼Ït(x0,x1)â¥v(t,x)â vÎ¸(t,x)â¥2. (5.5)For training, we sample (i) x0 â¼ Ï0, x1 â¼ Ï1, and t â¼ U(0, 1), (ii) interpolateaccording to Equation (5.2), (iii) add noise from a standard Gaussian, and(iv) minimize the loss defined in Equation (5.5). For sampling, we samplex0 â¼ Ï0 and integrate from t = 0 to t = 1 using Eulerâs method. At eachtime-step, the Euler solver iteratively predicts the vector field for xt andupdates its position xt+ât = xt + vÎ¸(t,x)ât. More details on the trainingand sampling algorithms are provided in Appendix A.3.5.2.2 AlignmentSeveral previous works [18, 109, 110] demonstrate that constructing a straighterpath between base distribution Ï0 and target distribution, Ï1 minimizes thetransport costs and improves performance. In our work, we reduce thetransport costs between samples from the harmonic prior Ï0 and samplesfrom the data distribution Ï1 by rotationally aligning them using the Kabsch425.2. Methodalgorithm [4] similar to [109, 110]. This approach leads to faster convergenceand reduces the path length between atoms by leveraging the similarity in"shape" of the samples as seen in Figure 5.1 and Figure 5.2 without incurringhigh computational costs.5.2.3 Stochastic SamplingWe employ a variant of the stochastic sampling technique inspired by [111].Specifically, we inject noise at each time step, evaluate the vector field fromthe intermediate state after adding noise, and then perform the determin-istic ODE step. The original method utilizes a second-order integration,which averages the denoiser output at the noisy intermediate state and thestate at the next time step after integration. In our experiment, we usethe stochastic sampler without this second-order correction term, whichempirically provided a performance boost comparable to the second-ordermethod. We apply stochastic sampling only during the final part of the inte-gration steps, specifically within the range t â [0.8, 1.0]. This helps preventdrifting towards overpopulated density regions and improves the quality ofthe samples, as noted by [111]. Stochastic sampling has improved both thediversity and accuracy of the generated conformers, measured by Coverageand Average Minimum RMSD (AMR), respectively, as shown in Table 5.1.Detailed information on the stochastic sampling algorithm is provided inAppendix A.3.5.2.4 Chirality CorrectionFigure 5.3: Two molecules that are isomorphic (one-to-one correspondencebetween atoms and bonds) can have different orientations in 3D determinedby the order of the tetrahedral chiral carbonâs neighbours.While generating conformations, it is necessary to take into account thestereochemistry of atoms bonded to four distinct groups, also referred to astetrahedral chiral centres. Molecules that only differ in the orientation ofthese chiral centres are enantiomers and are mirror images of each other. The435.2. Methodorientated volume (OV) of a chirality centre can be computed based on thevolume of the tetrahedron formed by the ordered coordinates of the distinctgroups (say p1,p2,p3,p4) as,OV (p1,p2,p3,p4) = signï£«ï£¬ï£¬ï£­â£â£â£â£â£â£â£â£1 1 1 1x1 x2 x3 x4y1 y2 y3 y4z1 z2 z3 z4â£â£â£â£â£â£â£â£ï£¶ï£·ï£·ï£¸ , (5.6)where sign is the sign function14 applied on the value of the determinant.We also have access to the required orientation for each chiral centre usingthe chirality tag provided by RDKit. In our work, we observe that ourmodel sometimes generates conformations with incorrect chirality. In orderto correct those errors, once we generate conformations, we do a post hoccorrection where we compute the orientation volume for each chiral centreand compare it with the chiral tag provided by RDKit. If there is a mismatch,we simply flip the conformation against the z-axis (corresponding to takingthe mirror image). This is similar to the orientation correction done for thegenerated local structure in GeoMol [80].5.2.5 ArchitectureFigure 5.4: Overall Architecture of ET-Flow consisting of 3 components.ET-Flow takes as input the molecular graph G, time-step t, relative positionsRâ and predicts the vector field vâÎ¸, (1) Embedding Layer for encoding theatomic positions, molecular graph information and time-step into invariantfeatures xi, equivariant features vâi (for atom i). (2) Update layers (withattention layers) updates the invariant features xi and equivariant featuresvâi at each layer and (3) Output Layer for predicting the vector field vâÎ¸ giveninvariant features xi and equivariant features vâi.14https://en.wikipedia.org/wiki/Sign_function445.2. MethodThe ET-Flow architecture (Figure 5.4) consists of 3 major components: anembedding layer, update layers and an output layer. We use a modifiedversion of the embedding and update layers from the equivariant transformerarchitecture of TorchMD-NET [19] whereas the output layer utilizes the gatedequivariant blocks from [21]. We highlight our modifications over the originalarchitectures with this colour. These modifications enable stabilized trainingsince we use a larger network than the one proposed in the TorchMD-NET[19]. Additionally, since our input structures are interpolations betweenstructures sampled from a prior and an actual conformation, it is importantto ensure our network is numerically stable when the interpolations containtwo atoms very close to each other.Embedding Layer: This layer learns an embedding for the ith atom usingthe atomic number (zi), atomic attributes (hi), time-step t, edge features lijand neighborhood embedding ni. Atomic numbers are embedded into twovectors, one for atomic number embedding and the other for neighbourhoodembedding ni. The atomic attributes hi are projected into a feature vectorusing a simple 2-layer Multi-Layer Perceptron (MLP). The neighbourhoodembedding for an atom ni is computed as,zi = embedint(zi) (5.7)hi = MLP(hi) (5.8)ni =Nâj=1embednbh(zj) · g(dij , lij) (5.9)where the second term g(dij , lij) is a edge-based feature, which is a linearprojection of a vector concatenating k exponential radial basis functionsmultiplied with a cutoff function Ï(Â·) acting on edge lengths dij and edgefeatures lij . It is computed as,g(dij , lij) =WF[Ï(dij)eRBF1 (dij)....Ï(dij)eRBFK (dij), lij](5.10)Ï(dij) ={12(cos(Ïdijdcutoff+ 1)), if dij â¤ dcutoff0, otherwise(5.11)The hyper-parameters for the radial basis functions are set as proposed inUnke and Meuwly [112]. Once the neighbourhood embedding ni is computed,the final atom-level embedding is computed via a linear projection as,xi =WC[embedint(zi), hi, t, ni](5.12)455.2. MethodAttention Mechanism: The multi-head dot-product attention operationuses atom features xi, atom attributes hi, time-step t and inter-atomicdistances rij to compute attention weights. The input atom-level features xiare mixed with the atom attributes hi and the time-step t using an MLP andthen further normalized using a LayerNorm [113]. To compute the attentionmatrix, the inter-atomic distances rij are projected into two dimensionalfilters DK and DV as,DK = Ï(WDKeRBF (rij) + bDK)DV = Ï(WDVeRBF (rij) + bDV)(5.13)The atom level features are then linearly projected along with a LayerNormoperation to derive the query Q and key K vectors. The value vector V iscomputed with only the linear projection of atom-level features. ApplyingLayerNorm on Q, K vectors (also referred to as QK-Norm) has proven tostabilize unnormalized values in the attention matrix [114, 115] when scalingnetworks to a large number of parameters. The Q and K vectors are thenused along with the distance filter DK for a dot-product operation over thefeature dimension,Q = LayerNorm(WQxi), K = LayerNorm(WKxi), V =W V xi (5.14)dot(Q,K,DK) =FâkQk Â·Kk Â·DKk . (5.15)The attention matrix is derived by passing the above dot-product operationmatrix through a non-linearity and weighting it using a cosine cutoff Ï(dij)(similar to Equation 5.11) which ensures the attention weights are non-zeroonly when two atoms are within a specified cutoff,A = SiLU(dot(Q,K,DK)) · Ï(dij). (5.16)Using the value vector V and the distance filter DV , we derive 3 equally sizedfilters by splitting along the feature dimension (s1ij , s2ij , s3ij = split(Vj Â·DVij )).A linear projection is then applied to combine the attention matrix and thevectors s3ij to derive an atom-level feature yi (= WO(âNj Aij · s3ij)). Theoutput of the attention operation is yi (an atom-level feature) and two scalarfilters s1ij and s2ij (edge-level features).Update Layer: The update layer computes interactions between atoms inthe attention block and uses the outputs to update the scalar feature xi and465.3. Experimentsthe vector feature vâi. First, the scalar feature output yi from the attentionmechanism is split into 3 features (q1i , q2i , q3i ), out of which q1i and q2i are usedfor the scalar feature update as,âxi = q1i + q2i · â¨U1vâi · U2vâiâ©, (5.17)where â¨U1vâi · U2vâiâ© is the inner product between linear projections of vectorfeatures vâi with matrices U1, U2.The edge vector update consists of two components. First, we compute avector wâi, which for each atom is computed as a weighted sum of vectorfeatures and a clamped-norm of the edge vectors over all neighbours. Finally,a sum over wâi and a linear projection of the vector feature gives us,wâi =Nâjs1ij · vâj + s2ij Â·râi â râjmax(â¥râi â râjâ¥, ϵ) , (5.18)âvâi = wâi + q3i · U3vâi (5.19)where U1 and U3 are projection vectors over the feature dimension of thevector feature vâi. Our main modifications in this layer are that we clamp theminimum value of the norm (to ϵ = 0.01) to prevent numerically large valuesin cases where positions of two atoms are sampled two close from the prior.Output Layer: The output layer consists of Gated Equivariant Blocks from[21]. Given atom scalar xi and vector features vâi, the updates in each blockis defined as,xi,updated, wâi = split(MLP([xi, U1vâi)])) (5.20)vâi,updated = (U2vâi) · wâi (5.21)Here, U1 and U2 are linear projection matrices that act along feature dimen-sions. Our modification is to use LayerNorm in the MLP to improve trainingstability.5.3 ExperimentsWe empirically evaluate ET-Flow by comparing the generated and ground-truth conformers in terms of distance-based RMSD (Section 5.3.2) andchemical property based metrics (Section 5.3.4). We present the generalexperimental setups in Section 5.3.1. The implementation details are providedin Appendix A.2.475.3. ExperimentsTable 5.1: Molecule conformer generation results on GEOM-DRUGS (δ =0.75Ã ). SS is ET-Flow with stochastic sampling. For both ET-Flow andET-Flow-SS, we sample conformations over 50 time steps.Recall PrecisionCoverage â AMR â Coverage â AMR âmean median mean median mean median mean medianGeoDiff 42.10 37.80 0.835 0.809 24.90 14.50 1.136 1.090GeoMol 44.60 41.40 0.875 0.834 43.00 36.40 0.928 0.841Torsional Diff. 72.70 80.00 0.582 0.565 55.20 56.90 0.778 0.729MCF - S (13M) 79.4 87.5 0.512 0.492 57.4 57.6 0.761 0.715MCF - B (62M) 84.0 91.5 0.427 0.402 64.0 66.2 0.667 0.605MCF - L (242M) 84.7 92.2 0.390 0.247 66.8 71.3 0.618 0.530ET-Flow (8.3M) 79.53 84.57 0.452 0.419 74.38 81.04 0.541 0.470ET-Flow - SS (8.3M) 79.62 84.63 0.439 0.406 75.19 81.66 0.517 0.4425.3.1 Experimental SetupEvaluation: Our evaluation methodology is similar to that of [12]. First,we look at RMSD-based metrics like Coverage and Average Minimum RMSD(AMR) between generated and ground truth conformer ensembles. For this,we generate 2K conformers for a molecule with K ground truth conformers.Second, we look at chemical similarity using properties like Energy (E),dipole moment (µ), HOMO-LUMO gap (âÏµ) and the minimum energy (Emin)calculated using xTB [34].Baselines: We benchmark ET-Flow against leading approaches outlined inSection 3.3. Specifically, we assess the performance of GeoMol [80], GeoDiff[77], Torsional Diff [12], and MCF [14]. The most recent among theseapproaches, MCF, has demonstrated superior performance across evaluationmetrics compared to its predecessors. Itâs worth mentioning that GeoDiffinitially utilized a limited subset of the DRUGS dataset; thus, for a faircomparison, we consider its re-evaluated performance as presented in [12].5.3.2 Ensemble RMSDAs shown in Table 5.1 and Table 5.2, ET-Flow outperforms all precedingmethodologies and demonstrates competitive performance with the previousstate-of-the-art, MCF [14]. Despite being significantly smaller with only8.3M parameters, ET-Flow shows a substantial improvement in the qualityof generated conformers, as evidenced by superior Precision metrics across allMCF models, including the largest MCF-L. When compared to MCF-S, which485.3. ExperimentsTable 5.2: Molecule conformer generation results on GEOM-QM9 (δ = 0.5Ã ).For both ET-Flow, we sample conformations over 50 time steps.Recall PrecisionCoverage â AMR â Coverage â AMR âmean median mean median mean median mean medianCGCF 69.47 96.15 0.425 0.374 38.20 33.33 0.711 0.695GeoDiff 76.50 100.00 0.297 0.229 50.00 33.50 1.524 0.510GeoMol 91.50 100.00 0.225 0.193 87.60 100.00 0.270 0.241Torsional Diff. 92.80 100.00 0.178 0.147 92.70 100.00 0.221 0.195MCF 95.0 100.00 0.103 0.044 93.7 100.00 0.119 0.055ET-Flow 94.99 100.00 0.083 0.035 91.00 100.00 0.116 0.047is closer in size, ET-Flow achieves markedly better Precision while the impacton Recall is less significant and limited to Recall Coverage. Notably, ourRecall AMR remains competitive with much bigger MCF-B, underscoring theinherent advantage of our method in accurately predicting overall structures.5.3.3 Coverage Threshold PlotsFigure 5.5: Recall and Precision Coverage result on GEOM-DRUGS as afunction of the threshold distance. ET-Flow outperforms TorsionDiff by alarge margin, especially in a lower threshold region.We show a breakdown of the performance on GEOM-DRUGS of ET-Flowagainst Torsional diffusion [12] as a function of the threshold distance inFigure 5.5. Across a wide range of threshold values, ET-Flow consistently495.3. Experimentsoutperforms Torsional Diffusion in terms of both recall and precision, exceptfor slight underperformance at very high threshold values. Particularly notableis ET-Flowâs superior performance at lower threshold regions, highlightingits ability for accurate conformer prediction. We were unable to replicatethe results of MCF [14] as their codebase is currently not publicly available;therefore, we did not include them in our comparison plot.5.3.4 Ensemble PropertiesRMSD provides a geometric measure for assessing ensemble quality, but itis also essential to consider the chemical similarity between generated andground truth ensembles. For a random 100-molecule subset of DRUGS,if a molecule has K ground truth conformers, we generate a minimum of2K and a maximum of 32 conformers per molecule. These conformers arethen relaxed using GFN2-xTB [34], and the Boltzmann-weighted propertiesof the generated and ground truth ensembles are compared. Specifically,using xTB [34], we compute properties such as energy (E), dipole moment(µ), HOMO-LUMO gap (âÏµ), and the minimum energy (Emin). Table 5.3illustrates the median errors for ET-Flow and the baselines, highlighting ourmethodâs capability to produce chemically accurate ensembles. Notably, weachieve significant improvements over both TorsionDiff and MCF across allevaluated properties.Table 5.3: Median averaged errors of ensemble properties between sampledand generated conformers (E, âÎµ, Emin in kcal/mol, and µ in debye).E µ âÏµ EminOMEGA 0.68 0.66 0.68 0.69GeoDiff 0.31 0.35 0.89 0.39GeoMol 0.42 0.34 0.59 0.40Torsional Diff. 0.22 0.35 0.54 0.13MCF 0.68±0.06 0.28± 0.05 0.63±0.05 0.04±0.00ET-Flow 0.18±0.01 0.18±0.01 0.35±0.06 0.02±0.005.3.5 Results on GEOM-XLWe now assess how well a model trained on GEOM-DRUGS generalises tounseen molecules with large numbers of atoms, using the GEOM-XL datasetcontaining a total of 102 molecules. This provides insights into the modelâs505.3. ExperimentsAMR-P â AMR-R â # molsmean median mean medianGeoDiff 2.92 2.62 3.35 3.15 -GeoMol 2.47 2.39 3.30 3.14 -Tor. Diff. 2.05 1.86 2.94 2.78 -MCF - S 2.22 1.97 3.17 2.81 102MCF - B 2.01 1.70 3.03 2.64 102MCF - L 1.97 1.60 2.94 2.43 102ET-Flow (ours) 2.31 1.93 3.31 2.84 102Tor. Diff. 1.93 1.86 2.84 2.71 77MCF - S 2.02 1.87 2.9 2.69 77MCF - B 1.71 1.61 2.69 2.44 77MCF - L 1.64 1.51 2.57 2.26 77ET-Flow (ours) 2.00 1.80 2.96 2.63 75Table 5.4: Generalization results on GEOM-XL.capacity to tackle larger molecules and out-of-distribution tasks. Uponexecuting the checkpoint provided by Torsional Diffusion, we encountered27 failed cases for generation, likely due to RDKit failures, similar to theobservations in MCF, albeit with slightly different exact numbers. In bothexperiments involving all 102 molecules and a subset of 75 molecules, ET-Flowachieves performance comparable to Torsional Diffusion and MCF-S, butfalls short of matching the performance of MCF-B and MCF-L. Itâs worthnoting that MCF-B and MCF-L are significantly larger models, potentiallyaffording them an advantage in generalization tasks. As part of our futurework, we plan to scale up our model and conduct further tests to explore itsperformance in this regard.5.3.6 Inference Steps AblationIn Table 5.1, our sampling process with ET-Flow utilizes 50 inference steps.To evaluate the methodâs performance under constrained computational re-sources, we conducted an ablation study by progressively reducing the numberof inference steps. Specifically, we sample for 5, 10 and 20 time steps. Theresults on GEOM-DRUGS are presented in Table 5.5. We observed minimalperformance degradation with a decrease in the number of steps. Notably,ET-Flow demonstrates high efficiency, maintaining performance across allprecision and recall metrics even with as few as 5 inference steps. Interestingly,ET-Flow with 5 steps still achieves superior precision metrics compared to all515.4. Discussion and Future DirectionsTable 5.5: Ablation over number of inference steps on GEOM-DRUGS(δ = 0.75Ã ). Performance of ET-Flow at 5 steps is competent across allmetrics while also retaining state-of-the-art performance on precision metricswhen compared with previous methods.Recall PrecisionCoverage â AMR â Coverage â AMR âmean median mean median mean median mean medianET-Flow (5 Steps) 77.84 82.21 0.476 0.443 74.03 80.8 0.55 0.474ET-Flow (10 Steps) 79.05 84.00 0.451 0.415 74.64 81.38 0.533 0.457ET-Flow (20 Steps) 79.29 84.04 0.449 0.413 74.89 81.32 0.531 0.454ET-Flow (50 Steps) 79.53 84.57 0.452 0.419 74.38 81.04 0.541 0.470existing methods. This underscores ET-Flowâs ability to generate high-qualityconformations while operating within limited computational budgets.5.4 Discussion and Future DirectionsIn this chapter, we present our simple and scalable method ET-Flow, whichutilizes an equivariant transformer with flow matching to achieve state-of-the-art performance on multiple molecular conformer generation benchmarks.We emphasize that incorporating inductive biases, such as equivariance, andenhancing probability paths with a harmonic prior and RMSD alignmentenables us to attain these results while maintaining parameter and sampleefficiency.Limitations: While ET-Flow demonstrates competitive performance inmolecular conformer generation, there are areas for improvement. Onesuch area is the recall metrics, which capture the diversity of generatedconformations. Another area is the use of an additional chirality correctionstep that is used to predict conformations with the desired chirality. Moreover,although our performance on the GEOM-XL dataset [116] is comparable toMCF-S [14] and TorsionDiff [12], there is still room for improvement.We propose three future directions to address the limitations,â¢ Introducing Stochasticity: We observe during experiments that awell-designed sampling process incorporating stochasticity can enhancethe quality and diversity of generated samples. An extension of ourcurrent approach could involve using Stochastic Differential Equations525.4. Discussion and Future Directions(SDEs), which utilize both vector field and score in the integrationprocess, potentially improving the diversity of samples.â¢ Model Scaling: We propose to scale the number of parameters ofET-Flow, which has not only been shown to be useful across differentdomains of deep learning but has also shown to be useful in molecularconformer generation for MCF [14].â¢ Incorporating Chirality Correction in Architecture: Instead ofan additional chirality correction step, we aim to incorporate this di-rectly into the architecture in the future via using an SO(3) equivariantnetwork instead of O(3) equivariance.53Chapter 6ConclusionMolecules are in constant movement due to the various covalent and non-covalent interactions at play resulting in the formation of a complex potentialenergy surface. Improving existing methods for sampling and estimatingthe potential energy surface can accelerate existing drug-discovery efforts.In this thesis, we explore two areas where machine learning can be used toimprove traditional approaches: (1) Improving Molecular Dynamics with MLPotentials and (2) Directly sampling of stable conformations via EquivariantFlow Matching.In Chapter 4, we observe that generating QM datasets for training ML po-tentials is a computationally intensive task, and existing datasets exhibittrade-offs between various dataset biases (Table 3.1). Our study exploresthe impact of data biases, specifically structural diversity, conformationaldiversity, and dataset size, on the generalization of ML potentials. We con-duct controlled experiments by training ML potentials under a fixed budgetand a fixed structural diversity setting and evaluate their performance oncarefully designed test sets. Our results emphasize the need for a carefulbalance between structural and conformational diversity to achieve optimalgeneralization. Moreover, we observe that current ML potentials only general-ize within the limits of the training distribution, highlighting the importanceof defining and recognizing the modelâs applicability domain. We hope ourstudy informs QM data generation efforts with an improved understanding ofthe trade-off between structural and conformational diversity. Additionally,we hope our work inspires the development of improved ML potentials thatgeneralize to molecules and conformations beyond the training distribution.In Chapter 5, we present a novel approach to molecular conformer generation,which involves sampling low-energy conformations from a molecular graph.Current state-of-the-art methods rely on complex modelling techniques toachieve competitive results, either through expensive local structure methods[12] or large architecture models without task-specific inductive biases [14].54Chapter 6. ConclusionWe challenge this status quo by developing a simple and scalable method,ET-Flow, which uses a competent equivariant architecture TorchMD-NET[19], flow-matching [3] instead of diffusion and an informative moleculargraph specific prior to achieve state-of-the-art results on various molecularconformer generation benchmarks. Notably, our method uses significantlyfewer parameters and inference steps than previous methods. While recentstudies [14, 117] suggest that inductive biases are not essential when usinglarge datasets and high-parameter count neural networks, we hope our workmotivates the use of inductive biases and modelling approaches like FlowMatching to develop efficient generative models for molecular conformergeneration.55Bibliography[1] Alexandre Duval, Simon V Mathis, Chaitanya K Joshi, Victor Schmidt,Santiago Miret, Fragkiskos D Malliaros, Taco Cohen, Pietro Liò, YoshuaBengio, and Michael Bronstein. A hitchhikerâs guide to geometric gnnsfor 3d atomic systems. arXiv preprint arXiv:2312.07511, 2023.[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion proba-bilistic models. Advances in neural information processing systems, 33:6840â6851, 2020.[3] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel,and Matt Le. Flow matching for generative modeling. arXiv preprintarXiv:2210.02747, 2022.[4] W. Kabsch. A solution for the best rotation to relate two sets ofvectors. Acta Crystallographica Section A, 32(5):922â923, Sep 1976.doi: 10.1107/S0567739476001873. URL https://doi.org/10.1107/S0567739476001873.[5] Richard P Feynman, Robert B Leighton, Matthew Sands, and Everett MHafner. The feynman lectures on physics; vol. i. American Journal ofPhysics, 33(9):750â752, 1965.[6] Frank Noé, Gianni De Fabritiis, and Cecilia Clementi. Machine learningfor protein folding and dynamics. Current opinion in structural biology,60:77â84, 2020.[7] Julie R Schames, Richard H Henchman, Jay S Siegel, Christoph ASotriffer, Haihong Ni, and J Andrew McCammon. Discovery of a novelbinding trench in hiv integrase. Journal of medicinal chemistry, 47(8):1879â1881, 2004.[8] Jacob D Durrant, Henrik Keränen, Benjamin A Wilson, and J AndrewMcCammon. Computational identification of uncharacterized cruzainbinding sites. PLoS neglected tropical diseases, 4(5):e676, 2010.56Bibliography[9] Jun Zhang, Yi Isaac Yang, and Frank Noé. Targeted adversarial learningoptimized sampling. The journal of physical chemistry letters, 10(19):5791â5797, 2019.[10] Maho Nakata and Toshiyuki Maeda. Pubchemqc b3lyp/6-31g*//pm6dataset: the electronic structures of 86 million molecules using b3lyp/6-31g* calculations. arXiv preprint arXiv:2305.18454, 2023.[11] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusionimplicit models. arXiv preprint arXiv:2010.02502, 2020.[12] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, andTommi Jaakkola. Torsional diffusion for molecular conformer generation.Advances in Neural Information Processing Systems, 35:24240â24253,2022.[13] Greg Landrum et al. Rdkit: A software suite for cheminformatics,computational chemistry, and predictive modeling. Greg Landrum, 8(31.10):5281, 2013.[14] Yuyang Wang, Ahmed A. Elhag, Navdeep Jaitly, Joshua M. Susskind,and Miguel Angel Bautista. Swallowing the bitter pill: Simplifiedscalable conformer generation, 2024.[15] Peiye Zhuang, Samira Abnar, Jiatao Gu, Alex Schwing, Joshua MSusskind, and Miguel Angel Bautista. Diffusion probabilistic fields. InThe Eleventh International Conference on Learning Representations,2022.[16] Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Do-ersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran,Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architec-ture for structured inputs & outputs. arXiv preprint arXiv:2107.14795,2021.[17] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochas-tic interpolants: A unifying framework for flows and diffusions. arXivpreprint arXiv:2303.08797, 2023.[18] Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang,Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio.Conditional flow matching: Simulation-free dynamic optimal transport.arXiv preprint arXiv:2302.00482, 2(3), 2023.57Bibliography[19] Philipp Thölke and Gianni De Fabritiis. Torchmd-net: Equivarianttransformers for neural network based molecular potentials. arXivpreprint arXiv:2202.02541, 2022.[20] C David Sherrill. The born-oppenheimer approximation. School ofChemistry and Biochemistry, Georgia Institute of Technology, 242, 2005.[21] Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, AlexandreTkatchenko, and K-R Müller. Schnetâa deep learning architecturefor molecules and materials. The Journal of Chemical Physics, 148(24),2018.[22] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariantmessage passing for the prediction of tensorial properties and molecularspectra. In International Conference on Machine Learning, pages9377â9388. PMLR, 2021.[23] Simon Axelrod and Rafael Gomez-Bombarelli. Geom, energy-annotatedmolecular conformations for property prediction and molecular genera-tion. Scientific Data, 9(1):185, 2022.[24] Mark McGann. Fred pose prediction and virtual screening accuracy.Journal of chemical information and modeling, 51(3):578â596, 2011.[25] Ashutosh Kumar and Kam YJ Zhang. Advances in the developmentof shape similarity methods and their application in drug discovery.Frontiers in chemistry, 6:315, 2018.[26] Erwin Schrödinger. An undulatory theory of the mechanics of atomsand molecules. Physical review, 28(6):1049, 1926.[27] Walter Kohn, Axel D Becke, and Robert G Parr. Density functionaltheory of electronic structure. The journal of physical chemistry, 100(31):12974â12980, 1996.[28] Douglas R Hartree. The wave mechanics of an atom with a non-coulomb central field. part i. theory and methods. In MathematicalProceedings of the Cambridge Philosophical Society, volume 24, pages89â110. Cambridge university press, 1928.[29] John F Stanton and Rodney J Bartlett. The equation of motion coupled-cluster method. a systematic biorthogonal approach to molecular exci-tation energies, transition probabilities, and excited state properties.The Journal of chemical physics, 98(9):7029â7039, 1993.58Bibliography[30] Bogumil Jeziorski and Hendrik J Monkhorst. Coupled-cluster methodfor multideterminantal reference states. Physical Review A, 24(4):1668,1981.[31] Hendrik J Monkhorst. Calculation of properties with the coupled-cluster method. International Journal of Quantum Chemistry, 12(S11):421â432, 1977.[32] Axel D Becke. Density functionals for static, dynamical, and strongcorrelation. The Journal of Chemical Physics, 138(7), 2013.[33] RHWJ Ditchfield, Warren J Hehre, and John A Pople. Self-consistentmolecular-orbital methods. ix. an extended gaussian-type basis formolecular-orbital studies of organic molecules. The Journal of ChemicalPhysics, 54(2):724â728, 1971.[34] Christoph Bannwarth, Sebastian Ehlert, and Stefan Grimme. Gfn2-xtbâan accurate and broadly parametrized self-consistent tight-bindingquantum chemical method with multipole electrostatics and density-dependent dispersion contributions. Journal of Chemical Theoryand Computation, 15(3):1652â1671, 2019. doi: 10.1021/acs.jctc.8b01176. URL https://doi.org/10.1021/acs.jctc.8b01176. PMID:30741547.[35] Marcus Elstner, Th Frauenheim, E Kaxiras, G Seifert, and S Suhai. Aself-consistent charge density-functional based tight-binding scheme forlarge biomolecules. physica status solidi (b), 217(1):357â376, 2000.[36] Michael Gaus, Qiang Cui, and Marcus Elstner. Density functionaltight binding: application to organic and biological molecules. WileyInterdisciplinary Reviews: Computational Molecular Science, 4(1):49â61, 2014.[37] Anders S Christensen, Tomas Kubar, Qiang Cui, and Marcus Elstner.Semiempirical quantum mechanical methods for noncovalent interac-tions for chemical and biochemical applications. Chemical reviews, 116(9):5301â5337, 2016.[38] Junmei Wang, Romain M Wolf, James W Caldwell, Peter A Kollman,and David A Case. Development and testing of a general amber forcefield. Journal of computational chemistry, 25(9):1157â1174, 2004.59Bibliography[39] Kenno Vanommeslaeghe, Elizabeth Hatcher, Chayan Acharya, Sib-sankar Kundu, Shijun Zhong, Jihyun Shim, Eva Darian, Olgun Gu-vench, P Lopes, Igor Vorobyov, et al. Charmm general force field: Aforce field for drug-like molecules compatible with the charmm all-atomadditive biological force fields. Journal of computational chemistry, 31(4):671â690, 2010.[40] David Van Der Spoel, Erik Lindahl, Berk Hess, Gerrit Groenhof, Alan EMark, and Herman JC Berendsen. Gromacs: fast, flexible, and free.Journal of computational chemistry, 26(16):1701â1718, 2005.[41] Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet:Universal directional graph neural networks for molecules. Advances inNeural Information Processing Systems, 34:6790â6802, 2021.[42] Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, RafaelGomez-Bombarelli, and Tommi Jaakkola. Forces are not enough: Bench-mark and critical evaluation for machine learning force fields withmolecular simulations. arXiv preprint arXiv:2210.07237, 2022.[43] Marco Gori, Gabriele Monfardini, and Franco Scarselli. A new modelfor learning in graph domains. In Proceedings. 2005 IEEE internationaljoint conference on neural networks, 2005., volume 2, pages 729â734.IEEE, 2005.[44] Alessandro Sperduti and Antonina Starita. Supervised neural net-works for the classification of structures. IEEE transactions on neuralnetworks, 8(3):714â735, 1997.[45] Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner,and Gabriele Monfardini. The graph neural network model. IEEEtransactions on neural networks, 20(1):61â80, 2008.[46] Christoph Goller and Andreas Kuchler. Learning task-dependent dis-tributed representations by backpropagation through structure. InProceedings of international conference on neural networks (ICNNâ96),volume 1, pages 347â352. IEEE, 1996.[47] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals,and George E Dahl. Neural message passing for quantum chemistry. InInternational conference on machine learning, pages 1263â1272. PMLR,2017.60Bibliography[48] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti,David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational in-ductive biases, deep learning, and graph networks. arXiv preprintarXiv:1806.01261, 2018.[49] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar VeliÄkoviÄ.Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.arXiv preprint arXiv:2104.13478, 2021.[50] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer,Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedingsof the IEEE/CVF conference on computer vision and pattern recognition,pages 11976â11986, 2022.[51] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weis-senborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani,Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An imageis worth 16x16 words: Transformers for image recognition at scale.arXiv preprint arXiv:2010.11929, 2020.[52] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Michael C Mozer, andYoram Singer. Identity crisis: Memorization and generalization underextreme overparameterization. arXiv preprint arXiv:1902.04698, 2019.[53] Yang Song and Stefano Ermon. Generative modeling by estimatinggradients of the data distribution. Advances in neural informationprocessing systems, 32, 2019.[54] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight andfast: Learning to generate and transfer data with rectified flow. arXivpreprint arXiv:2209.03003, 2022.[55] Diederik P Kingma and Max Welling. Auto-encoding variational bayes.arXiv preprint arXiv:1312.6114, 2013.[56] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion modelswith exponential integrator. arXiv preprint arXiv:2204.13902, 2022.[57] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David KDuvenaud. Neural ordinary differential equations. Advances in neuralinformation processing systems, 31, 2018.61Bibliography[58] Michael S. Albergo and Eric Vanden-Eijnden. Building normalizingflows with stochastic interpolants, 2023.[59] Peter Eastman, Pavan Kumar Behara, David L Dotson, RaimondasGalvelis, John E Herr, Josh T Horton, Yuezhi Mao, John D Chodera,Benjamin P Pritchard, Yuanqing Wang, et al. Spice, a dataset of drug-like molecules and peptides for training machine learning potentials.Scientific Data, 10(1):11, 2023.[60] Justin S Smith, Olexandr Isayev, and Adrian E Roitberg. Ani-1, a dataset of 20 million calculated off-equilibrium conformations for organicmolecules. Scientific data, 4(1):1â8, 2017.[61] Justin S Smith, Roman Zubatyuk, Benjamin Nebgen, Nicholas Lub-bers, Kipton Barros, Adrian E Roitberg, Olexandr Isayev, and SergeiTretiak. The ani-1ccx and ani-1x data sets, coupled-cluster and densityfunctional theory properties for molecules. Scientific data, 7(1):134,2020.[62] Zhao Xu, Youzhi Luo, Xuan Zhang, Xinyi Xu, Yaochen Xie, MengLiu, Kaleb Dickerson, Cheng Deng, Maho Nakata, and Shuiwang Ji.Molecule3d: A benchmark for predicting 3d geometries from moleculargraphs. arXiv preprint arXiv:2110.01717, 2021.[63] Anders S Christensen, Sai Krishna Sirumalla, Zhuoran Qiao, Michael BOâConnor, Daniel GA Smith, Feizhi Ding, Peter J Bygrave, AnimashreeAnandkumar, Matthew Welborn, Frederick R Manby, et al. Orbnetdenali: A machine learning potential for biological and organic chemistrywith semi-empirical cost and dft accuracy. The Journal of ChemicalPhysics, 155(20), 2021.[64] Clemens Isert, Kenneth Atz, José Jiménez-Luna, and Gisbert Schneider.Qmugs, quantum mechanical properties of drug-like molecules. ScientificData, 9(1):273, 2022.[65] Maho Nakata and Tomomi Shimazaki. Pubchemqc project: a large-scalefirst-principles electronic structure database for data-driven chemistry.Journal of chemical information and modeling, 57(6):1300â1308, 2017.[66] Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n)equivariant graph neural networks. In International conference onmachine learning, pages 9323â9332. PMLR, 2021.62Bibliography[67] Guillem Simeon and Gianni De Fabritiis. Tensornet: Cartesian tensorrepresentations for efficient learning of molecular potentials. Advancesin Neural Information Processing Systems, 36, 2024.[68] Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng,Bin Shao, and Tie-Yan Liu. Se (3) equivariant graph neural networkswith complete local frames. In International Conference on MachineLearning, pages 5583â5608. PMLR, 2022.[69] Thorben Frank, Oliver Unke, and Klaus-Robert Müller. So3krates:Equivariant attention for interactions on arbitrary length-scales inmolecular systems. Advances in Neural Information Processing Systems,35:29400â29413, 2022.[70] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun,Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learninglocal equivariant representations for large-scale atomistic dynamics.Nature Communications, 14(1):579, 2023.[71] Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, andGábor Csányi. Mace: Higher order equivariant message passing neu-ral networks for fast and accurate force fields. Advances in NeuralInformation Processing Systems, 35:11423â11436, 2022.[72] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se(3)-transformers: 3d roto-translation equivariant attention networks.Advances in neural information processing systems, 33:1970â1981, 2020.[73] Yi-Lun Liao, Brandon Wood, Abhishek Das, and Tess Smidt.Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. arXiv preprint arXiv:2306.12059, 2023.[74] Saro Passaro and C Lawrence Zitnick. Reducing so (3) convolutionsto so (2) for efficient equivariant gnns. In International Conference onMachine Learning, pages 27420â27438. PMLR, 2023.[75] Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant:Covariant molecular neural networks. Advances in neural informationprocessing systems, 32, 2019.[76] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li,Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-andtranslation-equivariant neural networks for 3d point clouds. arXivpreprint arXiv:1802.08219, 2018.63Bibliography[77] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and JianTang. Geodiff: A geometric diffusion model for molecular conformationgeneration. arXiv preprint arXiv:2203.02923, 2022.[78] Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang.Learning neural generative dynamics for molecular conformation gener-ation. arXiv preprint arXiv:2102.10240, 2021.[79] Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradientfields for molecular conformation generation. In International conferenceon machine learning, pages 9558â9568. PMLR, 2021.[80] Octavian Ganea, Lagnajit Pattanaik, Connor Coley, Regina Barzilay,Klavs Jensen, William Green, and Tommi Jaakkola. Geomol: Torsionalgeometric generation of molecular 3d conformer ensembles. Advancesin Neural Information Processing Systems, 34:13757â13769, 2021.[81] Yuyang Wang, Ahmed A Elhag, Navdeep Jaitly, Joshua M Susskind,and Miguel Angel Bautista. Generating molecular conformer fields.arXiv preprint arXiv:2311.17932, 2023.[82] Gregor NC Simm and José Miguel Hernández-Lobato. A gen-erative model for molecular distance geometry. arXiv preprintarXiv:1909.11459, 2019.[83] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density esti-mation using real nvp. arXiv preprint arXiv:1605.08803, 2016.[84] Leo Liberti, Carlile Lavor, Nelson Maculan, and Antonio Mucherino.Euclidean distance geometry and applications. SIAM review, 56(1):3â69, 2014.[85] Shitong Luo, Chence Shi, Minkai Xu, and Jian Tang. Predictingmolecular conformation via dynamic graph score matching. Advancesin Neural Information Processing Systems, 34:19784â19795, 2021.[86] Pascal Vincent. A connection between score matching and denoisingautoencoders. Neural computation, 23(7):1661â1674, 2011.[87] Sereina Riniker and Gregory A Landrum. Better informed distancegeometry: using what we know to improve conformation generation.Journal of chemical information and modeling, 55(12):2562â2574, 2015.64Bibliography[88] Jörg Behler and Michele Parrinello. Generalized neural-network repre-sentation of high-dimensional potential-energy surfaces. Physical reviewletters, 98(14):146401, 2007.[89] Alireza Khorshidi and Andrew A Peterson. Amp: A modular ap-proach to machine learning in atomistic simulations. Computer PhysicsCommunications, 207:310â324, 2016.[90] Oliver T Unke and Markus Meuwly. A reactive, scalable, and trans-ferable model for molecular energies from a neural network approachbased on local information. The Journal of chemical physics, 148(24),2018.[91] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers forneural network based molecular potentials. In International Conferenceon Learning Representations, 2021.[92] Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice EA Allen,Daniel J Cole, Christoph Ortner, and Gábor Csányi. Linear atomic clus-ter expansion force fields for organic molecules: beyond rmse. Journalof chemical theory and computation, 17(12):7696â7711, 2021.[93] Nongnuch Artrith, Alexander Urban, and Gerbrand Ceder. Efficientand accurate machine-learning interpolation of atomic energies in com-positions with many species. Physical Review B, 96(1):014112, 2017.[94] So Takamoto, Chikashi Shinagawa, Daisuke Motoki, Kosuke Nakago,Wenwen Li, Iori Kurata, Taku Watanabe, Yoshihiro Yayama, HirokiIriguchi, Yusuke Asano, et al. Towards universal neural network po-tential for material discovery applicable to arbitrary combination of 45elements. Nature Communications, 13(1):2991, 2022.[95] Johannes Gasteiger, Janek GroÃ, and Stephan Günnemann. Directionalmessage passing for molecular graphs. arXiv preprint arXiv:2003.03123,2020.[96] Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin,and Shuiwang Ji. Spherical message passing for 3d molecular graphs.In International Conference on Learning Representations, 2021.[97] Chi Chen and Shyue Ping Ong. A universal graph deep learninginteratomic potential for the periodic table. Nature ComputationalScience, 2(11):718â728, 2022.65Bibliography[98] Zhuoran Qiao, Anders S Christensen, Matthew Welborn, Frederick RManby, Anima Anandkumar, and Thomas F Miller III. Unite: Uni-tary n-body tensor equivariant network with applications to quantumchemistry. arXiv preprint arXiv:2105.14655, 3, 2021.[99] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Ana-tole Von Lilienfeld. Quantum chemistry structures and properties of134 kilo molecules. Scientific data, 1(1):1â7, 2014.[100] Maho Nakata, Tomomi Shimazaki, Masatomo Hashimoto, andToshiyuki Maeda. Pubchemqc pm6: Data sets of 221 million moleculeswith optimized molecular geometries and electronic properties. Journalof Chemical Information and Modeling, 60(12):5891â5899, 2020.[101] Johannes Hoja, Leonardo Medrano Sandonas, Brian G Ernst, AlvaroVazquez-Mayagoitia, Robert A DiStasio Jr, and Alexandre Tkatchenko.Qm7-x, a comprehensive dataset of quantum-mechanical propertiesspanning the chemical space of small organic molecules. Scientific data,8(1):43, 2021.[102] Alexander G Donchev, Andrew G Taube, Elizabeth Decolvenaere, CoryHargus, Robert T McGibbon, Ka-Hei Law, Brent A Gregersen, Je-LuenLi, Kim Palmo, Karthik Siva, et al. Quantum chemical benchmarkdatabases of gold-standard dimer interaction energies. Scientific data,8(1):55, 2021.[103] Ilyes Batatia, Simon Batzner, David Peter Kovacs, Albert Musaelian,Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky,and Gabor Csanyi. The design space of e3-equivariant atom-centeredinteratomic potentials. arXiv preprint arXiv:2205.06643, 2022.[104] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attentiontransformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990,2022.[105] Rajarshi Guha and Darrell Velegol. Harnessing shannon entropy-baseddescriptors in machine learning models to enhance the prediction ac-curacy of molecular properties. Journal of Cheminformatics, 15(1):54,2023.[106] Yang Zhang and Jeffrey Skolnick. Scoring function for automatedassessment of protein structure template quality. Proteins: Structure,Function, and Bioinformatics, 57(4):702â710, 2004.66Bibliography[107] Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, BonnieBerger, and Tommi Jaakkola. Eigenfold: Generative protein structureprediction with diffusion models. arXiv preprint arXiv:2304.02198,2023.[108] Hannes Stark, Bowen Jing, Regina Barzilay, and Tommi Jaakkola.Harmonic prior self-conditioned flow matching for multi-ligand dockingand binding site design. In NeurIPS 2023 AI for Science Workshop,2023.[109] Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow match-ing. Advances in Neural Information Processing Systems, 36, 2024.[110] Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meetsflow matching for generating protein ensembles, 2024.[111] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidatingthe design space of diffusion-based generative models, 2022.[112] Oliver T Unke and Markus Meuwly. Physnet: A neural network forpredicting energies, forces, dipole moments, and partial charges. Journalof chemical theory and computation, 15(6):3678â3693, 2019.[113] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layernormalization. arXiv preprint arXiv:1607.06450, 2016.[114] Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski,Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron,Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transform-ers to 22 billion parameters. In International Conference on MachineLearning, pages 7480â7512. PMLR, 2023.[115] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari,Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer,Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint arXiv:2403.03206, 2024.[116] Simon Axelrod and Rafael Gomez-Bombarelli. Molecular machinelearning with conformer ensembles. Machine Learning: Science andTechnology, 4(3):035025, 2023.[117] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green,Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J67Ballard, Joshua Bambrick, et al. Accurate structure prediction ofbiomolecular interactions with alphafold 3. Nature, pages 1â3, 2024.[118] Hadrien Mary, Emmanuel Noutahi, DomInvivo, Michel Moreau, Lu Zhu,Steven Pak, Desmond Gilmour, t, Valence-JonnyHsu, Honoré Houn-wanou, Ishan Kumar, Saurav Maheshkar, Shuya Nakata, Kyle M.Kovary, Cas Wognum, Michael Craig, and DeepSource Bot. datamol-io/datamol: 0.11.4, September 2023. URL https://doi.org/10.5281/zenodo.8357317.[119] Philipp Pracht, Stefan Grimme, Christoph Bannwarth, Fabian Bohle,Sebastian Ehlert, Gereon Feldmann, Johannes Gorges, Marcel Müller,Tim Neudecker, Christoph Plett, et al. Crestâa program for theexploration of low-energy molecular chemical space. The Journal ofChemical Physics, 160(11), 2024.68Appendix AAppendixA.1 Structural differences between GEOM-Drugsand GEOM-QM9 DistributionTo illustrate the structural differences between the drug-like molecules fromGEOM-Drugs and the small molecules from GEOM-QM9, we create finger-prints for each molecule using the fingerprint function from the datamollibrary [118]. Subsequently, we extracted two principal components fromthese fingerprints using Principal Component Analysis (PCA). The resultingprincipal components were then plotted, revealing a noticeable separationbetween clusters representing GEOM-Drugs and GEOM-QM9 molecules.A.2 ET-Flow Implementation DetailsA.2.1 Input FeaturizationAtomic features (or Node Features) are computed using RDKit [13] featuresas described in Table A.1. For computing edge features and edge index, weuse a combination of global (radius based edges) and local (molecular graphedges) similar to [12].Name Description Rangechirality Chirality Tag {unspecified, tetrahedral CW, tetrahedral CCW, other}degree Number of bonded neighbors {x : 0 â¤ x â¤ 10, x â Z}charge Formal charge of atom {x : â5 â¤ x â¤ 5, x â Z}num_H Total Number of Hydrogens {x : 0 â¤ x â¤ 8, x â Z}number_radical_e Number of Radical Electrons {x : 0 â¤ x â¤ 4, x â Z}hybrization Hybrization type {sp, sp2, sp3, sp3d, sp3d2, other}aromatic Whether on a aromatic ring {True, False}in_ring Whether in a ring {True, False}Table A.1: Atomic features included in ET-Flow.69A.2. ET-Flow Implementation DetailsFigure A.1: Structural differences between GEOM-Drugs (IID-S set) andGEOM-QM9 (OOD-S set) are evident from the distinct separation betweenthe two clusters. Each point in the plot represents a molecule.A.2.2 Evaluation MetricsFollowing the approaches of [12, 77, 80], we utilize Average Minimum RMSD(AMR) and Coverage (COV) to assess the performance of molecular conformergeneration. Here, Cg denotes the set of generated conformations, and Crdenotes the set of reference conformations. For both AMR and COV, wecalculate and report Recall (R) and Precision (P). Recall measures the extentto which the generated conformers capture the ground-truth conformers,while Precision indicates the proportion of generated conformers that areaccurate. The specific formulations for these metrics are detailed in thefollowing equations:AMR-R(Cg, Cr) =1|Cr|âRâCrminRËâCgRMSD(R, RË)COV-R(Cg, Cr) =1|Cr| |{R â Cr|RMSD(R, RË) < δ, RË â Cg}|AMR-P(Cr, Cg) =1|Cg|âRËâCgminRâCrRMSD(RË,R)70A.3. Training and Sampling AlgorithmCOV-P(Cr, Cg) =1|Cg| |{RË â Cg|RMSD(RË,R) < δ,R â Cr}|A lower AMR score signifies improved accuracy, while a higher COV scorereflects greater diversity in the generative model. Following [12], the thresholdδ is set to 0.5 for GEOM-QM9 and 0.75 for GEOM-DRUGS.A.2.3 Training Details and HyperparametersHyper-parameter ET-Flownum_layers 20hidden_channels 160num_heads 8neighbor_embedding Truecutoff_lower 0.0cutoff_higher 10.0node_attr_dim 8edge_attr_dim 1reduce_op Trueactivation SiLUattn_activation SiLU# param 8.3MTable A.2: Hyperparameters for ET-FlowFor GEOM-DRUGS, we train ET-Flow for a fixed 250 epochs with a batchsize of 64 and 5000 training batches per epoch per GPU on 8 A100 GPUs.For the learning rate, we use the Adam Optimizer with a cosine annealinglearning rate which goes from a maximum of 10â3 to a minimum 10â7 over250 epochs with a weight decay of 10â10. For GEOM-QM9, we train ET-Flowfor 200 epochs with a batch size of 128, and use all of the training dataset perepoch on 4 A100 GPUs. We use the cosine annealing learning rate schedulewith maximum of 8 · 10â4 to minimum of 10â7 over 100 epochs, post whichthe maximum is reduced by a factor of 0.05. We select checkpoints based onthe lowest validation error.A.3 Training and Sampling AlgorithmThe following algorithms go over the pseudo-code for the training and samplingprocedure. For each molecule, we use up to 30 conformations with the highestBoltzmann weights as provided by CREST [119] similar to that of [12]71A.4. VisualizationsAlgorithm 1: Training procedureInput: molecules [G0, ..., GN ] each with true conformers [CG,1, ...CG,KG ],the harmonic prior Ï0, learning rate α, number of epochs Ne,initialized vector field vθOutput: trained flow matching model vθfor iâ 1 to Ne dofor G in [G0, ..., GN ] doSample t â¼ U [0, 1] and C1 â [CG,1, ...CG,KG ];Sample prior C0 â¼ Ï0(G);Align C0 â RMSDAlign(C0, C1);Sample Ct â¼ N(tC1 + (1â t)C0, Ï2t(1â t));Construct vector field ut â x1 â x0 + 1â2t2ât(1ât)z, z â¼ N (0, I);Compute loss L â â¥vÎ¸(t, Ct)â utâ¥2;Take gradient step θ â θ â Î±âÎ¸L;A.4 VisualizationsFigure A.2 shows randomly selected examples of sampled conformers fromET-Flow for GEOM-DRUGS. The left column is the reference molecule fromthe ground truth, and the remaining columns are samples generated with 50sampling steps. Figure A.3 showcases the ability for ET-Flow to generatequality samples with fewer sampling steps.72A.4. VisualizationsAlgorithm 2: Inference procedureInput: molecular graph G, number conformers K, number of samplingsteps NOutput: predicted conformers [C1, ...CK ]for C in [C1, ...CK ] dosample prior CË â¼ Ï0(G);for n â 0 to N â 1 doSet tâ nN ;Set âtâ 1N ;Predict vË = vθ(t, CË);Update CË = CË + vËâtAlgorithm 3: Stochastic SamplerInput: molecular graph G, number conformers K, number of samplingsteps N , stochasticity level churn, stochastic sampling range[tmin, tmax]Output: predicted conformers [C1, ...CK ]for C in [C1, ...CK ] dosample prior CË â¼ Ï0(G);for n â 0 to N â 1 doSet tâ nN ;Set âtâ 1N ;Set γ â churnN ;if t â [tmin, tmax] thenSample ϵ â¼ N(0, I);âtËâ Î³(1â t);tËâ max(tââtË, 0);CË â CË +âtËât2 â tË2Ïµ;Predict vË = vÎ¸(tË, CË);Set âtâ ât+âtË;elsePredict vË = vθ(t, CË);Update CË = CË + vËât73A.4. VisualizationsFigure A.2: Examples of conformers of ground truth and ET-Flow for GEOM-DRUGS.74A.4. VisualizationsFigure A.3: Examples of conformers of ground truth and ET-Flow for differentnumber of sampling steps.75
Thesis/Dissertation
eng
Computer Science
Vancouver : University of British Columbia Library
Attribution-NonCommercial-NoDerivatives 4.0 International
Graduate
Ding, Jiarui
Investigating ML potentials and deep generative models for efficient conformational sampling
Text