Superficial White Matter Analysis: An Efficient Point-cloud-based Deep Learning Framework with Supervised Contrastive Learning for Consistent Tractography Parcellation across Populations and dMRI Acquisitions

Abstract

Diffusion MRI tractography is an advanced imaging technique that enables in vivo mapping of the brain’s white matter connections. White matter parcellation classifies tractography streamlines into clusters or anatomically meaningful tracts. It enables quantification and visualization of whole-brain tractography. Currently, most parcellation methods focus on the deep white matter (DWM), whereas fewer methods address the superficial white matter (SWM) due to its complexity. We propose a novel two-stage deep-learning-based framework, Superficial White Matter Analysis (SupWMA), that performs an efficient and consistent parcellation of 198 SWM clusters from whole-brain tractography. A point-cloud-based network is adapted to our SWM parcellation task, and supervised contrastive learning enables more discriminative representations between plausible streamlines and outliers for SWM. We train our model on a large-scale tractography dataset including streamline samples from labeled long- and medium-range (over 40mm) SWM clusters and anatomically implausible streamline samples, and we perform testing on six independently acquired datasets of different ages and health conditions (including neonates and patients with space-occupying brain tumors). Compared to several state-of-the-art methods, SupWMA obtains highly consistent and accurate SWM parcellation results on all datasets, showing good generalization across the lifespan in health and disease. In addition, the computational speed of SupWMA is much faster than other methods.

1. Introduction

Diffusion magnetic resonance imaging (dMRI) tractography is the only non-invasive method that can map the brain’s white matter connections (Basser et al., 2000). Tractography methods estimate trajectories of brain white matter connections and represent those trajectories using streamlines. Each streamline is a set of ordered points in 3D space (Zhang et al., 2022a). Performing whole-brain tractography generates streamlines throughout the entire white matter, including the deep white matter (DWM) that connects distant cortical regions, and the superficial white matter (SWM) that includes short-range association connections (u-fibers) connecting adjacent and nearby gyri (Guevara et al., 2020; Zhang et al., 2022a). Whole-brain tractography can produce hundreds of thousands of streamlines, which are not directly useful to clinicians and researchers for quantification or visualization. Therefore, tractography parcellation is needed to classify white matter streamlines into clusters or anatomically meaningful tracts (Zhang et al., 2022a). Fig. 1 provides a visualization of clusters in the entire SWM and several example SWM clusters from an anatomically curated atlas (Zhang et al., 2018b), where clusters are annotated with anatomical labels by a neuroanatomist. SWM parcellation is important to enable analyses of the SWM in neuroscientific studies in health and disease (Reginold et al., 2016; d’Albis et al., 2018; Malykhin et al., 2011; Ji et al., 2019) (see (Guevara et al., 2020) for a review).

Fig. 1.

Visualization of SWM clusters. (a) shows clusters from the entire SWM, where each cluster has a unique color, and (b) shows several example SWM clusters.

Parcellation of the SWM is a challenging task due to the small diameter of SWM fasciculi, their high variability across brains, and their unique position near the cortex (Guevara et al., 2020). (We note that we focus on parcellation of SWM tractography in this work. Additional challenges inherent to performing SWM tractography, which include high curvature and complex fiber crossings (Reveley et al., 2015), gyral bias (Van Essen, 2013; Schilling et al., 2018), and partial voluming effects (Alexander et al., 2001), can be addressed to some extent by high-resolution imaging (Song et al., 2014a) and development of advanced tractography methods (St-Onge et al., 2018).) Whereas most tractography parcellation methods currently focus on the DWM (Yendiki et al., 2011; Garyfallidis et al., 2018; Wasserthal et al., 2018; Zhang et al., 2020a; Chen et al., 2021), few methods can parcellate the SWM (Oishi et al., 2008; Román et al., 2017; Guevara et al., 2020; Zhang et al., 2018b). The existing SWM tractography parcellation methods use either region of interest (ROI)-based selection or streamline clustering. ROI-based methods parcellate SWM tractography based on the ROIs streamlines end in and/or pass through (Ouyang et al., 2016; Oishi et al., 2008; Malykhin et al., 2011; Schilling et al., 2022). These ROI-based methods are the most commonly used but highly depend on the ROI parcellation scheme. Streamline clustering methods group SWM streamlines based on the similarity of their geometric trajectories (Zhang et al., 2018b; Guevara et al., 2012; Román et al., 2017). Streamline clustering methods for SWM parcellation are automatic and can leverage SWM atlases (Guevara et al., 2017; Román et al., 2017; Zhang et al., 2018b; Román et al., 2022), but challenges remain to achieve consistent parcellation across subjects and reduce runtime.

In recent years, deep-learning-based methods (Wasserthal et al., 2018; Zhang et al., 2020a; Astolfi et al., 2020a; Xu et al., 2019; Chen et al., 2021, 2022) have been successful for fast and consistent tractography parcellation, but representing dMRI data in a way that can best take advantage of deep networks is still an open challenge. Voxel-based (Wasserthal et al., 2018; Lu et al., 2021; Wasserthal et al., 2019) white matter tract parcellation methods take volumetric image data (e.g., fiber orientation distribution function (FOD) peaks (Wasserthal et al., 2018)) and predict a tract’s presence and/or orientation for each voxel. Streamline-based tractography parcellation methods (Zhang et al., 2020a; Xu et al., 2019; Ngattai Lam et al., 2018) provide input to deep networks by encoding streamlines into different features. For example, DCNN+CL+ATT (Xu et al., 2019) uses a 1D feature descriptor containing streamline point spatial coordinates, and DeepWMA (Zhang et al., 2020a) represents streamlines as 3-channel 2D images. However, the ambiguity of streamline data (the points along a streamline can equivalently be represented in forward or reverse order) poses challenges when using these representations as deep network input. Point clouds, as an important geometric data format (Guo et al., 2020), can potentially enable efficient and discriminative representations for streamlines. Coordinates of streamline points can be the input for point-cloud-based deep networks, as used in (Astolfi et al., 2020b) for tractography filtering (binary classification of streamlines). There are two main advantages of using a point cloud representation for streamlines. First, streamline points are not on a regular grid and thus can be naturally represented using point clouds. Second, fiber tracking in tractography provides no directional information (points along a streamline can equivalently be represented in forward or reverse order). With the usage of a point cloud representation, streamlines with equivalent forward and reverse point orders (e.g., from cortex to subcortical structures or vice versa) can have the same global shape feature representation in the network. To our knowledge, no deep learning methods have focused on SWM parcellation, and point-cloud-based deep networks have not yet been used for white matter parcellation, in particular for SWM parcellation.

In this paper, we propose a novel two-stage deep learning framework, Superficial White Matter Analysis (SupWMA), for SWM parcellation from whole-brain tractography. A point cloud classification network is used in both stage one and stage two. In stage one, the network takes streamlines of the whole-brain tractography as input, and outputs streamlines classified as SWM and DWM. In this way, stage one filters out DWM streamlines. In stage two, the network takes SWM streamlines as input, and it classifies them into 198 SWM clusters and outliers. Supervised contrastive learning is employed in stage two to obtain more discriminative representations of SWM streamlines. In our study, we focus on long- and medium-range SWM streamline connections (over 40mm in length).

Our contributions are as follows. First, we decompose SWM parcellation into two stages, including a binary classification of SWM and DWM streamlines followed by a multi-class classification of 198 SWM clusters while performing removal of SWM outliers. Second, we modify the point-cloud-based network structure to preserve streamline pose and orientation information (because location in the brain is important for classification), and the modified network provides improvements to computational speed and accuracy. Third, we adapt supervised contrastive learning with a data augmentation technique (bilateral augmentation) for streamline classification. More discriminative representations between SWM plausible streamlines and outliers are obtained. Fourth, we evaluate our methods on six independently acquired datasets of 329 subjects with different ages and health conditions (including neonates and patients with space-occupying brain tumors), which are completely separate from the training set, and we obtain consistent results efficiently.

This investigation extends our previous conference publication (Xue et al., 2022) to improve performance of SWM parcellation by designing a two-stage framework and incorporating bilateral data augmentation. Furthermore, we demonstrate successful SWM parcellation on three additional populations including neonates, patients with brain tumors, and patients with psychiatric disorders. We also include additional qualitative and quantitative evaluations of parcellation results, and an investigation of point importance along the streamline for classification.

The remaining structure of this paper is as follows. Section 2 describes the datasets used, the proposed framework, and the model training and testing. Section 3 presents the experimental setup and results for the training dataset and testing datasets. Finally, the discussion, conclusions, and future work are given in Section 4.

2. Methodology

An overview of the SupWMA framework is given in Fig. 2. In the rest of this section, first we introduce relevant datasets (Section 2.1), then we present our two-stage framework and contrastive learning methodologies (Sections 2.2-2.4), and finally we describe training and testing procedures (Sections 2.5-2.6).

Fig. 2.

Overview of the SupWMA framework: (a) samples of training tractography data, (b) deep network structure and parcellation model training, (c) parcellation of unseen testing datasets.

2.1. dMRI datasets and tractography

Training dataset:

A high-quality large-scale tractography dataset with 1 million labeled streamlines was used for model training and validation. This dataset was derived from an anatomically curated white matter tractography atlas (Zhang et al., 2018b), which was provided by the O’Donnell Research Group (ORG) and is available to access online¹. The ORG atlas has been previously used for DWM streamline classification with deep learning (Zhang et al., 2020a). In brief, the atlas was generated by creating dense tractography maps (Malcolm et al., 2010) of 100 young healthy adults in the Human Connectome Project (HCP) (Van Essen et al., 2013) and applying a fiber clustering method (O’Donnell and Westin, 2007; O’Donnell et al., 2012; Zhang et al., 2018b) to group streamlines across subjects according to their similarity in shape and location. More specifically, this method performed spectral embedding of streamlines to learn a high-dimensional space where each streamline had a unique representation. In this space, a k-means clustering was performed to group the streamlines into multiple clusters. Clusters do not overlap in the high-dimensional spectral space, enabling unambiguous cluster definitions. For each cluster, a data-driven process was then applied to identify apparent outlier streamlines, defined as those that had a fiber similarity over 2 standard deviations from the cluster’s mean fiber similarity (as in (O’Donnell et al., 2017; Zhang et al., 2018a,b)). In this way, every cluster in the ORG atlas obtained a corresponding group of outlier streamlines. A neuroanatomist then annotated anatomical labels for clusters in the atlas. Atlas clusters are bilateral (clusters have similar shapes and cortical projections across both left and right hemispheres). In comparison with other publicly available SWM atlases (Guevara et al., 2020), the ORG atlas provides comprehensive coverage of the SWM and a fine parcellation scale. The ORG atlas also includes labeled outlier streamlines (Zhang et al., 2020a). The stage-one training dataset (D₁) includes all 1 million streamlines, where each streamline is labeled as either SWM or DWM. The stage-two training dataset (D₂) includes only SWM streamlines (237896 streamlines), where streamline labels include 198 SWM cluster classes and 198 corresponding SWM outlier categories. Each streamline has one label. Note that SWM streamlines in the ORG atlas are long- and medium-range SWM connections (over 40mm in length).

Testing datasets:

For experimental evaluation, we used data from 329 subjects from six independently acquired datasets (testing datasets) with different imaging protocols across the lifespan (neonates, children, young adults and older adults, ranging in age from 1 day to 82 years) and health conditions (healthy control, Parkinson’s patients, neuropsychiatric disorders, and neurosurgical patients with brain tumors). Tractography streamlines in these testing datasets do not have labels. Table 1 presents an overview of the demographics and acquisition protocols for the six testing datasets. HCP (Van Essen et al., 2013), dHCP (Edwards et al., 2022), ABCD (Volkow et al., 2018), CNP (Poldrack et al., 2016) and PPMI (Marek et al., 2011) datasets are publicly available. Image data from the BTP dataset was acquired at Brigham and Women’s Hospital. The Partners Healthcare Institutional Review Board approved the data usage, and all participants gave informed consent prior to scanning.

Table 1.

Demographics and dMRI acquisition of the six independently acquired datasets for experimental evaluation.

Dataset	Number	Age	Gender	Healthy/Disease	dMRI data
HCP	100	22 to 35 y (29.0 ± 3.5)	54 F 46 M	100 H	b=3000 s/mm2; 108 directions; TE/TR=89/5520 ms; resolution=1.25 mm3
dHCP	40	1 to 27 d (6.30 ± 7.47)	15 F 25 M	40 H	b=400/1000/2600 s/mm2; 300 directions; TE/TR=90/3800 ms; resolution=1.5 mm3
ABCD	50	9 to 11 y (10.02 ± 0.69)	25 F 25 M	50 H	b = 3000 s/mm2; 96 directions; TE/TR = 88/4100 ms; resolution = 1.7 mm3
CNP	50	21 to 50 y (32.36 ± 8.65)	23 F 27 M	14 H 12 ADHD 12 BD 12 SZ	b = 1000 s/mm2; 64 directions; TE/TR = 93/9000 ms; resolution = 2mm3
PPMI	50	45 to 80 y (62.52 ± 6.95)	13 F 37 M	25 H 25 PD	b = 1000 s/mm2; 64 directions; TE/TR = 88/7600 ms; resolution = 2 mm3
BTP	39	23 to 82 y (48.9 ± 15.3)	16 F 23 M	39 BTP	b=2000 s/mm2; 31 directions; TE/TR=98/12700ms; resolution=2.3 mm3

Abbreviations: Dataset: HCP - Human Connectome Project (Van Essen et al., 2013) (These subjects are different from those used in the atlas); dHCP - Developing Human Connectome Project (Edwards et al., 2022); ABCD - Adolescent Brain Cognitive Development (Volkow et al., 2018); CNP - Consortium for Neuropsychiatric Phenomics (Poldrack et al., 2016); PPMI - Parkinson’s Progression Markers Initiative (Marek et al., 2011); BTP - Brain Tumor Patient. Age: d - day; y - year. Gender: F - female; M - male. Healthy/Disease: H - healthy; ADHD - attention-deficit/hyperactivity disorder; BP - bipolar disorder; SZ - schizophrenia; PD - Parkinson’s disease; BTP - brain tumor patient.

Whole-brain tractography was performed on all testing datasets using the two-tensor Unscented Kalman Filter (UKF)² method (Malcolm et al., 2010; Reddy and Rathi, 2016). This method works robustly across acquisitions and different populations including neonates, children, adults and patients with peritumoral edema (Chen et al., 2015; Zhang et al., 2018b, 2020a). To perform tractography we used the same parameter settings used for generating tractography in the ORG atlas (Zhang et al., 2018b). In brief, tractography was seeded in all voxels within the brain mask where fractional anisotropy (FA) was greater than 0.1. Tracking stopped where the FA value fell below 0.08 or the normalized mean signal (the sum of the normalized signal across all gradient directions) fell below 0.06. Streamlines that were longer than 40 mm were retained (Guevara et al., 2012; Jin et al., 2014; Zhang et al., 2018b). On average across testing datasets, each subject’s tractography data contained approximately 0.65 million streamlines, with an average streamline length of about 92 mm. The number of points along a streamline depended on the streamline’s length (points were stored approximately every 0.9 mm along streamlines). Tractography quality control was performed using whitematteranalysis software in SlicerDMRI (Norton et al., 2017; Zhang et al., 2020b). Tractography data were registered into the space of the training dataset using an affine transform produced by registering the b0 image of each subject to the mean T2-weighted image of the ORG atlas using 3D Slicer³ (Fedorov et al., 2012).

2.2. Two-stage framework for SWM parcellation

The goal of our two-stage framework (Fig. 2(b)) is to decompose the complicated SWM parcellation problem into two subproblems: a binary classification of SWM and DWM streamlines (stage one), and a multi-class classification of 198 SWM clusters and 198 SWM outlier categories (stage two). Stage one handles the classification of SWM vs DWM streamlines (which is non-trivial, as these categories cannot be easily distinguished by a simple threshold on streamline length as shown in Supplementary Fig. S1). Stage two handles the actual parcellation of the SWM and enables removal of streamline outliers from each cluster. Multi-stage frameworks that divide the target task into subtasks have also been applied successfully in many other medical imaging tasks (Wu et al., 2020a; Al-masni et al., 2020; Panda et al., 2021; Zhang et al., 2017b; Jain et al., 2020). In our case, the two stages have the same network structure for point cloud classification (Section 2.3) but different training data and procedures (Section 2.5). The two-stage inference process for SWM parcellation is described in Section 2.6.

Stage one takes whole-brain tractography streamlines as input, and outputs streamlines classified as SWM and DWM. In this way, stage one filters out DWM streamlines, which are not of interest for SWM parcellation. Stage one is trained with cross-entropy (CE) loss on the training dataset D₁. The task in stage one is relatively straightforward: stage one differentiates SWM and DWM streamlines, which have relatively large distances in the latent global feature space due to the spatial differences between SWM and DWM anatomy.

Stage two takes SWM streamlines as input, and it outputs streamlines classified into 198 SWM clusters and outliers. Stage two is trained with both cross-entropy loss and supervised contrastive loss (Khosla et al., 2020) on the training dataset D₂ (Section 2.4). The task in stage two is more challenging than that of stage one, because plausible streamlines from different clusters, as well as outlier streamlines, can have similar trajectories (Fig. 2(a)). Therefore, we incorporate supervised contrastive learning in this stage to enhance the feature learning (Section 2.4).

2.3. Point cloud network structure

The proposed network (Fig. 3) is based on PointNet (Charles et al., 2017), which is widely used for point cloud classification and segmentation (Guo et al., 2020). PointNet includes a shared multi-layer perceptron (MLP), a symmetric aggregation function, and fully connected (FC) layers, as well as data-dependent transformation nets. The transformation nets are designed to perform affine transformation of input point clouds into a canonical space (Charles et al., 2017). However, in our framework we remove the transformation nets to preserve important information about the spatial position of streamlines in the brain while improving the computational speed of the network.

Fig. 3.

Diagram of the point-cloud-based network architecture.

The proposed network structure (Fig. 3) is used in both stage one and stage two. Our network includes two parts: the encoder Enc(·) that extracts a non-linear global feature for each streamline and the classifier Cla(·) that predicts the streamline label using the extracted global feature in latent space. Specifically, for Enc(·), the input is the RAS (Right, Anterior, Superior) coordinates of streamline points, denoted as X = (x₁, x₂, …, x_n), where x_i is the 3-D coordinate of point i. Therefore, the dimension of X is n×3 (where n = 15, an appropriate number for streamline representations (Zhang et al., 2018b; O’Donnell and Westin, 2007; Zhang et al., 2020a)). Each point x_i is individually encoded by a shared MLP of three layers that have 64, 128, and 1024 output dimensions, respectively. Each layer includes batch normalization and rectified linear units (ReLU) activation. Therefore, an output X′ with n×1024 dimensions is generated, where X’ = (MLP(x₁), MLP(x₂), …, MLP(x_n)). Then, a symmetric function (max-pooling (Charles et al., 2017)) aggregates the encoded features X′ to form a 1024-dimension global shape feature g of the streamline. g is invariant to the order of points along a streamline, so that streamlines with equivalent forward and reverse point orderings are allowed to have the same global shape feature (representation). Finally, the Cla(·), consisting of three ordinary FC layers with sizes of 512, 256, and k (number of output classes), is used for streamline class prediction. The first and second FC layers are both followed by batch normalization and ReLU activation. The third FC layer followed by a softmax layer is used to output the prediction result (the class with the highest score).

2.4. Supervised contrastive learning

Supervised contrastive learning (SCL) (Khosla et al., 2020) extends self-supervised contrastive learning (Chen et al., 2020) to a fully-supervised mode by proposing a supervised contrastive loss, which aims to pull global features (outputs of Enc(·)) with the same class label closer in the latent space and push apart global features with different class labels. SCL has been shown to successfully improve the performance of supervised learning in computer vision (Schiffer et al., 2021; Khosla et al., 2020; Kopuklu et al., 2021; Zhong et al., 2021) and natural language processing (Han et al., 2021; Gunel et al., 2021; Huang et al., 2021) tasks.

In the proposed framework, we employ SCL in stage two to obtain more distinguishable global features for SWM clusters and outliers. In SCL training for Enc(·) (Fig. 4), a projector head Proj(·) (Khosla et al., 2020; Chen et al., 2020) is added on top of Enc(·). Proj(·) has two additional FC layers of sizes 1024 and 128 followed by a normalization layer. Therefore, the contrastive feature z of input X for calculating contrastive loss is formed as z = Proj(g) = Proj(Enc(X)). Proj(·) may retain more instance-specific streamline information in the global feature g to benefit downstream tasks (Chen et al., 2020).

Fig. 4.

Training procedure in stage two: the process of supervised contrastive learning and downstream task training

The supervised contrastive loss is defined as:

$$ℒ=\sum _{i\in I}{ℒ}_i=\sum _{i\in I}\frac{-1}{\left|P\left(i\right)\right|}\sum _{p\in P\left(i\right)}\log \frac{\exp \left(z_i\cdot z_p/\tau \right)}{{\sum }_{a\in A\left(i\right)}\exp \left(z_i\cdot z_a/\tau \right)}$$

where I is the streamline set in a training batch (i ∈ I ≡ {1, …, M}); P(i) is the streamline set that has the same class label as streamline i (p ∈ P(i)); A(i) is the set of all other streamlines in I except for streamline i (a ∈ A(i) ≡ I\{i}); z_i, z_p and z_a are contrastive features obtained from Proj(·) for streamlines i, p and a; τ (temperature) is a pre-defined hyperparameter set to be 0.1 as suggested in (Chen et al., 2020).

Since our clusters in training are bilateral, with similar shapes and cortical projections across both left and right hemispheres (see details in Supplementary Material 2), we employed a data augmentation technique, called bilateral augmentation (O’Donnell and Westin, 2007; Zhang et al., 2018a), for improving the performance of SCL. Data augmentation techniques have been shown to be helpful in contrastive learning (Chen et al., 2020; Khosla et al., 2020), where feature representations of augmented data from meaningful transformations are encouraged to be invariant to representations of the original sample. Bilateral augmentation generates a symmetric streamline in the other hemisphere, and the generated streamline has the same class label as the original streamline. Therefore, it encourages SCL to obtain the same representation for symmetric streamlines. In addition, the number of streamlines is doubled by applying bilateral augmentation in each training batch.

2.5. Two-stage model training

In stage one, Enc(·) and Cla(·) are trained together with CE loss. The learning rate is 0.001 (selected from a grid search of 0.01, 0.001, and 0.0001) and batch size is 1024 (selected from a grid search of 256, 512, 1024, 2048). Relatively large batch sizes are preferred since the number of streamlines (1 million) in our training dataset is large. Therefore, we chose large numbers for the grid search of batch sizes. In stage two, training includes two phases: contrastive learning and downstream learning (Khosla et al., 2020; Chen et al., 2020). In the contrastive learning phase, Enc(·) and Proj(·) are trained with supervised contrastive loss. The learning rate is 0.01 (selected from a grid search of 0.01, 0.001, and 0.0001), and the batch size is 3072 (selected from fine-tuning based on suggestions in (Khosla et al., 2020; Chen et al., 2020)). In the downstream learning phase, the parameters of Enc(·) are frozen as in (Khosla et al., 2020), and Proj(·) is untouched. Cla(·) takes g (the output of Enc(·)) as the input and is trained with CE loss for predicting streamline labels. The learning rate is 0.001, and the batch size is 1024. The selection process of learning rate and batch sizes for Cla(·) is the same as stage one. All training processes utilize Adam (Kingma and Ba, 2015) as the optimizer with no weight decay (selected from a grid search of 0, 0.0001, 0.0001). Training and validation were performed with Pytorch (v1.7.0) on a NVIDIA GeForce RTX 2080 Ti GPU machine.

2.6. Inference for SWM parcellation

For inference, the trained two-stage model is applied for SWM parcellation of unlabeled tractography data. The two-stage prediction process with trained networks works as follows. In stage one, streamlines with predicted labels of DWM are filtered out, and streamlines with predicted labels of SWM are retained for input to stage two. In stage two, the model takes these predicted SWM streamlines from stage one and classifies them into 198 SWM clusters and SWM outliers. Therefore, the final overall SWM parcellation result includes 198 SWM clusters and one non-SWM cluster (which contains SWM outliers and DWM streamlines).

3. Experiments and results

In this section, we first introduce experiments on the training dataset where streamline labels are available (Section 3.1), followed by the experiments on independently acquired testing datasets (Section 3.2).

3.1. Experiments on the training dataset

We conducted three experiments using 5-fold crossvalidation on the training dataset. First, we compared Sup-WMA to two deep-learning-based state-of-the-art (SOTA) methods. Second, we assessed the performance of ablated versions of our proposed SupWMA framework. Third, we analyzed the importance of different points along a streamline, e.g., endpoints or center points, for streamline classification within our proposed SupWMA framework.

3.1.1. Comparison with SOTA deep learning methods

We selected two SOTA deep-learning-based tractography parcellation methods (DeepWMA⁴ (Zhang et al., 2020a) and DCNN+CL+ATT⁵ (Xu et al., 2019)) for comparison to the proposed SupWMA method.

Deep white matter analysis (DeepWMA) (Zhang et al., 2020a) is designed for DWM parcellation using a Convolutional Neural Network (CNN) and streamline spatial coordinate features. DeepWMA has eight CNN layers and three FC layers. It uses a “FiberMap” input feature that converts spatial coordinates of streamlines into 3-channel images that enable streamlines with forward and backward point orders to have nearly equivalent representations. DeepWMA achieves consistent DWM parcellation results across populations (Zhang et al., 2020a).

DCNN+CL+ATT (Xu et al., 2019) is also designed for DWM parcellation. DCNN (Deep CNN, inspired by (He et al., 2016) and adapted from (Xu et al., 2018)) is employed with soft spatial attention (ATT) modules (Xu et al., 2015). During training, this method employs two losses: focal loss (Lin et al., 2017) to help in unbalanced datasets and center loss (CL) (Wen et al., 2016) to assist the network to obtain better streamline representations. This framework obtains satisfactory accuracy for predicting functionally important white matter pathways that should be protected in the surgery of epilepsy patients (Xu et al., 2019).

For both of these approaches, we trained their networks and tuned hyperparameters based on the suggested settings in their papers and released codes. For evaluation, we computed the accuracy and macro F1-score metrics, which have been widely used for the evaluation of tractography parcellation (Liu et al., 2019; Ngattai Lam et al., 2018; Zhang et al., 2020a; Xu et al., 2019). For each cross-validation fold, the accuracy of streamline classification was calculated, and the mean and standard deviation of the macro F1-score across 199 streamline classes were also reported. Finally, the average of each metric across the five folds was computed and is presented in Table 2. We also include the floating point operations (FLOPs), which measure the number of required operations performed for model inference (Charles et al., 2017), for evaluating the efficiency of each method.

Table 2.

Quantitative comparisons of SOTA deep-learning-based methods on the training dataset.

Methods	Accuracy	F1-score	FLOPs/streamline
DeepWMA	92.74%	73.35±7.93%	40.87M
DCNN+CL+ATT	95.31%	83.01±4.09%	36.82M
SupWMA	96.79%	88.79±2.91%	5.68M

Compared to the two SOTA methods (Table 2), SupWMA achieves the highest mean accuracy (96.79%) and macro F1-score (88.79%) with the lowest standard deviation. SupWMA outperforms DeepWMA and DCNN+CL+ATT by 4.05% and 1.48% in accuracy. Also, the F1-score of SupWMA is 15.44% and 5.78% higher than DeepWMA and DCNN+CL+ATT, respectively. Furthermore, the FLOPs of SupWMA are much lower than other SOTA methods. Overall, these results demonstrate the benefits of our novel deep learning framework to enable accurate and efficient SWM parcellation.

3.1.2. Ablation study

To analyze the effectiveness of our framework design, we performed an ablation study including the following: PointNet_orignial that is the original PointNet implementation (with transformation nets equipped), PointNet_two-stage that employs the original PointNet with a two-stage framework, SupWMA_two-stage that removes transformation nets from SupWMA_two-stage, and the proposed SupWMA_{two-stage+SCL} that uses SCL and the two-stage SWM parcellation framework. In addition to accuracy, F1-score, and FLOPs per streamline, we also calculated the inference time per streamline in 5-fold cross validation.

The ablation study (Table 3) shows that two-stage design in PointNet_two-stage increases accuracy and macro F1-score compared to PointNet_orignial. Although the FLOPs per streamline are doubled using PointNet_two-stage, the computational speed for SWM parcellation is not greatly reduced due to the two-stage parcellation framework that first removes a large number of DWM streamlines. Inference time per streamline slightly increases from 13.07ms for PointNet_orignial to 15.41ms for PointNet_two-stage with a two-stage parcellation framework. SupWMA_two-stage further improves the model performance and reduces the FLOPs as well as the inference time compared to PointNet_two-stage, showing the advantages of removing the transformation nets for our application. Finally, the proposed SupWMA_{two-stage+SCL} achieves the best performance, i.e., the highest accuracy, the highest F1-score, the lowest FLOPs, and the shortest inference time. These demonstrate the benefits of our network structure, two-stage parcellation framework and designed SCL that can extract highly discriminative global features of streamlines for parcellation.

Table 3.

Ablation experiments on the training dataset.

Methods	Accuracy	F1-score	FLOPs/streamline	Inference time/streamline
PointNetorignial	96.11%	86.48±3.47%	9.58M	13.07ms
PointNettwo-stage	96.39%	87.15±3.37%	19.16M	15.41ms
SupWMAtwo-stage	96.48%	87.47±3.18%	5.68M	6.82ms
SupWMAtwo-stage+SCL	96.79%	88.79±2.91%	5.68M	6.77ms

3.1.3. Critical streamline point indices for SWM classification

To study how our model classifies SWM streamlines, we explore the importance scores of point indices (Charles et al., 2017) for SWM streamline classification. Importance scores of point indices are quantified by the max pooling operation with our trained model. Since we implement 5-fold cross-validation for evaluating performance, five SupWMA trained models with the best average F1-score on five folds are selected. We input the validation data (20% of the entire dataset) into its corresponding trained model for each fold to get importance scores for all streamline points in the training dataset. Importance scores are obtained using the stage two model to focus this experiment on SWM (after removal of DWM). Importance scores for each point index are obtained from the max-pooling operation in the trained encoder. To be specific, we have 15 points for each streamline, and each point is encoded as a 1024-dimension (d) feature vector. For each dimension in the 1024-d feature vector, the max-pooling operation returns the maximum value and its point index out of 15 point values. Therefore, we can obtain the number of times (N_max) a point index is returned by the max-pooling operation. The possible N_max for each point index ranges from 0 to 1024. The higher the N_max, the more important the point is.

Fig. 5(a) presents the average importance score (normalized N_max) and its standard deviation for each point index across streamlines in the entire training dataset. It can be seen that the two streamline endpoints are the most critical points for classification of SWM streamlines. Fig. 5(b) shows the total importance scores of the two endpoints (red bar) and the other points (blue bar). While the two endpoints are the most informative, the other points along the streamline also provide important information for SWM classification.

Fig. 5.

(a) The average importance score and its standard deviation for each point index across streamlines. (b) The total importance scores of the two streamline endpoints and all of the other streamline points.

3.2. Experiments on independently acquired testing datasets

We also performed experiments on six independently acquired testing datasets (Table 1), where streamline labels were not available. In addition to the two deep-learning-based SOTA methods (DeepWMA and DCNN+CL+ATT, introduced in Section 3.1), we also compared whitematteranalysis (WMA)⁶ (O’Donnell and Westin, 2007; Zhang et al., 2018b; O’Donnell et al., 2012). WMA performs spectral clustering (O’Donnell and Westin, 2007) and entropy-based registration (O’Donnell et al., 2012) for white matter parcellation. Several deep learning papers (Wasserthal et al., 2019; Zhang et al., 2020a; Chen et al., 2021) have employed WMA as a comparison method. We note that WMA performs tractography parcellation by applying the same anatomically curated atlas (see (Zhang et al., 2018b) for details) used for generating our training data. WMA was used with default parameters in our study. Evaluation metrics and results on testing datasets are as follows.

3.2.1. SWM cluster identification rate

For all compared methods, we quantified the SWM cluster identification rate (CIR) (Table 4). The CIR is a metric that measures the success of white matter cluster identification in the absence of streamline labels (Zhang et al., 2020a, 2018b; Chen et al., 2021, 2022). In our study, a cluster was considered to be successfully detected if there were at least 10 streamlines (Zhang et al., 2018b, 2020a).

Table 4.

Cluster identification rates across methods for six testing datasets.

Methods	HCP	dHCP	ABCD	CNP	PPMI	BTP
WMA	97.86±2.37%	60.13±17.19%	89.92±5.14%	94.58±5.12%	96.35±3.12%	83.88±9.85%
DeepWMA	98.49±1.82%	64.27±15.36%	89.55±3.83%	95.63±2.50%	96.86±2.15%	85.02±7.63%
DCNN+CL+ATT	99.02±1.64%	65.35±16.84%	90.85±3.88%	96.23±2.74%	96.97±2.39%	85.21±8.15%
SupWMA	99.28±1.41%	68.29±17.24%	93.61±3.31%	96.76±2.66%	97.35±2.39%	86.62±7.92%

Table 4 indicates that our proposed SupWMA method has the highest CIR on all six testing datasets, which demonstrates a good generalization of SupWMA to subjects across ages and health conditions. Although the training data only includes subjects from the healthy young adult population, CIRs for SupWMA are over 93% on four testing datasets (HCP: healthy young adults; ABCD: healthy children; CNP: patients with psychiatric disorders; PPMI: healthy older adults and older adults with Parkinson’s disease). CIRs of all methods decrease on two highly challenging datasets: BTP (patients with space-occupying brain tumors) and dHCP (neonates). CIRs are around 85% for the BTP dataset due to patient-specific peritumoral effects. CIRs are below 70% in the dHCP dataset, where the neonate brains do not yet have fully myelinated SWM (Zhang et al., 2019b; Grotheer et al., 2022).

3.2.2. Cluster distance to atlas

For all compared methods, we quantified the SWM cluster distance to atlas (CDA) (Table 5). The CDA aims to quantify the degree to which identified SWM clusters are geometrically similar to the corresponding atlas (training dataset) clusters, where a low CDA indicates high geometric similarity to the atlas. CDA leverages the well-known minimum average directflip (MDF) pairwise streamline distance, which is widely applied in white matter parcellation/clustering (Garyfallidis et al., 2012; Zhang et al., 2018b; Chen et al., 2021). For each streamline in the cluster of each testing subject, we calculated MDF distances between that streamline and all streamlines of the corresponding cluster of the atlas, and we recorded the minimum MDF distance. Then we computed the CDA as the average of these streamline-specific minimum MDF distances. For each method, Table 5 shows the average CDA across all identified SWM clusters of all testing subjects in all datasets.

Table 5.

Cluster distance to atlas (unit: mm) across methods for six testing datasets.

Methods	HCP	dHCP	ABCD	CNP	PPMI	BTP
WMA	4.096±0.632	5.462±1.223	4.457±0.735	4.630±0.846	4.414±0.760	4.619±0.912
DeepWMA	4.035±0.627	5.544±2.182	4.275±0.718	4.350±0.725	4.273±0.688	4.576±1.141
DCNN+CL+ATT	3.882±0.518	4.841±2.369	4.141±0.893	4.179±0.667	4.097±0.543	4.361±2.220
SupWMA	3.902±0.530	4.785±0.813	4.133±0.578	4.190±0.596	4.123±0.574	4.303±0.730

As shown in Table 5, the distance between identified SWM clusters and the corresponding clusters in the ORG atlas is low for all methods across all testing datasets (under 5.6 mm). This demonstrates that identified clusters are highly geometrically similar to clusters in the atlas in general. In addition, SupWMA and DCNN+CL+ATT are two best performers (lower CDA values than other methods) in all datasets, and SupWMA performs the best on three very challenging datasets: dHCP (neonates), ABCD (children) and BTP (tumor patients) datasets.

3.2.3. Inter-subject parcellation variability

For all compared methods, we quantified the inter-subject parcellation variability (ISPV) (Table 6). The ISPV is a metric for evaluating performance of tractography parcellation (Roberts et al., 2017; Zhang et al., 2017a, 2018b) that evaluates if a cluster has a similar number of streamlines across subjects. The ISPV is defined as the coefficient of variation (standard deviation divided by the mean) of the number of streamlines per cluster. Therefore, the lower the ISPV, the higher the consistency of the number of streamlines of the corresponding cluster across subjects.

Table 6.

Inter-subject parcellation variability across methods for six testing datasets.

Methods	HCP	dHCP	ABCD	CNP	PPMI	BTP
WMA	0.7046±0.1912	1.7164±0.7683	0.8659±0.2826	0.8633±0.2902	0.8081±0.2458	1.1275±0.3560
DeepWMA	0.6496±0.1732	1.5832±0.6934	0.8169±0.3120	0.8170±0.3313	0.7750±0.2652	1.0996±0.4066
DCNN+CL+ATT	0.6424±0.1500	1.5372±0.6062	0.8106±0.2762	0.8143±0.2740	0.7748±0.2439	1.0901±0.3570
SupWMA	0.6323±0.1399	1.4759±0.5422	0.7799±0.2398	0.8008±0.2669	0.7636±0.2334	1.0626±0.3453

As shown in Table 6, SupWMA achieves highly consistent SWM parcellation results overall. It has lower ISPV than other compared methods on all testing datasets, especially on challenging dHCP (neonates), ABCD (children) and BTP (tumor patients) datasets, which are from individuals whose white matter anatomy is largely different from the training data population (healthy young adults). In addition, all methods have much higher ISPV values on dHCP than other datasets potentially because of rapid brain development during the perinatal period (Gilmore et al., 2018).

3.2.4. Population cluster heatmap

A population cluster heatmap displays the percentage of subjects that have streamlines present in each voxel. Here, population cluster heatmaps are generated to visualize and assess the cluster spatial overlap across subjects, as in (Zhang et al., 2020a). For computing heatmaps, clusters from different subjects and methods were registered (described in Section 2.1) into the space of a white matter atlas (Zhang et al., 2018b).

As shown in Fig. 6, identified clusters are visually comparable across the four SWM parcellation methods in general.Most clusters are identified robustly, as indicated by orange and red colors showing high cluster identification rates and high spatial overlap of clusters across subjects. However, more challenging clusters are less robustly identified by some methods, resulting in light orange and yellow colors that indicate lower identification of clusters that results in lower spatial overlap across subjects. We can observe that SupWMA obtains more consistent parcellation results than other methods, as indicated by visually higher (more orange and red) voxel values in the heatmaps, especially on the dHCP dataset (neonates) and the clusters in the fourth column of each method. For this fourth column, it can be observed that WMA and DeepWMA almost fail to detect the clusters in the fourth column on the dHCP dataset.

Fig. 6.

Visual comparison of the voxel-based population cluster heatmaps of four example clusters from results of WMA, DeepWMA, DCNN+CL+ATT, and SupWMA. Heatmaps depict the percentage of subjects that have streamlines passing through each voxel (see colorbar). The population mean T2 image of the ORG atlas is selected as the background image. Note that as these heatmaps show each population as a whole, the BTP heatmaps are not highly affected by patient-specific individual tumor locations.

3.2.5. Cluster spatial overlap

Cluster spatial overlap (CSO) is used to quantify if clusters from different methods are spatially comparable. In our experiments, CSO was quantified by comparing population cluster heatmaps using the weighted Dice (wDice) score, which has been used in many other tractography parcellation studies (Cousineau et al., 2017; Zhang et al., 2019a, 2020a). It extends the Dice score (Taha and Hanbury, 2015) to non-binary maps and assigns higher weighting to voxels with higher values.

As Table 7 shows, the CSO is generally high between methods across all datasets. This demonstrates that the results of SupWMA have good spatial overlap with other methods, and thus the results are comparable. SupWMA has higher CSO with the two deep learning methods (DWMA and DCNN+CL+ATT) than with the clustering-based method (WMA) on all testing datasets. A possible reason is that SupWMA and the two other deep learning methods were trained on the same dataset.

Table 7.

Cluster spatial overlap (CSO) between compared methods and our method.

	WMA | SupWMA	DWMA | SupWMA	DCNN+CL+ATT | SupWMA
HCP	0.8290±0.1669	0.8740±0.1247	0.9232±0.0795
dHCP	0.5131±0.3636	0.6401±0.3463	0.7601±0.2915
ABCD	0.6834±0.2432	0.7762±0.2058	0.8320±0.1652
CNP	0.7421±0.2306	0.8488±0.1731	0.9035±0.1267
PPMI	0.8007±0.1992	0.8647±0.1592	0.9169±0.1142
CNP	0.6984±0.2901	0.7605±0.2681	0.8415±0.2220

3.2.6. Subject-specific cluster visualization

We provide a visualization of the identified SWM clusters in an example individual subject for each dataset and across all SWM parcellation methods (Fig. 7).

Fig. 7.

Visualization of example individual clusters. For each testing dataset, a representative subject (with mean CIR closest to the population mean) is selected for visualization. The average CIR (across the six representative subjects) is displayed for each method. Parcellated clusters and b0 images in the background were transferred into the atlas space for visualization.

As shown in Fig. 7, all methods have relatively good performance on clusters shown in blue, cyan and magenta. However, for yellow clusters, SupWMA identifies more streamlines that are visually plausible than other methods, especially on three datasets: dHCP (neonates), ABCD (children) and BTP (tumor patients). Also, it can be seen that SupWMA has better performance on the dHCP dataset (neonates) than other methods across the four selected clusters.

3.2.7. Performance evaluation for clusters in peritumoral regions

We provide a visualization of patient-specific SupWMA identified clusters in the peritumoral region, in selected patients from the BTP dataset (Fig. 8). SupWMA can successfully identify peritumoral SWM clusters. Compared to clusters in the atlas, apparently atypical individual geometries are seen in identified SWM clusters, especially in patients with relatively larger tumors that may cause mass effect (e.g., the first and second rows of Fig. 8).

Fig. 8.

Visualization of SupWMA identified clusters located in the peritumoral region and corresponding clusters in the atlas. Four subjects with different diagnoses and tumor locations are selected, and clusters that have a large amount of overlap with edema and/or tumor are visualized. Brain tumors are shown in green, while each SWM cluster is displayed in a unique color.

To quantify the performance of SupWMA in the peritumoral region, we performed an additional experiment to calculate the identification rate of SWM clusters in the region of patient-specific tumor (n=39 patients) or edema (n=28 patients), as well as the identification rate of SWM clusters within the normal appearing white matter (n=39 patients). First, expert segmentations of edema and tumor were performed and transformed to atlas space (using the transforms obtained in Section 2.1), and atlas clusters intersecting these segmentations were identified. This gave a list of expected peritumoral clusters (i.e., clusters likely to be proximal to tumor and/or edema) for each patient. Second, patient-specific identification rates in the peritumoral region were computed as the percentage of the expected peritumoral clusters that were identified in each patient (Table 8). The identification rate of all other clusters (in the normal-appearing white matter) was also computed for each patient (Table 8). Finally, for purposes of selecting peritumoral clusters for visualization, the patient-specific volume of overlap of each identified cluster with edema and tumor was computed.

Table 8.

SWM cluster identification rates (CIRs) obtained using SupWMA for patient-specific clusters in the normal appearing white matter and the peritumoral region.

	Cluster identification rate
Clusters in normal-appearing white matter	92.98±6.76%
Patient-specific peritumoral clusters (in/near edema)	74.60±15.18%
Patient-specific peritumoral clusters (in/near tumor)	67.08±22.35%

As shown in Table 8, SupWMA obtains a high SWM CIR (92.98%) in the normal-appearing white matter. Although SWM connections may fail to be traced by tractography in peritumoral regions due to presence of mass lesions, SupWMA still identifies SWM clusters robustly with SWM CIR of 74.6% for for clusters in/near peritumoral edema, and 67.08% for clusters nearest the tumor.

3.2.8. Computation time

Computation time was tested on a Linux workstation with CPU (AMD Ryzen 5 3600) and GPU (NVIDIA RTX 2080 Ti) using a randomly selected subject (0.44 million streamlines). Table 9 shows the computation time across all compared SWM parcellation methods. Our method has the shortest computation time using CPU only and using CPU plus GPU. Also, benefiting from the efficient network structure, our method has the smallest computation time increase when changing from CPU plus GPU mode to CPU only mode.

Table 9.

Comparisons of computation time across methods.

Methods	CPU only	CPU + GPU
WMA	101min	—
DeepWMA	4min17s	2min22s
DCNN+CL+ATT	4min51s	1min51s
SupWMA	1min52s	1min16s

4. Discussion and conclusion

In this study, we proposed SupWMA, a novel deep learning framework for SWM parcellation, with successful application on datasets across image acquisition protocols, ages and health conditions. SupWMA enables efficient and consistent SWM parcellation.

SWM parcellation is a challenging problem due to the small diameter of SWM fasciculi, their short fiber length, their high variability across subjects, and their unique position near the cortex. As the first deep learning method for SWM parcellation, SupWMA can efficiently generate consistent SWM parcellation results to benefit downstream neuroscientific studies. These studies rely on diffusion features (FA, MD, etc.) of tracts from SWM parcellation (Wu et al., 2014; Reginold et al., 2016; d’Albis et al., 2018; Ji et al., 2018; Hatton et al., 2014).

Unlike other deep learning methods (Xu et al., 2019; Liu et al., 2019; Ngattai Lam et al., 2018) for white matter parcellation, SupWMA encodes geometric features of points on streamlines using a point-cloud-based network, which allows streamlines with equivalent forward and reverse orders to have the same representation. The two-stage framework design improves performance by decomposing the complicated SWM parcellation into two easier subtasks and enabling the SCL training to focus on the SWM. The incorporation of SCL enhances the SWM parcellation performance by helping the encoder to output discriminative representations of streamlines. Also, the removal of transformation nets preserves significant spatial features and reduces the computation time of training and testing processes. Compared to existing SOTA methods, SupWMA is the top performer on the training dataset and on six independently acquired testing datasets.

There are also several limitations and directions for future work. First, SupWMA has leveraged an anatomically curated atlas that includes SWM clusters discovered in a data-driven way in a population of healthy adults. While the employed atlas provides more comprehensive coverage of the SWM compared to other atlases (Guevara et al., 2017; Román et al., 2017), it should be noted that it focuses on the long- and medium-range SWM connections (over approximately 40mm in length). Recent research has shown the potential of high-resolution dMRI scans to enable tractography of shorter, highly curved u-fibers of the SWM (Song et al., 2014b; Ramos-Llordén et al., 2020). Benefitting from that, future work can utilize new training data (SWM atlases) to identify shorter connections, which are nearer the cortex and hence more variable across subjects (Román et al., 2021, 2022). Second, the proposed method is designed to enable studies of brain symmetry and asymmetry, as the definition of the SWM clusters is performed in a bilateral fashion. However, while corresponding SWM clusters in the left and right hemispheres have highly similar geometry and generally intersect corresponding cortical regions (Supplementary Material 2), leveraging additional anatomical information can potentially improve clustering results for the creation of new SWM atlases (Chen et al., 2021). Third, a potential limitation of our method is that only one atlas dataset was used for training. This strategy of training models on one dataset (e.g., HCP) and applying trained models on other datasets has been widely used in deep white matter parcellation/segmentation (Wasserthal et al., 2018, 2019; Zhang et al., 2020a; Lu et al., 2022) and dMRI registration (Zhang et al., 2022b). Furthermore, related work in SWM parcellation using non-deep methods traditionally uses one atlas dataset to define the parcellation (Guevara et al., 2017; Román et al., 2021, 2022). In this work, we have used an atlas that provides comprehensive coverage of the SWM and a fine parcellation scale (Zhang et al., 2018b). Fourth, another limitation of our method is that it may fail to identify SWM connections that are not traced by tractography (due to issues such as the presence of mass lesions, ongoing neurodevelopment, or differences in scan protocol). However, SupWMA identified clusters robustly for all datasets overall. In the challenging BTP dataset, cluster identification was generally robust even in the peritumoral region (Fig. 8 and Table 8). In the worst-case cluster performance in the challenging dHCP dataset, where neurodevelopment is ongoing and scanning is challenging, the cluster with the minimum identification rate was still found in 10% of subjects (Supplementary Material 3). Fifth, performing streamline tractography near the cortex is challenging due to the crossing of long-range fibers through superficial fiber systems (Reveley et al., 2015) and to the bias where tractography algorithms preferentially terminate in the crowns of gyri (Schilling et al., 2018; Van Essen et al., 2014). While addressing these fundamental anatomical challenges in performing tractography is beyond the scope of the current work, we note that future tractography methods that increase cortical coverage of streamlines (Bastiani et al., 2017; St-Onge et al., 2018; Wu et al., 2020b) will affect streamline endpoints. As the endpoints were shown to be important for SWM parcellation, this may improve depiction and discrimination of small bundles near the cortex. Sixth, more advanced point-cloud-based network structures (Qi et al., 2017; Wang et al., 2019) can be applied to improve the performance.

1. Novel superficial white matter parcellation method using diffusion MRI tractography

2. Accurate and efficient parcellation with a two-stage deep learning framework

3. Improved point-cloud-based network and supervised contrastive learning

4. Good generalization to subjects across lifespan and health conditions