Generating predicted 3D structures of molecular systems using generative models with equivariant attention layers

Abstract

Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating predicted 3D structures of molecular systems using a generative machine learning model. In one aspect, a method comprises: generating a predicted 3D structure of a molecular system using a generative machine learning model by iteratively updating a current predicted 3D structure of the molecular system, comprising, at each of a plurality of time steps: processing data characterizing a current predicted 3D structure at the time step to update the current predicted 3D structure by (i) generating atom embeddings for the molecular system based using the current predicted 3D structure, (ii) updating the atom embeddings using an equivariant attention operation, and (iii) processing the updated atom embeddings to generate an updated predicted 3D structure for the molecular system.

Description

GENERATING PREDICTED 3D STRUCTURES OF MOLECULAR SYSTEMS USING GENERATIVE MODELS WITH EQUIVARIANT ATTENTION LAYERSCROSS-REFERENCE TO RELATED APPLICATIONS

[0001] This application claims the benefit of priority to U. S. Provisional Application No.63 / 715,278, filed on November 1, 2024, the contents of which are hereby incorporated by reference.BACKGROUND

[0002] This specification relates to processing data using machine learning models.

[0003] Machine learning models receive an input and generate an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate the output based on the received input and on values of the parameters of the model.

[0004] Some machine learning models are deep models that employ multiple layers of models to generate an output for a received input. For example, a deep neural network is a deep machine learning model that includes an output layer and one or more hidden layers that each apply a nonlinear transformation to a received input to generate an output.SUMMARY

[0005] This specification generally describes a system implemented as computer programs on one or more computers in one or more locations that can generate predicted 3D structures of molecular systems using a generative machine learning model.

[0006] According to on aspect, there is provided a method performed by one or more computers, comprising: obtaining data identifying a molecular system one or more molecules; processing the data identifying the molecular system using a generative machine learning model parametrized by a transformation neural network to generate data defining a predicted three-dimensional (3D) structure of the molecular system by iteratively updating data defining a current predicted 3D structure of the molecular system over a plurality of time steps, comprising, at each time step: obtaining data defining a current predicted 3D structure at the time step, processing a network input based on the current predicted 3D structure at the time step using the transformation neuralnetwork to generate a transformation output defining a modification to the current predicted 3D structure by (i) generating a respective atom embedding for each atom in the molecular system based at least in part on the current predicted 3D structure, (ii) updating the atom embeddings using an equivariant attention operation, and (iii) processing the updated atom embeddings to generate the transformation output; and updating the current predicted 3D structure based on the transformation output generated by the transformation neural network; wherein the predicted 3D structure of the molecular system is defined by the current predicted 3D structure after a last time step of the plurality of time steps.

[0007] In some implementations, the equivariant attention operation is equivariant to rotations of the current predicted 3D structure.

[0008] In some implementations, generating the respective atom embedding for each atom in the molecular system based at least in part on the current predicted 3D structure comprises, for each atom: generating a spatially-independent embedding of the atom based on properties of the atom that are invariant to a spatial position of the atom; and generating a spatially-dependent embedding of the atom based on at least a 3D spatial position of the atom; wherein the spatially-independent embedding and the spatially-dependent embedding collectively define the embedding of the atom.

[0009] In some implementations, for each atom, the spatially-independent embedding of the atom is based on one or more of: an elemental type of the atom, a charge of the atom, an atomic mass of the atom, an electron configuration of the atom, an electronegativity of the atom, an ionization energy of the atom, or an electrostatic potential of the atom.

[0010] In some implementations, for each atom, the spatially-dependent embedding of the atom comprises three spatial dimensions and a plurality of channel dimensions.

[0011] In some implementations, updating the atom embeddings using an equivariant attention operation comprises: generating, for each atom, respective query, key, and value embeddings representing the spatially-independent embedding of the atom; generating, for each atom, respective query, key, and value embeddings representing the spatially-dependent embedding of the atom using equivariant linear neural network layer operations, wherein the query, key, and value embeddings representing the spatially-dependent embedding of the atom each comprise three spatial dimensions and a plurality of channel dimensions; and updating the atom embeddings using the equivariant attention operation based on the query, key, and value embeddings representing the spatially-independent and spatially-dependent embeddings of the atoms.

[0012] In some implementations, updating the atom embeddings using the equivariant attention operation based on the query, key, and value embeddings representing the spatially-independent and spatially-dependent embeddings of the atoms comprises: generating an array of attention weights, wherein each attention weight corresponds to a respective pair of atoms and is generated by performing operations comprising: determining a spatially-independent attention weight based on a scalar inner product between: (i) a query embedding representing a spatially-independent embedding of a first atom in the pair of atoms, and (ii) a key embedding representing a spatially-independent embedding of a second atom in the pair of atoms, determining a spatially-dependent attention weight based on a Frobenius inner product between: (i) a query embedding representing a spatially-dependent embedding of the first atom in the pair of atoms, and (ii) a key embedding representing a spatially-dependent embedding of the second atom in the pair of atoms, and generating the attention weight for the pair of atoms by combining the spatially-independent attention weight and the spatially-dependent attention weight; and updating the atom embeddings by combining the value embeddings representing the spatially-independent embeddings and spatially-dependent embeddings of the atoms using the array of attention weights.

[0013] In some implementations, the method further comprises, after updating the atom embeddings using the equivariant attention operation, updating the atom embeddings using an equivariant gated feedforward neural network layer.

[0014] In some implementations, updating the atom embeddings using an equivariant gated feedforward neural network layer comprises, for each atom: generating an intermediate spatially-dependent embedding by combining: (i) the spatially-dependent embedding of the atom, and (ii) a result of processing the spatially-independent embedding of the atom using a feedforward neural network layer; and replacing the spatially-dependent embedding by a result of processing the intermediate spatially-dependent embedding by an equivariant linear neural network layer.

[0015] In some implementations, the method further comprises: generating a respective token embedding for each token in the molecular system based at least in part on the current predicted 3D structure, wherein one or more of the tokens represent: amino acids, nucleic acids, functional groups, or multi-atom fragments; updating the token embeddings using the equivariant attention operation; and processing both the updated atom embeddings and the updated token embeddings to generate the transformation output.

[0016] In some implementations, the equivariant attention operation over the token embeddings is a global attention operation and the equivariant attention operation over the atom embeddings is a local self-attention operation.

[0017] In some implementations, processing both the updated atom embeddings and the updated token embeddings to generate the transformation output comprises: updating the atom embeddings using an equivariant cross-attention operation over the token embeddings.

[0018] In some implementations, processing the updated atom embeddings to generate the transformation output comprises: processing only the spatially-dependent embeddings of the atoms to generate, for each atom, an output defining an update to a 3D spatial position of the atom.

[0019] In some implementations, the method further comprises processing only the spatially-independent embeddings of the atoms to generate, for each atom, a confidence of the generative machine learning model in a 3D spatial position of the atom the updated predicted 3D structure.

[0020] In some implementations, the generative machine learning model is a generative diffusion model, or a flow matching model, or a stochastic interpolant model.

[0021] In some implementations, outputting the predicted 3D structure of the molecular system comprises: storing the predicted 3D structure in a memory; or transmitting the predicted 3D structure over a data communications network; or displaying a visual rendering of the predicted 3D structure on a display of a user device.

[0022] In some implementations, the method further comprises selecting one or more molecules in the molecular system for physical synthesis based at least in part on the predicted 3D structure of the molecular system.

[0023] In some implementations, the method further comprises physically synthesizing the one or more molecules selected for physical synthesis.

[0024] According to another aspect, there is provided a method performed by one or more computers, comprising: obtaining data defining a protein and a set of candidate ligands; generating, for each candidate ligand in the set of candidate ligands, a predicted 3D structure of a molecular system comprising the protein and the candidate ligand by performing the operations of the method described above; and determining, for each candidate ligand in the set of candidate ligands, a predicted binding affinity of the candidate ligand for the protein based at least in part on the predicted 3D structure of the molecular system comprising the protein and the candidate ligand.

[0025] In some implementations, the method further comprises selecting one or more of the candidate ligands for targeting the protein based at least in part on the predicted binding affinities of the candidate ligands for the protein.

[0026] In some implementations, one or more candidate ligands for target the protein comprises selecting one or more candidate ligands for inclusion in a drug targeting the protein.

[0027] In some implementations, the method further comprises selecting one or more of the candidate ligands for physical synthesis based at least in part on the predicted binding affinities of the candidate ligands for the protein.

[0028] In some implementations, the method further comprises physically synthesizing the one or more candidate ligands selected for physical synthesis.

[0029] The subject matter described in this specification can be implemented in particular embodiments so as to realize one or more of the following advantages.

[0030] The system described in this specification uses a generative machine learning model to generate predicted 3D structures of molecular systems. More specifically, the generative machine learning model can generate a predicted 3D structure of a molecular system by refining the predicted 3D structure over a sequence of time steps using a neural network that is referred to for convenience in this specification as a “transformation” neural network. At each time step, the generative machine learning model can process the current predicted 3D structure using the transformation neural network to generate a transformation output that defines a modification to the current predicted 3D structure. The system can then update the current predicted 3D structure based on the transformation output generated by the transformation neural network.

[0031] The transformation neural network can implement equivariant attention operations to update embeddings representing atoms or “tokens” (e.g., multi-atom structures such as amino acids or nucleic acids) in the molecular system. More specifically, the system can generate atom / token embeddings that include: (i) a spatially-independent part that is based on properties of atoms / tokens that are invariant to spatial position, and (ii) a spatially-dependent part that is based at least on 3D spatial positions of atoms / tokens. The system can construct the attention weights using spatially-dependent attention weights (based on the spatially-dependent embeddings) and spatially-independent attention weights (based on the spatially-independent embeddings) that are computed using equivariant operations such as Frobenius inner products. The equivariance of the neural network operations has the effect that applying a rotation to apredicted 3D structure provided as an input to the transformation neural network will produce an equivalent rotation in the transformation neural network output.

[0032] The equivariance of the attention operations implemented by the transformation neural network encodes a powerful inductive bias in the generative machine learning model and enables reduced consumption of computational resources by the generative machine learning model during training and inference. More specifically, the equivariance property of the transformation neural network mitigates the risk of overfitting the training data and reduces the amount of training data and the number of training iterations required for the generative machine learning model to achieve an acceptable prediction accuracy. Training a conventional, non-equivariant generative machine learning model to generate equivariant predictions often requires using augmented training data that includes multiple different augmentations (e.g., translations, rotations, etc.) of training examples for the generative machine learning model to avoid overfitting. The architecturally ensured equivariance property of the transformation neural network enables the equivariant generative machine learning models described in this specification to be trained without using such augmented training data, which can reduce the amount of training data and the training costs required to train the described systems.

[0033] Additionally, the equivariance property of the transformation neural network reflects physical symmetries that occur in physical systems. The equivariance property of the transformation neural network can therefore improve a physical consistency, accuracy, and generalization of the predicted 3D structures generated using the transformation neural network. By improving the physical consistency, accuracy, and generalization of the predictions generated using the transformation neural network, the equivariance property of the transformation neural network can therefore also enable the transformation neural network to be implemented using a less complex neural network architecture, e.g., with fewer neural network parameters and fewer neural network layers, thus reducing consumption of computational resources (e.g., memory and computing power) during inference.

[0034] The transformation neural network can perform attention operations over a set of atom embeddings. However, the number of atoms in complex molecular systems can be very large, e.g., on the order of thousands, and the computational complexity of a global attention operation can increase quadratically with the number of embeddings participating in the attention operation. Therefore applying a global attention operation over the set of atom embeddings maybe computationally infeasible. However, applying only a local attention operation over the set of atom embeddings may prevent the transformation neural network from learning to recognize and exploit long-range, global interactions among atoms in the molecular system. To address this issue, the transformation neural network can implement a global attention operation over a set of token embeddings (which may be far fewer in number, e g., by one or more orders of magnitude, than the atom embeddings) and a local attention operation over the atom embeddings. The system can then enrich the atom embeddings with a cross-attention operation over the token embeddings. In this manner the system can significantly reduce consumption of computational resources during training and inference while maintaining the predictive accuracy of the model.

[0035] The details of one or more embodiments of the subject matter of this specification are set forth in the accompanying drawings and the description below. Other features, aspects, and advantages of the subject matter will become apparent from the description, the drawings, and the claims.BRIEF DESCRIPTION OF THE DRAWINGS

[0036] FIG. 1 illustrates an example equivariant generative machine learning model.

[0037] FIG. 2 illustrates an example equivariant transformation neural network.

[0038] FIG. 3 illustrates an example equivariant Transformer layer.

[0039] FIG. 4 illustrates an example equivariant gated feed-forward layer.

[0040] FIG. 5 illustrates a torsion angle between two groups of atoms for a molecule system.

[0041] FIG. 6 is a flow diagram of an example process for generating a predicted 3D structure for a molecule system using an equivariant generative machine learning model.

[0042] FIG. 7 is a flow diagram of an example process for training an equivariant generative machine learning model.

[0043] FIG. 8 is a flow diagram of an example process for screening a set of candidate ligands.

[0044] Like reference numbers and designations in the various drawings indicate like elements.DETAILED DESCRIPTION

[0045] FIG. 1 illustrates an example equivariant generative machine learning model 100. The equivariant generative machine learning model 100 is an example of a system implemented ascomputer programs on one or more computers in one or more locations in which the systems, components, and techniques described below are implemented.

[0046] The equivariant generative machine learning model 100 can be configured to process initial molecule data 102 for a molecule system to generate data characterizing a predicted 3D structure 104 for the molecule system.

[0047] The equivariant generative machine learning model 100 can be configured to process initial molecule data 102 and generate predicted 3D structures 104 for any of a variety of molecule systems that include one or more molecules. For example, the equivariant generative machine learning model 100 can process initial molecule data 102 and generate predicted 3D structures 104 for molecule systems that include any combination of, e.g., proteins, ligands, nucleic acids, molecule complexes such as protein-ligand complexes, and so on.

[0048] In general, the initial molecule data 102 for the example molecule system can include data in any appropriate format for representing the chemical composition of the molecule system. For example, the initial molecule data 102 for the molecule system can include data specifying atoms, charges on atoms, bonds, functional groups, amino acids, branches (e.g., from a molecule main chain), ring structures, and so forth for the molecules within the molecule system. As a further example, the initial molecule data 102 for an example molecule system can include, e.g., Simplified Molecular Input Line Entry System (SMILES) formatted data, International Chemical Identifier (InChi) formatted data, and so on specifying chemical structures of molecules and / or portions of molecules within the molecule system.

[0049] As another example, the initial molecule data 102 for the molecule system can include data specifying one or more amino acid sequences for the molecule system. Each amino acid sequence for the molecule system can specify an identity and an ordering for each of a plurality of amino acid residues within a molecule (e.g., a protein, a ligand, etc.) of the molecule system.

[0050] Similarly, the initial molecule data 102 for the molecule system can include data specifying one or more nucleic acid sequences for the molecule system. Each nucleic acid sequence for the molecule system can specify an identity and an ordering for each of a plurality of nucleotides within a molecule (e.g., a nucleic acid, a ligand, etc.) of the molecule system.

[0051] The initial molecule data 102 for the molecule system can include data characterizing initial noisy 3D structures for one or more molecules of the molecule system. The data characterizing the initial noisy 3D structures for the molecules of the molecule system can specify chemicalcompositions and 3D structures for the molecules by specifying a chemical element and an initial spatial position for each atom of the molecules within the molecule system. The equivariant generative machine learning model 100 can obtain the data characterizing the initial noisy 3D structures for the molecules of the molecule system by any of a variety of methods.

[0052] For example, in some cases, the equivariant generative machine learning model 100 can generate the initial noisy 3D structure for the molecule systems by sampling the initial noisy 3D structure for the molecule system from a distribution over a space of possible 3D structures for the molecule system. In particular, the equivariant generative machine learning model 102 can generate the initial noisy 3D structure for the molecule system by sampling an initial noisy spatial position for each atom of the example molecule system from a distribution of noisy spatial positions for the atom.

[0053] As another example, in some cases, the equivariant generative machine learning model 100 can generate the initial noisy 3D structure for the molecule system to represent an estimated or predicted 3D structures of the molecules within the molecule system. The estimated or predicted 3D structures for molecules within the molecule system can be determined by any of a variety of means, such as by physical experiment (e.g., x-ray crystallography), by computational simulation (e.g., by molecular dynamics simulations, quantum mechanical simulations, etc.), by machine learning techniques (e.g., using a molecule structure prediction machine learning model) and so on.

[0054] The equivariant generative machine learning model can generate the predicted 3D structure for the molecule system by iteratively updating (e.g. refining) the initial molecule data 102 for the molecule system. In particular, the equivariant generative machine learning model 100 can use an equivariant transformation neural network 106 to process the initial molecule data 102 for the molecule system over a sequence of time steps to generate the predicted 3D structure 104 for the molecule system.

[0055] At each time step of generating the predicted 3D structure 104 for the molecule system, the equivariant generative machine learning model 100 can process a network input that includes data representing a current 3D structure for the molecule system as of the time step using the equivariant transformation neural network 106 to generate a network output that characterizes a transformation of the current 3D structure for the molecule system for the time step. The equivariant generative machine learning model 100 can combine the network output from the transformation neuralnetwork 106 with the data representing the current predicted 3D structure for the molecule system to generate updated molecule data 108 that represents an updated 3D structure of the molecule system for the time step, which the equivariant generative machine learning model 100 can then provide as a network input to the equivariant transformation neural network 106 during a next time step of generating the predicted 3D structure 104 for the molecule system.

[0056] The equivariant generative machine learning model 100 can use the final updated molecule data 108 generated at a last time step to determine the predicted 3D structure 104 for the molecule system. By stochastically generating the predicted 3D structure 104 (e.g., by sampling the initial molecule data 102 from an initial distribution, by sampling the updated molecule data 108 from a distribution defined by the network output of the equivariant transformation neural network 106, etc.), the equivariant generative machine learning model 100 can therefore sample the predicted molecule structure 104 from a distribution over a space of possible 3D structures for the molecule system.

[0057] The equivariant generative machine learning model 100 can have any appropriate generative machine learning architecture for generating the predicted 3D structure for the molecule system by iteratively transforming (e.g., denoising, diffusing, etc.) data characterizing the predicted 3D structure 104 for the molecule system using the equivariant transformation neural network 106. For example, the equivariant generative machine learning model 100 can be implemented as, e.g., a generative diffusion model, a generative flow matching model, a stochastic interpolant model, and so on. As a further example, the equivariant generative machine learning model 100 can utilize a similar generative architecture to those described by, e.g., Abramson et al. in “Accurate Structure Prediction of Biomolecular Interactions with AlphaFold 3”, Watson et al. in “De Novo Design of Protein Structure and Function with RF diffusion”, and so on.

[0058] As an example, the equivariant generative machine learning model 100 can be a generative diffusion model configured to generate the predicted 3D structure 104 following a reverse diffusion process. When the equivariant generative machine learning model 100 is a generative diffusion model, the equivariant generative machine learning model 100 can be configured to generate the updated molecule data 108 at each time step based on the network output from the equivariant transformation neural network 106 by sampling (e.g., using a diffusion sampling technique) the updated molecule data 108 from a distribution defined by the network output from the equivariant transformation neural network 106 for the time step.

[0059] As another example, the equivariant generative machine learning model 100 can be a generative flow matching model or a stochastic interpolant model that is configured to generate the predicted 3D structure 104 in accordance with a set of transformation differential equations that characterize a transformation from a distribution (e.g., a noise distribution) for the initial molecule data 102 to a distribution for the final updated molecule data 108 representing the predicted 3D structure 104 for the molecule system. When the equivariant generative machine learning model 100 is a generative flow matching model or a stochastic interpolant model, the equivariant generative machine learning model 100 can be configured to generate the updated molecule data 108 at each time step based on the network output from the equivariant transformation neural network 106 for the time step by generating the updated molecule data 108 in accordance with a set of transformation differential equations that are defined by defined by the network output from the equivariant transformation neural network 106 for the time step (e g., by numerically integrating the set of transformation differential equations).

[0060] The network inputs to the equivariant transformation neural network 106 can represent (e.g., encode) a 3D structure for the molecule system using any appropriate scheme for encoding the positions of atoms within the molecule system and the equivariant transformation neural network 110 can have a neural network architecture suitable for processing (e.g., denoising, transforming, etc.) the network inputs representing 3D structures for the molecule system. The updated molecule data 108 can represent (e.g., encode) 3D structures for the molecule system using the same encoding scheme as used by network inputs to the equivariant transformation neural network 106. Example neural network architectures of the equivariant transformation neural network 106 are described in more detail below with reference to FIGs. 2-5.

[0061] In some implementations, the equivariant generative machine learning model 100 can include a structure decoding neural network 110 that is configured to determine the 3D positions of atoms within the molecule system specified by updated molecule data 108 for the molecule system. For example, the equivariant generative machine learning model can process the final updated molecule data 108 using 100 the structure decoding neural network 110 to generate the predicted 3D structure 104 for the molecule system.

[0062] In some implementations, the equivariant generative machine learning model 102 can receive conditioning data 112 and can conditionally process the initial molecule data 102 using the conditioning data 112 to generate the predicted 3D structure 104 for the molecule system. Theconditioning data 112 can include any appropriate data for conditional molecular structure generation. As an example, the conditioning data 112 can specify physical conditions (e.g., temperatures, pressures, salinities, acidities, etc.) for the molecule system and the equivariant generative machine learning model 100 can generate the predicted molecule structure 104 for the molecule system under the specified physical conditions. As another example, when the molecule system includes a protein and a ligand, the conditioning data 112 can include initial (e.g., unbound) predicted structures for the protein and the ligand and the equivariant generative machine learning model 100 can generate the predicted molecule structure 104 for a protein-ligand complex of the protein and the ligand as conditioned on the initial predicted structures for the protein and the ligand. In some implementations, the conditioning data 112 can include latent representations of the initial (unbound) structures for a protein and a ligand (e.g., latent representations that can be processed to determine spatial positions for each atom within the protein and the ligand) and the equivariant generative machine learning model 100 can generate the predicted molecule structure 104 for a protein-ligand complex of the protein and the ligand as conditioned on the initial predicted (unbound) structures for the protein and the ligand.

[0063] An example process by which the equivariant generative machine learning model can process the initial molecule data 102 (and, optionally, the conditioning data 112) to generate the predicted 3D structure 104 is described in more detail below with reference to FIG. 6.

[0064] A training system 114 can train the equivariant generative machine learning model 100 to optimize an objective function for the model 100 using a set of training data that includes initial molecule data and target 3D structures for a plurality of example molecule systems. In particular, as described in more detail below with reference to FIG. 7, the training system 114 train the equivariant generative machine learning model 100 by generating parameter updates 116 for the model 100 to optimize an objective function that measures an error between the target 3D structures and corresponding predicted 3D structures generated by the equivariant machine learning model 100 for the example molecule systems.

[0065] The predicted 3D structure 104 for the molecule system can be used as part of performing any of a variety of down-stream tasks regarding the molecule system. As one example, the predicted 3D structure 104 for the molecule system can be processed by the equivariant generative machine learning model 100 or another generative machine learning model to generate a predicted 3D structure for a molecule complex that includes the molecule system. As a further example, themolecule system can be a protein or a ligand included within a protein-complex and the predicted 3D structure 104 for the molecule system can be processed by the equivariant generative machine learning model 100 or another generative machine learning model to generate a predicted 3D structure for the protein-ligand complex that includes the molecule system.

[0066] As another example, the predicted 3D structure 104 for the molecule system can be processed to determine predicted molecule properties (e.g., a binding affinity, a free energy, atomic forces, an electron density, etc.) for the molecule system. For example, the predicted molecule properties for the molecule system can be determined by performing a computational simulation (e.g., a molecular dynamics simulation, a quantum mechanical simulation, etc.) of the predicted 3D structure 104 for the molecule system. As another example, the predicted molecule properties for the molecule system can be determined by processing the predicted 3D structure 104 using a predictive machine learning model configured (e.g., trained) to generate predicted molecule properties by processing 3D structures for molecule systems.

[0067] As described in more detail below with reference to FIGs. 2-5, the equivariant transformation neural network 106 can be configured to perform a sequence of equivariant neural network operations to process the network inputs to and generate the network outputs from the equivariant transformation neural network 106. The equivariance of the neural network operations performed by the equivariant transformation neural network 106 can ensure that processing initial molecule data 102 representing a spatial rotation of the initial, noisy 3D structure of the molecule system using the equivariant transformation neural network 100 will result in the equivariant generative machine learning model 100 generating updated molecule data 108 that represents an equivalent rotation of the predicted 3D structure for the molecule system.

[0068] The equivariance of the neural network operations performed by the equivariant transformation neural network 106 can encode a powerful inductive bias in the equivariant generative machine learning model 100 and can enable reduced consumption of computational resources by the equivariant generative machine learning model 100 during training and inference. More specifically, the equivariance property of the equivariant transformation neural network 106 can mitigate the risk of overfitting to training data and can reduce the amount of training data and the number of training iterations required for the equivariant generative machine learning model 100 to achieve an acceptable prediction accuracy. Further, the equivariance property of the equivariant transformation neural network 106 can enable theequivariant transformation neural network 106 to be implemented using a less complex neural network architecture, e.g., with fewer neural network parameters and fewer neural network layers as compared to non-equivariant transformation neural networks, which can further reduce consumption of computational resources (e.g., memory and computing power) during inference using the equivariant generative machine learning model 100.

[0069] As described in more detail below, the equivariant transformation neural network 106 can implement equivariant attention operations to update embeddings representing atoms or groups of atoms (e.g., multi-atom structures such as amino acids, nucleic acids, functional groups, etc.) within the molecular system. In some implementations, the equivariant transformation neural network 106 can implement a global attention operation over a set of embeddings representing respective groups of atoms within the molecule system (which may be far fewer in number, e.g., by one or more orders of magnitude, than the number of atoms within the molecule system) and a local attention operation over a set of atom embeddings representing each atom within the molecule system. This can enable the equivariant transformation neural network 106 to avoid the computational cost of performing global attention operations among atom embeddings representing the individual atoms of the molecule system (which can increase quadratically with the number of atoms within the molecule) while also enabling the equivariant transformation neural network 106 to still learn to recognize and exploit long-range, global interactions among atoms of the molecular system.

[0070] The equivariant generative machine learning model 100 can therefore accurately and efficiently generate the predicted 3D structure 104 for the molecule system, which can be used to perform any of a variety of downstream tasks, such as predicting properties of the molecule system, performing simulations of the molecule system, and so on. In particular, a prediction system 118 can use the predicted 3D structure 104 for the molecule system as generated by the equivariant generative machine learning model 100 to perform any of a variety of downstream prediction tasks regarding the molecule system. As a particular example, the prediction system 118 can use the equivariant generative machine learning model 100 to generate predicted 3D structures for a collection of candidate molecules (e.g., candidate ligands), which can be used to predict properties of the candidate molecules as part of screening a collection of candidate molecules for physical synthesis, for inclusion within a drug and so on, as described in more detail below with reference to FIG. 8.

[0071] FIG. 2 illustrates an example equivariant transformation neural network 106.

[0072] As described above with reference to FIG. 1, an equivariant generative machine learning model (e.g., the equivariant generative machine learning model 100 of FIG. 1) can use the equivariant transformation neural network 106 to process initial molecule data for a molecule system over a sequence of time steps to generate a predicted 3D structure for the molecule system. At each time step of generating the predicted 3D structure for the molecule system, the equivariant transformation neural network 106 can process a network input 202 that includes input molecule data representing a current 3D structure for the molecule system as of the time step to generate a network output 204 that characterizes a transformation of the current 3D structure for the molecule system for the time step. The equivariant generative machine learning model can combine the network output 204 from the equivariant transformation neural network with the input molecule data for the time step to generate updated molecule data that represents an updated 3D structure of the molecule system for the time step, which the equivariant generative machine learning model can then provide as a network input to the equivariant transformation neural network during a next time step of generating the predicted 3D structure for the molecule system.

[0073] The equivariant generative machine learning model can implement any of a variety of machine learning architectures for and the equivariant transformation neural network 106 can be configured to generate any appropriate network outputs for use by the equivariant generative machine learning model in generating updated 3D structures for molecule systems.

[0074] For example, the equivariant generative machine learning model can be a generative diffusion model configured to generate the predicted 3D structure for the molecule system following a reverse diffusion process and the equivariant transformation neural network 106 can be configured to generate the network output 204 at each time step to define a distribution for the updated molecule data for the molecule system at the time step (e.g., from which the generative machine learning model can use to sample the updated molecule data for the molecule system at the time step). As another example, the equivariant generative machine learning model can be a generative flow matching model or a stochastic interpolant model that is configured to generate the predicted 3D structure for the molecule system in accordance with a set of transformation differential equations and the equivariant transformation neural network 106 can be configured to generate the network output 204 at each time step to define the set of transformation differentialequations for the time step (e.g., which the generative machine learning model can numerically integrate to generate the updated molecule data for the molecule system at the time step).

[0075] At each time step of generating the predicted 3D structure for the molecule system, the network input 202 to the equivariant transformation neural network 106 can represent (e.g., encode) a 3D structure for the molecule system using any appropriate scheme for encoding the positions of atoms within the molecule system and the equivariant transformation neural network 106 can have any appropriate neural network architecture suitable for (equivariantly) processing (e g., denoising, transforming, etc.) the network inputs representing 3D structures for the molecule system. At each time step of generating the predicted 3D structure for the molecule system, the equivariant generative machine learning model can generate the updated molecule data to represent (e.g., encode) 3D structures for the molecule system using the same encoding scheme as used by network inputs to the equivariant transformation neural network 106.

[0076] In general, the network input 202 to the equivariant transformation neural network 106 can include a plurality of numerical features that specify 3D spatial positions for each atom within the molecule system. For example, the network input 202 to the equivariant transformation neural network 106 can include vectors of numerical values that specify 3D spatial positions for each atom within the molecule system. As another example, the network input 202 to the equivariant transformation neural network 106 can include a plurality of vectors of numerical values, each representing a respective atom or group of atoms (e.g., a respective amino acid residue, nucleotide, functional group, etc.) within the molecule system. For example, the network input 202 to the equivariant transformation neural network 106 can be sequences of embeddings, with each embedding representing a respective atom as an atom embedding for the atom or group of atoms (e.g., a respective amino acid residue, nucleotide, functional group, etc.) as a token embedding for the group of aotms. As a further example, the network input 202 to the equivariant transformation neural network 106 can be a sequence of embeddings that characterizes a graph representations of the molecule system that includes respective embeddings for each of a plurality of graph nodes of the graph that each represent respective atoms or groups of atoms (e.g., respective amino acid residues, nucleotides, functional groups, etc.) and respective embeddings for each of a plurality of graph edges for the graph that each represent respective bonds or interactions between the atoms or groups of atoms represented by respective pairs of graph nodes of the graph.

[0077] When a network input to the transformation neural network 110 includes token embeddings representing respective groups of atoms (e.g., respective amino acid residues, functional groups, etc.) within the molecule system, the network input 202 to the equivariant transformation neural network 106 can specify positions and orientations (e.g., rotations) of each of the groups of atoms within the molecule system. In particular, in some implementations, the network input 202 to the equivariant transformation neural network 106 can specify relative rotations between groups of atoms within the molecule system using torsion angles between the groups of atoms, as described in more detail below with reference to FIG. 5.

[0078] The spatial positions of each atom within a group of atoms (e.g., respective amino acid residues, functional groups, etc.) can be determined by the embedding for the group of atoms within the network input 202 to the equivariant transformation neural network 106 (and, when applicable, any numerical features of the network input 202 characterizing orientations of the group of atoms relative to other groups of atoms within the molecule system). As an example, the position of each atom within the group of atoms can be predefined and fixed and the position of each atom within the molecule system can be determined based on (i) the predefined location of the atom within the group of atoms and (ii) the location and orientation of the group of atoms specified by the embedding for the group of atoms. As another example, the embedding of the group of atoms can include numerical features specifying the position of each atom within the group of atoms and the position of each atom within the molecule system can be determined based on (i) the location of the atom within the group of atoms and (ii) the location and orientation of the group of atoms (e.g., both specified by the embedding for the group of atoms).

[0079] As described above, the equivariant transformation neural network 106 can be configured to perform a sequence of equivariant neural network operations to process the network input 202 to and generate the network output 204. The equivariance of the neural network operations performed by the equivariant transformation neural network 106 can ensure that, when processing a network input 202 that represents a spatially rotated 3D structure of the molecule system, the equivariant transformation neural network 100 will generate an appropriate network output 204 that can be processed by the equivariant generative machine learning model to generated updated molecule data that represents an equivalent rotation of an updated 3D structure for the molecule system.

[0080] In particular, the network input 202 can include spatially-independent scalar input molecule data 204-A and spatially-dependent input vector molecule data 204-B for the molecule system. Similarly, the network output 204 can include spatially-independent scalar input molecule data 204-A and spatially-dependent input vector molecule data 204-B for the molecule system.

[0081] For example, each input embedding of the network input 202 and each output embedding of the network output 204 that represents a respective atom or group of atoms can include: (i) a spatially-independent (e.g., scalar) part that is based on properties of atom or group of atoms that are invariant to an overall spatial position and rotation of the molecule system (e.g., chemical elements, torsion angles, and so on associated with the atom or group of atoms), and (ii) a spatially-dependent (e.g., vector) part that is based at least on 3D spatial positions of that atom or group of atoms. As an example, the spatially-dependent (e.g., vector) part of the embedding representing an atom or group of atoms can include an array of one or more 3D spatial coordinates associated with the atom or group of atoms, such as 3D spatial coordinates representing spatial positions associated with the atom or group of atoms, 3D spatial coordinates representing spatial displacements associated with the atom or group of atoms, 3D spatial coordinates representing embeddings of positions and / or displacements associated with the atom or group of atoms (e.g., as generated by processing data characterizing the positions and / or displacements associated with the atom or group of atoms using an embedding neural network), and so on.

[0082] As a further example, for each atom or group of atoms represented by the network input 202 and the network output 204, the respective embeddings for the atom or group of atoms can include (i) a spatially-independent scalar embedding for the atom or group of atoms that represents properties of the atom or group of atoms that are invariant to an overall spatial position and rotation of the molecule system and (ii) a spatially-dependent vector embedding for the atom or group of atoms that is based at least on 3D spatial positions of that atom or group of atoms. For example, an embedding, XL, for an / -th atom or group of atoms can be defined as the pair:

[0083] Where X is the spatially-independent scalar embedding for the z-th atom or group of atoms, Nsis a scalar dimension for the embedding X, (e.g., the number of scalar features within the embedding Xj), X is the spatially-dependent vector embedding for the i-th atom or group ofatoms, and Nvis a vector dimension for the embedding Xt(e g., the number of three dimensional vector features within the embedding Xi).

[0084] As described above, the equivariant transformation neural network 106 can be configured to perform equivariant network operations to process the input sequence embeddings, Xt,..., Xn, of the network input 202 representing n atoms or groups of atoms of the molecule system to generate the output sequence of embeddings, X' i,..., X'nfollowing:X'i. X'n= fQ(Xi. X

[0085] Where / e denotes the function defined by the network operations of the equivariant transformation neural network 106 as parameterized by the network parameters, 9, of the equivariant neural network 106.

[0086] The equivariance of the network operations of the equivariant transformation neural network 106 can ensure that, when processing a spatial rotations of the input sequence of embeddings, T(Xj),..., T(Xn), the equivariant transformation neural network 106 generates a sequence of equivalent spatial rotations of the output sequence of embeddings, T(X'j),..., T(X'n) where X'j,..., X'n= / g(X(-,..., Xri), following:

[0087] Where T(Xj) represents applying a same spatial rotation to each of the vector features of the vector embedding, X, of the embedding Xt, with T

[0088] In some implementations, the equivariant transformation neural network 106 can be configured to generate the network output 204 to be equivariant with respect to spatial translations of the vector features of the vector embeddings of the network input 202. For example, the equivariant transformation neural network 106 can be configured to ensure an equivariance of the network output 204 with respect to spatial translations of the vector features of the vector embeddings of the network input 202 by centering the vector features of the vector embeddings of the network input 202, e.g., by subtracting a position of a center of mass of the molecule system from each of the vector features of the vector embeddings of the network input 202.

[0089] T o perform such equivariant neural network operations, the equivariant neural network 106 can include any appropriate combination of any of a variety of equivariant processing layers 206-A through 206-N that are each configured to equivariantly process respective layer inputs to generate respective layer outputs 208-A through 208-N. For example, the equivariant processinglayers 206-A through 206-N can include one or more linear layers that are configured to generate respective layer outputs as linear transformations of their respective layer inputs. As another example, the equivariant processing layers 206-A through 206-N can include one or more equivariant Transformer layers that are configured to generate respective layer outputs by performing equivariant attention operations to process their respective layer inputs. As another example, the equivariant processing layers 206-A through 206-N can include one or more equivariant gated feed-forward layers that are configured to configured to generate respective layer outputs as non-linear transformations of their respective layer inputs.

[0090] As described above, in some implementations, the equivariant generative machine learning model can receive conditioning data 112 for performing conditional molecular structure generation. The conditioning data 112 can include one or more embeddings that represent respective conditioning features for conditionally generating the predicted 3D structure for the molecule system. In particular, similar to the embeddings of the network input 202 and the network output 204, each embedding for a conditioning feature of the conditioning data 112 can include (i) a spatially-independent scalar embedding for the conditioning feature that represents properties of the conditioning feature that are invariant to an overall spatial position and rotation of the molecule system and (ii) a spatially-dependent vector embedding for the conditioning feature that represents 3D spatial positions or displacements associated with the conditioning feature.

[0091] The equivariant transformation neural network 106 can be configured to process the conditioning data 112 by any of a variety of methods. As an example, in some implementations, the network input 202 to the equivariant transformation neural network 106 can include the conditioning data 112. As another example, when the equivariant transformation neural network 106 includes equivariant Transformer layers, some or all of the equivariant Transformer layers can be configured to perform (equivariant) cross-attention operations between sequences of input embeddings and the embeddings representing the conditioning features of the conditioning data 112. As another example, the equivariant transformation neural network 106 can be configured to process the conditioning data 112 using a ControlNet neural network (e.g., an implementation of a ControlNet neural network that utilizes equivariant processing layers such as the linear layers, equivariant Transformer layers, equivariant gated feed-forward layers as described above) to determine conditional adjustments to the network output 204 from the equivariant transformationneural network 106, as described by Zhang et al. in “Adding Conditional Control to Text-to-Image Diffusion Models”.

[0092] In some implementations, the transformation neural network 106 can generate, for each atom of the molecule system, a transformation output for the atom that defines an update to a 3D spatial of the atom by processing only the updated spatially-dependent (e.g., vector) embeddings for the molecule system.

[0093] In some implementations, the transformation neural network 106 can generate, for each atom, a confidence for a 3D spatial position of the atom within the updated predicted 3D structure by processing only the spatially-independent (e.g., scalar) embeddings of the atoms.

[0094] Example architectures of equivariant Transformer layers and equivariant gated feedforward neural networks that can be included as equivariant processing layers of the equivariant transformation neural network 106 are described in more detail next with reference to FIG. 3 and FIG. 4.

[0095] FIG. 3 illustrates an example equivariant Transformer layer 300. The equivariant Transformer layer 300 is an example of an equivariant processing layer that can be included within an equivariant transformation neural network (e.g., the equivariant transformation neural network 106 described in more detail above with reference to FIG. 2).

[0096] The equivariant Transformer layer 300 can be configured to perform equivariant attention operations to process their respective layer inputs to process input data 302 that includes input scalar data 302-A and input vector data 302-B to generate output data 304 that includes output scalar data 304-A and output vector data 304-B. In particular, as described above with reference to FIG. 2, the equivariant Transformer layer 300 can process input embeddings that represent atoms or groups of a molecule system to generate output embeddings representing the atoms or groups of the molecule system.

[0097] As described above, the input data 302 to the equivariant Transformer layer 300 can include a plurality of input embeddings that each represent respective atoms or groups of atoms of a molecule system. The embedding for each atoms or group of atoms within the molecule system can include (i) a spatially-independent scalar embedding for the atom or group of atoms that represents properties of the atom or group of atoms that are invariant to an overall spatial position and rotation of the molecule system and (ii) a spatially-dependent vector embedding for the atomor group of atoms that is based at least on the 3D spatial positions of that atom or group of atoms within the molecule system.

[0098] For example, the embedding, X for an z-th atom or group of atoms within the molecule system can be defined as the pair= (x. A ), where X is the spatially-independent scalar embedding and X is the spatially-dependent vector embedding for the i-th atom or group of atoms within the molecule system.

[0099] The equivariant Transformer layer 300 can process the input embeddings for the atoms or groups of atoms within the molecule system by performing equivariant attention operations using an equivariant attention layer 306 to generate an attention layer output 308 that includes output scalar data 308-A and output vector data 308-B. In particular, the equivariant attention layer 306 can generate, as the attention layer output 308, updated embeddings for each of the atoms or groups of atoms of the molecule system.

[0100] For example the equivariant attention layer 306 can generate the updated embedding, X't, for the z-th atom or group of atoms within the molecule system that is defined as the pair X'i = (X', X'f), where ' is an updated spatially-independent scalar embedding and ' is an updated spatially-dependent vector embedding for the i-th atom or group of atoms within the molecule system.

[0101] In general, the equivariant attention layer 306 can perform an equivariant query-key -value (QKV) attention operation to generate the updated embeddings for the atoms or groups of atoms within the molecule system.

[0102] As part of performing the equivariant query-key-value (QKV) attention operation, the equivariant attention layer 306 can determine an attention weight,for each pairing of the input embedding, XL, for the i-th atom or group of atoms of the molecule system with a j-th input embedding, Xj. The equivariant attention layer 306 can use the attention weights and a set of value embeddings to generate the updated embeddings for the atoms or groups of atoms within the molecule system. For example the equivariant attention layer 306 can generate the updated embedding, X' for the z-th atom or group of atoms within the molecule system following:

[0103] Where V is the scalar value embedding and is the vector value embedding for a value embedding, Vj, associated with the j'-th input embedding, Xj.

[0104] The equivariant attention layer 306 can generate each attention weight, a(y, to be invariant to rotations of the molecule system, each scalar value embedding, VjSto be invariant to rotations of the molecule system, and each vector value embedding,to be equivariant to or rotations of the molecule system, which enables the equivariant attention layer 306 to generate each updated embedding, X' as an equivariant transformation of the input embeddings to the attention layer 306.

[0105] For example, the equivariant attention layer 306 can generate each value embedding, Vj, by processing the associated input embedding, Xy, using an equivariant processing layer for generating the value embeddings, such as a linear layer or an equivariant gated feed-forward layer.

[0106] To determine invariant attention weights, the equivariant attention layer 306 can generate each attention weight, atj, as an invariant function of a query embedding, Q,associated with the input embedding, Xj, (e g., as generated by processing the input embedding Xj using an equivariant processing layer for generating the query embeddings, such as a linear layer of an equivariant gated feed-forward layer) and a key embedding, Kj = K, K, with a / '-th input embedding, Xy, (e g., as generated by processing the input embedding Xj using an equivariant processing layer for generating the key embeddings, such as a linear layer of an equivariant gated feed-forward layer).

[0107] For example, in some implementations, the equivariant attention layer 306 can generate the attention weight, atj, using a Frobenius inner product, following:

[0108] Where btj is a bias term that can specify an attention mask of the equivariant attention layer 306, Dsis a scalar dimension of the key and query embeddings, Dvis a vector dimension of the key and query embeddings, and where, for matrices A and B, < A, B >Fdenotes the Frobenius inner product of the matrices A and B (e g., with < A, B >F=.pq^-pq Bpq),.

[0109] As another example, in some implementations, the equivariant attention layer 306 can generate the attention weight, a^, using a logarithmic distance, following:

[0110] In some cases, generating the updated embeddings based solely on a linear combination of such value embeddings can limit an expressivity of the equivariant attention layer 306. In some implementations, the equivariant attention layer can be configured to generate more expressive outputs by generating equivariant basis expansions of the value embeddings and generating the updated embeddings using the equivariant basis expansions of the value embeddings.

[0111] For example, in some implementations, the equivariant attention layer 306 can generate an equivariant basis expansion Bj for each value embedding V, using a cross-product operation, following:

[0112] The equivariant attention layer 306 can then generate expanded value embeddings by processing the equivariant basis expansions and the value embeddings using linear layer. For example, the equivariant attention layer 306 can generate each updated value embedding, V'j, following:V'i = f ( / D

[0113] represents an output of the linear layer processing a concatenation of the value embedding Vj with the equivariant basis expansion Bj.

[0114] The equivariant attention layer 306 can then use the updated value embeddings, V'j, to generate each updated embedding,for the atoms and / or group of atoms within the molecule system following:

[0115] The equivariant attention layer 306 can be configured to perform any of a variety of equivariant attention operations by generating appropriate key and value embeddings. For example, the equivariant attention layer 306 can be configured to perform an equivariant selfattention operation to generate the updated embeddings, X't, by generating respective key and value embeddings associated with each atom or group of atoms within the molecule system basedon the input embeddings, Xt. As another example, the equivariant attention layer 306 can be configured to perform an equivariant cross-attention operation to generate the updated embeddings, X' based on input conditioning data 112 by generating respective key and value embeddings based on input embeddings of conditioning features within the input conditioning data 112.

[0116] As another example, the equivariant attention layer 306 can be configured (e.g., by using appropriately selected bias terms) to perform various equivariant cross-attention operations between the input embeddings, Xt. For example, the equivariant attention layer 306 can be configured to perform a cross attention operation to generate updated embeddings representing groups of atoms based on the input embeddings representing individual atoms by, e.g., generating and processing query embeddings based only on the input embeddings representing groups of atoms and generating and processing key and value embeddings based only on the input embeddings representing individual atoms. As another example, the equivariant attention layer 306 can be configured to perform a cross attention operation to generate updated embeddings representing individual atoms based on the input embeddings representing groups of atoms by, e.g., generating and processing query embeddings based only on the input embeddings representing individual and generating and processing key and value embeddings based only on the input embeddings representing groups of atoms.

[0117] In general, the bias termcan be selected to implement any appropriate attention mask for the equivariant attention layer 306. In particular, when performing an equivariant self-attention operation to generate the updated embeddings, the bias term b,: can be selected to implement local attention operations to update embeddings representing individual atoms within the molecule system and a global attention operations between embeddings representing groups of atoms within the molecule.

[0118] For example, to implement the local attention operations to update embeddings representing individual atoms, the bias term b(can be selected such that, when either the embedding Xtor the embedding Xj represents an individual atom of the molecule system, the attention weight, atj, (i) is non-zero only if a distance between the atoms or groups of atoms represented by the embeddings Xtand Xj is less a pre-determined threshold distance and (ii) is zero when the distance between the atoms or groups of atoms represented by the embeddings XLand Xj exceeds the pre-determined threshold distance. To implement global attention operations to update the embeddings representing groups of atoms within the molecule system, the bias term btj can be selected such that the attention weight, a^, can be non-zero for each embedding Xtand embedding Xj that both represent groups of atoms within the molecule system.

[0119] In some implementations, the equivariant attention layer 306 can implement local and global attention operations utilizing an attention mask that is not defined by the bias term bjj. For example, in some implementations, the equivariant attention layer 306 can implement local and global attention operations by determining the attention weights using a multiplicative attention mask, rriij.

[0120] Implementing local attention operations to update the embeddings representing individual atoms and global attention operations to update the embeddings representing groups of atoms within the molecule system can enable the equivariant attention layer 306 to avoid the computational cost of performing global attention operations among atom embeddings representing the individual atoms of the molecule system (which can increase quadratically with the number of atoms within the molecule) while also enabling the equivariant attention layer 306 to still learn to recognize and exploit long-range, global interactions among atoms of the molecular system.

[0121] In some implementations, the equivariant Transformer layer 300 can directly output the attention layer output 308 as the output data 304. In other implementations, the equivariant Transformer layer 300 can process the attention layer output 304 using an equivariant gated feedforward layer 310 to generate the output data 304. An example architecture of the equivariant gated feed-forward layer 310 is described in more detail below with reference to FIG. 4.

[0122] FIG. 4 illustrates an example equivariant gated feed-forward layer 400. The equivariant gated feed-forward layer 400 is an example of an equivariant processing layer that can be included within an equivariant transformation neural network (e.g., the equivariant transformation neural network 106 described in more detail above with reference to FIG. 2).

[0123] The equivariant gated feed-forward layer 400 can be configured to process a layer input that includes input scalar data 402-A and input vector data 402-B to generate a layer output that includes output scalar data 404-A and output vector data 404-B following a non-linear transformations of the input scalar data 402-A and input vector data 402-B.

[0124] In particular, the equivariant gated feed-forward layer 400 can process the input scalar data 402-A using a first feed-forward layer 406-A to generate the output scalar data 404-A. The equivariant gated feed-forward layer 400 can process the input scalar data 402-A using a second feed-forward layer 406-B and can combine the output of the second feed-forward layer 406-B with the input vector data 402-B by means of an outer product operation. The equivariant gated feedforward layer 400 can then process the combination of the output of the second feed-forward layer 406-B and the input vector data 402-B using a linear layer 408 to generate the output vector data 404-B.

[0125] As described above with reference to FIG. 2 and FIG. 3, the equivariant gated feed-forward layer 400 can process input embeddings that represent atoms or groups of a molecule system to generate output embeddings representing the atoms or groups of the molecule system. In particular, the equivariant gated feed-forward layer 400 can process an input embedding, Xt, for an z-th atom or group of atoms of the molecule system to generate an input embedding, X' for the z-th atom or group of atoms of the molecule system following:x" = Mf2(Xs-) ®Xr

[0126] Where Xsdenotes the input scalar embeddings for each of the atom or group of atoms, Xvdenotes the input vector embeddings for each of the atom or group of atoms,is the output scalar embedding for the z-th atom or group of atoms, X' is the output vector embedding for the f flz-th atom or group of atoms, gQdenotes the operation performed by the first feed-forward layer 406-A, g^2denotes the operation performed by the second feed-forward layer 406-B, and g denotes the operation performed by the linear layer 408.

[0127] For example, the operation, g^, performed by the linear layer 408 can be defined following:

[0128] Where A is a multi-dimensional array of multiplicative weights of the linear layer 408, Xj is the input scalar embedding for the 7 -th atom or group of atoms, and Xklis an Z-th vector feature of the input vector embedding for the: -th atom or group of atoms.

[0129] FIG. 5 illustrates a torsion angle between two groups of atoms for a molecule system.

[0130] As described above with reference to FIG. 1, an equivariant generative machine learning model (e.g., the equivariant generative machine learning model 100 of FIG. 1) can use an equivariant transformation neural network (e.g., the equivariant transformation neural network 106 of FIG. 1) to process a network input characterizing an initial, noisy molecular structure for a molecule system to generate a predicted molecule structure for the molecule system. In general, the equivariant generative machine learning model can use the equivariant transformation neural network to process and generate embeddings of the molecule structure that specify 3D spatial locations for each atom within the molecule system. In particular, the equivariant generative machine learning model can use the equivariant transformation neural network to process and generate multiple embeddings (e.g., feature vectors, atom embeddings, token embeddings, etc.) for the molecule system, with each embedding characterizing a position and orientation of a respective atom or group of atoms (e.g., a respective amino acid, functional group, etc.) within the molecule system.

[0131] As an example, the equivariant generative machine learning model can use the equivariant transformation neural network to process and generate respective embeddings for atom groups 504 and 506 of FIG. 5. The equivariant generative machine learning model can also use the transformation neural network process and generate a torsion angle 502 between the atom groups 504 and 506 within the molecule that specifies a relative orientation between the atom groups 504 and 506. The torsion angle 502 specifies an angle between a plane for the atom group 504 and a plane for the atom group 506. In particular, when atoms of the atom groups 504 and 506 share a chemical bond, the torsion angle 502 can specify an angle between the planes for the atom groups 504 and 506 and within a plane normal to the chemical bond between the atom groups 504 and 506.

[0132] The equivariant generative machine learning model can use the equivariant transformation neural network to process and generate a plurality of torsion angles for a molecule system. For example, the equivariant generative machine learning model can use the equivariant transformation neural network to process and generate side-chain torsion angles for a protein (e.g., torsion angles between side-chain atom groups and back-bone atom groups for the protein), backbone torsion angles for a protein (e.g., torsion angles between different back-bone atom groups for the protein), and so on.

[0133] As described above, the equivariant transformation neural network can be configured to generate updated embeddings for the molecule systems that are equivariant under spatial displacements and rotations of the molecule system. In particular, the equivariant transformation neural network can be configured to process and update embeddings for the molecule system that each include (i) scalar features that are invariant to spatial displacements and rotations of the molecule system and (ii) vector features that are equivariant to spatial displacements and rotations of the molecule system. The scalar features of the embeddings for the molecule system can represent properties of the molecule system that are independent of a position or a rotation of the molecule system, such as the torsion angle 502, the chemical compositions of the atom groups 504 and 506, the chemical elements of each of the atoms of the atom groups 504 and 506, and so on. The vector features of the embeddings for the molecule system can represent spatial relationships among the atoms of the molecule system, such as positions and orientations of the atom groups 504 and 506, the displacement between the atom groups 504 and 506, positions of atoms within the atom groups 504 and 506, displacements between the atoms of the atom groups 504 and 506, and so on.

[0134] FIG. 6 is a flow diagram of an example process 600 for generating a predicted 3D structure for a molecule system. For convenience, the process 600 will be described as being performed by a system of one or more computers located in one or more locations. For example, an equivariant generative machine learning model, e.g., the equivariant generative machine learning model 100 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 600.

[0135] The system can obtain initial molecule data for the molecule system (step 602). The molecule system can include one or more molecules, e.g., proteins, ligands, protein-ligand complexes, and so on. In general, the system can receive data specifying chemical compositions of the molecules within the molecule system and can generate the initial molecule data based on the specified chemical compositions of the molecules within the molecule system. For example, the system can generate the initial molecule data by receiving and processing data specifying atoms, charges on atoms, bonds, functional groups, amino acids, branches (e.g., from a molecule main chain), ring structures, and so forth for the molecules within the molecule system. As a further example, the system can generate the initial molecule data by receiving and processing, e g., Simplified Molecular Input Line Entry System (SMILES) formatted data, InternationalChemical Identifier (InChi) formatted data, and so on specifying chemical structures for the molecules within the molecule system.

[0136] As another example, the system can generate the initial molecule data by receiving and processing data specifying one or more amino acid sequences for one or more molecules of the molecule system. Each amino acid sequence for the molecule system can specify an identity and an ordering for each of a plurality of amino acid residues within a molecule of the molecule system.

[0137] As another example, the system can generate the initial molecule data by receiving and processing data specifying one or more nucleic acid sequences for one or more molecules of the molecule system. Each nucleic acid sequence for the molecule system can specify an identity and an ordering for each of a plurality of nucleotides within a molecule (e.g., a nucleic acid, a ligand, etc.) of the molecule system.

[0138] In some implementations, the system can receive and process data characterizing initial or estimated 3D structures for one or more molecules of the molecule system as part of generating the initial molecule data. The initial or estimated 3D structures of the one or more molecules of the molecule system can be determined by any of a variety of means, such as by physical experiment (eg., x-ray crystallography), by computational simulation (e.g., by molecular dynamics simulations, quantum mechanical simulations, etc.), and so on. As an example, the initial or estimated 3D structure of a protein can be an unfolded or a partially folded structure for the protein. As another example, when the molecule system includes a protein-ligand complex, the system can generate the initial molecule data for the protein-ligand complex by receiving and processing initial (unbound) 3D structures for a ligand and a protein of the protein-ligand complex. The data characterizing the initial or estimated 3D structures for the one or more molecules of the molecule system can specify chemical compositions and 3D structures for the one or more molecules by specifying a chemical element and an initial spatial position for each atom of the one or more molecules.

[0139] The system can generate the initial molecule data using the received data characterizing the chemical composition of the molecule system by any of a variety of methods. As an example, the system can sample an initial noisy spatial position for each atom of the molecule system from a distribution of noisy spatial positions for the atom. As another example, when the system receives data specifying initial or estimated 3D structures for one or more molecules of the molecule system, the system can generate the initial molecule data by determining the initial noisyspatial position for each atom of the molecule system as specified by the initial or estimated 3D structures for the one or more molecules of the molecule system.

[0140] The initial molecule data can include a plurality of numerical features that specify 3D spatial positions for each atom within the molecule system. For example, the initial molecule data can be a vector of numerical values that specifies 3D spatial positions for each atom within the molecule system. As another example, the initial molecule data can include a plurality of vectors of numerical values, each representing a respective atom or group of atoms (e.g., a respective amino acid residue, functional group, etc.) within the molecule system. For example, the initial molecule data can be a sequence of embeddings, with each embedding representing a respective atom within the molecule system as an atom embedding for the respective atom or a respective group of atoms (e.g., a respective amino acid residue, functional group, etc.) within the molecule system as a token embedding for the respective group of atoms.

[0141] When the initial molecule data includes token embeddings representing respective groups of atoms (e.g., respective amino acid residues, functional groups, etc.) within the molecule system, the initial molecule data can include structure parameters that specify positions and orientations (e.g., rotations) of each of the groups of atoms within the molecule system. In particular, in some implementations, the initial molecule data can include structure parameters specifying torsion angles (e.g., side chain torsion angles, back bone torsion angles, etc.) between groups of atoms of the molecule system. The spatial positions of each atom within a group of atoms can be determined by the token embedding for the group of atoms within the initial molecule data (and, when applicable, structural parameters within the initial molecule data characterizing orientations of the group of atoms relative to other groups of atoms within the molecule system). As an example, the position of each atom within the group of atoms can be predefined and fixed and the position of each atom within the molecule system can be determined based on (i) the predefined location of the atom within the group of atoms and (ii) the location and orientation of the group of atoms specified by the token embedding for the group of atoms. As another example, the token embedding of the group of atoms can include numerical features (e.g., structure parameters) specifying the position of each atom within the group of atoms and the position of each atom within the molecule system can be determined based on (i) the location of the atom within the group of atoms and (ii) the location and orientation of the group of atoms (e.g., both specified by the token embedding for the group of atoms).

[0142] In some implementations, the system can receive conditioning data for the molecule system (step 604). The conditioning data can include any appropriate data for conditional molecular structure generation. As an example, the conditioning data can specify physical conditions (e.g., temperatures, pressures, salinities, acidities, etc.) for the molecule system. As another example, when the molecule system includes a protein-ligand complex, the conditioning data can include initial (e.g., unbound) predicted structures for a protein and a ligand of the protein-ligand complex.

[0143] In general, the system can process the initial molecule data identifying the molecular system using an equivariant transformation neural network to generate data defining the predicted 3D structure of the molecular system. In particular, the system can use the equivariant transformation neural network to iteratively update data representing a current predicted 3D structure of the molecular system over a sequence of time steps. At each time step, the system can perform steps 606 through 610 as described below.

[0144] At each time step, the system can obtain data defining a current predicted 3D structure for the molecule system at the time step (step 606). The data defining the current predicted 3D structure for the molecule system at the time step can include atom embeddings representing respective atoms within the molecule system and / or token embeddings representing respective groups of atoms (e.g., functional groups, amino acid residues, nucleotides, etc.) within the molecule system. In some implementations, the data defining the current predicted 3D structure for the molecule system can include a respective atom embedding representing each atom within the molecule system.

[0145] Each atom embedding for an atom of the molecule structure and each token embedding fora group of atoms of the molecule structure can include a spatially-independent embedding (e.g., a scalar embedding) and a spatially-dependent embedding (e.g., a vector embedding) that collectively define the embedding for the atom or group of atoms. The system can generate the spatially-independent embedding for each atom or group of atoms based on properties of the atom or group of atoms that are invariant to a spatial position of the atom or group of atoms, such as chemical compositions, a charges, an atomic masses, electron configurations, electronegativities, ionization energies, electrostatic potentials, torsion angles, and so on associated with the atom or group of atoms. The system can generate the spatially-dependent embedding for each atom or group of atoms based at least in part on a 3D spatial position of the atom or group of atoms within the molecule system. The system can generate the spatially-dependent embedding for each atomor group of atoms to include three spatial dimensions and a plurality of channel dimensions (e.g., with each channel dimension representing a respective spatially-dependent vector feature for the atom or group of atoms).

[0146] The system can determine the current predicted structure for the molecule system at the first time step using the initial molecule data for the molecule system. At each time step after the first time step, the system determine the current predicted structure for the time step to be an updated 3D structure for the molecule system as generated during a previous time step (e.g., an updated 3D structure for the molecule system as generated following step 610 during the previous time step).

[0147] The system can process a network input that represents the current predicted 3D structure at the time step using the equivariant transformation neural network to generate a transformation output for the time step that defines a modification to the current predicted 3D structure at the time step (step 608).

[0148] As described in more detail above with reference to FIG. 2, the equivariant transformation neural network can process the atom embeddings representing the atoms of the molecule structure and / or the token embeddings representing the groups of atoms of the molecule structure using a sequence of equivariant processing layers to generate updated atom embeddings and / or token embeddings for the molecule structure. Each of the equivariant processing layers can be configured to perform a processing operation that is equivariant to rotations of the current predicted 3D structure for the molecule structure.

[0149] In particular, the equivariant transformation neural network can include one or more equivariant Transformer layers that are each configured to update the atom embeddings and / or the token embeddings for the molecule structure by performing respective equivariant attention operations. As described above, each equivariant Transformer layer can be configured to perform a respective equivariant query-key-value attention operation.

[0150] For example, some or all of the equivariant Transformer layers can be configured to update the atom embeddings and / or the token embeddings for the molecule structure by performing a respective equivariant self-attention operation. As part of performing an equivariant self-attention operation, an equivariant Transformer layer can generate, for each atom and / or group of atoms of the molecule structure, respective query, key, and value embeddings by processing the input atom embeddings and / or token embeddings for the molecule structure using an equivariant processinglayer. For each atom and / or group of atoms, the query, key, and value embeddings can include respective spatially-independent (e.g. scalar) query, key, and value embeddings and respective spatially dependent (e.g., vector) query, key, and value embeddings. The spatially dependent (e.g., vector) query, key, and value embeddings can each include three spatial dimensions and a plurality of channel dimensions.

[0151] The equivariant Transformation layer can use the query, key, and value embeddings for the atoms and / or groups of atoms of the molecule system to equivariantly generate the updated embeddings for the molecule system using query-key-value attention. In particular, the system can use the queries and value embeddings for the molecule system to generate an array of attention weights that each correspond to a respective pair of atoms and / or group of atoms within the molecule system and are each invariant to rotations of the current predicted 3D structure for molecule system.

[0152] The system can determine the attention weights for each pair of atoms or group of atoms by combining any of a variety of spatially-independent (e.g., scalar) attention weight terms and spatially-dependent (e.g., vector) attention weight terms.

[0153] In some implementations, the system can determine each attention weight based at least in part on a scalar inner product between: (i) a spatially independent (e.g., scalar) query embedding for a first atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight, and (ii) a spatially independent (e.g., scalar) key embedding for a second atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight.

[0154] In some implementations, the system can determine each attention weight based at least in part on a Frobenius inner product between: (i) a spatially dependent (e.g., vector) query embedding for the first atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight, and (ii) a spatially dependent (e.g., vector) key embedding for the second atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight.

[0155] In some implementations, the system can determine each attention weight based at least in part on a logarithmic distance between: (i) the spatially dependent (e.g., vector) query embedding for the first atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight, and (ii) the spatially dependent (e.g., vector) key embedding for the second atom or group of atoms of the pair of atoms and / or group of atoms for the attention weight.

[0156] The system can generate the updated atom embeddings and / or updated token embeddings for the molecule structure as a combination of the value embeddings generated for the atoms and / or groups of atoms, e.g., as weighted using the generated attention weights.

[0157] In some implementations, some or all of the equivariant Transformer layer can, after updating the embeddings for the molecule system using a respective equivariant attention operation, update the embeddings for the molecule system using an equivariant gated feedforward neural network layer. As described above with reference to FIG. 4, each equivariant gated feedforward layer can update each the embeddings for the molecule system by first generating an intermediate spatially-dependent embedding by combining: (i) the spatially-dependent (e.g., vector) embedding and (ii) a result of processing the spatially-independent (e.g., vector) embedding using a feedforward neural network layer and by then updating the spatially-dependent (e.g., vector) embedding by processing the intermediate spatially-dependent (e.g., vector) embedding by an equivariant linear neural network layer.

[0158] In some implementations, the equivariant Transformer neural networks can be configured to update the token embeddings for the molecule system representing groups of atoms within the molecule system following a global self-attention operation and to update the atom embeddings for the molecule system representing individual atoms within the molecule system following a local self-attention operation.

[0159] In some implementations, one or more of the equivariant Transformer neural networks can be configured to update the atom embeddings and / or the token embeddings for the molecule structure by performing equivariant cross-attentions operations. As an example, when the system obtains conditioning data, the system can one or more of the equivariant Transformer neural networks can be configured to update the atom embeddings and / or the token embeddings for the molecule structure by performing equivariant cross-attentions operations between the atom embeddings and / or the token embeddings for the molecule structure for the molecule system and embeddings of the conditioning data.

[0160] In some implementations, the system can generate, for each atom of the molecule system, a transformation output for the atom that defines an update to a 3D spatial of the atom by processing only the updated spatially-dependent (e.g., vector) embeddings for the molecule system.

[0161] In some implementations, the system generate, for each atom, a confidence of the generative machine learning model in a 3D spatial position of the atom within the updated predicted 3D structure by processing only the spatially-independent (e.g., scalar) embeddings for the molecule system.

[0162] After generating the updated atom embeddings and / or the token embeddings for the molecule structure, the transformation neural network can process the updated atom and / or token embeddings to generate the transformation output based on the updated atom and / or embeddings.

[0163] As described above with reference to FIG. 2, the equivariant transformation neural network can be configured to generate appropriate transformation outputs for any of a variety of generative machine learning architectures.

[0164] As an example, the system can implement a generative diffusion model configured to generate the predicted 3D structure for the molecule system following a reverse diffusion process., and the equivariant transformation neural network can be configured to generate a transformation output for each time step that defines a distribution of updated molecule structures for the molecule system at the time step.

[0165] As another example, system can implement a generative flow matching model or a stochastic interpolant model that is configured to generate the predicted 3D structure in accordance with a set of transformation differential equations that characterize a transformation from a distribution (e.g., a noise distribution) for the initial molecule data for the molecule system to a distribution for the final updated molecule data representing the predicted 3D structure for the molecule system. When the system implements a generative flow matching model or a stochastic interpolant model, the equivariant transformation neural network can be configured to generate a transformation output for each time step that defines the set of transformation differential equations for transforming the current 3D structure as of the time step to generate the updated molecule data for the molecule system at the time step.

[0166] The system can then update the current predicted 3D structure using the transformation output generated by the equivariant transformation neural network for the time step (step 610). As described above, the system can implement any of a variety of generative machine learning architectures and can update the current predicted 3D structure using the transformation output generated by the equivariant transformation neural network by any appropriate method.

[0167] For example, when the system implements a generative diffusion model, the system can generate the updated molecule data at each time step based on the network output from the equivariant transformation neural network for the time step by sampling (e.g., using a diffusion sampling technique) the updated molecule data for the molecule system from a distribution defined by the network output from the equivariant transformation neural network and the diffusion coefficient value for the time step.

[0168] As another example, when the system implements a generative flow matching model or a stochastic interpolant model, the system can to generate the updated molecule data for the molecule system at each time in accordance with a set of transformation differential equations that are defined by defined by the network output from the equivariant transformation neural network for the time step (e.g., by numerically integrating the set of transformation differential equations).

[0169] The system can output the generated predicted 3D structure for the molecule system (step 612). In particular, the system can output the predicted 3D structure for the molecule system as the current predicted 3D structure for the molecule system after the final time step.

[0170] The system can output the generated predicted 3D structure for the molecule system by any of a variety of means, e.g., by storing the predicted 3D structure in a memory of the system, by transmitting the predicted 3D structure over a data communications network, by displaying a visual rendering of the predicted 3D structure on a display of a user device, and so on.

[0171] The predicted 3D structure for the molecule system can be used as part of performing any of a variety of down-stream tasks.

[0172] As one example, a prediction system can process the predicted 3D structure for the molecule system to generate a predicted 3D structure for a molecule complex that includes molecules of the molecule system. As a further example, when the molecule system includes a protein or a ligand, a prediction system can process the predicted 3D structure for the molecule system to generate a predicted 3D structure for a protein-ligand complex that includes protein or the ligand of the molecule system.

[0173] As another example, a prediction system can process the predicted 3D structure for the molecule system to determine predicted molecule properties (e.g., a binding affinity, a free energy, atomic forces, an electron density, etc.) for the molecule system. For example, the prediction system can determine the predicted molecule properties for the molecule system by performing a computational simulation (e.g., a molecular dynamics simulation, a quantum mechanicalsimulation, etc.) of the predicted 3D structure for the molecule system. As another example, the prediction system can determine the predicted molecule properties for the molecule system by processing the predicted 3D structure using a molecule property prediction neural network that is configured (e.g., trained) to generate predicted molecule properties by processing 3D structures for molecule systems.

[0174] FIG. 700 is a flow diagram of an example process 700 for training an equivariant generative machine learning model to generate predicted 3D structures for molecule systems. For convenience, the process 700 will be described as being performed by a system of one or more computers located in one or more locations. For example, a training system, e.g., the training system 112 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 700.

[0175] The system can obtain training data for the equivariant generative machine learning model that includes training examples specifying molecule data for a plurality of example (e.g., training) molecule systems (step 702). Each training example can include data characterizing (i) a chemical composition of the example molecule system for the training example and (ii) a target 3D structure for the example molecule system for the training example.

[0176] The system can obtain the target 3D structures for the example molecule systems from any of a variety of sources. As an example, the training data can include target 3D structures for example molecule systems as determined by physical experiment (e.g., as determined using x-ray crystallography of physical samples of the example molecule systems, cryo-electron microscopy (cryo-EM) of physical samples of the example molecule systems, and so on). As another example, the training data can include target 3D structures for example molecule systems as determined by computational simulations (e.g., as determined using molecular dynamics simulations of the example molecule systems, quantum mechanical simulations of the example molecule systems, and so on). In particular, when the example molecule systems include protein-ligand complexes, the target 3D structures for the example molecule systems can be determined using molecular docking computational simulations of the example molecule systems.

[0177] In some implementations, to promote an equivariance of the generative machine learning model with respect to spatial translations of the example molecule systems, the system can generate multiple augmented training examples for each example molecule system. In particular, the system can generate each augmented training example for an example molecule system byspatially translating the target 3D structure for the example molecule system by a randomly sampled spatial offset.

[0178] The system can train the equivariant generative machine learning model over a sequence of training iterations. At each training iteration, the system can perform steps 704 through 710.

[0179] The system can generate noisy 3D structures for a plurality of training examples for the training iteration (step 704). The system can generate the noisy 3D structure for each training example for the training iteration by sampling and adding noise to data representing the target 3D structure for the training example.

[0180] In particular, for each training example, the system can select a time step for the training example and can sample the added noise from a noise distribution for the selected time step. For example, the system can sample the added noise for each training example from a noise distribution (e g., a multi-variate normal distribution, a multi-variate uniform distribution, etc.) that depends on the selected time step for the training example, such as noise distribution with zero mean and a variance that depends on the selected time step for the training example.

[0181] The system can process the noisy 3D structures using the equivariant generative machine learning model to generate predicted 3D structures for each of the training examples for the training iteration (step 706). As described above, the equivariant generative machine learning model can generate the predicted 3D structure for each training example by iteratively transforming the noisy 3D structure for the training example using an equivariant transformation neural network over a sequence of transformation time steps for the training example. At each transformation time step for a training example, the system can update (e.g., transform) a current predicted 3D structure for the training example based on a transformation output generated by the equivariant transformation neural network for the transformation time step. In particular, the system can iteratively update the noisy 3D structure for each training example using the equivariant transformation neural network to generate the predicted 3D structure for the training example following the process 600 described above with reference to FIG. 6.

[0182] The system can update parameters of the equivariant generative machine learning model to optimize the objective function for the generative machine learning model (step 708). In particular, the system can update neural network parameters of the equivariant transformation neural network to optimize the objective function of the equivariant generative machine learning model.

[0183] For each of the training examples, the objective function for the equivariant generative machine learning model can measure errors between the predicted 3D structures generated by the equivariant generative machine learning model over the transformation time steps for the training example and the target 3D structure for the training example. For example, for each training example, the objective function can measure an error between (i) the transformation outputs generated by the equivariant transformation neural network for the training example and (ii) a target transformation output. As a particular example, for each transformation time step for each training example, the objective function can be defined following:

[0184] Where t is a time associated with the transformation time step, xtis the current predicted 3D structure for the training example as of the transformation time step, ue(xt, t) is the transformation output from the equivariant transformation neural network at the transformation time step, x is the target 3D structure for the training example, and G is the added noise used to generate the noisy 3D structure for the training example.

[0185] In some implementations, the objective function for the equivariant generative machine learning model can scale the error between (i) the predicted transformation output and (ii) the target transformation output using a weighting value determined by processing data identifying the time step using a weighting function, w that defines a respective weighting value for each time step. For example, for each transformation time step for each training example, the objective function can be defined following:

[0186] Where w(t) is the weighting value for the transformation time step.

[0187] As part of updating the neural network parameters of the equivariant transformation neural network, the system can determine gradients of the objective function (e.g., gradients of the objective function with respect to the parameters of the equivariant transformation neural network). The system can update the parameters of the equivariant transformation neural network using the gradients of the objective function following any appropriate machine learning technique (e g., following stochastic gradient descent, ADAM, etc.).

[0188] The system can determine whether the training is complete (step 710). The system can use any of a variety of criteria to determine whether the training is complete. For example, the system can determine that training is complete after a pre-determined number of training iterations. Asanother example, the system can determine that training is complete when a value of the objective function for the equivariant generative machine learning model falls below a pre-determined threshold. As another example, the system can determine that training is complete when a difference between values of the training objective function for the current training iteration and a previous training iteration falls below a pre-determined threshold.

[0189] If the system determines that training is not complete, the system can continue to a next training iteration (e.g., return to step 704)

[0190] When the system determines that training is complete, the system can provide the trained equivariant generative machine learning model (step 712).

[0191] FIG. 8 is a flow diagram of an example process for screening a set of candidate ligands. For convenience, the process 800 will be described as being performed by a system of one or more computers located in one or more locations. For example, a prediction system, e.g., the prediction system 116 of FIG. 1, appropriately programmed in accordance with this specification, can perform the process 800.

[0192] The system can obtain data characterizing a protein and a set of candidate ligands (e.g., a set of candidate ligands for binding with the protein) (step 802). The set of candidate ligands can include, e.g., small molecule ligands (e.g., ligands that are non-polymer organic molecules with low molecular weights), poly-peptide ligands (e.g., proteins), nucleic acids (e.g., ribonucleic acids (RNA), deoxyribonucleic acids (DNA), etc.) and so on.

[0193] In general, the data characterizing the protein and the set of candidate ligands can specify chemical compositions and 3D structures of the protein and each of the candidate ligands. For example, the data characterizing the protein and the set of candidate ligands can specify respective chemical elements and spatial positions for individual atoms of the protein (or of a candidate ligand). As another example, the data characterizing the protein and the set of candidate ligands can specify respective chemical compositions, spatial positions, and spatial orientations for molecular components of the protein (or of a candidate ligand), such as functional groups, amino acid residues, nucleotides, and so on.

[0194] As an example, the data characterizing the protein and the set of candidate ligands can include data specifying atoms, charges on atoms, bonds, functional groups, amino acids, branches (e.g., from a molecule main chain), ring structures, and so forth within the protein and the candidate ligands. For example, the data characterizing the protein and the set of candidate ligands caninclude, e.g., Simplified Molecular Input Line Entry System (SMILES) formatted data, International Chemical Identifier (InChi) formatted data, and so on specifying some or all of the protein, the candidate ligands, or both (e.g., specifying compositions of molecular components of the protein, the candidate ligands, or both, such as amino acid residues, nucleotides, etc.).

[0195] The data characterizing the protein and the set of candidate ligands can specify the chemical composition of the protein by specifying an amino acid sequence for the protein. The amino acid sequence for the protein can specify an identity and an ordering for each of a plurality of amino acid residues within the protein. Similarly, when a candidate ligand is a protein itself, the data characterizing the protein and the set of candidate ligands can specify the chemical composition of the candidate ligand by specifying an amino acid sequence for the candidate ligand.

[0196] When a candidate ligand is a nucleic acid, the data characterizing the protein and the set of candidate ligands can specify the chemical composition of the candidate ligand by specifying a nucleotide sequence for the candidate ligand. The nucleotide sequence for the candidate ligand can specify an identity and an ordering for each of a plurality of nucleotides within the candidate ligand.

[0197] In general, the data characterizing the protein and the set of candidate ligands can include data characterizing a known or predicted 3D structure for the protein. The known or predicted 3D structure of the protein can be determined by any of a variety of means, such as by physical experiment (eg., x-ray crystallography), by computational simulation (e.g., by molecular dynamics simulations, quantum mechanical simulations, etc.), by machine learning techniques (eg., using a structure prediction machine learning model), and so on. Similarly, the data characterizing the protein and the set of candidate ligands can include data characterizing known or predicted 3D structures for each of the candidate ligands (e.g., as determined by physical experiment, computational simulation, machine learning techniques, etc.).

[0198] For each of the candidate ligands, the system can then process the data characterizing the protein and the candidate ligand using the generative machine learning model to generate a predicted joint 3D structure of a protein-ligand complex that includes the protein and the candidate ligand (804). In particular, the system generate the predicted joint 3D structures of the proteinligand complexes for the candidate ligands using an equivariant generative machine model following the process 600 described above with reference to FIG. 6.

[0199] In some cases, the system can generate the predicted 3D structures of the protein-ligand complexes for the candidate ligands substantially in parallel by loading the data characterizing the protein and the candidate ligands into a local memory of a computing unit (e.g., a graphics processing unit, a tensor processing unit, a central processing unit, an application specific integrated circuit, etc.) and performing the operations of the process 600 of FIG. 6 for each of the candidate ligands substantially in parallel using the computing unit.

[0200] The system can then use the predicted joint 3D structures of the protein-ligand complexes for the candidate ligands to predict molecule properties of the candidate ligands (step 806). As described above with reference to FIG. 1, the system can determine any of a variety of predicted molecule properties for the candidate ligands based, at least in part, on the predicted joint 3D structures of the protein-ligand complexes for the candidate ligands. For example, the system can determine predicted binding affinities for protein-ligand complexes of the protein with each of the candidate ligands, predicted energies (e g., predicted binding energies) of the protein-ligand complexes of the protein and with each of the candidate ligands, and so on.

[0201] The predicted three-dimensional (3D) structure of a protein-ligand complex can be used to estimate the binding affinity of the ligand for the protein by analyzing structural and energetic features of the complex. For example, the predicted structure can be used to compute physical or statistical interaction scores, such as hydrogen bonding patterns, hydrophobic contact areas, electrostatic complementarity, or binding free energy estimates derived from molecular mechanics or machine-learning models trained on known complexes.

[0202] The system can determine a ranking of the candidate ligands based on the predicted molecule properties for the candidate ligands (step 808). In particular, the system can use the predicted molecule properties for the candidate ligands to determine a respective binding score for each of the candidate ligands with the protein and can then rank the candidate ligands based on their respective binding scores. The binding scores for the candidate ligands can characterize, e.g., binding affinities, potencies, selectivities, and so on of the candidate ligands for the protein.

[0203] The system can use the ranking of the candidate ligands to screen the set of candidate ligands (step 810). For example, the system can use the ranking of the candidate ligands to screen the candidate ligands for further validation (e.g., by physical experimentation using synthesized samples of selected candidate ligands, computational validation using computational simulation of selected candidate ligands, etc ).

[0204] The system can select candidate ligands for the purpose of performing any of a variety of physical experiments and computational simulations for the selected candidate ligands. For example, the system can screen the candidate ligands for the purpose of physically synthesizing the selected candidate ligands to perform physical experiments regarding, e.g., the binding affinities, 3D structures, and so on, of the candidate ligands or of protein-ligand complexes that include the candidate ligands. As another example, the system can screen the candidate ligands for the purpose of physically synthesizing selected candidate ligands to perform physical experiments regarding, e.g., absorption properties, distribution properties, metabolism properties, excretion properties, toxicity properties, and so on of the selected candidate ligands.

[0205] As another example, the system can the system can use the ranking of the candidate ligands to select candidate ligands for inclusion within a drug. For example, the system can select a candidate ligand for inclusion within a drug targeting the protein based on the rankings of the candidate ligands determined based on binding scores of the candidate ligands with the protein.

[0206] This specification uses the term “configured” in connection with systems and computer program components. For a system of one or more computers to be configured to perform particular operations or actions means that the system has installed on it software, firmware, hardware, or a combination of them that in operation cause the system to perform the operations or actions. For one or more computer programs to be configured to perform particular operations or actions means that the one or more programs include instructions that, when executed by data processing apparatus, cause the apparatus to perform the operations or actions.

[0207] Embodiments of the subject matter and the functional operations described in this specification can be implemented in digital electronic circuitry, in tangibly-embodied computer software or firmware, in computer hardware, including the structures disclosed in this specification and their structural equivalents, or in combinations of one or more of them. Embodiments of the subject matter described in this specification can be implemented as one or more computer programs, i.e., one or more modules of computer program instructions encoded on a tangible non-transitory storage medium for execution by, or to control the operation of, data processing apparatus. The computer storage medium can be a machine-readable storage device, a machine-readable storage substrate, a random or serial access memory device, or a combination of one or more of them. Alternatively or in addition, the program instructions can be encoded on an artificially-generated propagated signal, e.g., a machine-generated electrical, optical, orelectromagnetic signal, that is generated to encode information for transmission to suitable receiver apparatus for execution by a data processing apparatus.

[0208] The term “data processing apparatus” refers to data processing hardware and encompasses all kinds of apparatus, devices, and machines for processing data, including by way of example a programmable processor, a computer, or multiple processors or computers. The apparatus can also be, or further include, special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application-specific integrated circuit). The apparatus can optionally include, in addition to hardware, code that creates an execution environment for computer programs, e.g., code that constitutes processor firmware, a protocol stack, a database management system, an operating system, or a combination of one or more of them.

[0209] A computer program, which may also be referred to or described as a program, software, a software application, an app, a module, a software module, a script, or code, can be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages; and it can be deployed in any form, including as a stand-alone program or as a module, component, subroutine, or other unit suitable for use in a computing environment. A program may, but need not, correspond to a file in a file system. A program can be stored in a portion of a file that holds other programs or data, e.g., one or more scripts stored in a markup language document, in a single file dedicated to the program in question, or in multiple coordinated files, e.g., files that store one or more modules, sub-programs, or portions of code. A computer program can be deployed to be executed on one computer or on multiple computers that are located at one site or distributed across multiple sites and interconnected by a data communication network.

[0210] In this specification the term “engine” is used broadly to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine will be implemented as one or more software modules or components, installed on one or more computers in one or more locations. In some cases, one or more computers will be dedicated to a particular engine; in other cases, multiple engines can be installed and running on the same computer or computers.

[0211] The processes and logic flows described in this specification can be performed by one or more programmable computers executing one or more computer programs to perform functions by operating on input data and generating output. The processes and logic flows can also beperformed by special purpose logic circuitry, e.g., an FPGA or an ASIC, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0212] Computers suitable for the execution of a computer program can be based on general or special purpose microprocessors or both, or any other kind of central processing unit. Generally, a central processing unit will receive instructions and data from a read-only memory or a random access memory or both. The essential elements of a computer are a central processing unit for performing or executing instructions and one or more memory devices for storing instructions and data. The central processing unit and the memory can be supplemented by, or incorporated in, special purpose logic circuitry. Generally, a computer will also include, or be operatively coupled to receive data from or transfer data to, or both, one or more mass storage devices for storing data, e.g., magnetic, magneto-optical disks, or optical disks. However, a computer need not have such devices. Moreover, a computer can be embedded in another device, e.g., a mobile telephone, a personal digital assistant (PDA), a mobile audio or video player, a game console, a Global Positioning System (GPS) receiver, or a portable storage device, e.g., a universal serial bus (USB) flash drive, to name just a few.

[0213] Computer-readable media suitable for storing computer program instructions and data include all forms of non-volatile memory, media and memory devices, including by way of example semiconductor memory devices, e.g., EPROM, EEPROM, and flash memory devices; magnetic disks, e.g., internal hard disks or removable disks; magneto-optical disks; and CD-ROM and DVD-ROM disks.

[0214] To provide for interaction with a user, embodiments of the subject matter described in this specification can be implemented on a computer having a display device, e.g., a CRT (cathode ray tube) or LCD (liquid crystal display) monitor, for displaying information to the user and a keyboard and a pointing device, e.g., a mouse or a trackball, by which the user can provide input to the computer. Other kinds of devices can be used to provide for interaction with a user as well; for example, feedback provided to the user can be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback; and input from the user can be received in any form, including acoustic, speech, or tactile input. In addition, a computer can interact with a user by sending documents to and receiving documents from a device that is used by the user; for example, by sending web pages to a web browser on a user’s device in response to requests received from the web browser. Also, a computer can interact with a user by sending text messages or other formsof message to a personal device, e.g., a smartphone that is running a messaging application, and receiving responsive messages from the user in return.

[0215] Data processing apparatus for implementing machine learning models can also include, for example, special-purpose hardware accelerator units for processing common and computeintensive parts of machine learning training or production, i.e., inference, workloads.

[0216] Machine learning models can be implemented and deployed using a machine learning framework, e.g., a TensorFlow framework, or a Jax framework.

[0217] Embodiments of the subject matter described in this specification can be implemented in a computing system that includes a back-end component, e.g., as a data server, or that includes a middleware component, e.g., an application server, or that includes a front-end component, e.g., a client computer having a graphical user interface, a web browser, or an app through which a user can interact with an implementation of the subject matter described in this specification, or any combination of one or more such back-end, middleware, or front-end components. The components of the system can be interconnected by any form or medium of digital data communication, e.g., a communication network. Examples of communication networks include a local area network (LAN) and a wide area network (WAN), e.g., the Internet.

[0218] The computing system can include clients and servers. A client and server are generally remote from each other and typically interact through a communication network. The relationship of client and server arises by virtue of computer programs running on the respective computers and having a client-server relationship to each other. In some embodiments, a server transmits data, e.g., an HTML page, to a user device, e.g., for purposes of displaying data to and receiving user input from a user interacting with the device, which acts as a client. Data generated at the user device, e.g., a result of the user interaction, can be received at the server from the device.

[0219] While this specification contains many specific implementation details, these should not be construed as limitations on the scope of any invention or on the scope of what may be claimed, but rather as descriptions of features that may be specific to particular embodiments of particular inventions. Certain features that are described in this specification in the context of separate embodiments can also be implemented in combination in a single embodiment. Conversely, various features that are described in the context of a single embodiment can also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as acting in certain combinations and even initially be claimed assuch, one or more features from a claimed combination can in some cases be excised from the combination, and the claimed combination may be directed to a subcombination or variation of a subcombination.

[0220] Similarly, while operations are depicted in the drawings and recited in the claims in a particular order, this should not be understood as requiring that such operations be performed in the particular order shown or in sequential order, or that all illustrated operations be performed, to achieve desirable results. In certain circumstances, multitasking and parallel processing may be advantageous. Moreover, the separation of various system modules and components in the embodiments described above should not be understood as requiring such separation in all embodiments, and it should be understood that the described program components and systems can generally be integrated together in a single software product or packaged into multiple software products.

[0221] Particular embodiments of the subject matter have been described. Other embodiments are within the scope of the following claims. For example, the actions recited in the claims can be performed in a different order and still achieve desirable results. As one example, the processes depicted in the accompanying figures do not necessarily require the particular order shown, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous.

[0222] What is claimed is:

Claims

CLAIMS1. A method performed by one or more computers, the method comprising:obtaining data identifying a molecular system one or more molecules;processing the data identifying the molecular system using a generative machine learning model parametrized by a transformation neural network to generate data defining a predicted three-dimensional (3D) structure of the molecular system, comprising:iteratively updating data defining a current predicted 3D structure of the molecular system over a plurality of time steps, comprising, at each time step:obtaining data defining a current predicted 3D structure at the time step; processing a network input based on the current predicted 3D structure at the time step using the transformation neural network to generate a transformation output defining a modification to the current predicted 3D structure, comprising:generating a respective atom embedding for each atom in the molecular system based at least in part on the current predicted 3D structure;updating the atom embeddings using an equivariant attention operation; andprocessing the updated atom embeddings to generate the transformation output; andupdating the current predicted 3D structure based on the transformation output generated by the transformation neural network;wherein the predicted 3D structure of the molecular system is defined by the current predicted 3D structure after a last time step of the plurality of time steps.

2. The method of claim 1, wherein the equivariant attention operation is equivariant to rotations of the current predicted 3D structure.

3. The method of any preceding claim, wherein generating the respective atom embedding for each atom in the molecular system based at least in part on the current predicted 3D structure comprises, for each atom:generating a spatially-independent embedding of the atom based on properties of the atom that are invariant to a spatial position of the atom; andgenerating a spatially-dependent embedding of the atom based on at least a 3D spatialposition of the atom;wherein the spatially-independent embedding and the spatially-dependent embedding collectively define the embedding of the atom.

4. The method of claim 3, wherein for each atom, the spatially-independent embedding of the atom is based on one or more of: an elemental type of the atom, a charge of the atom, an atomic mass of the atom, an electron configuration of the atom, an electronegativity of the atom, an ionization energy of the atom, or an electrostatic potential of the atom.

5. The method of any one of claims 3-4, wherein for each atom, the spatially-dependent embedding of the atom comprises three spatial dimensions and a plurality of channel dimensions.

6. The method of any one of claims 3-5, wherein updating the atom embeddings using an equivariant attention operation comprises:generating, for each atom, respective query, key, and value embeddings representing the spatially-independent embedding of the atom;generating, for each atom, respective query, key, and value embeddings representing the spatially-dependent embedding of the atom using equivariant linear neural network layer operations, wherein the query, key, and value embeddings representing the spatially-dependent embedding of the atom each comprise three spatial dimensions and a plurality of channel dimensions; andupdating the atom embeddings using the equivariant attention operation based on the query, key, and value embeddings representing the spatially-independent and spatially-dependent embeddings of the atoms.

7. The method of claim 6, wherein updating the atom embeddings using the equivariant attention operation based on the query, key, and value embeddings representing the spatially-independent and spatially-dependent embeddings of the atoms comprises:generating an array of attention weights, wherein each attention weight corresponds to a respective pair of atoms and is generated by performing operations comprising:determining a spatially-independent attention weight based on a scalar inner product between: (i) a query embedding representing a spatially-independent embedding of a first atom in the pair of atoms, and (ii) a key embedding representing a spatially-independentembedding of a second atom in the pair of atoms;determining a spatially-dependent attention weight based on a Frobenius inner product between: (i) a query embedding representing a spatially-dependent embedding of the first atom in the pair of atoms, and (ii) a key embedding representing a spatially-dependent embedding of the second atom in the pair of atoms; andgenerating the attention weight for the pair of atoms by combining the spatially-independent attention weight and the spatially-dependent attention weight; andupdating the atom embeddings by combining the value embeddings representing the spatially-independent embeddings and spatially-dependent embeddings of the atoms using the array of attention weights.

8. The method of any one of claims 3-7, further comprising, after updating the atom embeddings using the equivariant attention operation:updating the atom embeddings using an equivariant gated feedforward neural network layer.

9. The method of claim 8, wherein updating the atom embeddings using an equivariant gated feedforward neural network layer comprises, for each atom:generating an intermediate spatially-dependent embedding by combining: (i) the spatially-dependent embedding of the atom, and (ii) a result of processing the spatially-independent embedding of the atom using a feedforward neural network layer; and replacing the spatially-dependent embedding by a result of processing the intermediate spatially-dependent embedding by an equivariant linear neural network layer.

10. The method of any preceding claim, further comprising:generating a respective token embedding for each token in the molecular system based at least in part on the current predicted 3D structure, wherein one or more of the tokens represent: amino acids, nucleic acids, functional groups, or multi-atom fragments;updating the token embeddings using the equivariant attention operation; and processing both the updated atom embeddings and the updated token embeddings to generate the transformation output.

11. The method of claim 10, wherein the equivariant attention operation over the token embeddings is a global attention operation; andwherein the equivariant attention operation over the atom embeddings is a local selfattention operation.

12. The method of any one of claims 10-11, wherein processing both the updated atom embeddings and the updated token embeddings to generate the transformation output comprises:updating the atom embeddings using an equivariant cross-attention operation over the token embeddings.

13. The method of any one of claims 3-12, wherein processing the updated atom embeddings to generate the transformation output comprises:processing only the spatially-dependent embeddings of the atoms to generate, for each atom, an output defining an update to a 3D spatial position of the atom.

14. The method of any one of claims 3-13, further comprising processing only the spatially-independent embeddings of the atoms to generate, for each atom, a confidence of the generative machine learning model in a 3D spatial position of the atom the updated predicted 3D structure.

15. The method of any preceding claim, wherein the generative machine learning model is a generative diffusion model, or a flow matching model, or a stochastic interpolant model.

16. The method of any preceding claim, wherein outputting the predicted 3D structure of the molecular system comprises:storing the predicted 3D structure in a memory; ortransmitting the predicted 3D structure over a data communications network; or displaying a visual rendering of the predicted 3D structure on a display of a user device.

17. The method of any preceding claim, further comprising selecting one or more molecules in the molecular system for physical synthesis based at least in part on the predicted 3D structure of the molecular system.

18. The method of claim 17, further comprising physically synthesizing the one or more molecules selected for physical synthesis.

19. A method performed by one or more computers, comprising:obtaining data defining a protein and a set of candidate ligands;generating, for each candidate ligand in the set of candidate ligands, a predicted 3D structure of a molecular system comprising the protein and the candidate ligand by performing the method of any one of claims 1-16; anddetermining, for each candidate ligand in the set of candidate ligands, a predicted binding affinity of the candidate ligand for the protein based at least in part on the predicted 3D structure of the molecular system comprising the protein and the candidate ligand.

20. The method of claim 19, further comprising selecting one or more of the candidate ligands for targeting the protein based at least in part on the predicted binding affinities of the candidate ligands for the protein.

21. The method of claim 20, wherein selecting one or more candidate ligands for target the protein comprises selecting one or more candidate ligands for inclusion in a drug targeting the protein.

22. The method of any one of claims 19-21, further comprising selecting one or more of the candidate ligands for physical synthesis based at least in part on the predicted binding affinities of the candidate ligands for the protein.

23. The method of claim 22, further comprising physically synthesizing the one or more candidate ligands selected for physical synthesis.

24. A system comprising:one or more computers; andone or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the method of any one of claims 1-23.

25. One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations of the method of any one of claims 1-23.

Images