Any2Graph: Deep End-To-End Supervised Graph Prediction With An Optimal Transport Loss

Authors: paul  krzakala
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:01:15

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 properties (permutation invariance, differentiability and scalability) and is designed to handle any-sized graphs. Numerical experiments showcase the versatility of the approach that outperform 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).

Main paper: 
 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 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 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. 
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 = [f 1 , . . . , f m ] ⊤ ∈ R m×d encodes node features with f i ∈ R d labeling each node indexed by i, A ∈ R m×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 G m the set of attributed graphs of m nodes and G = M m=1 G m , 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, 1 m ∈ R m is the all one vector and we denoteσ m = {P ∈ {0, 1} m×m | P1 m = 1 m , P T 1 m = 1 m } the set of permutation matrices. Graph Isomorphism Two graphs g 1 = (F 1 , A 1 ), g 2 = (F 2 , A 2 ) ∈ G mare said to be isomorphic whenever there exists P ∈ σ m such that (F 1 , A 1 ) = (PF 2 , PA 2 P T ), in which case we denote g 1 ∼ g 2 . 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 ĝ = ( F, Â), 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 followingEQUATIONIn 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 | T1 m = 1 m , T T 1 m = 1 m } is the tightest of those relaxations as it isexactly 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] .T∈πm m i,j=1EQUATIONHowever, these two points of view differ in their interpretation of the FGW metric. From the GM perspective, FGW is cast as an approximation of the original problem, and the optimal transport plan is typically projected back to the space of permutation via Hungarian Matching [47] . From the OT perspective, FGW is used as a metric between distributions with interesting topological properties [44] . This raises the question of the tightness of the relaxation between GM and FGW. In the linear case, i.e., when ℓ A = 0, the relaxation is tight and this phenomenon is known in the OT literature as the equivalence between Monge and Kantorovitch formulation [31] . The quadratic case, however, is much more complex, and sufficient conditions under which the tightness holds have been studied in both fields [1, 35] .As seen above, both OT and GM perspectives offer ways to characterize the same objects. In the remainder of this paper, we adopt the OT terminology, e.g., we use the term transport plan in place of doubly stochastic matrix. We provide a quantitative analysis of the effect of the relaxation in F.2.Computing the FGW distance requires solving the optimization problem presented in Equation ( 2) whose objective rewrites ⟨T, U⟩ + ⟨T, L ⊗ T⟩ where U is a fixed matrix, L a fixed tensor and ⊗ the tensor matrix product. A standard way of solving this problem [47] is to use a conditional gradient (CG) algorithm which iteratively solves a linearization of the problem. Each step of the algorithm requires solving a linear OT/Matching problem of cost ⟨T, C (k) ⟩ where the linear cost C (k) = U + L ⊗ T (k) is updated at each iteration. The linear problem can be solved with a Hungarian solver with cost O(M 3 ) while the overall complexity of computing the tensor product L ⊗ T (k) is theoretically O(M 4 ). Fortunately, this bottleneck can be avoided thanks to a O(M 3 ) factorization proposed originally by Peyré et al. [32] .Comparing graphs of arbitrary size The metrics defined above cannot directly be used to compare graphs of different sizes. To overcome this problem, Vayer et al. [45] proposed a more general formulation that fully leverages OT to model weights on graph nodes and can be used to compare graphs of different sizes as long as they have the same total mass. However, this approach raises specific issues. In scenarios where masses are uniform, nodes in larger graphs receive lower mass which might not be suitable for practical applications. Conversely, employing non-uniform masses complicates interpretation, as decoding a discrete object from a weighted one becomes less straightforward. Those issues can be mitigated by leveraging Unbalanced Optimal Transport (UOT) [40] , which relaxes marginal constraints, allowing for different total masses in the graphs. Unfortunately, UOT introduces several additional regularization parameters that are difficult to tune, especially in scenarios like SGP, where model predictions exhibit wide variability during training.Another close line of work is Partial Matching (PM) [13] , which consists in matching a small graph g to a subgraph of the larger graph ĝ. In practice, this can be done by adding dummy nodes to g through some padding operator P after which one can directly compute PM(ĝ, g) = GM(ĝ, P(g)) [21] . However, PM is not suited to train a model as the learned model would only be able to predict a graph that includes the target graph. Partial Matching and its relationship with our proposed loss is discussed in more detail in Appendix B.3. 
optimal transport loss for supervised graph prediction: 
 A size-agnostic representation for graphs Our first step toward building an end-to-end SGP framework is to introduce a space Ŷ to represent any graph of size up to M .EQUATIONWe refer to the elements of Ŷ as continuous graphs, in opposition with discrete graphs of G. Here h i (resp. A i,j ) should be interpreted as the probability of the existence of node i (resp. edge [i, j] ). Any graph of G can be embedded into Ŷ with a padding operator P defined asEQUATIONWe denote Y = P(G) ⊂ Ŷ the space of padded graphs. For any padded graph in Y, the padding operator can be inverted to recover a discrete graph P -1 : Y → G. Besides, any continuous graph ŷ ∈ Ŷ can be projected back to padded graphs Y by a threshold operator T : Ŷ → Y. Note that Ŷ is convex and of fixed dimension which makes it ideal for parametrization with a neural network. Hence, the core idea of our work is to use a neural network to make a prediction ŷ ∈ Ŷ and to compare it to a target g ∈ G through some loss ℓ(ŷ, P(g)). This calls for the design of an asymmetric loss ℓ : Ŷ × Y → R + .An Asymmetric loss for SGP The Partially Masked Fused Gromov Wasserstein (PMFGW) is a loss between a padded target graph P(g) = (h, F, A) ∈ Y with real size m = ∥h∥ 1 ≤ M and a continuous prediction ŷ = ( ĥ, F, Â) ∈ Ŷ. We define PMFGW(ŷ, P(g)) as:EQUATIONLet us decompose this loss function to understand the extend to which it simultaneously takes into account each property of the graph. The first term ensures that the padding of a node is well predicted. In particular, this requires the model to predict correctly the number of nodes in the target graph. The second term ensures that the features of all non-padding nodes (h i = 1) are well predicted. Similarly, the last term ensures that the pairwise relationships between non-padded nodes (h i = h j = 1) are well predicted. The normalizations in front of the sums ensure that each term is a weighted average of its internal losses as T i,j = M , T i,j h j = m and T i,j T k,l h j h l = m 2 . Finally α = [α h , α f , α A ] ∈ ∆ 3 is a triplet of hyperparameters on the simplex balancing the relative scale of the different terms. For ℓ A and ℓ h we use the cross-entropy between the predicted value after a sigmoid and the actual binary value in the target. This is equivalent to a logistic regression loss after the OT plan has matched the nodes. For ℓ f we use the squared ℓ 2 or the cross-entropy loss when the node features are continuous or discrete, respectively.A key feature of this loss is its flexibility. Not only other ground losses can be considered but it is also straightforward to introduce richer spectral representations of the graph [4] . For instance, in Section 5, we explore the benefits of leveraging a diffused version of the nodes features.Finally, PMFGW translates all the good properties of FGW to the new size-agnostic representation. Proposition 1 (Complexity). The objective of the inner optimization can be evaluated in O(M 3 ). Proposition 2 (GI Invariance). If ŷ ∼ ŷ′ and g ∼ g ′ then PMFGW(ŷ, P(g)) = PMFGW(ŷ ′ , P(g ′ )). Proposition 3 (Positivity). PMFGW(ŷ, P(g)) ≥ 0 with equality if and only if ŷ ∼ P(g).See Appendix A for a toy example illustrating the behavior of the loss and Appendix B.1 and B.2 for formal statements and proofs of Proposition 1, 2 and 3.Relation to existing metrics PMFGW is an asymmetric extension of FGW [44] suited for comparing a continuous predicted graph with a padded target. The extension is achieved by adding (1) a novel term to quantify the prediction of node padding, and (2) the partial masking of the components of the second and third terms to reflect padding. It should be noted that in contrast to what is usually done in OT, the node masking vectors (h and ĥ) are not used as a marginal distribution but directly integrated into the loss. In that sense, the additional node masking term is very similar to the one of OTL p [39] that proposed to use uniform marginal weight and move the part that measures the similarity between the distribution weights in an additional linear term. However, OTL p is restricted to linear OT problems and does not use the marginal distributions as a masking for other terms as in PMFGW.PMFGW also relates to Partial GM/GW Chapel et al. [13] as both metrics compare graphs by padding the smallest one with zero-cost dummy nodes. The critical difference lies in the new vector ĥ which predicts which sub-graphs are activated, i.e., should be matched to the target. The exact relationship between Partial Fused Gromov Wasserstein (PFGW) and PMFGW is summarized below Proposition 4. If l h is set to a constant value, PMFGW is equal to PFGW (up to a constant).Remark 1. In that case, the vector ĥ disappears from the loss and cannot be trained. In particular, this would prevent the model from learning to predict the size of the target graph.The formal definition of PFGW and the proof of Proposition 4 are provided in Appendix B.3.4 Any2Graph: a framework for end-to-end SGPQuestion: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
any2graph problem formulation: 
 The goal of Supervised Graph Prediction (SGP) is to learn a function f : X → G using the training samples {(x i , g i )} n i=1 ∈ (X × G) n . In Any2Graph, we relax the output space and learn a function f : X → Ŷ that predicts a continuous graph ŷ := f (x) as defined in the previous section. Assuming f is a parametric model (in this work, a deep neural network) completely determined by a parameter θ, the Any2Graph objective writes as the following empirical risk minimization problem:EQUATIONAt inference time, we recover a discrete prediction by a straightforward decoding f (x) = P -1 •T (ŷ),where T is the thresholding operator with threshold 1 /2 on the edges and nodes and P -1 is the inverse padding defined in the previous section. In other words, the full decoding pipeline P -1 • T removes the nodes i (resp. edges (i, j)) whose predicted probability is smaller than 1 /2 i.e. ĥi < 1 /2 (resp. Âi,j < 1 /2)). Unlike surrogate regression methods, this decoding step is very efficient.The model fθ : X → Ŷ (left part of Figure 1 ) is composed of three modules , namely the encoder that extracts features from the input, the transformer that convert these features into M nodes embeddings, that are expected to capture both feature and structure information, and the graph decoder that predicts the properties of our output graph, i.e., ( ĥ, F, Â). As we will discuss later, the proposed architecture draws heavily on that of Relationformer [36] since the latter has been shown to yield to state-of-the-art results on the Image2Graph task.Encoder The encoder extracts k feature vectors in R de from the input. Note that k is not fixed a priori and can depend on the input (for instance sequence length in case of text input). This is critical for encoding structures as complex as graphs and the subsequent transformer is particularly apt at treating this kind of representation. By properly designing the encoder, we can accommodate different types of input data. In Appendix C, we describe how to handle images, text, graphs, and vectors and provide general guidelines to address other input modalities.Transformer This module takes as input a set of feature vectors and outputs a fixed number of M node embeddings. This resembles the approach taken in machine translation, and we used an architecture based on a stack of transformer encoder-decoders, akin to Shit et al. [36] .Graph decoder This module decodes a graph from the set of node embeddings Z = [z 1 , . . . , z M ] T using the following equation:EQUATIONwhere σ is the sigmoid function and MLP m , MLP f , MLP k s are multi-layer perceptrons heads corresponding to each component of the graph (mask, features, structure). The adjacency matrix is expected to be symmetric which motivate us to parameterize it as suggested by Zaheer et al. [56] .Positioning with Relationformer As discussed above, the architecture is similar to the one proposed in Relationformer [36] , with two modifications: (1) we use a symmetric operation with a sum to compute the adjacency matrix while Relationformer uses a concatenation that is not symmetric;(2) we investigate more general encoders to enable graph prediction from data other than images. However, as stated in the previous section, the main originality of our framework lies in the design of the PMFGW loss. Interestingly Relationformer uses a loss that presents similarities with FGW but where the matching is done on the node features only, before computing a quadratic-linear loss similar to PMFGW. In other words, they solve a bi-level optimization problem, where the plan is computed on only part of the information, leading to potentially suboptimal results on heterophilic graphs as demonstrated in the next section. 
numerical experiments: 
 This section is organized as follows. First, we describe the experimental setting (5.1) including baselines. Next, we showcase the state-of-the-art performances of Any2Graph for a wide range of metrics and datasets (5.2). Finally, we provide an empirical analysis of the key hyperparameters of Any2Graph (5.3). The code for Any2Graph and all the experiments will be released on GitHub.Datasets We consider 5 datasets thus covering a wide spectrum of different input modalities, graph types, and sizes. The first one, Coloring, is a new synthetic dataset that we proposed, inspired by the four-color theorem. The input is a noisy image partitioned into regions of colors and the goal is to predict the graph representing the regions as nodes (4 color classes) and their connectivity in the image. An example is provided in Figure 5 and more details are in Appendix D. Then, we consider four real-world benchmarks. Toulouse [5] is Sat2Graph datasets where the goal is to extract the road network from binarized satellite images of a city. USCities is also a Sat2Graph dataset but features larger and more convoluted graphs. Note that we leave aside the more complex RGB version of USCities as it was shown to require complex multi-level attention architecture [36] , which is beyond the scope of this paper. Finally, following Ucak et al. [42] , we address the Fingerprint2Graph task where the goal is to reconstruct a molecule from its fingerprint representation (list of tokens). We consider two widely different datasets for this tasks: QM9 [50] , a scarce dataset of tiny molecules (up to 9 nodes) and GBD13Blum and Reymond [7] , a large dataset2 featuring molecules with up to 13 heavy atoms. Additional details concerning the datasets (e.g. dataset size, number of edges, number of nodes) are provided in Appendix E.1.We compare Any2Graph, to our direct end-to-end competitor Relationformer [36] that has shown to be the state-of-the-art method for Image2Graph. For a fair comparison, we use the same architecture (presented in Figure 1 ) for both approaches so that the only difference is the loss. We conjecture that Any2Graph and Relationformer might benefit from feature diffusion, that is replacing the node feature vector F by the concatenation [F, AF] before training. We denote by "+FD" the addition of feature diffusion before training. Moreover, we also compare with a surrogate regression approach (FGW-Bary) based on FGW barycenters [9] . We test both the end-to-end parametric variant, FGWBary-NN, where weight functions α k , as well as K = 10 templates, are learned by a neural network and the non-parametric variant, FGWBary-ILE, where the templates are training samples and α k are learned by sketched kernel ridge regression [54, 18] with gaussian kernel. Both have been implemented using the codes provided by Brogat-Motte et al. [9] , modified to incorporate sketching. Hyperparameters regarding architectures and optimization are provided in Appendix E.2.The heterogeneity of our datasets, calls for task-agnostic metrics focusing on different fine-grained levels of the graph. At the graph level, we report the PMFGW loss between continuous prediction ŷ and padded target P(g) and the graph edit distance Gao et al. [19] between predicted graph P -1 T ŷ and target g. We also report the Graph Isomorphism Accuracy (GI ACC), a metric that computes the proportion of perfectly predicted graphs. At the edge level, we treat the adjacency matrix prediction as a binary prediction problem and report the Precision and Recall. Finally, at the node level, we report NODE, the fraction of well-predicted node features, and SIZE a metric reporting the accuracy of the predicted number of nodes. See Appendix E.3 for more details.Prediction Performances Table 1 shows the performances of the different methods on the five datasets. First, we observe that Any2Graph achieves state-of-the-art performances for all datasets and graph level metrics. On the Sat2Graph tasks (Toulouse and USCities) we note that Relationformer performs very close to Any2Graph. In fact, the features (2D positions) are enough to uniquely identify the nodes, making Relationformer's Hungarian matching sufficient. Moreover, both methods highly benefit from feature diffusion on Fingerprint2Graph tasks which we discuss further in Appendix F. Note that both barycenter methods struggle on Toulouse, possibly due to a lack of expressivity. On QM9 on the other hand, FGW-Bary-ILE performs close to Any2Graph but is still outperformed, even more so if we add feature diffusion. It should be stressed that FGW-Bary-ILE is placed here under the most favorable conditions possible, in particular with the use of a SOTA kernel and being given the ground truth number of nodes. Finally, on GDB13, the most challenging dataset, Any2Graph strongly outperforms all competitors. To summarize this part, we observe that Any2Graph and its variant Any2Graph+FD are consistently better by a large margin on three out of five benchmarks and tied for first place with RelationFormer on the remaining ones. Computational Performances Table 2 shows the cost of training and inference of all methods, expressed as the number of graphs processed per second. All values are computed on NVIDIA V100/Intel Xeon E5-2660. The heavy cost of barycenters computation in FGWBary makes it several orders of magnitude slower than Any2Graph. We note that Relationformer is faster at learning time because it avoids the need to solve a QAP. Overall, our proposed approach offers the best of both worlds, achieving SOTA prediction performance at all levels of the graph, at a very low computational inference cost.Qualitative Performances We display a sample of the graph predictions performed by the models in Figure 5 (left) along with a larger collection in appendix G. We observe that not only Any2Graph can adapt to different input modalities, but also to the variety of target graphs at hand. For instance, Coloring graphs tend to be denser, while Sat2Graph maps can contain multiple connected components and Fingerprint2Graph molecules exhibit unique substructures such as loops.Sensitivity to dataset size In figure 5 , we provide the test edit distance for different numbers of training samples in the explored datasets. Interestingly, we observe that a performance plateau is reached for all datasets. We also observe that Coloring/Toulouse are simpler than USCities/QM9 (in Scalability to larger graphs As stated in property 1, each iteration of the PMFGW solver scales with O(M 3 ). Denoting k(M ) the average number of iterations required for convergence, this means that the actual cost of computing the loss scales with O k(M )M 3 . We provide an empirical estimation of k(M ) in figure 2 which we obtain by computing PMFGW(P(g 1 ), P(g 2 )) for pairs of graphs g 1 , g 2 sampled from the Coloring dataset. We observe that k(M ) seems linear but can be made sub-linear using feature diffusion (FD). Still, the cubic cost prevents Any2Graph from scaling beyond a few tens of nodes. Choice of maximal graph size M The default value of M is the size of the largest graph in the train set. We explore whether or not overparametrizing the model with higher values bring substantial benefits. To this end, we train our model on Coloring for M between 10 (default value) and 25 and report the (test) edit distance. To quantify the effective number of nodes used by the model, we also record the number of active nodes, i.e., that are masked less than 99% of the time (see Figure 3 ). Interestingly, we observe that performances are robust w.r.t. the choice of M which can be explained by the number of useful nodes reaching a plateau. This suggests the model automatically learns the number of nodes it needs to achieve the required expressiveness. Sensitivity to the weights α We investigate the sensitivity of the proposed loss to the triplet of weight hyperparameters α. To this end, we train our model on Coloring for different values α on the simplex and report the (test) edit distance on Figure 4 . We observe that the performance is optimal for uniform α and robust to other choices as long as there is not too much weight on the structure loss term (corner α A = 1). Indeed, the quadratic term of the loss being the hardest, putting too much emphasis on it might slow down the training. This explains the efficiency of feature diffusion, as it moves parts of the structure prediction to the linear term. Further evidence backing this intuitive explanation is provided in F.1.h = 1 F = 1 A = 1 0 1 2 3 4 
conclusion and limitations: 
 We present Any2Graph, a novel deep learning approach to Supervised Graph Prediction (SGP) leveraging an original asymmetric Partially-Masked Fused Gromov-Wasserstein loss. To the best of our knowledge, Any2Graph stands as the first end-to-end and versatile framework to consistently achieve state-of-the-art performances across a wide range of graph prediction tasks and input modalities. Notably, we obtain excellent results in both accuracy and computational efficiency. Finally, we illustrate the adaptability of the proposed loss (using diffused features), and the robustness of the method to sensitive parameters such as maximum graph size and weights of the different terms in PMFGW.The main limitation of Any2Graph is its scalability to graphs of larger size. Considering the tasks at hand, in this paper, we limit ourselves to relatively small graphs of up to 20 nodes.For the future, we envision two working directions to address this issue. First, given the promising results with feature diffusion, we plan to introduce more tools from spectral graph theory [14, 20, 4] in Any2Graph, e.g., using diffusion on the adjacency matrix to capture higher-order interactions that may occur in large graphs. More generally, this extension could be useful even for smaller graphs. Secondly, we expect that the solver computing the optimal transport plan can be accelerated using approximation. In particular, the entropic regularization [32] might unlock the possibility of fully parallelizing the optimization on a GPU while low-rank OT solvers Scetbon et al. [34] could allow Any2Graph to scale to large output graphs. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We include a discussion regarding the computational cost of the different models in the main paper. Besides, we report the exact computing resources used to train Any2Graph and an estimate of the total computing cost of the research project in appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: On the one hand, we consider a target graph of size 2, g = (F 2 , A 2 ) whereF 2 = f 1 f 2 ; A 2 = 0 1 1 0for some node features f 1 and f 2 . For M = 3 the padded target is P(g) = (h, F, A) whereh = 1 1 0 ; F = f 1 f 2 - ; A = 0 1 - 1 0 - --- .On the other hand, we consider a predicted graph ŷa,h = ( ĥ, F, Â) that has the formĥ = 1 h 1 -h ; F = f 1 f 2 f 2 ; Â = 0 a 1 -a a 0 0 1 -a 0 0for some a, h ∈ [0, 1]. The loss between the prediction and the (padded) target isL train (a, h) = PMFGW(ŷ a,h , P 3 (g))We are interested in the landscape of this loss. First of all, it appears that ŷ1,1 and ŷ0,0 and P(g) are isomorphic, thus we get two global minima L train (1, 1) = L train (0, 0) = 0. Going into greater detail, it can be shown that for 2 we have the following expressionℓ h (a, b) = ℓ A (a, b) = (a -b)L train (a, h) = min (1 -a) 2 + 2 3 (1 -h) 2 ; a 2 + 2 3 h 2and the optimal transport plan is the permutation (1, 2, 3) when (1 -a) 2 + 2 3 (1 -h) 2 < a 2 + 2 3 h 2 and (1, 3, 2) otherwise. In this toy example, the optimal transport plan is always a permutation.At inference time, we could similarly be interested in the edit distance between the (discrete) prediction and the target L eval (a, h) = ED(P -1 3 T (ŷ a,h ), g).Once again, the expression can be computed explicitlyL eval (a, h) = 1[a < 0.5 and h > 0.5] + 1[a > 0.5 and h < 0.5]We provide in Figure 6 an illustration of the edit distance and the proposed loss that is clearly a continuous and smoothed version of the edit distance which allows for learning the NN parameters.In this section, we write PMFGW(ŷ, P(g)) = minT∈π M i,j T i,j ℓ h ( ĥi , h j )+ i,j T i,j ℓ f ( fi , f j )h j + i,j,k,l T i,j T k,l ℓ A ( Âi,k , A j,l )h j h l ,meaning that we absorb the normalization factors in the ground losses to lighten the notation.Alternatively, we also consider the matrix formulation:PMFGW(ŷ, P(g)) = min The red line represents the transition between the regime where the optimal transport plan is the permutation (1, 2, 3) and that where it is (1, 3, 2). In both cases, the optimal transport plan is a permutation.where C i,j = ℓ h ( ĥi , h j ) + ℓ f ( fi , f j )h j , L i,j,k,l = ℓ A ( Âi,k , A j,l )h j h l and ⊗ is the tensor/matrix product.The following results generalize the Proposition 1 of Peyré et al. [32] so that it can be applied to the computation of PMFGW.Proposition 5. Assuming that the ground loss than can decomposed as ℓ(a, b) = f 1 (a) + f 2 (b)h 1 (a)h 2 (b), for any transport plan T ∈ R n×m and matrices A, W ∈ R n×n and A ′ , W ′ ∈ R m×m , then the tensor product of the form(L ⊗ T) i,i ′ = j,j ′ T j,j ′ ℓ(A i,j , A ′ i ′ ,j ′ )W i,j W ′ i ′ ,j ′can be computed asL ⊗ T = U 1 TW ′T + WTU T 2 -V 1 TV T 2 , where U 1 = f 1 (A) • W, U 2 = f 2 (A ′ ) • W ′ , V 1 = h 1 (A) • W, V 2 = h 2 (A ′ ) • W ′ and [•] is the point-wise multiplication.Proof. Thanks to the decomposition assumption the tensor product can be decomposed into 3 terms:(L ⊗ T) i,i ′ = j,j ′ T j,j ′ f 1 (A i,j )W i,j W ′ i ′ ,j ′ + j,j ′ T j,j ′ f 2 (A ′ i ′ ,j ′ )W i,j W ′ i ′ ,j ′ - j,j ′ T j,j ′ h 1 (A i,j )h 2 (A ′ i ′ ,j ′ )W i,j W ′ i ′ ,j ′ . = j f 1 (A i,j )W i,j j ′ T j,j ′ W ′ i ′ ,j ′ + j ′ f 2 (A ′ i ′ ,j ′ )W ′ i ′ ,j ′ j T j,j ′ W i,j - j h 1 (A i,j )W i,j j ′ T j,j ′ h 2 (A ′ i ′ ,j ′ )W ′ i ′ ,j ′ .Introducing U 1 , U 2 , V 1 and U 2 as defined above, we write:(L ⊗ T) i,i ′ = j (U 1 ) i,j j ′ T j,j ′ W ′ i ′ ,j ′ + j ′ (U 2 ) i ′ ,j ′ j T j,j ′ W i,j - j (V 1 ) i,j j ′ T j,j ′ (V 2 ) i ′ ,j ′ , = j (U 1 ) i,j (T W ′T ) j,i ′ + j ′ (U 2 ) i ′ ,j ′ (W T ) i,j ′ - j (V 1 ) i,j (T j,j ′ V T 2 ) j,i ′ , which concludes that L ⊗ T = U 1 TW ′T + WTU T 2 -V 1 TV T 2 .Remark 2 (Computational cost). U 1 , U 2 , V 1 , V 2 can be pre-computed for a cost of O(n 2 + m 2 ), after which L ⊗ T can be computed (for any T) at a cost of O(mn 2 + nm 2 ).Remark 3 (Kullback-Leibler divergence decomposition). The Kullback-Leibler divergence KL(p, q) = q log q p + (1 -q) log (1-q) (1-p) , which we use as ground loss in our experiments satisfies the required decomposition given f 1 (p) = -log(p), f 2 (q) = q log(q) + (1 -q) log(1 -q), h 1 (p) = log( 1-p p ) and, h 2 (q) = 1 -q Remark 4. The tensor product that appears in PMFGW is a special case of this theorem that corresponds to n = m = M , W i,j = 1 and W ′ i ′ ,j ′ = h i ′ h j ′ . Thus, proposition 1 is a direct corollary.First, we provide below a more detailed version of Proposition 2Proposition 6 (GI Invariance). For any m ≤ M , ŷ, ŷ′ ∈ Ŷ and g, g ′ ∈ G m , we have thatŷ ∼ ŷ′ , g ∼ g ′ =⇒ PMFGW(ŷ, P(g)) = PMFGW(ŷ ′ , P(g ′ )).Proof. We denote ŷ = ( ĥ, F, Â) and P(g) = (h, F, A). Since ŷ and ŷ′ are isomorphic, there exist a permutation P ∈ σ M such that ŷ′ = (P ĥ, P F, P ÂP T ). Moreover, the fact that g and g ′ are isomorphic implies that P(g) and P(g ′ ) are isomorphic as well, thus there exist a permutation Q ∈ σ M such that P(g ′ ) = (Qh, QF, QAQ T ). Plugging into the PMFGW objective we get PMFGW(ŷ ′ , P(g ′ )) = minT∈π M i,jT i,j ℓ h ((P ĥ) i , (Qh) j ) + i,jT i,j ℓ f ((P F) i , (QF) j )(Qh) j + i,j,k,l T i,j T k,l ℓ A ((P ÂP T ) i,k , (QAQ T ) j,l )(Qh) j (Qh) lT∈π M i,j(P T TQ) i,j ℓ h ( ĥi , h j ) + i,j (P T TQ) i,j ℓ f ( fi , f j )h j + i,j,k,l (P T TQ) i,j (P T TQ) k,l ℓ A ( Âi,k , A j,l )h j h l .Denoting T = P T TQ we have that PMFGW(ŷ ′ , P(g ′ )) = minT∈π M i,j Ti,j ℓ h ( ĥi , h j ) + i,j Ti,j ℓ f ( fi , f j )h j + i,j,k,lTi,j Tk,l ℓ A ( Âi,k , A j,l )h j h l = PMFGW(ŷ, P(g)).We now provide a more detailed version of Proposition 3Definition 1. We say that ℓ : X × X → R is positive when for any x, y ∈ X × X , ℓ(x, y) ≥ 0 with equality if and only if x = y. Proposition 7 (Positivity). Let us assume thatℓ h : [0, 1] × {0, 1} → R, ℓ f : R d × R d → R and ℓ A : [0, 1] × {0, 1} → R are positive.Then we have that for any ŷ ∈ Ŷ, g ∈ G :• i) PMFGW(ŷ, P(g)) ≥ 0• ii) There is equality if and only if ŷ ∼ P(g)• iii) In that case P -1 T (ŷ) ∼ gProof. The direct implication of ii) is the only statement that is not trivial. First, let us show that if PMFGW(ŷ, P(g)) = 0, the optimal transport T * is a permutation. Recall that any transport plan is a convex combination of permutations [6] i.e. there exist λ 1 , ..., λ K ∈]0, 1] and P 1 , ...,EQUATIONThis is a sum of positive terms, thus all terms are null and in particular, for any kEQUATION)Thus all the P k are optimal transport plans. In the following, we chose one of them and denote it P.Moving back to the developed formulation of PMFGW we get that0 = PMFGW(ŷ, P(g)) = i,j P i,j ℓ h ( ĥi , h j )+ i,j P i,j ℓ f ( fi , f j )h j + i,j,k,l P i,j P k,l ℓ A ( Âi,k , A j,l )h j h l .Once again this is a sum of positive terms thus for all i, j, k, and l0 = P i,j ℓ h ( ĥi , h j ) = P i,j ℓ f ( fi , f j )h j = P i,j P k,l ℓ A ( Âi,k , A j,l )h j h l and thus 0 = ℓ h ((P T ĥ) j , h j ) = ℓ f ((P T F) j , F j )h j = ℓ A ((P T ÂP) j,l , A j,l )h j h l .And from the positivity of ℓ h , ℓ f and ℓ A we get that: P T ĥ = h, P T Following Chapel et al. [13] , we define an OT relaxation of the Partial Matching problem.For a large graph ĝ = ( F, Â) ∈ G M and a smaller graph g = (F, A) ∈ G m , the set of transport plan transporting a subgraph of ĝ to g can be defined asπ M,m = {T ∈ [0, 1] M ×m | T1 m ≤ 1 M , T T 1 M = 1 m , 1 T M T1 m = m}and the associated partial Fused Gromov Wasserstein distance is partialFGW(ĝ, g) = minT∈π M,m M i=1 m j=1 T i,j ℓ F ( fi , f j ) + M i,k=1 m j,l=1 T i,j T k,l ℓ A (A i,k , A j,l ).In the following, we show that partialFGW(ĝ, g) is equivalent to the padded Fused Gromov Wasserstein distance defined as paddedFGW(ĝ, g) = minT∈π M M i=1 m j=1 T i,j ℓ F ( fi , f j ) + M i,k=1 m j,l=1 T i,j T k,l ℓ A (A i,k , A j,l ).Lemma 5. Any transport plan T ∈ π M has the form T = (T p T 2 ) where T p is a partial transport plan i.e. T p ∈ π M,m .Proof. Let us check that T p is in π M .•1 M = T1 M = T p 1 m + T 2 1 M -m ≥ T p 1 m • 1 M = T T 1 M = T T p 1 M T T 2 1 M . And thus T T p 1 M = 1 m• From the previous we immediately get that 1 T M T p 1 m = m Lemma 6. For any partial transport plan T p ∈ π M,m there existT 2 ∈ R M ×(M -m) such that T = (T p T 2 ) ∈ π M .Proof. Let us define p = 1 M -T p 1 m . This is the mass of the larger graph that is not matched by T p .Note that since T p ∈ π M,m we have that p ≥ 0. Thus we can set T 2 = 1 M -m p1 T M -m i.e. we spread the remaining mass uniformly across the padding nodes. Let us check that T = (T p T 2 ) ∈ π M is indeed a valid transport plan.• T1 M = T p 1 m + T 2 1 M -m = T p 1 m + p = 1 m • T T 1 M = T T p 1 M T T 2 1 M = 1 m 1 M -m (p T 1 M )1 M -m = 1 m 1 M -m (1 T M 1 M -1 T m T T p 1 M )1 M -m = 1 m 1 M -m = 1 MProposition 8. paddedFGW and partialFGW are equal and any optimal plan T * of paddedFGW has the form T * = (T * p , T 2 ) where T * p is optimal for partialFGW.Proof. Follows directly from the two previous lemmas.Remark 7. Since paddedFGW is equivalent to the proposed PMFGW loss if and only if ℓ h is set to a constant, proposition 4 is a direct corollary of Proposition 8. Remark 8. The algorithm proposed to compute PMFGW can be applied to paddedFGW and thus to partialFGW. Hence, we have indirectly introduced an alternative to the algorithm of Chapel et al. [13] . Further comparisons are left for future work.Philosophy of the Any2Graph encoder Any2Graph is compatible with different types of inputs, given that one selects the appropriate encoder. The role of the encoder is to extract a set of feature vectors from the inputs x i.e. each input is mapped to a list of k feature vectors of dimension d e where k is not necessarily fixed. This is critically different from extracting a unique feature vector (k = 1).If k is set to 1, the rest of the architecture must reconstruct an entire graph from a single vector, and the architecture is akin to that of an auto-encoder. In Any2Graph, we avoid this undesirable bottleneck by opting for a richer (k > 1) and more flexible (k is not fixed) representation of the input. The k feature vectors are then fed to a transformer which is well suited to process sets of different sizes. Since the transformer module is permutation-invariant any meaningful ordering is lost in the process. To alleviate this issue, we add positional encoding to the feature vectors whenever the ordering carries information. Finally, note that the encoder might highly benefit from pre-training whenever applicable; but this goes beyond the scope of this paper.We now provide a general description of the encoders that can be used for each input modality. Images For Image2Graph task we use Convolutional Neural Networks (CNN) as suggested in [36] .From an input image of shape h × w × c the CNN outputs a tensor of shape H × W × C which is seen as H × W feature vectors of dimension C. The raw output of the CNN is reshaped and passed through a linear layer to produce the final output of shape H × W × d e . Since the ordering of the H × W features carries spatial information we add positional encoding accordingly.Fingerprint/text For tasks where the input is a list of tokens (e.g. Text2Graph or Fingerprint2Graph) we use the classical NLP pipeline: each token is transformed into a vector by an embedding layer and the list of vectors is then processed by a transformer encoder module. In text2graph the tokens ordering carries semantic meaning and positionnal encoding should be added. On the contrary, in Fingerprint2Graph, the fingerprint ordering carries no information and the permutation invariance of the transformer module is a welcomed property.Graph For a graph2graph task (not featured in this paper) we would suggest using a Graph Neural Network (GNN) [52] . A GNN naturally extracts k feature vectors from an input graph, where k is the number of nodes in the input graph. No positional encoding is required.Vector We explore a Vect2Graph task in Appendix D. The naive encoder we use is composed of k parallel MLPs devoted to the extraction of the k feature vectors. This approach is arguably simplistic and more suited encoders should be considered depending on the type of data.We introduce Coloring, a new synthetic dataset well suited for benchmarking SGP methods. The main advantages of Coloring are:• The output graph is uniquely defined from the input image.• The complexity of the task can be finely controlled by picking the distribution of the graph sizes, the number of node labels (colors) and the resolution of the input image. • One can generate as many pairs (inputs, output) as needed to explore different regimes, from abundant to scarce data.To generate a new instance of Coloring, we apply the following steps:• 0) Sample the number of nodes (graph size) m. In this paper, we sample uniformly on some interval [M min , M max ]. As highlighted above, Coloring is a flexible dataset. Beyond the default dataset simply referred as Coloring we also explore 2 variations in our experiments. ColoringBig is a more challenging dataset that features larger graphs. ColoringVect is a variation of Coloring where the input image is flattened and treated as a vector allowing us to explore a synthetic Vect2Graph task. The properties of these datasets, along with the performances of Any2Graph are reported in table 3 . We hope that Coloring will be used to benchmark future SGP methods.3: Summary of the properties of the 3 variations of Coloring considered in this paper. We also report the test edit distance achieved by the different models. For FGWBary we report the best performing variant that is ILE for ColoringVect and NN for Coloring. None scales to ColoringBig. In the following, we provide a detailed description of the metrics reported in table 1 .Graph Level First we report the PMFGW loss between continuous prediction ŷ and padded target P(g). For this computation, we set α to the values displayed in table 5 .PMFGW(ŷ, P(g))Then we report the graph edit distance Gao et al. [19] between predicted graph P -1 T ŷ and target g which we compute using Pygmtools Wang et al. [48] . All edit costs (nodes and edges) are set to 1. Note that for Toulouse and USCities, node labels are 2D positions and we consider two nodes features to be equal (edit cost of 0) whenever the L2 distance is smaller than 5% than the image width.EDIT(P -1 T ŷ, g)Finally, we report the Graph Isomorphism Accuracy GI ACC that isGI ACC(ŷ, g) = 1[EDIT(P -1 T ŷ, g) = 0]Node level Recall that for a prediction ŷ = ( ĥ, F, Â) the size of the predicted graph is m = || ĥ > 0.5|| 1 . Denoting m the size of the target graph we report the size accuracy:SIZE ACC(ŷ, g) = 1[ m = m].The remaining node and edge-level metrics need the graphs to have the same number of nodes. To this end, we select the m nodes with the highest probability ĥi , resulting in a graph ĝ = ( F, Ã) with ground truth size. This is equivalent to assuming that the size of the graph is well predicted. Then we use Pygmtools to compute a one-to-one matching σ between the nodes of ĝ and g that can be used to align graphs (we use the matching that minimizes the edit distance). In the following, we assume that g and ĝ have been aligned. We can now define the node accuracy NODE ACC asNODE ACC(ĝ, g) = 1 m m i=1 1[ Fi = F i ],which is the average number of node features that are well predicted.Edge level Since the target adjacency matrices are typically sparse, the edge prediction accuracy is a poorly informative metric. To mitigate this issue we report both Edge Precision and Edge Recall :EDGE PREC.(ĝ, g) = m i,j=1 1[ Ãi,j = 1, A i,j = 1] m i,j=1 1[ Ãi,j = 1] EDGE PREC.(ĝ, g) = m i,j=1 1[ Ãi,j = 1, A i,j = 1] m i,j=1 1[A i,j = 1]All those metrics are then averaged other the test set.We estimate the total computational cost of this work to be approximately 1000 hours of GPU (mostly Nvidia V100). We estimate that up to 70% of this computation time was used in preliminary work and experiments that did not make it to the paper.The PMFGW loss is composed of three terms, two of them are linear and account for the prediction of the nodes and their features, one is quadratic and accounts for the prediction of edges. The last term is arguably the harder to minimize for the model, as a consequence, we observe that the training performs best when the two first terms are minimized first which then guides the minimization of the structure term. In other words, the model must first learn to predict the nodes before addressing their relationship. Fortunately, this behavior naturally arises in Any2Graph as long as α A , the hyperparameter controlling the importance of the quadratic term, is not too large. This is illustrated in figure 9 . ). In the first scenario, the first two terms of the loss are learned very fast and the structure is optimized next. In the second scenario, setting α A = 10 prevents this desirable learning dynamic.For the datasets where many nodes in the graphs share the same features (QM9 and GDB13) the good prediction of the nodes and their features is not enough to guide the prediction of the edges and this desirable dynamic does not occur. This motivates us to perform Feature Diffusion (FD) before training. The diffused node features carry a portion of the structural information. This makes the node feature term slightly harder to minimize but in turn, the subsequent prediction of the structure is much easier and we recover the previous dynamic. This is illustrated in figure 10 . We observe that the feature loss decreases slightly slower with FD (the features are more complex) but the minimization of the structure term is largely accelerated.As stated in 2, we adopted the OT point of view when designing Any2Graph. In practice, this means that we do not project the OT plan back to the set of permutations with a Hungarian matcher before plugging it in the loss as in Simonovsky and Komodakis [37] . Testing the effect of adding this extra step we observed a 5% to 10% increase of the edit distance across datasets (table 6 ) along with a more unstable training curve (figure 11 ). This confirms that a continuous transport plan provides a slightly more stable gradient than a discrete permutation, which aligns with the findings of De Plaen et al. [16] on the similar topic of object detection.Coloring Toulouse USCities QM9 GDB13 We checked the robustness of Any2Graph to the seed used for training. For each dataset, we ran 5 times the whole training pipeline with 5 random seeds. The results are displayed in table 7 . Notably, we observe that all state-of-the-art performances observed in table 1 hold even if we had reported the worst seed (we reported the first seed).Coloring Toulouse USCities QM9 GDB13 We tested if, once trained on Toulouse dataset, the predictive model is able to cope with out-ofdistribution data. Figure 17 shows that this is the case on these toy images, that are not related to satellite images or road maps. We leave for future work the investigation of this property. NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Justification: To the best of our knowledge, this work is not harmful in any of the ways detailed in the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).All code is available at https://github.com/KrzakalaPaul/Any2Graph. 38th Conference on Neural Information Processing Systems (NeurIPS2024).For computational purposes we use the 'ABCDEFGH' subset (more than 1.3 millions molecules) . 


Decision: This paper presents Any2Graph, a versatile end-to-end deep learning framework designed for Supervised Graph Prediction (SGP). The core innovation of this framework is the introduction of a novel loss function, Partially-Masked Fused Gromov-Wasserstein (PMFGW), which extends the Fused Gromov-Wasserstein loss to accommodate graph prediction with arbitrary sizes and node orderings.  Along with a standard pipeline carefully modified from Relationformer, the framework demonstrates consistent and state-of-the-art performance across a diverse range of datasets, including real-world tasks like map construction from satellite images and molecule prediction from fingerprints. 

The reviewers praise the paper regarding its novelty, theorectical soundness, strong empirical performance, and clearity. The PMFGW loss is highlighted as a significant contribution, providing a differentiable, permutation-invariant, and size-agnostic approach to graph prediction, which is a meaningful extension of existing methods. Some minor issues were raised regarding the presentation, experimental evaluations, and complexity analyses. However, all reviewers were satisfied with the authors' explanations and provided positive support for the paper. The paper is recommended for acceptance.

Reviews: review: paper_summary: paper_summary: This work aims to design an end-to-end pipeline for structured graph prediction (SGP). The proposed loss, PMFGW, is a extension of Fused Gromov Wasserstein to generate graphs with bounded arbitrary sizes, along with a standard pipeline carefully modified from Relationformer. It is empirically verified on public tasks and a novel synthetic dataset, Coloring, with outstanding performance.
strength_weakness: strength_weakness: One may argue that a clearly successful paper could have more novelty, than proposing a loss with an additional term, the significance of which I am not sure about. The positive side is that the simple design has good empirical performance, and perhaps could inspire other researchers to simply add such a term to solve graph tasks with different sizes. 

Hence, I would like to provide a score of weak accept for now, and see how the other reviewers think about this potential argument.
questions: questions: 1. Could you please reveal more why the reported results about Relationformer in Table 1 are largely different from their original ones? I suppose the reason may be as written around line 278 ''use the same architecture for both approaches...'', but I still want to make sure the difference is well understood. Is there any other issue here, like data split?
limitations: limitations: Yes.
paper_summary: paper_summary: This work proposes Any2Graph, an end-to-end deep learning framework for Supervised Graph Prediction (SGP) leveraging a novel OT loss called PMFGW. The model consistently achieves state-of-the-art performances across multiple graph prediction tasks and input modalities.
strength_weakness: strength_weakness: - There are some inconsistent notations in the paper. For example, the loss function for node features is denoted as $l_F$ in Eq. (1), while the author uses $l_f$ in the next sentence.
- The author emphasizes the importance of the maximum node size $M$ that strongly influences efficiency and expressiveness of the model. However, the analysis in Section 5.3 shows $M=16$ is sufficient to tackle the problem, indicating that the datasets may be relatively simple and cannot fully reflect the prediction ability on more complex graphs with higher orders of magnitude of nodes.
questions: questions: - The author uses PMFGW as the evaluation metric at the graph level, which is also the training objective of the model. The comparison may be not fair in that case since the loss functions $l_h, l_f, l_A$ and weights $\alpha$ can be arbitrarily chosen. Please correct me if I am wrong.
- In Section 4, the author mentions that threshold 1/2 is chosen. I wonder whether and how the threshold influences the model's capability. Please show me some experimental results on one or more datasets.
- The datasets may be relatively simple since $M=16$ is sufficient to tackle the problem. Can you give some results on more complecated datasets (\emph{i.e.}, with higher orders of magnitude of nodes)?
limitations: limitations: - As mentioned in the paper, the main limitation is its scalability to graphs of larger size.
paper_summary: paper_summary: This paper presents Any2Graph, a generic framework for end-to-end supervised graph prediction (SGP) with an optimal transport loss. The framework handles various input modalities and output graphs of arbitrary size and node ordering. The novel Partially-Masked Fused Gromov-Wasserstein loss is differentiable and permutation invariant, making it suitable for SGP. Numerical experiments showcase the versatility and superior performance of Any2Graph on a synthetic dataset and several real-world tasks such as map construction and molecule prediction. The paper addresses a practical challenge in deep learning models and offers a promising solution.
strength_weakness: strength_weakness: 1. This paper addresses the SGP problem, characterized by size agnostic, node order insensitivity, and a vast search space. As the number of nodes increases, the computational load may significantly increase. However, the discussion of this issue in the paper is relatively limited.
2. Evaluation and Ablation Studies: The evaluation results are impressive but lack sufficient ablation studies. Additional analysis on the impact of different components (e.g., loss function) would enhance our understanding of the proposed method’s strengths and weaknesses.
questions: questions: 1. Can an analysis of computational complexity be provided? Can a computationally feasible solution be proposed for graph generation problems with a large number of nodes (e.g., graphs with more than 100 nodes, which are commonly encountered in various data sets)?
2. Have you considered performing ablation studies or comparisons against popular baselines? 
3. Can you discuss potential limitations or failure cases of your method? How might these impact real-world applications? Are there scenarios where traditional methods may outperform Any2Graph?
limitations: limitations: As Discussed In Sec. Weakness & Sec. Limitations
paper_summary: paper_summary: The authors propose a flexible framework for end-to-end Supervised Graph Prediction (SGP), called Any2Graph, capable of handling various types of input data. This framework leverages a novel, fully differentiable, and node permutation-invariant optimal transport-based loss called the Partially Masked Fused Gromov-Wasserstein (PMFGW) loss. Unlike the Fused Gromov-Wasserstein loss (FGW), which cannot compare graphs of different sizes, the PMFGW is size-agnostic. To satisfy this property, the discrete output space is relaxed to a continuous space. A discrete graph is mapped to a continuous graph using a padding operator. The PMFGW measures the discrepancy between the predicted continuous graph and the target padded graph. Essentially, PMFGW extends FGW by adding an additional term that ensures the padding of a node is well predicted, and the partial mapping of the second and third terms. The authors introduce three propositions demonstrating that PMFGW translates all properties of FGW to the new size-agnostic graph representation. The architecture of the end-to-end SGP framework is a modification of the Relationformer, ensuring versatility in terms of input data (not only images) and modifying the operation for computing the adjacency matrix. The authors showed that Any2Graph achieved state-of-the-art prediction performance at a very low computational inference cost on four real-world datasets characterized by different types of input data.
strength_weakness: strength_weakness: I think to improve the readability of of the paper I suggest to include a extend the description of the graph matching (section 2, paragraph 'Comparing graphs of the same size' by including a formal definition of the graph matching problem.
questions: questions: At the end of section 2, the authors state that Unbalanced OT introduces several additional regularization parameters (I believe three) that are difficult to tune, especially in scenarios like SGP, where model predictions exhibit wide variability during training. While I agree, I was wondering if the authors have conducted any preliminary experiments using a UOT-based loss.
limitations: limitations: The authors have adequately addressed the limitations of the Any2Graph framework. Their work does not appear to have any potential negative societal impact.
paper_summary: paper_summary: This paper presents a generic framework for end-to-end supervised graph prediction. The proposed framework can take different types of data as input and learn to output graphs. The core of the framework is a novel loss function (PMFGW) that enables generalizability in different scenarios. In the end, the paper demonstrates the capability and superior performance of Any2Graph over existing solutions on synthetic datasets as well as real-world datasets.
strength_weakness: strength_weakness: W1: The computational performance results could be more comprehensive. While the asymptotic time complexity is provided, it would be valuable to understand how training time scales with an increase in M. This information would help readers assess the practicality of using this approach for their specific use cases.

W2: It would be interesting to understand the failure cases produced by Any2Graph.

W3: Any2Graph produces continuous outputs; it would be interesting to know their distribution. How will the predicted graph affect the predicted graph if we use a different threshold, e.g., 0.1 or 0.9?
questions: questions: Please check the weaknesses section.
limitations: limitations: The paper's contributions are unlikely to have any negative societal impacts.

Comments: title: Reply to the rebuttal
Thanks!
comment: Thanks for your response. After reading the response, I will maintain my score for ''Technically solid, moderate-to-high impact paper''.
I appreciate the responses, especially the detailed answer regarding UOT as a loss. I will maintain the proposed score.
I appreciate the response and all generally makes sense. I'll maintain my score.
I'm quite appreciated for detailed explanations and additional experiments provided in the attachment. Results have demonstrated that the method is suitable for more complicated graphs with $M=50$ (though it could be better to show concrete values of the metric in the final version) and is robust to different threasholds as well. After reading the comments of other reviewers and the author's rebuttal, I'd like to change the rating to 7.

Wish you all the best.