The Multiscale Surface Vision Transformer

Authors: simon  dahan
logan z j williams
daniel  rueckert
emma c robinson
l z j dahan
d  williams
e c rueckert
  robinson
Venue: MIDL_io_2024_Conference
Published: 2024-01-29 17:18:26

Abstract: Surface meshes are a favoured domain for representing structural and functional information on the human cortex, but their complex topology and geometry pose significant challenges for deep learning analysis. While Transformers have excelled as domain-agnostic architectures for sequence-to-sequence learning, the quadratic cost of the self-attention operation remains an obstacle for many dense prediction tasks. Inspired by some of the latest advances in hierarchical modelling with vision transformers, we introduce the Multiscale Surface Vision Transformer (MS-SiT) as a backbone architecture for surface deep learning. The self-attention mechanism is applied within local-mesh-windows to allow for high-resolution sampling of the underlying data, while a shifted-window strategy improves the sharing of information between windows. Neighbouring patches are successively merged, allowing the MS-SiT to learn hierarchical representations suitable for any prediction task. Results demonstrate that the MS-SiT outperforms existing surface deep learning methods for neonatal phenotyping prediction tasks using the Developing Human Connectome Project (dHCP) dataset. Furthermore, building the MS-SiT backbone into a U-shaped architecture for surface segmentation demonstrates competitive results on cortical parcellation using the UK Biobank (UKB) and manually-annotated MindBoggle datasets. Code and trained models are publicly available at https://github.com/metrics-lab/surface-vision-transformers.

Main paper: 
 introduction: 
 In recent years, there has been an increasing interest in using attention-based learning methodologies in the medical imaging community, with the Vision Transformer (ViT) (Dosovitskiy et al., 2020) emerging as a particularly promising alternative to convolutional methods. The ViT circumvents the need for convolutions by translating image analysis to a sequence-to-sequence learning problem, using self-attention mechanisms to improve the modelling of long-range dependencies. This has led to significant improvements in many medical imaging tasks, where global context is crucial, such as tumour or multi-organ segmentation (Tang et al., 2021; Ji et al., 2021; Hatamizadeh et al., 2022) . At the same time, there has been a growing enthusiasm for adapting attention-based mechanisms to irregular geometries where the translation of the convolution operation is not trivial, but the representation of the data as sequences can be straightforward, for instance for protein modelling (Atz et al., 2021; Jumper et al., 2021; Baek et al., 2021) or functional connectomes (Kim et al., 2021) . Similarly, vision transformers (ViTs) have been recently translated to the study of cortical surfaces (Dahan et al., 2022) , by re-framing the problem of surface analysis on sphericalised meshes as a sequence-to-sequence learning task and by doing so improving the modelling of long-range dependencies in cortical surfaces. Transformer models have also emerged as a promising tool for modelling various cognitive processes, such as language and speech (Millet et al., 2022; Défossez et al., 2023) , vision (Tang et al., 2023) , and spatial encoding in the hippocampus (Whittington et al., 2021) .Despite promising results on high-level prediction tasks, one of the main limitations of the ViT remains the computational cost of the global self-attention operation, which scales quadratically with sequence length. This limits the ability of the ViT to capture fine-grained details and to be used directly for dense prediction tasks. Various strategies have been developed to overcome this limitation, including restricting the computation of self-attention to local windows (Fan et al., 2021; Liu et al., 2021) or implementing linear approximations (Wang et al., 2020; Xiong et al., 2021) . Among these, the hierarchical architecture of the Swin Transformer (Liu et al., 2021) has emerged as a particularly favoured candidate. This implements windowed local self-attention, alongside a shifted window strategy that allows cross-window connections. Neighbouring patch tokens are progressively merged across the network, producing a hierarchical representation of image features. This hierarchical strategy has shown to improve performance over the global-attention approach of the standard ViT, and has already found applications within the medical imaging domain (Hatamizadeh et al., 2022) . Cheng et al. (2022) attempted to adapt such windowed local attention to the study of cortical meshes. However, attention windows were defined as the vertices forming the hexagonal patches of a low-resolution grid, but not the patch features. This restricts the feature extraction with self-attention to a small number of vertices on the mesh and greatly limits the local feature extraction capabilities of the model.In this paper, we therefore introduce the Multiscale Surface Vision Transformer (MS-SiT) as a novel backbone architecture for surface deep learning. The MS-SiT takes inspiration from the Swin Transformers model and extends the Surface Vision Transformers (SiT) (Dahan et al., 2022) to a hierarchical version that can serve for any high-level or dense prediction task on sphericalised meshes. First, the MS-SiT introduces a local-attention operation between surface patches and within local attention-windows defined by the subdivisions of a high-resolution sampling grid. This allows for the modelling of fine-grained details of cortical features (with sequences of up to 20,480 patches). Moreover, to preserve the modelling of long-range dependencies between distant regions of the input surface, the MS-SiT adapts the shifted local-attention approach, introduced in (Liu et al., 2021) , by shifting the sampling grid across the input surface. This allows propagation of information between neighbouring attention-windows, achieving global attention at a reduced computational cost; however, it is challenging to implement due to the irregular spacing and sampling of vertices on native surface meshes. We evaluate our approach on neonatal phenotype prediction tasks derived from the Developing Human Connectome Project (dHCP), as well as on cortical parcellation for both UK Biobank (UKB) and manually-annotated MindBoggle datasets. Our proposed MS-SiT architecture strongly surpasses existing surface deep learning methods for predictions of cortical phenotypes and achieves competitive performance on cortical parcellation tasks, highlighting its potential as a holistic deep learning backbone and a powerful tool for clinical applications. 
methods: 
 Backbone The proposed MS-SiT adapts the Swin Transformer architecture (Liu et al., 2021) to the case of cortical surface analysis, as illustrated in Figure 1 . Here, input data X ∈ R |V 6 |×C (C channels) is represented on a 6th-order icospheric (ico6) tessellation: I 6 = (V 6 , F 6 ), with |V 6 | = 40962 vertices and |F 6 | = 40962 faces. This data is first partitioned into a sequence of |F 5 | = 20480 non-overlapping triangular patches:T 5 = {t 1 5 , t 2 5 , ..t |F 5 | 5 } (with t i 5 ⊂ V 6 , |t i 5 | = |t 5 | = 6), by patching the data with ico5:I 5 = (V 5 , F 5 ), |V 5 | = 10242, |F 5 | = 20480 (Figure 1 A.2).Imaging features for each patch are then concatenated across channels, and flattened to produce an initial sequence:X 0 = X 0 1 , ..., X 0 |F 5 | ∈ R |F 5 |×(C|t 5 |) (Figure 1A.3). Trainable positional embeddings, LayerNorm (LN) and a dropout layer are then applied, before passing it to the MS-SiT encoder, organised into l = {1, 2, 3, 4} levels.At each level of the encoder, a linear layer projects the input sequence X l to a 2 (l-1) ×Ddimensional embedding space:X l emb ∈ R |F 6-l |×2 (l-1) D .Local multi-head self-attention blocks (local-MHSA), described in section 2, are then applied, outputting a transformed sequence of the same resolution (X l M HSA ∈ R |F 6-l |×2 (l-1) D ). This is subsequently downsampled through a patch merging layer, which follows the regular downsampling of the icosphere, to merge clusters of 4 neighbouring triangles together (Figure 1B ), generating output:X l out ∈ R |F 6-l-1 |×2 (l+1) D .This process is repeated across several layers, with the spatial resolution of patches progressively downsampled from I 5 → I 4 → I 3 → I 2 , but the channel dimension doubling each time. In doing so, the MS-SiT architecture produces a hierarchical representation of patch features, with respectively |F 5 | = 20480, |F 4 | = 5120, |F 3 | = 1280, and |F 2 | = 320 patches. In the last level, the patch merging layer is omitted (see Figure 1 ) and the sequence of patches is averaged into a single token, and input to a final linear layer, for classification or regression (Figure 1A.5) . Inspired by previous work (Cao et al., 2021) , the segmentation pipeline employs a UNet-like architecture, with skip-connections between encoder and decoder layers, and patch partition instead of patch merging applied during decoding. An illustration of the pipeline is provided in Figure 3 , Appendix A.Local Multi-Head Self-Attention blocks are defined similarly to ViT blocks (Dosovitskiy et al., 2020): as successive multi-head self-attention (MHSA) and feed-forward (FFN) layers, with LayerNorm (LN) and residual layers in between (Figure 1C ). Here, a Window-MHSA (W-MHSA) replaces the global MHSA of standard vision transformers, applying self-attention between patches within non-overlapping local mesh-windows. To provide the model with sufficient contextual information, this attention window is defined by an icosahedral tessellation three levels down from the resolution used to represent the feature sequence. This means that at level l, while the sequence is represented by I 6-l , the attention windows correspond to the non-overlapping faces F 6-(l+3) defined by I 6-(l+3) . For example, at level 1 the features are input at ico5, and local attention is calculated between the subset of 64 triangular patches that overlap with each face of ico2 (F 2 ), see Figure 1 .B.1. Only in the last layer, is attention not restricted to local windows but applied globally to the I 2 grid, allowing for global sharing of information across the entire sequence. More details of the parameterisation of window attention grids is provided in the Appendix A, Table 3 . This use of local self-attention significantly reduces the computational cost of attention at levell, from O(|F 6-l | 2 ) to O(w l |F 6-l |) with w l << |F 6-l |.Self-Attention with Shifted Windows Cross-window connections are introduced through Shifted Window MHSA (SW-MHSA) modules, to improve the modelling power of the local self-attention operations. These alternate with the W-MHSA, and are implemented by shifting all the patches in the sequence I 6-l , at level l by w s positions, where w s is a fraction of the window size w l (typically w l = 64). In this way, a fraction of the patches of each attention window now falls within an adjacent window (see Figure 4 ). This preserves the cost of applying self-attention in a windowed fashion, whilst increasing the models representational power by sharing information between non-overlapping attention windows. The W-MHSA and SW-MHSA implementation can be summarised as follows:Xl = W-MSA(X l emb ) + X l emb Z l = FFN( Xl ) + Xl Ẑl = SW-MSA(Z l ) + Z l X l M HSA = FFN( Ẑl ) + Ẑl (1)Here X l emb and X l M HSA correspond to input and output sequences of the local-MHSA block at level l. Residual connections are referred to by the + symbol.Training details Augmentation strategies were introduced to improve regularisation and increase transformation invariance. This included implementing random rotational transforms, where the degree of rotation about each axis was randomly sampled in the range ∈ [-30 • , +30 • ] (for the regression tasks) and ∈ [-15 • , +15 • ] (for the segmentation tasks). In addition, elastic deformations were simulated by randomly displacing the vertices of a coarse ico2 grid to a maximum of 1/8th of the distance between neighbouring points (to enforce diffeomorphisms (Fawaz et al., 2021) ). These deformations were interpolated to the high-resolution grid of the image domain, online, during training. The effect of tuning the parameters of the SW-MHSA modules is presented in Table 4 and reveals that the best results are obtained while shifting half of the patches. 
experiments & results: 
 All experiments were run on a single RTX 3090 24GB GPU. The AdamW optimiser (Loshchilov and Hutter, 2017) with Cosine Decay scheduler was used as the default optimisation scheme, more details about optimisation and hyper-parameters tuning in Appendix B.2. A combination of Dice Loss and CrossEntropyLoss was used for the segmentation tasks and MSE loss was used for the regression tasks. Surface data augmentation was randomly applied with a probability of 80%. If selected, one random transformation is applied: either rotation (50%) or non-linear warping (50%). For all regression tasks, a custom balancing sampling strategy was applied to address the imbalance of the data distribution.Data from the dHCP comes from the publicly available third release1 (Edwards et al., 2022) and consists of cortical surface meshes and metrics (sulcal depth, curvature, cortical thickness and T1w/T2w myelination) derived from T1-and T2-weighted magnetic resonance images (MRI), using the dHCP structural pipeline, described by (Makropoulos et al., 2018) and references therein (Kuklisova-Murgasova et al., 2012; Schuh et al., 2017; Hughes et al., 2017; Cordero-Grande et al., 2018; Makropoulos et al., 2018) . In total 580 scans were used from 419 term neonates (born after 37 weeks gestation) and 111 preterm neonates (born prior to 37 weeks gestation). 95 preterm neonates were scanned twice, once shortly after birth, and once at term-equivalent age.Aug 1 , where the MS-SiT models were compared against several surface convolutional approaches and three versions of the Surface Vision Transformer (SiT) using different grid sampling resolutions. The MS-SiT model consistently outperformed all other models across all prediction tasks (PMA and GA) and data configurations (template and native). Specifically, on the PMA task, the MS-SiT model outperformed other models by over 54% compared to Spherical UNet (Zhao et al., 2019) , 13% to MoNet (Monti et al., 2016) , and 12% to the SiT (ico3) (average over both data-configurations), achieving a prediction error of 0.49 MAE on template data, which is within the margin of error of age estimation in routine ultrasound (typically, 5 to 7 days on average). On the GA task, the MS-SiT model achieved an even larger improvement with 49%, 43%, and 21% reduction in MAE relative to Spherical UNet, MoNet, and SiT (ico3), respectively. Importantly, the model demonstrated much greater transformation invariance, with only a 5% drop in performance between the template and native configurations, compared to 53% for Spherical UNet, and 10% for MoNet. Results2. (Zhao et al., 2019) 3. (Monti et al., 2016) also revealed a significant benefit to using the SW-MHSA with a 16% improvement over the vanilla version on GA predictions.Data & Tasks Cortical segmentation was performed using 88 manually labelled adult brains from the MindBoggle-101 dataset (Klein and Tourville, 2012) 4 , annotated into 31 regions using a modified version of the Desikan-Killiany (DK) atlas (Desikan et al., 2006) , which delineates regions according to features of cortical shape. Surface files were processed with the Ciftify pipeline (Dickie et al., 2019) , which implements HCP-style post-processing including file conversion to GIFTI and CIFTI formats, and MSM Sulc alignment (Robinson et al., 2014 (Robinson et al., , 2018) ) 5 . Separately, FreeSurfer annotation parcellations (based on a standard version of the DK atlas with 35 regions) were available for 4000 UK Biobank subjects, processed according to (Alfaro-Almagro et al., 2018) . These were used for pretraining. In both cases, datasets were split into 80%/10%/10% sets. As the annotations characterise folding patters, we used shape-based cortical metrics as input features: sulcal depth and curvature maps. Results are presented in Table 2 . The MS-SiT was compared against three other gDL approaches for cortical segmentation: Adv-GCN, a graph-based method optimized for alignment invariance (Gopinath et al., 2020) , SPHARM-net (Ha and Lyu, 2022) , a spherical harmonic-based CNN method, and MoNet, which learns filters by fitting mixtures of Gaussians on the surface (Monti et al., 2016) . MoNet achieved the best dice results overall, while MS-SiT superforms the two other gDL networks. However, a per region box plot (Fig 2 ) of its performance relative to the MS-SIT shows this is largely driven by improvements to two large regions. Overall, MoNet and the MS-SIT differ significantly for 10 out of 32 regions, with MS-SIT outperforming MoNet for 6 of these. We also evaluated the performance of the MS-SiT model by providing it with more inductive biases, via transfer learning from a model first trained on the larger UKB dataset (achieving 0.94 dice for cortical parcellation), increasing slightly the final performance. 
discussion: 
 The novel MS-SiT network presents an efficient and reliable framework, based purely on self-attention, for any biomedical surface learning task where data can be represented on sphericalised meshes. Unlike existing convolution-based methodologies translated to study general manifolds, MS-SiT does not compromise on filter expressivity, computational complexity, or transformation equivariance (Fawaz et al., 2021) . Instead, with the use of local and shifted attention, the model is able to effectively reduce the computational cost of applying attention on larger sampling grids, relative to (Dahan et al., 2022) , improving phenotyping performance, and performing competitively on cortical segmentation. Compared to convolution-based approach, the use of attention allows for the retrieval of attention maps, providing interpretable insights into the most attended cortical regions (Fig 5 ), and the methodology's robustness to transformations enables it to perform well on both registered and native space data, removing the need for spatial normalisation using image registration.6. Run on a different train/test split, (Gopinath et al., 2020) 7. (Ha and Lyu, 2022) 8. (Monti et al., 2016) http://www.developingconnectome.org.https://MindBoggle.info/data13 datasets failed due to missing files 
 
 Appendix:  


Decision: All reviewers found that this paper is novel and addresses an important topic that the MIDL community will be interested in. The paper is written well and demonstrates intensive results clearly.

Reviews: review: paper_summary: paper_summary: This paper presents a multi scale surface vision transformer to process cortical surface data in neuroimaging. This model reduces the computational cost of fulll-resolution self-attention by combining local-mesh-window for within region and shifted-window self-attention for between region information sharing. The authors test their MS-SiT backbone on a phenotyping prediction task on dHCP data and on cortical segmentation on MindBoggle and UK Biobank data.
main_review: main_review: Thanks a lot for addressing and clarifying the minor points I had mentioned. I think this is a well-written paper with a great presentation of methods and results. The method itself could have a great impact on the neuroimaging field.
strength_weakness: strength_weakness: I do not see a major weakness of the paper worth mentioning here. 
There are only a few minor points that come to my mind. These mostly concern some details in the experiments and a bit of background on the task settings. I will list those below under detailed comments.
questions: questions: See comments above.
paper_summary: paper_summary: This paper presents a method to perform brain surface analysis using deep learning. The main contribution is a network architecture implementing graph NN on a cortical surfaces, based on transformers. The authors describe a multi scale transformer implementation that achieves very good results.
main_review: main_review: Strong paper with an interesting method and good performance. The authors did a good job addressing the reviewers' comment. They have added details on the position encoding part, which is quite important.
strength_weakness: strength_weakness: Not many weaknesses, this is a good paper. 

Quite a few details are missing but lots is already crammed into the limited available space. Extra details are provided in supplementary materials which are very useful. 

One big detail that is missing is how the positional encoding is performed. I assume some sort of spherical angular encoding, but this ought to be explained.
questions: questions: Explain position encoding.
paper_summary: paper_summary: The authors propose Multiscale SiT (MS-SiT), a model that draws inspiration from Swin Transformers and builds on the Surface Vision Transformers (SiT) framework from their previous research. The method adapts the shifted local-attention technique by shifting the sampling grid across the surface allowing model to preserve global interactions between distant regions on the input surface.
main_review: main_review: The authors have impressively addressed my concerns and other comments, significantly improving the quality of the paper. Their attentive and thorough responses validate my initial assessment of their work.
strength_weakness: strength_weakness: Limitations of their work is not discussed. Baseline comparison are limited and does not justify the need for sophisticated methods. 

The lack of use of a widely used dataset, such as the HCP young adult, limits the study's comparability with existing literature. This is particularly crucial for phenotyping prediction.

Minor: 
Missing some key information, i.e. input data used for analysis. For cortical parcellation task, the authors clearly state which input features are used. However, it is unclear what is used for phenotypic prediction. All these metrics: sulcal depth, curvature, cortical thickness and T1w/T2w myelination? If so why not in the parcellation task.
questions: questions: Given the emphasis on "Multiscale" in the title and the inspiration drawn from Swin Transformers, could you provide a more in-depth explanation of the multiscale concept as it applies to your work? Specifically, how do hierarchical representations contribute to the objectives and findings of your paper? 

Could you provide further explanation on the following points:

1. What motivated the choice of a U-net like architecture?
2. Can you elaborate on the computational footprint of your approach?
3. What are the time and memory requirements to train these models?

Comments: comment: We would like to thank the reviewer for their careful review of our submission and their many constructive comments. 

**Points to address:**

*Multiscale concept*: The main motivation for introducing the ‘multiscale’ concept in the surface transformer architecture comes from the fact that understanding phenotypes requires a global understanding of the brain and therefore using deep learning architectures that can effectively connect distant areas of the brain although not adjacent. However, one of the main limitations that were pointed out in the SiT paper [Dahan et al. 2022, MIDL] is the cost of the self-attention operation that limits the resolution at which patches are sampled, and therefore the modelling of subtle variations across the cortex, specifically for dense prediction tasks. The multiscale framework proposed in the Swin transformer paper [Ze Liu et al 2021] (but also concurrent work such as [*Multiscale Vision Transformers*
, H.Fan et 2021]) allows to model the data across scales – i.e. at the highest resolution the model can learn to model high-resolution features and textures, but only apply attention within regional local windows. By merging neighbouring patches of extracted features across layers, the model downsamples the resolution at which the patches are represented but increases the receptive field of the self-attention operation and eventually learns to attend across the whole brain. This is essential for modelling subtle cortical patterns and, at the same time, long-range dependencies between cortical regions.    

*U-Net architecture*: The use of a U-Net architecture was motivated by a previous paper that implemented vision transformers for 2D medical image segmentation  (Hu Cao et al 2021). In general UNet-style architectures are widely used for such tasks as they balance robust localisation of objects with sharp segmentation of boundaries. They achieve this by inserting high-frequency information back into the network through skip connections, applied between the encoding and decoding paths. We added this reference in the method section. 

*Computational footprint, Time and Memory requirements*: We thank the reviewer for the opportunity to provide essential information regarding the computational cost and performance of our method. The MS-SiT in its default parametrisation as in Table 3 has 27.5M trainable parameters. This is significantly larger than the SiT-tiny used in [Dahan et al 2022, MIDL] -  with 5M params but smaller than the SiT-base model and its 86M parameters [Dahan et al 2022, MIDL] - and seems reasonable compared to today’s standard with transformer models. On a standard 24G memory NVIDIA GPU, we can use a relatively large batch size of 28 (taking 22G of memory), which allows for fast convergence. We added a figure in Appendix B to compare the training and validation losses of an MS-SiT trained on PMA prediction against a SiT-tiny (ico2). The MS-SiT achieves convergence after 450 epochs (corresponding to 2.5 hours of training) compared to the SiT model (around 800 epochs). Additionally, as pointed in the method section, the MS-SiT local self-attention operation reduces the complexity from $\mathcal{O}(|F_{6-l}|^2)$ to $\mathcal{O}(w_{l}|F_{6-l}|)$ with $w_{l} << |F_{6-l}|$, where $|F_{6-l}|$ corresponds to the length of the sequence and $w_{l}$ the size of the attention window. Therefore, where the SiT is limited in the sequence length because of the quadratic complexity of the self-attention operation, here the MS-SiT can be trained on an input sequence of 20480 tokens, in a similar time to an SiT-tiny (ico2) (trained with 320 patches) and improves the performance.

**Other comments**

*Missing information*: We apologise for the lack of clarity around the input data used. For phenotype predictions, we used four cortical metrics (sulcal depth, curvature, cortical thickness and myelin maps). We added a sentence in the main text to address the confusion. For the parcellation task, we trained the network purely on shape-based cortical metrics – specifically sulcal depth (which captures coarse-scale patterns of cortical folding – predominantly primary folds) and mean curvature (secondary and tertiary folds), as the DK atlas [Desikan et al. 2006] parcels the brain into regions that align with the primary folding patterns of the brain. We clarify this point in the main text section 3.2.

*HCP dataset*: The use of HCP data would be indeed essential and will be integrated in future work. We already have tried to experiment on cognitive score regression notably fluid intelligence prediction. Preliminary results seem to indicate that the MS-SiT can achieve a good performance using memory task contrast maps. We obtained a 0.42 Pearson correlation score (CV - 5 folds) for ~1000 HCP subjects. This can be compared with GNN methods trained on functional connectivity data derived from individual ICA cortical parcellation, achieving 0.33 Pearson correlation [Dahan et al 2021].
We would like to thank the reviewer for their careful review of our submission and their constructive comments. We corrected the typos that were present in our first submissions and added some missing details in the Appendix. 

**Points to address:**

*Explain positional embedding*: As done in the original SiT paper [Dahan et al 2022, MIDL] positional embeddings (PE) are learnt. This is implemented through initialising the PE tokens at random, to be optimised during the training. The latest literature in vision transformers seems to indicate that fixed positional embedding (notably in the form of sine-cosine embeddings) is becoming the default setting. This could be tested in future work. We updated the method section to add the mention of trainable PE and added some additional details regarding PE in Appendix A.2. Also, we have never thought about using spherical angular encoding, as suggested by the reviewer. This is a great suggestion and will be tested in future work.
We would like to thank the reviewer for their careful review of our submission and their constructive comments and feedback on the methodology's impact.  

**Points to address:**

*Explain the choice of optimizer and the scheduler*: We ran extensive experiments comparing the use of SGD and AdamW optimisers, across tasks, and with a wide range of different hyper-parameters. We found that AdamW with a learning rate and cosine decay (eta_min = 1e-5) with a linear learning rate warm-up was the best performing overall. This is consistent with standard practices in vision transformer papers that use AdamW and cosine warm-up scheduling [A. Gotmare et al 2018, Z. Liu et al 2021]. Briefly, for the regression tasks we found that AdamW (with LR=1e-5 and 1000 warm-up iterations) and SGD (with a small learning rate, typically 1e-5) were performing similarly. However, for the segmentation task, we found that SGD was unstable such that convergence was difficult to achieve even with a learning rate scheduling scheme (e.g. LR step decay). Overall, we selected the AdamW with cosine warm-up as the default optimisation setting. We added more details regarding the optimisation of the models as well as the training hyper-parameters used for the different tasks in Appendix B 2.2.  

*Details on the segmentation loss*: For optimisation of the segmentation task, we used the Dice Cross Entropy loss implementation of Monai. While developing the model, we compared the results against the use of dice or cross-entropy only and as expected it seems that the Dice cross-entropy loss was performing better in terms of final Dice. We did not use any particular weighting between the losses (default MONAI parameter), but we believe that this can be an important hyperparameter to finetune depending on the segmentation task. Some additional information was also added in Appendix B 2.2.

*Intuition behind the cortical segmentation task and details about the data*: The DK atlas [Desikan et al. 2006] parcels the brain into regions that align with the primary folding patterns of the brain. We therefore trained the network purely on shape-based cortical metrics – specifically sulcal depth (which captures coarse-scale patterns of cortical folding – predominantly primary folds) and mean curvature (secondary and tertiary folds). Both these metrics are output by default from the FreeSurfer surface processing pipeline. We clarify this point in the main text section 3.2.