Progressive Conditioned Scale-Shift Recalibration

• The paper demonstrates that integrating DSN and FGN modules enables online recalibration of QKV features, boosting accuracy on domain shift benchmarks up to 67.5%.
• The method applies a local affine transformation with per-channel scale and shift factors to efficiently adapt transformer attention without retraining the backbone.
• Empirical evaluations on ImageNet-C, ImageNet-R, and VisDA-2021 show PCSR achieving robust performance with less than 3% parameter overhead and significant improvements over prior methods.

Progressive Conditioned Scale-Shift Recalibration (PCSR) is a method designed to improve online test-time adaptation of transformer network models, particularly in the presence of significant domain shifts between source and target data distributions. PCSR addresses the challenge that the Query, Key, and Value (QKV) features of self-attention modules can change drastically under distribution shift, resulting in degraded model performance at test time. The approach introduces lightweight, layer-wise modules to recalibrate these features dynamically via per-channel scale and shift factors, determined by ongoing inference data, and adapts these modules online without updating the underlying transformer weights (Tang et al., 14 Dec 2025).

1. Mathematical Formulation

PCSR applies a simple local affine (scale-shift) transform to the Q, K, and V matrices in each transformer layer $\ell$. For $N$ tokens and hidden dimension $d$, the pre-attention features $Q, K, V \in \mathbb{R}^{N \times d}$ are recalibrated as: $$Q' = \gamma \odot Q + \beta,\quad K' = \gamma \odot K + \beta,\quad V' = \gamma \odot V + \beta$$ where $\odot$ denotes element-wise multiplication, and $\gamma,\beta \in \mathbb{R}^d$ are scale and shift parameters predicted for each layer. In practice, these scale and shift factors are shared across $Q, K, V$ to minimize parameter count and runtime.

Self-attention then proceeds as usual, but with recalibrated features: $$\text{Attention}(Q', K', V') = \text{Softmax}\Big( \frac{Q' {K'}^\top}{\sqrt{d}} \Big) V'$$ This mechanism allows for adaptation to the new domain by correcting the QKV statistics on-the-fly as distributional shifts occur.

2. Domain Separation Network (DSN)

To estimate the current domain shift at each layer, PCSR introduces a Domain Separation Network $\Phi^\ell$. The DSN distills a domain token $D^\ell \in \mathbb{R}^d$ from the set of $N$ patch tokens $\{P_i^\ell\}_{i=1}^N$: $$f_i^\ell = \Phi^\ell(P_i^\ell),\quad D^\ell = \frac{1}{N}\sum_{i=1}^N f_i^\ell$$ $\Phi^\ell$ is implemented as a linear projection or a small MLP.

DSN's objective is to ensure $D^\ell$ represents the "Fermat-Weber point" of $\{f_i^\ell\}$ by maximizing pairwise cosine similarity among $f_i^\ell$. The similarity matrix is defined as: $$M^{\ell}_{jk} = \frac{f_j^\ell \cdot f_k^\ell}{\|f_j^\ell\| \|f_k^\ell\|}$$ and the similarity loss is: $$L_s(\theta; x) = -\frac{1}{L}\sum_{\ell=1}^L \frac{1}{N^2}\sum_{j,k=1}^N M^\ell_{jk}$$ Higher similarity ensures feature clustering and a representative domain shift token.

3. Factor Generator Network (FGN)

The Factor Generator Network $\Psi^\ell$ predicts the concatenated vector $[\gamma^\ell; \beta^\ell] \in \mathbb{R}^{2d}$, conditioned on both the layer's Domain token $D^\ell$ and Class token $C^\ell$: $$[\gamma^\ell; \beta^\ell] = \Psi^\ell\big([D^\ell ; C^\ell]\big) = W^\ell [D^\ell ; C^\ell] + b^\ell$$ Here, $\Psi^\ell$ is a single fully connected layer with no additional nonlinearity or normalization, mapping $\mathbb{R}^{2d} \to \mathbb{R}^{2d}$.

This network enables the scale-shift recalibration to incorporate both current domain information and class context.

4. Progressive Domain Shift Separation and Adaptation

The overall adaptation process proceeds progressively at each transformer layer:

• The backbone transformer (e.g., ViT) is initialized with source weights $\theta_s$. Only DSN modules ($\Phi^\ell$) and FGN modules ($\Psi^\ell$) are adapted online.
• For each incoming test batch $B_t = \{x_i\}_{i=1}^B$, patch tokens $P_i^\ell$ and class token $C^\ell$ are extracted at each layer.
• DSN computes the current domain token $D^\ell$, and FGN predicts $[\gamma^\ell; \beta^\ell]$.
• QKV features are recalibrated, and attention is computed as described above.
• The process repeats for all $L$ layers, yielding final predictions $\hat{y}$.

Adaptation is governed by a combined loss: $$L(\theta; x) = L_e(\theta; x) + \lambda L_s(\theta; x)$$ where $$L_e(\theta; x) = \mathbb{1}[\text{Ent}(\hat{y}) < E_0] \cdot \text{Ent}(\hat{y})$$, with $\text{Ent}(\hat{y})$ denoting the output entropy, and $\lambda$ balancing the two terms depending on batch statistics. Updates are applied via SGD only to DSN and FGN; the backbone weights remain unchanged.

5. Architectural and Implementation Details

The DSN module $\Phi^\ell$ uses a linear d-to-d layer ($d^2 + d$ parameters per layer), while FGN $\Psi^\ell$ uses a $(2d)$ fully connected layer with $4d^2 + 2d$ parameters. For a 12-layer ViT-B/16 with $d=768$, the total parameter overhead is approximately 1.7M (<3% of the ViT).

Key hyperparameters include:

• Learning rates: $\eta_{\Phi} = 0.2$, $\eta_{\Psi} = 5 \times 10^{-4}$
• Batch size: 64
• Optimizer: SGD without momentum
• Three random seeds for repeatability

Pseudo-code for a single batch:

• Forward pass through layers, compute domain and class tokens, adapt QKV via affine recalibration
• Compute total loss, back-propagate to update DSN and FGN
• Repeat for subsequent batches; predictions at each stage are produced by the recalibrated model

Ablation studies confirm effective design choices, including:

• Conditioning on both class and domain tokens yields the best results (67.5% on ImageNet-C Level 5), outperforming single-token conditions
• Shared scaling across QKV achieves equivalent accuracy to independent scaling, with lower computational overhead

6. Empirical Evaluation

PCSR was evaluated on standard domain shift benchmarks:

• ImageNet-C (corruptions, levels 1–5; main results and breakdown on Level 5)
• ImageNet-R, ImageNet-A, VisDA-2021

Primary backbone: ViT-B/16, with scalability demonstrated on ViT-L/16.

Performance summary on ImageNet-C Level 5:

Method	Avg. Acc (%)
Source	51.0
TENT	62.8
SAR	63.6
PCSR	67.5

PCSR demonstrates an absolute improvement of 3.9% over the previous best (SAR/DePT-G) and 16.5% over the source model. Cross-dataset results show consistent improvements:

• ImageNet-R: 66.5% (PCSR) vs. 62.0% (SAR)
• ImageNet-A: 52.1% (PCSR) vs. 45.3% (SAR)
• VisDA-2021: 64.8% (PCSR) vs. 60.1% (TENT/SAR)
• ViT-L/16: 70.4% (PCSR) vs. 64.0% (SAR)

Sensitivity analysis confirms the robustness of hyperparameter choices and conditioning strategies.

7. Context and Significance

PCSR introduces a distinct paradigm for online test-time adaptation in transformer models by inserting two computationally lightweight modules per layer: a Domain Separation Network (DSN) and a Factor Generator Network (FGN). The method interprets domain adaptation as a progressive, layerwise domain shift separation process and leverages real-time adaptation without modifying the original transformer backbone parameters. The approach is notable for its parameter efficiency (<3% overhead on ViT-B/16), real-time feasibility, and substantial empirical gains over prior state-of-the-art approaches in challenging domain-shifted test scenarios (Tang et al., 14 Dec 2025).