AugLoss: A Robust Augmentation-based Fine Tuning Methodology

TL;DR

AugLoss is a novel methodology that combines data augmentation and robust loss functions to improve deep learning model robustness against both noisy labels during training and feature distribution shifts during testing.

Contribution

It introduces a unified approach, AugLoss, that simultaneously addresses label noise and distribution shifts, advancing robustness in deep learning models.

Findings

Outperforms previous methods on real-world corrupted datasets

Enhances robustness against both training noise and test shifts

Demonstrates significant accuracy improvements in experiments

Abstract

Deep Learning (DL) models achieve great successes in many domains. However, DL models increasingly face safety and robustness concerns, including noisy labeling in the training stage and feature distribution shifts in the testing stage. Previous works made significant progress in addressing these problems, but the focus has largely been on developing solutions for only one problem at a time. For example, recent work has argued for the use of tunable robust loss functions to mitigate label noise, and data augmentation (e.g., AugMix) to combat distribution shifts. As a step towards addressing both problems simultaneously, we introduce AugLoss, a simple but effective methodology that achieves robustness against both train-time noisy labeling and test-time feature distribution shifts by unifying data augmentation and robust loss functions. We conduct comprehensive experiments in varied settings of real-world dataset corruption to showcase the gains achieved by AugLoss compared to previous state-of-the-art methods. Lastly, we hope this work will open new directions for designing more robust and reliable DL models under real-world corruptions.