A framework for differentiable Supervised Graph Prediction

Paul Krzakala 1 2 Junjie Yang 1 Rémi Flamary 2 Florence d’Alché-Buc 1 Charlotte Laclau 1 Matthieu Labeau 1

Abstract
We introduce a general framework to train a deepneural network to output a graph from a variety
of input modalities. The framework is built using
a novel Optimal Transport loss that exhibits allnecessary properties (permutation invariance anddifferentiability) and allows for handling graphsof any size. We showcase the versatility and stateof-the-art performances of the proposed approachon various real-world tasks and a novel challenging synthetic dataset.

1. Introduction

Figure 1: A sample of predictions from Any2Graph (ourmodel) and its direct competitor Relationformer for different
SGP tasks. Any2Graph can adapt to a variety of input
modalities and output graphs. We extend Relationformer
so that it can process inputs that are not images, but its loss
limits its performances when the nodes are not uniquely
identified by their features (e.g. for molecules). We provide
more qualitative results in appendix G.

This work addresses Supervised Graph Prediction (SGP),a supervised task where the output of the predictive model
is a graph. In contrast to supervised tasks involving graphs
as input that benefit from the hegemony of Graph Neural
Networks, there is no broadly accepted framework to address SGP with deep learning. This can be explained by thediscrete nature of the output set which makes it challenging
to parameterize with a neural network. Consequently, mostexisting methods rely on domain-specific heuristics, suchas introducing an ad hoc ordering of the nodes, to simplify
the task [1] [2] [3]. Another line of work is to leverageenergy-based models [4] or surrogate regression methods[5] to circumvent the difficulty of directly predicting a graph.
However, these approaches typically suffer from an expensive decoding step at inference time. To the best of ourknowledge, the graph barycenter of Brogat-Motte et al. [6]and the Relationformer model [7] are the only approaches
that can tackle a variety of SGP tasks in a fully end-to-end
manner. Unfortunately, those methods suffer from some
limitations which we discuss in detail in section 3.

ages a novel, fully differentiable, OT-based loss. The framework is completed with a novel synthetic dataset suited tothe evaluation of SGP methods. We demonstrate that our approach is versatile and achieves state-of-the-art performancefor diverse real-world tasks, such as constructing maps fromsatellite images (Sat2Graph) or predicting molecules fromfingerprints (Fingerprint2Graph).
Notations. An attributed graph g with m nodes can be
represented by a tuple (F, A) where F ∈ Rm×d encodes
node features and A ∈ Rm×m is the (symmetric) adjacency matrix. Further, we denote
SM Gm the set of attributedgraphs of m nodes and G = m=1 Gm , the set of attributedgraphs of size up to M . In the following, 1m ∈ Rm is theall one vector and we denote σm = {P ∈ {0, 1}m×m |P1m = 1m , PT 1m = 1m } the set of permutation matrices. Two graphs g1 = (F1 , A1 ), g2 = (F2 , A2 ) ∈ Gm
are said to be isomorphic whenever there exists P ∈ σmsuch that (F1 , A1 ) = (PF2 , PA2 PT ), in which case wedenote g1 ∼ g2 . In this work, we consider all graphs to beunordered, meaning that all operations should be invariant
by Graph Isomorphism (GI).

Contributions In this paper, we introduce Any2Graph, ageneral framework for end-to-end SGP. Any2Graph lever1Télécom Paris, IP Paris, LTCI 2 Ecole polytechnique, IP
Paris, CMAP UMR 7641. Correspondence to: Krzakala Paul<paul.krzakala@telecom-paris.fr>.
Published at the 2 nd Differentiable Almost Everything Workshop
at the 41 st International Conference on Machine Learning, Vienna, Austria. July 2024. Copyright 2024 by the author(s).

1A framework for differentiable Supervised Graph Prediction
ENCODER

TRANSFORMERGRAPH DECODER

PM-FGW Loss

MLP
Input data
dependent

TRANSFORMERENCODER

TRANSFORMERDECODER

MLP

1
1
1
0

MLP

Input

Set Of Features

Nodes queries

Nodes Embeddings

Prediction

Padded Target

Figure 2: Illustration of the architecture for a target graph of size 3 and M = 4

2. An Optimal Transport loss for SGP

any graph of size up to M . We define Ŷ as:

2.1. Graph matching and optimal transport.

{(h, F, A) | h ∈ [0, 1]M , F ∈ RM ×d , A ∈ [0, 1]M ×M }.
(2)
We refer to the elements of Ŷ as ’continuous’ graphs, in
opposition with ’discrete’ graphs of G. Here hi (resp. Ai,j )should be interpreted as the probability of the existence ofnode i (resp. edge [i, j] ). Any graph g = (Fm , Am ) ∈ Gm
can be embedded into Ŷ with a Padding operator P that
adds m′ = M − m dummy nodes to g

 
 

A
0m′
Fm
1m
. (3)
, Tm
,
P(g) =0m′0m ′0m′ 0m′ ,m′Designing a loss function to compare graphs is a challenging task. 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. One of thesolutions is to solve a Graph Matching (GM) problem i.e.
find a one-to-one matching between the nodes of the graphs
before computing the pairwise errors between matchednodes/edges. Formally, the graph matching discrepancy
is defined as GM(ĝ, g) = minP∈σm E(ĝ, g, P) where thematching cost E(ĝ, g, P) is
m
X

i,j=1

Pi,j ℓF (f̂i , fj ) +mX

We denote Y = P(G) ⊂ Ŷ 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 7→ G. Besides, any
continuous graph ŷ ∈ Ŷ can be projected back to padded
graphs Y by a threshold operator T : Ŷ 7→ Y. Note that Ŷ
is convex and of fixed dimension which makes it ideal forparametrization with a neural network. Hence, the core ideaof our work is to use a neural network to make a predictionŷ ∈ Ŷ and to compare it to some target g ∈ G through someloss ℓ(ŷ, P(g)). This calls for the design of an asymmetricloss ℓ : Ŷ × Y 7→ R+ .

Pi,j Pk,l ℓA (Âi,k , Aj,l ) (1)i,j,k,l=1

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 [8]. To
mitigate this computational complexity, it has been suggested to replace the space of permutation matrices witha convex relaxation [9]. 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 asuitable choice [10]. Interestingly, the resulting metric
FGW(ĝ, g) = minT∈πm E(ĝ, g, T) is known in the Optimal Transport (OT) [11] field as a special case of the(Fused) Gromov-Wasserstein (FGW) distance proposed byMémoli [12]. FGW is differentiable, GI invariant and efficient solvers are available which make it an ideal candidate
for SGP [13]. Unfortunately, this formulation is limited
to graphs of the same size, and existing variations are notsuited for SGP as discussed in appendix A.1.

2.3. 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 acontinuous prediction ŷ = (ĥ, F̂, Â) ∈ Ŷ, we define
PMFGW(ŷ, P(g)) as:αh XTi,j ℓh (ĥi , hj )T∈πM Mi,j
αf X
Ti,j ℓf (f̂i , fj )hj
+
m i,j
αA X
+ 2
Ti,j Tk,l ℓA (Âi,k , Aj,l )hj hlmmin

2.2. A size-agnostic representation for graphs
Our first step toward building an End-To-End SGP framework is to introduce a space Ŷ that can be used to represent

i,j,k,l2(4)

A framework for differentiable Supervised Graph Prediction

Table 1: Graph level, edge level, and node level metrics reported on test for the different models and datasets. ∗ denotes
methods that use the actual size of the graph at inference time, hence the performance reported is a non-realistic upper-bound.
G RAPH L EVEL
GI ACC . ↑ PMFGW ↓E DGE L EVELP REC . ↑ R EC . ↑N ODE L EVEL ACC .
N ODE ↑S IZE ↑

0.93
0.32
0.03

72.17
80.39
98.63

83.81
86.34
98.87

79.15
92.68
99.99

N.A.
99.32
99.85

0.00
93.28
93.62

1.21
0.02
0.02

72.52
99.25
99.34

56.30
99.24
99.26

1.62
99.25
99.39

N.A.
98.30
98.81

2.09
1.86

55.00
58.10

0.13
0.12

92.96
92.91

87.98
90.85

95.18
95.70

79.80
78.95

QM9

FGWBARY-ILE∗
R ELATIONFORMERA NY 2G RAPH

2.84
3.80
2.13

28.95
9.95
29.85

0.28
0.22
0.14

82.96
86.07
90.19

79.76
73.31
88.08

92.99
99.34
99.77

N.A.
96.0
95.45

GDB13

R ELATIONFORMERA NY 2G RAPH

8.83
3.63

0.01
16.25

0.29
0.11

84.14
90.83

55.89
84.86

97.57
99.80

98.65
98.15

DATASET

M ODELC OLORING

FGWBARY-ILE∗
R ELATIONFORMERA NY 2G RAPH

7.60
5.47
0.33

0.90
18.14
84.44

T OULOUSEFGWBARY-ILE∗
R ELATIONFORMERA NY 2G RAPH

9.00
0.13
0.13

USC ITIESR ELATIONFORMERA NY 2G RAPHE DIT D IST. ↓2.4. Any2Graph framework

The loss takes simultaneously into account each propertyof the graph. More precisely, the first term ensures thatthe padding of a node is well predicted. In particular, thisrequires 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 (hi = 1) are well predicted.
Similarly, the last term ensures that the relationships between pairs of non-padding nodes (hi = hj = 1) are well
predicted. The normalization in front of the sums ensures
thatPeach term isP
a weighted averagePof its internal losses
as Ti,j = M , Ti,j hj = m and Ti,j Tk,l hj hl = m2 .
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 thenode features are continuous or discrete, respectively. A key
feature of this loss is its flexibility. Not only any ground
losses can be considered but it is also straightforward to
introduce richer spectral representations of the graph [14].
For instance, we discuss the benefits of leveraging a diffusedversion of the nodes features F in appendix D.2. Finally,
PMFGW translates all the good properties of FGW to the
new size-agnostic representation:
Proposition 1 (Complexity). Each step of the inner optimization can be performed in O(M 3 ).
Proposition 2 (GI Invariance). If ŷ ∼ ŷ ′ and g ∼ g ′ then
PMFGW(ŷ, P(g)) = PMFGW(ŷ ′ , P(g ′ )).
Proposition 3 (Positivity). PMFGW(ŷ, P(g)) ≥ 0 withequality if and only if ŷ ∼ P(g).

We now wrap up the tools introduced above to build
Any2Graph, a general framework for SGP.
Formally, the goal of SGP is to learn a function f : X → G
using the training samples {(xi , gi )}ni=1 ∈ (X × G)n . In
Any2Graph, we relax the output space and learn a functionfˆ : X → Ŷ that predicts a continuous graph ŷ := f (x) as
defined in the previous section. Assuming fˆ is a parametricmodel (in this work, a deep neural network) completely determined by a parameter θ, the Any2Graph objective writes
as the following empirical risk minimization problem:

n

min
θ

1X
PMFGW(fˆθ (xi ), P(gi ))
n i=1

(5)

At inference time, we recover a discrete prediction by astraightforward decoding f (x) = P −1 ◦ T (ŷ). Where
T is the thresholding operator with threshold 1/2 on theedges and nodes and P −1 is the inverse padding definedin 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. ĥi <1/2
(resp. Âi,j <1/2)). Unlike surrogate regression methods, thisdecoding step does not induce any computational overhead.
The architecture we propose for fˆθ is heavily inspired bythat of Relationformer [7] and illustrated in figure 2. A fulldescription is provided in appendix B.1. The main novelty
is that we investigate more general encoders to enable graphprediction from data other than images. We discuss how tobuild encoders adapted to text, images, graphs and vectors
in appendix B.2 .

We discuss PMFGW in greater detail in appendix A. The
proofs of the propositions are provided in Appendix C.
3

A framework for differentiable Supervised Graph Prediction

3. Numerical experiments

Table 2: Speed of the different methods on QM9 (in graphper second). Training of FGWBary has a closed-form expression computed at once on CPU.

3.1. Experimental setting
We consider 5 datasets that cover a wide spectrum of different input modalities, graph types and sizes. The 4 real-world
datasets Toulouse, USCities, QM9 and GDB13 are completed with Coloring, a novel synthetic dataset described inappendix F. We compare Any2Graph to its direct competitor Relationformer [7] and a surrogate regression approach(FGW-Bary) based on FGW barycenters [6]. For a fair comparison, we use the same architecture for Relationformer
and Any2Graph and augment both with feature diffusion as
discussed in appendix D.2. Given the heterogeneity of thedatasets considered, we report task-agnostic metrics suchas average edit distance [15] or GI Accuracy (the fractionof graphs that are perfectly predicted). The detailed experimental setting is provided in appendix E.

M ETHODT RAIN . ↑P RED . ↑

N.A.
2K
1K

1
10 K
10 K

FGWBARYR ELATIONFORMERA NY 2G RAPH

Active Nodes

0.4

13
0.3

12
0.2

11
0.1

10

0.0

10 12 14 16 18 20

25

Edit Distance

0.5

14

F=1

A=1

4

h=1

0

Figure 3: Effect of the choice of M (left) and α (right) on
the performances (test edit distance) for Coloring.

3.2. Prediction Performances

quantify the effective number of nodes used by the model,
we also record the number of active nodes i.e. that aremasked less than 99% of the time. Results are reported in
Figure 3 (left). Interestingly, we observe that performances
are robust w.r.t. the choice of M which can be explained by
the number of active nodes reaching a plateau. This suggeststhat the model automatically learns how many nodes it needs
to reach the required expressiveness. Similarly, we checkthat our model is robust to the choice of parameter α whichbalances the terms of the loss. To this end, we train ourmodel on Coloring for different values α on the simplex
and report the (test) edit distance on Figure 3 (right). We
observe that performance is optimal for uniform α androbust to other choices as long as there is not too muchweight on the structure loss term (corner αA = 1). We
further discuss this property in appendix D.1.

Table 1 shows the performances of the different methods onthe five datasets. First, we observe that Any2Graph achieves
state-of-the-art performances for all datasets and graph-level
metrics. Relationformer is similar to PMFGW, except its
loss relies on a bipartite matching taking only the node features into account. The consequence is that Relationformer
performs very close to Any2Graph on the Sat2Graph tasks(Toulouse and USCities) but much poorly on any task where
nodes are not uniquely identified by their features. FGWBary relies on a computationally heavy barycenter whichin its current implementation cannot scale to the datasets
featuring the larger graphs (GDB13 and USCities). The
SOTA kernel it leverages is very efficient for QM9, despite
this, it is still outperformed by Any2Graph given that we
use feature diffusion.
3.3. Computational Performances

4. Conclusion and limitations

Table 2 shows the ”speed” of each method both in training
and inference. Speed is expressed in terms of the number ofgraph processes per second. Because of the barycenter computation, FGWBary is several orders of magnitude slower
than Any2Graph. Note that Relationformer is faster at training time because it does not require solving a QAP. Overall
our proposed approach strikes the best of both worlds byachieving SOTA prediction performances at all levels of
the graph, at a very low computational inference cost. All
values are computed on NVIDIA V100/Intel Xeon E5-2660.

We present Any2Graph, a novel end-to-end deep learning
approach to SGP based on an original asymmetric PartiallyMasked Fused Gromov-Wasserstein loss. To the best of ourknowledge, it is the first complete and versatile frameworkto achieve SOTA performance on various graph prediction
tasks and input modalities, both in terms of accuracy andcomputational cost, which makes it a good candidate to be
used as a novel baseline in SGP.
The main limitation of Any2Graph is its scalability to large
graphs. We envision two approaches to address this issue.
First, we plan to explore more general diffusion schemeson the adjacency matrix to capture higher-order interactions
that may occur in large graphs. Secondly, we wish to accelerate the OT plan computation with entropic regularization[16] to fully parallelize the solver on a GPU.

3.4. Analysing the proposed approach
First, we explore the effect of M (the maximum number ofnodes) whose default value is that of the largest graph in the
train set. We train our model on Coloring for M between 10(default value) and 25 and report the (test) edit distance. To
4

A framework for differentiable Supervised Graph Prediction

References
[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]

[11]
[12]

[13]

[14]

[15]

[16]

[17]

[18]

X. Bresson and T. Laurent, “A two-step graph convolutional decoder for molecule generation,” arXiv preprint
arXiv:1906.03412, 2019.
D. Belli and T. Kipf, “Image-conditioned graph generation for road network extraction,” arXiv preprint
arXiv:1910.14388, 2019.
A. Babu, A. Shrivastava, A. Aghajanyan, A. Aly, A. Fan,and M. Ghazvininejad, “Non-autoregressive semantic parsing for compositional task-oriented dialog,” arXiv preprint
arXiv:2104.04923, 2021.
M. Suhail, A. Mittal, B. Siddiquie, et al., “Energy-basedlearning for scene graph generation,” in Proceedings ofthe IEEE/CVF conference on computer vision and patternrecognition, 2021, pp. 13 936–13 945.
C. Brouard, H. Shen, K. Dührkop, F. d’Alché-Buc, S.
Böcker, and J. Rousu, “Fast metabolite identification with
input output kernel regression,” Bioinformatics, vol. 32,
no. 12, pp. i28–i36, 2016.
L. Brogat-Motte, R. Flamary, C. Brouard, J. Rousu, andF. d’Alché-Buc, “Learning to predict graphs with fused
gromov-wasserstein barycenters,” in International Conference on Machine Learning, PMLR, 2022, pp. 2321–2335.
S. Shit, R. Koner, B. Wittmann, et al., “Relationformer:
A unified framework for image-to-graph generation,” inEuropean Conference on Computer Vision, Springer, 2022,
pp. 422–439.
E. M. Loiola, N. M. M. De Abreu, P. O. Boaventura-Netto,
P. Hahn, and T. Querido, “A survey for the quadratic assignment problem,” European journal of operational research,
vol. 176, no. 2, pp. 657–690, 2007.
Y. Aflalo, A. Bronstein, and R. Kimmel, “On convex relaxation of graph isomorphism,” Proceedings of the National
Academy of Sciences, vol. 112, no. 10, pp. 2942–2947,
2015.
S. Gold and A. Rangarajan, “A graduated assignment algorithm for graph matching,” IEEE Transactions on patternanalysis and machine intelligence, vol. 18, no. 4, pp. 377–
388, 1996.
C. Villani, Optimal transport : old and new. Berlin:
Springer, 2009.
F. Mémoli, “Gromov-wasserstein distances and the metricapproach to object matching.,” Foundations of Computational Mathematics, vol. 11, no. 4, pp. 417–487, 2011.
T. Vayer, L. Chapel, R. Flamary, R. Tavenard, and N.
Courty, “Fused gromov-wasserstein distance for structured
objects,” Algorithms, vol. 13 (9), p. 212, 2020.
A. Barbe, M. Sebban, P. Gonçalves, P. Borgnat, and R. Gribonval, “Graph diffusion wasserstein distances,” in Joint
European Conference on Machine Learning and Knowledge Discovery in Databases, Springer, 2020, pp. 577–
592.
X. Gao, B. Xiao, D. Tao, and X. Li, “A survey of graphedit distance,” Pattern Analysis and applications, vol. 13,
pp. 113–129, 2010.
G. Rioux, Z. Goldfeld, and K. Kato, “Entropic gromovwasserstein distances: Stability and algorithms,” arXiv
preprint arXiv:2306.00182, 2023.
A. Thual, H. Tran, T. Zemskova, et al., “Aligning individual brains with fused unbalanced gromov-wasserstein,” inNeural Information Processing Systems (NeurIPS), 2022.

[19]

[20]

[21]

[22]

[23]

[24]

[25]
[26]

[27]

[28]

[29]

[30]

[31]

[32]

[33]

5

L. Chapel, M. Z. Alaya, and G. Gasso, “Partial optimal tranport with applications on positive-unlabeled learning,” Advances in Neural Information Processing Systems, vol. 33,
pp. 2903–2913, 2020.
T. Vayer, L. Chapel, R. Flamary, R. Tavenard, and N.
Courty, “Optimal transport for structured data with application on graphs,” in International Conference on Machine
Learning (ICML), 2019.
M. Thorpe, S. Park, S. Kolouri, G. K. Rohde, and D.
Slepčev, “A transportation lˆ p l p distance for signal analysis,” Journal of mathematical imaging and vision, vol. 59,
pp. 187–210, 2017.
J. T. Vogelstein, J. M. Conroy, V. Lyzinski, et al., “Fastapproximate quadratic programming for large (brain) graphmatching,” arXiv preprint arXiv:1112.5507, 2011.
G. Peyré, M. Cuturi, and J. Solomon, “Gromov-wassersteinaveraging of kernel and distance matrices,” in International
conference on machine learning, PMLR, 2016, pp. 2664–
2672.
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R.
Salakhutdinov, and A. J. Smola, “Deep sets,” Advances inneural information processing systems, vol. 30, 2017.
K. Xu, W. Hu, J. Leskovec, and S. Jegelka, “Howpowerful are graph neural networks?” arXiv preprint
arXiv:1810.00826, 2018.
G. Birkhoff, “Three observations on linear algebra,” Univ.
Nac. Tacuman, Rev. Ser. A, vol. 5, pp. 147–151, 1946.
M. Simonovsky and N. Komodakis, “Graphvae: Towards
generation of small graphs using variational autoencoders,”
in Artificial Neural Networks and Machine Learning–
ICANN, Springer, 2018, pp. 412–422.
H. De Plaen, P.-F. De Plaen, J. A. Suykens, M. Proesmans,
T. Tuytelaars, and L. Van Gool, “Unbalanced optimal transport: A unified framework for object detection,” in Proceedings of the IEEE/CVF Conference on Computer Vision
and Pattern Recognition, 2023, pp. 3198–3207.
Z. Wu, B. Ramsundar, E. N. Feinberg, et al., “Moleculenet:
A benchmark for molecular machine learning,” Chemical
science, vol. 9, no. 2, pp. 513–530, 2018.
L. C. Blum and J.-L. Reymond, “970 million druglike small
molecules for virtual screening in the chemical universe
database gdb-13,” Journal of the American Chemical Society, vol. 131, no. 25, pp. 8732–8733, 2009.
U. V. Ucak, I. Ashyrmamatov, and J. Lee, “Reconstructionof lossless molecular representations from fingerprints,”
Journal of Cheminformatics, vol. 15, no. 1, pp. 1–11, 2023.
G. Landrum et al., “Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling,”
Greg Landrum, vol. 8, no. 31.10, p. 5281, 2013.
Y. Yang, M. Pilanci, and M. J. Wainwright, “Randomized
sketches for kernels: Fast and optimal nonparametric regression,” The Annals of Statistics, vol. 45, no. 3, pp. 991–
1023, Jun. 2017. DOI: 10.1214/16- AOS1472. [Online]. Available: https://doi.org/10.1214/16AOS1472.
T. El Ahmad, P. Laforgue, and F. d’Alché-Buc, “Fast Kernel
Methods for Generic Lipschitz Losses via \p\-Sparsified
Sketches,” Transactions on Machine Learning Research,2023.

A framework for differentiable Supervised Graph Prediction
[34]

[35]

[36]

[37]

[38]
[39]

[40]

R. Wang, Z. Guo, W. Pan, et al., “Pygmtools: A python
graph matching toolkit,” Journal of Machine Learning Research, vol. 25, no. 33, pp. 1–7, 2024. [Online]. Available:
https : / / jmlr . org / papers / v25 / 23 - 0572 .
html.
K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in Proceedings of the IEEE
conference on computer vision and pattern recognition,
2016, pp. 770–778.
R. Xiong, Y. Yang, D. He, et al., “On layer normalization in
the transformer architecture,” in International Conference
on Machine Learning, PMLR, 2020, pp. 10 524–10 533.
A. Vaswani, N. Shazeer, N. Parmar, et al., “Attention is
all you need,” Advances in neural information processing
systems, vol. 30, 2017.
D. P. Kingma and J. Ba, “Adam: A method for stochastic
optimization,” arXiv preprint arXiv:1412.6980, 2014.
A. Okabe, B. Boots, K. Sugihara, and S. N. Chiu, “Spatial
tessellations: Concepts and applications of voronoi diagrams,” 2009.
R. M. Karp, Reducibility among combinatorial problems.
Springer, 2010.

6

A framework for differentiable Supervised Graph Prediction

A. Additional details about PMFGWThis section is organised as follows. First we discuss existing discrepancy to compare graphs of different sizes and why
we believe that they are not suited for SGP A.1. Then, we highlight the relationship between PMFGW and other existing
metrics A.2. Next, we describe the solver we use to compute the PMFGW optimal transport plan A.3. Finally, we introduce
a toy example illustrating the loss landscape of PMFGW on a toy example A.4
A.1. Comparing graphs of arbitrary size
The GM and FGW described in section 2.1 cannot directly be used to compare graphs of different sizes but several lines ofwork have been proposed to address this limitation.
The first approach is to fully leverage Optimal Transport by introducing weights on the graph nodes. The Fused-GromovWasserstein distance can then be used to compare graphs of different sizes as long as they have the same total mass [13].
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 complicatesinterpretation, as decoding a discrete object from a weighted one becomes less straightforward. Those issues can be mitigatedby leveraging Unbalanced Optimal Transport (UOT) [17], 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) [18], which consists in matching a small graph g to a subgraph of thelarger graph ĝ. 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̂, g) = GM(ĝ, P(g)) [10]. However, PM is not suited to train a model as the learned model
would only predict a graph that includes the target graph, with no indication of which subgraph is actually the correct
prediction. We discuss the relationship between PMFGW and Partial Matching in the next section.
A.2. Relation to existing metrics
PMFGW is an asymmetric extension of FGW [19] 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 maskingof 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 ĥ) are not used as a marginal distribution but directly integrated into theloss. In that sense, the additional node masking term is very similar to the one of OTLp [20] 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, OTLp 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 [18] as both metrics compare graphs by padding the smallest one with zero-cost
dummy nodes. The critical difference lies in the new vector ĥ which predicts which sub-graphs are activated, i.e., should bematched to the target. To be precise, PMFGW and Partial Fused Gromov Wasserstein (PFGW) are equal if and only if lh is
set to 0. We prove this statement and provide a formal definition of PFGW in C.3. Note that setting lh = 0 would obviously
be undesirable since the vector ĥ would disappears from the loss and the model would not be trained to predict it correctly.
In particular, this would prevent the model from learning to predict the size of the target graph.
A.3. PMFGW solver
Computing PMFGW requires solving the inner optimization problem presented in Equation (4) 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 [21]is to use a conditional gradient (CG) algorithm which iteratively solves a linearization of the problem. Each step of thealgorithm 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 (Proposition 1) that extend a result from Peyré et al. [22].
7A framework for differentiable Supervised Graph Prediction

A.4. Illustrating PMFGW On A Toy Example
On the one hand, we consider a target graph of size 2, g = (F2 , A2 ) where


 
0 1
f1
; A2 =
F2 =f2
1 0
for some node features f1 and f2 . For M = 3 the padded target is P(g) = (h, F, A) where

 
 
0
f1
1
h = 1 ; F = f2  ; A =  1
−
0
−


−
− .
−

1
0
−

On the other hand, we consider a predicted graph ŷa,h = (ĥ, F̂, Â) that has the form

 

0
f1
1
ĥ =  h  ; F̂ = f2  ; Â =  a
f2
1−a
1−h


a
0
0


1−a
0 
0

for some a, h ∈ [0, 1]. The loss between the prediction and the (padded) target is
Ltrain (a, h) = PMFGW(ŷa,h , P3 (g))
We are interested in the landscape of this loss. First of all, it appears that ŷ1,1 and ŷ0,0 and P(g) are isomorphic,
thus we get two global minima Ltrain (1, 1) = Ltrain (0, 0) = 0. Going into greater detail, it can be shown that forℓh (a, b) = ℓA (a, b) = (a − b)2 we have the following expression
2
2
Ltrain (a, h) = min (1 − a) + (1 − h)2 ; a2 + h2
3
3
2



and the optimal transport plan is the permutation (1, 2, 3) when (1 − a)2 + 32 (1 − h)2 < a2 + 32 h2 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
Leval (a, h) = ED(P3−1 T (ŷa,h ), g).
Once again, the expression can be computed explicitly
Leval (a, h) = 1[a < 0.5 and h > 0.5] + 1[a > 0.5 and h < 0.5]
We provide in Figure 4 an illustration of the edit distance and the proposed loss that is clearly a continuous and smoothedversion of the edit distance which allows for learning the NN parameters.
1.0

1.0

1.0
0.6

0.6

0.4

0.5

h

h

0.8

0.5
0.4

0.2
0.2
0.0
0.0

0.5

1.0

0.0

0.0

0.0

0.0
0.5

1.0

Figure 4: Heatmap of Ltrain (left) and Leval (right). The red line represents the transition between the regime where theoptimal plan is the permutation (1, 2, 3) and that where it is (1, 3, 2). In both cases, the optimal plan is a permutation.
8

A framework for differentiable Supervised Graph Prediction

B. Additional details about the architecture
B.1. Overview
The architecture that we use to parameterize f : X → Ŷ (left part of Figure 2) is composed of three modules , namely theencoder that extracts features from the input, the transformer that convert these features into M nodes embeddings, thatare expected to capture both feature and structure information, and the graph decoder that predicts the properties of our
output graph, i.e., (ĥ, F̂, Â). As we will discuss later, the proposed architecture draws heavily on that of Relationformer [7]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 Rde 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. We describe how to handle images, text, graphs, and vectors and provide
general guidelines to address other input modalities in the next section.
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 [7].
Graph decoder
equation:

This module decodes a graph from the set of node embeddings Z = [z1 , . . . , zM ]T using the followingĥi = σ(MLPm (zi )),F̂i = MLPf (zi ),Âi,j = σ(MLP2s (MLP1s (zi ) + MLP1s (zj )))(6)

where σ is the sigmoid function and MLPm , MLPf , MLPks are multi-layer perceptrons heads corresponding to eachcomponent of the graph (mask, features, structure). The adjacency matrix is expected to be symmetric which motivate us toparameterize it as suggested by [23].
Positioning with Relationformer As discussed above, the architecture is similar to the one proposed in Relationformer [7],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 doneon 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.
B.2. Encoding any input to graphs
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 ismapped to a list of k feature vectors of dimension de where k is not necessarily fixed. This is critically different fromextracting a unique feature vector (k = 1). If k is set to 1, the rest of the architecture must reconstruct an entire graph from asingle 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 wheneverapplicable; 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 [7]. From an input image ofshape 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 × de .
Since the ordering of the H × W features carries spatial information we add positional encoding accordingly.
9

A framework for differentiable Supervised Graph Prediction

A

Piece

Convolutionnal
Neural Network(CNN)

Graph Neural
Network(GNN)

Self-Attention
Of

Text

Input (Text)

Set Of Features

Set of Features

Input (graph)

Set Of Features

Input (Image)

A
Piece
Of

Convolutionnal
Neural Network(CNN)

Graph Neural
Network(GNN)

Self-Attention

Text
But
LongerInput (Text)

Set Of Features

Set of Features

Input (graph)

Input (Image)Set Of Features

Figure 5: Illustration of encoders extracting k features vectors for different input modalities. For text/fingerprint, k is thenumber of input tokens. For graphs, k is the size of the input graph. For images, k depends on the resolution of the image
and the CNN kernel size.
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 atransformer 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 invarianceof 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) [24].
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 F. The naive encoder we use is composed of k parallel MLPs devoted tothe extraction of the k feature vectors. This approach is arguably simplistic and more suited encoders should be considereddepending on the type of data.

C. Formal Statements And Proofs
In this section, we write
PMFGW(ŷ, P(g)) = min

T∈πM

X

Ti,j ℓh (ĥi , hj ) +i,jXTi,j ℓf (f̂i , fj )hj +i,jXTi,j Tk,l ℓA (Âi,k , Aj,l )hj hl ,i,j,k,lmeaning that we absorb the normalization factors in the ground losses to lighten the notation.
Alternatively, we also consider the matrix formulation:
PMFGW(ŷ, P(g)) = min ⟨T, C⟩ + ⟨T, L ⊗ T⟩,
T∈πM

where Ci,j = ℓh (ĥi , hj ) + ℓf (f̂i , fj )hj , Li,j,k,l = ℓA (Âi,k , Aj,l )hj hl and ⊗ is the tensor/matrix product.
10A framework for differentiable Supervised Graph Prediction

C.1. PMFGW fast computation
The following results generalize Proposition 1 of [22] so that it can be applied to the computation of PMFGW.
Proposition 4. Assuming that the ground loss than can decomposed as ℓ(a, b) = f1 (a) + f2 (b) − h1 (a)h2 (b), for any
transport plan T ∈ Rn×m and matrices A, W ∈ Rn×n and A′ , W′ ∈ Rm×m , then the tensor product of the form
X
(L ⊗ T)i,i′ =Tj,j ′ ℓ(Ai,j , A′i′ ,j ′ )Wi,j Wi′′ ,j ′j,j ′can be computed as
L ⊗ T = U1 TW′T + WTUT2 − V1 TV2T ,where U1 = f1 (A) · W, U2 = f2 (A′ ) · W′ , V1 = h1 (A) · W, V2 = h2 (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′ =X

Tj,j ′ f1 (Ai,j )Wi,j Wi′′ ,j ′ +

j,j ′

=

X
j

f1 (Ai,j )Wi,j

X

Tj,j ′ f2 (A′i′ ,j ′ )Wi,j Wi′′ ,j ′ −

j,j ′

X

Tj,j ′ Wi′′ ,j ′ +X

Tj,j ′ h1 (Ai,j )h2 (A′i′ ,j ′ )Wi,j Wi′′ ,j ′ .

j,j ′f2 (A′i′ ,j ′ )Wi′′ ,j ′X

Tj,j ′ Wi,j −

j

j′

j′

X

X

h1 (Ai,j )Wi,jj

X

Tj,j ′ h2 (A′i′ ,j ′ )Wi′′ ,j ′ .

j′Introducing U1 , U2 , V1 and U2 as defined above, we write:(L ⊗ T)i,i′ =(U1 )i,jX(U1 )i,j (T W ′T )j,i′ +

j

=

X

X

Tj,j ′ Wi′′ ,j ′ +

j′

j

X

(U2 )i′ ,j ′Tj,j ′ Wi,j −(U2 )i′ ,j ′ (W T )i,j ′ −X(V1 )i,j

j

j

j′

X

X

X

X

Tj,j ′ (V2 )i′ ,j ′ ,j′(V1 )i,j (Tj,j ′ V2T )j,i′ ,jj′which concludes that L ⊗ T = U1 TW′T + WTUT2 − V1 TV2T .
Remark 1 (Computational cost). U1 , U2 , V1 , V2 can be pre-computed for a cost of O(n2 + m2 ), after which L ⊗ T canbe computed (for any T) at a cost of O(mn2 + nm2 ).
Remark 2 (Kullback-Leibler divergence decomposition). The Kullback-Leibler divergence KL(p, q) = q log pq + (1 −(1−q)q) log (1−p), which we use as ground loss in our experiments satisfies the required decomposition given f1 (p) = − log(p),f2 (q) = q log(q) + (1 − q) log(1 − q), h1 (p) = log( 1−pp ) and, h2 (q) = 1 − q
Remark 3. The tensor product that appears in PMFGW is a special case of this theorem that corresponds to n = m = M ,Wi,j = 1 and Wi′′ ,j ′ = hi′ hj ′ . Thus, proposition 1 is a direct corollary.
C.2. PMFGW divergence properties
First, we provide below a more detailed version of Proposition 2
Proposition 5 (GI Invariance). For any m ≤ M , ŷ, ŷ ′ ∈ Ŷ and g, g ′ ∈ Gm , we have thatŷ ∼ ŷ ′ , g ∼ g ′ =⇒ PMFGW(ŷ, P(g)) = PMFGW(ŷ ′ , P(g ′ )).
Proof. We denote ŷ = (ĥ, F̂, Â) and P(g) = (h, F, A). Since ŷ and ŷ ′ are isomorphic, there exist a permutation P ∈ σMsuch that ŷ ′ = (Pĥ, PF̂, PÂPT ). 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, QAQT ). Plugging into the
PMFGW objective we get
11A framework for differentiable Supervised Graph Prediction

PMFGW(ŷ ′ , P(g ′ )) = min

T∈πM

X

Ti,j ℓh ((Pĥ)i , (Qh)j ) +Ti,j ℓf ((PF̂)i , (QF)j )(Qh)j

i,j

i,j

X

+

X

Ti,j Tk,l ℓA ((PÂPT )i,k , (QAQT )j,l )(Qh)j (Qh)li,j,k,l

= min

T∈πM

X

(PT TQ)i,j ℓh (ĥi , hj ) +

i,j

+

X

(PT TQ)i,j ℓf (f̂i , fj )hj

i,j

X

T

T

(P TQ)i,j (P TQ)k,l ℓA (Âi,k , Aj,l )hj hl .

i,j,k,lDenoting T̃ = PT TQ we have that

PMFGW(ŷ ′ , P(g ′ )) = min
T̃∈πM

X

T̃i,j ℓh (ĥi , hj ) +

i,j

+

X

X

T̃i,j ℓf (f̂i , fj )hji,jT̃i,j T̃k,l ℓA (Âi,k , Aj,l )hj hl

i,j,k,l= PMFGW(ŷ, P(g)).

We now provide a more detailed version of Proposition 3
Definition 1. We say that ℓ : X̂ × X 7→ R is positive when for any x, y ∈ X̂ × X , ℓ(x, y) ≥ 0 with equality if and only if
x = y.
Proposition 6 (Positivity). Let us assume that ℓh : [0, 1] × {0, 1} 7→ R, ℓf : Rd × Rd 7→ R and ℓA : [0, 1] × {0, 1} 7→ R
are positive. Then we have that for any ŷ ∈ Ŷ, g ∈ G :• i) PMFGW(ŷ, P(g)) ≥ 0• ii) There is equality if and only if ŷ ∼ P(g)• iii) In that case P −1 T (ŷ) ∼ g
Proof. The direct implication of ii) is the only statement that is not trivial. First, let us show that if PMFGW(ŷ, P(g)) = 0,
the optimal transport T∗ is a permutation. Recall that any transportof permutations [25] i.e.
PK plan is a convex combination
PK
there exist λ1 , ..., λK ∈]0, 1] and P1 , ..., PK ∈ σM such that k=1 λk = 1 and T∗ = k=1 λk Pk . Thus
0 = ⟨T∗ , C⟩ + ⟨T∗ , L ⊗ T∗ ⟩
=

K
X

k=1

λk ⟨Pk , C⟩ +

K
X

(7)

λ2k ⟨Pk , L ⊗ Pk ⟩ +KXλk λl ⟨Pk , L ⊗ Pl ⟩.

(8)

k̸=l

k=1

This is a sum of positive terms, thus all terms are null and in particular, for any k

0 = ⟨Pk , C⟩ + ⟨Pk , L ⊗ Pk ⟩.

(9)

Thus all the Pk are optimal transport plans. In the following, we chose one of them and denote it P. Moving back to thedeveloped formulation of PMFGW we get that12

A framework for differentiable Supervised Graph Prediction

0 = PMFGW(ŷ, P(g)) =X

Pi,j ℓh (ĥi , hj ) +X

Pi,j ℓf (f̂i , fj )hj +

i,j

i,j

X

Pi,j Pk,l ℓA (Âi,k , Aj,l )hj hl .

i,j,k,lOnce again this is a sum of positive terms thus for all i, j, k, and l
0 = Pi,j ℓh (ĥi , hj ) = Pi,j ℓf (f̂i , fj )hj = Pi,j Pk,l ℓA (Âi,k , Aj,l )hj hl
and thus
0 = ℓh ((PT ĥ)j , hj ) = ℓf ((PT F̂)j , Fj )hj = ℓA ((PT ÂP)j,l , Aj,l )hj hl .
And from the positivity of ℓh , ℓf and ℓA we get that: PT ĥ = h, PT F̂[: m] = F[: m] and PT ÂP[: m, : m] = A[: m, : m].
Since the nodes i > m are not activated, by abuse of notation we simply write PT F̂ = F and PT ÂP = A. This concludes
that ŷ ∼ P(g).

C.3. PMFGW and Partial Fused Gromov Wasserstein
Following [18], we define an OT relaxation of the Partial Matching problem.
For a large graph ĝ = (F̂, Â) ∈ GM and a smaller graph g = (F, A) ∈ Gm , the set of transport plan transporting a subgraphof ĝ to g can be defined as
πM,m = {T ∈ [0, 1]M ×m | T1m ≤ 1M , TT 1M = 1m , 1TM T1m = m}and the associated partial Fused Gromov Wasserstein distance ispartialFGW(ĝ, g) =minT∈πM,mm
M X
X

Ti,j ℓF (f̂i , fj ) +i=1 j=1

M X
m
X

Ti,j Tk,l ℓA (Ai,k , Aj,l ).

i,k=1 j,l=1

In the following, we show that partialFGW(ĝ, g) is equivalent to the padded Fused Gromov Wasserstein distance defined as

paddedFGW(ĝ, g) = min

T∈πM

m
M X
X

Ti,j ℓF (f̂i , fj ) +i=1 j=1

M X
m
X

Ti,j Tk,l ℓA (Ai,k , Aj,l ).

i,k=1 j,l=1

Lemma 4. Any transport plan T ∈ πM has the form T = Tp


T2 where Tp is a partial transport plan i.e. Tp ∈ πM,m .

Proof. Let us check that Tp is in πM .
• 1M = T1M = Tp 1m + T2 1M −m ≥ Tp 1m
 T

Tp 1M
T
• 1M = T 1M =. And thus TTp 1M = 1m
TT2 1M• From the previous we immediately get that 1TM Tp 1m = m

Lemma 5. For any partial transport plan Tp ∈ πM,m there exist T2 ∈ RM ×(M −m) such that T = Tp
13
T2 ∈ πM .

A framework for differentiable Supervised Graph Prediction

Proof. Let us define p = 1M − Tp 1m . This is the mass of the larger graph that is not matched by Tp . Note that since1
p1TM −m i.e. we spread the remaining mass uniformly acrossTp ∈ πM,m we have that p ≥ 0. Thus we can set T2 = M −m
the padding nodes. Let us check that T = Tp T2 ∈ πM is indeed a valid transport plan.
• T1M = Tp 1m + T2 1M −m = Tp 1m + p = 1m
 T
 

 
 
1m
1m
Tp 1M
1m
T
• T 1M =
=
=
=
1
1
T
T
T
T
1M −m = 1M
TT2 1MM −m (p 1M )1M −m
M −m (1M 1M − 1m Tp 1M )1M −m

Proposition 7. paddedFGW and partialFGW are equal and any optimal plan T∗ of paddedFGW has the form T∗ =(Tp∗ , T2 ) where Tp∗ is optimal for partialFGW.
Proof. Follows directly from the two previous lemmas.
Remark 6. The proposed PMFGW is equivalent to paddedFGW (and thus to partialFGW) if and only if ℓh is set to a
constant.
Remark 7. 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 [18]. Further comparisons are left for future work.

D. Training Any2Graph
D.1. Learning dynamics

loss features
loss mask
loss structure

1.0

Losses (test set)Losses (test set)

The PMFGW loss is composed of three terms, two of them are linear and account for the prediction of the nodes and theirfeatures, one is quadratic and accounts for the prediction of edges. The last term is arguably the harder to minimize for themodel, as a consequence, we observe that the training performs best when the two first terms are minimized first whichthen 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 6.

0.5
0.0

1.0
0.5
0.0

0 1 2 3 4 5 6 7 8 9 10
Epochsloss features
loss mask
loss structure

0 1 2 3 4 5 6 7 8 9 10
EpochsFigure 6: First epochs of training for Coloring. The test values of the 3 components of the loss are reported. On the left
(resp. right) α is set to [1, 1, 1] (resp. [1,1,10]). In the first scenario, the first two terms of the loss are learned very fast andthe structure is optimized next. In the second scenario, setting αA = 10 prevents this desirable learning dynamic.
D.2. Feature Diffusion
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 the desirable dynamic introduced above does notoccur. This motivates us to perform Feature Diffusion (FD) that is replacing the node feature vector F by the concatenation[F, AF]14

A framework for differentiable Supervised Graph Prediction

Table 3: Performances of the Relationformer and Any2Graph with (+FD) and without feature diffusion.
G RAPH L EVEL
GI ACC . ↑ PMFGW ↓DATASET

M ODELQM9

R ELATIONFORMERR ELATIONFORMER + FD
A NY 2G RAPH
A NY 2G RAPH + FD

9.15
3.80
3.44
2.13

0.05
9.95
7.50
29.85

GDB13

R ELATIONFORMERR ELATIONFORMER + FD
A NY 2G RAPH
A NY 2G RAPH + FD

11.40
8.83
7.45
3.63

0.00
0.01
0.05
16.25

E DIT D IST. ↓

0.3
0.2
0.1
0.0

0

2

4

6
8
Epochs

10

0.48
0.22
0.21
0.14

21.42
86.07
86.21
90.19

4.77
73.31
77.27
88.08

99.28
99.34
99.26
99.77

91.80
96.0
93.65
95.45

0.43
0.29
0.22
0.11

81.96
84.14
87.20
90.83

31.49
55.89
60.41
84.86

97.77
97.57
99.41
99.80

97.45
98.65
96.15
98.15

0.5

loss features
loss mask
loss structure

0.4

N ODE L EVEL ACC .
N ODE ↑S IZE ↑

Loss (test set)

Loss (test set)

0.5

E DGE L EVELP REC . ↑ R EC . ↑

0.3
0.2
0.1
0.0

12

loss features
loss mask
loss structure

0.4

0

2

4

6
8
Epochs10

12Figure 7: First epochs of training for GDB13. The test values of the 3 components of the loss are reported. On the left, we
perform FD before training, on the right, we leave node features unchanged. 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.

before training. The diffused node features carry a portion of the structural information. This makes the node feature termslightly 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 7. We observe that both Any2Graph and Relationformer benefits from such procedure,as reported in table 3
D.3. Optimal Transport Relaxation
We adopted the OT point of view when designing Any2Graph. In practice, this means that we do not project the OT planback to the set of permutations with a Hungarian matcher before plugging it in the loss as in [26]. Testing the effect ofadding this extra step we observed a 5% to 10% increase of the edit distance across datasets (table 4) along with a moreunstable training curve (figure 8). This confirms that a continuous transport plan provides a slightly more stable gradientthan a discrete permutation, which aligns with the findings of [27] on the similar topic of object detection.
DatasetED without Hungarian
ED with Hungarian

Coloring

Toulouse

USCities

QM9

GDB13

0.20
0.23

0.13
0.15

1.86
2.03

2.13
2.08

3.63
3.86

Table 4: Effect of adding Hungarian Matching on the performances evaluated with the test edit distance. We observe, that
Hungarian Matching slightly decreases the performances on all datasets but QM9.

15

A framework for differentiable Supervised Graph Prediction

With Hungarian Matching

Loss (test set)

0.30

Without Hungarian Matching
0.25

0.20

0.15

5

10

15

20

25

30

EpochsFigure 8: First epochs of training for GDB13 with and without projection of the optimal transport plan to the set of
permutations with Hungarian matching. Hungarian matching slightly decreases the performances and induces more
oscillations of the loss, which could be explained by a less stable gradient.

E. Experimental setting
E.1. Datasets
In this paper, we consider five datasets for which we provide a variety of statistics in Table 5. Coloring is a new synthetic
dataset which we describe in detail in Appendix F. Toulouse (resp. USCities) is a Sat2Graph dataset [2] where the inputs
are images of size 64 × 64 (resp. 128 × 128). QM9 [28] and GDB13 [29] are datasets of small molecules which we
use to address the Fingerprint2Graph task. Here, we compute a fingerprint representation of the molecule and attempt toreconstruct the original molecule from this loss representation. Following [30] we use the Morgan Radius-2 fingerprint[31] which represents a molecule by a bit vector of size 2048, where each bit represents the presence/absence of a given
substructure. Finally, we feed our model with the list of non-zeros bits, i.e. the list of substructures (tokens) present in themolecule. The list of substructures has a min/average/max length of 2/21/27 for QM9 and 7/29/36 for GDB13.
Table 5: Table summarizing the properties of the datasets considered.
DATASETS IZE( TRAIN / TEST / VALID )N ODES( MIN / MEAN / MAX )E DGES( MIN / MEAN / MAX )I NPUT M ODALITYN ODE FEATURES

Coloring
Toulouse
USCities
QM9GDB13100 K /10 K /10 K
80 K /10 K /10 K
130 K /10 K /10 K
120 K /10 K /10 K
1300 K /70 K /70 K

4/7.0/10
3/6.1/9
2/7.5/17
1/8.8/9
5/12.7/13

3/10.9/22
2/5.0/14
1/5.8/20
0/9.4/13
5/15.15/18

RGB I MAGESG REY I MAGESG REY I MAGESL IST OF TOKENSL IST OF TOKENS

4 C LASSES ( COLORS )2D POSITIONS
2D POSITIONS
4 CLASSES ( ATOMS )5 CLASSES ( ATOMS )E.2. Competitors
We compare Any2Graph, to our direct end-to-end competitor Relationformer [7] 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 2) for both approaches
so that the only difference is the loss. We apply feature diffusion both to Any2Graph and Relationformer. Moreover, we also
compare with a surrogate regression approach (FGW-Bary-ILE) based on FGW barycenters [6] whose prediction function
writes as:
KX
αk (x; W )FGW22 (y, ȳk ),(10)
f (x) = arg min
y∈Ŷm

k=1

th

where each ȳk denotes the k template graph and αk the k th weight function. The templates are the training samples andweight function αk are learned by sketched kernel ridge regression [32], [33] with gaussian kernel. For fair comparison, we
also integrated feature diffusion to FGW-Bary-ILE.
16

A framework for differentiable Supervised Graph Prediction

E.3. Metrics
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 ŷ and padded target P(g). For thiscomputation, we set α to the values displayed in table 6.
PMFGW(ŷ, P(g))
Then we report the graph edit distance [15] between predicted graph P −1 T ŷ and target g which we compute using Pygmtools[34]. 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.
E DIT(P −1 T ŷ, g)Finally, we report the Graph Isomorphism Accuracy GI ACC that is
GI ACC(ŷ, g) = 1[E DIT(P −1 T ŷ, g) = 0]
Node level Recall that for a prediction ŷ = (ĥ, F̂, Â) the size of the predicted graph is m̂ = ||ĥ > 0.5||1 . Denoting m thesize of the target graph we report the size accuracy:
SIZE ACC(ŷ, 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 mnodes with the highest probability ĥi , resulting in a graph ĝ = (F̃, Ã) with ground truth size. This is equivalent to assumingthat the size of the graph is well predicted. Then we use Pygmtools to compute a one-to-one matching σ between the nodesof ĝ 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 ĝ have been aligned. We can now define the node accuracy NODE ACC asm

NODE ACC(ĝ, g) =1 X
1[F̃i = Fi ],m i=1

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 P REC .(ĝ, g) =Pm

EDGE P REC .(ĝ, g) =

Pm

i,j=1

1[Ãi,j = 1, Ai,j = 1]
i,j=1 1[Ãi,j = 1]

Pm

i,j=1

1[Ãi,j = 1, Ai,j = 1]
i,j=1 1[Ai,j = 1]Pm

All those metrics are then averaged other the test set.
E.4. Training And Architecture Hyperparameters
Encoder We follow the guidelines established in B.2 for the choice of the encoder. In particular, all encoders for Coloring,
Toulouse and USCities are CNNs. The encoder for Coloring is a variation of Resnet18 [35], where we remove the firstmax-pooling layer and the last two blocks to accommodate for the low resolution of our input image. We proceed similarly17

A framework for differentiable Supervised Graph Prediction

for Toulouse except that we only remove the last block. For USCities we keep the full Resnet18. For the Fingerprint2Graphdatasets, we use a transformer encoder. In practice, this transformer encoder and that of the encoder-decoder module are
merged to avoid redundancy. All encoders end with a linear layer projecting the feature vectors to the hidden dimension de .
Transformer We use the Pre-LN variant [36] of the transformer encoder-decoder model as described in [37]. To reduce
the number of hyperparameters, encoder and decoder modules both consist of stacks of Nτ layers, with Nh heads and thehidden dimensions of all MLP is set to 4 × de .
Decoder

All MLPs in the decoder module have one hidden layer.

Optimizer We train all neural networks with the Adam optimizer [38], learning rate η, 8000 warm-up steps and all otherhyperparameters set to default values. We also use gradient clipping with a max norm set to 0.1.
All hyperparameters are given in Table 6.
Table 6: Hyperparameters used to train our models. We also report the total training time on a NVIDIA V100.
DATASET
Coloring
Toulouse
USCities
QM9
GDB13

αH
1
1
2
1
1

PMFGW
αF
αA
1
5
5
1
1

1
1
0.5
1
1

de

A RCHITECTURE
Nτ Nh M

256
256
256
128
512

3
4
4
3
5

8
8
8
4
8

12
12
20
12
15

η
3 E -4
1 E -4
1 E -43 E -43 E -4

O PTIMIZATIONBATCHSIZE S TEPS
128
128
128
128
256

75 K
100 K
150 K
150 K
150 K

T IME
4H
8H
14 H
6H
24 H

F. Coloring: a new synthetic dataset for benchmarking SGP
We introduce Coloring, a new synthetic dataset well suited for benchmarking SGP methods. The main advantages ofColoring 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 nodelabels (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:

Figure 9: Illustration of the five steps to follow to generate a new instance of Coloring.
• 0) Sample the number of nodes (graph size) m. In this paper, we sample uniformly on some interval [Mmin , Mmax ].
• 1) Sample m centroids on [0, 1] × [0, 1]. In this paper, we sample the centroids as uniform i.i.d. variables.
• 2) Partition [0, 1] × [0, 1] (the image) in a Voronoi diagram fashion [39]. In this paper, we use the L1 distance. and an
image of resolution H × H.
18

A framework for differentiable Supervised Graph Prediction• 3) Create the associated graph i.e. each node is a region of the image and two nodes are linked by an edge whenever
the two associated regions are adjacent.
• 4) Color the graph with K > 4 colors. In this paper, we use K = 4. A coloring is said to be valid whenever no adjacentnodes have the same color. Note that graph coloring is known to be NP-complete [40].
• 5) Color the original image accordingly.
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 avariation of Coloring where the input image is flattened and treated as a vector allowing us to explore a synthetic Vect2Graphtask. The properties of these datasets, along with the performances of Any2Graph are reported in table 7. We hope that
Coloring will be used to benchmark future SGP methods.
Table 7: 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 forColoring. None scales to ColoringBig.
DATASET

MMIN

MMAX

H

N UMBER O F S AMPLESA NY 2G RAPHR ELATIONFORMERFGWBARY-NN

ColoringVect
Coloring
ColoringBig

4
6
10

6
10
15

16
32
64

100 K
100 K
200 K

0.46
0.20
1.01

1.40
5.47
8.91

2.09
6.73
N.A.

G. Additional Qualitative Results
G.1. Out of distribution performances
We tested if, once trained on Toulouse dataset, the predictive model is able to cope with out-of-distribution data. Figure 10
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.

Figure 10: Any2Graph trained on Toulouse performing on out-of-distribution inputs. Input images are displayed on top row
and prediction in the bottom row.

19

A framework for differentiable Supervised Graph Prediction

G.2. Qualitative results on COLORING

Input

TargetAny2Graph

Relationformer FGWBary-ILE

Figure 11: Graph prediction on the Coloring dataset.

20

FGWBary-NN

A framework for differentiable Supervised Graph Prediction

G.3. Qualitative results on QM9

Input

TargetAny2Graph

Relationformer FGWBary-ILE

1 10
298 409 26
1677 9 725 95 34
6 357 2747 3 48 1002 47
109 1409 260
33 6491 436
312 7 12669 102 30
2 29 393 4 8
131 464 181161 521 1658
1045 36

1 663
10 187 144
184 764 9
7 17 427
34 2 20
3 131 53
61 50 582
599 645 154
1221 389 11

1 1001024 490 151297 12 1803
107 30 29
252 4 1692
13 50 1013
5 143 280
203 1323 640
11 127

1 529
1717 145 15
25 45 170
6 452 214 1092 3
738 8 311643 545 304
167 605 11
404

1 27
158 10 85
187 62 124
537 877 24
6 723 43
149 4 1545
725

1 94
32 89 689
38 7 78455 1084 2
14 285 3
21 1157 482
1567 5 1409
82 49 54
157

Figure 12: Graph prediction on the QM9 dataset.

21

FGWBary-NN

A framework for differentiable Supervised Graph Prediction

G.4. Qualitative results on GDB13

Input

TargetAny2Graph

Relationformer

2048 146
226 309 522
527 584 610
655 678 856
911 925 926
973 1018 1024
1049 1054 1059
1161 1291 1358
1372 1466 1689
1728 1858 18811922 2039

2048 24
107 171 411
439 522 584
613 655 678
685 738 925
934 1018 1059
1070 1161 1186
1223 1238 1264
1291 1324 1543
1704 1725 1880
1909 1922

2048 30
49 165 330
377 407 547
622 655 674
757 786 874
884 934 1040
1199 1232 1359
1379 1393 1451
1467 1566 1633
1690 1695 1749
1751 1872 1883
1943

2048 70
161 174 200
316 377 454
469 559 655
694 821 925
984 1018 1026
1033 1113 1128
1229 1291 1379
1506 1513 1753
1818 1826 1872
1878 1892 19202003

2048 96
253 300 313
317 377 502
610 649 839
872 925 949
1018 1043 1113
1249 1324 1379
1430 1602 1637
1644 1652 1716
1763 1872 1949
2040

2048 421 143 293
378 516 522
655 728 738
757 890 925
934 949 1018
1059 1128 1161
1223 1365 1440
1479 1507 1572
1607 1634 1684
2006 2039

2048 14113 292 425
486 522 549
566 589 640
655 677 678
705 806 909
925 949 1018
1027 1059 1113
1223 1256 1324
1385 1467 1563
1634 1847 1872
1882 1958

Figure 13: Graph prediction on the GDB13 dataset.

22

A framework for differentiable Supervised Graph Prediction

G.5. Qualitative results on TOULOUSE

Input

TargetAny2Graph

Relationformer FGWBary-NN

Figure 14: Graph prediction on the Toulouse dataset.

23

FGWBary-ILE

A framework for differentiable Supervised Graph Prediction

G.6. Qualitative results on USCities

Input

TargetAny2Graph

Relationformer

Figure 15: Graph prediction on the USCities dataset.

24