Surface Vision Transformers: Attention-Based Modelling applied to Cortical Analysis

TL;DR

This work addresses learning long-range, surface-based brain signals on non-Euclidean geometries by translating Vision Transformer architecture to genus-zero surfaces. The Surface Vision Transformer (SiT) patches cortical data projected onto an icosphere and processes the patch sequence with multi-head self-attention, preserving spatial resolution. Evaluations on dHCP cortical metrics for PMA and GA prediction show SiT achieving competitive performance versus geometric DL baselines and demonstrating robustness to registration. Attention maps provide interpretable insight into perinatal cortical development, and self-supervised pretraining further improves performance, highlighting practical benefits for medical surface data analysis.

Abstract

The extension of convolutional neural networks (CNNs) to non-Euclidean geometries has led to multiple frameworks for studying manifolds. Many of those methods have shown design limitations resulting in poor modelling of long-range associations, as the generalisation of convolutions to irregular surfaces is non-trivial. Motivated by the success of attention-modelling in computer vision, we translate convolution-free vision transformer approaches to surface data, to introduce a domain-agnostic architecture to study any surface data projected onto a spherical manifold. Here, surface patching is achieved by representing spherical data as a sequence of triangular patches, extracted from a subdivided icosphere. A transformer model encodes the sequence of patches via successive multi-head self-attention layers while preserving the sequence resolution. We validate the performance of the proposed Surface Vision Transformer (SiT) on the task of phenotype regression from cortical surface metrics derived from the Developing Human Connectome Project (dHCP). Experiments show that the SiT generally outperforms surface CNNs, while performing comparably on registered and unregistered data. Analysis of transformer attention maps offers strong potential to characterise subtle cognitive developmental patterns.

Figures (6)

• Figure 1: Surface Vision Transformer (SiT) architecture. The cortical data is first resampled, using barycentric interpolation, from its template resolution (32492 vertices) to a sixth order icosphere (mesh of 40962 equally spaced vertices). The regular icosphere is divided into triangular patches of equal vertex count (b, c) that fully cover the sphere (not shown), which are flattened into feature vectors (d), and then fed into the transformer model.
• Figure 2: Example of attention maps for a single preterm neonate (born at 26.7 weeks GA) scanned at multiple time points (32.4 and 41.9 weeks PMA). T1w/T2w ratio and sulcal depth maps are shown in (a), and the attention maps for the 32 week and 42 week scans are shown in (b) and (c), respectively. (d) and (e) represent the thresholded attention maps for the 32 week and 42 week scans, respectively. The attention maps capture differences in the T1w/T2w ratio of the somatosensory (orange box), auditory (white box) and visual (green arrow) cortices, and changes in sulcal depth in the frontal (purple box) and temporo-occipital (magenta box) cortices. Additional illustrations are provided in Appendix \ref{['appendix:illustrations']}.
• Figure 3: Illustration of the Surface Vision Transformer model implementation a deconfouning strattegy of PMA for the GA prediction task.
• Figure 4: Multi-Head Self-Attention module (MSA) in purple, Self-attention layer in yellow and Feed Forward Network (FFN) Layer in pink. FFN layer expands the sequence dimension to $4D$ (e), then reduces it to $D$ after activation and dropout (f).
• Figure 5: Average attention maps in the PMA prediction task across heads for term and preterm neonates (a) when training from scratch and (b) training after self-supervision. With self-supervision, attention shifts away from the medial wall cut (an artefact of cortical surface processing), towards the lateral cortical surface. Attention is complementary across heads, and is highest in association cortical areas, which is consistent with the primary-to-association trajectory of cortical development in the perinatal period. Moreover, attention maps are similar between preterm and term neonates.