Relevance-driven Input Dropout: an Explanation-guided Regularization Technique

TL;DR

This work tackles overfitting in data-scarce settings by introducing RelDrop, an explanation-guided input regularization that occludes the most attribution-sensitive input regions to force the model to leverage a broader set of features. By computing attribution maps and applying a balanced mask to inputs in both 2D images and 3D point clouds, RelDrop acts as an informed regularizer requiring only one extra backward pass per batch. Empirically, RelDrop improves generalization and robustness to occlusion across CIFAR, ImageNet, ModelNet40, and ShapeNet, with notable gains in zero-shot performance on distribution-shifted ImageNet variants and qualitative evidence of more distributed, semantically meaningful feature usage. The method offers a practical pathway to tighter generalization by integrating attribution signals into data augmentation and can be extended to other architectures and modalities in future work.

Abstract

Overfitting is a well-known issue extending even to state-of-the-art (SOTA) Machine Learning (ML) models, resulting in reduced generalization, and a significant train-test performance gap. Mitigation measures include a combination of dropout, data augmentation, weight decay, and other regularization techniques. Among the various data augmentation strategies, occlusion is a prominent technique that typically focuses on randomly masking regions of the input during training. Most of the existing literature emphasizes randomness in selecting and modifying the input features instead of regions that strongly influence model decisions. We propose Relevance-driven Input Dropout (RelDrop), a novel data augmentation method which selectively occludes the most relevant regions of the input, nudging the model to use other important features in the prediction process, thus improving model generalization through informed regularization. We further conduct qualitative and quantitative analyses to study how Relevance-driven Input Dropout (RelDrop) affects model decision-making. Through a series of experiments on benchmark datasets, we demonstrate that our approach improves robustness towards occlusion, results in models utilizing more features within the region of interest, and boosts inference time generalization performance. Our code is available at https://github.com/Shreyas-Gururaj/LRP_Relevance_Dropout.

Figures (6)

• Figure 1: Schematic overview of workflow, illustrated for image data. As shown by the relevance heatmaps, Panel (a) depicts an Unregularized Model that, without data augmentation, overfits by heavily depending on a few highly relevant features. To nudge the model into learning from a larger set of informative features, repeatedly computes attributions during training, and then augments the data by selectively masking the (currently) most relevant parts of the input (Panel (b)). The impact of training with is clear in Panel (c), where heatmaps at the bottom reflect how the Regularized Model utilizes a larger set of features to make predictions.
• Figure 2: Regularization effect of on ResNet50 overfitting. The curves show the augmentation-free baseline (baseline), (), and two variations of (), with attribution hyperparameters $\varepsilon=0.8$ (performing worst) and $\varepsilon=0.001$ (performing best) on the CIFAR-100 dataset. A moving average (more solid lines) with window size $5$ is visualized along with the raw data for both the training (left) and test (right) curves. -trained model generalizes better to the test data compared to both and the baseline. Note that the y-axis scale of training (left) and test (right) figures are adjusted to the respective ranges of training and test accuracies.
• Figure 3: Qualitative effects of on model decision-making. attributions are visualized (deeper red indicates higher importance). Columns show, from left to right, the ImageNet pre-trained ResNet50 model, the baseline finetuned without input augmentation, and the models finetuned with and , respectively. The class labels for the input images are shown on the left. We observe increased reliance on within-object features for the finetuned model. Although this effect does not hold for all samples equally (last row), we confirm an overall improvement via the quantitative evaluation of in Table \ref{['table:ResNet_results_table_Image']}.
• Figure 4: Robustness of -trained models to ordered point removal on ModelNet40 WuSKYZTX15. -trained models are more robust against point removal, especially when a balance between random and relevance-guided input data augmentation is employed ($\alpha=0.5$ and $\beta=0.5$).
• Figure A.1: Attribution maps indicate the dropout strategy for different dropout parameters during training. The points highlighted in red indicate the $(x, y, z)$ coordinates of the points being dropped and replaced by $(0, 0, 0)$, and the ones highlighted in blue are unaltered and retained for the next epoch of training. An "Aeroplane" sample is considered for illustration with a constant value of $\alpha=0.5$ and varying values of $\beta=(0.15, 0.5, 0.85)$. Left: represents the dropout with parameters, $\alpha=0.5$, $\beta=0.15$, Center: represents the dropout with parameters, $\alpha=0.5$, $\beta=0.15$ and the, Right: represents the dropout with parameters, $\alpha=0.5$, $\beta=0.85$.