CTA: Cross-Task Alignment for Better Test Time Training

TL;DR

This paper introduces CTA, a novel cross-task alignment method that enhances test-time training robustness by aligning supervised and self-supervised models without requiring specialized architectures, leading to improved generalization.

Contribution

CTA is a new approach that aligns representations of supervised and self-supervised models during test-time training, improving robustness without changing model architecture.

Findings

Significant robustness improvements on benchmark datasets.

Effective alignment reduces gradient interference.

Enhances generalization under distribution shifts.

Abstract

Deep learning models have demonstrated exceptional performance across a wide range of computer vision tasks. However, their performance often degrades significantly when faced with distribution shifts, such as domain or dataset changes. Test-Time Training (TTT) has emerged as an effective method to enhance model robustness by incorporating an auxiliary unsupervised task during training and leveraging it for model updates at test time. In this work, we introduce CTA (Cross-Task Alignment), a novel approach for improving TTT. Unlike existing TTT methods, CTA does not require a specialized model architecture and instead takes inspiration from the success of multi-modal contrastive learning to align a supervised encoder with a self-supervised one. This process enforces alignment between the learned representations of both models, thereby mitigating the risk of gradient interference, preserving the intrinsic robustness of self-supervised learning and enabling more semantically meaningful updates at test-time. Experimental results demonstrate substantial improvements in robustness and generalization over the state-of-the-art on several benchmark datasets.