Shortcut Invariance: Targeted Jacobian Regularization in Disentangled Latent Space

TL;DR

This paper introduces a novel latent-space regularization technique that reduces neural network reliance on shortcut features by injecting targeted anisotropic noise, improving out-of-distribution generalization without needing shortcut labels.

Contribution

The authors propose a Jacobian regularization method in disentangled latent space that preserves full representations while reducing shortcut reliance, enhancing OOD robustness.

Findings

Achieves state-of-the-art OOD performance on shortcut-learning benchmarks.

Does not require shortcut labels or conflicting samples for training.

Effectively flattens decision boundaries along shortcut axes.

Abstract

Deep neural networks are prone to learning shortcuts, spurious correlations present in the training data that undermine out-of-distribution (OOD) generalization. Most prior work mitigates shortcut learning through input-space reweighting, either relying on explicit shortcut labels or inferring shortcut structure from heuristics such as per-sample loss. Moreover, these approaches typically assume the presence of some shortcut-conflicting examples in the training set, an assumption that is often violated in practice, particularly in medical imaging where data is aggregated across institutions with different acquisition protocols. We propose a latent-space method that views shortcut learning as over-reliance on shortcut-aligned axes. In a disentangled latent space, we identify candidate shortcut-aligned axes via their strong correlation with labels and reduce classifier reliance on them by injecting targeted anisotropic noise during training. Unlike prior latent-space based approaches that remove, project out, or adversarially suppress shortcut features, our method preserves the full representation and instead impose functional invariance by regularizing the classifier's sensitivity along those axes. We show that injecting anisotropic noise induces targeted Jacobian and curvature regularization, effectively flattening the decision boundary along shortcut axes while leaving core feature dimensions largely unaffected. Our method achieves state-of-the-art OOD performance across standard shortcut-learning benchmarks without requiring shortcut labels or shortcut-conflicting samples.