Generating prediction outputs using dynamic graphs

1 . A method performed by one or more data processing apparatus for generating prediction outputs characterizing a set of entities, the method comprising:

obtaining data defining a graph, comprising: (i) a set of nodes, wherein each node in the set of nodes represents a respective entity from the set of entities, (ii) a current set of edges, wherein each edge in the current set of edges connects a pair of nodes, and (iii) a respective current embedding of each node;

at each of a plurality of time steps:

updating the respective current embedding of each node of the graph, comprising processing data defining the graph using a graph neural network to update the respective current embedding of each node of the graph using feature representations of neighboring nodes identified based on the current set of edges at the time step;

updating the current set of edges of the graph based at least in part on the updated embeddings of the nodes of the graph, comprising:

processing the updated embeddings of the nodes in the graph, using a set of neural network parameters that have been trained by a machine learning training technique, to generate (i) a respective relevance score between each of a plurality of pairs of nodes in the graph, and (ii) a respective masking output for each node that characterizes whether to update edges connected to the node in the current set of edges; and

updating the current set of edges of the graph based on the relevance scores and the masking outputs, wherein updating the current set of edges comprises:

determining nodes with edges designated for replacement based on the masking outputs;

updating edges belonging to each node designated for replacement by adding one or more edges to the graph, or removing one or more edges from the graph, or both based on the set of relevance scores; and

after updating the current set of edges of the graph, providing the updated graph for processing at a next time step; and

at one or more of the plurality of time steps:

generating a prediction output characterizing the set of entities based on the current embeddings of the nodes of the graph.

2 . The method of claim 1 , wherein processing the updated embeddings of the nodes in the graph, using the set of neural network parameters that have been trained by a machine learning training technique, to generate the respective relevance score between each pair of nodes in the plurality of pairs of nodes in the graph comprises:

for first node and a second node in the pair of nodes, determining the relevance score between the first node and the second node based on the updated embeddings of the first node and the second node.

3 . The method of claim 2 , wherein the set of neural network parameters that have been trained by a machine learning training technique comprises a set of query parameters and a set of key parameters, and wherein determining the relevance score between the first node and the second node in the pair of nodes based on the updated embeddings of the first node and the second node comprises:

processing the updated embedding of the first node in accordance with the set of query parameters to generate a query embedding of the first node;

processing the updated embedding of the second node in accordance with the set of key parameters to generate a key embedding of the second node; and

determining the relevance score between the first node and the second node based on a similarity measure between: (i) the query embedding of the first node, and (ii) the key embedding of the second node.

4 . The method of claim 3 , wherein the similarity measure between: (i) the query embedding of the first node, and (ii) the key embedding of the second node, comprises an inner product of: (i) the query embedding of the first node, and (ii) the key embedding of the second node.

5 . The method of claim 1 , wherein the current set of edges includes a predefined set of static edges and a current set of dynamic edges, and wherein updating the current set of edges of the graph based on the relevance scores and the masking output comprises comprises updating only the current set of dynamic edges for nodes with edges designated for replacement based on the relevance scores between pairs of nodes in the graph.

6 . The method of claim 5 , wherein updating the current set of dynamic edges for nodes with edges designated for replacement based on the relevance scores between pairs of nodes in the graph comprises, for each given first node of the graph:

determining that a relevance score between the given first node and a given second node is higher than a relevance score between the given first node and any other node other than the given second node; and

adding an edge connecting the given first node to the given second node to the current set of dynamic edges.

7 . The method of claim 6 , further comprising removing one or more edges from the current set of dynamic edges prior to adding any edges to the current set of dynamic edges.

8 . The method of claim 5 , wherein at each of the plurality of time steps, the updated set of dynamic edges includes a predefined number of dynamic edges.

9 . The method of claim 1 , wherein updating the respective current embedding of each node of the graph comprises:

for each node in the graph, processing an input comprising the current embedding of the node using an encoder neural network to generate a feature representation of the node; and

providing the respective feature representation of each node as an input to the graph neural network.

10 . The method of claim 9 , further comprising, for each node in the graph:

receiving, at each time step, respective input features corresponding to the node;

wherein the encoder neural network processes an input comprising both: (i) the current embedding corresponding to the node, and (ii) the input features corresponding to the node, to generate the feature representation corresponding to the node.

11 . The method of claim 9 , wherein processing data defining the graph using the graph neural network to update the respective current embedding of each node of the graph using feature representations of neighboring nodes identified based on the current set of edges at the time step comprises, for one or more nodes of the graph:

processing respective feature representations of: (i) the node, and (ii) one or more other nodes that are connected to the node, in accordance with a plurality of graph neural network parameters to update the current embedding of the node.

12 . The method of claim 1 , wherein generating a prediction output characterizing the set of entities based on the current embeddings of the nodes of the graph comprises:

generating a pooled embedding by pooling current embeddings of the nodes of the graph; and

processing the pooled embedding to generate the prediction output.

13 . The method of claim 1 , wherein the current embeddings of the nodes of the graph and the current set of edges of the graph are updated in accordance with values of a set of parameters, and further comprising:

updating the values of the set of parameters using gradients of an objective function that measures respective errors between:

(i) the current set of edges of the graph, and (ii) a target set of edges of the graph, at each of the plurality of time steps; and

(i) the masking outputs for each node in the graph, and (ii) target masking outputs for each node in the graph, at each of the plurality of time steps.

14 . The method of claim 13 , wherein the objective function further measures an error between: (i) the prediction outputs characterizing the set of entities, and (ii) target outputs characterizing the set of entities.

15 . The method of claim 13 , wherein the target set of edges and the target masking outputs at each of the plurality of time steps are generated using a disjoint-set union data structure.

16 . The method of claim 1 , wherein at a first time step of the plurality of time steps:

the current set of edges includes a respective edge connecting each node in the graph to itself; and

the respective current embedding of each node is a default embedding.

17 . The method of claim 1 , wherein each entity in the set of entities is a respective atom in a molecule, and wherein the prediction output characterizes the energy required to break up the molecule.

18 . The method of claim 1 , wherein each entity in the set of entities is an object in a physical system, and wherein the prediction output characterizes a respective predicted future position of each of the objects in the physical system.

19 . A system comprising:

one or more computers; and

one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for generating prediction outputs characterizing a set of entities, the operations comprising:

obtaining data defining a graph, comprising: (i) a set of nodes, wherein each node in the set of nodes represents a respective entity from the set of entities, (ii) a current set of edges, wherein each edge in the current set of edges connects a pair of nodes, and (iii) a respective current embedding of each node;

at each of a plurality of time steps:

updating the respective current embedding of each node of the graph, comprising processing data defining the graph using a graph neural network to update the respective current embedding of each node of the graph using feature representations of neighboring nodes identified based on the current set of edges at the time step;

updating the current set of edges of the graph based at least in part on the updated embeddings of the nodes of the graph, comprising:

processing the updated embeddings of the nodes in the graph, using a set of neural network parameters that have been trained by a machine learning training technique, to generate (i) a respective relevance score between each of a plurality of pairs of nodes in the graph, and (ii) a respective masking output for each node that characterizes whether to update edges connected to the node in the current set of edges; and

updating the current set of edges of the graph based on the relevance scores and the masking outputs, wherein updating the current set of edges comprises:

 determining nodes with edges designated for replacement based on the masking outputs;

 updating edges belonging to each node designated for replacement by adding one or more edges to the graph, or removing one or more edges from the graph, or both based on the set of relevance scores; and

after updating the current set of edges of the graph, providing the updated graph for processing at a next time step; and

at one or more of the plurality of time steps:

generating a prediction output characterizing the set of entities based on the current embeddings of the nodes of the graph.

20 . One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for generating prediction outputs characterizing a set of entities, the operations comprising:

obtaining data defining a graph, comprising: (i) a set of nodes, wherein each node in the set of nodes represents a respective entity from the set of entities, (ii) a current set of edges, wherein each edge in the current set of edges connects a pair of nodes, and (iii) a respective current embedding of each node;

at each of a plurality of time steps:

updating the respective current embedding of each node of the graph, comprising processing data defining the graph using a graph neural network to update the respective current embedding of each node of the graph using feature representations of neighboring nodes identified based on the current set of edges at the time step;

updating the current set of edges of the graph based at least in part on the updated embeddings of the nodes of the graph, comprising:

processing the updated embeddings of the nodes in the graph, using a set of neural network parameters that have been trained by a machine learning training technique, to generate (i) a respective relevance score between each of a plurality of pairs of nodes in the graph, and (ii) a respective masking output for each node that characterizes whether to update edges connected to the node in the current set of edges; and

updating the current set of edges of the graph based on the relevance scores and the masking outputs, wherein updating the current set of edges comprises:

determining nodes with edges designated for replacement based on the masking outputs;

updating edges belonging to each node designated for replacement by adding one or more edges to the graph, or removing one or more edges from the graph, or both based on the set of relevance scores; and

after updating the current set of edges of the graph, providing the updated graph for processing at a next time step; and

at one or more of the plurality of time steps:

generating a prediction output characterizing the set of entities based on the current embeddings of the nodes of the graph.