DAFT: Distilling Adversarially Fine-tuned Models for Better OOD Generalization

TL;DR

DAFT is a novel method that distills adversarially robust features from a teacher model to improve out-of-distribution generalization, outperforming existing methods especially on smaller networks.

Contribution

The paper introduces DAFT, a new distillation approach that leverages adversarial training to enhance OOD robustness of models, with modifications to standard adversarial training for better teacher guidance.

Findings

DAFT outperforms state-of-the-art OOD methods on DomainBed benchmarks.

DAFT improves robustness especially for smaller neural networks.

Significant gains up to 6% over ERM and distillation baselines.

Abstract

We consider the problem of OOD generalization, where the goal is to train a model that performs well on test distributions that are different from the training distribution. Deep learning models are known to be fragile to such shifts and can suffer large accuracy drops even for slightly different test distributions. We propose a new method - DAFT - based on the intuition that adversarially robust combination of a large number of rich features should provide OOD robustness. Our method carefully distills the knowledge from a powerful teacher that learns several discriminative features using standard training while combining them using adversarial training. The standard adversarial training procedure is modified to produce teachers which can guide the student better. We evaluate DAFT on standard benchmarks in the DomainBed framework, and demonstrate that DAFT achieves significant improvements over the current state-of-the-art OOD generalization methods. DAFT consistently out-performs well-tuned ERM and distillation baselines by up to 6%, with more pronounced gains for smaller networks.