Abstract: Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a neural network that includes one or more graph neural network layers. In one aspect, a method comprises: generating data defining a graph, comprising: generating a respective final feature representation for each node, wherein, for each of one or more of the nodes, the respective final feature representation is a modified feature representation that is generated from a respective feature representation for the node using 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 of each node; and processing, for each of one or more of the nodes having modified feature representations, the updated node embedding of the node to generate a respective de-noising prediction for the node.

Abstract: This specification describes a simulation system that performs simulations of physical environments using a graph neural network. At each of one or more time steps in a sequence of time steps, the system can process a representation of a current state of the physical environment at the current time step using the graph neural network to generate a prediction of a next state of the physical environment at the next time step. Some implementations of the system are adapted for hardware GLOBAL acceleration. As well as performing simulations, the system can be used to predict physical quantities based on measured real-world data. Implementations of the system are differentiable and can also be used for design optimization, and for optimal control tasks.