Any2Graph: Supervised Graph Prediction Poster

Paper context

Paper title: Any2Graph: Deep End-To-End Supervised Graph Prediction With An Optimal Transport Loss Abstract: This poster introduces Any2Graph, a framework for supervised graph prediction using optimal transport loss to handle permutation invariance. It details the architecture, graph representation via masking vectors, and presents qualitative and quantitative results on datasets like Coloring and USCities. Paper body (method & results): Any2Graph: Deep End-To-End Supervised Graph Prediction With An Optimal Transport Loss Paul Krzakala Télécom Paris, IP Paris, LTCI & CMAP Junjie Yang Télécom Paris, IP Paris, LTCI Rémi Flamary Ecole polytechnique, IP Paris, CMAP Florence d’Alché-Buc Télécom Paris, IP Paris, LTCI Charlotte Laclau Télécom Paris, IP Paris, LTCI Matthieu Labeau Télécom Paris, IP Paris, LTCI Abstract We propose Any2Graph, a generic framework for end-to-end Supervised Graph Prediction (SGP) i.e. a deep learning model that predicts an entire graph for any kind of input. The framework is built on a novel Optimal Transport loss, the Partially-Masked Fused Gromov-Wasserstein, that exhibits all necessary proper- ties (permutation invariance, differentiability) and is designed to handle any-sized graphs. Numerical experiments showcase the versatility of the approach that out- performs existing competitors on a novel challenging synthetic dataset and a variety of real-world tasks such as map construction from satellite image (Sat2Graph) or molecule prediction from fingerprint (Fingerprint2Graph). 1 1 Introduction This work focuses on the problem of Supervised Graph Prediction (SGP), at the crossroads of Graph-based Learning and Structured Prediction. In contrast to node and graph classification or link prediction widely covered in recent literature by graph neural networks, the target variable in SGP is a graph and no particular assumption is made about the input variable. Emblematic applications of SGP include knowledge graph extraction [28] or dependency parsing [17] in natural language processing, conditional graph scene generation in computer vision [53], [12], or molecule identification in chemistry [10, 55], [49], to name but a few. Moreover, close to SGP is the unsupervised task of graph generation notably motivated by de novo drug design [8, 15, 41]. SGP raises some specific issues related to the complexity of the output space and the absence of widely accepted loss functions. First, the non-Euclidean nature of the output to be predicted makes both inference and learning challenging while the size of the output space is extremely large. Second, the arbitrary size of the output variable to predict requires a model with a flexible expressive power in the output space. Third, graphs are characterized by the absence of natural or ground-truth ordering of their nodes, making comparison and prediction difficult. This particular issue calls for a node permutation invariant distance to predict graphs. Scrutinizing the literature through the lens of these issues, we note that existing methodologies circumvent the difficulty of handling output graphs in various ways. A first body of work avoids end-to-end learning by relying on some relaxations. For instance, energy-based models (see for instance [38]) convert the problem into the learning of an 1All code is available at https://github.com/KrzakalaPaul/Any2Graph. 38th Conference on Neural Information Processing Systems (NeurIPS 2024). arXiv:2402.12269v4 [cs.LG] 15 Oct 2024 energy function of input and output while surrogate regression methods [10] implicitly embed output graphs into a given Hilbert space where the learning task boils down to vector-valued regression. Note that these two families of approaches generally involve a rather expensive decoding step at inference time. In what follows, we focus on methods that directly output graphs or close relaxations, enabling end-to-end learning. One strategy to overcome the need for a permutation invariant loss is to exploit the nature of the input data to determine a node ordering, with the consequence that application to new types of data requires similar engineering. For instance, in de novo drug generation SMILES representations [8] are generally used to determine atom ordering. In semantic parsing, the target graph is a tree that can be serialized [3] while in text-to-knowledge-graph, the task is re-framed into a sequence-to-sequence problem, often addressed with large autoregressive models. Finally, for road map extraction from satellite images, one can leverage the spatial positions of the nodes to define a unique ordering [5]. Another line of research proposes to address this problem more directly by seeking to solve a graph-matching problem, i.e., finding the one-to-one correspondence between nodes of the graphs. Among approaches in this category, we note methods dedicated to molecule generation [25] where the invariant loss is based on a characterization of graphs, ad-hoc to the molecule application. While being fully differentiable their loss does not generalize to other applications. In the similar topic of graph generation, Simonovsky and Komodakis [37] propose a more generic definition of the similarity between graphs by considering both feature and structural matching. However, they solve the problem using a two-step approach by using first a smooth matching approximation followed by a rounding step using the Hungarian algorithm to obtain a proper one-to-one matching, which comes with a high computational cost and introduces a non-differentiable step. For graph scene generation, Relationformer [36] is based on a bipartite object matching approach solved using a Hungarian matcher [11]. The main shortcoming of this approach is that it fails to consider structural information in the matching process. The same problem is encountered by Melnyk et al. [28]. We discuss Relationformer in more detail later in the article. Finally, another way to approach end-to-end learning is to leverage the notion of graph barycenter to define the predicted graph. Relying on the Implicit Loss Embedding (ILE) property of surrogate regression, Brogat-Motte et al. [9] have exemplified this idea by exploiting an Optimal Transport loss, the Fused Gromov-Wasserstein (FGW) distance [45] for which barycenters can be computed efficiently [32, 44]. They proposed two variants, a non-parametric kernel-based one and a neural network-based one, referred to as FGW-Bary and FGW-BaryNN, respectively. However, to calculate the barycenter, the size must be known upstream, leaving the challenge of arbitrary size unresolved. In addition, prediction accuracy is highly dependent on the expressiveness of the barycenter, i.e. the nature and number of graph templates, resulting in high training and inference costs. In contrast to existing works, our goal is to address the problem of supervised graph prediction in an end-to-end fashion, for different types of input modalities and for output graphs whose size and node ordering can be arbitrary. Main contributions This paper presents Any2Graph, a versatile framework for end-to-end SGP. Any2Graph leverages a novel, fully differentiable, OT-based loss that satisfies all the previously mentioned properties, i.e., size agnostic and invariant to node permutation. In addition, the encoder part of Any2Graph allows us to leverage inputs of various types, such as images or sets of tokens. We complete our framework with a novel challenging synthetic dataset which we demonstrate to be suited for benchmarking SGP models. The rest of the paper is organized as follows. After a reminder and a discussion about the relation between graph matching and optimal transport (Section 2), we introduce in Section 3, a size-agnostic graph representation and an associated differentiable and node permutation invariant loss. This loss, denoted as Partially Masked Fused Gromov Wasserstein (PMFGW) is a novel and necessary adaptation of the FGW distance [45]. This loss is then integrated into Any2Graph, an end-to-end learning framework depicted in Figure 1 and presented in Section 4. We express the whole framework objective as an ERM problem and highlight the adaptations necessary for extending existing deep learning architectures [36] to more general input modalities. Section 5, presents a thorough empirical study of Any2Graph on various datasets. We evaluate our approach on four real-world problems with different input modalities as well as Coloring, a novel synthetic dataset. As none of the existing approaches could cover the range of input modalities, nor 2 scale to very large-sized datasets, we adapted them for the purpose of fair comparison. The numerical results showcase the state-of-the-art performances of the proposed method in terms of prediction accuracy and ability to retrieve the right size of target graphs as well as computational efficiency. 2 Background on graph matching and optimal transport Graph representation and notations An attributed graph g with m nodes can be represented by a tuple (F, A) where F = [f1, . . . , fm]⊤∈Rm×d encodes node features with fi ∈Rd labeling each node indexed by i, A ∈Rm×m is a symmetric pairwise distance matrix that describes the graph relationships between the nodes such as the adjacency matrix or the shortest path matrix. Further, we denote Gm the set of attributed graphs of m nodes and G = SM m=1 Gm, the set of attributed graphs of size up to M, where the size refers to the number of nodes in a graph and the largest size M is an important hyperparameter. In the following, 1m ∈Rm is the all one vector and we denote σm = {P ∈{0, 1}m×m | P1m = 1m, PT 1m = 1m} the set of permutation matrices. Graph Isomorphism Two graphs g1 = (F1, A1), g2 = (F2, A2) ∈Gm are said to be isomorphic whenever there exists P ∈σm such that (F1, A1) = (PF2, PA2PT ), in which case we denote g1 ∼g2. In this work, we consider all graphs to be unordered, meaning that all operations should be invariant by Graph Isomorphism (GI). Comparing graphs of the same size Designing a discrepancy to compare graphs is challenging, for instance, even for two graphs of the same size ˆg = (ˆF, ˆA), g = (F, A), one cannot simply compute a point-wise comparison as it would not satisfy GI invariance. A solution is to solve a Graph Matching (GM) problem, i.e., to find the optimal matching between the graphs and compute the pairwise errors between matched nodes and edges. This problem can be written as the following GM(ˆg, g) = min P∈σm m X i,j=1 Pi,jℓF (ˆfi, fj) + m X i,j,k,l=1 Pi,jPk,lℓA( ˆAi,k, Aj,l). (1) In particular, with the proper choice of ground metrics ℓf and ℓA, this is equivalent to the popular Graph Edit Distance (GED) [33]. The minimization problem however is a Quadratic Assignment Problem (QAP) which is known to be one of the most difficult problems in the NP-Hard class [27]. To mitigate this computational complexity, Aflalo et al. [2] suggested to replace the space of permutation matrices with a convex relaxation. The Birkhoff polytope (doubly stochastic matrices) πm = {T ∈[0, 1]m×m | T1m = 1m, TT 1m = 1m} is the tightest of those relaxations as it is exactly the convex hull of σm which makes it a suitable choice [21]. Interestingly, the resulting metric is known in OT [46] field as a special case of the (Fused) Gromov-Wasserstein (FGW) distance proposed by [29]. FGW(ˆg, g) = min T∈πm m X i,j=1 Ti,jℓF (ˆfi, fj) + m X i,j,k,l=1