ShapKO: Shapley-Adaptive Modality Knockout for Robust Multimodal Learning
2026-07-10 • Computer Vision and Pattern Recognition
Computer Vision and Pattern RecognitionMachine Learning
AI summaryⓘ
The authors address a problem where medical AI models struggle when some types of input data are missing, which happens often in real life. They note that sometimes the models rely too much on one type of data, ignoring others, making them less reliable when data is incomplete. To fix this, the authors created ShapKO, a method that adaptively decides which types of data to temporarily hide during training, based on how important each type is. This helps the model learn to use all available data better and perform well even when some data is missing. They tested ShapKO on multiple medical tasks and found it improved performance without needing to change the model itself.
multimodal learningmodality dominancemissing data robustnesstraining-time maskingShapley valuesadaptive knockoutmedical AIclinical classificationsurvival predictioncancer detection
Authors
Nusrat Binta Nizam, Fengbei Liu, Sunwoo Kwak, Minh Nguyen, Ruining Deng, Mert R. Sabuncu
Abstract
Multimodal medical models often degrade when inputs are missing, a common scenario in real-world clinical workflows. Separately, even when all modalities are present, modality dominance is observed during training, where optimization over-relies on a highly predictive modality and undertrains complementary sources, resulting in poor robustness under partial availability. While training-time modality knockout improves missing-modality robustness, existing approaches use static masking rates that cannot adapt to evolving modality utility during training. We introduce ShapKO (Shapley-Adaptive Modality Knockout), a dynamic training strategy that learns modality-specific knockout probabilities based on validation utility. ShapKO periodically evaluates performance across modality subsets, estimates modality importance via Shapley values, and updates masking probabilities to suppress dominant modalities more frequently. This adaptive process promotes complementary representations, while requiring no architectural modifications. We evaluate ShapKO on three datasets covering multitask clinical classification, survival prediction, and cancer detection. ShapKO consistently improves performance under modality absence and yields interpretable trajectories of learned masking behavior. Code is available at: https://github.com/sumona00/ShapKO