IT3: Idempotent Test-Time Training

2024-10-05

arXiv.org

AI Summary

In real-world scenarios, deep learning models often suffer performance degradation due to distribution shifts between training and test data. This paper proposes a general test-time adaptation (TTA) method grounded in the principle of idempotence: leveraging output stability—i.e., $f(x, f(x, 0)) approx f(x, 0)$—to implicitly detect distribution shifts and refine representations, without auxiliary tasks, domain priors, or label supervision. Our key contribution is the first formulation of idempotence as a unified TTA objective, optimized via iterative forward passes and implicit projection to minimize the idempotent loss $|f(x, f(x, 0)) - f(x, 0)|$. The approach is architecture-agnostic and fully unsupervised, enabling online adaptation. Extensive experiments across diverse tasks—including image corruption classification, aerodynamic prediction, tabular imputation, facial age estimation, and aerial image segmentation—demonstrate substantial robustness gains. The method seamlessly supports MLPs, CNNs, and GNNs.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Adversarial Attacks & Robustness
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Abstract

This paper introduces Idempotent Test-Time Training (IT$^3$), a novel approach to addressing the challenge of distribution shift. While supervised-learning methods assume matching train and test distributions, this is rarely the case for machine learning systems deployed in the real world. Test-Time Training (TTT) approaches address this by adapting models during inference, but they are limited by a domain specific auxiliary task. IT$^3$ is based on the universal property of idempotence. An idempotent operator is one that can be applied sequentially without changing the result beyond the initial application, that is $f(f(x))=f(x)$. At training, the model receives an input $x$ along with another signal that can either be the ground truth label $y$ or a neutral"don't know"signal $0$. At test time, the additional signal can only be $0$. When sequentially applying the model, first predicting $y_0 = f(x, 0)$ and then $y_1 = f(x, y_0)$, the distance between $y_0$ and $y_1$ measures certainty and indicates out-of-distribution input $x$ if high. We use this distance, that can be expressed as $||f(x, f(x, 0)) - f(x, 0)||$ as our TTT loss during inference. By carefully optimizing this objective, we effectively train $f(x,cdot)$ to be idempotent, projecting the internal representation of the input onto the training distribution. We demonstrate the versatility of our approach across various tasks, including corrupted image classification, aerodynamic predictions, tabular data with missing information, age prediction from face, and large-scale aerial photo segmentation. Moreover, these tasks span different architectures such as MLPs, CNNs, and GNNs.

Problem

• Adapting deep learning models to test-time distribution shifts without auxiliary tasks
• Enforcing idempotence to replace domain-specific auxiliary tasks in adaptation
• Improving out-of-distribution performance across diverse domains and architectures

Innovation

• Enforces idempotence for test-time adaptation
• Uses only current test instance
• Eliminates need for auxiliary tasks