# Improving Generalization in MRI-Based Deep Learning Models for Total Knee Replacement Prediction

arXiv: https://arxiv.org/abs/2504.19203 · HTML: https://mlsanity.com/paper/2504.19203/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 31 references that are in the map.
- First posted on arXiv 27 Apr 2025.

## Builds on (most cited first, 19 of 19)

- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Improved Regularization of Convolutional Neural Networks with Cutout](https://mlsanity.com/paper/1708.04552/) (2017)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Deep Domain Confusion: Maximizing for Domain Invariance](https://mlsanity.com/paper/1412.3474/) (2014)
- [Self-Challenging Improves Cross-Domain Generalization](https://mlsanity.com/paper/2007.02454/) (2020)
- [Learning to Optimize Domain Specific Normalization for Domain Generalization](https://mlsanity.com/paper/1907.04275/) (2019)
- [Causality-inspired Single-source Domain Generalization for Medical Image Segmentation](https://mlsanity.com/paper/2111.12525/) (2021)
- [Adversarially Adaptive Normalization for Single Domain Generalization](https://mlsanity.com/paper/2106.01899/) (2021)
- [Minimal-Entropy Correlation Alignment for Unsupervised Deep Domain Adaptation](https://mlsanity.com/paper/1711.10288/) (2017)
- [MixStyle Neural Networks for Domain Generalization and Adaptation](https://mlsanity.com/paper/2107.02053/) (2021)
- [PyTorch Metric Learning](https://mlsanity.com/paper/2008.09164/) (2020)
- [Realistic Adversarial Data Augmentation for MR Image Segmentation](https://mlsanity.com/paper/2006.13322/) (2020)
- [Adversarial Consistency for Single Domain Generalization in Medical Image Segmentation](https://mlsanity.com/paper/2206.13737/) (2022)
- [Three New Validators and a Large-Scale Benchmark Ranking for Unsupervised Domain Adaptation](https://mlsanity.com/paper/2208.07360/) (2022)
- [Disentangling A Single MR Modality](https://mlsanity.com/paper/2205.04982/) (2022)
- [Deep learning predicts total knee replacement from magnetic resonance images](https://mlsanity.com/paper/2002.10591/) (2020)
- [Predicting Knee Osteoarthritis Progression from Structural MRI using Deep Learning](https://mlsanity.com/paper/2201.10849/) (2022)
- [Unsupervised Domain Adaptation for Automated Knee Osteoarthritis Phenotype Classification](https://mlsanity.com/paper/2212.07023/) (2022)
- [Segmentation of Knee Bones for Osteoarthritis Assessment: A Comparative Analysis of Supervised, Few-Shot, and Zero-Shot Learning Approaches](https://mlsanity.com/paper/2403.08761/) (2024)

## Abstract (arXiv)

Knee osteoarthritis (KOA) is a common joint disease that causes pain and mobility issues. While MRI-based deep learning models have demonstrated superior performance in predicting total knee replacement (TKR) and disease progression, their generalizability remains challenging, particularly when applied to imaging data from different sources. In this study, we show that replacing batch normalization with instance normalization, using data augmentation, and applying contrastive loss improves generalization. For training and evaluation, we used MRI data from the Osteoarthritis Initiative (OAI) database, considering sagittal fat-suppressed intermediate-weighted turbo spin-echo (FS-IW-TSE) images as the source domain and sagittal fat-suppressed three-dimensional (3D) dual-echo in steady state (DESS) images as the target domain. The results demonstrated a statistically significant improvement in classification metrics across both domains by replacing batch normalization with instance normalization in the baseline model, generating augmented input views using the Global Intensity Non-linear (GIN) augmentation method, and incorporating a supervised contrastive loss alongside the classification loss to align representations of samples with the same label. The GIN method with contrastive loss performed better than all evaluated single-source domain generalization methods when using 3D instance normalization. Comparing GIN with and without contrastive loss (for both normalization types) showed that adding contrastive loss consistently led to better performance.
