PaperlayerView exact revision on arXiv

Source-linked AI summary

TTN: A Domain-Shift Aware Batch Normalization in Test-Time Adaptation

Hyesu Lim, Byeonggeun Kim, Jaegul Choo, Sungha Choi

arXiv:2302.05155v2cs.CV

TL;DR

Test-time adaptation must address performance degradation under source-to-target domain shift, while prior TBN-based approaches have largely been evaluated with large batches and stationary shifts. TTN mixes source and test-batch statistics according to layer sensitivity, and improves existing TTA methods across varied batch sizes and realistic scenarios.

  • Problem

    TTA must address source-to-target domain shift, but prior TBN-based evaluations have mainly assumed large test batches and a single stationary shift.

    Source: §1 · p. 1

  • Method

    TTN interpolates source and test-batch statistics according to each BN layer's sensitivity to domain shift, with interpolation weights optimized before test time.

    Source: §1 · p. 2

  • Results

    TTN outperforms existing TTA methods across batch sizes and single, mixed, and continuously changing shifts in image classification and semantic segmentation.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    TTN can complement other optimization-based TTA methods without extra test-time computation beyond per-batch statistics.

    Source: §1 · p. 2

    The stated setup mainly considers varied test batch sizes and batches drawn from single, multiple, or mixed target domains.

    Source: §2.1 · p. 3

Abstract

from arXiv · show

This paper proposes a novel batch normalization strategy for test-time adaptation. Recent test-time adaptation methods heavily rely on the modified batch normalization, i.e., transductive batch normalization (TBN), which calculates the mean and the variance from the current test batch rather than using the running mean and variance obtained from the source data, i.e., conventional batch normalization (CBN). Adopting TBN that employs test batch statistics mitigates the performance degradation caused by the domain shift. However, re-estimating normalization statistics using test data depends on impractical assumptions that a test batch should be large enough and be drawn from i.i.d. stream, and we observed that the previous methods with TBN show critical performance drop without the assumptions. In this paper, we identify that CBN and TBN are in a trade-off relationship and present a new test-time normalization (TTN) method that interpolates the statistics by adjusting the importance between CBN and TBN according to the domain-shift sensitivity of each BN layer. Our proposed TTN improves model robustness to shifted domains across a wide range of batch sizes and in various realistic evaluation scenarios. TTN is widely applicable to other test-time adaptation methods that rely on updating model parameters via backpropagation. We demonstrate that adopting TTN further improves their performance and achieves state-of-the-art performance in various standard benchmarks.

1 INTRODUCTION

Domain shift degrades target-domain performance, while conventional and transductive batch normalization have complementary weaknesses under practical test conditions. TTN interpolates source and test-batch statistics according to layer sensitivity and improves existing test-time adaptation methods across varied batch sizes and scenarios.

1 INTRODUCTION

  • Domain shift degrades target performance, while TBN-based methods are mainly assessed with large batches and stationary shifts and can drop sharply at small batch sizes.CBN and TBN have a trade-off: each can be useful where the other fails.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • TTN interpolates source and test-batch statistics, adjusting their importance according to each BN layer's sensitivity to domain shift.The interpolation weight is optimized after pretraining and before test time.Source: §1 · p. 2
  • TTN trains only newly added interpolation weights, leaving pretrained model weights unchanged.The method is presented as compatible with other optimization-based TTA methods without additional test-time computation beyond per-batch statistics.Source: §1 · p. 2
  • Across image classification and semantic segmentation, TTN outperforms existing TTA methods across batch sizes and single, mixed, and continuously changing shifts.The evaluation covers batch sizes from 200 to 1 and three realistic scenarios.Source: §1 · pp. 2–3

2 METHODOLOGY

The method frames test-time adaptation around distribution shift and practical test settings, including varying batch sizes and multiple target domains. TTN mixes source and test-batch statistics using weights informed by domain-shift sensitivity and optimized during post-training.

2.1 PROBLEM SETUP

  • TTA is framed as managing test batches under covariate shift with a shared label space, considering varied batch sizes and single, multiple, or mixed target domains.Source: §2.1 · p. 3

2.3 POST TRAINING

  • Larger gradient-distance scores between clean and augmented inputs assign greater importance to test-batch statistics for the corresponding BN layers and channels.The prior is formed by comparing affine-parameter gradients from clean and augmented inputs.Source: Figure 3 · p. 4, §2.3 · pp. 4–5

2.2 TEST-TIME NORMALIZATION LAYER

  • TTN linearly combines current-batch and source statistics using learnable α weights while retaining BN affine parameters; α varies by layer and channel.Source: §2.2 · p. 4

2.3 POST TRAINING

  • During post-training, all parameters except α are frozen; augmented source data optimizes α with cross-entropy and an MSE constraint toward prior A.The prior initializes α, and labeled source data is available during post-training.Source: §2.3 · pp. 4–5

3 EXPERIMENTS

Experiments test TTN across corruption benchmarks, batch sizes, domain-shift scenarios, and semantic segmentation. They report robust performance, benefits for other TTA methods, and component and layer-wise analyses of α.

3.1 EXPERIMENTAL SETUP

  • The evaluation covers corruption robustness in classification and natural domain shifts in semantic segmentation, using varied test batch sizes and domain-adaptation scenarios.The image-classification setup includes CIFAR-10/100-C and ImageNet-C; segmentation uses Cityscapes, BDD-100K, Mapillary, GTAV, and SYNTHIA.Source: §3 · p. 5, §3.1 · pp. 5–6

3.2 EXPERIMENTS ON IMAGE CLASSIFICATION

  • TTN performs robustly across batch sizes and single, continuously changing, and mixed corruption settings, outperforming constant-α normalization baselines and favorably comparing with optimization-based methods.The reported comparisons cover Tables 1–3.Source: §3.2 · pp. 6–7, Table 1 · p. 6, Table 2 · p. 6
  • Adding TTN reduces other TTA methods’ dependence on large test batches and improves their performance across evaluation scenarios.The comparison includes TENT, SWR, and CoTTA.Source: §3.2 · p. 7
  • TTN and TTN-adapted methods preserve source-domain knowledge better than TBN-based methods, including when test batches are large.On the source-domain test, TBN-based baselines lose performance even with batch size 200.Source: §3.2 · p. 7, Table 4 · p. 7
  • TTN is more robust than TBN under class imbalance, where skewed test-batch statistics can impair adaptation.The evaluation sorts test images by class label, producing batches that are often dominated by one class.Source: §3.2 · pp. 7–8, Table 5 · p. 7

3.3 EXPERIMENTS ON SEMANTIC SEGMENTATION

  • In semantic segmentation on unseen domains with batch size 2, TTN helps use test-batch statistics where TBN, TENT, and SWR perform worse than the source model.The evaluation uses a Cityscapes-trained model on real-world and synthetic datasets.Source: §3.3 · p. 8, Table 6 · p. 8

3.4 ABLATION STUDIES

  • MSE regularization helps optimized α remain effective across batch sizes, whereas CE optimization without MSE can overfit to large batches.Initializing α with prior A does not significantly change results, though it gives a better starting point than random initialization in one comparison.Source: §3.4 · p. 8, Table 7 · p. 8

3.5 VISUALIZATION OF α

  • Optimized α decreases from shallow to deep layers, indicating greater TBN weighting in shallow layers and greater CBN weighting in deeper layers.The authors interpret this pattern as domain-shift information diminishing in deeper layers, where content information remains.Source: §3.5 · p. 8, Figure 4 · p. 9

4 RELATED WORK

The related work distinguishes test-time adaptation from domain-adaptation settings with access to larger target datasets, and reviews normalization and parameter-update approaches. TTN’s distinction is channel-wise sensitivity for post-training normalization weights, rather than layer-wise sensitivity for test-time model updates.

4 RELATED WORK

  • Unlike UDA and SFDA, TTA uses test data online rather than relying on a sufficiently large unlabeled target dataset for a particular domain.The related work identifies adapting normalization statistics and adapting model parameters as two key factors in recent TTA approaches.Source: §4 · p. 9
  • Prior normalization approaches use test-batch statistics, combine source and test statistics, or adapt them with moving averages and pseudo-batches.The paper names prediction-time BN as TBN and distinguishes TTN by optimizing interpolating weights using channel-wise domain-shift sensitivity.Source: §4 · p. 9
  • TENT, SWR, and CoTTA update model parameters using entropy or consistency objectives, while other approaches use contrastive learning.CoTTA also stochastically restores parts of the pretrained model.Source: §4 · p. 9
  • TTN uses channel-wise sensitivity to optimize α during post-training, unlike SWR’s layer-wise sensitivity for regularizing test-time model updates.Both approaches measure domain-shift sensitivity by comparing gradients.Source: §4 · p. 9

5 CONCLUSIONS

TTN combines CBN and TBN using a mixing weight derived from domain-shift sensitivity. It performs strongly in realistic settings and can complement optimization-based TTA, though its test-time mixing ratio remains fixed.

5 CONCLUSIONS

  • TTN mixes CBN and TBN using an optimized weight derived from each layer’s sensitivity to domain shift.The weight is optimized before deployment and remains fixed during test time.Source: §5 · p. 9
  • TTN significantly outperforms other normalization techniques across various realistic evaluation settings.Source: §5 · p. 9
  • TTN can complement optimization-based TTA methods, while a fixed test-time mixing ratio may leave room for improvement when domain gaps vary.The paper identifies adapting the ratio to distribution shift during test time as future work.Source: §5 · p. 9

A.1 IMPLEMENTATION DETAILS

The implementation optimizes TTN’s mixing weight using source training data, with datasets, augmentations, and training settings varying by task. Test-time adaptation settings follow the respective baseline methods, with scenario-specific adjustments.

A.1 IMPLEMENTATION DETAILS

  • CIFAR-10/100-C and ImageNet use augmented training data to optimize α on pretrained WRN-40-2 and ResNet-50 models, respectively.ImageNet uses 64,000 randomly sampled training instances per epoch.Source: §A.1 · p. 13
  • α optimization uses method-specific augmentations; ImageNet uses the same augmentation for obtaining the prior and optimizing α.For other datasets, Gaussian noise is excluded from α-optimization augmentation to avoid overlap with common-corruption types.Source: §A.1 · p. 13
  • The post-training procedure uses Adam for 30 epochs, with task-specific learning rates and batch sizes, and stops when the MSE loss saturates.For semantic segmentation, TTN is trained on Cityscapes with batch size 2; images are resized to height 800 while preserving aspect ratio.Source: §A.1 · p. 13
  • AdaptiveBN’s test-time interpolation weight is set using its suggested N values for each test batch size.The reported settings pair batch sizes from 200 to 1 with α values from 0.44 to 0.06.Source: §A.1 · p. 13
  • At test time, optimization-based methods follow TENT, SWR, and CoTTA defaults, with learning rates and update schedules adjusted for scenario and task.Segmentation uses test batch size 2 and learning rate 1e-6; all TTA methods use a constant learning rate.Source: §A.1 · p. 13

A.2 EVALUATION SCENARIO DETAILS

The appendix describes a class-imbalanced evaluation setup and provides pseudocode for obtaining TTN’s prior, optimizing α, and applying the resulting normalization at test time.

A.2 EVALUATION SCENARIO DETAILS

  • In the class-imbalanced setting, test examples are sorted by class before batching, so most batches contain one class and have biased statistics.The authors report higher error rates with larger batches than smaller ones in this setting, attributing this pattern to more extremely skewed statistics.Source: §A.2 · p. 13

A.3 PSEUDOCODE

  • The prior A is obtained by collecting gradients on clean and augmented inputs, computing gradient-distance scores, and defining A from those scores.The pseudocode makes BN-layer weight and bias parameters trainable when collecting gradients.Source: §A.3 · p. 14
  • Post-training initializes α with prior A, replaces BN layers with TTN layers, and updates α on source-data minibatches.At inference, the pretrained model uses the optimized α in TTN layers while processing target minibatches.Source: §A.3 · p. 14

B.1 ABLATION STUDIES

The ablations examine α granularity and regularization, then compare constant-weight and other TTN variants. They report that channel-wise α performs best, lower test-statistic weighting helps at small batch sizes, and online TTN improves on offline TTN across three scenarios.

B.1 ABLATION STUDIES

  • Channel-wise optimized α performs best among the compared granularity levels on CIFAR-10-C with WRN-40.The comparisons include channel-wise, layer-wise, and constant α.Source: §B.1 · p. 15, Table 8 · p. 15
  • The MSE regularizer keeps α nearer its prior, but excessive strength raises error rates while too little can overfit to batch size 200.The authors empirically set λ to 1; too-small regularization may reduce generalizability to smaller test batches.Source: §B.1 · p. 16, Table 9 · p. 16

B.2 MORE COMPARISONS ON CIFAR10-C

  • Lower constant α, which uses less test-batch statistics, performs better for small test batch sizes.The comparison spans constant α values from 0 (CBN) to 1 (TBN).Source: §B.2 · p. 16, Table 10 · p. 16

B.3 VARIANTS OF TTN FOR SMALL TEST BATCH SIZE

  • Online TTN further optimizes post-trained α with entropy minimization and an MSE penalty, and improves over offline TTN in single, continuously changing, and mixed scenarios.Source: §B.3 · p. 16, Table 11 · p. 17
  • Dynamic post-training batch sizes improve performance at test batch sizes 2 or 1 but slightly worsen it at 200 or 64; scaled TTN uses a scale value of 0.4.The dynamic setting samples training batch sizes from 4 to 200 for each iteration.Source: §B.3 · p. 16

B.4 STUDY ON AUGMENTATION TYPE

The augmentation-type ablations indicate that TTN is robust to the augmentations used in post-training, with the prior largely invariant to augmentation choice.

B.4 STUDY ON AUGMENTATION TYPE

  • The prior is largely invariant to augmentation choice, with row-wise standard deviations below 0.5 in Figure 5(a).The broader ablation reports row-wise standard deviations below 1 in most cases.Source: §B.4 · p. 17
  • Average error rates were 11.74% for the prior ablation and 11.67% for the alpha-optimization ablation, matching TTN at 11.67%.The diagonal elements in Figure 5(b), where augmentation and test corruption align, tend to have the lowest error rates.Source: §B.4 · p. 17

B.5 RESULTS ON IMAGENET-C

On ImageNet-C, TTN leads the normalization-based methods, while adding TTN improves optimization-based TTA methods.

B.5 RESULTS ON IMAGENET-C

  • TTN performs best among the compared normalization-based methods on ImageNet-C and improves TENT and SWR when combined with them.Table 13 reports error rates averaged over 15 corruption types at severity level 5 for each test batch size.Source: §B.5 · p. 17, Table 13 · p. 18
  • SWR updates the entire model and is sensitive to learning rate at small test batch sizes, requiring careful tuning.The reported setup uses smaller learning rates for SWR at batch sizes 16, 4, 2, and 1 than at 64.Source: §B.5 · pp. 17–18
  • The ImageNet-C comparison uses ResNet-50 and evaluates multiple test batch sizes on 15 corruption types at severity level 5.The reported metric is error rate.Source: §B.5 · p. 17, Table 13 · p. 18

B.6 ERROR RATES OF EACH CORRUPTION

This section reports per-corruption CIFAR-10-C error rates for TTN using WRN-40 and points to averaged corruption results elsewhere.

B.6 ERROR RATES OF EACH CORRUPTION

  • Table 14 reports TTN error rates for each CIFAR-10-C corruption using the WRN-40 backbone.The averaged results over corruptions are reported in Table 1.Source: §B.6 · p. 18, Table 14 · p. 18
  • The per-corruption results use CIFAR-10-C and WRN-40.The section identifies Table 14 as the source of these individual corruption results.Source: §B.6 · p. 18
  • The aggregate results over corruptions are reported in Table 1 rather than in the per-corruption table.Table 14 is titled as results of each corruption.Source: §B.6 · p. 18, Table 14 · p. 18

B.7 SEGMENTATION RESULTS

In semantic segmentation, TTN maintains more consistent performance across test batch sizes than TBN, with the largest reported accuracy advantage at batch size 2.

B.7 SEGMENTATION RESULTS

  • TTN outperforms TBN by 3.77% in average accuracy at B=2, 2.54% at B=4, and 1.99% at B=8.Table 15 reports mIoU on four unseen domains for test batch sizes 2, 4, and 8.Source: §B.7 · p. 18, Table 15 · p. 18
  • TTN’s performance difference across batch sizes 2, 4, and 8 is small, whereas TBN drops at smaller batch sizes.The authors assume TBN statistics at B=8 are sufficiently reliable because the backbone was trained with batch size 8.Source: §B.7 · p. 18, Table 15 · p. 18
  • Figure 5 encodes augmentation type by column and test corruption type by row, with error rates normalized within each row.It provides a comparison of corruption-alignment effects, not segmentation outcomes.Source: Figure 5 · p. 19

Related papers

Loading 2302.05155v2…