Aligning individual brains with Fused Unbalanced Gromov-Wasserstein
Anonymous Author(s)
AffiliationAddress email
Abstract
Individual brains vary in both anatomy and functional organization, even within agiven species.
Inter-individual variability is a major impediment when trying todraw generalizable conclusions from neuroimaging data collected on groups ofsubjects.
Current co-registration procedures rely on limited data, and thus lead tovery coarse inter-subject alignments.
In this work, we present a novel method forinter-subject alignment based on Optimal Transport, denoted as Fused Unbalanced
Gromov Wasserstein (FUGW).
The method aligns cortical surfaces based onthe similarity of their functional signatures in response to a variety of stimulationsettings, while penalizing large deformations of individual topographic organization.
We demonstrate that FUGW is well-suited for whole-brain landmark-free alignment.
The unbalanced feature allows to deal with the fact that functional areas vary insize across subjects.
Our results show that FUGW alignment significantly increasesbetween-subject correlation of activity for independent functional data, and leadsto more precise mapping at the group level.
Introduction
The availability of millimeter or sub-millimeter anatomical or functional brain images has openednew horizons to neuroscience, namely that of mapping cognition in the human brain and detectingmarkers of diseases.
Yet this endeavour has stumbled on the roadblock of inter-individual variability:while the overall organization of the human brain is largely invariant, two different brains (evenfrom monozygotic twins [32]) may differ at the scale of centimeters in shape, folding pattern andfunctional responses.
The problem is further complicated by the fact that functional images arenoisy, due to imaging limitations and behavioral differences across individuals that cannot be easilyovercome.
The status quo of the field is thus to rely on anatomy-based inter-individual alignmentthat approximately matches the outline of the brain [4] as well as its large-scale cortical foldingpatterns [12, 14].
Existing algorithms thus coarsely match anatomical features with diffeomorphictransformations, by warping individual data to a simplified template brain.
Such methods lose muchof the original individual detail and blur the functional information that can be measured in brainregions (see Figure 1).
In order to improve upon the current situation, a number of challenges have to be addressed: (i) Thereexists no template brain with functional information, which by construction renders any corticalmatching method blind to function.
This is unfortunate, since functional information is arguably themost accessible marker to identify cortical regions and their boundaries [17].
(ii) When comparingtwo brains – coming from individuals or from a template – it is unclear what regularity should beimposed on the matching [41].
While it is traditional in medical imaging to impose diffeomorphicity
[Figure]
- subsub- average of all subjects z-score i n d i v i d u a l f s a v e r a g e m e s h
In this experiment contrasting areas of the brain which respond to mathematical tasks against other that don’t, we observe great variability in locations and strength of brain activations across subjects (row 1).
The classical approach consists in wrapping this data to a common surface template (row 2), where they can be averaged, often resulting in loss of individual details and detection power.
These images were generated using Nilearn software [1].
[4], such a constrain does not match the frequent observation that brain regions vary across individualsin their fine-grained functional organization [17, 37].
(iii) Beyond the problem of aligning humanbrains, it is an even greater challenge to systematically compare functional brain organization in twodifferent species, such as humans and macaques [25, 22, 45, 13].
Such inter-species comparisonsintroduce a more extreme form of variability in the correspondence model.
Related work Several attempts have been made to constrain the brain alignment process by usingfunctional information.
The first one consists in introducing functional maps into the diffeomorphicframework and search for a smooth transformation that matches functional information [36, 46, 34],the most popular framework being arguably Multimodal Surface Matching (MSM) [34, 17].
A second family of less constrained functional alignment approaches have been proposed, basedon heuristics, by matching information in small, possibly overlapping, cortical patches [20, 39, 6].
This popular framework has been called hyperalignment [20, 19], or shared response models [10].
Yet these approaches lack a principled framework and cannot be considered to solve the matchingproblem at scale.
Neither do they allow to estimate a group-level template properly [44].
An alternative functional alignment framework has followed another path [18], considering functionalsignal as a three-dimensional distribution, and minimizing the transport cost.
However, this frameworkimposes unnatural constraints of non-negativity of the signal and only works for one-dimensionalcontrasts, so that it cannot be used to learn multi-dimensional anatomo-functional structures.
Animportant limitation of the latter two families of methods is that they operate on a fixed spatialcontext (mesh or voxel grid), and thus cannot be used on heterogeneous meshes such as between twoindividual human anatomies or, worse, between a monkey brain and a human brain.
Contributions Following [5], we use the Wasserstein distance between source and target functionalsignals – consisting of contrast maps acquired with fMRI – to compute brain alignments.
Wecontribute two notable extensions of this framework: (i) a Gromov-Wasserstein (GW) term topreserve global anatomical structure – this term introduces an anatomical penalization againstimprobably distant anatomical matches, yet without imposing diffeomorphic regularity – as well as(ii) an unbalanced correspondence that allows mappings from one brain to another to be incomplete,for instance because some functional areas are larger in some individuals than in others, or maysimply be absent.
We show that this approach successfully addresses the challenging case of differentcortical meshes, and that derived brain activity templates are sharper than those obtained with standardanatomical alignment approaches.
Methods
Optimal Transport yields a natural framework to address the alignment problem, as it seeks to derivea plan – a coupling – that can be seen as a soft assignment matrix between cortical areas of a sourceand target individual.
As discussed previously, there is a need for a functional alignment method thatrespects the rich geometric structure of the anatomical features, hence the Wasserstein distance aloneis not sufficient.
By construction, the GW distance [23, 24] can help preserve the global geometryunderlying the signal.
The more recent fused GW distance [43] goes one step further by making itpossible to integrate functional data simultaneously with anatomical information.
2.1 Fused Unbalanced Gromov-Wasserstein
We leverage [43, 38] to present a new objective function which interpolates between a loss preservingthe global geometry of the underlying mesh structure and a loss aligning source and target features,while simultaneously allowing not to transport some parts of the source and target distributions.
Weprovide an open-source solver that minimizes this loss 1 .
Formulation We denote F s ∈ R n,c the matrix of features per vertex for the source subject.
In theproposed application, they correspond to c functional activation maps, sampled on a mesh with nvertices representing the source subject’s cortical surface.
Let D s ∈ R n,n + be the matrix of pairwisegeodesic distances 2 between vertices of the source mesh.
Moreover, we assign the distributionw s ∈ R n + on the source vertices.
Comparably, we define F t ∈ R p,c , D t ∈ R p,p + and w t ∈ R p + for thetarget subject, whose individual anatomy is represented by a mesh comprising p vertices.
Eventually,w s and w t set the transportable mass per vertex, which, without prior knowledge, we choose to beuniform for the source and target vertices respectively: w s ≜ ( 1 n , ..., 1 n ) , w t ≜ ( 1 p , ..., 1 p ) .
Given a tuple of hyper-parameters θ ≜ ( ρ, α, ε ) , where ρ, ε ∈ R + and α ∈ [0 , 1] , for any coupling
P ∈ R n,p , we define the fused unbalanced Gromov-Wasserstein loss as
L θ ( P ) ≜ (1 − α ) (cid:80) 0 ≤ i<n 0 ≤ j<p ||
F s i − F t
j || 22 P
i,j + α (cid:80) 0 ≤ i,k<n 0 ≤ j,l<p | D s i,k − D t j,l | 2 P
i,j P k,l + ρ (cid:0) KL ( P #1 ⊗ P #1 | w s ⊗ w s ) + KL ( P # [Equation]
⊗ P #2 | w t ⊗ w t ) (cid:1) + [Equation]
E ( P ) (1)
Wasserstein loss L W ( P )
Gromov-Wasserstein loss L GW ( P )
Marginal constraints L U ( P )
Entropywhere L W ( P ) matches vertices with similar features, L GW ( P ) penalizes changes in geometry and
L U ( P ) fosters matching all parts of the source and target distributions.
Throughout this paper,we refer to relaxing the hard marginal constraints of the underlying OT problem into soft onesas unbalancing .
Here, P #1 ≜ ( (cid:80)
j P i,j ) 0 ≤ i<n denotes the first marginal distribution of P , and
P #2 ≜ ( (cid:80) i P i,j ) 0 ≤ j<p the second marginal distribution of P .
The notation ⊗ represents the
Kronecker product between two vectors or two matrices.
KL ( ·|· ) denotes the Kullback Leiblerdivergence, which is a typical choice to measure the discrepancy between two measures in the contextof unbalanced optimal transport [21].
The last term E ( P ) ≜ KL (cid:0) P ⊗ P | ( w s ⊗ w t ) ⊗ ( w s ⊗ w t ) (cid:1)is mainly introduced for computational purposes, as it helps accelerate the approximation scheme ofthe optimisation problem.
Typically, it is used in combination with a small value of ε , so that theimpact of other terms is not diluted.
On the other hand, the parameters α and ρ offer control over two other aspects of the problem: while α realizes a trade-off between the impact of different features and different geometries in the resulting alignment, ρ controls the amount of mass transported by penalizing configurations such that the marginal distributions of the transportation plan P are far from the prior weights w s and w t .
This potentially helps adapting the size of areas where either the signal or the geometry differs too much between source and target.
[Figure]
transported particular, less
Eventually, we define X s ≜ ( F s , D s , w s ) and X t ≜ ( F t , D t , w t ) , and seek to derive an optimalcoupling P ∈ R n,p minimizing
[Equation]
This can be seen as a natural combination of the fused GW [43] and the unbalanced GW [38] distances.
To the best of our knowledge, it has never been considered in the literature.
Following [38], weapproximate FUGW via a lower bound. It involves solving a minimization problem with respect totwo independent couplings: Using a Block-Coordinate Descent (BCD) scheme, we fix a coupling andminimize with respect to the other. This allows us to always be dealing with linear problems insteadof a quadratic one. Eventually, each BCD iteration consists in alternatively solving two entropicunbalanced OT problems, whose solutions can be approximated using the scaling algorithm
Details concerning the lower bound as well as the corresponding BCD iteration can be found in the Appendix (see Alg. S1).
Toy example illustrating the unbalancing property As exemplified in Figure 1, brain responseselicited by the same stimulus vary greatly between individuals.
Figure 2 illustrates a similar yetsimplified version of this problem, where the goal is to align two different signals supported on thesame spherical meshes.
In this example, for each of the n = p = 3200 vertices, the feature is simplya scalar.
On the source mesh, the signal is constituted of two von Mises density functions that differby their concentration (large and small), while on the target mesh, only the large one is present, butat a different location.
We use the optimal coupling matrix P obtained from Eq.
2 to transport thesource signal on the target mesh.
As shown in Figure 2.B, the parameter ρ allows to control the masstransferred from source to target.
When ρ = 100 , we approach the solution of the fused GW problem.
Consequently, we observe the second mode on the target when transporting the source signal.
Whenthe mass control is weaker ( ρ = 1 ), the smaller blob is partly removed because it has no counterpartin the target configuration, making the transport ill-posed.
Barycenters Barycenters represent common patterns across samples.
Their role is instrumental inidentifying a unique target for aligning a given group of individuals.
As seen in Fig.
1, the vertex-wisegroup average does not usually provide well-contrasted maps.
Inspired by the success of the GWdistance when estimating the barycenter of structured objects [29, 43], we use FUGW to find thebarycenter ( F B , D B )
∈
R k,c
× R k,k of all subjects s ∈ S , as well as the corresponding couplings P s,B from each subject to the barycenter.
More precisely, we solve
[Equation]
where we set the weights w B to be the uniform distribution.
By construction, the resulting barycenterbenefits from the advantages of FUGW, i.e. equilibrium between geometry-preserving and feature- matching properties, while not forcing hard marginal constraints.
The FUGW barycenter is estimatedusing a Block-Coordinate Descent (BCD) algorithm that consists in alternatively (i) minimizing the
OT plans P s,B for each FUGW computation in (3) with fixed X B and (ii) updating the barycenter
X B through a closed form with fixed P s,B .
See Alg. S4 for more details.
Numerical experiments
We design three experiments to assess the performance of FUGW.
In Experiments 1 and 2, we areinterested in assessing if aligning pairs of individuals with FUGW increases correlation betweensubjects compared to a baseline correlation.
We also compare the ensuing gains with those obtainedwhen using the competing method MSM [34, 35] to align subjects.
In Experiment 3, we derive abarycenter of individuals and assess its ability to capture fine-grained details compared to classicalmethods.
Dataset In all three experiments, we leverage data from the Individual Brain Charting dataset [30].
It is a longitudinal study on 12 human subjects, comprising 400 fMRI maps per subject collected ona wide variety of stimuli (motor, visual, auditory, theory of mind, language, mathematics, emotions,and more), movie-watching data, T1-weighted maps, as well as other features such as retinotopywhich we don’t use in this work.
We leverage these 400 fMRI maps.
The training, validation and testsets respectively comprise 326, 43 and 30 contrast maps acquired for each individual of the dataset.
Tasks and MRI sessions differ between each of the sets.
More details, including preprocessing, areprovided in Supplementary Materials.
Baseline alignment correlation For each pair of individuals ( s, t ) under study, and for each fMRIcontrast c in the test set, we compute the Pearson correlation corr ( F s · ,c , F t · ,c ) after these maps havebeen projected onto a common surface anatomy (in this case, fsaverage5 mesh).
Throughout thiswork, such computations are made for each hemisphere separately.
Experiment 1 - Aligning pairs of humans with the same anatomy For each pair ( s, t ) under study
,we derive an alignment P s,t ∈ R n × p using FUGW on a set of training features.
In this experiment,source and target data lie on the same anatomical mesh ( fsaverage5 ), and n = p = 10 240 for eachhemisphere.
Since each hemisphere’s mesh is connected, we align one hemisphere at a time.
Computed couplings are used to align contrast maps of a the validation set from the source subjectonto the target subject.
Indeed, one can define ϕ s → t : X ∈ R n × q (cid:55)→ (cid:0) ( P s,t ) T X (cid:1) ⊘ P s,t #2 ∈ R p × qwhere ⊘ represents the element-wise division.
ϕ s → t transports any matrix of features from the sourcemesh to the target mesh.
We measure the Pearson correlation corr (cid:0) ϕ s → t ( F s ) , F t (cid:1) between eachaligned source and target maps.
We run a similar experiment for MSM and compute the correlation gain induced on a test set by
FUGW and MSM respectively.
For both models, we selected the hyper-parameters maximizingcorrelation gain on a validation set.
In the case of FUGW, in addition to gains in correlation,hyper-parameter selection was influenced by three other metrics that help us assess the relevance ofcomputed couplings:
Transported mass For each vertex i of the source subject, we compute {+(cid:80) 
+}
[-(cid:80) 
-]
≤ j<p P s,t i,j
Vertex displacement Taking advantage of the fact that the source and target anatomies are the same,we define D
= D s = D t and compute for each vertex i of the source subject the quantity(cid:80)
j P s,t i,j · D i,j / (cid:80)
j P s,t i,j , which measures the average geodesic distance on the corticalsheet between vertex i and the vertices of the target it has been matched with
Vertex spread Large values of ε increase the entropy of derived couplings.
To quantify this effect,and because we don’t want the matching to be too blurry, we assess how much a vertex wasspread .
Considering ˜ P
i
=
P s,t i / (cid:80) j P s,t i,j ∈ R p as a probability measure on target vertices, we estimate the anatomical variance of this measure by sampling q pairs ( j q , k q ) of ˜ P
i andcomputing their average geodesic distance 1 q (cid:80) j q ,k q
D j q ,k q
Experiment 2 - Aligning pairs of humans with individual anatomies We perform a second alignment experiment, this time using individual meshes instead of an anatomical template.
Importantly,in this case, there is no possibility to compare FUGW with baseline methods, since those cannothandle this case.
However, individual meshes are significantly larger than the common anatomical template used in
Experiment 1 ( n ≈ m ≈ 160k vs. 10k previously), resulting in couplings too large to fit on GPUs– for reference, a coupling of size 10k × 10k already weights 400Mo on disk.
We thus reduce thesize of the source and target data by clustering them into 10k small connected clusters using Ward’salgorithm [40].
More details are given in supplementary section A.4.
Experiment 3 - Comparing FUGW barycenters with usual group analysis Since it is verydifficult to estimate the barycentric mesh, we force it to be equal to the fsaverage5 template. Empirically, this we force the distance matrix D B to be equal to that of fsaverage5 , and only estimate thefunctional barycenter F B .
We initialize it with the mean of ( F s ) s ∈ S and derive F B and ( P s,B ) s ∈ Sfrom problem 3.
Then, for a given stimulus c , we compute its projection onto the barycenter for each subject.
We usethese projections to compute two maps of interest: (i) M B,c
the mean of projected contrast mapsacross subjects and (ii) T B,c the t-statistic (for each vertex) of projected maps.
We compare thesetwo maps with their unaligned counterparts M 0 ,c and T 0 ,c respectively.
[Equation]
The first map helps us to qualitatively evaluate the precision of FUGW alignments and barycenter.
The second one is classically used to infer the existence of areas of the brain that respond to specificstimuli.
We assess whether FUGW helps find the same clusters of vertices.
Eventually, we quantifythe number of vertices significantly activated or deactivated with and without alignment respectively.
Results
4.1 Experiment 1 - Template anatomy
Aligning subjects on a fixed mesh We set α = 0 .
5 , ρ = 1 and ε = 10 − 3 .
Pearson correlationbetween source and target contrast maps is systematically and significantly increased when alignedusing FUGW, as illustrated in Figure 3 where correlation grows by almost 40% from 0 .
258 to 0 .
We also varied training sets by selecting subsets of training contrasts and find that similar performanceon the test set can be achieved regardless of the training data (see Supplementary section A.5 and inparticular Supplementary Table S1).
Hyper-parameters selection Hyper-parameters used to obtain these results were chosen afterrunning a grid search on α , ε and ρ and evaluating it on the validation dataset.
Computation tookabout 100 hours using 4 Tesla V100-DGXS-32GB GPUs. More precisely, it takes about 4 minutes tocompute one coupling between a source and target 10k-vertex hemisphere on a single GPU, whenthe solver was set to run 10 BCD and 400 Sinkhorn iterations. In comparison, MSM takes about thesame time on Intel(R) Xeon(R) CPU E5-2698 v4 @
2.20GHz CPUs.
Results are reported in Figure 4and provide multiple insights concerning FUGW.
Firstly, without anatomical constraint ( α = 0 ), source vertices can be matched with target vertices thatare arbitrarily far on the cortical sheet.
Even though this can significantly increase correlation, it alsoresults in very high vertex displacement values (up to 100 mm ).
Such couplings are not anatomicallyplausible.
Secondly, without functional information ( α = 1 ), couplings recover a nearly flawlessmatching between source and target meshes, so that, when ε = 10 − 5 (ie when we force couplingsto find single-vertex-to-single-vertex matches), vertex displacement and spread are close to 0 andcorrelation is unchanged.
Fusing both constraints ( 0 < α < 1 ) yields the largest gains in correlationwhile allowing to compute anatomically plausible reorganizations the cortical sheet between subjects.
The impact of ρ (controlling marginal penalizations) on correlation seems modest, with a slighttendency of increased correlation in unbalanced problems (low ρ ).
Finally, it is worth noting that a relatively wide range of α and ρ yield comparable gains.
The fact that
FUGW performance is weakly sensitive to hyper-parameters makes it a good off-the-shelf tool forneuroscientists who wish to derive inter-individual alignments.
However, ε is of dramatic importancein computed results and should be chosen carefully.
Vertex spread is a useful metric to choosesensible values of ε ; for human data one might consider that it should not exceed 20 mm .
Mass redistribution in unbalanced couplings Unbalanced couplings provide additional information about how functional areas might differ in size between pairs of individuals.
This is illustrated in
Figure 5, where we observe variation in size of the auditory area between a given pair of individuals.
This feature is indeed captured by the difference of mass between subjects (although the displayedcontrast was not part of the training set).
4.2 Experiment 2 - Individual anatomies
is significantly better after alignment on individual anatomies than after projecting subjects onto a common anatomical template
As shown in Figure 6, we obtain correlationgains which are comparable to that of Experiment 1 (about 35% gain) while working on individual meshes.
This tends to show that FUGWcan compute meaningful alignments betweenpairs of individuals without the use of an anatomical template, which helps bridge most conceptual impediments listed in Section 1.
Moreover, this opens the way for computationof simple statistics in cohorts of individuals inthe absence of a template.
Indeed, one can pickan individual of the cohort and use it as a reference subject on which to transport all other individuals.
We give an example in Figure S4, showing that FUGW correctly preserved idiosyncrasies of eachsubject while transporting their functional signal in an anatomically sound way.
[Figure]
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Figure]
We show a contrast map from the test set which displays areas showing stronger activation during auditory tasks versus equivalent visual tasks.
It shows much more anterior activations on the target subject compared to the source subject.
This is consistent with the observation that more mass is present in anterior auditory areas of the source subject than in the target subject (Panel B).
4.3 Experiment 3 - Barycenter
In the absence of a proper metric to quantify the correctness of a barycenter, we first qualitativelycompare the functional templates obtained with and without alignment.
In Figure 7.A, we do sousing brain maps taken from the test set.
We can see that the barycenter obtained with FUGW yieldssharper contrasts and more fine-grained details than the barycenter obtained by per-vertex averaging.
We also display in Figure 7.B the result of a one-sample test for the same contrast, which can readilybe used for inference.
The one-sample test map obtained after alignment to the FUGW templateexhibits the same supra-threshold clusters as the original approach, but also some additional spotswhich were likely lost due to inter-subject variability in the fsaverage5 space.
This approach is thusvery useful to increase power in group inference.
We quantify this result by counting the number
A mean of all subjects mean of all subjects aligned to barycenter group analysis on all subjects group analysis on all subjects aligned to barycenter
B
C t-statistic z-score 0.
of supra-threshold vertices with and without alignment for each contrast map of the test set.
Ouralignment method significantly finds more such vertices of interest, as shown in Figure 7.C.
Discussion
FUGW can derive meaningful couplings between pairs of subjects without the need of a pre-existinganatomical template.
It is well-suited to computing barycenters of individuals, even for small cohorts.
In addition, we have shown clear evidence that FUGW yields gains that cannot be achieved bytraditional diffeomorphic registration methods.
These methods impose very strong constraints tothe displacement field, that may prevent reaching optimal configurations.
More deeply, this findingsuggests that brain comparison ultimately requires lifting hard regularity constraints on the alignmentmodels, and that two human brains differ by more than a simple continuous surface deformation.
However, current results have not shown a strong correlation gain of unbalanced OT compared tobalanced OT, likely because the cohort under study is too small.
Leveraging datasets such as HCP[42] with a larger number of subjects will help lower the standard error on correlation gain estimates.
In this work, we decided to rely on a predefined anatomical template ( fsaverage5 ) to derive functionalbarycenters.
It would be interesting to investigate whether more representative anatomical templatescan be learned during the process.
This would in particular help to customize templates to differentpopulations or species.
Additionally, using an entropic solver introduces a new hyper-parameter εthat has a strong effect, but is hard to interpret.
Future work may replace the scaling algorithm
[11]used here by the majorization-minimization one [9], which does not require entropic smoothing.
Thissolution can yield sparse couplings while being orders of magnitude faster, which will prove usefulwhen computing barycenters on large cohorts.
Finally, we plan to make use of FUGW to derive alignments between human and non-humanprimates without anatomical priors.
Indeed, the understanding of given brain mechanisms will benefitfrom more detailed invasive measurements made on other species only if brains can be matchedacross species; moreover, this raises the question of features that make the human brain unique, byidentifying patterns that have no counterpart in other species.
By maximizing the functional alignmentbetween areas, but also allowing for some regions to be massively shrunk or downright absent in onespecies relative to the other, the present tool could shed an objective light on the important issue ofwhether and how the language-related areas of the human cortical sheet map onto the architecture ofnon-human primate brains.