A Layer Selection Approach to Test Time Adaptation
This work addresses the problem of model collapse in TTA for machine learning practitioners, offering an incremental improvement to enhance robustness in real-world deployment scenarios.
The paper tackles performance degradation in Test Time Adaptation (TTA) under challenging distribution shifts by proposing GALA, a layer selection criterion that identifies beneficial updates and filters noisy gradients, which improves existing TTA methods across multiple datasets, domain shifts, and architectures.
Test Time Adaptation (TTA) addresses the problem of distribution shift by adapting a pretrained model to a new domain during inference. When faced with challenging shifts, most methods collapse and perform worse than the original pretrained model. In this paper, we find that not all layers are equally receptive to the adaptation, and the layers with the most misaligned gradients often cause performance degradation. To address this, we propose GALA, a novel layer selection criterion to identify the most beneficial updates to perform during test time adaptation. This criterion can also filter out unreliable samples with noisy gradients. Its simplicity allows seamless integration with existing TTA loss functions, thereby preventing degradation and focusing adaptation on the most trainable layers. This approach also helps to regularize adaptation to preserve the pretrained features, which are crucial for handling unseen domains. Through extensive experiments, we demonstrate that the proposed layer selection framework improves the performance of existing TTA approaches across multiple datasets, domain shifts, model architectures, and TTA losses.