Learning to Generate Gradients for Test-Time Adaptation via Test-Time Training Layers

TL;DR

This paper introduces a learned optimizer called Meta Gradient Generator (MGG) for test-time adaptation, which effectively utilizes historical gradient information to improve model robustness and speed in adapting to new data distributions.

Contribution

The paper proposes a novel learned optimizer with a gradient memory layer for test-time adaptation, outperforming prior methods in accuracy and efficiency.

Findings

Surpasses state-of-the-art on ImageNet-C, R, Sketch, and A.

Achieves 7.4% accuracy improvement over previous SOTA.

Provides faster adaptation with fewer data and iterations.

Abstract

Test-time adaptation (TTA) aims to fine-tune a trained model online using unlabeled testing data to adapt to new environments or out-of-distribution data, demonstrating broad application potential in real-world scenarios. However, in this optimization process, unsupervised learning objectives like entropy minimization frequently encounter noisy learning signals. These signals produce unreliable gradients, which hinder the model ability to converge to an optimal solution quickly and introduce significant instability into the optimization process. In this paper, we seek to resolve these issues from the perspective of optimizer design. Unlike prior TTA using manually designed optimizers like SGD, we employ a learning-to-optimize approach to automatically learn an optimizer, called Meta Gradient Generator (MGG). Specifically, we aim for MGG to effectively utilize historical gradient information during the online optimization process to optimize the current model. To this end, in MGG, we design a lightweight and efficient sequence modeling layer -- gradient memory layer. It exploits a self-supervised reconstruction loss to compress historical gradient information into network parameters, thereby enabling better memorization ability over a long-term adaptation process. We only need a small number of unlabeled samples to pre-train MGG, and then the trained MGG can be deployed to process unseen samples. Promising results on ImageNet-C, R, Sketch, and A indicate that our method surpasses current state-of-the-art methods with fewer updates, less data, and significantly shorter adaptation iterations. Compared with a previous SOTA method SAR, we achieve 7.4% accuracy improvement and 4.2 times faster adaptation speed on ImageNet-C.