Abstract:Deep neural networks have repeatedly been shown to be non-robust to the uncertainties of the real world, even to naturally occurring ones. A vast majority of current approaches have focused on data-augmentation methods to expand the range of perturbations that the classifier is exposed to while training. A relatively unexplored avenue that is equally promising involves sanitizing an image as a preprocessing step, depending on the nature of perturbation. In this paper, we propose to use control for learned models to recover from distribution shifts online. Specifically, our method applies a sequence of semantic-preserving transformations to bring the shifted data closer in distribution to the training set, as measured by the Wasserstein distance. Our approach is to 1) formulate the problem of distribution shift recovery as a Markov decision process, which we solve using reinforcement learning, 2) identify a minimum condition on the data for our method to be applied, which we check online using a binary classifier, and 3) employ dimensionality reduction through orthonormal projection to aid in our estimates of the Wasserstein distance. We provide theoretical evidence that orthonormal projection preserves characteristics of the data at the distributional level. We apply our distribution shift recovery approach to the ImageNet-C benchmark for distribution shifts, demonstrating an improvement in average accuracy of up to 14.21% across a variety of state-of-the-art ImageNet classifiers. We further show that our method generalizes to composites of shifts from the ImageNet-C benchmark, achieving improvements in average accuracy of up to 9.81%. Finally, we test our method on CIFAR-100-C and report improvements of up to 8.25%.

Ambient Diffusion: Learning Clean Distributions from Corrupted Data

Consistent Diffusion Meets Tweedie: Training Exact Ambient Diffusion Models with Noisy Data

Integrating Amortized Inference with Diffusion Models for Learning Clean Distribution from Corrupted Images

An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

Learning Diffusion Model from Noisy Measurement using Principled Expectation-Maximization Method

Ambient Diffusion Posterior Sampling: Solving Inverse Problems with Diffusion Models trained on Corrupted Data

GSURE-Based Diffusion Model Training with Corrupted Data

Slight Corruption in Pre-training Data Makes Better Diffusion Models

How much is a noisy image worth? Data Scaling Laws for Ambient Diffusion

DiracDiffusion: Denoising and Incremental Reconstruction with Assured Data-Consistency

Extracting Training Data from Diffusion Models

Unlearnable Examples for Diffusion Models: Protect Data from Unauthorized Exploitation

A High-Quality Robust Diffusion Framework for Corrupted Dataset

Efficient Diffusion-Driven Corruption Editor for Test-Time Adaptation

Improving Diffusion Models's Data-Corruption Resistance using Scheduled Pseudo-Huber Loss

Mask, Stitch, and Re-Sample: Enhancing Robustness and Generalizability in Anomaly Detection through Automatic Diffusion Models

Exploiting Diffusion Prior for Out-of-Distribution Detection

DC4L: Distribution Shift Recovery via Data-Driven Control for Deep Learning Models

Patch-Based Diffusion Models Beat Whole-Image Models for Mismatched Distribution Inverse Problems

Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images

Steerable Conditional Diffusion for Out-of-Distribution Adaptation in Medical Image Reconstruction