Training a Graph Neural Network Using a Denoising Objective

Abstract

A method, system, and apparatus, including a computer program encoded on a computer storage medium, for training a neural network including one or more graph neural network layers. In one aspect, the method includes generating data defining a graph including generating a respective final feature representation for each node, where for each of one or more of the nodes, the respective final feature representation is a modified feature representation generated from the respective feature representation for the node using a respective noise, processing the data defining the graph using one or more of the graph neural network layers of the neural network to generate a respective updated node embedding for each node, and for each of one or more of the nodes having the modified feature representation, processing the updated node embedding for the node to generate a respective denoised prediction for the node.

Description

[Technical field]

[0001] CROSS-REFERENCE TO RELATED APPLICATIONS This application claims the benefit of the filing date of U.S. Provisional Patent Application No. 63 / 194,851, entitled “TRAINING GRAPH NEURAL NETWORKS USING A DE-NOISING OBJECTIVE,” filed May 28, 2021, which is incorporated by reference in its entirety herein.

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

[0003] A machine learning model receives input and generates an output, e.g., a predicted output, based on the received input. Some machine learning models are parametric models and generate an output based on the received input and the 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. [Prior art documents] [Non-patent literature]

[0005] [Non-Patent Document 1] Chanussot et al., “The Open Catalyst 2020 (OC20) Dataset and Community Challenges,” ACS Catalysis, 6059~6072, 2020, arXiv:2010.09990) [Non-Patent Document 2] Ramakrishnan et al., “Quantum chemistry structures and properties of 134 kilo molecules,” Sci Data 1, 140022 (2014) [Non-Patent Document 3] Hu et al., “Open Graph Benchmark: Datasets for Machine Learning on Graphs,” arXiv:2005.00687 Summary of the Invention [Means for solving the problem]

[0006] This specification describes a training system, generally implemented as a computer program on one or more computers at one or more locations, that trains a neural network that includes one or more graph neural network layers.

[0007] As used throughout this specification, a "graph" refers to a data structure that includes at least (i) a set of nodes and (ii) a set of edges. Each edge in the graph can connect a respective pair of nodes in the graph. A graph can be a "directed" graph, i.e., each edge connecting a pair of nodes is defined as pointing from a first node to a second node or vice versa, or an "undirected" graph, i.e., edges (or pairs of edges directed in opposite directions) are not associated with a direction.

[0008] In general, data defining a graph may include data defining the nodes and edges of the graph and may be represented in any suitable numerical format. For example, a graph may be defined by data including a listing of tuples {(i,j)}, where each tuple (i,j) represents an edge in the graph connecting node i and node j. Moreover, each edge in the graph may be associated with a set of one or more edge features, and each node in the graph may be associated with a set of one or more node features.

[0009] In one aspect, a method for training a neural network including one or more graph neural network layers is described. The method includes generating data defining a graph comprising (i) a set of nodes, (ii) a node embedding for each node, and (iii) a set of edges each connecting a respective pair of nodes. In an implementation, this includes obtaining a respective initial feature representation for each node, generating a respective final feature representation for each node, where for each of one or more of the nodes, the respective final feature representation is a modified feature representation generated from the respective feature representation for the node using a respective noise, and generating data defining the graph using the respective final feature representation of the nodes. In an implementation, the node embedding for each node is generated from the respective final feature representation of the node.

[0010] The method processes data defining the graph using one or more of the graph neural network layers of the neural network to generate a respective updated node embedding for each node. The method processes the updated node embeddings of the nodes to generate, for each of one or more of the nodes having a modified feature representation, a respective denoised prediction for the node that characterises a denoised representation for the node that does not include the noise used to generate the modified feature representation for the node. The method determines updates to current values ​​of neural network parameters of the neural network to optimise an objective function measuring error in the denoised prediction for the node, in particular to optimise the respective denoised prediction for the node.

[0011] In some implementations, for each of one or more of the nodes having a modified representation, the respective denoising prediction for the node predicts the noise that was used to generate the modified representation of the node.

[0012] In some implementations, for each of one or more of the nodes having a modified feature representation, a respective denoising prediction for the node predicts a respective initial feature representation for the node.

[0013] In some implementations, for each of one or more of the nodes having a modified representation, a respective denoised prediction for the node characterizes the target representation of the node.

[0014] In some implementations, for each of one or more of the nodes having a modified feature representation, the respective denoising prediction for the node predicts an incremental feature representation for the node that, when added to the modified feature representation for the node, results in the target feature representation for the node.

[0015] In some implementations, the method further comprises processing the updated node embeddings of the nodes to generate a task prediction, where the objective function also measures an error in the task prediction.

[0016] In some implementations, both (i) the updated node embedding of the node and (ii) the original node embedding of the node, before being updated using the graph neural network layer, are processed to generate a task prediction.

[0017] In some implementations, the graph represents a molecule and the task prediction is a prediction of the equilibrium energy of the molecule.

[0018] In some implementations, the objective function measures, for each of multiple graph neural network layers of the neural network, a respective error in denoising predictions for the nodes based on updated node embeddings generated by the graph neural network layers.

[0019] In some implementations, for each of one or more of the nodes having a modified feature representation, processing the updated node embedding of the node to generate a respective denoised prediction for the node includes processing the updated node embedding of the node using one or more neural network layers to generate a respective denoised prediction for the node.

[0020] In some implementations, determining updates to current values ​​of neural network parameters of the neural network to optimize the objective function includes backpropagating gradients of the objective function through the neural network parameters of the graph neural network layer.

[0021] In some implementations, for each of one or more of the nodes, a respective final feature representation for the node is generated by adding a respective noise to the respective feature representation for the node.

[0022] In some implementations, generating data defining the graph using the final feature representations of each of the nodes includes, for each pair of nodes comprising a first node and a second node, determining a respective distance between the final feature representation for the first node and the final feature representation for the second node, and determining that each pair of nodes corresponding to a distance less than a predefined threshold is connected by an edge in the graph.

[0023] In some implementations, the graph further comprises a respective edge embedding for each edge.

[0024] In some implementations, generating data defining the graph includes generating an edge embedding for each edge in the graph based at least in part on differences between the final feature representations of each of the nodes connected by the edge.

[0025] In some implementations, the graph represents a molecule, with each node in the graph representing a respective atom in the molecule, and generating the data defining the graph includes generating a node embedding for each node based on the type of atom represented by the node.

[0026] In some implementations, the neural network includes at least 10 graph neural network layers.

[0027] In some implementations, each graph neural network layer of the graph neural network is configured to receive a current graph and update the current graph according to current neural network parameter values ​​of the graph neural network layer, including updating a current node embedding of each of the one or more nodes in the graph based on (i) the node's current node embedding, and (ii) a respective current node embedding of each of the node's one or more neighbors in the graph.

[0028] In some implementations, the current graph comprises an edge embedding for each edge, and updating the current node embedding of each of the one or more nodes in the graph further includes updating the node embedding of the node based at least in part on the respective edge embedding of each of the one or more edges coupled to the node.

[0029] In some implementations, the graph represents a molecule, each node in the graph represents a respective atom in the molecule, the initial feature representation for each node represents an initial spatial position of the corresponding atom in the molecule, and for each of one or more of the nodes having a modified feature representation, the target feature representation for the node represents a final spatial position of the corresponding atom after atomic relaxation.

[0030] According to another aspect, one or more non-transitory computer storage media are provided that store instructions that, when executed by one or more computers, cause the one or more computers to perform operations of the methods described herein.

[0031] According to another aspect, a system is provided comprising one or more computers and one or more storage devices communicatively coupled to the one or more computers, the one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the methods described herein.

[0032] As used throughout this specification, an "embedding" refers to an ordered collection of numbers, for example a vector, matrix, or other tensor of numbers.

[0033] The subject matter described herein can be implemented in particular embodiments to realize one or more of the following advantages.

[0034] The systems described herein can train a graph neural network (i.e., a neural network including one or more graph neural network layers) to generate denoised predictions for nodes in the graph. In particular, before training the graph neural network on the graph, the system can use noise to modify feature representations of nodes in the graph, e.g., by adding noise to the feature representations of the nodes. The denoised predictions can, for example, predict the value of the noise that modified the feature representations of the nodes, or can predict a reconstruction of the original (i.e., before the feature representations were modified using noise) feature representations of the nodes in the graph.

[0035] Training a graph neural network to generate denoised predictions can normalize the training of the graph neural network, and in particular can enable effective training of graph neural networks with a large number of graph neural network layers, e.g., more than 100 graph neural network layers. In contrast, many conventional systems are limited to training graph neural networks with a much smaller number of graph neural network layers (e.g., fewer than 10 layers) before the performance of the graph neural network saturates or even degrades with the addition of more graph neural network layers. Deeper graph neural networks (i.e., with more graph neural network layers), when trained by the system described herein, can achieve higher prediction accuracy for more complex prediction tasks than would be achievable using shallower graph neural networks (i.e., with fewer graph neural network layers).

[0036] Generating denoised predictions requires each node embedding to encode unique information to denoise the node representation, which can, for example, mitigate the effects of "over-smoothing," where node embeddings become nearly identical after being processed through several graph neural network layers. Moreover, training a graph neural network to generate denoised predictions can reduce the likelihood of "over-fitting," for example, because noise added to the node representations in the graph prevents the graph neural network from memorizing the original node representations. Training a graph neural network to generate denoised predictions also encourages the graph neural network to implicitly learn the distribution of the "real world," i.e., graphs whose node representations are not modified, and the graph neural network can exploit this implicit knowledge to achieve higher accuracy in "task" prediction. The described technique works differently from other techniques that involve missing node or edge features, and can be combined with these other techniques. The system described herein thus enables more efficient use of computational resources (e.g., memory and computing power) by enabling effective training of deeper graph neural networks that achieve higher accuracy while mitigating the effects of over-smoothing and over-fitting.

[0037] The details of one or more embodiments of the subject matter herein 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, drawings, and claims. [Brief description of the drawings]

[0038] [Figure 1] FIG. 1 is a block diagram of an example training system for training a graph neural network. [Diagram 2]FIG. 2 illustrates an example of operations that may be performed by a training system. [Diagram 3] FIG. 1 is a flow diagram of an example process for training a graph neural network using the training system. [Figure 4] FIG. 1 illustrates exemplary experimental results. [Diagram 5] FIG. 1 illustrates exemplary experimental results. [Figure 6] FIG. 1 illustrates exemplary experimental results. DETAILED DESCRIPTION OF THE PREFERRED EMBODIMENTS

[0039] Like reference numbers and designations in the various drawings indicate like elements.

[0040] 1 is a block diagram of an example training system 100 for training a graph neural network 150, e.g., a neural network including one or more graph neural network layers. System 100 is an example of a system implemented as a computer program on one or more computers at one or more locations in which the systems, components, and techniques described below are implemented.

[0041] The system 100 can train the graph neural network 150 by using the graph neural network 150 to generate denoised predictions 108 for nodes in the graph 104. In general, a "graph" refers to a data structure that includes at least (i) a set of nodes and (ii) a set of edges. For example, a graph can be represented as, e.g., G=(V,E), where V is a set of nodes and E is a set of edges. Each edge in the graph can connect a pair of nodes in the graph. In some implementations, the graph G can additionally be represented using graph-level properties g, e.g., G=(V,E,g), where the graph-level properties g can include any suitable aspect of the overall system represented by the graph 104. As a particular example, the graph 104 can represent a physical system, e.g., a molecule, and each node in the graph 104 can represent, e.g., an atom in the molecule, and the graph-level properties can include, e.g., the energy of the molecule. In general, the graph 104 may represent any suitable type of system, such as a collection of particles, a point cloud, or a social network. Examples of systems that may be represented by the graph 104 are described in more detail below.

[0042] The training system 100 may include a noise engine 160 that may be configured to process the initial feature representation 102 for each node in the graph 104 and generate a respective final feature representation 112 for each node in the graph 104. In general, a "feature representation" for a node may characterize any suitable aspect of the element represented by the node. For example, in the case of a graph representing molecules, the initial feature representation 102 for a node may include, for example, the spatial positions (e.g., x, y, and z coordinates) of atoms in the molecule represented by the node. In some implementations, the initial feature representation 102 for a node may include an ordered set of features of the element represented by the node. For example, the initial feature representation for a node may include an ordered set of features of the element represented by the node (e.g.,

[0043]

number

[0044] , the spatial position of the atom in the molecule) and the type of atom represented by the node from a set of possible atom types (e.g., including one or more of carbon, oxygen, nitrogen, etc.).

[0045] For one or more of the nodes, the noise engine 160 may generate the final feature representation 112 by modifying (at least a portion of) the initial feature representation 102 with a respective noise, e.g., Gaussian noise. In some implementations, the noise engine 160 may randomly sample a respective noise value for each of the one or more nodes from a distribution, e.g., a Gaussian distribution, and may add the respective noise value to the respective initial feature representation 102 for the node to generate the final feature representation 112. For the remaining nodes in the graph 104, the final feature representation may be the same as the initial feature representation. As a particular example, the noise engine 160 may modify the initial feature representation 102 for some or all of the nodes in the graph 104 as follows:

[0046]

number

[0047] However, v i is the initial feature representation for node i, and σ i is the noise for node i,

[0048]

number

[0049] is the final feature representation for node i. As a particular example, if the graph 104 represents a molecule, each node in the graph 104 represents an atom in the molecule, and the initial feature representation 102 for the nodes in the graph 104 includes (i) the spatial location of the atom represented by the node, and (ii) the type of the atom represented by the node, the noise engine 160 can generate the final feature representation 112 for the node by adding (or otherwise synthesizing) noise with features representing the spatial location of the atom represented by the node. That is, for each node in the graph, the noise engine 160 can use noise to perturb the features representing the spatial location of the corresponding atom, such that the final feature representation for the node defines a perturbed spatial location for the corresponding atom. Each initial feature representation for each node can further include, for example, a feature that defines the type of atom represented by the node, and the noise engine 160 can optionally refrain from synthesizing (e.g., adding) noise to the features representing the atom type.

[0050] In some implementations, the noise engine 160 may generate the final feature representation 112 for a node by scaling the initial feature representation 102 for the node using a respective noise, e.g., a noise value sampled from a Gaussian distribution. In some implementations, the noise engine 160 may modify the initial feature representation 102 for the node using a noise having the same dimensionality as the initial feature representation 102. For example, if the initial feature representation is an N-dimensional vector, the noise may also be an N-dimensional vector. In general, the noise engine 160 may generate the final feature representation 112 for a node in the graph 104 using the respective noise in any suitable manner.

[0051] In an implementation in which the initial feature representation 102 is modified for only a portion of the nodes in the graph 104 (e.g., rather than all nodes in the graph 104), the noise engine 160 may randomly select nodes in the graph 104 for which the initial feature representation 102 is modified.

[0052] The graph neural network 150 may include (i) an encoder 110, (ii) an updater 120, and (iii) a decoder 130, each of which is described in more detail below.

[0053] The noise engine 160 may provide the final feature representations 112 for the nodes in the graph 104 to the encoder 110. The encoder 110 may be configured to generate data defining the graph 104 using the respective final feature representations 112 for the nodes. For example, the encoder 110 may assign respective nodes in the graph 104 to each element in the system represented by the graph 104. The encoder 110 may then instantiate edges between pairs of nodes in the graph 104. In general, the encoder 110 may instantiate edges between pairs of nodes in the graph 104 in any suitable manner.

[0054] In some implementations, the encoder 110 may instantiate edges between pairs of nodes in the graph 104 by determining, for each pair of nodes, a respective distance between the final representations of these nodes. The encoder 110 may then determine that each pair of nodes corresponding to a distance less than a predefined threshold is joined by an edge in the graph 104. The threshold distance may be any suitable numerical value. In some implementations, the encoder 110 may instantiate edges between pairs of nodes in the graph 104 based on the type of system being represented by the graph 104.

[0055] As a particular example, if the graph 104 represents a molecule and each node in the graph 104 represents an atom in the molecule, the encoder 110 may assign edges in the graph 104 between pairs of nodes that correspond to bonds between atoms in the molecule represented by the pair of nodes. In some implementations, the distance between the final feature representations may characterize local interactions between the atoms represented by the nodes. For example, the threshold distance may represent a connectivity radius (R) such that edges connecting pairs of nodes within the connectivity radius represent local interactions of adjacent atoms in the molecule. The search for adjacent nodes in the graph 104 may be performed via any suitable search algorithm, for example, a kd-tree algorithm.

[0056] In addition to assigning nodes and instantiating edges, the encoder 110 may generate a respective node embedding for each node in the graph 104. In general, an "embedding" of an entity may refer to a representation of the entity as an ordered set of numerical values, e.g., a vector or matrix of numerical values.

[0057] The encoder 110 can generate a node embedding for each node by using a node embedding sub-network. The node embedding sub-network of the encoder 110 can process the final feature representation 112 for each node in the graph 104 and can generate a node embedding for each node in the graph 104. As a particular example, if the graph 104 represents a molecule, the node embedding sub-network can generate a node embedding for the node based on, for example, the type of atom represented by the node. As another particular example, the node embedding sub-network can generate a node embedding based on whether the atom is part of an adsorbate or a catalyst, for example, the node embedding can include a 1 for an adsorbate and a 0 for a catalyst.

[0058] In some implementations, in addition to generating a node embedding for each node in the graph 104, the encoder 110 may generate an edge embedding for each edge in the graph 104 using an edge embedding sub-network of the encoder 110. For example, the edge embedding sub-network of the encoder 110 may process the final feature representations 112 for the nodes in the graph 104 and generate an edge embedding for each edge in the graph 104 based at least in part on the difference between the respective final feature representations 112 for the nodes connected by the edge. As a particular example, an embedding e for an edge k connecting a pair of nodes may be generated. k can be expressed as follows:

[0059]

number

[0060] where d is the vector displacement for an edge connecting a pair of nodes, |d| is the distance,

[0061]

number

[0062] are radial Bessel basis functions defined below by equation (3), and concat denotes the concatenation operation.

[0063]

number

[0064] In this manner, the encoder 110 can generate data defining the graph 104 that includes (i) a set of nodes, (ii) a set of edges that each connect a respective pair of nodes, (iii) a node embedding for each node, and optionally (iv) an edge embedding for each edge.

[0065] After generating data defining the graph 104, the encoder 110 may provide the data to the updater 120. The updater 120 may update the graph 104 over multiple internal update iterations to generate the final graph 106. "Updating" a graph refers to performing a step of message passing (e.g., a step of information propagation) between nodes and edges included in the graph, for example, by updating node embeddings and / or edge embeddings for some or all nodes and edges in the graph based on node embeddings and / or edge embeddings of adjacent nodes in the graph. The updater 120 may include one or more graph neural network layers, each of which may be configured to receive a current graph and update the current graph according to current parameters of the graph neural network layer. The updater 120 may include any number of graph neural network layers, for example, 1 layer, 10 layers, 100 layers, or any other suitable number of graph neural network layers. In some implementations, the updater 120 includes at least 10 graph neural network layers.

[0066] In particular, each graph neural network layer may be configured to update the current node embedding of each node in the graph 104 based on (i) the node's current node embedding and (ii) the respective current node embedding of each of the node's one or more neighbors in the graph 104. A pair of nodes in the graph 104 are "adjacent" nodes if they are connected to each other by an edge. In implementations in which the graph 104 additionally includes an edge embedding for each edge, each graph neural network layer may update the node embedding of the node based on the respective edge embedding of each of the one or more edges also connected to the node in the graph 104.

[0067] As a specific example, at each update iteration, each graph neural network layer calculates the current message vector for the edge connecting node u to node v.

[0068]

number

[0069] may be configured to determine as follows:

[0070]

number

[0071] however,

[0072]

number

[0073] is the node embedding of node u in the previous update iteration,

[0074]

number

[0075] is the node embedding of node v in the previous update iteration,

[0076]

number

[0077] and

[0078]

number

[0079] is the message vector for each determined edge in the previous update iteration, and ψt+1 is a message function implemented by the graph neural network layer, for example as a fully connected neural network layer (e.g., the same for each edge). After determining the message vector, in each update iteration, the graph neural network layer calculates the current node embeddings for nodes u connected to node v by edges

[0080]

number

[0081] can be updated as follows:

[0082]

number

[0083] however,

[0084]

number

[0085] is the updated node embedding for the update iteration, and the update function φ t+1 is implemented by a graph neural network layer, for example as a fully connected neural network layer (e.g., the same for each node), and the first sum is the sum of the neighbors N of node v v The second sum is over the total number of neighbors N u This covers a total of .

[0086] The final update iteration of the updater 120 generates data that defines the final graph 106. The final graph 106 may have the same structure as the initial graph 104 but a different node embedding (e.g., the final graph 106 may have the same number of nodes and the same number of edges as the initial graph 104). In some implementations, the final graph 106 may additionally include a different edge embedding.

[0087] The updater 120 may provide data defining the final graph 106 to the decoder 130. The decoder 130 may be configured to process the data defining the final graph 106 to generate a denoised prediction 108 for each of the one or more nodes having a modified feature representation. In particular, for each of the one or more nodes having a modified feature representation, the decoder 130 may process the updated node embeddings for the node using one or more neural network layers to generate a respective denoised prediction 108 for the node. The denoised prediction 108 may characterize a denoised representation for the node, without the noise used to generate the modified feature representation for the node.

[0088] In some implementations, the denoising prediction 108 for a node may predict the noise used to generate the modified feature representation for the node.

[0089] In some implementations, the denoising prediction for a node 108 may predict an initial feature representation for the node. For example, if the graph 104 represents a molecule, the denoising prediction for a node 108 may predict the initial spatial positions of atoms in the molecule represented by the node before the initial spatial positions are modified by using noise.

[0090] In some implementations, the denoised prediction 108 for a node may characterise a target feature representation for the node. The "target feature representation" for a node may characterise any suitable aspect of the element represented by the node in the graph 104. As a particular example, if the initial feature representation 102 for the node includes spatial positions of atoms in a molecule, the target feature representation for the node may include different spatial positions of the atoms in the molecule, e.g., the spatial positions of the atoms after atomic relaxation of the molecule. As another example, if the graph represents a social network and each node in the graph represents a respective user in the social network, the target embedding for each node may, for example, characterise the amount of time (e.g., in minutes) that the corresponding user interacts with the social network over a specified time period (e.g., one day). In one example, the denoised prediction 108 for each node may include an output feature representation that is an estimate of the target feature representation for the node. In another example, the denoised prediction 108 for each node may include a prediction for an incremental feature representation for the node that, when added to the revised feature representation for the node, results in the target feature representation of the node.

[0091] Because generating the denoised prediction 108 requires that each node embedding encodes unique information to denoise the feature representation of the node, this can mitigate the effects of “over-smoothing”, where node embeddings become nearly identical after being processed through several graph neural network layers, for example.

[0092] In some implementations, the decoder 130 can process the updated node embeddings of the nodes to generate a task prediction 109. The task prediction can be, for example, a single output for the input graph 104 or a respective output for each node in the input graph 104. In general, the task prediction 109 can be any suitable prediction that characterizes one or more of the elements represented by the nodes in the graph 104. The task prediction 109 can be, for example, a classification prediction or a regression prediction. The classification prediction can include a respective score for each class in a set of possible classes, where the score for a class can define the likelihood that the set of elements represented by the graph 104 is included in that class. The regression prediction can include one or more numerical values, each drawn from a continuous range of values, that characterize the set of elements represented by the graph 104.

[0093] In one example, to generate a task prediction, the decoder 130 may process (i) the updated node embedding of the node and (ii) the original node embedding of the node before it was updated using the graph neural network layers. As a particular example, the decoder 130 may generate a task prediction y as follows:

[0094]

number

[0095] however,

[0096]

number

[0097] is the updated node embedding of node i,

[0098]

number

[0099] is the original node embedding of node i, |V| is the total number of nodes in the graph, and MLP Update and MLPs Enc are, for example, the fully connected neural network layers of the updater and encoder, respectively, and b Update is the bias term of the updater, and b Enc is the bias term of the encoder, and W Update is the linear neural network layer of the updater, and W Enc is the linear neural network layer of the encoder.

[0100] As a particular example, if the graph 104 represents a molecule, the task prediction 109 may be a prediction of one or more of an equilibrium energy, an internal energy, or a highest occupied molecular orbital (HOMO) energy of the molecule represented by the graph 104. In some implementations, the decoder 130 may process updated node embeddings of less than all nodes to generate the task prediction 109, e.g., in some cases, the decoder 130 may generate the task prediction 109 by processing an updated node embedding of a single node in the graph 104. Examples of task predictions are described in more detail below.

[0101] The encoder 110, the updater 120, and the decoder 130 may have any suitable neural network architecture that enables them to perform their specified functions. For example, the encoder 110, the updater 120, and the decoder 130 may have any suitable neural network layers (e.g., convolutional layers, fully connected layers, recurrent layers, attention layers, graph neural network layers, etc.) with any suitable number of layers (e.g., 2 layers, 5 layers, or 10 layers) and coupled in any suitable configuration (e.g., as a linear arrangement of layers).

[0102] The system 100 may further include a training engine 140 that may train the neural network 150 using the denoised predictions 108. The training engine 150 may evaluate an objective function for one or more of the graph neural network layers of the neural network 150 that measures a respective error in the denoised predictions 108 for the nodes based on the updated node embeddings generated by the graph neural network layers. More specifically, the graph neural network 150 may generate a respective denoised prediction for the nodes in the graph 104 at each of the one or more graph neural network layers, i.e., based on the updated node embeddings generated by the graph neural network layers. In some implementations, the objective function may additionally measure an error in the task predictions 109, for example, using a cross-entropy error measure, a squared error measure, or any other suitable error measure. As a particular example, the objective function L may be expressed as follows: L = λL de-noising +L task (7) However, L de-noising measures the respective error in the denoising prediction for a node, and L task �� measures the error in task prediction, and λ is a weighting factor. The weighting factors and noise vary according to the application and can be optimized using a hyper-parameter sweep; just as an example, the weighting factor can be of order 0.1 and the noise standard deviation can be of order 0.01.

[0103] To optimize the objective function, the training engine 140 may determine the gradient of the objective function with respect to the current values ​​of the neural network parameters, e.g., using a backpropagation technique. The training engine 140 may then use the gradient to update the current values ​​of the neural network parameters, e.g., using any suitable gradient descent optimization technique, e.g., the RMSprop or Adam gradient descent optimization technique. In particular, the training engine 140 may backpropagate the gradient of the objective function through the neural network parameters of the graph neural network layer.

[0104] In some implementations, the training engine 140 first generates a denoised prediction (e.g., L de-noising ), and then training the neural network 150 to optimize the objective function based on both the denoised predictions and the task predictions (e.g., L as defined in equation (7)). For example, training the neural network 150 to generate the denoised predictions 108 can reduce the likelihood of “overfitting” because noise added to the representations of the nodes in the graph 104 prevents the neural network 150 from memorizing the initial node representations.

[0105] In some implementations, the training engine 140 can pre-train the neural network 150 to optimize an objective function based solely on the denoising predictions, and then train the neural network 150 to optimize an objective function based solely on the task predictions.

[0106] In some implementations, the training engine 140 can pre-train the neural network 150 to optimize an objective function based on the denoising prediction and the first task prediction, and then train the neural network 150 to optimize the objective function based only on the second task prediction. The second task prediction may be different from the first task prediction. For example, the first task prediction may include predicting a HOMO energy of a molecule, while the second task prediction may include predicting an equilibrium energy of the molecule.

[0107] Optionally, after pre-training the neural network 150 to optimize the objective function based on the denoising prediction, the training engine 140 can "freeze" some of the parameters of the neural network and then train only the unfrozen parameters of the neural network to the objective function based on the task prediction. (Freezing the parameters of the neural network can refer to designating the current values ​​of the parameters as fixed, static values ​​that are not further modified during training.) For example, the training engine 140 can pre-train the neural network to optimize the objective function based on the denoising prediction, freeze the parameters of the encoder 110 and the updater 120, and then train only the parameters of the decoder to the objective function based on the task prediction.

[0108] Exemplary applications of the system 100 and the trained graph neural network 150 are described in more detail below. In general, using the trained graph neural network 150 involves obtaining feature representations for the nodes, using the feature representations for the nodes to generate data defining the graph 104, and processing the data defining the graph 104 using the graph neural network 150 to generate respective updated node embeddings for each node. Output from the graph neural network may then comprise one or more of features decoded from the updated node embeddings of the graph 104, denoising predictions 108, and task predictions 109, depending on the application.

[0109] In some implementations, the graph 104 may represent one or more molecules, where a "molecule" includes a large slab of atoms, such as the surface of a catalyst. Each node in the graph 104 may then represent a respective atom in the molecule. In general, the feature representation of a node defines the type of atom represented by the node. It may include other features such as the atomic number, whether the atom is an electron or proton donor or acceptor, whether the atom is part of an aromatic system, the degree of hybridization, for example for carbon atoms, and the number of hydrogens bonded to the atom if hydrogen atoms are not explicitly represented. In some implementations, the feature representation of a node may include the (3D) spatial location of the atoms in the molecule. In some implementations, the feature representation of a node does not include the (3D) spatial location of the atoms in the molecule, i.e., the molecule may be defined by bonds and atom types. If the representation does not include spatial location, noise may be added by randomly altering one or more features of the nodes (and optionally one or more features of the edges as well), and the denoised prediction may then comprise a reconstruction of the features. In some implementations, for example, if two or more interacting entities are modeled, the node features may indicate to which of the entities an atom belongs. Thus, in general, the feature representation of the nodes defines the structure and properties (e.g., types of atoms) of the molecule.

[0110] In some implementations, the neural network is trained to identify the resulting structure of the atoms from the initial structure of the atoms in the molecule. The structure may be decoded from the node embedding, e.g., as the (3D) spatial position of the atoms decoded from the node embedding, or as bond and atom type, e.g., in some implementations, derived from the respective denoising predictions for the nodes. Also or alternatively, the neural network is trained to generate a task prediction, where the task prediction can characterize one or more predicted properties of the molecule, e.g., the equilibrium energy of the molecule, the energy required to break the molecule, or the charge of the molecule. For example, the task may be to predict one or more properties of the molecule, such as a bond state prediction, e.g., a measure of the energy required to break one or more of the molecules, or a measure of how tightly the atoms are bonded, such as a bond angle or bond length measure, or one or more HOMO or LUMO energies of the molecule, or a property of the distribution of electrons in the molecule, such as size, charge, dipole moment, or static polarizability.

[0111] Local random distortions of the geometry of one or more molecules at a local energy minimum are almost certain to have larger energy configurations. Thus, in implementations where the denoising prediction for a node predicts the initial spatial positions of atoms before they are modified by using noise, the neural network is implicitly trained to determine the equilibrium or relaxed structure from the initial structure.

[0112] As mentioned above, depending on the application, the trained graph neural network 150 may be used by using features decoded from the node embeddings of the graph 104, or denoising prediction, or task prediction.

[0113] One exemplary application involves using the trained graph neural network 150 to obtain catalytic molecules or molecules that interact with a catalyst. In this application, the feature representation for the nodes comprises features that define the structure and properties of the catalytic molecules or molecules that interact with the catalyst. The output from the graph neural network may then comprise features decoded, for example, from the updated node embedding of the graph or from a denoising prediction that represents the resulting structure when the molecules interact, and / or a task prediction that characterizes the resulting state of the molecules, for example, the equilibrium energy of the molecules, or the change in energy resulting from the interaction, or the energy required to split the molecule. The prediction that characterizes the resulting structure or the resulting state of the molecules may be used to obtain catalytic molecules or molecules that interact with a catalyst, for example, by screening multiple candidate molecules. The screening may be, for example, to identify molecules that interact particularly strongly in a desired manner, or to screen out unsuitable molecules, or to identify catalytic molecules that interact with multiple other molecules, or molecules that interact with multiple different catalytic molecules (which may be either useful or unwanted).

[0114] The screening process may involve, for example, using output from the graph neural network to determine a score for each of a plurality of candidate catalytic molecules and / or catalyst-interacting molecules, and using the score to select one or more of the candidates. The method may further involve creating the catalytic molecules or catalyst-interacting molecules obtained by the method, and optionally testing the interactions in the real world.

[0115] In particular examples of the present application, the catalytic molecule comprises an enzyme or a receptor portion of an enzyme, and the molecule that interacts with the catalyst is a ligand of the enzyme, e.g., an agonist or antagonist of the receptor or enzyme. The ligand may be, for example, a drug, or a ligand of an industrial enzyme. One or both of the molecules may comprise a protein molecule.

[0116] A further related application involves using the trained graph neural network 150 to identify drug molecules that inhibit the replication of pathogens, i.e., to obtain drug molecules that interact with pathogen molecules. Pathogen molecules are molecules associated with pathogens, where the replication of the pathogen is inhibited when the drug molecule interacts with the pathogen molecule. Thus, in the method described above, pathogen molecules are used instead of catalytic molecules. Thus, the feature representation for the nodes may then comprise the features that define the molecule, and the output from the graph neural network is used to screen candidate drug molecules and / or pathogen molecules to obtain the drug molecule. The method may also involve creating drug molecules and optionally testing the drug molecules against pathogens in the real world.

[0117] Another exemplary application involves using the trained graph neural network 150 to determine a reaction mechanism of a chemical reaction to create a product involving two or more molecules interacting. One or more of the molecules may then be modified to modify the reaction mechanism, e.g., to increase the rate or product yield of the reaction. The reaction mechanism may then be used to create the product. The feature representation for the node may then comprise features that define the molecule. Output from the graph neural network may comprise, for example, features decoded from an updated node embedding of the graph or from a denoising prediction that represents the resulting structure when the molecules interact, and / or a task prediction that characterizes the resulting state of the molecule. For example, the output may predict one or more of the energy state, binding state, and fit of one or more transition states of the molecule along a reaction coordinate.

[0118] In the above applications, the feature representation for the node may then comprise features determined from one or more measurements made on the real-world molecule using, for example, electron microscopy, to characterize the structure or properties of the molecule. The features thus obtained may then be processed by the trained graph neural network to obtain a graph neural network output, for example a task prediction output for characterizing one or more properties of the molecule, such as the equilibrium energy, bonding states, bond angle or bond length measures, HOMO or LUMO energy, or size, charge, dipole moment, or static polarizability.

[0119] Some exemplary training datasets that can be used to train the graph neural network 150 to perform the above tasks are the OC20 dataset (Chanussot et al., "The Open Catalyst 2020 (OC20) Dataset and Community Challenges," ACS Catalysis, 6059-6072, 2020, arXiv:2010.09990), the QM9 dataset, Ramakrishnan et al., "Quantum chemistry structures and properties of 134 kilo molecules," Sci Data 1, 140022 (2014), the OGBG-PCQM4M dataset from the Open Graph Benchmark, Hu et al., "Open Graph Benchmark: Datasets for Machine Learning on Graphs," arXiv:2005.00687, and the OGBG-MOLPCBA dataset, also from the Open Graph Benchmark.

[0120] In some implementations, the graph 104 can represent a physical system, each node in the graph 104 can represent a respective object in the physical system, and the task prediction can characterize a predicted respective future state of one or more objects in the physical system, e.g., a respective position and / or velocity of each of the one or more objects in the physical system at a future time.

[0121] One exemplary application involves using the trained graph neural network 150 to predict a state of or control a physical system. The feature representation for the nodes may comprise features determined from the objects. Such features may comprise mass or moment of inertia of the objects, position, orientation, linear or angular velocity, or acceleration, and the edges may represent connected or interacting objects, e.g., objects connected by joints. Outputs from the graph neural network, e.g., features decoded from the updated node embeddings, denoising predictions 108, or task predictions, may define predictions of future states of objects in the physical system, e.g., for a single time step or for a rollout over multiple time steps. The outputs may be used to provide action control signals for controlling objects that depend on the future states. For example, the trained graph neural network 150 may be included in a Model Predictive Control (MPC) system to predict states or trajectories of a physical system, e.g., for use by a control algorithm in controlling the physical system to maximize rewards or minimize costs predicted from the future states.

[0122] In some implementations, the graph 104 can represent a point cloud (e.g., generated by a lidar or radar sensor), each node in the graph 104 can represent a respective point in the point cloud, and the task prediction can predict the class of an object represented by the point cloud.

[0123] In some implementations, the graph 104 can represent a portion of a text, where each node in the graph 104 can represent a respective word in that portion of text, and the task prediction can, for example, predict the sentiment expressed in that portion of text, e.g., positive, negative, or neutral.

[0124] In some implementations, the graph 104 can represent an image, each node in the graph 104 can represent a respective portion of the image (e.g., a pixel or region of the image), and the task prediction can, for example, characterize a class of object depicted in the image.

[0125] In some implementations, the graph 104 may represent an environment in the vicinity of a partial or complete autonomous vehicle, where each node in the graph may represent a respective agent (e.g., pedestrian, bicyclist, vehicle, etc.) in the environment or an element of the environment (e.g., traffic light, traffic sign, road lane, etc.), and the task prediction may, for example, predict a future respective trajectory of one or more of the agents represented by the nodes in the graph. For example, the prediction output may characterize a respective likelihood of the vehicle agent represented by the node in the graph making one or more possible driving decisions, e.g., going straight, changing lanes, turning left, or turning right. In this example, to predict a future trajectory of the agent represented by the node in the graph, the system may process updated node embeddings only for the node representing the agent, i.e., without processing updated node embeddings for other nodes in the graph. The edges of the graph may represent, for example, physical proximity or connectivity of the agents or elements, where connectivity may be defined as the existence of a route, such as a road or path, connecting the agents or elements. For example, the trained graph neural network may be used to control a mechanical agent in a real-world environment. The trained graph neural network may process feature representations for nodes comprising features representing other agents or elements of the environment, e.g., by agent or element type, and position, configuration, orientation, linear or angular velocity, or acceleration of the other agents or elements, to generate a graph neural network output for controlling the agent.

[0126] In some implementations, the graph 104 can represent a social network (e.g., on a social media platform), where each node in the graph can represent a respective person in the social network, and each edge in the graph can represent, for example, a relationship (e.g., a “follower” or “friend” relationship) between two corresponding people in the social network, and the task prediction can predict, for example, which people in the social network are likely to perform a particular action in the future (e.g., purchase a product or attend an event).

[0127] In some implementations, the graph 104 can represent a road network, each node in the graph can represent a route segment in the road network, each edge in the graph can represent that two corresponding route segments are connected in the road network, and the task prediction can predict, for example, the time required to follow a specified path through the road network, or the traffic volume on a specified path through the road network.

[0128] In some implementations, the graph 104 may be a computation graph that represents, for example, computational operations performed by a neural network model, where each node in the graph may represent a group of one or more related computations (e.g., operations performed by a group of one or more neural network layers), and each edge in the graph may represent that an output of one group of computations is provided as an input to another group of computations. In these implementations, the task prediction may, for example, predict a respective computational unit (i.e., from a set of available computational units) that should perform an operation corresponding to each node in the graph, for example, to minimize the time required to perform the operation specified by the graph. Each computational unit may, for example, be a respective thread, a central processing unit (CPU), or a graphics processing unit (GPU). Thus, the trained graph neural network may be used to perform a task of assigning computational operations to physical or logical computational units. The trained graph neural network may process the feature representations for the nodes, which comprise features representing groups of computations, to generate graph neural network outputs (e.g., features decoded from updated node embeddings, denoising predictions, or task predictions) to identify respective computational units that should perform operations corresponding to each node in the graph.

[0129] In some implementations, as described above, the graph 104 can represent a protein, each node in the graph can represent a respective amino acid in the protein's amino acid sequence, and each edge in the graph can represent that two corresponding amino acids in the protein are separated in the protein's structure by less than a threshold distance (e.g., 8 angstroms). In these implementations, the task prediction can, for example, predict protein stability or protein function.

[0130] In some implementations, the graph 104 can represent a knowledge base, where each node in the graph can represent a respective entity in the knowledge base, and each edge in the graph can represent a relationship between two corresponding entities in the knowledge base. In these implementations, the task prediction can, for example, predict missing features associated with one or more entities in the knowledge base.

[0131] Figure 2 illustrates an example of an operation that may be performed by the training system 100 in Figure 1. The system 100 may train the graph neural network 150 by using the neural network 150 to generate a denoising prediction 234. In some implementations, the system 100 may additionally train the neural network 150 by using the neural network 150 to generate a task prediction 232.

[0132] As described above with reference to FIG. 1, the system 100 can train the neural network by using the objective function defined in equation (7). In any training iteration, the objective function can be: (i) a denoising prediction (e.g., L de-noising ),(ii) task prediction (e.g., L task ), or (iii) a term that measures the error in both. In some implementations, the system 100 first calculates a denoised prediction (e.g., L de-noising ), and then the neural network 150 can be trained to optimize the objective function based on both the denoising predictions and the task predictions.

[0133] The system can include a noise engine that can be configured to generate a final feature representation for the nodes in the graph. In some implementations, the graph can represent a molecule 202, and each node in the graph can represent a respective atom in the molecule 202. In such a case, the noise engine can process an initial feature representation for each node in the graph, where the initial feature representation for the node represents an initial spatial position of the corresponding atom in the molecule 202. The noise engine can generate a final feature representation for the nodes, where the final feature representation for some or all of the nodes is a modified feature representation that is generated using the respective noise. For example, the noise engine can generate the final feature representation for the node by adding noise to the initial spatial positions of the atoms in the molecule 202 represented by the node.

[0134] The neural network 150 may include (i) an encoder 210 , (ii) an updater 220 , and (iii) a decoder 230 .

[0135] The noise engine can provide the final feature representations for the nodes in the graph to the encoder 210. The encoder 210 can be configured to generate data defining the graph. For example, the encoder 210 includes a node embedding sub-network that can generate node embeddings for each node, for example, based on the final feature representations for each node.

[0136] After generating the data defining the graph, the encoder 210 can provide the data defining the graph to the updater 220. The updater can include one or more graph neural network layers, e.g., N graph neural network layers. Each graph neural network layer can be configured to update a current node embedding of one or more nodes in the graph based on a current node embedding of the node and a respective current node embedding of each of the node's one or more neighbors in the graph. For example, as shown in FIG. 2, the graph neural network layer 215 can update the node embedding of a node indicated by a filled circle based on the current node embedding of adjacent nodes.

[0137] The decoder 230 may be configured to process the updated node embeddings of the nodes for each of one or more of the nodes having modified feature representations to generate a respective denoised prediction 234 for the node. The denoised prediction 234 may characterise a denoised feature representation for the node that does not include the noise used to generate the modified feature representation for the node. As described above with reference to FIG. 1, the denoised prediction 234 may predict the noise used to generate the modified feature representation for the node. In some implementations, the denoised prediction 234 may predict an initial feature representation for the node before it is perturbed with noise, e.g., an initial spatial location of an atom in the molecule 202 represented by the node.

[0138] In some implementations, the denoising prediction 234 for a node can characterize a target feature representation for the node. For example, as shown in FIG. 2, the target feature representation can specify final spatial positions of atoms in the molecule 202 represented by the node after atomic relaxation. In this manner, the system 100 can map initial spatial positions of atoms in the molecule 202 (e.g., specified by an initial feature representation for the node) to final spatial positions of atoms in the molecule 202.

[0139] In some implementations, the decoder 230 can generate a task prediction 232. In general, the task prediction 232 can be any suitable prediction that characterizes one or more of the elements represented by the nodes in the graph. As shown in FIG. 2, the task prediction 232 can be, for example, an equilibrium energy of the molecule 202 after atomic relaxation.

[0140] Training the neural network 150 to generate denoised predictions 234 can encourage the neural network 150 to implicitly learn the “real-world” distributions of the graph, i.e., where the node feature representations are not modified, and the neural network 150 can leverage this implicit knowledge to achieve higher accuracy in task predictions 232.

[0141] An exemplary process for training the neural network 150 using the training system 100 will now be described in more detail.

[0142] 3 is a flow diagram of an exemplary process 300 for training a graph neural network using a training system. For convenience, process 300 is described as being performed by one or more computer systems located at one or more locations. For example, a training system suitably programmed in accordance with this specification, such as training system 100 of FIG. 1, may perform process 300.

[0143] The system generates 302 data defining a graph including (i) a set of nodes, (ii) a node embedding for each node, and (iii) a set of edges each connecting a respective pair of nodes. As described above with reference to FIG. 1, the system may generate the data defining the graph based on a feature representation for each node. For example, the system may obtain a respective initial feature representation for each node and may generate a respective final feature representation for each node. For each of one or more of the nodes, the respective final feature representation may be a modified feature representation generated from the respective feature representation for the node using the respective noise, e.g., by adding the respective noise to the respective feature representation for the node, as defined by Equation (1).

[0144] The system can generate data defining the graph using the final feature representations of each of the nodes. For example, the system can determine, for each pair of nodes including a first node and a second node, a respective distance between the final feature representation for the first node and the final feature representation for the second node. The system can then determine that each pair of nodes corresponding to a distance less than a predefined threshold is connected by an edge in the graph. As shown in FIG. 2, the graph can represent a molecule, and each node in the graph can represent a respective atom in the molecule. In this case, the initial feature representation for each node can represent, for example, an initial spatial location of the corresponding atom in the molecule. The system can generate data defining the graph by generating a node embedding for each node based on the type of atom represented by the node. In general, the graph can represent any suitable physical system.

[0145] In some implementations, the graph may further include a respective edge embedding for each edge. In such a case, the system may generate the graph by generating an edge embedding for each edge in the graph based at least in part on the difference between the respective final feature representations of the nodes connected by the edge, e.g., as defined by equations (2) and (3) above.

[0146] The system processes the data defining the graph using one or more graph neural network layers of the neural network to generate updated node embeddings for each node (304). In some implementations, the neural network includes at least 10 graph neural network layers. Each graph neural network layer can be configured to update the current graph.

[0147] For example, as described above with reference to FIG. 1, each neural network layer can receive a current graph and can update the current graph according to the current neural network parameter values ​​of the graph neural network layer. This can include, for example, updating a current node embedding of each of one or more nodes in the graph based on (i) the current node embedding of the node and (ii) the respective current node embedding of each of the node's one or more neighbors in the graph, e.g., as defined by Equation (4) and Equation (5) above. In some implementations, the current graph can further include an edge embedding for each edge. In such a case, each neural network layer can update the node embedding of the node based at least in part on the respective edge embedding of each of the one or more edges coupled to the node.

[0148] For each of one or more of the nodes having a modified representation, the system processes the updated node embeddings of the nodes to generate a respective denoised prediction for the node (306). For example, the system may process the updated node embeddings of the nodes using one or more neural network layers to generate a respective denoised prediction for the node. The denoised predictions may characterise a denoised representation for the node that does not include the noise used to generate the modified representation of the node.

[0149] In one example, the denoising prediction can predict the noise used to generate the modified feature representation of the node. In another example, the denoising prediction can predict an initial feature representation of the node, e.g., the initial spatial positions of atoms in a molecule before being modified by using noise. In yet another example, the denoising prediction can characterize a target feature representation of the node, e.g., the new positions of atoms in a molecule after atomic relaxation. In some implementations, the target feature representation for a node can be an incremental feature representation for the node that, when added to the modified feature representation for the node, results in the target feature representation of the node.

[0150] The system determines updates to current values ​​of neural network parameters of the neural network to optimize an objective function that measures the error in the denoising prediction for the nodes (308). For example, the system may back-propagate the gradient of the objective function through the neural network parameters of the graph neural network layer. The objective function may measure, for each of multiple graph neural network layers of the neural network, a respective error in the denoising prediction for the nodes based on the updated node embeddings generated by the graph neural network layer.

[0151] In some implementations, the system can process the updated node embedding of the node to generate a task prediction, where the objective function also measures the error in the task prediction, e.g., as defined in Equation (7). In general, the error in the task prediction can be determined using a set of training data appropriate for the task, i.e., the system can be trained using supervised learning. In such a case, the system can process both (i) the updated node embedding of the node and (ii) the original node embedding of the node before being updated using the graph neural network layer, to generate a task prediction. As a specific example, the graph can represent a molecule, e.g., as shown in FIG. 2, and the task prediction can be a prediction of the equilibrium energy of the molecule.

[0152] Exemplary experimental results achieved using the system for training neural networks are now described in more detail.

[0153] FIG. 4 illustrates exemplary experimental results 400 achieved using the system 100 for training neural networks, described above with reference to FIGS.

[0154] The system 100 can train a neural network by using the neural network to generate a denoised prediction. As explained above, generating a denoised prediction requires each node embedding to encode unique information to denoise the feature representation of the node, which can mitigate the effects of "over-smoothing", for example, where node embeddings become nearly identical after being processed through several graph neural network layers. In FIG. 4, "MAD" is a measure of node embedding diversity that can quantify "over-smoothing", where a higher number indicates a higher level of diversity in node embeddings. As shown in FIG. 4, the system described herein can maintain a higher level of node embedding diversity throughout the neural network compared to other available techniques (e.g., "DropEdge" and "DropNode"). This is especially evident at neural network layer 15, where the measure of node embedding diversity of the system described herein is much larger. Thus, the system described herein may be superior to other available systems in mitigating the effects of "over-smoothing".

[0155] Figure 5 shows exemplary experimental results 500 achieved using the system for training neural networks 100 described above with reference to Figures 1 and 2. In Figure 5, "previous SOTA" refers to the state of the art performance achieved using other available systems, and "ev MAE" refers to the prediction error in task prediction.

[0156] The graph on the left shows that even after three message passing steps (e.g., involving three graph neural network layers in the neural network), the system described herein exceeds the state of the art performance achieved using other available systems. The graph on the right shows that even with fewer neural network parameters (e.g., having shared weights between graph neural network layers), the performance of the state of the art can be exceeded by the system described herein.

[0157] FIG. 6 shows exemplary experimental results comparing the performance of various neural networks in the task of predicting the HOMO energy of a molecule. In particular, the horizontal axis of graph 600 represents the number of gradient steps used to train each neural network, and the vertical axis of graph 600 represents the prediction accuracy of each neural network. The best performing neural network (labeled "Pre-Trained GNS-TAT" in FIG. 6) is a graph neural network that is pre-trained to optimize an objective function based on denoising predictions (as described herein) before being trained to perform the HOMO energy prediction task. Thus, these experimental results show the advantages that can be achieved by pre-training a graph neural network to optimize an objective function based on denoising predictions.

[0158] This specification uses the term "configured" with respect to systems and computer program components. A system of one or more computers configured to perform a particular operation or action means that the system has software, firmware, hardware, or a combination thereof installed thereon that, when operated, causes the system to perform the operation or action. A computer program or programs configured to perform a particular operation or action means that one or more programs contain instructions that, when executed by a data processing device, cause the device to perform the operation or action.

[0159] The subject matter and functional operations described herein may be implemented in digital electronic circuitry, tangibly embodied computer software or firmware, computer hardware, or a combination of one or more of them, including the structures disclosed herein and their structural equivalents. The subject matter described herein may 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 a data processing device or for controlling the operation of a data processing device. The computer storage medium may 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 additionally, the program instructions may be encoded in an artificially generated propagated signal, e.g., a machine-generated electrical, optical, or electromagnetic signal, which is generated to encode information for transmission to a suitable receiver device for execution by a data processing device.

[0160] 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. An apparatus may also be or further include special purpose logic circuitry, e.g., an FPGA (field programmable gate array) or an ASIC (application specific integrated circuit). An apparatus may optionally include, in addition to hardware, code that creates an execution environment for a computer program, e.g., code that constitutes a processor firmware, a protocol stack, a database management system, an operating system, or one or more combinations thereof.

[0161] A computer program, which may also be referred to or described as a program, software, software application, app, module, software module, script, or code, may be written in any form of programming language, including compiled or interpreted languages, or declarative or procedural languages, and may 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 correspond to a file in a file system, but need not. A program may be stored in a file that holds one or more scripts stored among other programs or data, for example, in a markup language document, in a single file dedicated to the program, or in multiple cooperating files, for example, in a file that stores one or more modules, subprograms, or portions of code. A computer program may be deployed to run on one computer, or on multiple computers located at one site or distributed across multiple sites and interconnected by a data communication network.

[0162] The term "engine" is used broadly herein to refer to a software-based system, subsystem, or process that is programmed to perform one or more specific functions. Generally, an engine is implemented as one or more software modules or components installed on one or more computers at one or more locations. In some cases, one or more computers are dedicated to a particular engine, and in other cases, multiple engines can be installed and run on the same computer or computers.

[0163] The processes and logic flows described herein may 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 may also be performed by special purpose logic circuitry, e.g., FPGAs or ASICs, or by a combination of special purpose logic circuitry and one or more programmed computers.

[0164] A computer suitable for executing a computer program can be based on a general-purpose or special-purpose microprocessor, or both, or any other type of central processing unit. In general, the central processing unit receives 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 memory can be augmented by or incorporated in special-purpose logic circuitry. In general, the computer also includes one or more mass storage devices, e.g., magnetic disks, magneto-optical disks, or optical disks, for storing data, or is operatively coupled to receive data from or transfer data to such mass storage devices, or both. However, a computer need not have such devices. Moreover, a computer may be incorporated within another device, such as a mobile phone, 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, such as a Universal Serial Bus (USB) flash drive, to name just a few.

[0165] 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.

[0166] To interact with a user, embodiments of the subject matter described herein may 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 pointing device, e.g., a mouse or trackball, by which the user can provide input to the computer. Other types of devices may also be used to interact with the user, e.g., feedback provided to the user may be any form of sensory feedback, e.g., visual feedback, auditory feedback, or tactile feedback, and input from the user may be received in any form, including acoustic input, speech input, or tactile input. In addition, the computer may interact with the user by sending documents to and receiving documents from devices used by the user, e.g., by sending web pages to a web browser on the user's device in response to a request received from the web browser. The computer may also interact with the user by sending text messages or other forms of messages to a personal device, e.g., a smartphone running a messaging application, and receiving responsive messages from the user in return.

[0167] A data processing device for implementing machine learning models may also include, for example, dedicated hardware accelerator units for handling the common and computationally intensive parts of machine learning training and production, i.e., estimation, workloads.

[0168] The machine learning model may be implemented and deployed using a machine learning framework, for example, the TensorFlow framework.

[0169] An embodiment of the subject matter described herein may be implemented in a computing system, including back-end components, e.g., as a data server, or including middleware components, e.g., an application server, or including front-end components, e.g., a client computer having a graphical user interface, web browser, or app through which a user may interact with an implementation of the subject matter described herein, or any combination of one or more such back-end, middleware, or front-end components. The components of the system may 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.

[0170] A computing system may include clients and servers. Clients and servers 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., HTML pages, to a user device, e.g., to display data to and receive user input from a user interacting with the device acting as a client. Data generated at the user device, e.g., results of user interaction, may be received from the device at the server.

[0171] Although this specification contains many specific implementation details, these should not be interpreted 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 a particular embodiment of a particular invention. Some features described in the context of a separate embodiment may also be implemented in a single embodiment in combination. Conversely, various features described in the context of a single embodiment may also be implemented in multiple embodiments separately or in any suitable subcombination. Moreover, although features may be described above as working in some combinations and may even be initially claimed as such, one or more features from a claimed combination may, in some cases, be deleted from the combination, and the claimed combination may be directed to a subcombination or a variation of the subcombination.

[0172] Similarly, although operations are illustrated in the figures and claimed in a particular order, this should not be understood as requiring that such operations be performed in the particular order or sequential order illustrated, or that all of the illustrated operations be performed, to achieve desirable results. In some situations, 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 program components and systems described may generally be integrated together in a single software product or packaged in multiple software products.

[0173] Specific embodiments of the present 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. By way of example, the processes depicted in the accompanying figures do not necessarily require the particular order depicted, or sequential order, to achieve desirable results. In some cases, multitasking and parallel processing may be advantageous. [Explanation of symbols]

[0174] 100 Training System 102 Initial feature representation 104 Graphs, Initial Graphs 106 Final Graph 108 Noise Removal Prediction 109 Task Prediction 110 Encoder 112 Final feature representation 120 Updater 130 Decoder 140 Training Engine 150 Graph Neural Networks 160 Noise Engine 202 molecules 210 Encoder 215 Graph Neural Network Layers 220 Updater 230 Decoder 232 Task Prediction 234 Noise Removal Prediction

Claims

1. 1. One or more computer-implemented methods for training a neural network including one or more graph neural network layers, comprising: generating data defining a graph comprising (i) a set of nodes, (ii) a node embedding for each node, and (iii) a set of edges each connecting a respective pair of nodes; obtaining a respective initial feature representation for each node; generating a respective final feature representation for each node, where for each of one or more of the nodes, the respective final feature representation is a modified feature representation generated from the respective initial feature representation for the node using a respective noise; and generating said data defining said graph using said final representations of each of said nodes. and processing the data defining the graph using one or more of the graph neural network layers of the neural network to generate respective updated node embeddings for each node; For each of one or more of the nodes having a modified representation, processing the updated node embeddings of the node to generate a respective denoised prediction for the node that characterises a denoised representation for the node that does not include the noise used to generate the modified representation of the node; determining updates to current values ​​of neural network parameters of the neural network to optimize an objective function that measures the error in the denoising prediction for the node; A method comprising:

2. 2. The method of claim 1 , wherein for each of one or more of the nodes having a modified representation, the respective denoised prediction for the node predicts the noise used to generate the modified representation of the node.

3. 2. The method of claim 1, wherein for each of one or more of the nodes having a modified representation, the respective denoised prediction for the node predicts the respective initial representation of the node.

4. 2. The method of claim 1, wherein for each of one or more of the nodes having a modified representation, the respective denoised prediction for the node characterizes a target representation for the node.

5. 5. The method of claim 4, further comprising: predicting, for each one or more of the nodes having a modified representation, an incremental representation for the node that, when added to the modified representation for the node, results in the target representation for the node.

6. processing the updated node embeddings of the nodes to generate a task prediction; The method of claim 1 , wherein the objective function also measures an error in the task prediction.

7. 7. The method of claim 6, wherein both (i) the updated node embedding of the node and (ii) the original node embedding of the node before being updated using the graph neural network layer are processed to generate the task prediction.

8. the graph represents a molecule; 8. The method of claim 6 or 7, wherein the task prediction is a prediction of the equilibrium energy of the molecule.

9. 2. The method of claim 1 , wherein the objective function measures, for each of a plurality of graph neural network layers of the neural network, a respective error in a denoising prediction for the node that is based on an updated node embedding generated by the graph neural network layer.

10. for each of one or more of the nodes having a modified feature representation, processing the updated node embedding of the node to generate the respective denoised prediction for the node, 2. The method of claim 1 , comprising processing the updated node embeddings of the nodes using one or more neural network layers to generate the respective denoised predictions for the nodes.

11. determining the updates to the current values ​​of the neural network parameters of the neural network to optimize the objective function; backpropagating the gradient of the objective function through the neural network parameters of the graph neural network layer.

2. The method of claim 1, comprising:

12. 2. The method of claim 1, wherein for each of one or more of the nodes, the respective final feature representation for the node is generated by adding the respective noise to the respective initial feature representation for the node.

13. generating the data defining the graph using the final representations of each of the nodes, for each pair of nodes comprising a first node and a second node, determining a respective distance between the final feature representation for the first node and the final feature representation for the second node; determining that each pair of nodes corresponding to a distance less than a predefined threshold is connected by an edge in the graph; 2. The method of claim 1, comprising:

14. The method of claim 1 , wherein the graph further comprises a respective edge embedding for each edge.

15. generating the data defining the graph, generating an edge embedding for each edge in the graph based at least in part on the differences between the respective final feature representations of the nodes connected by the edge.

15. The method of claim 14, comprising:

16. the graph represents a molecule; each node in the graph represents a respective atom in the molecule; generating the data defining the graph, generating a node embedding for each node based on the type of atom represented by the node.

2. The method of claim 1, comprising:

17. The method of claim 1 , wherein the neural network comprises at least 10 graph neural network layers.

18. Each graph neural network layer of the graph neural network layers Receiving a current graph; updating the current graph according to current neural network parameter values ​​of the graph neural network layer; updating the current node embedding of each of one or more nodes in the graph based on (i) the current node embedding of the node, and (ii) a respective current node embedding of each of one or more neighbors of the node in the graph. Including, updating and The method of claim 1 , configured to:

19. the current graph comprises an edge embedding for each edge; updating the current node embeddings of each of one or more nodes in the graph; updating the node embedding of the node based at least in part on a respective edge embedding of each of one or more edges coupled to the node.

20. The method of claim 18, further comprising:

20. the graph represents a molecule; each node in the graph represents a respective atom in the molecule; the initial feature representation for each node represents an initial spatial position of a corresponding atom in the molecule; 6. The method of claim 5, wherein for each one or more of the nodes having a modified representation, the target representation for the node represents a final spatial position of the corresponding atom after atomic relaxation.

21. One or more non-transitory computer storage media having instructions stored thereon that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective method of claim 1.

22. 1. A system comprising: the one or more computers; one or more storage devices communicatively coupled to the one or more computers; 20. A system comprising: one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations of the respective methods recited in claim 1.

Images