Test-time adaptation using Transporting Flow

Test-time adaptation (TTA) considers a pretrained model on a source domain and aims to adapt it to a target domain without labels, where the target is subject to distribution shift. It has been found that models rely on both causal and spurious correlations (non-causal features) to predict output. While most existing TTA methods based on feature alignment perform alignment along directions with maximum variance, they do not explicitly look for discrepancy in non-causal dimensions. We believe non-causal dimensions can be relevant for prediction task and aligning those dimensions may result in improved performance. We introduce Transporting Flow (TOW), a regression-TTA method that finds non-causal dimensions in target feature space and adapts with the source ones. TOW leverages Optimal Transport on k-dimensional subspaces to compute the 2-Wasserstein distance between target and source feature representations. For adaptation, TOW uses Normalizing Flow to learn a bijective transporting function that maps target features on source. TOW consistently performs better than baselines on LIBS data.