Contrastive Test-Time Adaptation

In plain words

Imagine a classifier trained on one collection of photos that now has to work on a different collection, with no labels and no access to the original training photos. That situation is test-time adaptation. AdaContrast lets the model keep learning while it sees the new photos.

It guesses a label for each new image and then checks those guesses in two ways. First, it looks at recent images that look similar in the model's feature space and averages their guesses, which smooths out individual mistakes. Second, it learns from two differently cropped and coloured copies of the same image, pulling their features together, and it pushes apart only images whose guessed labels differ.

The authors report that this keeps the guesses reliable and works across many settings without much tuning. The numbers are their own and come from two image benchmarks.

TL;DR

AdaContrast adapts a source-trained ResNet to an unlabelled target set, with no source data at adaptation time. It trains self-training and a contrastive loss jointly, and pseudo labels are refined per mini-batch by nearest-neighbour voting over a memory queue. Author-reported gains over SHOT: 86.8 vs 83.0 per-class average on VisDA-C and 67.8 vs 67.1 on the DomainNet-126 seven-shift average. The official code has no LICENSE file, and it omits or changes several details the paper states, so read Paper vs code and Pitfalls before implementing.

Problem

Test-time adaptation (TTA) adapts a model trained on a labelled source domain to an unlabelled target domain, and source data is not available during adaptation §1. The paper addresses closed-set image classification, where source and target share one label set §3. Two problems follow: learning target features without labels, and building a target classifier when only the source classifier is available §1.

The authors name three weaknesses of existing methods. Entropy minimisation ignores relations between samples and disturbs calibration on the target §2 §4.3. Pseudo labels are often refreshed only once per epoch, so they lag behind the model §2 §3.1. Prior work keeps a global bank that is refreshed every epoch or every few epochs, which the authors argue suits streams of data poorly, since a stream cannot be revisited §3.1 §4.2.

Method

Source model. A feature extractor f_s and classifier h_s are trained on labelled source data with cross-entropy and label smoothing α = 0.1 §3. The bottleneck and weight-norm details are in §4.1.

Target model. The target model g_t = h_t(f_t) starts from the source weights θ_s §3. A momentum copy g_t' also starts from θ_s and follows the target model as an exponential moving average at each mini-batch step §3.1 eq. 2.

Per mini-batch, the method uses one weak view and two strong views of each target image fig. 2 §3.

1. Weak-view feature and refined pseudo label §3.1 eq. 3 eq. 4. The weak view is passed through the target model to get feature w. Its cosine distances to the N nearest entries of queue Q_w are computed. The probability vectors of those N entries are averaged, and the argmax of the average is the refined pseudo label ŷ. The refinement runs at every mini-batch step, not once per epoch.
2. Queue update with the momentum model §3.1 eq. 1. The weak view is also passed through the momentum model to produce a feature w' and probability p'. These are enqueued into Q_w, and the oldest entries are dropped. Q_w has M entries §3.1.
3. Contrastive term §3.2 eq. 5 eq. 6. Two strong views give a query q from the target model and a key k from the momentum model. The key is enqueued into a second queue Q_s of length P, together with its pseudo label. The InfoNCE loss uses k as the positive. Negatives are all entries of Q_s whose pseudo label differs from the current ŷ. No projection head is used fig. 2 §4.1.
4. Weak-strong cross-entropy §3.3 eq. 7. The pseudo label from the weak view supervises the prediction on the strong query view. No confidence threshold is applied.
5. Diversity term §3.3 eq. 8 eq. 9. The mean prediction over the mini-batch is p̄. The loss is the sum of p̄ log p̄, which is minimised when p̄ is spread over classes.
6. Total loss §3.3 eq. 10. L = γ1 L_ce + γ2 L_ctr + γ3 L_div, with γ1 = γ2 = γ3 = 1 and no tuning in any experiment §3.3.

Supplementary Algorithm 1 alg. 1 lists the same steps. The pseudocode below follows the paper's equations, not the official code. Where the code differs, see Paper vs code.

# AdaContrast, one target mini-batch  [alg. 1] [§3]
# Input:  images (B, 3, H, W) from the unlabelled target set
# Output: updated target model, updated queues Q_w and Q_s

def adapt_step(images, target, momentum, queue_w, queue_s):
    # Draw one weak view and two strong views of each image  [§3.1] [§3.2]
    view_weak = weak_augment(images)
    view_query = strong_augment(images)
    view_key = strong_augment(images)

    # Weak view through the target model gives the lookup feature  [§3.1]
    feat_weak, logits_weak = target(view_weak)    # (B, d), (B, C)

    # Average the probabilities of the N nearest queue entries  [eq. 3]
    nearest = top_n_cosine(feat_weak, queue_w.feats, N)     # (B, N)
    voted = mean(queue_w.probs[nearest], dim=1)             # (B, C)

    # Refined pseudo label is the argmax of the voted probabilities  [eq. 4]
    pseudo = argmax(voted, dim=1)                           # (B,)

    # Query from the target model, key from the momentum model  [§3.2]
    feat_query, logits_query = target(view_query)    # (B, d), (B, C)
    # The key gets no gradient
    key = normalize(momentum.feature(view_key))      # (B, d)

    # Contrastive loss with same-pseudo-label negatives removed  [eq. 5-6]
    loss_ctr = info_nce(normalize(feat_query), key, queue_s, pseudo, tau)

    # Cross-entropy of the strong view against the weak-view label  [eq. 7]
    loss_ce = cross_entropy(logits_query, pseudo)

    # Diversity term on the mini-batch mean prediction  [eq. 8-9]
    mean_prob = mean(softmax(logits_query), dim=0)          # (C,)
    loss_div = sum(mean_prob * log(mean_prob))

    # Weighted sum of the three terms; all weights are 1.0  [eq. 10]
    loss = g1 * loss_ce + g2 * loss_ctr + g3 * loss_div

    # SGD step on the target model  [§3.2]
    sgd_step(target, loss)

    # Momentum model tracks the target model slowly  [eq. 2]
    momentum.params = m * momentum.params + (1 - m) * target.params

    # Enqueue weak-view feature and probability from the momentum model  [eq. 1]
    feat_mom, probs_mom = momentum(view_weak)
    queue_w.push(feat_mom, probs_mom)

    # Enqueue the key with its pseudo label for the contrastive term  [§3.2]
    queue_s.push(key, pseudo)

# InfoNCE over a queue with same-class negatives removed  [eq. 5-6]
# query, key: (B, d) unit vectors; queue_keys: (P, d); queue_labels: (P,)

def info_nce(query, key, queue_keys, queue_labels, pseudo, tau):
    # Positive logit is column 0 and is never masked  [eq. 6]
    positive = sum(query * key, dim=1, keepdim=True)        # (B, 1)

    # Negative logits against every queued key  [eq. 5]
    negatives = query @ queue_keys.T                        # (B, P)

    # Drop negatives that share the current pseudo label  [eq. 6]
    same_class = pseudo[:, None] == queue_labels[None, :]   # (B, P)
    negatives = where(same_class, -inf, negatives)

    # Temperature scaling, then softmax over the positive and negatives  [eq. 5]
    logits = concat([positive, negatives], dim=1) / tau     # (B, 1 + P)
    return cross_entropy(logits, zeros(B))

Initialisation of the queues. Q_w is filled with features and probabilities of M randomly chosen target samples §3.1. Q_s is filled with random features and no labels are stated for it alg. 1. Both queues are then updated on the fly.

Augmentations. The weak view resizes to 256 × 256, takes a random 224 crop and applies a random horizontal flip. The strong view uses the MoCo v2 recipe: random resized crop with scale 0.2 to 1.0, colour jitter (0.4, 0.4, 0.4, 0.1) with probability 0.8, grayscale with probability 0.2, Gaussian blur with sigma in [0.1, 2.0] with probability 0.5, and a random horizontal flip [supp. §2.1]. Both use ImageNet mean and standard deviation.

Hyperparameters

Values marked "paper" are stated in the paper. "Not stated" means the paper gives no value. Code values are shown where the official code sets them.

Results

Author-reported throughout. Setup: closed-set adaptation on VisDA-C (train to validation, 12 classes) and DomainNet-126 (7 domain shifts, 126 classes). Baselines are UDA methods that use source data and source-free methods §4.1. Table numbers follow LaTeX source order tab. 1 to tab. 4.

VisDA-C train to validation, per-class accuracy (%) tab. 1. The table caption says all methods use ResNet-101 except the on-target rows, which use a ResNet-18 student tab. 1.

Overall accuracy on VisDA-C is 84.5% for AdaContrast, against 78.3% for SHOT §4.2 tab. 3. The overall-accuracy figure appears only in the text and in tab. 3.

DomainNet-126, accuracy (%) per domain shift, ResNet-50 tab. 2.

Learning-rate sensitivity tab. 3. Scale applied to the main learning rate.

Ablation tab. 4. Row 1 is the simplest pseudo labelling: one prediction pass per epoch, used as labels for that epoch.

Other reported results §4.2 §4.3:

• Knowledge distillation to a ResNet-18 student, initialised from a MoCo v2 pretrained model: 87.2% per-class average on VisDA-C tab. 1. The distillation procedure is not in the released code.
• Online adaptation, where each image is seen once and refinement is switched on after X samples: 78.7% VisDA-C per-class average and 62.6% DomainNet-126 average tab. 1 tab. 2.
• Calibration on VisDA-C validation after adaptation, against SHOT: ECE 0.65% vs 2.97% and MCE 8.20% vs 39.16% §4.3 fig. 3.
• Insensitivity to the queue size M and to the neighbour count N across the range shown in the paper. The strongest VisDA-C setting uses M = 55388 and N = 11 §4.3 fig. 4.

Paper vs code

Repository: DianCh/AdaContrast at c3c8b880131f2658d6fd0d5ed14d71f326174d57, owned by the first author. The paper's project page links it. The repository has no LICENSE file, so no code is quoted here. The code is cited by path and line only.

Searched: the paper and its arXiv source, the project page, and a GitHub repository search for "adacontrast". The search returned two repositories. The other, y9z9/AdaContrast, is not linked from the paper or the project page, so it was not used.

Differences:

1. Neighbour count. The paper uses N = 11 for its strongest VisDA-C result §4.3. The default in the code is 10 target.yaml. The released scripts do not override it.
2. Queue initialisation. The paper initialises Q_w from weak-augmented samples through the momentum model §3.1 eq. 1. The code takes features from the target model in eval mode on the deterministic test transform (resize 256, centre crop 224), not a weak augmentation target.py target.py utils.py. Both start from identical weights, so the difference is in the augmentation and BatchNorm mode, not the weights.
3. Labels of the initial dataset. The initial pseudo labels are stored in the training dataset target.py, but the training loop discards them and uses the per-batch refined labels target.py target.py. The paper describes only per-batch refinement, so this matches it in effect.
4. Queue used for the contrastive negatives. The code enqueues the current batch before computing the contrastive mask, while the logits were computed against the queue before the enqueue. See Pitfalls.
5. Shuffle BatchNorm. The key encoder shuffles the batch across GPUs before the momentum model and unshuffles afterwards builder.py builder.py. The paper does not mention this step.
6. Momentum update timing. The code applies the momentum update at the start of the key forward, before that step's SGD update builder.py. The paper's Algorithm 1 applies it after the update alg. 1.
7. Momentum buffers. The momentum update copies only parameters, not BatchNorm running statistics builder.py. The momentum model runs in train mode, so its BatchNorm uses batch statistics. The paper says nothing about this.
8. Cosine schedule. The paper gives η = η0 · 0.5 (cos(a·π/2) + 1) and says the rate decays to zero §4.1. The code uses 0.5 (1 + cos(π · t/T)), which reaches zero at the end, where t is the global step and T the total number of steps utils.py target.py. With a in [0, 1], the paper's formula stops at half the initial rate, so the paper's a must mean 2t/T for the stated behaviour.
9. Instance mask uses updated labels. Covered under Pitfalls.
10. Hyperparameters not in the paper: Q_s length 16384, momentum 0.999, temperature 0.07, Nesterov momentum moco.yaml sgd.yaml.
11. Modes not in the code. The released code has no online mode, no distillation step and no separate learning-rate sweep config. The sweep can be reproduced by overriding optim.lr on the command line sgd.yaml.
12. Source seeds. Scripts train three source seeds (2020, 2021, 2022) and run each target adaptation for those seeds train_VISDA-C_target.sh. The paper does not say whether tables average over seeds. Not stated.

Items that match the paper: the loss equations and weights target.py, the class-aware negative mask target.py, the weak-strong pairing target.py, the diversity term on strong logits target.py, the augmentations utils.py, the source label smoothing source.py, the source epochs, and the target epochs.

Pitfalls

• Queue order in the contrastive mask. The code enqueues the current batch's keys and pseudo labels target.py before computing the mask target.py. The logits were computed earlier, against the old keys builder.py. For the B slots overwritten in this step, the mask therefore uses new labels against logits from the old keys. The paper's equations imply the negatives are the queue before the update. Implement the order the paper states, or check the effect before copying the code's order.
• Random initial labels. The label buffer of Q_s is filled with random integers from the class range builder.py. Those slots can exclude real negatives until overwritten. The paper says only that Q_s is randomly initialised alg. 1.
• Distributed-only code. The entry point sets distributed mode when multiprocessing_distributed is on, which is the default root.yaml main_adacontrast.py. The default backend is NCCL root.yaml. The key encoder and queue updates call collective operations unconditionally builder.py utils.py. Expect an NCCL process group even for one GPU.
• Full sort in neighbour search. The code sorts every row of the full distance matrix to get the top N target.py. With M = 55388 this is wasteful; a top-N selection returns the same neighbours, up to tie order. Cost per step grows with M, so the queue size is the main compute knob.
• Diversity term on a mini-batch. The paper defines the term with an expectation over the target set §3.3 eq. 8-9. The code computes the mean over the current mini-batch of 128 target.py. The batch estimate is noisier and is what is actually optimised.
• Stated cosine formula. The formula in the paper, read with a in [0, 1], does not reach zero (see Paper vs code item 8). The code's version reaches zero.
• Text and table disagreements in the paper:
  • The text calls 78.7% "overall accuracy" and reports +0.5 over SHOT §4.2. Table 1 lists 78.7 under the per-class average. SHOT's overall figure is 78.3 tab. 3, so the difference is 0.4 and the metric is mislabelled.
  • The ablation text reports +0.9% on VisDA-C at 10× learning rate for the regularisation row §4.4. Table 4 gives 84.3 → 86.6, which is +2.3.
  • The ablation text gives the diverged accuracy as 10.0%. Table 4 gives 9.9.
  • The ablation text quotes VisDA-C numbers for the same-class exclusion comparison (67.7, 83.6, 81.5). These appear in the text only, not in the live table, so they cannot be checked from the paper's tables.
• Knowledge distillation and online mode are not reproducible from the repository, and the paper gives no distillation loss or schedule §4.2.
• Table and figure numbers follow LaTeX source order and were not checked against a compiled PDF. Equation numbers were counted in order from the source.
• Queue memory. The VisDA-C queue at M = 55388 holds 256-d float features and 12-class probabilities, about 57 MB of features in fp32. Every mini-batch scores its weak features against the whole queue, so the per-step cost grows with M.
• Pseudo labels are used for every sample. There is no confidence threshold §3.3, so wrong labels reach the loss. The diversity term is the paper's only guard against collapse.

Lineage

• Builds on MoCo for the contrastive setup, with a source-initialised encoder, a memory queue and an exclusion rule for same-label negatives §3.2 builder.py. The MoCo arXiv identifier is 1911.05722 as cited in the code builder.py.
• Builds on SHOT for the source model recipe, the bottleneck and the weight-norm classifier §3 §4.1.
• Uses weak-strong consistency from FixMatch and drops its confidence threshold §3.3.
• Uses the soft k-nearest-neighbour voting idea from the cited soft-kNN work §3.1.
• Compared against: MA, BAIT, TENT, SHOT and On-target for TTA, and MCC, CAN and related methods for UDA §4.1.
• Related to test-time training, which adapts with a rotation-prediction task §2. AdaContrast does not alter source training, and the paper contrasts itself with approaches that do §1.
• Superseded by: none known from the content in this wiki.