09. Test-Time Adaptation

• Primary References:

1. 1
   Li, Y., Wang, N., Shi, J., Hou, X., & Liu, J. (2018). Adaptive Batch Normalization for Practical Domain Adaptation (AdaBN). Pattern Recognition, 80, 109–117.
2. 2
   Wang, D., Shelhamer, E., Liu, S., Olshausen, B., & Darrell, T. (2021). Tent: Fully Test-Time Adaptation by Entropy Minimization. In International Conference on Learning Representations (ICLR 2021).
3. 3
   Niu, S., Wu, J., Zhang, Y., Wen, Z., Chen, Y., Zhao, P., & Tan, M. (2023). Towards Stable Test-Time Adaptation in Dynamic Wild World (SAR). In ICLR 2023.
4. 4
   Sun, Y., Wang, X., Liu, Z., Miller, J., Efros, A., & Hardt, M. (2020). Test-Time Training with Self-Supervision for Generalization under Distribution Shifts (TTT). In ICML 2020.

• Why I Wrote This Note: In our work on chest X-ray lung segmentation (TIG-UNet++, ICICT 2025) and tuberculosis screening (ICEEICT 2024), pixel intensity distributions shift whenever a model is evaluated on radiographs from a new hospital scanner. I wrote this note to understand when online normalization adaptation recovers performance and when unsupervised entropy minimization dangerously collapses into predicting a single class.

• Why Normalization Layers Hold the Key to Covariate Shift: First- and second-order channel statistics \((\mu_l, \sigma_l^2)\) in Batch Normalization encode domain style/lighting/contrast, while convolutional weights \(W_l\) encode semantic shape structure. Under covariate shift \(p_T(x) \neq p_S(x)\), source running statistics \((\hat{\mu}_S, \hat{\sigma}_S^2)\) become mismatched, shifting activations out of the linear regime of downstream nonlinearities.
• AdaBN (Gradient-Free): Simply replacing frozen source statistics \((\hat{\mu}_S, \hat{\sigma}_S^2)\) with test-batch statistics \((\mu_T, \sigma_T^2)\) reduces corruption error on ImageNet-C by \(10\%\text{–}15\%\) with zero backward passes.
• TENT (Gradient-Based Entropy Minimization): Wang et al. (2021) observed that higher test-time predictive Shannon entropy \(\mathcal{H}(\hat{y})\) strongly correlates with classification error under corruption. Updating only the affine scale and shift parameters \((\gamma_l, \beta_l)\) of normalization layers to minimize \(\mathcal{H}[p_\theta(y\mid x_T)]\) aligns target features back toward high-density decision regions.
• The 3 Deadly Traps of Naive TENT: TENT collapses toward a trivial constant prediction if (1) test batch size \(B < 8\), (2) test stream has strong label imbalance/temporal correlation, or (3) high-entropy OOD outliers generate massive noisy gradients (fixed by Reliable Entropy Minimization / SAR, Niu et al., 2023).

4.1 Step 1: Gradient-Free Statistics Recalibration (AdaBN)

In a standard Batch Normalization layer at channel \(c\), an activation \(z \in \mathbb{R}^{B \times C \times H \times W}\) is transformed via:

During standard model.eval(), PyTorch uses exponential moving averages \((\hat{\mu}_S, \hat{\sigma}_S^2)\) accumulated on the source training set. Under covariate shift, \(\mathbb{E}_{p_T}[z_c] \neq \hat{\mu}_{S,c}\). AdaBN computes batch statistics \((\mu_{T,c}, \sigma_{T,c}^2)\) directly on the incoming test batch (or updates a test-time EMA with momentum \(m \in [0.05, 0.1]\)):

4.2 Step 2: Test-Time Entropy Minimization (TENT)

While AdaBN centers and scales each channel independently, it cannot account for rotated cross-channel correlations under severe corruptions (e.g., Gaussian blur, fog, or low-dose X-ray scatter). TENT (Wang et al., ICLR 2021) freezes all convolutional/attention weights \(\theta_{\text{backbone}}\) (comprising \(>99\%\) of parameters) to prevent catastrophic forgetting, and optimizes only the channel-wise affine transformation parameters \(\boldsymbol{\Theta}_{\text{aff}} = \{\gamma_l, \beta_l\}_{l=1}^{L}\) by minimizing the Shannon entropy of the model's predictions on test batch \(\mathcal{B}_T\):

flowchart LR
    XT["Unlabeled Test Batch X_T ~ p_T(x)"] --> BN["Normalization Layers<br/>Compute μ_T, σ²_T + Affine (γ, β)"]
    BN --> CONV["Frozen Backbone Weights W_l<br/>(Requires_Grad = False)"]
    CONV --> LOG["Predicted Probabilities p̂(y | X_T)"]
    LOG --> ENT["Shannon Entropy H(p̂)<br/>Filter High-Entropy Outliers H(x) < E₀"]
    ENT -- "∇_{γ,β} H(p̂) (1 Step Update)" --> BN

4.3 Preventing Collapse: Entropy Filtering & Marginal Diversity (EATA / SAR)

Why can minimizing \(\mathcal{H}(\hat{p})\) collapse? Notice that a network outputting \(\hat{p}(y=1\mid x) = 1.0\) for every input \(x\) achieves the global minimum of Shannon entropy \(\mathcal{H}(\hat{p}) = 0\)! To prevent this degenerate solution without ground-truth labels:

1. 1
   Reliable Sample Filtering (EATA / SAR, Niu et al., 2022/2023): High-entropy samples (\(\mathcal{H}(\hat{p}(x)) > E_0\), where \(E_0 = 0.4 \times \ln C\)) are highly uncertain and produce conflicting, high-norm gradients that destabilize \(\gamma, \beta\). We mask out high-entropy samples from the backward pass:

1. 1
   Information Maximization (IM / Marginal Entropy Regularization, Liang et al., SHOT 2020): Subtract the entropy of the batch-mean prediction \(\bar{p} = \frac{1}{|\mathcal{B}_T|}\sum_{x \in \mathcal{B}_T} \hat{p}(y\mid x)\) to force predictions to remain diverse across classes:

4.4 PyTorch Implementation of Safeguarded TENT with Reliable Sample Filtering

import math
import torch
import torch.nn as nn

class SafeguardedTentAdapter:
    # Online Test-Time Adaptation updating only Norm affine parameters (gamma, beta) with entropy filtering
    def __init__(self, model: nn.Module, num_classes: int, lr: float = 1e-3, e_margin_ratio: float = 0.4):
        self.model = model
        self.e_margin = e_margin_ratio * math.log(max(num_classes, 2))
        self.affine_params = self._configure_norm_layers()
        self.optimizer = torch.optim.AdamW(self.affine_params, lr=lr, weight_decay=0.0)

    def _configure_norm_layers(self) -> list[nn.Parameter]:
        self.model.eval()
        self.model.requires_grad_(False)
        params = []
        for module in self.model.modules():
            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
                # For BatchNorm, use batch statistics at test time; LayerNorm/GroupNorm are batch-size agnostic
                if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)):
                    module.train()
                    module.track_running_stats = False
                for p in module.parameters(recurse=False):
                    p.requires_grad_(True)
                    params.append(p)
        return params

    def forward_and_adapt(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
        logits = self.model(x)
        probs = torch.softmax(logits, dim=-1)
        entropies = -(probs * probs.clamp_min(1e-12).log()).sum(dim=-1)  # (B,)

        # Filter out unreliable high-entropy OOD outliers (SAR / EATA criterion)
        reliable_mask = entropies < self.e_margin
        reliable_ratio = float(reliable_mask.float().mean().item())

        if reliable_mask.any():
            loss = entropies[reliable_mask].mean()
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            # Re-evaluate forward pass after 1-step adaptation
            with torch.no_grad():
                logits = self.model(x)

        return logits.detach(), {
            "mean_entropy": float(entropies.mean().item()),
            "reliable_ratio": reliable_ratio,
        }

1. 1
   TIG-UNet++ Chest X-Ray Lung Segmentation (ICICT 2025): Notice why our Thresholded Input Guided (TIG) architectural design is complementary to Test-Time Adaptation: by deriving a structural thresholded prior directly from each input radiograph to guide the UNet++ encoder-decoder skip pathways, TIG-UNet++ performs an explicit, gradient-free input-adaptive conditioning at batch size \(B=1\), avoiding the small-batch instability of BatchNorm TENT in clinical workflows where radiographs are processed one patient at a time!
2. 2
   Combining TTA with Selective Abstention: My implementation of SafeguardedTentAdapter above reveals a powerful synergy: the exact samples rejected by the reliable_mask (\(\mathcal{H}(\hat{p}(x)) \ge E_0\)) during TTA adaptation are the prime candidates for Selective Abstention in Risk-Constrained AI Decision-Making.

• BatchNorm vs. LayerNorm/GroupNorm in TTA: If you know your model will require \(B=1\) online Test-Time Adaptation at the edge, do not use BatchNorm in the architecture—use GroupNorm or LayerNorm (as in ConvNeXt and Vision Transformers). As Niu et al. (SAR, ICLR 2023) proved, BatchNorm's batch-level mean/variance dependency is the #1 culprit behind single-sample TTA collapse.
• Episodic Resetting (Continual TTA): In long-running streams (e.g., 24/7 factory defect cameras or grid monitors), adapting \(\gamma, \beta\) over \(100,000\) steps without anchor regularization causes gradual parameter drift. Stochastically restoring \(1\%\) of weights back toward source initialization \(\theta_0\) (CoTTA, Wang et al., CVPR 2022) prevents long-horizon forgetting.

• Implement SafeguardedTentAdapter with norm-layer isolation and reliable entropy threshold filtering (\(E_0 = 0.4 \ln C\)).Verified
• Benchmark SafeguardedTentAdapter with episodic exponential moving average anchor restoration on corrupted chest X-ray splits.Pending
• Compare pixel-wise entropy minimization on segmentation logits in TIG-UNet++ against AdaBN statistics recalibration.Pending

"Minimizing entropy on confident-but-slightly-shifted samples sharpens the decision boundary; minimizing entropy on completely unknown outliers teaches the network to hallucinate with 100% confidence."

Reliable & Diversity-Regularized TTA Objective:

• Connected Research Notes:
• Distribution Shift: A Practical Introduction — Explains why TTA works under covariate shift (\(p(x)\) changes, \(p(y\mid x)\) invariant) and why it fails under concept drift.
• Understanding Predictive Uncertainty — Provides the Shannon entropy objective \(\mathcal{H}(\hat{p})\) minimized by TENT.
• Risk-Constrained AI Decision-Making — Routing samples that fail the TTA reliability filter (\(\mathcal{H}(\hat{p}) \ge E_0\)) to human review.