Contrastive Representation Learning on TabTransformer Latent Features for Imbalanced Post-Stroke mRS Classification

Abstract

The increasing availability of electronic clinical records has enabled new opportunities for predictive modeling in healthcare. However, clinical data are characterized by heterogeneous patient information, limited sample availability, and highly imbalanced outcome distributions, which may hinder the ability of predictive models to capture complex patterns and generalize across underrepresented groups. This work investigates deep learning-based representation learning strategies for predicting the final post-stroke functional state derived from the modified Rankin Scale (mRS) using clinical tabular data. The proposed methodology learns informative patient representations while addressing class imbalance without altering the original data distribution, avoiding limitations of conventional resampling strategies. Specifically, the proposed framework combines a TabTransformer encoder with Supervised Contrastive Learning and Focal Loss to jointly optimize discriminative representation learning and imbalance-aware classification. By obtaining more separable latent representations and reducing the bias toward majority-class predictions, the approach aims to improve the identification of patients with unfavorable functional outcomes. Experimental results show that the proposed model achieves the highest accuracy (0.92) and macro-averaged F1-score (0.80), with balanced minority-class precision (0.67) and recall (0.64), unlike the other models, which showed a trade-off between these metrics. These results demonstrate improved minority-class identification while maintaining overall predictive performance in clinical outcome prediction.