The Multiscale Surface Vision Transformer

TL;DR

The paper tackles the challenge of applying powerful transformer architectures to cortical surface data without incurring prohibitive computation from global self-attention. It introduces the Multiscale Surface Vision Transformer (MS-SiT), a Swin-transformer–inspired backbone that uses local window self-attention and a shifted window strategy to build hierarchical, high-resolution surface representations and a UNet-like segmentation path. Across neonatal phenotyping tasks on the dHCP dataset and cortical parcellation on UKB and MindBoggle, MS-SiT demonstrates superior phenotype prediction accuracy and competitive segmentation performance, while offering interpretable attention maps and robustness to space alignment. The work provides a practical, scalable backbone for surface deep learning with potential impact on clinical neuroimaging analyses and cortical mapping tasks, with code and models publicly available.

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.

Figures (6)

• Figure 1: [A] MS-SiT pipeline. The input cortical surface is resampled from native resolution (1) to an ico6 input mesh and partitioned (2). The sequence is then flattened (3) and passed to the MS-SiT encoder layers (4). The head (5) can be adapted for classification or regression tasks. [B] illustrates the patch merging operation (here from $I_4$ to $I_3$ grid). High-resolution patches are grouped by 4 to form patches of lower-resolution sampling grid [C] A Local-MHSA block is composed of two attention blocks: Window-MHSA and Shifted Window-MHSA.
• Figure 2: Top: Inflated surface showing mean dice scores shown for each of the DKT regions, for both MoNet and the pre-trained MS-SIT. Bottom: Boxplots comparing regional parcellation results. Asterisks denote statistical significance for one-sided paired t-test (pink: MS-SiT $>$ MoNet; green MoNet $>$ MS-SIT; ****: $p<0.0001$, **: $p<0.01$, *: $p<0.05$ ).
• Figure 3: MS-SiT segmentation pipeline. Input data is resampled and partitioned as in Figure \ref{['fig:swin_model']}. The $l=\{1,2,3,4\}$ levels of the segmentation pipeline are similar to the MS-SiT encoder levels (Figure \ref{['fig:swin_model']}). The patch partition layers reverse the patch merging procedure of the MS-SiT encoder, upsampling the spatial resolution of patches from $I_2 \rightarrow I_3 \rightarrow I_4 \rightarrow I_5$. Skip connections between levels are used. Finally, a spherical resampling layer resamples the final embeddings to an ico6 tessellation (40962 vertices), before the final segmentation prediction.
• Figure 4: [1] W-MHSA applies self-attention within a local window, defined by a fixed regular icosahedral partitioning grid. Two local windows are show here, delimited by the yellow and blue colours here. [2] SW-MHSA shifts patches such that local attention is computed between patches originally in different local windows.
• Figure 5: Comparison of normalised attention maps from the last layers of a SiT model Fawaz2021 (applying global attention for all layers) and an MS-SiT model, both trained for GA-template prediction. MS-SiT maps display highly specific attention patterns, compared to the SiT counterparts, focusing on characteristic landmarks of cortical development such as the sensorimotor cortex with low myelination in preterm (pink arrows) and high myelination in term (blue arrows).