MIMIR: Masked Image Modeling for Mutual Information-based Adversarial Robustness

TL;DR

This paper introduces MIMIR, a novel self-supervised adversarial training method for Vision Transformers that leverages mutual information constraints to improve robustness against various attacks and corruptions.

Contribution

It provides a theoretical MI analysis for ViT autoencoders and proposes MIMIR, a new MI-penalized adversarial pre-training approach tailored for ViTs.

Findings

MIMIR improves natural and robust accuracy on multiple datasets.

MIMIR outperforms state-of-the-art adversarial training methods on ImageNet-1K.

MIMIR demonstrates robustness against unforeseen and adaptive attacks.

Abstract

Vision Transformers (ViTs) have emerged as a fundamental architecture and serve as the backbone of modern vision-language models. Despite their impressive performance, ViTs exhibit notable vulnerability to evasion attacks, necessitating the development of specialized Adversarial Training (AT) strategies tailored to their unique architecture. While a direct solution might involve applying existing AT methods to ViTs, our analysis reveals significant incompatibilities, particularly with state-of-the-art (SOTA) approaches such as Generalist (CVPR 2023) and DBAT (USENIX Security 2024). This paper presents a systematic investigation of adversarial robustness in ViTs and provides a novel theoretical Mutual Information (MI) analysis in its autoencoder-based self-supervised pre-training. Specifically, we show that MI between the adversarial example and its latent representation in ViT-based autoencoders should be constrained via derived MI bounds. Building on this insight, we propose a self-supervised AT method, MIMIR, that employs an MI penalty to facilitate adversarial pre-training by masked image modeling with autoencoders. Extensive experiments on CIFAR-10, Tiny-ImageNet, and ImageNet-1K show that MIMIR can consistently provide improved natural and robust accuracy, where MIMIR outperforms SOTA AT results on ImageNet-1K. Notably, MIMIR demonstrates superior robustness against unforeseen attacks and common corruption data and can also withstand adaptive attacks where the adversary possesses full knowledge of the defense mechanism. Our code and trained models are publicly available at: https://github.com/xiaoyunxxy/MIMIR.