Test-time adaptation using Transporting Flow

Test-time adaptation (TTA) considers a pretrained model on source domain and aims to adapt it to target domain without labels, where the target is subject to distribution-shift. Most existing TTA methods are based on entropy-minimization and feature alignment, where the former are classification-specific. While the latter can be extended to regression, they either naively align the entire feature space or use principal component analysis (PCA) to find dimensions with maximum variance and perform alignment along those dimensions. These maximum variance dimensions however may not be those that are important for domain adaptation, as they represent spurious correlations between domains. A natural question then arises as to how we can identify features with maximal geometric displacements if they exhibit low variance. We use Optimal Transport and propose Transporting Flow (TOW), a framework that actively optimizes for a k-dimensional subspace that incurs maximal displacement cost for feature alignment. TOW computes 2-Wasserstein distance on the pre-learned density of source features and the target activations, and leverages Normalizing Flows to implicitly find an optimal transport map that minimizes the transport cost during adaptation. Our proposed method can be extended to both continuous regression and classification tasks in standard settings with any loss function. We show that TOW outperforms baselines on four real-world datasets.