Test-Time Augmentation for Cross-Domain Leukocyte Classification via OOD Filtering and Self-Ensembling

Area of Science:

• Machine Learning
• Computer Vision
• Medical Imaging

Background:

• Domain shift, caused by variations in data acquisition, presents a significant challenge in machine learning, especially in medical applications.
• Test-time augmentation (TTA) is a technique used to enhance model robustness by aggregating predictions from multiple augmented data samples.
• Standard TTA can inadvertently introduce out-of-distribution (OOD) samples, negatively impacting prediction accuracy.

Purpose of the Study:

• To develop a novel method for filtering OOD samples generated during TTA.
• To improve the robustness and accuracy of machine learning models in the presence of domain shift.
• To introduce a lightweight self-ensemble strategy for enhanced prediction fusion.

Main Methods:

• A filtering procedure was implemented to identify and remove OOD samples from TTA images based on their distance from the training data distribution.
• Retained TTA images were weighted inversely proportional to their distance from the training data distribution.
• A Self-Ensemble with Confidence strategy was employed, fusing predictions from original and filtered TTA samples using weighted soft voting.

Main Results:

• The proposed method demonstrated consistent improvements over standard TTA and baseline inference on cross-domain leukocyte classification tasks.
• Effectiveness was particularly notable under conditions of strong domain shift.
• Ablation studies and statistical analyses confirmed the significant contribution of each component of the proposed method.

Conclusions:

• The developed filtering and weighting strategy effectively mitigates the negative impact of OOD samples in TTA.
• The Self-Ensemble with Confidence approach provides a model-agnostic and computationally efficient way to improve prediction accuracy.
• This method offers broad applicability across various machine learning domains facing domain shift challenges.