CARformer: Class-Aware Representation Learning for Small-Cohort, Imbalanced Psychiatric Disorder Classification in Brain MRI

Structural MRI (sMRI) has revealed morphological correlates associated with psychiatric disorders, motivating neuroimaging-based diagnostic support tools. However, multi-class classification from sMRI remains challenging in small-cohort psychiatric settings, where site-level datasets often contain only a few hundred subjects and exhibit class imbalance, and where subtle inter-class differences further complicate discrimination. We propose a class-aware representation learning framework (CARformer) that adapts domain-pretrained encoders for 3D sMRI multi-class psychiatric classification. The framework includes (1) a class-aware Bag-of-Queries (CA-BoQ) module that learns class-aligned query representations to aggregate class-specific dis-criminative evidence under inter-class differences, and (2) a Class-Aware Contrastive (CAC) learning strategy that strengthens class separation under imbalance. We evaluated the framework across two sites from a public brain MRI dataset to assess robustness under differences in diagnostic composition, class imbalance, and imaging acquisition protocols. Our results demonstrate consistent improvements in macro AUC and AUPRC over existing standard 3D Convolutional Neural Network (CNN) backbones and domain-pretrained fine-tuning frameworks across both sites, supporting the practical utility of class-aware adaptation of domain pre-training in small-cohort, imbalanced psychiatric imaging tasks. Source code is available at: https://github.com/xyfu33/CARformer.

Lecture Notes in Computer Science

MICCAI 2026

16885

0302-9743

N/A

N/A

10