CONGA:Confidence-and-Gradient-Aware Learning Rate Schedule for Test Time Adaptation
Abstract
Test-time adaptation (TTA) adapts pretrained models to test data on-the-fly. Current TTA methods have focused on what to adapt: lightweight domain-aware components (prompts, normalization statistics) updated with consistency-aware self-supervised losses. This work investigates the more fundamental yet underexplored optimization process, providing insights and guidelines on how to appropriately update models for TTA. By analyzing the optimization error during TTA, we identify a pivotal stability-plasticity trade-off: the model should adapt to novel distributions while retaining learned knowledge, which motivates our design of a CONfidence-and-Gradient-Aware scheduler (CONGA) to constrain model learning rate (LR) within an adaptive exploration interval. For each iteration, the lower bound encourages model exploration on informative confident samples, while the upper bound prevents aggressive overfitting to noisy optimization gradients. Based on our theoretical findings, an adaptation-progress-conditioned cosine decay function decides the specific LR within the interval. As an LR scheduler, CONGA is naturally applicable on existing TTA methods as a plug-in module, introducing little computation overheads. Extensive experiments and analysis demonstrate the superiority and validness of CONGA.
Lay Summary
Machine learning systems often perform worse when the data they see in the real world differs from the data they were originally trained on, such as changes in weather, lighting, or image quality. A common solution is to let models adapt themselves while they are being used, but adapting too aggressively can cause them to forget what they already learned or become confused by noisy data. In this work, we study how to make this adaptation process more reliable and stable. We design a simple method called CONGA that automatically adjusts how fast a model learns during adaptation: it encourages the model to learn from trustworthy new information while preventing harmful updates from unreliable signals. Because CONGA can be added to many existing adaptation methods with little extra cost, it provides a practical way to make AI systems more robust when deployed in changing real-world environments.