Abstract: Methods and systems for conditioning graph neural networks on affinity features. One of the methods includes obtaining graph data representing an input graph that comprises a set of nodes and a set of edges that each connect a respective pair of nodes, the graph data comprising respective node features for each of the nodes, edge features for each of the edges, and a respective weight for each of the edges; generating one or more affinity features, each affinity feature representing a property of one or more random walks through the graph guided by the respective weights for the edges; and processing the graph data using a graph neural network that is conditioned on the one or more affinity features to generate a task prediction for a machine learning task for the input graph.

Abstract: Methods, systems, and apparatus, including computer programs encoded on computer storage media, for performing persistent message passing using graph neural networks.

Abstract: Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for generating prediction outputs characterizing a set of entities. In one aspect, a method comprises: obtaining data defining a graph, comprising: (i) a set of nodes, wherein each node represents a respective entity from the set of entities, (ii) a current set of edges, wherein each edge 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, comprising processing data defining the graph using a graph neural network; and updating the current set of edges based at least in part on the updated embeddings of the nodes; 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.