Training with Explanations Alone: A New Paradigm to Prevent Shortcut Learning

TL;DR

This paper introduces Training with Explanations Alone (TEA), a novel method where a student model learns to ignore biases by matching explanation heatmaps from a teacher, improving bias resistance and generalization in AI models.

Contribution

The paper proposes TEA, a new training paradigm that uses explanation heatmaps for bias mitigation without needing a segmenter or direct output loss.

Findings

TEA outperforms 14 state-of-the-art bias mitigation methods.

TEA improves generalization to unseen hospitals and datasets.

TEA effectively reduces background and foreground bias influence.

Abstract

Application of Artificial Intelligence (AI) in critical domains, like the medical one, is often hampered by shortcut learning, which hinders AI generalization to diverse hospitals and patients. Shortcut learning can be caused, for example, by background biases -- features in image backgrounds that are spuriously correlated to classification labels (e.g., words in X-rays). To mitigate the influence of image background and foreground bias on AI, we introduce a new training paradigm, dubbed Training with Explanations Alone (TEA). TEA trains a classifier (TEA student) only by making its explanation heatmaps match target heatmaps from a larger teacher model. By learning from its explanation heatmaps, the TEA student pays attention to the same image features as the teacher. For example, a teacher uses a large segmenter to remove image backgrounds before classification, thus ignoring background bias. By learning from the teacher's explanation heatmaps, the TEA student learns to also ignore backgrounds -- but it does not need a segmenter. With different teachers, the TEA student can also resist bias in the image foreground. Surprisingly, by training with heatmaps alone the student output naturally matches the teacher output -- with no loss function applied to the student output. We compared the TEA student against 14 state-of-the-art methods in 5 datasets with strong background or foreground bias, including Waterbirds and an X-Ray dataset for COVID-19 and pneumonia classification. The TEA student had better resistance to bias, strongly surpassing state-of-the-art methods, and generalizing better to hospitals not seen in training.