Teaching Invariance Using Privileged Mediation Information

TMLR Paper5812 Authors

Abstract: The performance of deep neural networks often deteriorates in out-of-distribution settings due to relying on easy-to-learn but unreliable spurious associations known as shortcuts. Recent work attempting to mitigate shortcut learning relies on a priori knowledge of the shortcuts and invariance penalties, which are difficult to enforce in practice. To address these limitations, we study two causally-motivated methods that efficiently learn models that are invariant to shortcuts by leveraging privileged mediation information. We first adapt concept bottleneck models (CBMs) to incorporate mediators -- intermediate variables that lie on the causal path between input features and target labels -- resulting in a straightforward extension we call Mediator Bottleneck Models (MBMs). One drawback of this method is that it requires two potentially large models at inference time. To address this issue, we propose Teaching Invariance using Privileged Mediation Information (TIPMI), a novel approach which distills knowledge from a counterfactually invariant teacher trained using privileged mediation information to a student predictor that uses non-privileged, easy-to-collect features. We analyze the theoretical properties of both estimators, showing that they promote invariance to multiple unknown shortcuts and can result in better finite-sample efficiency compared to commonly used regularization schemes. We empirically validate our theoretical findings by showing that TIPMI and MBM outperform several state-of-the-art methods on one language and two vision datasets.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

Added a discussion about the front-door criterion in Section 3.1, as requested by reviewer GaHH.
Added a discussion about why the assumption that mediators are fully recoverable from the input features are realistic, as requested by reviewer GaHH.
Added additional experiments in Appendix F where the shortcut effects the mediator, as requested by reviewer dzPR.
Added to the discussion of why overfitting is a problem for TIPMI without cross-fitting, as requested by reviewer rRJz.
Clarified why TIPMI and MBM perform better than baselines, why superfluous mediators can impact MBMs performance, and our comparison of model sizes, as requested by reviewer rRJz.
Added links to proofs in the appendix, added suggestions to strengthen the paper, and fixed typos, as requested by reviewer rRJz.

Assigned Action Editor: ~Francesco_Locatello1

Submission Number: 5812