Get a GRIP on Test Time Adaptation! - Group Robust Inference-Time Policy Optimization for Vision Models

Deep neural networks commonly suffer significant performance drops when the test data distribution deviates from the training domain—a frequent challenge in real-world deployments. This paper addresses test-time adaptation, wherein a model must autonomously adapt to distribution shifts using unlabeled test data only. We propose a new robust inference-time policy optimization (GRIP) framework that robustly couples self-supervised learning with a policy optimization objective, enabling the model to iteratively realign its parameters without relying on source data or labels. GRIP’s core innovation lies in a loss function that resists noisy pseudo-labels and stabilizes updates, bolstered by an augmentation pipeline and a student-teacher learning strategy. We show that GRIP converges under standard assumptions on the loss space, ensuring consistent adaptation in the presence of nontrivial domain shifts. Empirically, GRIP establishes a new state-of-the-art on CIFAR-10-C and CIFAR-100-C, outperforming competing methods (e.g., TENT, LAME, CoTTA) across all 15 benchmark corruption types. These performance gains showcase GRIP’s effectiveness in mitigating extreme domain shifts and highlight its broad applicability wherever real-time adaptation is essential for reliable, resilient AI systems.