Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs
Authors: Sonia Mazelet
Rémi Flamary
Bertrand Thirion
Venue: NeurIPS 2025 poster ()
Keywords: Non-Convex Optimization
Unsupervised Learning
optimal transport
gnn
Primary area: deep_learning

Abstract: Optimal transport between graphs, based on Gromov-Wasserstein and
  other extensions, is a powerful tool for comparing and aligning
  graph structures.  However, solving the associated non-convex
  optimization problems is computationally expensive, which limits the
  scalability of these methods to large graphs. In this work, we
  present Unbalanced Learning of Optimal Transport (ULOT), a deep
  learning method that predicts optimal transport plans between two
  graphs. Our method is trained by minimizing the fused unbalanced
  Gromov-Wasserstein (FUGW) loss. We propose a novel neural
  architecture with cross-attention that is conditioned on the FUGW
  tradeoff hyperparameters. We evaluate ULOT on synthetic stochastic
  block model (SBM) graphs and on real cortical surface data obtained
  from fMRI. ULOT predicts transport plans with competitive loss up to
  two orders of magnitude faster than classical solvers.  Furthermore,
  the predicted plan can be used as a warm start for classical solvers
  to accelerate their convergence. Finally, the predicted transport
  plan is fully differentiable with respect to the graph inputs and
  FUGW hyperparameters, enabling the optimization of functionals of
  the ULOT plan.

Unsupervised Learning for Optimal Transport plan prediction between unbalanced graphs
Sonia Mazelet CMAP, Ecole Polytechnique
Rémi Flamary CMAP, Ecole Polytechnique
Palaiseau, France
Palaiseau, France
sonia. www.eeeeettpllyyeeciniiuuuueeeee edu
remi f ylwlaayy@lllyyeeeniiiuueeeeeeeee edu
Bertrand Thirion
Mind, Inria-Saclay
Palaiseau, France
bertrand. thirion@inria.ft
Abstract
Optimal transport between graphs, based on Ceommo--asseertein and other exten- sions, a powerful tool for comparing and aligning graph structures. However, solving the associated non-convex optimization problems is computationally ex- pensive, which limits the scalability these methods to large graphs. this work, we present Unbalanced Learning of Optimal Transport (ULOT), deep learning method that predicts optimal transport plans between two graphs. Our method is (FUGW) loss. We propose a novel neural architecture with cross- attention that is conditioned on the FUGW tradeoff hyperparameters. We evaluate ULOT on synthetic stochastic block model (SBM) graphs and on real cortical surface data obtained from fMRL. ULOT predicts transport plans with competitive loss up to orders of magnitude faster than classical solvers. Furthermore, the predicted plan can be used as a warm start for classical solvers to accelerate their convergence. Finally, the predicted transport plan is fully differentiable with respect to the graph inputs and FUGW hyperparameters, enabling the optimization of functionals of the ULOT plan.
trained by minimizing the fused unbalanced
Introduction
Graph alignment In many graph data applications, aligning matching nodes between two graphs necessary. Examples include object detection, where semantic correspondences between objects in two images from different domains can make the model more adaptive [12); graph edit distance
[22 18), where the goal is to compute the distance between two graphs, which requires finding the best correspondence between their nodes; shape matching 20 or brain alignment across subjects
30
But the problem of graph matching challenging because of the combinatorial nature of the problem. which can often be reformulated as Quadratic Assignment Problem (QAP) known to be Na-hard [16. One strategy that has been proposed recently is use deep learning learn the matching in
supervised setting (3) [25 These methods typically rely on graph neural network (GNN) to learn a representation of the nodes and then use a matching algorithm compute the correspondence between the nodes. However, the problem is even more difficult when the graphs are unbalanced, i.e.
when they have different numbers of nodes or when some nodes have noisy features or connections.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Optimal transport between graphs In recent years, Optimal Transport (OT) has emerged as powerful tool for solving the graph matching problem. can be seen as a continuous relaxation of the QAP with the Gromov- Washoov-tsssrtteein (GW) distance [7 21 40 which is a generalization of the classical Wasserstein distance distributions in different metric spaces. Extensions of the GW distance to labeled graphs have been proposed, such as the Fused
(FGW) distance [3 But those OT problems put strong constraints on the transport plan, which makes them very sensitive to noise, outliers and local deformations. This is why Unbalanced GW [27 and Fused Unbalanced GW (FUGW) [30] were proposed to generalize the GW and FGW distances to unbalanced settings with application in positive unlabeled learning and brain alignment.
Complexity of classical optimal transport solvers Quadratic OT problems such as FGW and FUGW are non-convex. Classical solvers rely on a block coordinate descent algorithm to iteratively solve linearized versions of the problem. This linearization of complexity (n, + nε31), where n1 and n2 are the number of nodes in the two graphs [21]. This makes the method unscalable for large graphs, for applications where we need to compute the transport plan for many pairs of graphs or when validation of the hyperparameters is necessary. This is especially problematic in the unbalanced case where the solution very sensitive to the parameters.
The ML community has recently proposed to use deep learning for accelerating or solving OT problems. For example [1]] proposed to estimate the OT mapping using a neural network and extensions to the GW have been proposed in [19 (38. But the most relevant work for purpose is Meta OT [2] which proposed to learn to predict the dual potentials of the entropic OT problem with neural network. But this approach is limited classical entropy regularized OT and not directly applicable to the quadratic FUGW problem which motivates our proposed approach detailed below.
Contributions We propose in this paper new method to learn a neural network that predicts the transport plan of the FUGW problem denoted as Unsupervised Learning of Optimal Transport plan prediction (ULOT). We propose a novel architecture based on graph neural networks (GNN) and cross attention mechanisms to predict the OT plan with a complexity of ((nr7i)), which significantly faster than classical solvers. In addition the neural network is conditioned by the parameters of the FUGW problem, which allows to efficiently predict OT plan for all the possible values of the parameters. This is particularly useful in the unbalanced case where the solution is very sensitive to the parameters and validation is often necessary. We show in our experiments, on simulated and real life data, that our method outperforms classical solvers in terms of speed by two orders of magnitude while providing OT plans with competitive loss. The predicted transport plan can also be used as a warm start for classical solvers, which reduces the number of iterations and the overall time of the algorithm. We also show that our method provides a smooth estimation of the transport plan that can be used for numerous applications such as parameter validation for label propagation or gradient descent of a FUGW loss. Our code is available at ptttps://github.... com/ smazelet/ULOT
2 Learning to predict OT plans between graphs
In this section we first introduce the FUGW optimal transport problem and its associated loss function. Next we present our amortized optimization strategy called Unbalanced Learning of Optimal Transport (ULOT) and detail the architecture of the neural network. Finally we discuss the related works in deep learning for optimal transport and graph matching.
2.1 Fused Unbalanced Gromov Wasserstein (FUGW)
Definition of the FUGW loss Consider the two graphs {G+ (Fg. DE. (w))k=(112)) with nı and ng nodes respectively. For lo E (1, they are characterized by their node features F, E Reexd their connectivity matrices D₁ E [R (usually adjacency matrix shortest path distance matrix) and their node weights Wa E An, A [(w]) [(w],..,, (.) Σne w = }} that characterize the node's relative importance 34.. Note that in the following we will assume that these weights are uniform, i.e. w/ 1/ng for i I,..., Flk. The goal of FUGW [30] is to learn positive transport plan P 12

---

between the nodes of G1 and G2 that minimizes the following loss function:
I. I. C2; P) -1111-0
Σ
(R),
(Fa),
P.j
(1)
1-=1
B2. 22
Σ
1(DI), - P
(2)
(KL(P #1 @ PH1 w1) (L @ Pe2lw2 @ w2)).
(3)
The FUGW loss a combination of the Wasserstein distance 1 that measures the preservation of node features, the Gromov Wasserstein distance 2 that measures the conservation of local geometries and a penalization of the violation of the marginals for the OT plan with the KL. divergence (3). The
terms ①1) and (2) weighed by the trade-off parameters is weighted by p. Unbalanced OT is a very general and robust framework that can adapt to differences in the geometry of the graph vertices.
[0, 1] and the marginal penalization (3)
Complexity of solving minp L L" (2, In order to solve the FUGW problem one needs to minimize the FUGW loss with respect to the optimal transport plan P. Because of the Gromov Wasserstein term, the complexity of the FUGW loss its gradient for given plan P is theoretically quartic O (nfr3)) In the case of the square loss, Peyré et al [21] showed that the complexity can be reduced to cubic complexity Onn,ng + mfnz). which remains computationally intensive for large graphs with typically more than 10k nodes.
Existing methods to minimize the FUGW loss a block coordinate descent scheme lower bound of the objective that consists in solving at each iteration linearization of the quadratic problem 81. 30 and requires each iteration to compute the cubic gradient. In the following will compare ULOT to three different types of inner-solvers for the linearized problem: the Majorization- minimization (MM) algorithm (6. the inexact Bregman Proximal Point (IBPP) algorithm [39] and a more classical L-BFGS-B algorithm [4]. We will also compare our approach to the entropic regularization of the FUGW loss for which the inner problem can be solved using the Sinkhorn algorithm. All those methods are iterative and require O(n, + nf gradient/inearrizaton computation at each iteration. This is a major bottleneck for the FUGW problem, especially when the graphs are large or FUGW has to be solved multiple times, for instance when computing FUGW barycenter while selecting the parameters (a,, p) the method.
2.2 ULOT optimization problem and architecture
Learning to predict FUGW OT plans We propose train a model p (G1, G2), parametrized by 0, that, given two graphs G1, G2 and (p, a) parameters can predict FUGW transport plan, or at least good solution the FUGW problem, between them. We want to avoid training the model in supervised way where for each pair of graphs (G1, G2) in the training set we need to pre-compute the corresponding transport plan P. This why we propose train the model in an unsupervised way using amortized optimization [3]. We do this by sampling pairs of graphs (G1, G2) from a training dataset 'D and parameters (p, or) and minimizing the expected FUGW loss. The ULOT model trained to minimize:
min B,, G.-DP [L (Li PB (Gi G2)]
(4)
Where P is the distribution of the parameters (p, a) chosen in the experiment as Pp. Pa. where Pp is the log uniform distribution between 10-7 and 1, and Pa follows the Beta distribution with parameters (0.5, 0.5). The fact that we optimize the loss some intervals of the parameters (p. a) allows the model to generalize to different values of the parameters to explore or even optimize them.
Encoding the parameters (p.a) Since we want the model (G1, G2) depend on the param- eters (p. a), we need to encode them in a way that can be used by the neural network. To do that, we add them to the node features of the graphs at each layer of the network. p is a positive sealar that impacts the mass of the OT plan and can be included as such. But is a scalar in [0, 1] fixing the tradeoff between the Wasserstein and Gromov-Wasserrter terms that can have large impact close to 0 and 1. In order to facilitate its use in the neural network, propose to encode it in more
3

---

attention layers
Unbalanced layer
Linear
MLP
koftmax
Linear
Linear sigmoid
Gr Gz)
G2
MLP
(columa)
Linear
HLP
Linear
F2
Linear
Figure 1: ULOT architecture for OT plan prediction The architecture consists of two parts. a node embedding layer repeated N times that relies on cross graph attention and self node updates (GCN), and the final transport plan prediction layer that predicts node weights and the output transport plan.
expressive way with positional encoding technique. We chose to use the same technique as in 32 using the Fourier basis for encoding time flow matching applications as follows:
á (cos (k=.)))==,....d (sin kn (1 k))k==,....dd
where we set d = 10 in our experiments. á and are concatenated to the node features at each layer of the network as detailed next.
2.3 Proposed cross-attention neural architecture
We now detail the proposed neural network architecture that takes as input two graphs Gi and G2 and the parameters (p. a) and predicts the transport plan P& (G1, G2). The proposed ULOT architecture summarized in Figure I consists of two main parts. The first part contains N layers of node embedding and attention-bran graph. The final layer predicts the transport plan from the learned node features and interactions with node reweighting. The design choices are further motivated in the ablation studies in Section
cross-graphe interactions that incorporate the geometry of each
Node embedding with cross attention The first block is an adaptation of the Graph Matching Network from [33 that outputs node features that are relevant in a matching context. Intuitively the cross attention mechanism computes node features that are similar for nodes that can be matched and dissimilar for nodes that should not. The node embeddings are computed each layer using two paths. One path is called the self path that simply consists of a GCN treating each graph independently. The use of GCNs is motivated by the Gromov term in the FUGW loss that requires a model taking into account the graph topology. For each graph Gş for k E {1,2} self node features are computed with
Feelf
GCN(FF)..
The second path is the cross path that computes similarity matrix between the node features of both graphs and learns new features that characterize their interactions. In parallel, crass node features Frme, Frmm are computed with an attention block, for k, k' E {1,2), & # k', learned the following way:
1. Compute embeddings:
Feross MLP p.
(7)
2. Compute similarity matrix and the row/column attention matrices SI, S2 with, for s cosine similarity and m] j n2]:
S softmax (a²
Sij s ((F (Fgros)))) and
S softmaXcolamn (a2
where R is hyperparameter.
3. Compute the updated node features Feoms and Femss.
F Femms
Linear
[FETORS
FCruS)
F Linear (Fenss FF (9)

---

At the end of the layer, for k. k E 1, 1,2, the learned node features from the self (GCN) and cross- attention paths are merged with the input node features with the following equation:
Frfinl MLP (Linear(Fs) Fx), Linear( (F F (iinear(F**00) P. &)
(10)
where the inputs to the MLP are concatenated. The use of those two paths allows to learn node features that take into account both the graph geometry and the cross-interactions between the graphs.
Transport plan prediction with cross attention and node scaling The optimal transport plan prediction block is designed so as to predict an unbalanced OT plan that must contain pairwise relationships between nodes but also that can discard (reweight) nodes that cost too much from an OT perspective. This is done by separating the two aspects. We chose to use for the pairwise relationship in the OT plan the average of the cross attention matrices corresponding to row-wise softmax and column-wise softmax respectively. This matrix gives good starting point for the plan, which is then refined in the unbalanced optimal transport plan block, where the scaling per node makes no longer balanced. This is why we also provide in the last block two node reweighting heads that allow scaling on the OT plan and enable the network to discard mass. These heads predict weight vectors 2 E R tM k {1,2} the nodes of the graphs that represent how much weight is transported from each node with
1 sigmoid(L.nearr Fffina - G)), 22 = sigmoid(L.nnarf Ffinal PP k))
(11)
We learn the transport plan from the similarity matrices Si, S2 of the last layer
S (1
(G1, G2) =
diag (1ag(
(12)
n2
This allows us to learn a transport plan that is unbalanced and separates the problem of finding the node interactions across graphs, done with the cross attention, and the problem of learning the individual node weights that is specific to unbalanced OT.
2.4 Related works
Deep optimal learning for transport Several methods have been proposed for accelerating the resolution of optimal transport problems with deep learning. For instance (26) proposed to model the dual potentials of the entropic regularized OT problem with a neural network. Neural OT [TT] learns the classical OT mapping between two distributions using neural network, Recent Neural OT extensions have also been proposed to solve the GW problem in [19 [8]. This usually allows for solving large OT problems but the resulting neural network is solution for a specific pair of distributions and needs to be optimized again for new distribution pairs.
Meta OT [2 uses strategy called amortized optimization 3 learn MLP neural network that predicts on the samples the dual potentials of the entropy regularized OT problem. Their model can be used to predict the entropic transport plan between two new distributions using the primal-dual relationship. However, Meta OT' cannot be used for Quadratic OT problems such as the Gromov- Wasserstein FUGW problems because the optimization problem is not convex and the primal-dual relationship is much more complicated 42). In fact, in order to use an approach similar to Meta OT, one would need to linearization of the problem that is O(n²) which would cancel part of the advantage of using an efficient neural network.
ULOT has been designed to perform OT plan prediction with O(r²) complexity. Our approach also focuses on graph data. There is a need to go beyond MLP in order to design neural network that can use the graph structure. In this sense, ULOT can be seen as a generalization of Meta OT to the case of unbalanced OT between graphs. Finally we learn a model that can predict the transport plan conditioned on the parameters the OT problem, in our case the (a, p). which is particularly novel and has not been done before to the best of our knowledge.
Deep learning for graph matching Various neural architectures have been proposed for deep graph matching, which refers to the problem of finding structural correspondence between graphs. However these methods often face notable limitations. Some approaches are limited to predicting global
similarity score between pairs of graphs without predicting node correspondence [33 14] Other
methods that do provide node-level matching typically rely on supervised training (36 25 which
limits their applications to domains with available ground truth correspondences such images.
5

---

OT plan with respect to for (1, 2) (2,
Nodes transported mass for increasing p values
ULOT
Node transported mass
cluster
0.00
0.04
cluster 2
cluster
pm0 0010
pm 0.002
p= 0.007
Figure 3: Illustration of OT plans for different values of (left) OT plan (red lines) predicted by ULOT (top) and estimated by the IBPP solver (bottom). The red lines opacity proportional the amount of transported mass. (right) Marginals on the nodes of both graphs for different values of colored by the cluster they belong to.
Unsupervised methods, on the other hand, rely on task-specific objectives such as enforcing cycle consistency 33 37 or adopt contrastive learning approaches by matching graphs to their augmented copies, for which they know the ground truth [15].
To improve the robustness of the matching between graphs of different sizes, some methods discard subset of correspondences. Various methods have been proposed such selecting the top-k most confident matches (35) or adding dummy nodes [T0]. However, while these methods have a similar goal to ULOT they perform hard selection of nodes, which is useful in the case of outliers but when we need more continuous way to reduce some node importance. Finally, while our method bears some similarities with neural networks for graph matching, it is important note that the objective of Unbalanced OT is fundamentally different from graph matching.
Numerical experiments
In this section, we evaluate ULOT and compare it to classical solvers on both a simulated dataset of Stochastic Block Models (SBMs) with different numbers of clusters and on the Individual Brain Charting (IBC) dataset
23 of functional MRI activations on brain surfaces.
3.1 Illustration and interpretation simulated graphs
Dataset and training setup We first train ULOT on a simulated dataset of Stochastic Block Models (SBMs) with 3 linearly connected clusters and 3D node features that are a one hot encoding of the cluster classes 1,2 or with additive centered Gaussian noise. To investi- gate the properties of the transport plans learned by ULOT in comparison to those estimated by classical solvers, we construct three types of graphs with different clusters. The first type in- cludes graphs where all three clusters 21, are present, the second type of graphs has clus- ters (1, 2) and the third has clusters 3). All graphs in the dataset have a random number of nodes ranging from 30 to 60 and the train- ing dataset consists of 50000 simulated pairs (G1, G2, P. a).
plan with respect for -1,
ULOT
Solver
a-0.0
a=0.5
2-1.0
Figure 2: Examples of transport plans (red lines) predicted by ULOT (top) and estimated by the IBPP solver (bottom) for different values. The lines opacity is proportional to the amount of transported mass.
With those three types of graphs, unbalanced OT should be able to find a transport plan that matches the clusters if they are present in both graphs, but discard clusters that are not shared (when o, p are properly selected). We compare ULOT the numerical solver IBPP, which provides good trade off between performance and speed.

---

Regularization path with respect to the parameters p and a We first illustrate the effect of p on the predicted transport plan between graphs containing clusters (1, 2) and (2, 3). We see in Figure 3 (left) that the transported mass increases with for both ULOT and the solver but while the solver is stuck in 0 mass local minimum, ULOTT is able to find a plan with mass and lower cost for 0.002. We also provide in Figure (right) the regularization path for the marginals on the nodes of both graphs which shows that the nodes from cluster 2 in green are the first to receive mass when p increases finding proper alignment of the clusters.
Next we illustrate the OT plans between two graphs of type 21, for different values of a in Figure Because the node features noisy, both the Gromov Wasserstein and Wasserstein terms in the loss are needed to predict accurate plans. We see that while the transport plans from ULOT and the solver are comparable for a < 1, the solver wrongly matches the clusters for does not use the node feature information and can permute classes. ULOT does not permute the clusters for 1 because it is continuous w.r.t. a and has learned to use the node features to find better plan.
because
Optimizing the hyperparameter for a prediction task We now consider the task of label propagation between graphs, where node labels are known on source graph but partly missing on a target graph. To infer the missing labels, we transport the one hot encoding of the source graph node labels onto the target graph, producing label probabilities for each node similarly to what was proposed in [28 224. A key challenge of this approach, when the graph types can differ, lies in selecting the appropriate FUGW parameters (p, () to ensure that the plan is relevant for label propagation.
Optim. (a, p) with respect to KL for (1. -+ (1, 2)
Initialization
Convveggence
0.4
0.0
10
10-
10
can easily compute and visualize the accuracy of the label propagation task as a function of the parameters (p,a) for different pairs of graphs. Due to lack of space this is provided in the supplementary material in Figure 11 We find that the accuracy surfaces are relatively smooth and that the optimal parameters greatly depend on the types of graphs transported.
Thanks to efficient ULOT framework,
Optim. la, with respect
to KL for (1. 3) - (2,
1.0
0.6
Since ULOT OT plans are by construction fully differentiable with respect p and a, we propose optimize them, but taking as objective classical smooth proxy for the accuracy: the Kullback-Leibe divergence between the one-hot encoded target classes and the predicted label scores. We show the optimization trajectories of parameters (p. a) for different simulated pairs of graphs in Figure 4 We see that the trajectories vary significantly between the two types of graph pairs. Indeed between different types, p has to be small to avoid mass transfer between clusters that are not present in both graphs, while for the same types, p can be larger to allow mass transfer between clusters that are present in both graphs. Interestingly, we see that for the same graph pair types, the general trend is similar but the trajectories converge to different values, underlining the necessity of parameter validation for individual pairs. This proof of concept shows that ULOT parameters can be optimized for a given task, allowing for efficient bi-level optimization when the inner optimization problem is a FUGW.
0.2
Initialication Convergence
10
10-1
10"
Figure 4: (a,, p) optimization trajectories for (top) different types (1, 2) - (1, 2), (bot- tom) same type (1, 2)..
Optimizing a graph wrt the FUGW loss One very interesting aspect of our method is that the predicted transport plan is fully differentiable with respect to the graph structure and features. We illustrate this by optimizing functional of graph. Given target graph G*. we optimize the function F(G) IG G Po (G, G))) where we expect the graph G to converge to or close to the graph G*. Starting from Go G, at each time step t, we predict the transport plan Py (G., G*) between G, and G*, compute the associated FUGW loss and update the nodes features and shortest path distance matrix of G, using backpropagation. We recover the adjacency matrix at time step I ++ 1 by thresholding the shortest path distance matrix. The influence of the tradeoff parameters is shown on Figure [ The trajectory for a = 0.5 shows that the graph converges to two-cluster graph with proper labels. In contrast, the trajectory for 1 shows that the graph converges to two-cluster
7

---

a 0.5
==1
step 0
target
Figure 5: Gradient descent steps of the minimization of of the ULOT FUGW loss between source graph (step 0) and target graph for = 0.5 (top) and (bottom).
Similarity matrix
MDS on the ULOT plan mass
(1,2)
ULO1 plan masss
(2.3)
0.8
(12,3)
(1.2) clusters
(2.3)
(1,2,3) clusters
(1.2)
(2.31
(1,,31
Figure 6: (left) Similarity matrix of the ULOT transport mass between simulated SBM graphs with
respective clusters (1, (2,3) and (1, 2,3) sorted by type. (left) MDS the similarity matrix.
graph with wrong labels. This occurs because only the Kttooov-aaseesee term is optimized, so the node features are not updated.
Using the ULOT transport plan mass The FUGW loss computed with the ULOT plan provides meaningful distance between graphs, but its computation is On nεn1), where n, n2 are the number of graph nodes. In contrast, computing the ULOT transport plan only has quadratic time complexity. unbalanced OT, the total mass of the OT plan m En, P.,j will decrease when the two graphs are very different, since in this case the marginal violation will less than the transport cost. This is why we propose to use the ULOT transport plan mass as a graph similarity measure (positive and between 0 and 1).
We evaluate this approach synthetic SBM drawn from three cluster configurations: (1, (2,3) and (1, 2,3).
compute the ULOT transport plan using fixed hyperparameters a 0.5 and 0.01. As shown in Figure 6 (left), the resulting similarity matrix reveals the structure of the dataset. This similarity matrix can naturally be used for (spectral) graph clustering, or even for dimensionality reduction and visualization illustrated in Figure (right) with multidimensional scaling of the similarity matrix where the relation between the types of graphs is clearly recovered.
Solving FUGW for Functional MRI brains
Dataset description We now evaluate ULOT on the Individual Brain Charting (IBC) dataset 23 which is a dataset of functional MRI activations on brain surfaces. The dataset is made of surface meshes with 160k vertices associated with different fMRI activations. As the dataset focuses on individual information, includes subject-specific brain geometrical models, associated with individual fMRI activations. This justifies the necessity of the unbalanced framework to adapt mass between regions that can vary in size between subjects [30]
Experimental setup The high dimensionality of the meshes in the dataset makes it particularly interesting for data augmentation. order to obtain more training graphs, we perform a parcellation using Ward algorithm (29), and reduce the graph size to 1000 nodes, where the node features are the fMRT activations averaged over the grouped vertices. Data augmentation consists of generating 10 different graphs for each of the 12 subjects, the results of the Ward clustering using randomly

---

FUGW loss for ULOT and solver
FUGW loss error and plan prediction time for ULOT and salvers
1.0
0.04
0.06
correlation
ULOT
0.03
0.8
0.02
MM
LBFGS6
0.03
Loss
0.01
Entropic le-
0.02
Entropic
0.01
0.2
0.00
0.0
0.00
0.02
0.04
10
10a
10'
10'
10*
Time
Figure 7: (left) Comparison of the loss obtained with ULOTT and IBPP solver, the dashed lines correspond to equality. (right) Plot of loss error VS time for ULOT and other solvers. Colored squares correspond to the 20-80% quantiles for both measures.
sampled activations. The geometric information matrix D is the shortest path distance matrix and the 3D node positions are concatenated with the node activation features to provide a positional encoding. We construct dataset of the 14400 different graph pairs and train the network using a 60/20/ train/val/test split. We provide more details on the experimental setup in Section
Comparison to solvers in terms loss and computational time loss of OT plans predicted by ULOT and the IBPP solver. We find in Figure (left) that both losses are very close and highly correlated with Pearson correlation of 0.99.
We first compare the FUGW
Next we compare the loss error (wrt the best among all solvers) of ULOT and the other solvers introduced in section 2.1 namely IBPP [39, MM [6) LBFGSB and Sinkhorn for the entropic regularized FUGW 30 using the Python library POT [9] We find in figure 7 (right) that even though ULOT makes errors, it up to 100 times faster than classical solvers and 10 times faster than Sinkhorn for smaller error.
This computational gain on graphs of size 1000 very important as the solvers have a cubic time complexity with respect to the number of nodes, while ULOT has quadratic time complexity as shown in Figure 8 (left) where computation time is plotted against the number of nodes.
ULOT as warmstart. Finally when high precision needed, we can use ULOT as a very efficient warm start for the IBPP solver. We find in Figure (right) that using ULOT as warm start allows the solver to converge much faster. This means that high precision required, using ULOTT as warmstart for solver is an efficient alternative.
IMRI activations prediction using ULOT plans We now illustrate the use of ULOT transport plans fMRI data for activation prediction between brain graphs. We use the same experimental setup as in 30 where fMRI activations from source brain graph G₁ are transported to predict activations target brain graph G2. First, compute the ULOT transport plan pouP between the two graphs using the model trained above. Then, we predict for new mental task, colorless colorless_auditoryy auditory contrast from MathLanguage, the fMRI activations Fy on the nodes of G2 by transporting the activations Fj from graph G1:
F = diag (PD..)) #2
(PO.P) Fi.
(13)
We visualize Figure the activations on the parcellated brain regions and observe that the general [MRI trend is conserved through the transportation. While this is only a qualitative experiment, this opens doors for future use of large scale experiments on fMRI data.
4 Conclusion, limits and future work
We have introduced ULOT, a new unsupervised deep learning approach for predicting optimal transport plans between graphs, trained by minimizing the FUGW loss. We have shown that ULOT is able to predict transport plans with low error on both simulated and real datasets, up 100 times faster than classical solvers. Its low complexity and differentiaaiilty make it naturally efficient for

---

Plan prediction time for different graph sizes
Solver iterations with and without ULOT warmstart
10"
ULOT
solver
Salver
0.005
without warmstart
selver with warmstart
Time t53
10a
OL
0.003
10
0.002
0.001
0.000
10
30
10a
10
10-1
100
101
10a
Graph size
(5)
Figure 8: (left) FUGW transport plan prediction time for ULOT and IBPP solver for different graph sizes. (right) IBPP solver loss along iterations with and without ULOT warmstart, reported with the 20% 80% quantiles.
Source brain
Transported activations
Target brain
0
Figure 9: Example of transporting colorless __tussoryrr activations from the MathLanguage contrast between two subjects from the IBC dataset. Each brain is equipped with its individual geometry
minimizing functionals of optimal transport plans and performing FUGW parameter selection. ULOT also allows for discovering novel ways use FUGW OT plans such as using their total mass measure of similarity of complexity (n²²) between graphs of different sizes.
While we believe that ULOT is very promising step towards the use of deep learning for optimal transport, we also acknowledge its limitations and propose research directions for addressing them. The very fast prediction comes the cost of small error in the predieted transport plans, which can limit its applications in a context where high precision is needed. While this can be avoided by using ULOT as a warmstart for solver, there is still room for improvement for directly predicting even more accurate plans. Also we were limited in our experiments to graphs of size n S 10000 due GPU memory constraints and going further might require dedicated developments such as lazy tensors or other memory-efficient techniques [7] for cross-attention. In the future, we plan to apply this method large-scale applications such as activation prediction high-resolution brain surfaces and computation of graph barycenters. While these applications require large training datasets, which is not the norm for fMRI data, we plan to further investigate our random parcellation data augmentation technique to train a more general model that can be effective across subjects.
Acknowledgments and Disclosure of Funding
This work was granted access to the HPC resources of IDRIS under the allocation made by GENCI. This work is supported by Hi! PARIS and ANR/France 2030 program (ANR-23 IACI-0005). This research was also supported in part by the French National Research Agency (ANR) through the MATTER project (ANR-23 2ANR-23-ERCC-OO 220-----CC-0000---1) This work benefited from state aid managed by the Agence Nationale de la Recherche under the France 2030 programme, reference ANR-22-PESN-ONN and from the European Union's Horizon 2020 Framework Programme for Research and Innovation under the Specific Grant Agreement HORIZON-INFRA 202---2222-2----000 grant agreement number 10. 30.3000110 147319. Finally, it received funding from the Fondation de I' l'Éccole polytechnique. We thank Quang Huy Tran for providing code for the FUGW solvers and for helpful discussions.
1016350
10

---

References
[1] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Op- tuna: A next-generation hyperparameter optimization framework. In The 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages (623-2631, 2019.
[2] Brandon Amos, Samuel Cohen, Giulia Luise, and levgen Redko. Meta optimal transport. arXiv preprint arXiv:2206.0055 05262, 2022.
[3] Brandon Amos et al. Tutorial on amortized optimization. Foundations and Trends in Machine Learning, 16(5):92--732 2023.
[4] Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. A limited memory algorithm for bound constrained optimization. SIAM Journal on scientific computing, 16(5):1199) 1900-128,, 1995.
[5] Laetitia Chapel, Mokhtar Z Alaya, and Gilles Gasso. Partial gromov-wasserstiin with appli- cations on positive-unlabeled learning. Advances in Neural Information Processing Systems, 2020.
[6] Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced opti- mal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270-23288 2021.
[7] Benjamin Charlier, Jean Feydy, Joan Alexis Glaunes, François-Dave -avid Collin, and Ghislain Durif. Kernel operations on the gpu, with autodiff, without memory overflows. Journal of Machine Learning Research, 22(74):1-66 2021.
[8] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xave Xaviir Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11):3000-3123, 2018.
[9] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamoniy, levgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
[10] Zheheng Jiang, Hossein Rahmani, Plamen Angelov, Sue Black, and Bryan M Williams. Graph- context attention networks for size-varied deep graph matching. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages
2022.
[11] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. arXiv preprint arXiv:2201. 12220, 2022.
[12] Wuyang Li, Xinyu Liu, and Yixuan Yuan. Sigma: Semantic-complete graph matching for domain adaptive object detection. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 5291-5300, 2022.
[13] Yujia Li, Chenjie Gu, Thomas Dullien, Oriol Vinyals, and Pushmeet Kohli. Graph matching networks for learning the similarity of graph structured objects. In International conference on machine learning. pages 3835-3845. PMLR, 2019.
[14] Xiang Ling, Lingfei Wu, Saizhuo Wang. Tengfei Ma, Fangli Xu, Alex X Liu, Chunming Wu, and Shouling Ji. Multilevel graph matching networks for deep graph similarity learning. JEEE Transactions on Neurul Networks and Learning Systems, 34(2):799-813, 2021.
[15] Chang Liu, Shaofeng Zhang, Xiaokang Yang, and Junchi Yan. Self-supervised learning of visual graph matching. In European Conference on Computer Vision, pages 370-388. Springer, 2022.
[16] Eliane Maria Loiola, Nair Maria Maia De Abreu, Paulo Oswaldo Boaventura-Netto. Peter Hahn, and Tania Querido. survey for the quadratic assignment problem. European journal of operational research, 17662):657-690, 2007.

---

[17] Facundo Mémoli. Gromov-wasserieee wasserstein distances and the metric approach to object matching. Foundations of computational mathematics, 11:417-487, 2011.
[18] Aldo Moscatelli, Jason Piquenot, Maxime Bérar, Pierre Héroux, and Sébastien Adam. Graph node matching for edit distance. Pattern Recognition Letters, 184:14-20, 2024.
[19] Maksim Nekrashevich, Alexander Korotin, and Evgeny Burnaev. Neural grosov-wasserslei optimal transport. arXiv e-prints, pages arXiv-2303, 2023.
[20] Maks Ovsjanikov, Mirela Ben-Chen, Justin Solomon, Adrian Butscher, and Leonidas Guibas. Functional maps: flexible representation of maps between shapes. ACM Transactions on Graphics (ToG), (1):1111,, 2012.
[21] Gabriel Peyré, Marco Cuturi, and Justin Solomon. Gromov-wassersten averaging of kernel and distance matrices. In International conference on machine learning, pages 2664-2672. PMLR, 2016.
[22] Chengzhi Piao, Tingyang Xu, Xiangguo Sun, Yu Rong, Kangfei Zhao, and Hong Cheng. Computing graph edit distance via neural graph matching. Proceedings of the VLDB Endowment, 16(8):1817-1829, 2023.
[23] Ana Luísa Pinho, Alexis Amadon, Torsten Ruest, Murielle Fabre, Elvis Dohmatob, Isabelle Denghien, Chantal Ginisty, Séverine Besmui- Séverine Roger, Laurence Laurier, et al. Individual brain charting, high-resolution firr dataset for cognitive mapping. Scientific data, 5(1):1-15, 2018
[24] levgen Redko, Nicolas Courty, Rémi Flamary, and Devis Tuia. Optimal transport for multi- source domain adaptation under target shift. In The 22nd International Conference on artificial intelligence and statistics, pages 849-858. PMLR, 2019.
[25] Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. Superglue: Learning feature matching with graph neural networks. In Proceedings of the IEÉE/CVF conference on computer vision and pattern recognition, pages 4938 4948,4977,45 2020.
[26] Vivien Seguy, Bharath Bhushan Damodaran, Rémi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large-scale optimal transport and mapping estimation. arXiv preprint arXiv: 2017.
[27] Thibault Séjoumé, François- -avier Vialard, and Gabriel Peyré. The unbalanced gromov wasserstein distance: Conic formulation and relaxation. Advances in Neural Information Processing Systems, 34:8766-87777 8779, 2021.
[28] Justin Solomon, Raif Rustamov, Leonidas Guibas, and Adrian Butscher. Wasserstein propa- gation for semi-supervised learning. In International Conference on machine learning, pages 306-314, PMLR, 2014.
[29] Bertrand Thirion, Gaël Varoquaux, Elvis Dohmatob, and Jean Baptiste Poline. Which fmri clustering gives good brain parcellations? Frontiers in neuroscience, 8: 8:167, 2014.
[30] Alexis Thual, Quang Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced gromov wasserstein. Advances in neural information processing systems, 35:21792-21804, 2022.
[31] Vayer Titouan, Nicolas Courty, Romain Tavenard, and Rémi Flamary. Optimal transport for structured data with application on graphs. In International Conference on Machine Learning, pages 6275-6284. PMLR, 2019.
[32] Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. arXiv preprint arXiv:2302.0044 2023.
[33] Siddharth Tourani, Muhammad Haris Khan, Carsten Rother, and Bogdan Savchynskyy. Discrete cycle-consistency based unsupervised deep graph matching. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pages 5252-5260, 2024.
12

---

[34] Titouan Vayer, Lactitia Chapel, Rémi Flamary, Romain Tavenard, and Nicolas Courty. Fused gromov-wasserstein distance for structured objects. Algorithms, 13(9):212, 2020.
[35] Runzhong Wang, Ziao Guo, Shaofei Jiang, Xiaokang Yang, and Junchi Yan. Deep learning of partial graph matching via differentiable top-k. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 6272-6281, 2023.
[36] Runzhong Wang, Junchi Yan, and Xiaokang Yang. Learning combinatorial embedding networks for deep graph matching. In Proceedings of the IEEE/CVF international conference computer vision, pages 3056-3065, 2019.
[37] Runzhong Wang, Junchi Yan, and Xiaokang Yang. Graduated assignment for joint multi-graph matching and clustering with application unsupervised graph matching network learning. Advances in neural information processing systems, 33:19908-19900
19919999889999,--111199919999999 19919, 2020.
(38) Tao Wang and Ziv Goldfeld. Neural entropic gromov-wasserstein alignment. arXiv preprint arX.v:231..07397. 2023.
[39] Yujia Xie, Xiangfeng Wang. Ruijia Wang, and Hongyuan Zha. A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence, pages 433-453. PMLR, 2020.
[40] Hongteng Xu, Dixin Luo, Hongyuan Zha, and Lawrence Carin Duke. Gromov-wasserstein learning for graph matching and node embedding. In International conference machine learning, pages 6932-6941. PMLR, 2019.
[41] Andrei Zanfir and Cristian Sminchisescu. Deep learning of graph matching. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages
2018.
[42] Zhengxin Zhang, Ziv Goldfeld, Youssef Mrouch, and Bharath K Sriperumbudur. Gromov- wasserstein distances: Entropic regularization, duality and sample complexity. The Annals of Statistics, 1614)1661, 2024.
13

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction accurately reflect the contributions by presenting the experimental results explained throughout the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in Section and the conclusion.
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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
14

---

Justification:
The paper is experimental and does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The experimental setup and hyperparameters and presented in the main paper and Section[A]
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
15

---

Answer:[Yor Yes]
Justification:
The code will be provided in the supplemental materials and at https
//github. com/ smazelet/ULOT] The data public.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips. //nips.cc/pil cc/public/guu http//uddde//CdSSubn
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experimental details and hyperparameters and presented in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: Quantitative results are reported with medians and 20% 80% quantiles
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
16

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 96% CI, if the hypothesis of Normality 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 out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computer resources are presented in Section
Guidelines:
The answer NA means that the paper does 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 failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https::/neurips ctcpippp.ccp/ciiie http:/rpp....pcuuooocccccccccaae
Answer: [Yes]
Justification: The fMRI IBC dataset respects the NeurIPS code of conduct in terms of privacy and content. To the best of our knowledge, this paper does not have potential harmful consequences.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The paper investigates fast approximation of solutions of optimal transport problems. There is no societal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
17

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We are the creators of all new assets in the paper, and use code under MIT license. We use the fMRI IBC dataset which is under license CCO.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt
package should be provided. For popular datasets, paperswi thcode. cww/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
18

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
The code details are provided reproduce the simulated SBM dataset.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/coodemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involved crowdsourcing nor new research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending 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 NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
19

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: /ttppr////euips..... wwwparippssc.coooeeeeeeeee for what should or should be described.
20

---

FUGW loss for ULOT and salver
Transport plan values for increasing values
10
0.05
pearson correlatio=11000
P.j for
0.04
0.6
velume the leamee
0.06
A, for and different clusters
0.05
ULOT
0.03
0.6
Pavave
0.02
0.92
0.01
0.00
0.0
0.00
0.02
0.04
10-
10
10
Selver
Figure 10: (left) Comparison of the loss obtained with ULOT and IBPP solver for the simulated graphs, the dashed line corresponds to the equality. (right) Transport plan values for increasing values of p, colored by whether they link common different clusters.
A Training setup details
Compute resources We trained our network on an NVIDIA V100 GPU for 100 hours on the IBC dataset and a few hours on the smaller simulated dataset. Note that while the network is O(n²) with n the number of graph nodes, the main training bottleneck comes from the need to compute the (nn)) FUGW loss for each transport plan at every epoch.
Hyperparameters The hyperparameters used for training ULOT on both the simulated graphs and the IBC dataset are reported in [ All the MLP and GMN have one hidden layer, with weights shared between the two graph branches. Moreover, the first MLP the node embedding layer preserves the dimensionality of the input node features. Hyperparameter optimization was performed on subset of the training data using the Optuna library [. The code is available in the supplementary materials and will be released github upon publication. We will also share the pre-trained model weights for both datasets.
Table 1: ULOT hyperparameters
Hyperparameter
Simulated dataset
IBC dataset
Learning rate
0.001
0.0001
Batch size
256
64
Optimizer
Adam
Adam
Number of node embedding layers N Embedding dimension for
5
10
10
Node embedding layer final out dimension
256
256
MLP hidden dimension
64
256
GCN hidden dimension Temperature value a
16
128
B Additional experiments on the simulated graphs
Comparison of FUGW loss for ULOT and solver on the simulated graphs We train ULOT on the dataset of simulated SBMs introduced in Section 3.1 and test on new pairs of simulated graphs sampled from the same distribution. We find in Figure I0 (left) that similarly to the IBC dataset ULOT finds transport plans that have FUGW loss perfectly correlated with the FUGW loss obtained with the IBPP solver.
Regularization path of the transport plan component Pa,j with respect to p We predict ULOT transport plans between pair of graphs with cluster configurations (1, and (2, 3) for different values and show the regularisation path of each transport plan entry (Prt.) (., for E [1, m1) and j E na] in Figure 10 (right). We observe that transport plan values corresponding to nodes in the shared cluster 2 increase more rapidly with p compared to values between nodes in non-overlapping clusters. Moreover, we find that there exists optimal value around p ~ 0.05, for which transport
21

---

Accuracy for (1, (1,
Accuracy for (1. 3) (1,
1.0
1.0
0.92
0.8
0.92
0.8
0.73
0.73
Accurait
0.6
0.55
0.55
0.37
0.4
0.37
0.4
0.18
0.18
0.00
0.00
0.004
0.025
0.159
1.00
0.0
0.004
0.025
0.159
1.00
0.0
Figure 11: Label propagation accuracy of the ULOT FUGW transport w.r.t (p. between (left) different types 21, - (1, 2), (right) same type (1, ()).
plan mass is predominantly assigned to entries corresponding to the common cluster, effectively discarding irrelevant correspondences.
Visualization of the accuracy for label propagation We visualize the accuracy surfaces on Figure I1 for the label propagation task introduced Section 3.1 for a range of (a, p) values. We consider two different types of pairs: pairs with clusters 2, 3) and (1, 2) and pairs with clusters (1, 2, 3). We see that the accuracy surfaces are smooth and that the optimal parameter values differ across pair types. We optimize p and using gradient descent on each pair of graphs by minimizing the KI. divergence between the predicted class probabilities obtained with the ULOT transport plan and the ground truth target one hot encodings of the classes on 50% of the nodes. We obtain an accuracy of 0.87 + 0.092 on pairs with similar clusters (1, 2,3) and 0.74 0.11 on the pairs with different clusters.
C Ablation studies
We provide ablation study of the impact of each step of the OT plan block the SBM dataset in T As first ablation we remove the node scaling in the OT plan head, as second ablation we replace the whole OT plan head by much simpler nonlinear transformation of the node embedings to recover an plan with (Pla.P) )s, = (+T")) ""F"""")"F""")))) "Fl where Fftinal and Ffinal are the outputs of the GCN and cross attention block for the first and second graphs respectively.
Table 2: Ablation studies
Ablation
Relative error
Pearson correlation
None
0.19 0.29
1.0
Remove node scaling
4.1 + 9.5
0.89
Replace OT plan block by nonlinearity
27 49
0.48
We see that node scaling essential for recovering unbalanced OT plans especially for small p parameters (unbalanced plans that do not sum to one), and removing them increases the error (but remains surprisingly well correlated). Using simpler output head for the OT plan leads failure to predict OT plans.
D Robustness to distribution shifts
We tested the following distribution shift: we trained ULOT on dataset of Stochastic Block models with balanced clusters (every cluster has same number of nodes), and tested the model on dataset of SBMs with unbalanced clusters, where the cluster ratio is sampled following Dirichlet distribution. We report the results Table and find that even though the relative loss error increases, the Pearson correlation between the ground truth losses and the ULOT losses remains stable.
22

---

Table 3: Distribution shifts
Distribution
Relative error
Pearson correlation
Balanced clusters (train distribution)
0.19 + 0.29
1.0
Unbalanced clusters (distribution shift)
0.27 + 0.69
0.98
E
Illustration of the fMRI alignement
Experimental details on the IBC dataset We use brain cortical surfaces from the IBC dataset consisting of approximately 160k vertices, each associated with fMRI contrasts obtained as subjects perform specific tasks. To train ULOT, we use the 239 contrasts from the tasks in Table 4 The predicted activation visualized Figure 9 is selected from the 30 left contrasts from the LathLanguage task. For each subject, we perform a parcellation of the surface to form 1000 brain regions, using the Ward's hierarchical clustering algorithm. Each contrast is then averaged over the vertices in every brain region. From this parcellation, we construct 1000 node graphs, where edges connect spatially adjacent regions. Each node has 242 dimensional feature vector composed of the region's ss contrasts concatenated to its 3D node position.
To augment the dataset, we generate multiple parcellations for each subject by randomly selecting 20% to 40% of the tasks, which produces variability in the geometries and the activations across all generated brains. The final dataset constructed from all possible graph pairs, which randomly split into 60% training, 20% validation, and 20% test sets.
Table 4: IBC tasks used for alignment
Task
ArchiEmotional
ArchiSocial
ArchiSpatial
ArchiStandard
HcpEmotion
HepGambling
HepLanguage
HepMotor
HepRelational
HepSocial
HcpWm
RSVPLanguage
PreferenceFacees
PreferenceHouses
Preference
PreferencePaintingg
MCSE
Moto
Visu
Audi
MVEB
MVIS
Lecl
Lec2
TheoryOfMind
PainMovie
EmotionalPain
Enumeration
VSTM
Self
23


Reviews:
[{"review_id": "1vu6Cj72xO", "reviewer": "Reviewer_bmor", "summary": "The paper proposes a deep learning method called Unbalanced Learning of Optimal Transport (ULOT) for predicting optimal transport plans between unbalanced graphs, addressing the computational inefficiency of classical solvers. The method minimizes the Fused Unbalanced Gromov-Wasserstein loss using a novel neural network architecture that integrates cross-attention and graph neural networks. This architecture is conditioned on the FUGW hyperparameters, allowing ULOT to generalize across different parameter values.", "questions": "See weaknesses.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1.\tUsing neural networks to approximate the transportation plan seems a new task and the inference time can be much less than traditional solvers.\n2.\tThe paper is well organized. The method is relatively simple, but the experiments are comprehensive.\n3. Results are quite good. It is easy to understand ULOT has better speed but it is surprising to see it even has lower error than other solvers.\n\n\n\nWeaknesses:\n1. The motivation of the model design is somehow not very clear. Why does the designed neural network can approximate the transportation plan? how well can it approximate the transportation plan (is there any theoretic bound)?\n2. Some training details are also lacking, e.g. what solvers are used as ground truth and what is the loss function?\n3. I assume one of the major drawbacks of using neural network for OT approximation is generalization. It would be good to see if the method can perform well when the test data is very different from training data.", "was_revised": true, "final_justification": "The rebuttal clarified the generalization experiments and somehow addressed my concerns. So I raised the rating to 4.", "rebuttal": "Thank you for accurately summarizing our paper and providing a fair evaluation of its strengths and weaknesses. We appreciate your questions and have done our best to answer them below. \n\n > Results are quite good. It is easy to understand ULOT has better speed but it is surprising to see it even has lower error than other solvers.\n\nWe would like to clarify why the ULOT loss is sometimes lower than the solver loss, which is explained by two factors. \n\nFirst, ULOT can outperform solvers when they converge to poor local minima. This happens for instance in scenarios such as the one presented in Figure 3 (left) where the constant plan equal to zero is a local minima. Because of its design that implies continuity with respect to its input, ULOT often avoids these local minima, which we believe is a crucial asset. \n\nSecond, in cases such as the one depicted in Figure 7 (right), ULOT has lower error than the entropic solver with high regularization, which is expected. Indeed, the entropic regularization is introduced to speed up the optimization at the cost of some non-optimality in the transport plan. The intuition is that lower entropic regularization leads to lower error but slower convergence.\n\n\n> The motivation of the model design is somehow not very clear. Why does the designed neural network can approximate the transportation plan?\n  \nWe appreciate your question, as it is fundamental that the model design be clearly motivated and we will improve this aspect in the final version. We designed the model so that it can learn transport plans that respect the conditions for minimizing the FUGW loss (unbalanceness, feature and structure integration). We originally tried other simpler architectures, but their limited performance led to the current architecture that we detail below.\n\n* The first block outputs node features that are relevant in a matching context, and is an adaptation of [13]. Intuitively the cross attention mechanism computes node features that are similar for nodes that can be matched and dissimilar for nodes that should not. The use of GCNs is motivated by the Gromov term in the FUGW loss that enforces that the model should take into account the graph topology (features are already used in MLP and cross-attention). \n* The optimal transport plan prediction block is designed so as to predict an unbalanced OT plan that must contain pairwise relationships between nodes but also that can discard (reweight) nodes that cost too much from an OT perspective. This is done by separating the two aspects:\n    *  We chose to use for the pairwise relationship in the OT plan the average of the cross attention matrix to which we apply a row-wise softmax and a column-wise softmax respectively. \n    * The pairwise matrix above gives a good starting point for the plan, which is then refined in the unbalanced optimal transport plan block, where the scaling per node makes it no longer balanced. This is why we also provide in the last block two node reweighting heads (Linear+sigmoid and multiplication of the plans by the diagonal weight matrices) that allow scaling on the OT plan and enable the network to discard mass.\n     \nWe provide an ablation study of the impact of each step of the OT plan block on the SBM dataset in the table below. As a first ablation we remove the node scaling in the OT plan head, as a second ablation we replace the whole OT plan head by a much simpler nonlinear transformation of the node embedings to recover an OT plan with  ($P_\\theta^{(\\alpha, \\rho)})_{i,j} = \\frac{1}{1+||(F^{\\text{final}}_1)_i -(F^{\\text{final}}_2)_j||^2_2},$ where $F^{\\text{final}}_1$ and $F^{\\text{final}}_2$ are the outputs of the GCN and cross attention block for the first and second graphs respectively.\n   \n\n| Ablation |   Relative error   |  Pearson correlation |  \n| -------- | -------- | -------- | \n|  None   |   $0.19 \\pm 0.29$  |  $1.0$ | \n| Remove node scaling | $4.1 \\pm 9.5$    |  $0.89$  | \n| Replace OT plan block by nonlinearity | $27 \\pm 49$   |  $0.48$  | \n\nWe can see that node scaling is essential for recovering unbalanced OT plans especially for small $\\rho$ parameters (unbalanced plans that do not sum to one), and removing them increases the error (but remains surprisingly well correlated). Using a simpler output head for the OT plan leads to a failure to predict OT plans.\n\nThis explanation and ablation study will be added to the paper.\n\n\n> how well can it approximate the transportation plan (is there any theoretic bound)?\n\nWe provide empirical results about the quality of the OT plans and of their approximation. To the best of our knowledge there is no theoretical guarantee about the estimation quality of amortized optimization [3], so this is still an open question.\n\n> Some training details are also lacking, e.g. what solvers are used as ground truth and what is the loss function?\n\nThe loss function is the FUGW loss and is introduced in equation (4). The solvers used as ground truth are presented in subsection 2.1. We used the IBPP [39] solver for all the results except for Figure 7 (right), where we compare many different solvers, namely the MM [6], LBFGSB [30] and Sinkhorn solvers [8]. When using several numerical solvers, the ground truth is taken as the solution with the smallest objective.\n\nWe will provide in the final version more details about all those solvers in the supplementary of the paper to be more self-contained.\n\n> I assume one of the major drawbacks of using neural network for OT approximation is generalization. It would be good to see if the method can perform well when the test data is very different from training data.\n\nThis is a very interesting point that is crucial in understanding the potential of ULOT for applications. We would like to clarify that all the performances reported are generalizations on test data, i.e. graph pairs that were not used for training, which testifies for the ability of the model to generalize to different pairs within the same distribution. \n\nThe question of robustness to data shift is indeed important. We tested the following distribution shift: we trained ULOT on a dataset of Stochastic Block models with balanced clusters (every cluster has the same number of nodes), and tested the model on a dataset of SBMs with unbalanced clusters, where the cluster ratio is sampled following a Dirichlet distribution. We report the results in the table below and find that even though the relative loss error increases, the Pearson correlation between the ground truth losses and the ULOT losses remains stable.\n\n\n|  | Relative loss error | Pearson correlation |\n| -------- | -------- | -------- |\n| Balanced clusters (train data)    |   $0.19 \\pm 0.29$  |  $1$   |\n| Unbalanced clusters (distribution shift) | $0.27 \\pm 0.69$    |  $0.98$  |\n\n\n### References\n\n[6] Laetitia Chapel, Rémi Flamary, Haoran Wu, Cédric Févotte, and Gilles Gasso. Unbalanced opti- mal transport through non-negative penalized linear regression. Advances in Neural Information Processing Systems, 34:23270–23282, 2021.\n\n[39]Yujia Xie, Xiangfeng Wang, Ruijia Wang, and Hongyuan Zha. A fast proximal point method for computing exact wasserstein distance. In Uncertainty in artificial intelligence, pages 433–453. PMLR, 2020.\n\n[8] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulations. Journal of Functional Analysis, 274(11):3090–3123, 2018.\n\n[30] Alexis Thual, Quang Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced gromov wasserstein. Advances in neural information processing systems, 35:21792–21804, 2022.\n\n[13] Yujia Li, Chenjie Gu, Thomas Dullien, Oriol Vinyals, and Pushmeet Kohli. Graph matching networks for learning the similarity of graph structured objects. In International conference on machine learning, pages 3835–3845. PMLR, 2019.\n\n[3] Amos, B. (2023). Tutorial on amortized optimization. Foundations and Trends® in Machine Learning, 16(5), 592-732."}, {"review_id": "eue8pPt1f0", "reviewer": "Reviewer_evnK", "summary": "In this paper, the authors propose a method for fast approximation of Fused Unbalanced Gromov Wasserstein OT plan with neural nets, trained in an unsupervised way. Their main idea is to directly plug the neural architecture within the FUGW cost, which has been proposed previously for regular OT but, to my knowledge, not for FUGW. The authors argue that one of the main advantage is the computational gain in handling the notoriously difficult quadratic problem associated with GW. Another notable advantage is that everything is differentiable, including with respect to the regularization parameter of the OT problem, or the input graph. Their architecture starts with regular GNNs to compute node representations, then the computation of an OT plan from a mix of cross-attention and fully connected layers. This includes the learning of node weights, as is required by unbalanced OT. The authors then test extensively their architecture on synthetic SBM data, where node features are the community labels. They include (sometimes supervised) optimization with respect to the different hyperparameters, or even the input graph. They also examine real MRI data, where they also test the use of the proposed method as warm-start for other solvers.", "questions": "See above", "limitations": "Limitations are adressed properly.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\n- the paper is very well-written, clear, and quite complete.\n- The proposed method is natural but novel, and handled in a clever way despite being quite non-trivial.\n- Experimental results are generally convincing, especially the computational advantages of the proposed method, which is the main strength of the proposed method here\n\nWeaknesses\n- I still have one interrogation on the architecture itself, which is not really alleviated by the experiments. It seems the role of node features is rather important compared to the graph structure: the latter only intervene in the early \"GNN\" part of the architecture, and the computed representations are then buried under several layers of attention, similarities, and so on. This is only an intuition, but node features seem preeminent here. Is the architecture performant when node features are not \"useful\" ? In the experiments, the node features contain the (noisy) labels, so a great deal of immediate information. The authors do test $\\alpha=1$, which would completely ignore the node features when solving FUGW directly, but here the node features are still in the parameterized transport plan, so are not ignored! It would be useful to test the architecture without node features altogether (or with more complex features that noisy labels). In other words, can the proposed architecture handle non-fused GW?\n- The authors test a lot of aspects of their architecture, so some small experiments are not necessarily useful, and may take up some space. For instance, the \"graph learning\" part (differentiability wrt the graph) only tries to learn the identity in the proposed experiments, which is very limited. Maybe it's better to put it in a more developed form in the appendix, or not at all and develop more the other experiments.", "was_revised": true, "final_justification": "I thank the authors for their answers, and maintain my positive score.", "rebuttal": "Thank you for your positive review about our paper. We also appreciate your remarks, as we think they raise fundamental questions regarding the model design.\n\n> I still have one interrogation on the architecture itself, which is not really alleviated by the experiments. It seems the role of node features is rather important compared to the graph structure: the latter only intervene in the early \"GNN\" part of the architecture, and the computed representations are then buried under several layers of attention, similarities, and so on. This is only an intuition, but node features seem preeminent here. Is the architecture performant when node features are not \"useful\" ? In the experiments, the node features contain the (noisy) labels, so a great deal of immediate information. The authors do test , which would completely ignore the node features when solving FUGW directly, but here the node features are still in the parameterized transport plan, so are not ignored! It would be useful to test the architecture without node features altogether (or with more complex features that noisy labels). In other words, can the proposed architecture handle non-fused GW?\n\nThis is a very interesting question, and you are completely right about the fact that the node features are used even for $\\alpha=1$. We think that this is actually a strength of our model, which makes it competitive in situations where the Gromov problem is hard to solve. Indeed, the Gromov problem has many invariants (in many cases, the featureless graphs are invariant to some permutations of nodes), and using the node features breaks these invariants and helps predict a deterministic solution.\n\nIn a case where the features are absent or not relevant, one should not use FUGW and ULOT but rather solve pure GW which is a distinct task that needs to be designed differently, probably with stronger graph architectures such as the Evoformer architecture from [B], which has a stronger focus on the graph connectivity matrices. This is an interesting discussion that we will add to the paper.\n\n\n> The authors test a lot of aspects of their architecture, so some small experiments are not necessarily useful, and may take up some space. For instance, the \"graph learning\" part (differentiability wrt the graph) only tries to learn the identity in the proposed experiments, which is very limited. Maybe it's better to put it in a more developed form in the appendix, or not at all and develop more the other experiments.\n\nWe agree that this is not a crucial part of the paper and we will move it to the appendix. The motivation behind it was to demonstrate one use case of the differentiability of the ULOT loss which, to the best of our knowledge, has not been previously explored. It is also a classical problem studied in recent papers about distribution flows to illustrate the trajectory and the geometry of the space [A] on Euclidean distributions and more recently in the GW space [C]. \n\n\n### References\n\n\n[A] Bonet, C., Vauthier, C., & Korba, A. (2025). Flowing Datasets with Wasserstein over Wasserstein Gradient Flows. arXiv preprint arXiv:2506.07534.\n\n[B] Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., ... & Hassabis, D. (2021). Highly accurate protein structure prediction with AlphaFold. nature, 596(7873), 583-589.\n\n[C] Zhang, Zhengxin, Ziv Goldfeld, Kristjan Greenewald, Youssef Mroueh, and Bharath K. Sriperumbudur. \"Gradient Flows and Riemannian Structure in the Gromov-Wasserstein Geometry.\" arXiv preprint arXiv:2407.11800 (2024)."}, {"review_id": "qBSiIT9w46", "reviewer": "Reviewer_3cZ4", "summary": "the paper proposes a   deep learning pipeline  (named Unsupervised Learning of Optimal Transport plan prediction (ULOT)) that predicts the  transport plan as a solution to the fused unbalanced Gromov-Wasserstein loss. The proposed pipeline works faster than SOTA solvers, with getting good accuracy. The proposed pipeline can also provide initial guesses to speed up  traditional solvers.", "questions": "- line 117: why are $\\alpha$ $\\beta$ encoded *\"at each layer of the network\"*? does it mean each layer of the GCN (FIg.1)? why is it not encoded only at the input layer?\n\n - lines 108+343: can a mix approach (supervised/unsupervised) be used for optimisation? \n e.g. it seems that training sets with G2 being a  sub-graph of G1 (e.g. jackknife/bootstrap resampling) would provide access to ground truth transport plans? would that help with requirement of \"large training datasets\"?\n\n - line 257: why is KL used and not cross entropy?", "limitations": "yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n- the paper is very well presented and clear.\n- contributions of the proposed pipeline are well articulated and well validated with experimental results\n\n**Weaknesses:**\n$\\emptyset$", "was_revised": true, "final_justification": "NA", "rebuttal": "We are grateful for your recognition of the contributions of the paper and hope we can clarify the questions you raised in the following.\n   \n   \n> line 117: why are alpha and beta encoded \"at each layer of the network\"? does it mean each layer of the GCN (FIg.1)? why is it not encoded only at the input layer?\n\n\nWe designed the model so that the information about $\\rho$ and $\\alpha$ is not lost along the layers, which is a common issue when models are conditioned on a real value, such as time in diffusion models. Inspired by the method in [A], which has become a standard practice for diffusion models, we concatenate the value of $\\rho$ and the encoding for alpha to the node features at the entry of each MLP, GCN and Linear layer.  This design ensures that all layers have direct access to the conditioning parameters and can learn themselves at which layer this information is the most pertinent. \n\n> lines 108+343: can a mix approach (supervised/unsupervised) be used for optimisation? e.g. it seems that training sets with G2 being a sub-graph of G1 (e.g. jackknife/bootstrap resampling) would provide access to ground truth transport plans? would that help with requirement of \"large training datasets\"?\n\nThis is an excellent idea that should be investigated further in future work. The difficulty with the approach is that a ground truth matching (such that the partial matching we would get if one graph is a subgraph of the other) is not always an optimal transport plan. Indeed, for high values of the  $\\rho$ parameter, all the mass must be transported, so the ground truth matching in the subgraph setting is no longer the optimal OT plan. This means that we still need to rely on solvers for ground truth plans, which makes training computationally intensive. Also note that all existing solvers only provide a local minimum and training in a supervised way to predict those minimums could lead to a sub-optimal solution. \n\n\n> line 257: why is KL used and not cross entropy?\n\nGiven that the ground truth distribution is a one-hot encoding, the KL divergence and the cross entropy are equivalent. We will clarify this in the final version, thanks for your question.\n\n\n### References\n\n[A] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33, 6840-6851."}, {"review_id": "P1OuHwnIxr", "reviewer": "Reviewer_HUXF", "summary": "This paper proposed a deep learning framework to solve the problem of optimal transport between graphs. Instead of using traditional solution methods, the proposed work can use deep neural network to predict the transportation faster than the existing method. They also showed their model's efficiency on both simulated and real-world (fmri) data.", "questions": "1. D1, D2 input are the same for different GCN and cross attention layers or they are updated?\n2. What is the training time/computational complexity of the model?\n3. What is the transferability of the model? For instance, in your simulated experiment, if you train a model with node number in cluster 1, n1, and node number in cluster 2, n2, does the model still work when you change the number in cluster 1 and the number in cluster 2 relatively but keep the total number the same (e.g., n1+2, n2-2)? Or you need to train the model every time you are facing a new environment?\n4. \\alpha being 0 or 1 are two extreme numbers in the results. What about 0.2 and 0.8? I would like to know whether features or short distance path are the dominant factor in the deep learning model, and how the results different from the ground truth?\n5. The similarity matrix overall had a quite high correlation value (> 0.6). Can this be caused by the sparsity? A diagnal-block high value correlation can be easily affected by the zero connection between different clusters. Also not clear about what the reason of showing fig. 6. MDS (fig 6 right) is not defined in the paper.\n6. An example of transportation plan of real-world data from ULOT and solver will help verify the model, especially if authors can also correlate the results with the biological reasoning.", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength: \nUse a deep neural network to approximate the complex FUGW solver.\nThe computational complexity in the testing phase is much less than the traditional methods.\nWith this warm-start of the model, the traditional solver can reach much lower-error results.\n\nWeakness:\nThe explanation of the model is not enough and unclear. The whole section 2.3 can be compressed into an algorithm format, and authors should explain why they had this structure or information flow, e.g., can author explain why the formation of unbalanced OT plan layer gives the transportation plan?\nMore results should be added to verify the model. Please see questions.", "was_revised": true, "final_justification": "The main revision of the paper should focus on the motivation, idea and logics of the proposed model and its structure. In the current version, simplifying the optimal transport calculation is important, but it is also important to let people understand why and how it works. With additional explanation and experiments shown by the authors, I will increase my rating to 4.", "rebuttal": "We thank you for your thoughtful review and the important questions you raised.  \n\n> The explanation of the model is not enough and unclear. The whole section 2.3 can be compressed into an algorithm format, and authors should explain why they had this structure or information flow, e.g., can author explain why the formation of unbalanced OT plan layer gives the transportation plan? More results should be added to verify the model. Please see questions.\n\n\nWe appreciate your question, as it is fundamental that the model design be clearly motivated and we will improve this aspect in the final version. We designed the model so that it can learn transport plans that respect the conditions for minimizing the FUGW loss (unbalanceness, feature and structure integration). We originally tried other simpler architectures, but their limited performance led to the current architecture that we detail below.\n\n* The first block outputs node features that are relevant in a matching context, and is an adaptation of [13]. Intuitively, the cross attention mechanism computes node features that are similar for nodes that can be matched and dissimilar for nodes that should not. The use of GCNs is motivated by the Gromov term in the FUGW loss that enforces that the model should take into account the graph topology (features are already used in MLP and cross-attention). \n* The optimal transport plan prediction block is designed so as to predict an unbalanced OT plan that must contain pairwise relationships between nodes but also that can discard (reweight) nodes that cost too much from an OT perspective. This is done by separating the two aspects:\n    *  We chose to use for the pairwise relationship in the OT plan the average of the cross attention matrix to which we apply a row-wise softmax and a column-wise softmax respectively. \n    * The pairwise matrix above gives a good starting point for the plan, which is then refined in the unbalanced optimal transport plan block, where the scaling per node makes it no longer balanced. This is why we also provide in the last block two node reweighting heads (Linear+sigmoid and multiplication of the plans by the diagonal weight matrices) that allow scaling on the OT plan and enable the network to discard mass.\n     \nWe provide an ablation study of the impact of each step of the OT plan block on the SBM datasets in the table below. As a first ablation we remove the node scaling in the OT plan head, as a second ablation we replace the whole OT plan head by a much simpler nonlinear transformation of the node embedings to recover an OT plan with  ($P_\\theta^{(\\alpha, \\rho)})_{i,j} = \\frac{1}{1+||(F^{\\text{final}}_1)_i -(F^{\\text{final}}_2)_j||^2_2},$ where $F^{\\text{final}}_1$ and $F^{\\text{final}}_2$ are the outputs of the GCN and cross attention block for the first and second graphs respectively.\n   \n\n| Ablation |   Relative error   |  Pearson correlation |  \n| -------- | -------- | -------- | \n|  None   |   $0.19 \\pm 0.29$  |  $1.0$ | \n| Remove node scaling | $4.1 \\pm 9.5$    |  $0.89$  | \n| Replace OT plan block by nonlinearity | $27 \\pm 49$   |  $0.48$  | \n\nWe can see that node scaling is essential for recovering unbalanced OT plans especially for small $\\rho$ parameters (unbalanced plans that do not sum to one), and removing them increases the error (but remains surprisingly well correlated). Using a simpler output head for the OT plan leads to a failure to predict OT plans.\n\nThis explanation and ablation study will be added to the paper.\n\n\n> * D1, D2 input are the same for different GCN and cross attention layers or they are updated? What is the training time/computational complexity of the model?\n\n+ D1 and D2 are the inputs to the first GCN and cross attention layer and are indeed not updated along layers. We will clarify this in the text and in Fig. 1\n+ The model has worst case (for fully connected graphs) forward computational complexity of $O(n_1n_2+n_1^2+n_2^2)$, but training the model requires computing the FUGW loss for each batch, which is $O(n_1n_2^2+n_2n_1^2)$. So the overall complexity is $O(n_1n_2^2+n_2n_1^2)$. As stated in the Appendix A first paragraph, training the model on the IBC dataset (fMRI) takes about 100 hours on a V100 GPU.\n    \n> What is the transferability of the model? For instance, in your simulated experiment, if you train a model with node number in cluster 1, n1, and node number in cluster 2, n2, does the model still work when you change the number in cluster 1 and the number in cluster 2 relatively but keep the total number the same (e.g., n1+2, n2-2)? Or you need to train the model every time you are facing a new environment?\n\nThis is a very interesting point. We do not expect our model to be robust to big distribution shifts but it is still relatively robust to small shifts. For the transferability experiment you propose, we test the model trained on balanced SBMs on unbalanced SBMs with cluster ratios sampled following a Dirichlet distribution and report the performances in the table below. We find that even though the relative loss error increases, the Pearson correlation between the ground truth losses and the ULOT losses remains stable.\n\n\n|  | Relative loss error | Pearson correlation |\n| -------- | -------- | -------- |\n| Balanced clusters (train data)    |   $0.19 \\pm 0.29$  |  $1$   |\n| Unbalanced clusters (distribution shift) | $0.27 \\pm 0.69$    |  $0.98$  |\n\nWe will add a figure in the supplemental material showing the error as a function of the distance between the test graph cluster ratios and the uniform ratios to further understand the robustness of the model to data shift.\n\n> \\alpha being 0 or 1 are two extreme numbers in the results. What about 0.2 and 0.8? I would like to know whether features or short distance path are the dominant factor in the deep learning model, and how the results different from the ground truth?\n\nDuring training and testing, for each graph pair the values of $\\alpha$ are sampled following a Beta distribution with parameters $(0.5,0.5)$. This implies that the $\\alpha$ values span the whole [0,1] range. Consequently, the results of Figure 7 are reported for many different values of $\\alpha$ between 0 and 1, and we find that the error is nearly independent of $\\alpha$.\n\nTo study the impact of the GCN layer, we perform an ablation study of the GCNs on the fMRI dataset. We observe that it drastically decreases the performance from a relative loss error of 19% to 146%, which highlights the fact that the structure information actually helps the neural network find the right solution.\n\nSince our network is differentiable, we plan on adding to the supplemental material a sensitivity analysis of the parts of our model to further investigate the importance of each component, thanks for this very interesting question.\n\n\n\n> The similarity matrix overall had a quite high correlation value (> 0.6). Can this be caused by the sparsity? A diagnal-block high value correlation can be easily affected by the zero connection between different clusters. Also not clear about what the reason of showing fig. 6. MDS (fig 6 right) is not defined in the paper.\n\nThe similarity matrix $S$ in Figure 6 shows the transport plan mass $S_{i,j}$ of the ULOT plan between graph $i$ and graph $j$. The values are not correlations but mass, so it makes sense that the values are high, with maximal values at $1$ when all the mass is transported.\n\nWe agree that we should add the reference for MDS, which stands for Multidimensional scaling [A] and is a method for visualizing data points in low dimension such that their similarities are preserved. We will add the reference and explanation to the paper.\n\n> An example of transportation plan of real-world data from ULOT and solver will help verify the model, especially if authors can also correlate the results with the biological reasoning.\n\nThis is indeed a very interesting question as ULOT can help solve important problems in neuroscience and genomics. Since the loss error is low on fMRI data, we expect the transport plans to be close to the solver transport plans and thus meaningful for real-life applications as was already demonstrated in [30]. We will provide such an illustration in the supplemental material of the final version.\n\n### References\n    \n[A] Kruskal, J. B. (1964). Multidimensional scaling by optimizing goodness of fit to a nonmetric hypothesis. Psychometrika, 29(1), 1-27.\n[30] Alexis Thual, Quang Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced gromov wasserstein. Advances in neural information processing systems, 35:21792–21804, 2022."}, {"review_id": "qXt5ckOJmG", "reviewer": "Reviewer_j6iH", "summary": "This paper proposes a new approach for predicting optimal transport plans between graphs, designed to improve the computational complexity (in time) with respect to the number of nodes of the graphs, in comparison to standard solvers.  It proposes to use a GNN with attention mechanism as model, and the FUGW distance as loss function. The method exhibits a computational complexity of $\\mathcal{O}(n^2)$ instead of the $\\mathcal{O}(n^3)$ obtained by standard solvers when computing the OT plan between two graphs with n nodes. The authors evaluated the methods on both synthetic and real-world problems, confirming computational gains while maintaining comparable FUGW losses.", "questions": "1) How do the authors believe their model compares to simpler neural network models, or how does it relate to other deep learning approaches for optimal transport? If such comparisons are not provided, how should one assess the architecture's performance? It might help if the authors included comparisons or discussions in this regard.\n\n2) While the authors demonstrate various features of their model through synthetic experiments, could the authors elaborate on potential real-world applications?\n\n3) The approach of encoding the weights, specifically rho and alpha, as node features utilizing the Fourier basis would benefit from additional clarification. Could the authors offer further justification or elaboration on this methodology?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "This paper introduces a novel approach for predicting FUGW transport plans. The authors present compelling experimental results using a real-world dataset of fMRI activations. Additionally, the method is  thorough examined through a synthetic problem, offering valuable insights into its behavior. When benchmarked against standard solvers, the proposed method achieves results of comparable accuracy but with significantly reduced computational time. The paper is clearly presented. In essence, there is no major weaknesses. The paper effectively combines deep learning techniques with conclusive experimental results. However, while the paper successfully conveys the potential of using neural networks for predicting FUGW OT plans, it does not compare its approach to other neural network architectures, which makes it challenging to assess the relevance and advantages of the proposed architecture.", "was_revised": true, "final_justification": "The paper presents a technically sound approach that is relevant to the field. While additional comparisons could strengthen the empirical evaluation and improve the overall clarity of the contribution, I appreciate the authors' thoughtful and detailed responses, which address my concerns with seriousness and care. In light of this, I am maintaining my positive evaluation.", "rebuttal": "We thank you for your thoughtful review. We appreciate your questions and have done our best to answer them below.  \n\n\n> How do the authors believe their model compares to simpler neural network models, or how does it relate to other deep learning approaches for optimal transport? If such comparisons are not provided, how should one assess the architecture's performance? It might help if the authors included comparisons or discussions in this regard.\n\n\n\nTo the best of our knowledge, there are no other deep learning approaches for predicting the FUGW optimal transport plan. Existing neural OT methods (except Meta OT [2]) use a neural network that is optimized to solve a unique OT problem but cannot be used on new data without a new optimization. We originally tried other simpler designs, but their limited performances led to the current architecture. \n\nWe provide an ablation study of the impact of each step of the OT plan block in the table below. As a first ablation we remove the node scaling in the OT plan head, as a second ablation we replace the whole OT plan head by a much simpler nonlinear transformation of the node embedings to recover an OT plan with  ($P_\\theta^{(\\alpha, \\rho)})_{i,j} = \\frac{1}{1+||(F^{\\text{final}}_1)_i -(F^{\\text{final}}_2)_j||^2_2},$ where $F^{\\text{final}}_1$ and $F^{\\text{final}}_2$ are the outputs of the GCN and cross attention block for the first and second graphs respectively.\n   \n\n| Ablation |   Relative error   |  Pearson correlation |  \n| -------- | -------- | -------- | \n|  None   |   $0.19 \\pm 0.29$  |  $1.0$ | \n| Remove node scaling | $4.1 \\pm 9.5$    |  $0.89$  | \n| Replace OT plan block by nonlinearity | $27 \\pm 49$   |  $0.48$  | \n\nWe can see that node scaling is essential for recovering unbalanced OT plans especially for small $\\rho$ parameters (unbalanced plans that do not sum to one), and removing them increases the error (but remains surprisingly well correlated). Using a simpler output head for the OT plan leads to a failure to predict OT plans.\n\n\n> While the authors demonstrate various features of their model through synthetic experiments, could the authors elaborate on potential real-world applications?\n\nOur model can be applied to any scenario where there is structured data, such as genomics [A] or neuroimaging [30]. In particular, we are interested in real-world applications in neuroimaging where FUGW-based methods have proven to be effective [30] but are still limited by their computational time.\n\nMore precisely, we will use in future works ULOT to predict fMRI activations faster, in scenarios where one subject has many recorded fMRI activations and another has few. As shown in [30] we can transport the activations from the first to the latter. Moreover, we can compute the barycenters of populations of subjects in order to get a typical fMRI recording. This is useful in practice but unscalable using solvers because it requires solving many FUGW problems along the optimization process. \n\nBeyond computational speed, ULOT enables new capabilities for neuroimaging applications such as efficient FUGW hyperparameter selection, which are hard to tune in practice and highly influence the performance.\n\n\n> The approach of encoding the weights, specifically rho and alpha, as node features utilizing the Fourier basis would benefit from additional clarification. Could the authors offer further justification or elaboration on this methodology?\n \nWe designed the model so that the information about $\\rho$ and $\\alpha$ is not lost along the layers, which is a common issue when models are conditioned on a real value, such as time in diffusion models. Inspired by the method in [B], which has become a standard practice for diffusion models, we concatenate the value of rho and the encoding for alpha to the node features at the entry of each MLP, GCN and Linear layer. We found that concatenating to the node features is a simple and effective way of keeping the parameters information.  This design ensures that all layers have direct access to the conditioning parameters and can learn themselves at which layer this information is the most pertinent.\n\nBesides, parameters $\\alpha$ and $\\rho$ have very different influences on the optimal transport plan: $\\rho$ influences the amount of mass that can be discarded, while $\\alpha$ has a more subtle role by controlling the tradeoff between the Wasserstein and Gromov Wasserstein terms. They also have very different theoretical dynamics with $\\rho\\in[0,+\\infty]$ and $\\alpha \\in[0,1]$. This is why we encode $\\alpha$ in a more complex way, using a method strongly inspired by the standard positional encodings in Transformers [B]. Here, we find that using $(\\\\cos(k\\\\pi\\\\alpha))\\_{k=1,\\\\dots,d}$ and $(\\\\sin(k\\\\pi(1-\\\\alpha))\\_{k=1,\\\\dots,d}$  improves training while capturing key properties. We also encode both $\\alpha$ and $1-\\alpha$ to ensure that the dynamics are the same close to $0$ and $1$.\n\n\n\n### References\n    \n   [30] Alexis Thual, Quang Huy Tran, Tatiana Zemskova, Nicolas Courty, Rémi Flamary, Stanislas Dehaene, and Bertrand Thirion. Aligning individual brains with fused unbalanced gromov wasserstein. Advances in neural information processing systems, 35:21792–21804, 2022.\n\n   [A] Huizing, G. J., Peyré, G., & Cantini, L. (2022). Optimal transport improves cell–cell similarity inference in single-cell omics data. Bioinformatics, 38(8), 2169-2177.\n\n   [B] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.\n\n   [C] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. Advances in neural information processing systems, 33, 6840-6851.\n\n   [2] Amos, B., Cohen, S., Luise, G., & Redko, I. (2022). Meta optimal transport. arXiv preprint arXiv:2206.05262."}]

Author rebuttal:

Meta-review: This paper introduce ULOT, a deep learning method that predicts optimal transport plans between graphs using a cross-attention architecture trained on FUGW loss. It is scalable, differentiable, and allows graph alignment with applications in synthetic and brain network. The reviewers were already quite positive about this work before rebuttal, and the discussion led to a consensus on the qualities of this submission. Hence, I recommend acceptance to NeurIPS 2025.
Decision: Accept (poster)