JOSA: Joint surface-based registration with atlas construction enables accurate alignment of the brain geometry and function

Abstract

Surface-based cortical registration is an important topic in medical image analysis and facilitates many downstream applications. Current approaches for cortical registration are mainly driven by geometric features, such as sulcal depth and curvature, and often assume that registration of folding patterns leads to alignment of brain function. However, functional variability of anatomically corresponding areas across subjects has been widely reported, particularly in higher-order cognitive areas. In this work, we present JOSA, a novel cortical registration framework that jointly models the mismatch between geometry and function while simultaneously learning an unbiased population-specific atlas. Using a semi-supervised training strategy, JOSA achieves superior registration performance in both geometry and function without requiring functional data at inference. This learning framework can be extended to any auxiliary data to guide spherical registration that is available during training but is difficult or impossible to obtain during inference, such as parcellations, architectonic identity, transcriptomic information, and molecular profiles.

1. Introduction

Image registration is a key research topic in medical image analysis. Deformable image registration establishes spatial correspondence between a pair of images via a nonlinear spatial transformation. This transformation is often obtained through an optimization procedure that maximizes a similarity measure. Approaches for registration of the brain are either directly based on 3D images from scans of the brain (volume-based) or based on features resampled onto the brain surface (surface-based) (Maintz and Viergever, 1998) or a combination of these (Postelnicu et al., 2009; Joshi et al., 2009). Surface-based registration of the human cerebral cortex, often referred to as cortical registration, extracts a variety of representative information from the brain images and solves the registration problem based on a surface matching framework.

Cortical registration methods achieve high accuracy in aligning the complex folding patterns of the cerebral cortex by matching the geometric features of the brain (Davatzikos and Bryan, 1996; Fischl et al., 1999b). It has been shown that inter-subject cortical registration can improve not only the statistical power of group functional analyses (van Atteveldt et al., 2004; Frost and Goebel, 2012), but also the predictability of cytoarchitecture of the brain (Yeo et al., 2010b; Fischl et al., 2008). Due to the geometric complexity of the cortex and the large variability among individuals, cortical registration remains a challenging task. Inter-subject surface alignment is commonly driven by geometric features that describe measures of cortical folding patterns, such as sulcal depth and mean curvature (Fischl et al., 1999b; Yeo et al., 2010a; Conroy et al., 2013). It is often assumed that an accurate registration of folding patterns will also successfully align brain function. However, the functional variability of anatomically corresponding areas across subjects has been widely reported, particularly in regions implicated in higher-order cognition (Fischl et al., 2008; Frost and Goebel, 2012). This variability implies that regions with different functional specializations may not be optimally aligned even when a perfect anatomical registration is achieved. There may exist many equally-good solutions for geometric registration with substantially different performance in registration of function. Another possible cause is the mismatch between geometry and function within each subject, i.e., different individuals may use a slightly different region to process the same function.

Traditional model-based deformable registration has been extensively studied (Fischl et al., 1999b; Yeo et al., 2010a; Sabuncu et al., 2010; Guntupalli et al., 2016; Robinson et al., 2014; Vercauteren et al., 2009; Nenning et al., 2017; Avants et al., 2008; Beg et al., 2005; Christensen et al., 1997; Van Essen et al., 1998). Typical strategies employ an iterative approach that seeks an optimal deformation field to warp a moving image to a fixed image. Methods usually involve optimization of a similarity measure between two feature maps, e.g., minimizing mean squared error (MSE) or maximizing normalized cross correlation, while regularizing the deformation field to have some desired property, such as smoothness and/or diffeomorphism. Widely used techniques for cortical surface registration map the surface onto the unit sphere and establish correspondence between feature maps in the spherical space (Fischl et al., 1999a). Conventional approaches, such as FreeSurfer (Fischl et al., 1999b), register an individual subject to a probabilistic population atlas by minimizing the geometry MSE weighted by the inverse variance of the atlas convexity, in a maximum a posteriori formulation. These anatomical registration methods have been adapted to functional registration by minimizing MSE on functional connectivity computed from functional magnetic resonance imaging (fMRI) data (Sabuncu et al., 2010). Spatial correspondence can also be maximized in a non-diffeomorphic manner by finding local orthogonal transforms that linearly combine features around each local neighborhood (Guntupalli et al., 2016, 2018). Recent discrete optimization approaches iteratively align local features using spherical meshes from low-resolution to high-resolution (Robinson et al., 2014, 2018). To encourage invertibility of the deformation field and preserve anatomical topology, diffeomorphism has been introduced into registration using an exponentiated Lie algebra, most often assuming a stationary velocity field (SVF) (Ashburner, 2007; Vercauteren et al., 2009). These strategies were extended to the sphere by regularizing the deformation using spherical thin plate spline interpolation (Yeo et al., 2010a). Several methods align functional regions, for example using Laplacian eigen embeddings computed from fMRI data (Nenning et al., 2017). These methods are successful but solve an optimization problem for each image pair, resulting in a high computational cost.

With the increased of popularity of deep learning techniques, many learning-based registration methods have been proposed. They can be categorized into supervised registration (Krebs et al., 2017; Sokooti et al., 2017; Yang et al., 2016; Cao et al., 2017) and unsupervised registration (Dalca et al., 2018a; Balakrishnan et al., 2019; Cheng et al., 2020a; Niethammer et al., 2019; Krebs et al., 2019; de Vos et al., 2019). Supervised registration aims to predict the deformation fields directly and evaluate the performance by comparing the predicted deformation with a “ground truth” deformation. However, due to the lack of the “ground truth”, deformation generated by iterative methods or simulation is often used as the target, which fundamentally limits the performance of supervised approaches (Sokooti et al., 2017; Yang et al., 2016; Cao et al., 2017). Further, the iterative registration methods are relatively slow which makes it impractical to generate large numbers of samples for supervised training, often resulting in overfitting. In contrast, Unsupervised registration methods predict a deformation field that is used to warp the moving image to the fixed image. The core idea of these methods is to employ a classical loss in the image space, thus forming an end-to-end training pipeline (de Vos et al., 2019; Balakrishnan et al., 2019). Semi-supervised methods employ additional information, such as segmentation or parcellation maps, to guide registration without requiring them during inference (Balakrishnan et al., 2019). Recent methods extend this strategy to the spherical domain by parameterizing the brain surfaces in a 2D grid that accounts for distortions (Cheng et al., 2020a) or directly on the sphere using spherical kernels (Zhao et al., 2021). These learning-based registration methods substantially improved registration speed by orders of magnitude at inference while achieving a superior or comparable registration accuracy relative to the iterative methods. However, these methods do not directly model structure-function mismatches within a subject.

In this work, we present a diffeomorphic cortical registration framework that 1) explicitly uses functional data to drive the geometric registration to optimally align function, and 2) explicitly models mismatches between geometry and function. We build on recent unsupervised spherical registration strategies (Balakrishnan et al., 2019; Cheng et al., 2020a) and use a joint deformation field shared by geometry and function to capture the relatively large differences among subjects. We introduce deformation fields that describe relatively small variations between geometry and function within each subject. To avoid the biases of existing anatomical templates, we simultaneously learn a population-specific atlas during training (Dalca et al., 2019). We develop a semi-supervised training strategy that uses task fMRI data to improve functional registration but without requiring task fMRI data during inference. However, in contrast to the term “unsupervised” commonly used in the literature (Balakrishnan et al., 2019; Cheng et al., 2020a), here we borrow the term “semi-supervised” to describe our strategy to highlight that auxiliary information, such as functional data, can be incorporated during training but are not required during inference. This semi-supervised training framework can also be easily extended to any auxiliary data that could be helpful to guide spherical registration but is difficult or impossible to obtain or undesirable to use during inference, such as parcellations, architectonic identity, transcriptomic information, and molecular profiles. Therefore, the registration at inference is solely based on the geometry of the subject surface, which further avoids the potential circularity issue in any subsequent analyses of functional or other auxiliary data. We demonstrate experimentally that the proposed framework yields improved registration performance in both anatomical and functional domains.

Some of the preliminary results have been presented in (Li et al., 2023). The current paper provides a more detailed description of the method, novel experimental results with a comparison to an extensive list of the state-of-the-art methods, as well as detailed ablation studies.

2. Methods

We describe Joint Spherical registration and Atlas building (JOSA), a method for surface-based cortical registration with simultaneous atlas construction that explicitly models the cross-subject variation in the relationship between anatomy and functional properties in each subject.

2.1. Generative model

Fig. 1 shows the graphical representation for the proposed generative model. Let A be an unknown population atlas with all geometric and functional cortical features of interest. We propose a generative model that describes the formation of the subject geometric I_g and functional I_f features by first warping the atlas A by a subject deformation field ϕ_j. This model characterizes the differences between subjects and results in a joint multi-feature image I_j. Geometric feature I_g is formed given an additional field ϕ_g that deforms the geometric features in I_j, and similarly for I_f and ϕ_f.

Figure 1:

Graphical representation of the generative model. Circles are random variables. Rounded squares indicate parameters. Shaded quantities are observations. The big plate represents replication. A represents the global atlas, I the input image, ϕ deformation field. The subscript j, g, and f stand for joint, geometry, and function, respectively.

2.1.1. Deformation priors

Let ϕⁱ_j, ϕⁱ_g, ϕⁱ_f be the joint, geometric, and functional deformation fields for each subject i, respectively. All variables in the model are subject-specific, except for the global atlas A, and we omit i for our derivation. We impose the deformation priors

$$\begin{array}{l}p\left({\varphi }_j\right)~\text{exp}\left\{-\left({\lambda }_j\parallel \nabla u_j{\parallel }^2+{\alpha }_j\parallel {\overline{u}}_j{\parallel }^2\right)\right\}\hfill \\ p\left({\varphi }_g\right)~\text{exp}\left\{-{\lambda }_g\parallel \nabla u_g{\parallel }^2\right\}\hfill \\ p\left({\varphi }_f\right)~\text{exp}\left\{-{\lambda }_f\parallel \nabla u_f{\parallel }^2\right\}\hfill \end{array}$$ (1)

where u_j is the spatial displacement for ϕ_j = Id + u_j, ∇u_j is its spatial gradient, and ū_j = 1/N i uⁱ_j, N is the number of subjects, and similarly for u_g and u_f. The gradient term encourages smooth deformations, while the mean term encourages an unbiased atlas A by penalizing the average deformation over the entire dataset and hence favouring atlases that are “close” to every subject in the training set. (Dalca et al., 2019).

2.1.2. Data likelihood

We treat the latent joint image I_j as a noisy warped atlas,

$$p\left({\mathit{\text{I}}}_j\mid {\varphi }_j;\mathit{\text{A}}\right)=𝒩\left({\mathit{\text{I}}}_j;{\varphi }_j\circ \mathit{\text{A}},{\sigma }^2\mathbb{I}\right)$$ (2)

where N (·; μ, Σ) is the multivariate Gaussian distribution with mean μ and covariance Σ, ◦ represents spatial transformation, σ represents additive noise, and I is the identity matrix. The geometric feature image I_g is then a noisy observation of a further-moved joint image I_j:

$$p\left({\mathit{\text{I}}}_g\mid {\varphi }_g,{\mathit{\text{I}}}_j\right)=𝒩\left({\mathit{\text{I}}}_g;{\varphi }_g\circ {\mathit{\text{I}}}_j,{\sigma }^2\mathbb{I}\right).$$ (3)

Therefore, the complete geometric image likelihood is given by:

$$\begin{array}{l}p\left({\mathit{\text{I}}}_g\mid {\varphi }_g,{\varphi }_j;\mathit{\text{A}}\right)={\int }_{{\mathit{\text{I}}}_j}p\left({\mathit{\text{I}}}_g\mid {\varphi }_g,{\mathit{\text{I}}}_j\right)p\left({\mathit{\text{I}}}_j\mid {\varphi }_j;\mathit{\text{A}}\right)\\ ={\int }_{{\mathit{\text{I}}}_j}𝒩\left({\mathit{\text{I}}}_g;{\varphi }_g\circ {\mathit{\text{I}}}_j,\sigma \mathbb{I}\right)𝒩\left({\mathit{\text{I}}}_j;{\varphi }_j\circ \mathit{\text{A}},\sigma \mathbb{I}\right)\\ ={\int }_{{\mathit{\text{I}}}_j}𝒩\left({\mathit{\text{I}}}_j;{\varphi }_g^{-1}\circ {\mathit{\text{I}}}_g,\sigma \mathbb{I}\right)𝒩\left({\mathit{\text{I}}}_j;{\varphi }_j\circ \mathit{\text{A}},\sigma \mathbb{I}\right)\\ \stackrel{*}{=}{\int }_{{\mathit{\text{I}}}_j}𝒩\left({\mathit{\text{I}}}_j;{\mu }_c,{\mathbf{\Sigma }}_c\right)𝒩\left({\varphi }_g^{-1}\circ {\mathit{\text{I}}}_g;{\varphi }_j\circ \mathit{\text{A}};2{\sigma }^2\mathbb{I}\right)\\ =𝒩\left({\varphi }_g^{-1}\circ {\mathit{\text{I}}}_g;{\varphi }_j\circ \mathit{\text{A}};2{\sigma }^2\mathbb{I}\right)\\ =𝒩\left({\mathit{\text{I}}}_g;{\varphi }_g\circ {\varphi }_j\circ \mathit{\text{A}};2{\sigma }^2\mathbb{I}\right)\end{array}$$ (4)

where in * we used an identity of the product of two Gaussian distributions, and μ_c, Σ_c are constants. We use a similar model for the functional image I_f.

2.1.3. Learning

Let Φ = {ϕ_j, ϕ_g, ϕ_f} and I = {I_g, I_f}. We estimate Φ by minimizing the negative log likelihood,

$$\begin{gathered} ℒ(\Phi \mid \mathit{\text{I}};\mathit{\text{A}})=-\text{log}p(\Phi \mid \mathit{\text{I}};\mathit{\text{A}})=-\text{log}p(\mathit{\text{I}}\mid \Phi ;\mathit{\text{A}})-\text{log}p(\Phi ) \\ =-\text{log}\prod _{k\in \left\{g,f\right\}}p\left({\mathit{\text{I}}}_k\mid {\varphi }_j,{\varphi }_k,\mathit{\text{A}}\right)-\text{log}\prod _{k\in \left\{j,g,f\right\}}p\left({\varphi }_k\right) \\ =\frac{1}{2{\sigma }^2}\left(\parallel {\mathit{\text{I}}}_g-{\varphi }_g\circ {\varphi }_j\circ \mathit{\text{A}}{\parallel }^2+\parallel {\mathit{\text{I}}}_f-{\varphi }_f\circ {\varphi }_j\circ \mathit{\text{A}}{\parallel }^2\right) \\ +{\lambda }_j\parallel \nabla {\mathit{\text{u}}}_j{\parallel }^2+{\lambda }_g\parallel \nabla {\mathit{\text{u}}}_g{\parallel }^2+{\lambda }_f\parallel \nabla {\mathit{\text{u}}}_f{\parallel }^2+{\alpha }_j\parallel {\overline{\mathit{\text{u}}}}_j{\parallel }^2 \\ +\text{ const}. \end{gathered}$$ (5)

2.2. Neural network approach with semi-supervised training strategy

We use a neural network to approximate the function h_θ,A(I) = Φ, where θ are network parameters. Fig. 2 shows the proposed network architecture. To work with surface-based data, the cortical surface of each subject is inflated into a sphere and then rigidly registered to an average space using FreeSurfer (Fischl, 2012). Geometric and functional features are parameterized onto a 2D grid using a standard conversion from Cartesian coordinates to spherical coordinates, resulting in a 2D image for each input (Cheng et al., 2020a).

Figure 2:

Network architecture and preprocessing pipeline. The network takes the geometric features from the subject and outputs one joint and two separate deformation fields. The joint and geometric deformation are composted for registration of folding patterns, and likewise for function. Task fMRI data were used for evaluating the functional loss only in a semi-supervised manner, hence, functional task data are not required during inference.

The network takes such a parameterized geometric image as input and outputs three velocity fields in the parameterized coordinates, each followed by an integration layer generating the corresponding deformation field. The joint deformation ϕ_j, which models the relatively large inter-subject variance, is shared among and composed with individual deformations ϕ_g and ϕ_f. We note that the separation of ϕ_g from ϕ_f enables us to explicitly model structure-function variability across subjects. This is of critical importance from a neuroscientific perspective because it has been shown that some functional areas are better predicted by folding patterns than others (Fischl et al., 2008). The losses ℒgeom and L_func shown in Fig. 2 represents the geometric part and the functional part of the data fidelity terms in Eq.(5). They are evaluated in the atlas and subject space, which also helps avoid atlas drift (Aganj et al., 2017) during atlas construction. L_reg represents the regularization and centrality terms which encourages smooth deformations and an unbiased estimation of the atlas.

In this study, we learn network parameters and use task fMRI data in a semi-supervised manner. As shown in the green block in Fig. 2, the task fMRI data and the corresponding functional atlas are not input into the neural network. Rather they are used only for evaluating the functional terms in the loss function (5). This obviates the need for functional data during inference, as the deformation fields can be inferred using only geometric features. The proposed framework is flexible in the sense that auxiliary data, perhaps from different modalities, can be integrated into the framework for simultaneous multi-modality registration.

2.3. Implementation

We implemented a Unet-like (Ronneberger et al., 2015) network based on the core architecture in VoxelMorph (https://voxelmorph.net) (Balakrishnan et al., 2019; Dalca et al., 2019). We used a 5-layer encoder with [128, 256, 384, 512, 640] filters and a symmetric decoder followed by 2 more convolutional layers with [64, 32] filters. Each layer involves convolution, max-pooling/down-sampling, and LeakyReLU activation. The spherical parameterization leads to denser sampling grids for regions at higher latitudes. To account for the difference in sampling density, we performed prior and distortion corrections identical to that described in (Cheng et al., 2020a). In short, weights proportional to sin(θ), where θ is the elevation, were used to correct the distortion. (Cheng et al., 2020a) also found that varying the locations of the poles in the projection had little impact on the resulting registration. The parameterized images were standardized identically but separately for structural and functional features, where the median was subtracted for each feature image followed by a division by the standard deviation.

During training, we randomly sampled the training data into mini-batches with a batch size of 8. For each batch, we augmented the data by adding Gaussian random deformations with a maximum σ = 8 and a proper spherical topology and distortion correction at each spatial location. We further augmented the data by adding Gaussian noise with σ = 1 for geometric features and σ = 6 for functional features. We used the Adam optimizer (Kingma and Ba, 2014) with an initial learning rate of 10⁻³. The learning rate was scheduled to decrease linearly to 10⁻⁴ within the first 500 epochs and then reduced by a factor of 0.9 if the validation loss does not decrease after every 100 epochs. The relative weights between functional loss and geometric loss were set to 0.7:0.3 empirically. We set the regularization hyperparameter λ_j in Eq. (5) for the joint (large) deformation to be 0.1 and λ_g and λ_f for the individual (small) deformations to be 0.2.

The atlases, as part of the network parameters, were initialized using Gaussian random noise and automatically learned during training. We exploited a three-stage training strategy to further avoid atlas expansion or drift (Aganj et al., 2017): 1) we performed a training run using a single deformation with geometric features only, so that the network learned a geometric atlas; 2) we used and fixed the geometric atlas obtained in step one and trained a JOSA network to learn the functional atlas; and 3) we fixed both the geometric and the functional atlas from the previous two steps and performed a final training for the composite deformation fields. We used TensorFlow (Abadi et al., 2016) with Keras front-end (Chollet, 2018) and the Neurite package (Dalca et al., 2018b), and all experiments were conducted in a Dell Workstation with dual Intel Xeon Silver 6226R CPUs and an Nvidia RTX6000 GPU. The source code and the trained model of JOSA will be released to the public at https://voxelmorph.net as well as integrated as part of FreeSurfer (Fischl, 2012) package.

3. Experiments

3.1. Language task fMRI data

We used task fMRI data from a large-scale language mapping study (Lipkin et al., 2022) (N = 800), where the language network was functionally localized using a task that contrasts reading/listening of sentences versus a perceptually matched control (such as strings of nonwords or degraded speech), collected using a standard blocked design. Ten different versions of the language localizer task were used (Table 2 in Lipkin et al. (2022)), with the vast majority of subjects (~77%) completing a version that contrasted reading of sentences with nonwords strings. In this localizer version, each stimulus was presented one word/nonword at a time at the rate of 450 ms per word/nonwords (12 words/nonwords per stimulus). Each stimulus was preceded by a 100 ms blank screen and followed by a 400 ms screen showing a picture of a finger pressing a button, and a blank screen for another 100ms, for a total trial duration of 6s. Experimental blocks lasted 18s, and fixation blocks lasted 14s. Each run (consisting of 5 fixation blocks and 16 experimental blocks) lasted 358s. Subjects completed 2 runs. Subjects were instructed to read attentively and press a button whenever they saw the finger-pressing picture on the screen. Structural and functional data were collected on a 3T Siemens Trio scanner at the Athinoula A. Martinos Imaging Center at the McGovern Institute for Brain Research at MIT. T1-weighted images were collected in 176 sagittal slices with 1 mm isotropic resolution. Functional data (BOLD) were acquired using an EPI sequence with 4 mm thick near-axial slices, 2.1 mm × 2.1 mm in-plane resolution, TR = 2,000 ms, and TE = 30 ms. We preprocessed the data using FreeSurfer v6.0.0 as described in (Lipkin et al., 2022). In brief, data were sampled onto the FreeSurfer average space (fsaverage), motion corrected, registered using the middle time point of each run, and spatially smoothed with a 4mm FWHM Gaussian filter. The subjects’ surfaces were reconstructed from the T1 images (default recon-all parameters). A “sentence vs. nonword” contrast t-map was originally generated for each subject using first-level GLM analysis based on the blocked design, and sampled on the reconstructed surface in the fsaverage space (Lipkin et al., 2022). The t-maps were resampled back onto the individual surfaces using FreeSurfer’s ‘mri_surf2surf’ utility as functional features for registration in this work. We randomly split the data into a training set with 600 subjects, a validation set with 100 subjects, and a test set with the remaining 100 subjects.

3.2. Baseline

We compare the registration performance of JOSA to FreeSurfer (Fischl et al., 1999b), Multimodel Surface Matching (MSM) (Robinson et al., 2014), the improved version of MSM with higher-order constraint (MSM-HOC) (Robinson et al., 2018), and SphereMorph (Cheng et al., 2020a) as surface registration baselines. For all baseline methods, we used three geometric features that quantify the folding patterns of the cortex to drive the registration, namely the mean curvature of the inflated surface (?h.inflated.H), the sulcal depth map (?h.sulc) and the mean curvature (?h.curv) of the non-inflated white matter surface, all available as part of the FreeSurfer standard outputs. For FreeSurfer registration, we ran mris_register (https://surfer.nmr.mgh.harvard.edu/fswiki/mris_register) to register each subject to fsaverage with the configurations used in their 2014 and 2018 papers, respectively. For a fair comparison, both versions of MSM were executed in a hierarchical manner such that ?h.inflated.H was used for initial coarse alignment, then refined by ?h.sulc and ?h.curv, consecutively. For SphereMorph (https://voxelmorph.net), we used fsaverage as the registration target and trained the network to predict a single deformation field purely based on the three geometric features. The predicted deformation is then used in turn to warp each subject’s functional data to the fsaverage space at the inference time.

3.3. Evaluation

Qualitatively, we computed the group mean images of both the geometric and the functional data in the test set for each registration method. We then visualized them by superimposing the functional group mean map with the curvature group mean map using Freeview (Fischl, 2012).

We quantified the geometric registration accuracy using the correlation between the registered individual data and the group mean (Cheng et al., 2020a). Specifically, we computed the Pearson correlation c_k = corr(I_k,Ī₎ between the individual image I_k to the group mean image Ī = 1/N k I_k for subject k, where N is the number of subjects. We quantified the functional registration accuracy using the size of the overlapped regions between the suprathreshold individual functional maps and the suprathreshold group mean. We expect a better registration will have a larger surface area in the individual functional active regions as well as a larger overlap with the group mean for a reasonably conservative threshold. This also helps avoid a potential evaluation bias when a single-task contrast is used where only a fraction of the surface area is active during tasks. Specifically, we computed $|s_k=\left|I_k^{th}\cap {\overline{I}}^{th}\right|,$ where $I_k^{th}$ is the suprathreshold individual image for subject k at threshold th and similarly for the group image Īth, and is the measure of cardinality. We used a threshold of 3 in this evaluation, where the suprathreshold regions for the group mean are shown in Fig. 3.

Figure 3:

Group average curvature map of 100 test subjects overlaid with the group average functional language activation map. (a) Rigid; (b) FreeSurfer; (c) MSM; (d) MSM-HOC; (e) SphereMorph; (f) JOSA. The language activation map is defined based on the contrast between well-formed sentences and a perceptually matched control stimulus. In the curvature maps, dark gray indicates sulci and light gray indicates gyri. The color bar for functional data (t-values) is shown on the right. White arrows highlight the key regions in functional alignment where large functional variability across subjects is often observed. JOSA achieved a substantially better alignment in both folding patterns and function with stronger functional response and sharper transition from task-active to non-active regions.

To better demonstrate the consistency in registration improvement across subjects, we assessed the pair-wise geometric correlation difference c^k_JOSA − c^k_Rigid between our proposed method JOSA and the rigid alignment, and similarly for FreeSurfer, MSM, MSM-HOC, and SphereMorph (i.e., the correlation difference for the same subject after registration in comparison to the rigidly aligned version). We similarly assessed the pair-wise functional difference s^k_JOSA − s^kRigid.

In functional evaluations, to further investigate the impact of the selected threshold th, we computed the sum of the superthreshold t-values in the group mean ${\sum }_j{\overline{I}}_j^{th}$ as a function of the threshold th ∈ {3, …, 5} We also looked into the two individual factors that have impact on the sum of the suprathreshold t-values, namely the number of vertices that have t-values above the threshold $M=\left|{\overline{I}}^{th}\right|$ as well as the average of the suprathreshold t-values within those regions $1/M{\sum }_j{\overline{I}}_j^{th}$ The former reflects the size of the functional active regions and the latter reflects the strength of the functional response.

We also compared the computational efficiency of the registration methods by measuring the run time the registration procedure for each of the test subjects. Iterative methods (FreeSurfer, MSM, and MSM-HOC) were run on a single CPU thread, while the learning-based methods (SphereMorph and JOSA) were evaluated using both CPU and a single GPU. See Section 2.3 for hardware details.

3.4. Stability of Atlas Construction

Bias towards the initial selected target is a well known issue in atlas construction (Fischl et al., 1999b; Dickie et al., 2017; Cheng et al., 2020b). To investigate the stability and the potential bias of atlas construction using JOSA, we repeated the geometric atlas learning procedure 6 times, each time with randomly selected batches of subjects and their orders through the training process. We qualitatively visualized the difference between each run to the average of all runs. To quantitatively measure the stability of the learned atlases, for each pair of the learned atlases with the 6 trials (hence 15 pairs in total), we computed Pearson correlation between the pairs of atlases.

3.5. Ablation studies

We conducted the following two ablation studies: 1) To illustrate benefit of learning the population-specific and unbiased atlas, we trained two JOSA networks with the only difference being in the registration target, where we used rigid average and FreeSurfer average for comparison to our learned atlas; 2) To illustrate the power of modeling the difference between geometry and function, we trained an additional JOSA network with a single deformation field that was used to align both geometry and function for comparison to our two separate deformations.

4. Results

4.1. Qualitative Comparison

Fig. 3 shows the group mean curvature maps (one of the geometric features) of the 100 test subjects overlaid with the group mean functional language activation maps for each of the registration methods. A better registration leads to a stronger average response (larger yellow areas) to language stimuli and a sharper transition from task-positive to task-negative regions (narrower margin from yellow to transparent regions). FreeSurfer, MSM, MSM-HOC, and SphereMorph significantly improved the registration of the folding patterns compared to the rigid alignment, albeit with some noticeable but minor differences in the average curvature patterns. However, all these methods only yield marginal improvement of the functional alignment. In contrast, JOSA achieved a substantially better alignment in both folding patterns and function. In particular, the predominant language region in the superior temporal gyrus shows a substantially stronger response and clearer functional boundaries. We also find additional/stronger language-responsive regions in the frontal eye field, the area in the inferior frontal gyrus (BA 44/45 near the classic Broca’s area), as well as the posterior section of the inferior parietal lobule as indicated by the white arrows, all consistent with the probabilistic language atlas in the earlier work (Lipkin et al., 2022).

4.2. Quantitative Evaluation

Quantitatively, Fig. 4 (a) shows the pair-wise geometric correlation difference between each of the registration methods to the rigid alignment, and Fig. 4 (b) shows the pair-wise functional suprathreshold overlap difference. Values higher than zero indicate consistent registration improvement over rigid alignment. Fig. S1 in the supplementary material shows the original correlation/overlap values without taking pair-wise differences. Iterative methods such as FreeSurfer and MSM substantially improved geometric correlations over rigid alignment (by ~0.3). Learning-based methods achieved a further significant improvement (by ~0.43) and were comparable between SphereMorph and JOSA. In contrast, only JOSA substantially improved functional registration, while all other methods did not improve or only marginally improved functional registration compared to rigid alignment.

Figure 4:

Violin plot of pairwise difference between each registration method to rigid alignment. (a) Geometric difference measured by correlation; (b) Functional difference measured by the suprathreshold (t-values >= 3) functional overlaps. JOSA yielded substantial higher correlation/larger overlap against rigid alignment compared to baseline methods and the improvement is consistent across subjects as shown by the pairwise difference measures.

Fig. 5 shows the total suprathreshold t-values of the group mean functional image for each registration method as a function of the threshold. The total suprathreshold t-values are measured as the sum of the t-scores over the regions that have values above the threshold. As shown in Fig. 5, all geometry-based registration methods yielded similar functional total t-values to that using only rigid alignment, while JOSA, in contrast, produces substantially higher functional total t-values. This significant difference between JOSA and baseline methods were consistent regardless of the choice of the threshold. We also note that the higher total t-values may come from either a larger region where the functional values exceeded the threshold or a stronger response or both. To further illustrate the effect, we plotted the curves for (a) the number of vertices within regions that have functional values above the threshold and (b) the averaged functional values within those regions in the Supplementary Fig. S2. We can see that JOSA has both larger functional regions as well as a stronger functional response within those regions, and the former contributes more than the latter.

Figure 5:

Functional total suprathreshold t-values of the group means as a function of the threshold. The total suprathreshold t-values in JOSA is substantially higher than that using baseline methods regardless of the selected threshold.

4.3. Run time

Table 1 summarizes the run time for each registration method. For a single subject registration, all methods started with the spherical cortical surface, and the output is the warped spherical surface. The default FreeSurfer takes approximately 9 min to complete the spherical registration procedure. MSM and MSM-HOC operated directly on the sphere without parameterization but took significantly longer to execute. In contrast, the learning-based methods take, on average, only 0.097 min (≈ 5.8 sec!), providing a speed up by two orders of magnitude compared to the iterative methods. The registration speed can be further accelerated by ~ 0.2 sec per subject with a GPU implementation. We note that this 0.2 sec CPU-to-GPU improvement is significant because, in both SphereMorph and JOSA, the parameterization and resampling back to sphere take most of the time during registration, and the actual inference time through the convolutional neural network is ~ 0.06 sec using GPU. That means if the parameterization and resampling can be parallelized and precomputed on a computer cluster, registering 1000 subjects can be done in under 1 minute.

Table 1:

Comparison of the run time for individual subject registration.

Method	CPU run time (min)	GPU run time (min)
FreeSurfer	8.82 ± 2.47
MSM	21.05 ± 1.87
MSM-HOC	57.09 ± 8.86
SphereMorph	0.0969 ± 0.0081	0.0947 ± 0.0085
JOSA	0.0972 ± 0.0089	0.0945 ± 0.0083

4.4. Atlas Stability

Fig. 6 shows the geometric atlases learned from the 6 repeated training runs, which are identical in all respects except that the subjects and their orders are randomly selected as the input to the training. Fig. 6 (a) shows the mean averaged over the 6 runs and (b) - (g) shows the absolute difference between each run and the mean. These atlases are highly consistent across runs with all major folds nearly identical to each other and subtle differences in the prefrontal, inferior parietal, and middle temporal regions as reflected in the difference images. This strong consistency is quantitatively confirmed by the Pearson correlations 0.94 ± 0.01 among all 15 possible pairs of these atlases. The results demonstrate that the atlas construction and its learning procedure in our proposed JOSA method are stable and not biased towards any specific subject.

Figure 6:

Learned geometric atlases from 6 repeated training with randomly selected subjects. (a) The mean atlas averaged over the 6 runs. (b) - (g) The absolute difference between each run to the mean. Red indicates sulci and blue indicates gyri. The learned atlases are highly consistent across runs all major folds nearly identical to each other and subtle difference in prefrontal, inferior parietal, and middle temporal areas.

4.5. Ablation Study

The results for the ablation studies are shown in Fig. 7. The first and the second column shows the result using separate deformations for geometry and function as with JOSA but varying the atlas used as the registration target. The third column shows the result using a single deformation for both geometry and function but the same learned atlas as with JOSA. The last column shows the full JOSA result again for easy reference. Again these are pair-wise differences from each registration method to the rigid result. From these ablation results, we can see that the substantial improvement in geometry is mainly attributable to a better atlas, whereas the improvement in the registration of function is primarily due to the separate modeling of deformation between geometry and function. We highlight the importance of this separate modeling of deformations. It is evident from the third column that the registration performance is sub-optimal for both geometry and function if they share the same deformation, which is commonly assumed (anatomy predicts function).

Figure 7:

Violin plot of pair-wise difference between each registration method to rigid alignment for ablation studies. (a) Geometric difference measured by correlation; (b) Functional difference measured by the suprathreshold overlaps. Geometric improvement is mainly attributable to a better atlas where as functional improvement is primarily due to the separate modeling of the deformation fields.

5. Conclusion and discussion

In this work, we presented JOSA, a novel registration framework that jointly models the relationship between folding patterns and cortical function. In the early 1900s, the correspondence between the two was discovered in primary motor and sensory regions where a cortical homunculus was drawn to illustrate the distorted mapping of human body parts to sensorimotor functions (Penfield and Boldrey, 1937). Similarly, the functional properties of primary visual cortex have been well studied, where the central and peripheral visual fields are located in the posterior and anterior part of the calcarine sulcus, respectively (Dougherty et al., 2003; Sereno et al., 1995; Tootell et al., 1998; Engel, 1997; Wandell et al., 2007; Engel et al., 1994), and the topographic mapping of the visual field to the visual cortex follows the complex-logarithm rule (Schwartz, 1980; Polimeni et al., 2010). However, the anatomo-functional mapping in higher-order functional areas, e.g., language, emotion, etc., is highly variable among individuals compared with the primary regions (Steinmetz and Seitz, 1991; Fischl et al., 2008; Frost and Goebel, 2012; Robinson et al., 2014). This variability is mainly attributable to the non-unique optimal solutions for geometric registration, where, without the guidance from functional data, their performance can be sub-optimal in registration of function. For example, there are several secondary folds around the active language region near the frontal eye field. Aligning any of them may yield equally good geometric registration, but only one of them may be optimal for functional registration. This is not obviously without access and visualization of the functional data of each individual subject. By using a semi-supervised training strategy with functional data as the guidance, JOSA is able to choose the “right” deformation from many valid geometric solutions that produces optimal functional registration. As evidenced in Fig. 4, it is clear that both JOSA and SphereMorph (no functional data were used) achieved comparable performance in geometric alignment, but only JOSA substantially improved functional registration. Moreover, a structure-function mismatch may exist and be different across individuals, i.e., different subjects may use slightly different regions to process similar function. By allowing small deviations of the functional alignment from geometric alignment, JOSA yields superior performance in registration of both folding patterns and functional properties relative to other traditional or learning-based methods. The result in the ablation study shown in Fig. 7, the 3rd column, illustrated this point, where the both the geometric and functional registration are sub-optimal when a single (shared between geometry and function) deformation is used.

In contrast to many other learning-based methods, JOSA also estimates a population-specific atlas as the registration target during training. Fig. 8 shows the learned atlas in (c) compared to atlases constructed by averaging over rigidly registered training subjects (a) and averaged FreeSurfer registered training subjects in (b). We find that the JOSA atlas provides more anatomical definition, supporting registration with higher resolution and finer details, which also contributes to the improved performance of JOSA (particularly in geometric registration).

Figure 8:

Comparison of atlases. (a) Rigid atlas; (b) FreeSurfer atlas; (c) JOSA atlas. The learned atlas from JOSA provides more anatomical definition than rigid or FreeSurfer atlas, thus supporting registration with higher resolution and finer details.

Atlas or template construction has been widely studied in classical iterative approaches (Collins et al., 1995; Fischl et al., 1999b; Desikan et al., 2006; Destrieux et al., 2010; Joshi et al., 2022). These methods build atlases by repeatedly registering subjects to an estimated atlas, and estimating a new atlas by averaging the registered subjects (Dickie et al., 2017). Inherited from this iterative construction approach, a bias towards the initial subject may be introduced to the atlas. JOSA, as one of the very recent learning-based methods, not only facilitates faster atlas construction (Lee et al., 2022b,a), but also potentially avoids the bias to any individual subject by random selection of subjects as inputs during training (Cheng et al., 2020b), as illustrated in Fig. 6. In this work, JOSA learned an atlas specific to the population that was used for training (Lipkin et al., 2022). It is, however, straightforward to extend this framework to learn any population-specific atlases, optionally conditioned on clinical attributes if desired (Dalca et al., 2019; Dey et al., 2021; Ding and Niethammer, 2022).

Using a semi-supervised training strategy, the deformation fields can be solely determined from geometric features, which can be derived from a single T1-weighted MR scan. Therefore, JOSA not only lifts the burden of acquiring functional data during inference, but also provides capacity for extension to other comprehensive data modalities that are related to geometry but difficult or even impossible to acquire, such as architectonics, transcriptomics, or proteomics. This geometry-only-based registration model further avoids the circularity issue in any subsequent analyses of functional or any other auxiliary data. In another words, if the functional data is used as one of the inputs to drive the registration, then any evaluation or analysis of the functional data after registration will be potentially biased as the network has observed the data before.

The current approach for better registration of brain function is limited by the use of one task fMRI localizer, the language localizer, used in this study. The partial coverage of the brain by the language-active regions may not provide useful registration guidance for functions that use different parts of the cerebral cortex. We therefore plan to extend JOSA to training with multi-modal datasets, enabling accurate registration of an array of brain functional responses. Although the language dataset used in this study (Lipkin et al., 2022) is one of the largest task fMRI datasets publicly available, another caveat in the training of the registration network is the limited number of samples (subjects). Because the functional data were used only in loss evaluation through the semi-supervised block (i.e., they were not seen by the convolutional layers directly), it poses a significant challenge for the network to learn the anatomo-functional relationships, hence often resulting in an overfitting issue when the sample size is small. We also acknowledge the difficulty in data acquisition. One potential solution in the future is to use a pooled dataset from different task contrasts or even resting-state fMRI data, where subjects may have different functional data available, and the network can be trained dynamically using the corresponding available features, atlases, and loss functions.