Matching High-Dimensional Geometric Quantiles for Test-Time Adaptation of Transformers and Convolutional Networks Alike

TL;DR

This work tackles test-time adaptation (TTA) under covariate shift by introducing an architecture-agnostic decorruptor that preprocesses test inputs without altering the pre-trained classifier. Training uses a novel quantile-loss objective based on high-dimensional geometric quantiles, aligning the feature marginals of corrupted inputs with those of the source data via a frozen classifier. The authors prove that, under a set of reasonable conditions and a good initialization, minimizing the quantile loss is equivalent to learning an optimal adapter that would be obtained from paired clean-corrupted data, effectively aligning class-conditionals up to transport-based notions. Empirically, the method yields substantial improvements across CIFAR10C, CIFAR100C, and TinyImageNet-C for both CNNs and transformer architectures, demonstrating strong architecture-independence and practical scalability with a memory-bank variance-reduction strategy.

Abstract

Test-time adaptation (TTA) refers to adapting a classifier for the test data when the probability distribution of the test data slightly differs from that of the training data of the model. To the best of our knowledge, most of the existing TTA approaches modify the weights of the classifier relying heavily on the architecture. It is unclear as to how these approaches are extendable to generic architectures. In this article, we propose an architecture-agnostic approach to TTA by adding an adapter network pre-processing the input images suitable to the classifier. This adapter is trained using the proposed quantile loss. Unlike existing approaches, we correct for the distribution shift by matching high-dimensional geometric quantiles. We prove theoretically that under suitable conditions minimizing quantile loss can learn the optimal adapter. We validate our approach on CIFAR10-C, CIFAR100-C and TinyImageNet-C by training both classic convolutional and transformer networks on CIFAR10, CIFAR100 and TinyImageNet datasets.

Figures (11)

• Figure 1: Comparing t-SNE plots of raw ResNet18 features versus Quantile-Loss corrected ResNet18 features on CIFAR10C with Gaussian distortions. ResNet18 is trained on CIFAR10. Observe that the corrected features maintain the clusters well at higher levels of distortions while the raw features lose the class structure as severity increases.
• Figure 2: Quantile-Matching pipeline - The black arrows denote the forward propagation. The dotted red arrows denote the backpropagation.
• Figure 3: An illustrative example where marginal distributions align perfectly but the class-conditionals are flipped. If the de-corruption operator class were chosen from rotations about the origin, it is possible to obtain both the correct operator as well as the one that flips class-conditionals perfectly.
• Figure 4: The errors are plotted on the de-corrupted features of Gaussian distorted CIFAR10C images at severity level $5$ (using ResNet18) as a function of epochs. The mean-squared error is computed using the test images of CIFAR10. Fig \ref{['fig:PairwiseQuantileScatter']} is a scatter plot between raw quantile loss and means-squared error losses. In Fig \ref{['fig:PairwiseQuantileEpoch']} Quantile loss versus means-squared error loss is plotted as a function of epochs after normalizing them to accommodate on the same scale. Observe that the two losses are positively correlated and the losses reduce with the number of training steps.
• Figure 5: Geometric Visualization of a Quantile Index of a Quantile.

Theorems & Definitions (14)

• Definition 1
• Theorem 1
• Theorem 2
• Theorem 3
• Proposition 1
• Theorem 4
• Definition 2
• Lemma 5
• proof
• proof