IT$^3$: Idempotent Test-Time Training

## Abstract

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 ⁢ ( 𝐱 ) ) = f ⁢ ( 𝐱 ) 𝑓 𝑓 𝐱 𝑓 𝐱 f(f(\mathbf{x}))=f(\mathbf{x}) italic_f ( italic_f ( bold_x ) ) = italic_f ( bold_x ) .
At training, the model receives an input 𝐱 𝐱 \mathbf{x} bold_x along with another signal that can either be the ground truth label 𝐲 𝐲 \mathbf{y} bold_y or a neutral ”don’t know” signal 𝟎 0 \mathbf{0} bold_0 . At test time, the additional signal can only be 𝟎 0 \mathbf{0} bold_0 . When sequentially applying the model, first predicting 𝐲 0 = f ⁢ ( 𝐱 , 𝟎 ) subscript 𝐲 0 𝑓 𝐱 0 \mathbf{y}_{0}=f(\mathbf{x},\mathbf{0}) bold_y start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = italic_f ( bold_x , bold_0 ) and then 𝐲 1 = f ⁢ ( 𝐱 , 𝐲 0 ) subscript 𝐲 1 𝑓 𝐱 subscript 𝐲 0 \mathbf{y}_{1}=f(\mathbf{x},\mathbf{y}_{0}) bold_y start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = italic_f ( bold_x , bold_y start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT ) , the distance between 𝐲 0 subscript 𝐲 0 \mathbf{y}_{0} bold_y start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT and 𝐲 1 subscript 𝐲 1 \mathbf{y}_{1} bold_y start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT measures certainty and indicates out-of-distribution input 𝐱 𝐱 \mathbf{x} bold_x if high.
We use this distance, that can be expressed as ‖ f ⁢ ( 𝐱 , f ⁢ ( 𝐱 , 𝟎 ) ) − f ⁢ ( x , 𝟎 ) ‖ norm 𝑓 𝐱 𝑓 𝐱 0 𝑓 𝑥 0 ||f(\mathbf{x},f(\mathbf{x},\mathbf{0}))-f(x,\mathbf{0})|| | | italic_f ( bold_x , italic_f ( bold_x , bold_0 ) ) - italic_f ( italic_x , bold_0 ) | | as our TTT loss during inference. By carefully optimizing this objective, we effectively train f ⁢ ( 𝐱 , ⋅ ) 𝑓 𝐱 ⋅ f(\mathbf{x},\cdot) italic_f ( bold_x , ⋅ ) 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.

## 1 Introduction

Supervised learning methods, while powerful, typically assume that training and test data come from the same distribution. Unfortunately, this is rarely true in practice. Data encountered by systems operating in the real world often differs substantially from what they were trained on due to data distribution shifts over time or other changes in the environment. This inevitably degrades performance, even in state-of-the-art models (Recht et al., ; Hendrycks et al., ; Yao et al., ). Modern deployed machine learning systems not only need to adapt to distribution shifts but also must do so on-the-fly using very limited data.

The problem setup this work addresses requires adapting to distribution shifts on-the-fly using only the current test instance or batch, without access to any additional labeled or unlabeled data during inference. During training, the model has access solely to the base distribution training data, with no prior knowledge of the test distribution. Adversarial robustness and domain adaptation address related challenges, they typically require additional data either during training or inference, and sometimes rely on specific assumptions about the nature of the shift. While effective in their contexts, they are not designed for immediate, instance-level adaptation and do not solve our problem setup. Test-Time Training (TTT) (Sun et al., ) offers an alternative by adapting the model during inference using an auxiliary self-supervised task on each test sample. This dynamic, on-the-fly adaptation allows the model to handle corrupted and Out-of-Distribution (OOD) data using only the current test instance or batch, without access to any other data. However, TTT employs an auxiliary task specific to the data modality (e.g., orientation prediction or inpainting for imagery data) (Gandelsman et al., ).

In this paper, we argue that enforcing idempotence can profitably replace the auxiliary tasks in TTT and results in an approach we dub IT^3 that is a versatile and powerful while generalizing well across domains and architectures.

More specifically, an operator $f$ is said to be idempotent if it can be applied sequentially without changing the result beyond the initial application, namely: $f(f(\mathbf{x}))=f(\mathbf{x})$. This can be understood as a generalization of orthogonal projection in linear spaces to non-linear settings. At training, IT^3 receives an input $\mathbf{x}$ along with another signal that can either be the ground truth label $y$ or a neutral ”don’t know” signal $\mathbf{0}$. Durasov et al. () sequentially applied a model trained with this policy such that $y_{0}=f(\mathbf{x},\mathbf{0})$ and $y_{1}=f(\mathbf{x},y_{0})$. The distance $||y_{0}-y_{1}||$ was shown to indicate both the prediction uncertainty and whether $\mathbf{x}$ is out-of-distribution (OOD). What if, at test time, we could actively minimize this distance whenever we encounter an instance? Could we “pull it” into the distribution? IT^3 uses this this distance as a loss for TTT sessions. When we unfold $y_{0}$ and $y_{1}$ in such a loss term we obtain: $||f(\mathbf{x},f(\mathbf{x},\mathbf{0}))-f(\mathbf{x},\mathbf{0})||$. A more detailed analysis of this term reveals an important insight: the optimization objective effectively guides the model towards making $f(\mathbf{x},\cdot)$ idempotent. While not trivial, it has been demonstrated that, with appropriate adjustments, training a model to achieve idempotence is feasible (Shocher et al., ). This ties everything together: idempotence, seen as a generalization of projection, suggests the existence of a subset onto which the model maps the internal representation of the input. In our case, this subset exists in the joint $X\times Y$ space, and corresponds roughly to the distribution of correctly paired $\mathbf{x},\mathbf{y}$ examples.

The result is a global method that does not rely on any specific domain properties. This is in contrast to prior TTT methods that rely on a domain specific auxiliary task. By leveraging the universal property of idempotence, IT^3 can adapt OOD test inputs on-the-fly across various domains, tasks and architectures.
This includes image classification with corruptions, aerodynamic predictions for airfoils and cars, tabular data with missing information, age prediction from faces, and large-scale aerial photo segmentation, using MLPs, CNNs, or GNNs.

## 2 Related Work

IT^3 relies on the notion of idempotence to globalize TTT. We briefly review these two fields.

### 2.1 Test-Time Training

The idea of leveraging test data for model adaptation dates back to methods like transductive learning  (Gammerman et al., ). Early approaches, such as transductive SVMs (Collobert et al., ) and local learning (Bottou & Vapnik, ), aimed to adapt predictions for specific test samples rather than generalizing across unseen data.

Training neural networks solely on single test instances, without pre-training, has been demonstrated in the ”deep internal learning” line of works, for many image enhancement tasks (Shocher et al., ; Gandelsman et al., ) and single image generative models (Shocher et al., ; Shaham et al., ).

Test-Time Training (TTT) has emerged as a solution to the problem of generalization under distribution shifts. Using a pre-trained network and at test-time refining on a single instance each time. In the foundational work of Sun et al. (), the model is adjusted in real-time by solving an auxiliary self-supervised task, such as predicting image rotations, on each test sample. This on-the-fly adaptation has proven effective in improving robustness on corrupted and Out-Of-Distribution (OOD) data. As the self-supervised learning methods became more efficient (He et al., ), they could be exploited for TTT (Gandelsman et al., ).
Extensions such as TTT++ (Liu et al., ) assume access to the entire test set. TENT (Wang et al., ) adapts during inference in the batch level, based on the batch entropy, but cannot be applied to single instances or very small batches. Moreover, it relies on updating the model’s normalization layers, making it architecture dependent.

### 2.2 Idempotence in Deep Learning

Idempotence, a concept rooted in mathematics and functional programming, refers to an operation where repeated application yields the same result as a single application. Mathematically, for a function $f$, being idempotent means

$$ $f(f(x))=f(x),\quad\forall x\;.$ (1) $$

In other words, applying the function multiple times has no effect beyond the first application.
In the context of linear operators, idempotence corresponds to orthogonal projection. Over $\mathbb{R}^{n}$, these are matrices $A$ that satisfy $A^{2}=A$, with eigenvalues of either 0 or 1; they can be interpreted as geometrically preserving certain components while nullifying others.
This concept was recently applied in generative modeling. Idempotent Generative Network (IGN) (Shocher et al., ) is a generative model that maps data instances to themselves, $f(x)=x$, and maps latents to targets that also map to themselves, $f(f(z))=f(z)$. It was further shown to ’project’ corrupted images onto the data manifold, effectively removing the corruptions without prior knowledge of the degradation.

Energy-Based Models (EBMs; Ackley et al. ()) offer a related perspective by defining a function $f$ that assigns energy scores to inputs, with higher energy indicating less desirable or likely examples, and lower energy indicating those that fit the model well. IGN introduces a similar concept but frames it differently: instead of $f$ directly serving as the energy function, the energy is implicitly defined via the difference $\delta(y)=D(f(y),y)$, where $D$ measures the distance between the model’s prediction and its input. In this framework, training $f$ to be idempotent minimizes $\delta(f(z))$, pushing the model toward a low-energy configuration where its outputs remain stable under repeated applications. Thus, $f$ can be interpreted as a transition operator that drives high-energy inputs toward a low-energy, stable domain, reducing the need for separate optimization procedures to find the energy minimum.

In concurrent work, the ZigZag method has first been proposed and then extended to recursive networks (Durasov et al., ; ). It introduces idempotence as a means to assess uncertainty in neural network predictions, ZigZag operates by recursively feeding the model’s predictions back as inputs, allowing the model to refine its outputs. The consistency between successive predictions acts as an uncertainty metric, where stable, unchanged outputs indicate higher confidence, while divergent predictions signal uncertainty or out-of-distribution (OOD) data. Unlike popular sampling-based uncertainty estimation methods (Gal & Ghahramani, ; Lakshminarayanan et al., ; Wen et al., ; Durasov et al., ), ZigZag does not require many forward passes or complex sampling, making it more computationally efficient for real-time applications.

## 3 Method

### 3.1 Initial Training

Let $f_{\theta}$ be a generic model that we wish to deploy in an environment where the statistical distribution of the samples it receives may change over time. Fig 1 depicts the initial training phase, we perform a standard supervised training with a slight modification inspired by the ZigZag method Durasov et al. (): We modify the first layer of the network implementing $f_{\theta}$ so that it can accept a second argument in addition to the data sample $\mathbf{x}$ that it is normally takes as input. This additional argument can be either $\mathbf{y}$, the desired output of the network given input $\mathbf{x}$, or a neutral ”don’t know” signal $\mathbf{0}$. During training, we minimize the supervised loss

$$ $\displaystyle\mathcal{L}_{\text{train}}=\|f_{\theta}(\mathbf{x},\mathbf{y})- \mathbf{y}\|+\|f_{\theta}(\mathbf{x},\mathbf{0})-\mathbf{y}\|\;,$ (2) $$

where$f(\mathbf{x},\cdot)$ is the model’s prediction given input $\mathbf{x}$ and the additional input.
when $\mathcal{L}_{\text{train}}$ is minimized, we can write

$$ $\displaystyle\mathbf{y}_{0}=f_{\theta}(\mathbf{x},\mathbf{0})\approx\mathbf{y} \;,\mathbf{y}_{1}=f_{\theta}(\mathbf{x},\mathbf{y}_{0})\approx f_{\theta}( \mathbf{x},\mathbf{y})\approx\mathbf{y}\;\Rightarrow\;f_{\theta}(\mathbf{x},f_ {\theta}(\mathbf{x},\mathbf{0}))\approx f_{\theta}(\mathbf{x},\mathbf{0})\;.$ (3) $$

Of course, this can only be expected to hold when $\mathbf{x}$ is within the training distribution. When $\mathbf{x}$ is out-of-distribution, $\mathbf{y}_{0}$ and $\mathbf{y}_{1}$ can be very different. In Durasov et al. (), this is exploited to estimate uncertainty: the greater the deviation from Eq. 3, the more uncertain the prediction is taken to be. In contrast, in this paper, we use Test-Time Training to enforce the constraint of Eq. 3 during inference.

### 3.2 Test-Time Training

At test time, IT^3 introduces a dynamic adaptation process
to refine the model’s predictions for new inputs,
particularly those that may be OOD or corrupted.
The goal is to make the model idempotent with respect to its second input, that is, ensuring that $f_{\theta}(\mathbf{x},\cdot)$
is an idempotent function for a given $\mathbf{x}$. A naive way to enforce this would be to minimize the loss function

$$ $\displaystyle\mathcal{L}_{\text{TTT}}=\|f_{\theta}(\mathbf{x},f_{\theta}( \mathbf{x},\mathbf{0}))-f_{\theta}(\mathbf{x},\mathbf{0})\|\;.$ (4) $$

However, directly minimizing this loss can produce undesirable side effects.
For instance, if $\mathbf{y}_{0}$ is an incorrect prediction,
minimizing the distance $\|\mathbf{y}_{0}-\mathbf{y}_{1}\|$ may cause $\mathbf{y}_{1}$ to be pulled toward the incorrect $\mathbf{y}_{0}$, thereby magnifying the error.
Another potential failure mode is due to the fact that $identity$ is idempotent. For $f=identity$, we get $\mathcal{L}_{\text{TTT}}=0$.

To prevent such a collapse, we modify the test-time training procedure as depicted in Fig, 1:
We keep a copy of the model as it was at the end of the training phase, denoted as $F=f_{\Theta}$, where $\Theta$ are the weights obtained after the initial training of Sec. 3.1, which will not be updated further.
We then take the test-time loss to be

$$ $\displaystyle\mathcal{L}_{\text{IT}^{3}}=\|F(\mathbf{x},f_{\theta}(\mathbf{x}, \mathbf{0}))-f_{\theta}(\mathbf{x},\mathbf{0})\|\;,$ (5) $$

where $f_{\theta}$ is the model being updated at test-time. Here, the first prediction $\mathbf{y}_{0}=f_{\theta}(\mathbf{x},\mathbf{0})$ is computed as before, but the second one, $\mathbf{y}_{1}=F(\mathbf{x},f_{\theta}(\mathbf{x},\mathbf{0}))$, is made using the frozen model $F$.
By updating only $f_{\theta}$ and keeping $F$ fixed,
we ensure that $\mathbf{y}_{0}$ is adjusted to minimize the discrepancy with $\mathbf{y}_{1}$,
without pulling $\mathbf{y}_{1}$ toward an incorrect $\mathbf{y}_{0}$. A similar idea was employed in the IGN approach (Shocher et al., )
meaningful predictions are required. After each TTT optimization iteration, the dynamic model $f_{\theta}$ is initialized with $\Theta$, ready for the next input.

### 3.3 Online IT 3

For streaming data scenarios, where the distribution shifts continuously over time, we modify IT^3 to operate in an online mode by not resetting $f_{\theta}$ back to $F$ after each TTT episode, as we did in Section 3.2.
We essentially assume that the distribution mostly shifts smoothly and, thus, there is a good reason to believe that the current state of $f_{\theta}$ is a better
initialization for the next TTT episode than the original $F$. This makes the model evolve over time.
In this scheme, it can happen that the performance of the model on data
from its original training decreases significantly, a phenomenon known as catastrophic
forgetting (Kirkpatrick et al., ). This is acceptable as the goal is to perform well on data at the present moment, rather than on past examples.

Another modification in the online setup is for the second sequential application of the model,
i.e., the $F$ in $F(\mathbf{x},f_{\theta}(\mathbf{x},\mathbf{0}))$.
Since the data keeps shifting, there is no reason to retain the frozen $F$ as an anchor indefinitely.
Over time, $f_{\theta}$ may diverge far from $F$, making it irrelevant. Relying on the old state of the model
would prevent the model from evolving efficiently. Replacing it with the current state of $f_{\theta}$ is out of
the question, as it causes collapse. We need an anchor that is influenced by a reasonable amount
of data, yet evolves over time. Our solution is to replace $F$ with an Exponential Moving Average (EMA) of the model $f_{\theta}$, denoted as $f_{\text{EMA}}$. This means $f_{\text{EMA}}$ is a smoothed version of $f_{\theta}$ over time.
The test-time loss in the online setting then becomes

$$ $\displaystyle\mathcal{L}_{\text{online}}=\|f_{\text{EMA}}(\mathbf{x},f_{\theta }(\mathbf{x},\mathbf{0}))-f_{\theta}(\mathbf{x},\mathbf{0})\|\;.$ (6) $$

By updating both $f_{\theta}$ and $f_{\text{EMA}}$ incrementally,
with $f_{\text{EMA}}$ serving as a stable reference that changes more slowly,
the model adapts to gradual shifts without overfitting to noise or temporary anomalies.

## 4 Experiments

We evaluate our approach across a diverse set of data types and tasks, including age prediction, image classification, and road segmentation in the visual domain, as well as aerodynamics prediction using 3D data and tabular data experiments. In all these scenarios, we first train the model using the supervised approach of Section 3.1 and then perform the test-time training of Section 3.2. For each task, we design an OOD data subset for evaluation. The OOD data is divided into several levels, with higher levels representing data that is progressively further from the training distribution. In each experiment, we observe how quickly the model’s performance degrades as the level of OOD-ness increases. We have not found test-time adaptation baselines matching the TTT problem setup for any of the tasks except image classification. So we provide the comparative numbers in this case and, for all cases, we compare against the performance of the vanilla non adaptive model. In all cases we used published, common, strong models that are state-of-the-art or close to it. Across all scenarios, our method degrades slower than the vanilla network baseline.

### 4.1 Tabular Data

Tabular data consists of numerical features and corresponding continuous target values for regression tasks from the UCI tabular datasets (Bay et al., ). They are widely used in machine learning research for benchmarking regression models. In our case, we use The Boston Housing dataset describes housing prices in the suburbs of Boston, Massachusetts. It includes various features related to socioeconomic and geographical factors that influence housing prices. We take a test set and gradually apply random feature zeroing with increasing probabilities of 5%, 10%, 15%, and 20% (4 mentioned levels of OOD). This random feature dropping simulates out-of-distribution (OOD) data by progressively altering the input features, making the data less similar to the original training distribution. As the probability of feature dropping increases, the data becomes more OOD, which lowers the model’s accuracy. The trained model is a simple Multi-Layer Perceptron (MLP) optimized using the Adam optimizer, and we observe that IT^3 consistently degrades less compared to the vanilla baseline across all OOD levels as depicted in Fig. 2.

Figure: Figure 2: UCI Results on OOD inputs: The plots illustrate the performance of IT^3 compared to a vanilla model across different OOD levels. Left: The mean absolute error (MAE) shows that ITTT outperforms the vanilla model, retaining performance better as the data shifts further from the training distribution. Right: The box plot for car data shows the distribution of MAE at various OOD levels, where ITTT with different batch sizes ([batch=1, batch=4, batch=8]) degrades less compared to the Not optimized baseline. Larger batch sizes preserve performance more effectively.
Refer to caption: x1.png

### 4.2 CIFAR

Figure: Figure 3: Test error (%) on CIFAR-10-C with level 5 corruptions. We compare our approaches, $IT^{3}$, with object recognition without self-supervision. $IT^{3}$ improves over the baseline and higher batch size improves even further. The comparison with TTT (Sun et al., ) is provided for context, but it is not a direct comparison, as TTT uses a batch size of 1. Augmentations over this single instance create a batch of 32, yet only one instance is accessed at a time.
Refer to caption: x3.png

We conducted similar experiments using the CIFAR-10 (Krizhevsky et al., ) dataset, selecting CIFAR-C (Hendrycks & Dietterich, ) as the out-of-distribution (OOD) data. CIFAR-C contains the same images as CIFAR-10 but with various common corruptions, such as Gaussian noise, blur, and contrast variations, simulating real-world conditions. These corruptions are applied at different severity levels, allowing us to evaluate how the model’s performance degrades as the data shifts further from the original CIFAR-10 distribution. For this experiment, we used the Deep Layer Aggregation (DLA) (Yu et al., ) network, known for its strong performance in image classification and robustness to overfitting. We trained the model according to the guidelines from the original DLA paper to ensure optimal results. Fig.2 shows the evaluation error on CIFAR-C at severity level 5 for different types of corruptions, following (Sun et al., ). As shown, IT^3 outperforms the vanilla model, with higher batch sizes yielding the best results.
In our basic setup, batch size of 1 does not work well. We did not explore the possibility of following Sun et al. () creating a batch of augmented copies, as this would be a domain specific element, hurting the purity of our general method.
For context, we add TTT results, while fully acknowledging that this is not a fair comparison as they have access to a single instance, making the batch size effectively 1, although augmented to a batch of size 32.

### 4.3 Age Prediction

Figure: Figure 4: Face Samples. The (top) row shows training images of middle-aged individuals, while (middle) and (bottom) display images of older and younger individuals (OOD).
Refer to caption: extracted/5896912/images/age_viz/age_img_test_1.png

To experiment with image-based age prediction from face images, we use the UTKFace dataset (Zhang et al., ), a large-scale collection containing tens of thousands of face images annotated with age information. The model is trained on face images of individuals aged between 20 and 60, while individuals younger or older than this range are considered out-of-distribution (OOD) (Fig.6). The further the age is from the 20-60 interval, the higher the OOD level we assign to it. We use a ResNet-152 backbone with five additional linear layers and ReLU activations. This architecture delivers strong accuracy, outperforming the popular ordinal regression model CORAL (Cao et al., ) and matching other state-of-the-art methods (Berg et al., ). We train our model on the UTKFace training set (limited to individuals aged 20-60) and then run inference on faces at different OOD levels. Once again, IT^3 significantly outperforms the non-optimized model, as shown in Figs.6 and6.

### 4.4 Road Segmentation

Figure: Figure 7: Road Samples. RoadTracer dataset (left) covers urban areas of six different countries while Massachusetts dataset (right) primarily features rural neighborhoods along with some urban areas.
Refer to caption: x6.png

Figure: Figure 10: Qualitative effect of IT^3 on Road Segmentation. From left to right: (1) Original aerial image, (2) Not optimized output, (3) IT^3 output at the 5th iteration, (4) IT^3 output at the 15th iteration, and (5) Ground truth label. The segmentation quality improves significantly with IT^3 iterations, as observed in the progressively refined outputs at the 5th and 15th iterations.
Refer to caption: x9.png

Our method can be generalized to segmentation tasks as well. To demonstrate this, we consider the problem of road segmentation in aerial imagery using the RoadTracer dataset (Bastani et al., ). We train a DRU-Net (Wang et al., ), on the RoadTracer dataset.

We perform OOD experiments using Massachusetts Road dataset (Mnih, ) that primarily comprises rural neighborhoods, as depicted in Fig. 9. We sample 450 images, each with dimensions of $1500\times 1500$ pixels and divide them into four groups based on the Mean Squared Error (MSE) of the segmentation outputs, effectively creating different levels of distributional shift within the sampled set. We then further train the network on these OOD subsets using the ZigZag method (Durasov et al., ).

We evaluate road segmentation performance by using Correctness, Completeness and Quality (CCQ) metric (Wiedemann et al., ) which is a popular metric to evaluate delineation performance. The Correctness, Completeness and Quality are equivalent to precision, recall and intersection-over-union, where the definition of a true positive has been relaxed from spatial coincidence of prediction and annotation to co-occurrence within a distance of 5 pixels. As shown in Fig. 9 and 9, IT^3 significantly improves the performance on OOD images.

### 4.5 Aerodynamics Prediction

Figure: Figure 11: Airfoil Samples. Training and testing profiles (left) show reasonable aerodynamics, while OOD samples (right) feature rare, high lift-to-drag shapes. Black arrows indicate pressure, and red lines show lift and drag.
Refer to caption: extracted/5896912/images/profiles.png

#### Wings.

Our method is versatile and can handle various types of data. To illustrate this, we generated a dataset of 2,000 wing profiles, as depicted in Fig.13, by sampling the widely used NACA parameters (Jacobs & Sherman, ). We used the XFoil simulator (Drela, ) to compute the pressure distribution along each profile and estimate its lift-to-drag coefficient, a crucial indicator of aerodynamic performance. The resulting dataset consists of wing profiles $\mathbf{x}_{i}$, represented by a set of 2D nodes, and the corresponding scalar lift-to-drag coefficient $\mathbf{y}_{i}$ for $1\leq i\leq 2000$.

We selected the top 5% of shapes, based on their lift-to-drag ratio, as out-of-distribution (OOD) samples. The OOD levels were determined using the ground truth lift-to-drag ratio, where higher OOD levels correspond to more aerodynamically streamlined shapes. The training set includes shapes with lift-to-drag values ranging from 0 to 60, with anything beyond this threshold considered OOD and excluded from training. We then trained a Graph Neural Network (GNN) composed of 25 GMM (Monti et al., ) layers, featuring ELU activations (Clevert et al., ) and skip connections (He et al., ), to predict the lift-to-drag coefficient $\mathbf{y}_{i}$ from the profile $\mathbf{x}_{i}$, following the approach of (Remelli et al., ; Durasov et al., ). As with previous experiments, IT^3 significantly improves performance on OOD shapes and provides more accurate predictions compared to the vanilla model, as shown in Figs. 13 and 13.

Figure: Figure 14: Car Samples. The car dataset comprises many regular vehicles (left) and a few streamlined ones (right), which we treat as being out-of-distribution. Red and blue denote high and low pressures respectively.
Refer to caption: extracted/5896912/images/cars.png

#### Cars.

As for wings, we experimented with 3D car models from a subset of the ShapeNet dataset (Chang et al., ), which contains car meshes suitable for CFD simulations, as depicted in Fig. 16. The experimental protocol was the same as for the wing profiles, except we used OpenFOAM (Jasak et al., ) to estimate drag coefficients and employed a more sophisticated network to predict them from the triangulated 3D car meshes.

To predict drag associated to a triangulated 3D car, we utilize similar model to airfoil experiments but with increased capacity. Instead of twenty five GMM layers, we use thirty five and also apply skip-connections with ELU activations.
Final model is being trained for 100 epochs with Adam optimizer and $10^{-3}$ learning rate. As with airfoils experiments, IT^3 significantly improves performance on OOD shapes and provides more accurate predictions compared to the vanilla model, as shown in Figs. 16 and 16.

### 4.6 Online IT 3

We test our proposed online variation on several tasks. Naturally, when the distribution remains constant (although shifted from the training distribution) we expect superior results w.r.t. the offline setup, as our model keeps being trained on the new distribution.
A more effective way to test constant adaptation over time is to use a continuously changing distribution. We test IT^3 on an increasing corruption / OOD level. We see in all cases that the online variation of IT^3 performs significantly better than the basic anchored variation.

**Table 1: Qualitative result for Online IT^3. We report evaluation metrics for the road segmentation task (left), airfoils lift-to-drag prediction (middle), and car drag prediction (right).The results suggest Online IT^3 enhances the performance compared to the original model. Additionally, online IT^3 significantly outperforms offline IT^3.**
| Method | Corr | Comp | Quality |
| --- | --- | --- | --- |
| Not Optimized | $55.7$ | $44.3$ | $39.5$ |
| IT^3 (batch=4) | $55.7$ | $49.1$ | $46.4$ |
| IT^3 (batch=8) | $58.1$ | $52.0$ | $48.5$ |
| IT^3 (batch=16) | $57.3$ | $52.7$ | $48.7$ |
| IT^3 (online) | 77.5 | 79.8 | 69.8 |

Road segmentation: Building upon our previous road segmentation experiments, we further evaluate the effectiveness of online IT^3. In the online IT^3 setup, OOD samples are ranked based on their mean squared error (MSE) loss when passed through the vanilla network. We begin by selecting the samples with low MSE loss, as these are closer to the training distribution given the network’s strong performance on them. Gradually, we introduce samples with progressively higher MSE loss, smoothly shifting between distributions and thereby allowing the model to adapt effectively to a range of OOD samples. As in previous experiment, we use DRU-Net trained on the RoadTracer dataset as vanilla model and 890 images are sampled from Massachusetts dataset as OOD images.

Firstly, the vanilla network is tested on the Massachusetts dataset without any additional fine-tuning. We then apply online IT^3 during inference to adapt the model to the OOD distribution as new data is presented. We evaluate the segmentation performance using the Correctness, Completeness, and Quality metrics, as described previously. Table 1 (left) summarizes the results. The application of IT^3 improved the performance over the initial network and the online IT^3 method significantly outperforms the offline IT^3.

Aerodynamics: Similarly, we conducted online experiments for airfoils lift-to-drag prediction and for car drag prediction.
We set the data stream such that OOD shapes appear with increasing aerodynamic properties, modeling a continuous domain shift in the data. As with the segmentation results, the online version significantly outperforms both the offline version and the original network, as shown in Tabs.1 (middle and right).

## 5 Limitations:

While global, IT^3 lacks domain expertise. Within the domains we experimented with, we are aware only of computer-vision algorithms that adhere to the restrictive problem setup. However, it is likely possible to implement domain specific methods based on self-supervision that can outperform IT^3.

In addition, we found that for some domains it is hard to apply IT^3 on single instances without also using additions that require domain expertise or access to training data. This is most common for domains where the information within a single input is limited. The reason for that is that in contrast to self-supervised auxiliary tasks, our TTT objective is based on predictions without independent information on the input data.

## 6 Conclusion

We have proposed an approach to test-time-training that relies on enforcing idempotence as new samples are being considered to effectively handle domain shifts. The method is generic and we have demonstrated that it is effective in a wide range of domains without requiring domain-specific knowledge, which sets it apart from other state-of-the-art methods.

In future work we plan to pursue the challenge of realistic online continual learning, where there is no pre-training at all and the data arrives in streams, sometimes with labels and sometimes not. We believe IT^3 can be adapted to such a setup across many different streaming modalities, which would make it extremely useful in real-world scenarios.

## Acknowledgments

We would like to thank Yossi Gandelsman and Yoad Tewel for their valuable advice and discussions.

## Appendix A Appendix

### A.1 Additional Road Segmentation Experiments

In order to further evaluate our method, we perform an additional OOD experiments using the test set of the RoadTracer dataset itself. We select cities from the RoadTracer test dataset that are not part of the RoadTracer training set and treat them as OOD samples. We divide the selected set into four groups based on the Mean Squared Error (MSE) of the segmentation outputs, effectively creating different levels of distributional shift within the test set. We then further train the network on these OOD subsets. In line with our other experiments, we demonstrate that applying IT^3 significantly improves segmentation performance on these OOD samples. The quantitative results of this experiment can be seen in Fig. 19 and 19.

Figure: Figure 17: Road Samples. RoadTracer train dataset (left) includes urban areas of cities in US. From RoadTracer test dataset, we selected images of cities that are not included in train dataset as OOD samples (right).
Refer to caption: x14.png

Figure: Figure 20: Additional qualitative results of IT^3 on Road Segmentation. From left to right: (1) Original aerial image, (2) Not optimized output, (3) IT^3 output at the 5th iteration, (4) IT^3 output at the 15th iteration, and (5) Ground truth label. The segmentation quality improves significantly with IT^3 iterations, as observed in the progressively refined outputs at the 5th and 15th iterations.
Refer to caption: x17.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 4
questions: - Does the test time objective result in networks that are idempotent on OOD samples? Presumably once the function drifts from the fixed anchor function that test time loss no longer reflects idempotence? It would be good to see measures of idempotence on training and test samples.

- How does the intuition about idempotence being a generalization of orthogonal projection hold up? The proposed method considers idempotence only in the y-variable, not the entire function.
soundness: 2
strengths: - Idempotent training is an interesting and novel approach for addressing the out-of-distribution generalization problem.

- The paper shows successful application of the considered approach on a large number of experimental settings, including in online-learning settings.
summary: - This paper proposes a test-time-training based approach to address the distribution shift or OOD generalization problem by learning models that are idempotent.

- In particular, this paper proposes a method where models $f_{\theta}: \mathcal{X} \times \mathcal{Y} \to \mathcal{Y}$ are trained by minimizing both the difference between $f_{\theta}(x, y)$ and $y$ as well as the difference between $f_{\theta}(x, 0)$ and $y$, resulting in a model that is idempotent on the training set. Then at test time, models are optimized on the test data before running inference to make them idempotent on test inputs, by minimizing the difference between $f_{\theta}(x, 0)$ and $f_{\theta}(x, f_{\theta}(x, 0))$ for an OOD input $x$.

- The paper shows empirically for several different settings that idempotent test-time-training improves classification accuracy on out-of-distribution samples.
weaknesses: - Test-time training requires undesirable extra compute for test-time optimization of the whole model. How much additional compute is needed compared to running inference on the base model? How does this method scale with model size?

- The paper lacks analysis that explains how or why idempotent training is expected to improve out-of-distribution analysis. Further investigation and ablations that provide intuition for how this method works would be valuable.

- The proposed method has not been evaluated on standard real-world out-of-distribution generalization benchmarks such as DomainBed [Gulrajani and Lopez-Paz 2020] and WILDS [Koh et al 2020]. The presented experiments are on smaller models/datasets.

References:

Gulrajani, Ishaan, and David Lopez-Paz. "In search of lost domain generalization." arXiv preprint arXiv:2007.01434 (2020).

Koh, Pang Wei, et al. "Wilds: A benchmark of in-the-wild distribution shifts." International conference on machine learning. PMLR, 2021.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. Do you use a special encoding or placeholder value instead of 0 for regression tasks?
2. If the purpose of the neural "don't know" zero is purely for contrast and its specific value is not important, will it be more computationally efficient to use a representative value such as the median of $y_i$ where $i \in \text{training set}$?
3. In EMA (Morales-Brotons et al., 2024), when $\alpha = 1$, the online $IT^3$ aligns with the offline version, and when $\alpha = 0$, it encounters the collapse issue described in Section 3.2. Could you provide guidance on selecting the value of $\alpha$ or share experimental results demonstrating performance across different $\alpha$ values?
4. The right panel of Figure 2 and Figure 16 both appear to represent the car data. Could there be a potential mistake or duplication here?

Morales-Brotons, D., Vogels, T., & Hendrikx, H. (2024). Exponential moving average of weights in deep learning: Dynamics and benefits. Transactions on Machine Learning Research.
soundness: 2
strengths: The authors extend the concept of TTT (Sun et al., 2020) by incorporating idempotence, offering a simple yet elegant solution. The approach is intuitive and Figure 1 is particularly helpful for quickly grasping the core idea, even for those who are not experts in distribution shift.
summary: This paper introduces $IT^3$, a generic method for addressing distribution shift challenges. By enforcing idempotence, $IT^3$ sequentially adapts to potential distribution shifts during inference. The authors show the performance of $IT^3$ across diverse scenarios.
weaknesses: 1. (cf Figure 1) During training, in addition to feeding the standard $(x, y)$ pair to train the model, the authors also input the $(x, 0)$ pair to ensure the model satisfies the property of idempotence, referring to the zero as a neutral “don't know” signal. While this approach may work in classification tasks where zero could represent a new “don’t know” class, in regression tasks, it is unclear how the model differentiates this zero from an actual label of 0.
2. (Line 211) Online $IT^3$ appears to rely significantly on the use of Exponential Moving Average (EMA). However, the authors did not provide a citation for this technique. 
3. In the first experiment (Section 4.1 on tabular data), the method of randomly substituting features with zeros in the test set may resemble a missing data scenario rather than a true distribution shift. In other experiments, the authors simulate distribution shift by splitting the training and testing data based on a specific covariate to create distinct distributions. It is unclear why the same method was not applied for the first experiment. If the authors prefer using the zeroing method, they should include figures or statistical tests to substantiate that the training and testing data are distributionally different, rather than relying solely on intuition.
4. There are unnecessary abbreviations throughout the manuscript. For instance, “such that” is shortened to “s.t.” in line 490. The proposed method, $IT^3$, is inconsistently referred to as ITTT in parts of the manuscript, such as in Figure 13.
5. Figures 14 to 16 are not mentioned or referred to in the main text. This omission is unusual and may confuse readers as to the purpose or relevance of these figures.
6. Figures 12 and 15 have the exact same title.

Some minor improvements and spelling corrections for clarity:
1. (Line 22) $x$ (not $\mathbf{x}$) is not mentioned before.
2. (Line 77) $y_2$ is not mentioned before.
3. (Line 85) "th input" should be corrected to "the input".
4. (Line 131) Missing right parenthesis: $f(f(z))$.
5. (Line 490) “s.t. that” should be replaced with “such that”.
6. (Line 495) Remove the extra colon (":").
7. In Table 1, the title references “qualitative” results, but the data presented are numerical and should be described as “quantitative” results.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: As mentioned before, I find this work quite interesting but the lack of proper baselines push me towards a weak reject opinion. As for specific questions and ways that the authors can improve their work:
*  More baselines are needed on all tasks; even if the method does not translate exactly to the setup considered, the authors could perform minor adaptations so that it does. For example, why not consider additional self-supervised tasks as discussed at Sun et al. (2020)? In the case where simple rotation prediction might not apply, something simple like denoising an artificially corrupted image could still work as a self-supervised task. Another example would be on the online setup; there, methods from the TTA literature could be applied, such as the work of [1] which works even on the instance level and without batch normalisation.
* Apart from the CIFAR-C case, most other distribution shifts are generated by just partitioning the datasets according to some rule and then training on a subset while considering the other as OOD. This is a bit of a constrained setting and I would encourage the authors to consider more diverse shifts, as that would better highlight the usefulness of IT$^3$. For example, why not add noise to the road segmentation task, in the same way that it was done for CIFAR 10? This could be plausible real world setting where there is a fault in the sensor, thus the images come out corrupted. 
* How is the label encoded in the input in the various settings considered? This is an important information for reproducibility. Furthermore, is the loss at Eq. 2 used for all settings (even when it might not make much sense, such as classification)? 
* In the online setting, the authors consider a smooth transition between the distribution shifts, which might not be practically realistic. How does the method behave when the transition between distribution shifts is not smooth? 
* How many update steps on each datapoint do the authors do? Does the test time optimization reach a fully idempotent model and does “more idempotence” provide better performance? 


[1] Towards Stable Test-Time Adaptation in Dynamic Wild World, Niu et al., ICLR 2023
soundness: 2
strengths: * The paper is mostly well written; the ideas are explained clearly and reasonable intuitions are being given. Same goes for the experimental evaluation and discussion of the settings. 
* The idea is simple and I find it quite interesting. While the main architecture of the method relies heavily on the prior work Durasov et al., the application on the TTT setting is novel, since, as far as I know, idempotence hasn’t been explored in such a scenario. 
* The tasks considered in the experiments are quite diverse, which is nice to see. They range from simple tabular data, to standard image classification, to regression and dense prediction with various architectures. 
* The authors demonstrate improvements with IT$^3$ on all settings considered upon (admittedly very simple) baselines.
summary: This work presents IT$^3$ a method for test time training that relies on the concept of idempotence. During training, the authors train the neural network to be able to predict the ground truth label by either conditioning on it (by concatenating it to the input) or by assuming a placeholder value for it (by concatenating a $\mathbf{0}$ value to the input). At test time, the authors then fine-tune the model on unlabelled data by matching the predictions of the model when using $\mathbf{0}$ as the label with that of the model when conditioning on its own output at the $\mathbf{0}$ label. This essentially leads to a soft idempotence constraint which, according to the authors, allows the model to move out-of-distribution data closer to in-distribution ones and thus improve performance when distribution shifts happen at test time. The authors then evaluate IT$^3$ on a variety of tasks and architectures.
weaknesses: * The main weakness I see in this work is the (almost complete) lack of proper baselines. For example, on only the CIFAR task there is another TTT baseline but on all of the others the authors just compare against not doing anything. This make it hard to gauge significance of the method against other prior art.
* The work makes claims about how idempotence can be seen as a generalisation of a projection and that it allows to map the OOD data to the in-distribution. Thus, while the authors do spend some time to explain why would intuitively their method works, they do not have any ablation studies to verify that these exact intuitions hold in practice. 
* IT$^3$ as a method requires two sequential forward passes through the model, so in practice it can be slow and the authors do not discuss the efficiency of IT$^3$ relative to other works.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: Please refer to the weaknesses section (especially point 2). I am happy to raise my score if my concerns are resolved.
soundness: 2
strengths: The main strengths of this paper:

1- Simplicity of the approach: the method is easy to understand and to implement. 

2- The paper is generally well-written and easy to follow.

3- The breadth of the experiments: The authors are commended for the variety of different tasks that they tested their proposed methods on.
summary: This paper presents a new method for test-time training: at training time the model is trained to be Idempotent, and at test-time two copies of the models are leveraged (one frozen and one adapted) to encourage the model to be idempotent under distribution shifts. Experiments are carried out on different tasks to demonstrate how versatile the proposed method is.
weaknesses: The main weaknesses of this paper are:

1- The method section, while simple, is a bit confusing. Why does the paper present two variants of IT$^3$? The experiments in section 4.6 and Table 1 clearly favors the online version of the method over the one with frozen predictor. This questions why not leveraging the online version in all the experiments? Is there an experimental setup where the offline one is much better than the online one? Why doesn't the paper present one method and treat the other as a special case/variant that is less powerful?

2- The main weakness of this paper the lack of baselines in the experimental section. In all of the presented results, a very suboptimal version of TTT is compared against in **one experiment** only. This really questions how strong is IT$^3$/IT$^3$-online is when compared against strong baselines. Here are some suggestions for necessary experiments:

2a. Since TTT/TTT++[A] are directly comparable with IT$^3$, then I suggest to *at least* have them in all of the experiments with one-to-one comparison in terms of batch size and model architecture. Also, consider adding the performance of IT$^3$ under batch size=1 in Figure 3 results.

2b. Another strong baseline that is suitable for both classification and regression tasks is ActMad [B]. A direct comparison against this baseline is also necessary is all the presented experiments.

2c. Another line of work that is directly comparable is Test-Time Adaptation (TTA). TTA works under a more conservative setup where no control over the training process is assumed. It is also important to compare against the current SOTA TTA method EATA [C] or more closely the dataset distillation method from [D] to further demonstrate the superiority if the proposed method.

2d. Since this is a 'test-time' method, a discussion on its computational requirements is necessary. How would the performance be when evaluated under computational constraints [E]?

2e. Benchmarks used in this work are somewhat small scale. Experiments on larger benchmarks such as ImageNet-C [F] and ImageNet-3DCC [G] in the classification setting are necessary. Similar arguments follow for regression tasks where for example one can follow the object detection experiments from ActMAD. 

3- In section 4.1, I am not sure about the Distribution Shift introduced in this experiment. For instance, the performance of non-optimized does not constantly degrade. Why zeroing out features is a good way of modeling distribution shift instead of adding random noise to the features? Can you please comment on this and provide justification for their choice of distribution shift.

4- Missing references: [B, C, D, E, F, G].

[A] TTT++: When Does Self-Supervised Test-Time Training Fail or Thrive?, NeurIPS 2021

[B] ActMAD: Activation Matching to Align Distributions for Test-Time-Training, CVPR 2023

[C] Efficient Test-Time Model Adaptation without Forgetting, ICML 2022

[D] Leveraging Proxy of Training Data for Test-Time Adaptation, ICML 2023

[E] Evaluation of Test-Time Adaptation Under Computational Time Constraints, ICML 2024

[F] Benchmarking Neural Network Robustness to Common Corruptions and Perturbations, ICLR 2019

[G] 3D Common Corruptions and Data Augmentation, CVPR 2022

Meta-review: This paper addresses the test time training problem to handle distribution shifts. The paper proposes a method to learn a model that is idempotent. This is an interesting idea which the reviewers appreciated. The authors' rebuttal was considered and some of the reviewers also engaged in discussions with the authors.

However, after the discussions the reviewers still expressed several concerns, such as lack of analysis regarding why idempotent training is helpful for out of distribution settings as pointed out by Reviewer hxXL.
 
Reviewer ALWg raised some concerns about the extra inference time and the authors reported some additional experiments and promised to include these results in the revised manuscript.

Reviewer x3JD expressed concerns regarding lack of strong baselines and comparisons with TTA methods.

I appreciate the authors efforts to report some additional results in response to the reviewers' comments. However, the reviewers still had their reservations and, in the current state even with these results considered, the paper does not appear ready to be accepted. The paper's idea is interesting and incorporating the suggestions from the reviewers will definitely strengthen the paper. The authors are advised to address the issues and consider resubmitting to another venue.

Decision: Reject