Triamese-ViT for Brain Age Estimation

• Triamese-ViT is a 3D-aware adaptation of Vision Transformers that integrates axial, coronal, and sagittal MRI views to accurately estimate brain age.
• It employs separate Transformer branches for each anatomical plane and generates attention maps that localize age- and disease-relevant neuroanatomical structures.
• The model achieves state-of-the-art performance with improved MAE and interpretability, demonstrating clinical relevance over 3D CNN baselines.

Triamese-ViT is a 3D-aware adaptation of the Vision Transformer (ViT) architecture designed to improve brain age estimation from Magnetic Resonance Imaging (MRI) scans. By integrating orientation-specific ViT branches for axial, coronal, and sagittal views, Triamese-ViT captures 3D spatial context while producing interpretable attention maps that localize age- and disease-relevant neuroanatomical structures. This architecture demonstrates advances in predictive accuracy, model interpretability, and clinical relevance compared to prevalent 3D convolutional neural network (CNN) baselines (Zhang et al., 2024).

1. Model Architecture and Workflow

Triamese-ViT operates on a pre-processed 3D MRI volume $I \in \mathbb{R}^{H \times W \times D}$ by extracting three orthogonal stacks: axial ($I_x$), coronal ($I_y$), and sagittal ($I_z$), corresponding to axial, coronal, and sagittal anatomical planes, respectively. For each view $I_d$ ($d \in \{x, y, z\}$), each 2D slice is partitioned into non-overlapping $P \times P$ patches, yielding $N = \frac{H \times W}{P^2}$ patches per slice. Each patch is projected to a $D_{\rm emb}$-dimensional embedding via a learnable linear map $E$. The resulting sequence of patch embeddings is prepended with a learnable “class” token $I_x$0 and summed with a position embedding $I_x$1 to form $I_x$2.

Each branch is encoded by a dedicated $I_x$3-layer Transformer encoder, with non-shared parameters across views, following the standard LayerNorm–Multi-Head Attention–residual and LayerNorm–MLP–residual pipeline. Multi-head self-attention operates with $I_x$4 heads of dimension $I_x$5, using projection matrices $I_x$6 such that $I_x$7 and $I_x$8. The output is processed through residual and MLP blocks with GELU activation and hidden dimension $I_x$9.

Following the final Transformer layer, the class token $I_y$0 is input to a 2-layer MLP head to yield a scalar prediction $I_y$1 for each view. The predictions from all three branches are concatenated into $I_y$2 and fused via a 2-layer “Triamese” MLP: $I_y$3 where $I_y$4 denotes GELU activation.

2. Attention Map Synthesis and Interpretability

Each orientation-specific branch provides 2D attention maps for every slice, derived from the final-layer Transformer attention matrix $I_y$5 by extracting the class token row. These per-patch attention weights are reshaped to the 2D slice grid and upsampled to the original resolution.

To construct a comprehensive 3D attention map, 2D attention maps from all three views are back-projected and averaged at the voxel level: $I_y$6 where $I_y$7 denotes the 2D attention map for slice index along view $I_y$8. This synthesized $I_y$9 highlights volumetric regions associated with brain age prediction. Interpretation identifies attention focalization in regions such as basal ganglia, thalamus, and midbrain, recapitulating clinical patterns of neurodegeneration and age-sensitive change.

3. Optimization Objectives and Regularization

The primary training objective is Mean Absolute Error (MAE) between predicted age $I_z$0 and reference $I_z$1: $I_z$2 The Brain Age Gap (BAG) is monitored as $I_z$3. Optionally, a Spearman correlation loss ensures monotonic association between predictions and ground truth: $I_z$4 Weight decay $I_z$5 and dropout (0.1) within ViT blocks serve as regularization.

The total loss may be formulated as: $I_z$6

4. Training Protocol and Dataset Characteristics

Triamese-ViT is evaluated on a dataset of $I_z$7 T1-weighted MRIs from the IXI and ABIDE cohorts, covering ages 6–90. Preprocessing (FSL 5.10) includes nonlinear registration to MNI space, skull stripping, intensity normalization, and resampling to $I_z$8 voxels at 2 mm isotropic. For each MRI, three non-overlapping orthogonal views are extracted.

Data augmentation is applied with 50% probability per operation: 3D translation by up to $I_z$9 voxels, rotation in $I_d$0, and random axis flips. Model hyperparameters: patch size $I_d$1, embedding dimension $I_d$2, $I_d$3 heads, $I_d$4 Transformer layers, dropout=0.1, ViT head MLP dim=3072, and Triamese MLP with hidden sizes $I_d$5. Optimization uses Adam ($I_d$6, $I_d$7), initial learning rate $I_d$8, weight decay $I_d$9, batch size 100, for approximately 200 epochs with early stopping at 15% validation split. Hardware configuration uses 2 × NVIDIA V100 (32 GB) GPUs.

5. Performance Evaluation and Comparative Results

Triamese-ViT achieves a Mean Absolute Error of 3.84 years (95% CI [3.65, 4.03]), outperforming a 3D ResNet baseline (MAE: 4.01). The Spearman correlation between predicted and chronological age, $d \in \{x, y, z\}$0, is 0.90 (±0.02), superior to ResNet (0.83). The Spearman correlation between Brain Age Gap and age, $d \in \{x, y, z\}$1, is –0.29 (±0.03) for Triamese-ViT, compared to –0.31 for ResNet. Absolute error improvements are statistically significant by paired $d \in \{x, y, z\}$2-test ($d \in \{x, y, z\}$3) (Zhang et al., 2024).

6. Clinical Relevance and Model Interpretability

The synthesized 3D attention $d \in \{x, y, z\}$4 reliably localizes areas implicated in age and neurodegenerative pathology: basal ganglia (associated with disorders such as Parkinson’s and Huntington’s disease), thalamus (sensory and attentional hub), and midbrain (notably the substantia nigra, important for dopamine production). Occlusion analysis validates these interpretive findings — masking a seven-voxel cube in these regions increases MAE by over 20%. Slightly elevated attention in left-hemisphere structures is observed, plausibly reflecting lateralization and right-handed dominance in the cohort. Such region-specific explainability provides mechanistic insight for linking accelerated brain aging with early Alzheimer’s disease and related conditions.

7. Methodological Advantages, Limitations, and Prospective Directions

Triamese-ViT advances the field by capturing 3D context using three orientation-specific ViT branches, combining this with end-to-end interpretability through composite attention maps, while yielding state-of-the-art metrics compared to 3D CNN baselines. The architecture scales with only marginal cost over a single-view ViT, but is subject to certain constraints: lack of diagonal plane awareness and subtle 3D feature coupling beyond the orthogonal planes, greater memory usage from maintaining three separate Transformer backbones, and slower training speeds than some single-ViT or 3D CNN alternatives.

Future research directions include integrating oblique or learned 3D patching schemes for improved off-axis representation, extending the Triamese architecture to multimodal MRIs (e.g., diffusion or functional MRI), and leveraging contrastive pretraining on large, unlabeled MRI datasets to enhance initialization (Zhang et al., 2024).