Improved Test-Time Adaptation for Domain Generalization

Liang ChenYong ZhangYibing SongYing ShanLingqiao Liu

article2023CVPR72 citations

Proposes an improved test-time adaptation framework that optimizes a learnable consistency loss aligned with the primary prediction task and updates only dedicated adaptive parameters to prevent performance degradation on unseen domains.

Listen

Deep learning models frequently suffer severe performance drops when deployed in real-world environments that differ from their training conditions. This challenge, known as distribution shift, is critical in high-stakes visual applications such as autonomous driving and automated object recognition. While domain generalization aims to make models resilient across unseen environments, existing strategies primarily adjust training data and fail to leverage incoming test samples. Test-time adaptation offers a promising fix by updating models on the fly during deployment, but standard test-time methods often rely on handcrafted auxiliary tasks or arbitrary parameter updates that can inadvertently degrade model accuracy.

The article demonstrates that test-time adaptation can be substantially improved through two coordinated mechanisms: a learnable consistency objective that automatically aligns with the main prediction goal, and the insertion of lightweight adaptive parameters tuned exclusively during deployment.

To establish these improvements, the researchers developed the Improved Test-Time Adaptation method and evaluated it across five standard image classification benchmarks encompassing diverse environments, including PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet. The evaluation applied a rigorous protocol across 60 trials per domain setting, using a standard four-block residual network architecture to compare the approach against more than twenty existing domain generalization techniques in both multi-source and single-source scenarios.

The findings show that the proposed method consistently outperforms conventional techniques across diverse test environments. First, in multi-source benchmarks, the method achieved an overall average accuracy of 60.2%, placing first across three of the five tested benchmarks and outperforming all twenty-two competing models on aggregate. Second, in more challenging single-source settings where training information is constrained, the approach led the field with an average accuracy of 68.4%, surpassing the standard empirical baseline by over 3 percentage points and the next-best alternative by more than 2 percentage points. Third, ablation experiments revealed that the learnable consistency objective outperformed common handcrafted auxiliary tasks, such as rotation prediction and entropy minimization. Finally, updating exclusively newly introduced adaptive blocks proved superior to modifying original feature extractor layers or batch normalization statistics, avoiding the performance degradation seen in previous methods.

These results demonstrate that dynamic model adjustment during inference is an effective route to reliable computer vision deployment. Rather than relying on rigid models or manually guessing which internal parameters to update, systems can safely adapt to novel operating conditions in real time without destabilizing core representations. This capability mitigates operational failure risks when vision algorithms encounter unfamiliar physical environments.

Organizations developing computer vision systems for shifting environments should adopt test-time training frameworks equipped with learnable alignment objectives and modular adaptive parameters rather than standard fixed-model approaches. When implementing this method, teams should budget for increased computational requirements during the training phase, as updating the auxiliary weight network introduces additional derivative calculations.

The primary operational limitation is this added training overhead, which requires extra forward and backward processing passes. However, confidence in the performance benefits remains high due to the exhaustive multi-trial validation conducted across varied image domains.

Fudan University Tencent University of Adelaide test application time test-time adaptation domain generalization domain adaptation

arXiv: 2304.04494 liangchen527/ITTA

Paper: Test-Time Training with Self-Supervision for Generalization under Distribution Shifts, Yu Sun et al. (2019). Introduces test-time training with self-supervised auxiliary tasks on unlabeled test inputs, establishing the core inference-time adaptation paradigm that the source improves upon.
Paper: Tent: Fully Test-Time Adaptation by Entropy Minimization, Dequan Wang et al. (2021). Pioneers fully test-time adaptation via entropy minimization and selective parameter updates, which serves as a primary baseline and comparison point for the source's learnable consistency and modular adaptive parameter approach.
Paper: In Search of Lost Domain Generalization, Ishaan Gulrajani et al. (2020). Establishes the standardized DomainBed benchmark and multi-domain evaluation protocols that define how modern domain generalization methods, including the source paper, are assessed.
Paper: Domain Generalization: A Survey, Kaiyang Zhou et al. (2021). Provides a comprehensive taxonomy and analysis of domain generalization paradigms, highlighting the limitations of fixed training-time methods that motivate test-time adaptation.
Paper: Continual Test-Time Domain Adaptation, Qin Wang et al. (2022). Analyzes the failure modes of standard test-time adaptation such as catastrophic forgetting and error accumulation under distribution shifts, motivating the source's need for stable, modular parameter updates.
Paper: Contrastive Test-Time Adaptation, Dian Chen et al. (2022). Presents contrastive and consistency-based objectives for adapting models to unlabeled target data at test time without source data access.
Paper: Do We Really Need to Access the Source Data? Source Hypothesis Transfer for Unsupervised Domain Adaptation, Jian Liang et al. (2020). Demonstrates source-free hypothesis transfer by freezing classification components and adapting feature extractors, laying the groundwork for test-time adaptation without accessing source data.
Paper: Deeper, Broader and Artier Domain Generalization, Da Li et al. (2017). Introduces the PACS multi-domain benchmark that forms a core part of the experimental evaluation in the source.
Paper: Moment Matching for Multi-Source Domain Adaptation, Xingchao Peng et al. (2018). Introduces DomainNet, one of the primary multi-source domain adaptation and generalization benchmark datasets evaluated by the source.

Paper: Feature Alignment and Uniformity for Test Time Adaptation, Shuai Wang et al. (2023). Extends test-time adaptation principles by analyzing feature alignment and uniformity to refine batch-wise feature distributions during inference.
Paper: Align Your Prompts: Test-Time Prompting with Distribution Alignment for Zero-Shot Generalization, Jameel Abdul Samadh et al. (2023). Applies test-time adaptation concepts to large foundation models by aligning multi-modal prompt tokens with source feature statistics on the fly.

Cover for Improved Test-Time Adaptation for Domain Generalization

Abstract

The main challenge in domain generalization (DG) is to handle the distribution shift problem that lies between the training and test data. Recent studies suggest that test-time training (TTT), which adapts the learned model with test data, might be a promising solution to the problem. Generally, a TTT strategy hinges its performance on two main factors: selecting an appropriate auxiliary TTT task for updating and identifying reliable parameters to update during the test phase. Both previous arts and our experiments indicate that TTT may not improve but be detrimental to the learned model if those two factors are not properly considered. This work addresses those two factors by proposing an Improved Test-Time Adaptation (ITTA) method. First, instead of heuristically defining an auxiliary objective, we propose a learnable consistency loss for the TTT task, which contains learnable parameters that can be adjusted toward better alignment between our TTT task and the main prediction task. Second, we introduce additional adaptive parameters for the trained model, and we suggest only updating the adaptive parameters during the test phase. Through extensive experiments, we show that the proposed two strategies are beneficial for the learned model (see Figure 1), and ITTA could achieve superior performance to the current state-of-the-art methods on several DG benchmarks. Code is available at https://github.com/liangchen527/ITTA.

Topics

test application time

Keywords

test-time adaptation domain generalization domain adaptation multi-source generalization learnable consistency loss contrastive learning adaptive risk minimization single-source DG tasks self-supervised learning adversarial training

1. Introduction
2. Related Works
2.1. Domain Generalization.
2.2. Test-Time Training and Adaptation
3. Methodology
3.1. A Learnable Consistency Loss for TTT
3.2. Including Additional Adaptive Parameters
4. Experiments
4.1. Settings
4.2. Multi-Source Generalization
4.3. Single-Source Generalization
5. Analysis
5.1. Compared with Other TTT-Based Models
5.2. Effectiveness of the Learnable Consistency Loss
5.3. Effectiveness of the Adaptive Parameters
5.4. Limitation
6. Conclusion
References

Knowls

Knowl 1 — Learnable Consistency Loss via Gradient Alignment for Test-Time Training
model/method
In test-time training (TTT), heuristically selected self-supervised auxiliary tasks may misalign with the primary prediction objective and degrade adaptation performance. Improved Test-Time Adaptation (ITTA) resolves this by parameterizing an auxiliary consistency loss with a weight subnetwork $f_w$ and optimizing $f_w$ so that the gradients of the auxiliary consistency loss align directly with the gradients of the primary classification task.

Given a feature representation $z \in \mathbb{R}^D$ of an input image $x$ produced by feature extractor $f_\theta$ , and its augmented counterpart $z' \in \mathbb{R}^D$ produced via intermediate feature-level augmentation (such as MixStyle), the weighted consistency loss $\mathcal{L}_{wcont}$ is defined as:
$\mathcal{L}_{wcont} = \| f_w(z - z') \|_2$
where $\| \cdot \|_2$ denotes the $L_2$ norm. The weight subnetwork $f_w: \mathbb{R}^D \to \mathbb{R}^D$ is parameterized by $w$ and consists of stacked dimension-wise piecewise-linear mapping layers: $\text{ReLU}(a \odot h + b)$ , where $a, b \in \mathbb{R}^D$ are learnable weight and bias vectors.

During training on source domain data with classification labels $y$ , the feature extractor $f_\theta$ and classifier $f_\phi$ are trained on the joint objective:
$\min_{\{\theta, \phi\}} \mathcal{L}_{main} + \alpha \mathcal{L}_{wcont}$
where $\mathcal{L}_{main} = \text{CE}(f_\phi(z), y) + \text{CE}(f_\phi(z'), y)$ , $\text{CE}(\cdot)$ denotes cross-entropy loss, and $\alpha > 0$ is a trade-off hyperparameter (set to $\alpha = 1$ ).

To align the optimization trajectory of the auxiliary consistency task with the main classification task, $f_w$ is updated by minimizing the gradient alignment loss:
$\min_w \mathcal{L}_{align} = \| \hat{\mathbf{g}}_{main} - \hat{\mathbf{g}}_{wcont} \|_2$
where $\mathbf{g}_{main} = \nabla_\theta \mathcal{L}_{main}$ , $\mathbf{g}_{wcont} = \nabla_\theta \mathcal{L}_{wcont}$ , and $\hat{\mathbf{g}} = \frac{\mathbf{g} - \mathbb{E}[\mathbf{g}]}{\sigma_{\mathbf{g}}}$ denotes the mean-subtracted, variance-normalized gradient vector.
Knowl 2 — Modular Adaptive Parameter Insertion for Test-Time Adaptation
model/method
Updating entire deep neural networks or heuristically chosen sub-modules (such as batch normalization layers) during test-time adaptation can lead to unstable adaptation or degrade representations under distribution shifts. Improved Test-Time Adaptation (ITTA) freezes all parameters of the pretrained backbone feature extractor $f_\theta$ and classifier $f_\phi$ at test time, and introduces dedicated, lightweight adaptive blocks $f_\Theta$ positioned sequentially after each backbone block.

For a feature extractor containing $m$ sequential blocks $f_\theta^1, \dots, f_\theta^m$ , an adaptive block $f_\Theta^i$ parameterized by $\Theta^i$ is appended after each block $i \in \{1, \dots, m\}$ . Forward propagation for an input $x$ is defined as:
$z^i = f_\Theta^i(f_\theta^i(z^{i-1})), \quad \text{with } z^0 = x$
Each adaptive module $f_\Theta^i$ preserves the spatial and channel dimensions of the intermediate feature representation and is implemented with 5 layers of piecewise-linear transformations $\text{ReLU}(a \odot h + b)$ , initialized with weights $a = \mathbf{1}$ and biases $b = \mathbf{0}$ .

During test-time adaptation on an unseen target domain $D_t$ , only the parameters $\Theta = \{\Theta^1, \dots, \Theta^m\}$ are updated online using the learned consistency objective:
$\min_\Theta \| f_w(z - z') \|_2, \quad \text{s.t. } \{z, z'\} = f_\Theta(f_\theta(x))$
Following test adaptation on sample $x$ , the adapted representation $z = z^m$ is passed to the classifier $f_\phi$ to output the final prediction $f_\phi(z)$ .
Knowl 3 — ITTA Training Phase Algorithm
algorithm
The training phase of Improved Test-Time Adaptation (ITTA) alternates between updating the feature extractor $f_\theta$ and classifier $f_\phi$ on the combined task loss, and updating the consistency weight subnetwork $f_w$ using normalized gradient matching.
Input: Minibatches of source domain samples $(x, y)$ , feature extractor $f_\theta$ , classifier $f_\phi$ , weight subnetwork $f_w$ , balance parameter $\alpha$
Output: Trained parameters $\theta, \phi, w$
for each minibatch $(x, y)$ in training loader do
    Forward pass through feature extractor to obtain original and augmented representations:
     $z, z' \leftarrow f_\theta(x)$

    Compute main loss and consistency loss:
     $\mathcal{L}_{main} \leftarrow \text{CrossEntropy}(f_\phi(z), y) + \text{CrossEntropy}(f_\phi(z'), y)$
     $\mathcal{L}_{wcont} \leftarrow \| f_w(z - z') \|_2$

    Zero gradients of $\theta, \phi$
    Backpropagate joint objective: $\nabla_{\theta, \phi} (\mathcal{L}_{main} + \alpha \mathcal{L}_{wcont})$
    Update parameters $\theta$ and $\phi$ via SGD

    Zero gradients of $\theta$
    Backpropagate $\mathcal{L}_{main}$ and extract $\mathbf{g}_{main} \leftarrow \nabla_\theta \mathcal{L}_{main}$
    Standardize gradient: $\hat{\mathbf{g}}_{main} \leftarrow (\mathbf{g}_{main} - \mathbb{E}[\mathbf{g}_{main}]) / \sigma_{\mathbf{g}_{main}}$

    Zero gradients of $\theta$
    Backpropagate $\mathcal{L}_{wcont}$ and extract $\mathbf{g}_{wcont} \leftarrow \nabla_\theta \mathcal{L}_{wcont}$
    Standardize gradient: $\hat{\mathbf{g}}_{wcont} \leftarrow (\mathbf{g}_{wcont} - \mathbb{E}[\mathbf{g}_{wcont}]) / \sigma_{\mathbf{g}_{wcont}}$

    Compute alignment objective $\mathcal{L}_{align} \leftarrow \text{MSE}(\hat{\mathbf{g}}_{main}, \hat{\mathbf{g}}_{wcont})$
    Zero gradients of $w$
    Backpropagate $\nabla_w \mathcal{L}_{align}$
    Update parameters $w$ of weight subnetwork $f_w$ via SGD
end for
return $f_\theta, f_\phi, f_w$
Knowl 4 — ITTA Test-Time Adaptation Phase Algorithm
algorithm
During inference on unseen target domains, ITTA keeps the pretrained feature extractor $f_\theta$ , classifier $f_\phi$ , and weight subnetwork $f_w$ fixed, and performs test-time online adaptation exclusively on the newly added adaptive blocks $f_\Theta$ .
Input: Test batches $x$ from target domain $D_t$ , pretrained feature extractor $f_\theta = \{f_\theta^1, \dots, f_\theta^m\}$ , pretrained classifier $f_\phi$ , pretrained weight subnetwork $f_w$ , adaptive blocks $f_\Theta = \{f_\Theta^1, \dots, f_\Theta^m\}$
Output: Predictions $\hat{y}$ for test samples $x$
for each test batch $x$ in test loader do
    Forward pass through the feature extractor interleaved with adaptive blocks:
     $z^1, z'^1 \leftarrow f_\Theta^1(f_\theta^1(x)), f_\Theta^1(f_\theta^1(x'))$
    for $i = 2$ to $m$ do
         $z^i, z'^i \leftarrow f_\Theta^i(f_\theta^i(z^{i-1})), f_\Theta^i(f_\theta^i(z'^{i-1}))$
    end for
     $z \leftarrow z^m, z' \leftarrow z'^m$

    Compute learned consistency loss:
     $\mathcal{L}_{wcont} \leftarrow \| f_w(z - z') \|_2$

    Zero gradients of adaptive parameters $\Theta$
    Backpropagate $\nabla_\Theta \mathcal{L}_{wcont}$
    Update $\Theta$ via SGD

    Forward pass adapted representation for prediction:
     $z \leftarrow f_\Theta(f_\theta(x))$
     $\hat{y} \leftarrow f_\phi(z)$
end for
return $\hat{y}$

Knowl 5 — Multi-Source Domain Generalization Performance on DomainBed

data/table

Under the DomainBed benchmark evaluation protocol with a ResNet-18 backbone and 60 trials per unseen domain using training-domain validation selection, ITTA achieves superior average classification accuracy across 5 standard DG benchmark datasets compared to 21 competitive baseline algorithms.

Method	PACS	VLCS	OfficeHome	TerraInc	DomainNet	Avg.
MMD	81.3 $\pm$ 0.8	74.9 $\pm$ 0.5	59.9 $\pm$ 0.4	42.0 $\pm$ 1.0	7.9 $\pm$ 6.2	53.2
RSC	80.5 $\pm$ 0.2	75.4 $\pm$ 0.3	58.4 $\pm$ 0.6	39.4 $\pm$ 1.3	27.9 $\pm$ 2.0	56.3
IRM	80.9 $\pm$ 0.5	75.1 $\pm$ 0.1	58.0 $\pm$ 0.1	38.4 $\pm$ 0.9	30.4 $\pm$ 1.0	56.6
ARM	80.6 $\pm$ 0.5	75.9 $\pm$ 0.3	59.6 $\pm$ 0.3	37.4 $\pm$ 1.9	29.9 $\pm$ 0.1	56.7
DANN	79.2 $\pm$ 0.3	76.3 $\pm$ 0.2	59.5 $\pm$ 0.5	37.9 $\pm$ 0.9	31.5 $\pm$ 0.1	56.9
GroupDRO	80.7 $\pm$ 0.4	75.4 $\pm$ 1.0	60.6 $\pm$ 0.3	41.5 $\pm$ 2.0	27.5 $\pm$ 0.1	57.1
CDANN	80.3 $\pm$ 0.5	76.0 $\pm$ 0.5	59.3 $\pm$ 0.4	38.6 $\pm$ 2.3	31.8 $\pm$ 0.2	57.2
VREx	80.2 $\pm$ 0.5	75.3 $\pm$ 0.6	59.5 $\pm$ 0.1	43.2 $\pm$ 0.3	28.1 $\pm$ 1.0	57.3
CAD	81.9 $\pm$ 0.3	75.2 $\pm$ 0.6	60.5 $\pm$ 0.3	40.5 $\pm$ 0.4	31.0 $\pm$ 0.8	57.8
CondCAD	80.8 $\pm$ 0.5	76.1 $\pm$ 0.3	61.0 $\pm$ 0.4	39.7 $\pm$ 0.4	31.9 $\pm$ 0.7	57.9
MTL	80.1 $\pm$ 0.8	75.2 $\pm$ 0.3	59.9 $\pm$ 0.5	40.4 $\pm$ 1.0	35.0 $\pm$ 0.0	58.1
ERM	79.8 $\pm$ 0.4	75.8 $\pm$ 0.2	60.6 $\pm$ 0.2	38.8 $\pm$ 1.0	35.3 $\pm$ 0.1	58.1
MixStyle	82.6 $\pm$ 0.4	75.2 $\pm$ 0.7	59.6 $\pm$ 0.8	40.9 $\pm$ 1.1	33.9 $\pm$ 0.1	58.4
MLDG	81.3 $\pm$ 0.2	75.2 $\pm$ 0.3	60.9 $\pm$ 0.2	40.1 $\pm$ 0.9	35.4 $\pm$ 0.0	58.6
Mixup	79.2 $\pm$ 0.9	76.2 $\pm$ 0.3	61.7 $\pm$ 0.5	42.1 $\pm$ 0.7	34.0 $\pm$ 0.0	58.6
Fishr	81.3 $\pm$ 0.3	76.2 $\pm$ 0.3	60.9 $\pm$ 0.3	42.6 $\pm$ 1.0	34.2 $\pm$ 0.3	59.0
SagNet	81.7 $\pm$ 0.6	75.4 $\pm$ 0.8	62.5 $\pm$ 0.3	40.6 $\pm$ 1.5	35.3 $\pm$ 0.1	59.1
SelfReg	81.8 $\pm$ 0.3	76.4 $\pm$ 0.7	62.4 $\pm$ 0.1	41.3 $\pm$ 0.3	34.7 $\pm$ 0.2	59.3
Fish	82.0 $\pm$ 0.3	76.9 $\pm$ 0.2	62.0 $\pm$ 0.6	40.2 $\pm$ 0.6	35.5 $\pm$ 0.0	59.3
CORAL	81.7 $\pm$ 0.0	75.5 $\pm$ 0.4	62.4 $\pm$ 0.4	41.4 $\pm$ 1.8	36.1 $\pm$ 0.2	59.4
SD	81.9 $\pm$ 0.3	75.5 $\pm$ 0.4	62.9 $\pm$ 0.2	42.0 $\pm$ 1.0	36.3 $\pm$ 0.2	59.7
Ours (ITTA)	83.8 $\pm$ 0.3	76.9 $\pm$ 0.6	62.0 $\pm$ 0.2	43.2 $\pm$ 0.5	34.9 $\pm$ 0.1	60.2

ITTA achieves the highest overall multi-source generalization accuracy of 60.2%, exceeding empirical risk minimization (ERM, 58.1%) by 2.1 percentage points and top prior domain generalization techniques (e.g., SD at 59.7%). It obtains the best individual accuracy on PACS (83.8%), VLCS (76.9%), and TerraInc (43.2%).

Knowl 6 — Single-Source Domain Generalization Accuracy on PACS

data/table

In single-source domain generalization on PACS (training on 1 domain and evaluating zero-shot on the remaining 3 domains across 12 source-to-target evaluation splits with 60 trials each using ResNet-18), ITTA achieves state-of-the-art performance.

Method	A $\to$ C	A $\to$ P	A $\to$ S	C $\to$ A	C $\to$ P	C $\to$ S	P $\to$ A	P $\to$ C	P $\to$ S	S $\to$ A	S $\to$ C	S $\to$ P	Avg.
RSC	66.3	88.2	57.2	65.8	82.4	68.7	60.5	41.3	53.1	53.8	65.9	48.4	62.6
Fish	67.1	89.2	57.0	66.7	85.6	64.5	55.1	33.9	51.2	59.1	67.1	58.4	62.9
CDANN	66.5	92.2	65.0	70.6	82.9	67.7	60.6	42.2	46.9	51.4	60.7	51.9	63.2
SelfReg	63.9	90.1	56.8	70.2	85.4	70.2	60.9	38.8	50.5	54.5	66.2	51.7	63.3
DANN	67.5	91.2	67.5	70.6	81.4	66.6	54.1	33.5	52.8	53.8	64.4	58.9	63.5
CAD	67.1	89.6	60.2	67.7	83.7	70.2	60.6	38.3	53.8	50.7	65.8	54.4	63.5
GroupDRO	66.5	90.5	58.9	70.8	85.7	69.7	62.3	41.1	48.2	54.8	65.2	53.9	64.0
MTL	67.3	90.1	58.9	70.2	84.2	71.9	58.3	38.5	52.8	55.4	66.1	55.2	64.1
IRM	67.5	93.0	62.9	67.6	83.8	68.9	63.7	39.9	49.0	54.9	63.1	54.9	64.1
ARM	66.0	91.2	58.7	70.6	84.2	69.1	59.2	42.1	52.1	60.0	62.9	53.8	64.2
Mixup	65.5	87.8	57.2	71.4	83.1	68.0	59.6	37.2	56.5	55.0	66.2	62.7	64.2
CORAL	66.8	90.3	61.5	67.9	85.4	70.4	55.9	40.4	49.8	55.8	67.6	58.9	64.2
SD	67.1	91.7	63.7	70.3	84.4	69.4	57.5	42.6	47.7	55.9	65.7	55.8	64.3
MMD	67.1	88.0	63.6	70.0	83.6	70.2	58.8	40.3	52.3	57.4	68.7	52.7	64.4
MLDG	67.3	90.8	64.4	70.8	84.2	69.7	61.6	41.3	50.4	49.9	66.8	58.7	64.7
CondCAD	66.9	92.3	60.8	71.0	84.7	72.6	61.2	40.7	55.7	52.3	64.2	55.3	64.8
ERM	67.3	91.7	60.1	70.4	82.3	68.1	59.6	44.7	56.5	52.8	68.1	58.4	65.0
VREx	67.1	91.0	62.6	71.1	84.1	71.7	62.4	37.7	53.6	60.6	66.7	57.5	65.5
Fishr	67.9	92.7	62.4	71.2	83.4	70.2	60.0	42.7	57.1	55.7	68.4	62.0	66.1
SagNet	67.6	92.3	59.5	71.8	82.8	69.9	62.5	45.2	64.1	55.8	65.7	55.9	66.1
MixStyle	68.5	91.2	65.1	73.2	85.0	71.7	63.6	46.3	51.6	54.2	67.0	58.3	66.3
Ours (ITTA)	68.9	92.4	62.5	75.3	85.9	70.2	66.5	52.2	63.8	57.6	68.0	57.9	68.4

(Domain abbreviations: Art [A], Cartoon [C], Photo [P], Sketch [S]).

ITTA obtains an average single-source generalization accuracy of 68.4%, outperforming ERM (65.0%) by 3.4% and the previous top model MixStyle (66.3%) by 2.1%, while leading in 5 out of the 12 single-source transfer tasks.

Knowl 7 — Comparison of ITTA with Test-Time Training Baselines on PACS

data/table

Using an identical baseline model (ResNet-18 backbone with MixStyle feature augmentation), ITTA is compared on the PACS leave-one-out benchmark against test-time adaptation and training approaches:

Model	Art	Cartoon	Photo	Sketch	Avg.
Baseline	79.9 $\pm$ 0.5	75.4 $\pm$ 1.1	94.4 $\pm$ 0.5	75.8 $\pm$ 1.2	81.4 $\pm$ 0.5
TTT	81.5 $\pm$ 0.8	77.6 $\pm$ 0.6	94.3 $\pm$ 0.2	78.4 $\pm$ 0.7	83.0 $\pm$ 0.2
MT3	82.0 $\pm$ 1.0	76.5 $\pm$ 1.0	94.1 $\pm$ 0.2	77.7 $\pm$ 1.3	82.6 $\pm$ 0.6
TENT	80.2 $\pm$ 0.9	77.2 $\pm$ 0.8	94.4 $\pm$ 0.2	77.4 $\pm$ 0.1	82.3 $\pm$ 0.5
Ours (ITTA)	84.7 $\pm$ 0.4	78.0 $\pm$ 0.4	94.5 $\pm$ 0.4	78.2 $\pm$ 0.3	83.8 $\pm$ 0.3

TTT (rotation prediction), MT3 (MAML-based contrastive loss), and TENT (test-time entropy minimization on batch normalization parameters) improve the baseline from 81.4% to 82.3%–83.0%. ITTA attains 83.8% average accuracy, outperforming all baseline test-time training methods.

Knowl 8 — Ablation of TTT Objectives and Parameter Selection Strategies

empirical result

An ablation study on PACS using a ResNet-18 backbone evaluates the performance of different auxiliary TTT objectives (entropy minimization [Ent.], rotation estimation [Rot.], naive consistency loss without $f_w$ , and the proposed learnable consistency loss $\mathcal{L}_{wcont}$ ) alongside parameter updating strategies (updating all backbone parameters [All], batch normalization layers only [BN], and the proposed adaptive parameters [Ada.]):

Model	Ent.	Rot.	$\mathcal{L}_{wcont}$	All	BN	Ada.	Art	Cartoon	Photo	Sketch	Avg.
Ours (Full ITTA)	–	–	✓	–	–	✓	84.7 $\pm$ 0.4	78.0 $\pm$ 0.4	94.5 $\pm$ 0.4	78.2 $\pm$ 0.3	83.8 $\pm$ 0.3
Ours w/o $f_w$	–	–	–	–	–	✓	83.1 $\pm$ 0.4	74.6 $\pm$ 0.6	94.0 $\pm$ 0.5	78.0 $\pm$ 0.8	82.5 $\pm$ 0.1
Ours w/ Ent.	✓	–	–	–	–	✓	79.9 $\pm$ 2.4	77.3 $\pm$ 0.3	94.8 $\pm$ 0.8	77.6 $\pm$ 0.4	82.4 $\pm$ 0.8
Ours w/ Rot.	–	✓	–	–	–	✓	81.1 $\pm$ 1.0	75.2 $\pm$ 0.5	94.9 $\pm$ 0.3	77.3 $\pm$ 0.6	82.1 $\pm$ 0.3
Ours w/o TTT	–	–	✓	–	–	–	83.3 $\pm$ 0.5	76.0 $\pm$ 0.5	94.4 $\pm$ 0.5	76.7 $\pm$ 1.4	82.8 $\pm$ 0.3
Ours w/ All	–	–	✓	✓	–	–	83.0 $\pm$ 0.7	77.0 $\pm$ 1.4	94.5 $\pm$ 0.7	77.4 $\pm$ 0.9	83.0 $\pm$ 0.2
Ours w/ BN	–	–	✓	–	✓	–	81.8 $\pm$ 0.5	75.6 $\pm$ 0.3	94.4 $\pm$ 0.3	77.9 $\pm$ 1.1	82.4 $\pm$ 0.5

The experiments show:

Removing the learnable weight subnetwork $f_w$ (naive consistency loss) drops accuracy from 83.8% to 82.5%, matching the performance of unaligned auxiliary tasks (entropy minimization at 82.4% and rotation prediction at 82.1%).
Updating only BN layers under the learned consistency loss results in 82.4% accuracy (worse than no test adaptation at 82.8%), and updating all backbone parameters reaches 83.0%, whereas updating the dedicated adaptive blocks $f_\Theta$ achieves 83.8%.

Knowl 9 — Computational Overhead of Second-Order Gradient Alignment
limitation
While standard joint training with a fixed auxiliary loss incurs 1 forward pass and 1 backward pass per iteration, updating the weight subnetwork $f_w$ in ITTA requires optimizing the alignment objective $\mathcal{L}_{align} = \| \hat{\mathbf{g}}_{main} - \hat{\mathbf{g}}_{wcont} \|_2$ . Because $\mathbf{g}_{main}$ and $\mathbf{g}_{wcont}$ are themselves first-order gradients of feature extractor parameters $\theta$ , computing $\nabla_w \mathcal{L}_{align}$ requires evaluating second-order derivatives, requiring 1 additional forward pass and 3 additional backward passes per training step, thereby adding computational overhead during training.

Coverage note — None was omitted; all key algorithmic mechanisms, mathematical formulations, benchmark results, ablations, and stated limitations are included.

References

1.Martin Arjovsky, Leon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.
2.Yogesh Balaji, Swami Sankaranarayanan, and Rama Chellappa. Metareg: Towards domain generalization using meta-regularization. In NeurIPS, 2018.
3.Alexander Bartler, Andre Buhler, Felix Wiewel, Mario Dobler, and Bin Yang. Mt3: Meta test-time training for self-supervised test-time adaption. In AISTATS, 2022.
4.Sara Beery, Grant Van Horn, and Pietro Perona. Recognition in terra incognita. In ECCV, 2018.
5.Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In NeurIPS, 2006.
6.Gilles Blanchard, Aniket Anand Deshmukh, Urun Dogan, Gyemin Lee, and Clayton Scott. Domain generalization by marginal transfer learning. arXiv preprint arXiv:1711.07910, 2017.
7.Gilles Blanchard, Gyemin Lee, and Clayton Scott. Generalizing from several related classification tasks to a new unlabeled sample. In NeurIPS, 2011.
8.Chaoqi Chen, Jiongcheng Li, Xiaoguang Han, Xiaoqing Liu, and Yizhou Yu. Compound domain generalization via meta-knowledge encoding. In CVPR, 2022.
9.Chaoqi Chen, Luyao Tang, Feng Liu, Gangming Zhao, Yue Huang, and Yizhou Yu. Mix and reason: Reasoning over semantic topology with data mixing for domain generalization. In NeurIPS, 2022.
10.Dian Chen, Dequan Wang, Trevor Darrell, and Sayna Ebrahimi. Contrastive test-time adaptation. In CVPR, 2022.
11.Liang Chen, Yong Zhang, Yibing Song, Lingqiao Liu, and Jue Wang. Self-supervised learning of adversarial example: Towards good generalizations for deepfake detection. In CVPR, 2022.
12.Liang Chen, Yong Zhang, Yibing Song, Jue Wang, and Lingqiao Liu. Ost: Improving generalization of deepfake detection via one-shot test-time training. In NeurIPS, 2022.
13.Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020.
14.Sungha Choi, Seunghan Yang, Seokeon Choi, and Sungrack Yun. Improving test-time adaptation via shift-agnostic weight regularization and nearest source prototypes. In ECCV, 2022.
15.Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
16.Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker. Domain generalization via model-agnostic learning of semantic features. In NeurIPS, 2019.
17.Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes (voc) challenge. IJCV, 88(2):303–338, 2010.
18.Chen Fang, Ye Xu, and Daniel N Rockmore. Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias. In ICCV, 2013.
19.Li Fei-Fei, Rob Fergus, and Pietro Perona. Learning generative visual models from few training examples: An incremental bayesian approach tested on 101 object categories. In CVPR worksho, 2004.
20.Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In ICML, 2017.
21.Francois Fleuret et al. Uncertainty reduction for model adaptation in semantic segmentation. In CVPR, 2021.
22.Yossi Gandelsman, Yu Sun, Xinlei Chen, and Alexei A Efros. Test-time training with masked autoencoders. In NeurIPS, 2022.
23.Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, Francois Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. JMLR, 17(1):2096–2030, 2016.
24.Muhammad Ghifary, David Balduzzi, W Bastiaan Kleijn, and Mengjie Zhang. Scatter component analysis: A unified framework for domain adaptation and domain generalization. IEEE TPAMI, 39(7):1414–1430, 2016.
25.Muhammad Ghifary, W Bastiaan Kleijn, Mengjie Zhang, and David Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In ICCV, 2015.
26.Jean-Bastien Grill, Florian Strub, Florent Altche, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. In NeurIPS, 2020.
27.Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. In ICLR, 2021.
28.Sivan Harary, Eli Schwartz, Assaf Arbelle, Peter Staar, Shady Abu-Hussein, Elad Amrani, Roei Herzig, Amit Alfassy, Raja Giryes, Hilde Kuehne, et al. Unsupervised domain generalization by learning a bridge across domains. In CVPR, 2022.
29.Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
30.Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
31.Shoubo Hu, Kun Zhang, Zhitang Chen, and Laiwan Chan. Domain generalization via multidomain discriminant analysis. In UAI, 2020.
32.Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normalization. In ICCV, 2017.
33.Zeyi Huang, Haohan Wang, Eric P Xing, and Dong Huang. Self-challenging improves cross-domain generalization. In ECCV, 2020.
34.Daehee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jaekoo Lee. Selfreg: Self-supervised contrastive regularization for domain generalization. In ICCV, 2021.
35.Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In ICML, 2021.
36.David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In ICML, 2021.
37.Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In ICCV, 2017.
38.Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Learning to generalize: Meta-learning for domain generalization. In AAAI, 2018.
39.Da Li, Jianshu Zhang, Yongxin Yang, Cong Liu, Yi-Zhe Song, and Timothy M Hospedales. Episodic training for domain generalization. In ICCV, 2019.
40.Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot. Domain generalization with adversarial feature learning. In CVPR, 2018.
41.Pan Li, Da Li, Wei Li, Shaogang Gong, Yanwei Fu, and Timothy M Hospedales. A simple feature augmentation for domain generalization. In ICCV, 2021.
42.Xiaotong Li, Yongxing Dai, Yixiao Ge, Jun Liu, Ying Shan, and Ling-Yu Duan. Uncertainty modeling for out-of-distribution generalization. In ICLR, 2022.
43.Yizhuo Li, Miao Hao, Zonglin Di, Nitesh Bharadwaj Gundavarapu, and Xiaolong Wang. Test-time personalization with a transformer for human pose estimation. In NeurIPS, 2021.
44.Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In ECCV, 2018.
45.Yuejiang Liu, Parth Kothari, Bastien van Delft, Baptiste Bellot-Gurlet, Taylor Mordan, and Alexandre Alahi. Ttt++: When does self-supervised test-time training fail or thrive? In NeurIPS, 2021.
46.Krikamol Muandet, David Balduzzi, and Bernhard Scholkopf. Domain generalization via invariant feature representation. In ICML, 2013.
47.Hyeonseob Nam, HyunJae Lee, Jongchan Park, Wonjun Yoon, and Donggeun Yoo. Reducing domain gap by reducing style bias. In CVPR, 2021.
48.Prashant Pandey, Mrigank Raman, Sumanth Varambally, and Prathosh Ap. Generalization on unseen domains via inference-time label-preserving target projections. In CVPR, 2021.
49.Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In ICCV, 2019.
50.Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron C Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks. In NeurIPS, 2021.
51.Alexandre Rame, Corentin Dancette, and Matthieu Cord. Fishr: Invariant gradient variances for out-of-distribution generalization. In ICML, 2022.
52.Yangjun Ruan, Yann Dubois, and Chris J Maddison. Optimal representations for covariate shift. In ICLR, 2022.
53.Bryan C Russell, Antonio Torralba, Kevin P Murphy, and William T Freeman. Labelme: a database and web-based tool for image annotation. IJCV, 77(1):157–173, 2008.
54.Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In ICLR, 2020.
55.Steffen Schneider, Evgenia Rusak, Luisa Eck, Oliver Bringmann, Wieland Brendel, and Matthias Bethge. Improving robustness against common corruptions by covariate shift adaptation. In NeurIPS, 2020.
56.Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In ICCV, 2017.
57.Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. In ICLR, 2021.
58.Baochen Sun and Kate Saenko. Deep coral: Correlation alignment for deep domain adaptation. In ECCV, 2016.
59.Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In ICML, 2020.
60.Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 1999.
61.Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan. Deep hashing network for unsupervised domain adaptation. In CVPR, 2017.
62.Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell. Tent: Fully test-time adaptation by entropy minimization. In ICLR, 2021.
63.Jianxiong Xiao, James Hays, Krista A Ehinger, Aude Oliva, and Antonio Torralba. Sun database: Large-scale scene recognition from abbey to zoo. In CVPR, 2010.
64.Zehao Xiao, Xiantong Zhen, Ling Shao, and Cees GM Snoek. Learning to generalize across domains on single test samples. In ICLR, 2022.
65.Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yanfeng Wang, and Qi Tian. A fourier-based framework for domain generalization. In CVPR, 2021.
66.Zhenlin Xu, Deyi Liu, Junlin Yang, Colin Raffel, and Marc Niethammer. Robust and generalizable visual representation learning via random convolutions. In ICLR, 2021.
67.Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren. Improve unsupervised domain adaptation with mixup training. arXiv preprint arXiv:2001.00677, 2020.
68.Fu-En Yang, Yuan-Chia Cheng, Zu-Yun Shiau, and Yu-Chiang Frank Wang. Adversarial teacher-student representation learning for domain generalization. In NeurIPS, 2021.
69.Nanyang Ye, Kaican Li, Haoyue Bai, Runpeng Yu, Lanqing Hong, Fengwei Zhou, Zhenguo Li, and Jun Zhu. Ood-bench: Quantifying and understanding two dimensions of out-of-distribution generalization. In CVPR, 2022.
70.Fuming You, Jingjing Li, and Zhou Zhao. Test-time batch statistics calibration for covariate shift. arXiv preprint arXiv:2110.04065, 2021.
71.Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: A meta-learning approach for tackling group distribution shift. arXiv preprint arXiv:2007.02931, 2020.
72.Marvin Zhang, Henrik Marklund, Nikita Dhawan, Abhishek Gupta, Sergey Levine, and Chelsea Finn. Adaptive risk minimization: Learning to adapt to domain shift. NeurIPS, 2021.
73.Tao Zhong, Zhixiang Chi, Li Gu, Yang Wang, Yuanhao Yu, and Jin Tang. Meta-dmoe: Adapting to domain shift by meta-distillation from mixture-of-experts. In NeurIPS, 2022.
74.Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang. Domain generalization with mixstyle. In ICLR, 2021.

Citation

MLA

Chen, L., et al. “Improved Test-Time Adaptation for Domain Generalization”. arXiv, 2023, http://arxiv.org/abs/2304.04494v2.

APA

Chen, L., Zhang, Y., Song, Y., Shan, Y., & Liu, L. (2023). Improved Test-Time Adaptation for Domain Generalization. arXiv. http://arxiv.org/abs/2304.04494v2

Chicago

Chen, L., Y. Zhang, Y. Song, Y. Shan, and L. Liu. 2023. “Improved Test-Time Adaptation for Domain Generalization”. arXiv. http://arxiv.org/abs/2304.04494v2.

Harvard

Chen, L. et al. (2023) “Improved Test-Time Adaptation for Domain Generalization”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2304.04494v2.

Vancouver

1. Chen L, Zhang Y, Song Y, Shan Y, Liu L (2023) Improved Test-Time Adaptation for Domain Generalization. arXiv

BibTeX

@article{chen2023improved,
  title = {Improved Test-Time Adaptation for Domain Generalization},
  author = {Chen, Liang and Zhang, Yong and Song, Yibing and Shan, Ying and Liu, Lingqiao},
  year = {2023},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2304.04494v2},
  eprint = {2304.04494}
}

Metadata:arXiv

Source Code

This paper has an official code repository available. Click below to access the source code.

View Repository

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: IEEE

Improved Test-Time Adaptation for Domain Generalization

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Learnable Consistency Loss via Gradient Alignment for Test-Time Training

Knowl 2 — Modular Adaptive Parameter Insertion for Test-Time Adaptation

Knowl 3 — ITTA Training Phase Algorithm

Knowl 4 — ITTA Test-Time Adaptation Phase Algorithm

Knowl 5 — Multi-Source Domain Generalization Performance on DomainBed

Knowl 6 — Single-Source Domain Generalization Accuracy on PACS

Knowl 7 — Comparison of ITTA with Test-Time Training Baselines on PACS

Knowl 8 — Ablation of TTT Objectives and Parameter Selection Strategies

Knowl 9 — Computational Overhead of Second-Order Gradient Alignment

References

Citation

BibTeX

Source Code

Access the Paper