Fused Unbalanced Gromov-Wasserstein (FUGW)

• FUGW is an optimal transport framework that fuses feature similarity and geometric structure with relaxed marginal constraints, enabling flexible alignment of cortical surfaces and graphs.
• It employs entropic regularization and unbalanced transport to accommodate variability in node masses, thereby producing sharper activation patterns and improved statistical consistency.
• The method leverages a block-coordinate descent solver with Sinkhorn-type updates, balancing computational efficiency with scalability in complex neuroimaging and graph matching tasks.

Fused Unbalanced Gromov-Wasserstein (FUGW) is an optimal-transport formulation that combines feature matching, geometry matching, relaxed marginal constraints, and, in entropic variants, smoothing for algorithmic tractability. It was introduced for inter-subject brain alignment as a method that aligns cortical surfaces based on the similarity of their functional signatures in response to a variety of stimulation settings while penalizing large deformations of individual topographic organization; the unbalanced feature allows the transport mass at each vertex to shrink or expand, accommodating functional-area size variability across subjects (Thual et al., 2022). Subsequent work has also presented FUGW as a graph-alignment loss between unbalanced graphs and as a building block for amortized prediction of transport plans, multi-marginal formulations, and transfer-operator constructions on labelled or noisy data (Mazelet et al., 21 May 2025).

1. Origin, scope, and data models

In its original neuroimaging setting, inter-subject brain alignment is posed as the search for a soft-matching between the cortical surfaces of two subjects so that functionally and anatomically corresponding regions are brought into register (Thual et al., 2022). Each brain is represented by a weighted graph

$\mathcal X=(F,D,w),$

where $F\in\mathbb R^{n\times c}$ encodes $c$ functional features at each of $n$ vertices, $D\in\mathbb R^{n\times n}_+$ is the matrix of pairwise geodesic distances on the mesh, and $w\in\Delta^{n-1}$ is the prior mass on each vertex, often uniform $1/n$. In this formulation, the alignment is a nonnegative coupling $P\in\mathbb R^{n\times p}_+$ between source and target cortical surfaces.

A closely related graph-theoretic formulation considers two graphs

$$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$

with node-feature matrices $F_k\in\mathbb R^{n_k\times d}$, intra-graph distance matrices $F\in\mathbb R^{n\times c}$0, and node-mass vectors $F\in\mathbb R^{n\times c}$1, typically uniform (Mazelet et al., 21 May 2025). This reformulation makes explicit that FUGW is not restricted to cortical meshes: it is a general quadratic OT objective for matching structured objects when both node attributes and internal geometry are informative.

The scope of the method follows from this dual representation. In the neuroimaging paper, FUGW is described as well-suited for whole-brain landmark-free alignment and as working on heterogeneous meshes, with no common template required (Thual et al., 2022). In the graph-learning paper, the same loss is used to compare synthetic stochastic block model graphs and real cortical-surface graphs obtained from fMRI (Mazelet et al., 21 May 2025). This suggests that the defining feature of FUGW is not a specific application domain, but the simultaneous treatment of feature similarity, relational structure, and mass mismatch.

2. Variational formulation

The pairwise FUGW objective combines a Wasserstein-type feature term, a Gromov-Wasserstein structural term, and unbalanced penalties on the induced marginals. In the cortical-alignment notation, the functional cost matrix is

$F\in\mathbb R^{n\times c}$2

and the fourth-order anatomical discrepancy tensor is

$F\in\mathbb R^{n\times c}$3

The Gromov-Wasserstein discrepancy is

$F\in\mathbb R^{n\times c}$4

while the fused feature term is

$F\in\mathbb R^{n\times c}$5

With trade-off parameter $F\in\mathbb R^{n\times c}$6, marginal-relaxation parameter $F\in\mathbb R^{n\times c}$7, and entropic regularizer $F\in\mathbb R^{n\times c}$8, the full objective is

$F\in\mathbb R^{n\times c}$9

where

$c$0

The KL terms relax the hard marginal constraints $c$1 and $c$2, and thereby allow the transport mass at each vertex to shrink or expand (Thual et al., 2022).

A graph-matching presentation uses the notation $c$3 and $c$4 instead of $c$5 and $c$6: $c$7 In this notation, $c$8 trades off feature-matching against structural matching, and $c$9 penalizes marginal mismatch via the Kullback-Leibler divergences (Mazelet et al., 21 May 2025).

The interpretation of the three principal components is stable across these formulations. The fused term encourages matching nodes whose attributes lie close in feature space. The GW term compares intra-graph distances and aligns local geometry without requiring the two graphs to live in the same ambient space. The unbalanced term softly enforces the marginals. As $n$0, one recovers the balanced FGW problem; as $n$1, one enforces exact marginal matching (Mazelet et al., 21 May 2025, Beier, 2023). A common misconception is that “unbalanced” means unconstrained transport. The actual formulation retains explicit KL penalties on the induced marginals; what changes is that equality is softened rather than removed.

3. Optimization and computational profile

Direct minimization of the FUGW objective is non-convex. In the cortical-alignment formulation, the proposed solver follows a Block-Coordinate Descent scheme that lower-bounds FUGW by introducing two couplings $n$2 and alternately minimizing in $n$3 and in $n$4; each subproblem becomes an entropic, unbalanced OT problem, solvable by Sinkhorn-type matrix scaling (Thual et al., 2022). The high-level iteration fixes one coupling, computes a linearized cost from the surrogate objective, solves an unbalanced Sinkhorn problem for the other coupling, rescales total mass, and alternates until $n$5.

The corresponding dual-variable updates are of Sinkhorn type: $n$6 followed by reconstruction

$n$7

Each BCD step decreases the surrogate objective, and Sinkhorn updates converge linearly under mild conditions (Thual et al., 2022).

A tutorial-style derivation in the transfer-operator paper presents the same strategy as a block-coordinate proximal scheme. One linearizes the GW term around the current coupling $n$8 to obtain a linear cost matrix $n$9, then solves a strictly convex entropic-unbalanced OT subproblem by generalized Sinkhorn updates with kernel $D\in\mathbb R^{n\times n}_+$0 and scaling vectors $D\in\mathbb R^{n\times n}_+$1 (Beier, 2023). Under mild compactness and boundedness assumptions on cost and features, the scheme converges to a stationary point of the nonconvex FUGW problem (Beier, 2023).

The computational bottleneck is the quadratic structural term. In the brain-alignment setting, complexity per Sinkhorn iteration is $D\in\mathbb R^{n\times n}_+$2, but the GW-term linearization requires $D\in\mathbb R^{n\times n}_+$3 precomputations; in practice meshes are reduced to $D\in\mathbb R^{n\times n}_+$4 vertices via spatial clustering (Thual et al., 2022). In the graph-learning setting, classical quadratic OT solvers for FGW/FUGW require $D\in\mathbb R^{n\times n}_+$5 per iteration, while the learned predictor discussed below runs in $D\in\mathbb R^{n\times n}_+$6 time and memory per forward pass (Mazelet et al., 21 May 2025). Typical settings reported for cortical alignment are $D\in\mathbb R^{n\times n}_+$7, $D\in\mathbb R^{n\times n}_+$8, and $D\in\mathbb R^{n\times n}_+$9 (Thual et al., 2022).

4. Brain alignment and empirical behavior

The reference application of FUGW is inter-subject functional alignment on the Individual Brain Charting dataset (Thual et al., 2022). The reported data comprise 12 subjects and approximately 400 fMRI contrast maps per subject covering motor, visual, auditory, language, mathematics, social cognition, and related paradigms, with splits of 326 training, 43 validation, and 30 test contrasts per subject. Mesh resolution is 10,240 vertices per hemisphere on fsaverage5, and experiments on individual anatomies additionally use meshes clustered to 10,000 nodes (Thual et al., 2022).

Evaluation focuses on between-subject Pearson correlation of unseen test contrasts before and after alignment, together with anatomical plausibility measures such as average geodesic displacement and coupling spread (Thual et al., 2022). The baselines are projection to fsaverage5 with no functional alignment, Multimodal Surface Matching (MSM), and balanced GW, implemented as the $w\in\Delta^{n-1}$0 limit (Thual et al., 2022).

On fsaverage5 meshes, the baseline correlation is $w\in\Delta^{n-1}$1. MSM yields a modest average gain of about $w\in\Delta^{n-1}$2, while FUGW with $w\in\Delta^{n-1}$3, $w\in\Delta^{n-1}$4, and $w\in\Delta^{n-1}$5 raises correlation to $w\in\Delta^{n-1}$6, corresponding to $w\in\Delta^{n-1}$7 (Thual et al., 2022). On individual anatomies, FUGW still yields an approximately $w\in\Delta^{n-1}$8 correlation gain (Thual et al., 2022). The paper further reports that FUGW barycenter maps show sharper activation patterns than simple vertex-wise averaging; one-sample $w\in\Delta^{n-1}$9-maps on aligned data reveal additional supra-threshold clusters, with peak $1/n$0-values up to $1/n$1 versus $1/n$2, and a $1/n$3–$1/n$4 increase in the count of significantly activated vertices at $1/n$5 (Thual et al., 2022).

These results are interpreted in the source paper as evidence that FUGW integrates both anatomical and functional information without requiring diffeomorphic deformations, handles area-size variability across individuals via unbalanced OT, and produces sharper group-level statistics with increased functional-correlation consistency (Thual et al., 2022). A plausible implication is that the method is most useful in settings where coarse anatomical co-registration underutilizes rich task-evoked or contrast-based functional signatures.

5. Learned plan prediction and graph-level extensions

FUGW has also been used as a training loss for amortized prediction of transport plans between unbalanced graphs. The ULOT method learns a map

$1/n$6

via a graph-attention neural network conditioned on $1/n$7, rather than solving the nonconvex optimization problem independently for each pair $1/n$8 (Mazelet et al., 21 May 2025). The architecture alternates self-updates by a small GCN on each graph, cross-attention computed from a “cross” embedding MLP$1/n$9, and a fusion MLP applied to the concatenation $P\in\mathbb R^{n\times p}_+$0. At the last layer, the network predicts two nonnegative node-volume vectors $P\in\mathbb R^{n\times p}_+$1 and $P\in\mathbb R^{n\times p}_+$2, and assembles the transport plan as

$P\in\mathbb R^{n\times p}_+$3

which is differentiable in $P\in\mathbb R^{n\times p}_+$4 and costs $P\in\mathbb R^{n\times p}_+$5 to evaluate (Mazelet et al., 21 May 2025).

Training is unsupervised: the objective minimizes the expected FUGW loss of the predicted plan over graph pairs and sampled hyperparameters, with $P\in\mathbb R^{n\times p}_+$6 and $P\in\mathbb R^{n\times p}_+$7, optimized with Adam and learning-rate scheduling (Mazelet et al., 21 May 2025). No ground-truth plans are needed. The predicted plan can also be used as a warm start for classical block-coordinate solvers such as IBPP or MM (Mazelet et al., 21 May 2025).

Empirically, on graphs of size around $P\in\mathbb R^{n\times p}_+$8, ULOT is reported to be up to $P\in\mathbb R^{n\times p}_+$9 faster than state-of-the-art unbalanced OT solvers and $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$0 faster than an entropic-regularized Sinkhorn approach, while incurring only a few percent loss increase (Mazelet et al., 21 May 2025). On 14,400 cortical graph-pairs across subjects and contrasts from the IBC dataset, ULOT’s FUGW loss matches IBPP to within $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$1 on average at a $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$2 speedup, and using ULOT as a warm start halves the number of solver iterations (Mazelet et al., 21 May 2025). The paper also reports that the network smoothly interpolates between feature-dominated matchings for $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$3 and structure-dominated matchings for $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$4, and automatically drops unmatched clusters when $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$5 is small (Mazelet et al., 21 May 2025). This makes explicit, in a learned setting, the same parameter semantics already present in the variational formulation.

6. Generalizations, related formulations, and limitations

Several papers place FUGW inside a broader family of GW-based methods. The transfer-operator work discusses fused and unbalanced variants of GW transport for labelled and noisy data, respectively, and uses transport plans to define a transfer operator $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$6 whose spectral clustering extracts coherent structures (Beier, 2023). Its numerical examples include rotating disks, interacting rotating disks, and a vorticity field of 2D Navier–Stokes; in the vorticity example, labels encode sign of vorticity and fused unbalanced GW with $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$7 is used, with recursive spectral clustering into $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$8 regions yielding nested coherent vortices (Beier, 2023). In this context, FUGW inherits the GW property of handling isometries such as rotations while adding label consistency and robustness to noise or mass mismatch.

A different extension is multi-marginal. The multi-marginal GW paper introduces entropic, unbalanced, fused multi-marginal costs on a family of labelled metric-measure spaces $$G_1=(F_1,D_1,\omega_1), \qquad G_2=(F_2,D_2,\omega_2),$$9, with structure cost $F_k\in\mathbb R^{n_k\times d}$0, label cost $F_k\in\mathbb R^{n_k\times d}$1, and fusion parameter $F_k\in\mathbb R^{n_k\times d}$2 (Beier et al., 2022). Its entropic unbalanced fused objective

$F_k\in\mathbb R^{n_k\times d}$3

adds Csiszár-divergence penalties on the marginals and an entropic KL penalty on the multi-marginal plan, and is solved by alternating minimization in two plans $F_k\in\mathbb R^{n_k\times d}$4, where each subproblem is a tree-structured multi-marginal unbalanced Sinkhorn solve (Beier et al., 2022). The paper relates this multi-marginal construction to unbalanced and fused GW barycenters and reports fused-barycenter examples on labelled digits, camel shapes, and noisy rotating particle clouds (Beier et al., 2022).

Related work on Fused Partial Gromov-Wasserstein distinguishes a partial-transport relaxation from the KL-based unbalanced relaxation used by FUGW. In the notation of that paper, a general fused unbalanced GW objective can be written with total-variation penalties on the marginal product measures, and the resulting partial formulation is evaluated for graph classification and clustering under outlier corruption (Bai et al., 14 Feb 2025). The reported comparison states that FPGW demonstrates markedly greater robustness to spurious nodes and unequal masses than FGW or FUGW, at comparable runtime (Bai et al., 14 Feb 2025). This does not negate the role of FUGW; rather, it identifies a separate relaxation mechanism—partial transport instead of KL-penalized marginal mismatch—for settings dominated by outlier mass.

The principal limitations reported for FUGW itself are algorithmic and hyperparameter-related. In brain alignment, the entropic smoothing parameter $F_k\in\mathbb R^{n_k\times d}$5 strongly affects coupling blur, selecting $F_k\in\mathbb R^{n_k\times d}$6 may require expert tuning, high-resolution meshes demand mesh-reduction strategies or GPU-memory optimizations, and the current solver is relatively slow, on the order of minutes per hemisphere on a GPU (Thual et al., 2022). The same paper suggests that majorization–minimization OT solvers could accelerate computation and yield sparser couplings, and identifies future work on learning an anatomical template within the barycenter loop and applying FUGW to cross-species alignment to identify uniquely human brain regions (Thual et al., 2022). More generally, practical guidance from the graph-learning formulation states that $F_k\in\mathbb R^{n_k\times d}$7 near $F_k\in\mathbb R^{n_k\times d}$8 emphasizes feature similarity, $F_k\in\mathbb R^{n_k\times d}$9 near $F\in\mathbb R^{n\times c}$00 enforces geometric consistency, intermediate $F\in\mathbb R^{n\times c}$01 often balances the two, small $F\in\mathbb R^{n\times c}$02 promotes discarding unmatched nodes, and large $F\in\mathbb R^{n\times c}$03 recovers a nearly balanced matching (Mazelet et al., 21 May 2025).