In-Model Merging for Enhancing the Robustness of Medical Imaging Classification Models

TL;DR

This paper introduces InMerge, a novel in-model merging technique that enhances medical imaging model robustness by selectively merging similar kernels during training, outperforming traditional training methods across multiple datasets.

Contribution

The paper proposes the first in-model merging method that improves CNN robustness in medical imaging by merging kernels within a single model during training, with analytical insights and empirical validation.

Findings

InMerge improves model robustness significantly.

It outperforms normally trained models on multiple datasets.

The approach is effective across different CNN architectures.

Abstract

Model merging is an effective strategy to merge multiple models for enhancing model performances, and more efficient than ensemble learning as it will not introduce extra computation into inference. However, limited research explores if the merging process can occur within one model and enhance the model's robustness, which is particularly critical in the medical image domain. In the paper, we are the first to propose in-model merging (InMerge), a novel approach that enhances the model's robustness by selectively merging similar convolutional kernels in the deep layers of a single convolutional neural network (CNN) during the training process for classification. We also analytically reveal important characteristics that affect how in-model merging should be performed, serving as an insightful reference for the community. We demonstrate the feasibility and effectiveness of this technique for different CNN architectures on 4 prevalent datasets. The proposed InMerge-trained model surpasses the typically-trained model by a substantial margin. The code will be made public.