Modeling physical systems with large language machine-learned models
A predictive system may access a set of physical structures corresponding to a physical system. Each physical structure representative of a configuration. The predictive system may encode the accessed physical structures to produce a set of encoded physical structures by encoding, for each accessed physical structure, a position of each constituent unit of the physical system within the accessed physical structure. The predictive system may train a machine-learned model using the encoded physical structures. The predictive system may retrain the machine-learned model by iteratively: accessing a set of two or more candidate physical structures, determining a first energy difference among the set of candidate physical structures, obtaining a second energy difference between a set of physical structures corresponding to the set of candidate physical structures using a method to calculate reference energy values.
1 . A method comprising:
accessing, by a predictive system, a set of physical structures corresponding to a physical system, each physical structure representative of a configuration of the physical system;
encoding, by the predictive system, the accessed physical structures to produce a set of encoded physical structures by encoding, for each accessed physical structure, a position of each constituent unit of the physical system within the accessed physical structure;
training, by the predictive system, a machine-learned transformer model, the machine-learned transformer model comprising one or more attention blocks configured to model interactions of constituent units by generating activation outputs as relative energy values that correspond to a physical structure distribution; and
retraining, by the predictive system, the machine-learned transformer model by iteratively:
accessing a set of candidate physical structures for the physical system in a training sample;
converting, through forward propagation of the machine-learned transformer model, the candidate physical structures into a plurality of token sequences and corresponding activation outputs, wherein each token sequence represents a candidate physical structure of the physical system,
determining a first energy difference between two candidate physical structures, the first energy difference being an aggregation of the corresponding activation outputs from the attention blocks;
comparing the first energy difference to a second energy difference between a set of physical structures corresponding to the set of candidate physical structures, wherein the second energy difference is derived from a ground truth distribution in the training sample using an energy function;
backpropagating the second energy difference to the transformer model to adjust one or more parameters in the attention blocks of the machine-learned transformer model; and
retraining the machine-learned transformer model using the set of candidate physical structures, the first energy difference, and the second energy difference; and
receiving a sequence of a target molecule; and
using the one or more attention blocks in the trained machine-learned transformer model to determine energy values over a set of physical structures of the target molecule.
2 . The method of claim 1 , wherein encoding a particular accessed physical structure of the accessed physical structures comprises:
converting one or more structural expressions of the physical system into a sequence string representation of the physical system;
tokenizing the sequence string representation of the physical system to produce a tokenized structural expression;
tokenizing, for each constituent unit in the physical system, coordinates for the constituent unit within the particular accessed physical structure to produce tokenized coordinates for the constituent unit; and
combining the tokenized structural expression and the tokenized coordinates for each constituent unit to produce a tokenized physical structure corresponding to the particular accessed physical structure.
3 . The method of claim 2 , wherein tokenizing coordinates for the constituent unit within the particular accessed physical structure comprises:
pixelating a rendered sphere to produce a set of pixels each corresponding to a location on the surface of the rendered sphere;
tokenizing, for a first constituent unit in the physical system, coordinates at a center of the sphere; and
tokenizing, for each additional constituent unit in the physical system, coordinates corresponding to a pixel selected from the set of pixels based on a location of the additional constituent unit relative to the center of the sphere.
4 . The method of claim 2 , wherein tokenizing coordinates for the constituent unit within the particular accessed physical structure comprises using a Cartesian coordinate system, an xyz coordinate system, an octree coordinate system, a polar coordinate system, a cylindrical coordinate system, or a barycentric coordinate system to generate the coordinates.
5 . The method of claim 2 , wherein the sequence string representation of the physical system includes an ordered set of constituent units, and wherein the coordinates for a constituent unit comprise coordinates relative to a center of a coordinate sphere.
6 . The method of claim 1 , wherein the machine-learned transformer model is iteratively retrained until one or more retraining criteria is satisfied.
7 . The method of claim 6 , wherein the retraining criteria is satisfied when performance measurement of a holdout set starts to decrease.
8 . The method of claim 1 , wherein retraining the machine-learned transformer model comprises modifying weights of one or more layers of the machine-learned transformer model to minimize a difference between the first energy difference and the second energy difference over subsequent iterations.
9 . The method of claim 1 , wherein retraining the machine-learned transformer model comprises modifying weights of one or more layers of the machine-learned transformer model by:
generating a set of token sequences representing two or more physical structures;
store activation outputs corresponding to the tokens in the token sequences;
determining an aggregated energy state for each token sequence; and
determining backpropagated gradients based on comparing the aggregated energy states for the token sequences.
10 . The method of claim 1 , wherein the first energy difference is determined based at least in part on a Boltzmann probability ratio between the set of candidate physical structures.
11 . The method of claim 1 , wherein the machine-learned transformer model is configured to:
determine, for a candidate physical structure, a force associated with each constituent unit in the physical system based on a relative energy associated with a plurality of neighboring positions for the constituent unit.
12 . The method of claim 11 , wherein the machine-learned transformer model is retrained further based on a difference in forces associated with the candidate physical structures.
13 . The method of claim 11 , wherein the force associated with each constituent unit in the machine-learned transformer model is based on a collective gradient of the relative energy associated with the plurality of neighboring positions for the constituent unit, wherein the collective gradient is determined based on a vector to a neighboring position from which a finite difference gradient is calculated.
14 . The method of claim 11 , wherein a determined force associated with a constituent unit in the physical system is based on forces associated with one or more preceding constituent units in a sequence string representation of the physical system.
15 . The method of claim 11 , wherein the relative energies associated with each of the plurality of neighboring positions for the constituent unit are determined using a softmax output of the machine-learned transformer model.
16 . The method of claim 11 , wherein a determined force associated with a constituent unit in the physical system is based on, for each of the plurality of neighboring positions for the constituent unit, the relative energy associated with the neighboring position and a distance between the constituent unit and the neighboring position.
17 . The method of claim 11 , wherein a determined force generated by the machine-learned transformer model associated with a constituent unit in the physical system is based additionally on reference forces applied from a solvent or solution on the constituent unit in the physical system.
18 . The method of claim 11 , wherein a determined force generated by the machine-learned transformer model associated with a constituent unit in the physical system is based additionally on inherent forces in a physical system comprising the physical system.
19 . The method of claim 1 , wherein a constituent unit in an encoded physical structure is encoded with multiple tokens, each token representing decreasing significance of the constituent unit's coordinates.
20 . The method of claim 1 , further comprising fine tuning, by the predictive system, the machine-learned transformer model by using experimentally determined structures of the physical system, optionally wherein the experimentally determined structures comprises a crystallographic structure of the physical system.
21 . The method of claim 1 , further comprising:
generating, by the predictive system, a probability distribution of structures for a physical system or a portion thereof.
22 . The method of claim 21 , further comprising:
comparing a target probability distribution and the generated probability distribution to predict a binding affinity between a first physical system or a portion thereof and a second physical system or a portion thereof.
23 . The method of claim 1 , wherein obtaining the second energy difference comprises performing a quantum mechanical energy calculation.
24 . The method of claim 1 , wherein encoding an accessed physical structure comprises using a roto-translational transformation to encode spatial data of one or more constituent units in the physical system.
25 . The method of claim 1 , wherein the machine-learned transformer model is a pretrained machine-learned language model and retraining the machine-learned transformer model comprises fine tuning the machine-learned transformer model using training samples of physical systems.
26 . The method of claim 1 , further comprising:
computing importance weights based on the first and second energy differences, determining confidence levels, or selecting configurations based on the first and second energy differences.
27 . The method of claim 1 , wherein encoding the accessed physical structures comprises generating a series of tokens, and generating the series of tokens comprises:
generating, for a first constituent unit, a first set of tokens; and
generating, for a second constituent unit, a second set of tokens based on the first set of tokens.
28 . The method of claim 1 , further comprising:
generating, using the machine-learned transformer model, a molecular sequence that is predicted to be binding to the target molecule.