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

TL;DR

This work tackles the problem of unreliable gradients in test-time adaptation (TTA) by introducing a learning-to-optimize framework that automatically generates reliable gradients. The authors design a Meta Gradient Generator (MGG) built on a Gradient Memory Layer (GML) that memorizes historical gradient information and uses it to produce refined gradients for online updates. The method is pretrained on a small unlabeled set and then deployed to unseen data, updating only lightweight components and enabling fast, data-efficient adaptation. Empirical results across ImageNet-C, ImageNet-R, Sketch, and ImageNet-A show state-of-the-art accuracy with shorter adaptation times and smaller data budgets, demonstrating practical viability for real-world distribution shifts. The approach combines end-to-end L2O training with a memory-augmented gradient refinement strategy to stabilize online optimization under noisy unsupervised signals.

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.

Figures (3)

• Figure 1: Method Differences. We devise an automatically learned meta gradient generator to optimize the original gradients produced by a TTA loss to be more reliable.
• Figure 2: An overall illustration of MGTTA, in which we design a two-step meta gradient generator (MGG) to generate optimized gradients for TTA. Given a trained model $f(\cdot;\theta)$, for each batch of test samples, we first calculate predictions and obtain gradients by backpropagation. Then, in Step 1 MGG first memorizes gradients and then in Step 2 MGG generates optimized gradients based on the historical gradient information. Finally, the model parameters $\theta$ are updated using the optimized gradients. Here, the learnable parameters within $\theta$ only involve norm layers and the rest are kept frozen during adaptation.
• Figure 3: Convergence speed comparisons on ImageNet-R.