Parameter-free Online Test-time Adaptation

In plain words

A trained image classifier often meets new conditions after training: other weather, other camera settings, other kinds of objects. Test-time adaptation means adjusting the model while it runs on this new, unlabelled data. Most earlier methods change the network's internal weights. In this paper's experiments those methods were very sensitive to their settings, and some made the model much worse when the settings did not match the situation.

The paper takes a more cautious route. It leaves the network untouched and only corrects the class probabilities the network outputs. The idea is that images which look alike should probably get the same label. For a batch of incoming images, each image's probabilities are nudged toward those of its look-alikes in the same batch, while staying close to what the original model said. The adjustment is computed with a short iterative procedure.

The authors report higher average accuracy than the unadapted model and the other methods across their test scenarios, with lower run time and memory use. They also report that it does not noticeably help in standard i.i.d., class-balanced settings.

TL;DR

• Online, unsupervised test-time adaptation with a frozen source classifier. Each batch of N target samples is handled on its own, and the output is a corrected probability matrix for that batch §5.
• The objective (eq. 3) is a KL term to the source predictions minus a Laplacian term over feature affinities. It is minimised by a concave-convex procedure with a closed-form update (eq. 5) §5 app. B.
• Author-reported: mean accuracy over 7 test scenarios rises from 51.0 to 55.0 (three RN-50 training recipes) and from 56.9 to 60.7 (five backbones) tab. 1. No gain on i.i.d. likelihood-shift scenarios §7.
• The paper's stated affinity is kNN with k in {1, 3, 5} app. C. The public code defaults to an RBF affinity instead (see Paper vs code).
• Use it when the test conditions are unknown and per-scenario tuning is not possible. Do not expect gains on i.i.d., class-balanced test data §8.

Problem

The setting is fully test-time adaptation: a source model trained on labelled data that cannot be accessed, unlabelled target data arriving as a possibly non-i.i.d. stream, and adaptation that must be unsupervised and independent of the model and training procedure §1.

Existing methods that adapt parameters (TENT, SHOT-IM, PseudoLabel, AdaBN) were shown to need scenario-specific hyperparameters. When tuned on one scenario, TENT can degrade accuracy by up to an absolute 66 points relative to the non-adapted model in other scenarios §4 fig. 3. Entropy minimisation on a low-diversity batch stream can also degrade the model silently, with no signal visible without labels §4 fig. 1 fig. 2.

The paper assumes covariate shift, meaning the label-given-input distribution is unchanged while the input distribution moves §3. It studies two shift types: prior shift (class frequencies change) and likelihood shift (input appearance given class changes) §3. Target classes may be superclasses of the source ImageNet classes, mapped through a hierarchy §3 app. A.

Method

Setting: a batch of N target samples with features phi(x_i) in R^d and source probabilities q_i over K classes, where q_ik = q(k | x_i) §5. The source network and its features are fixed. The output for each sample is an assignment vector z_i on the probability simplex §5.

The objective combines a KL term that keeps assignments near the source predictions with a Laplacian term that rewards agreement between neighbouring samples:

\mathcal{L}^{\text{LAME}}(\tilde{\mathbf{Z}}) = \sum_i \mathrm{KL}(\tilde{\mathbf{z}}_i \| \mathbf{q}_i) - \sum_{i,j} w_{ij}\, \tilde{\mathbf{z}}_i^T \tilde{\mathbf{z}}_j eq. 3

with w_ij = w(phi(x_i), phi(x_j)) an affinity that is large for nearby feature vectors §5. The KL form comes from maximising a regularised data likelihood over assignments under the simplex constraint §5 eq. 1 eq. 2.

Optimisation: the Laplacian term is concave when W is positive semi-definite and the KL term is convex. The concave-convex procedure replaces the concave part by its linear approximation at the current iterate, giving a tight upper bound §5 eq. 4. Solving the KKT conditions of that bound gives the update below, iterated until convergence; the derivation is in app. B.

\tilde{z}_{ik}^{(n+1)} = \frac{q_{\theta}(k|\mathbf{x}_i)\, \exp\!\big(\sum_j w_{ij}\, \tilde{z}_{jk}^{(n)}\big)}{\sum_{k'} q_{\theta}(k'|\mathbf{x}_i)\, \exp\!\big(\sum_j w_{ij}\, \tilde{z}_{jk'}^{(n)}\big)} eq. 5

Affinity choices app. C:

• kNN: w_ij = 1 if phi(x_i) is among the k nearest neighbours of phi(x_j), else 0. The paper says it follows a choice from Ziko et al. and searches k over {1, 3, 5}.
• RBF: w_ij = exp(-||phi(x_i) - phi(x_j)||^2 / (2 sigma^2)), with sigma the average distance of each point to its k-th neighbour.
• Linear: w_ij = phi(x_i)^T phi(x_j).

The paper reports that the three kernels perform roughly similarly app. C fig. 8.

# LAME on one batch  [§5] [eq. 3] [eq. 5] [app. C]
# Input:  source probabilities q (N, K); features phi (N, d)
# Output: corrected assignments z (N, K)

def lame_batch(q, phi, max_iters):
    # Affinity between all pairs of samples in the batch  [app. C]
    affinity = build_affinity(phi)              # (N, N)
    # Starting point; not stated in [§5], see Paper vs code
    z = q
    for step in range(max_iters):
        # Affinity-weighted sum of current neighbour assignments  [eq. 5]
        neighbour_sum = affinity @ z            # (N, K)
        # Source prior times the exponential of that sum  [eq. 5]
        scores = q * exp(neighbour_sum)         # (N, K)
        # Normalise each row so it sums to one  [eq. 5]
        z = scores / scores.sum(axis=1, keepdims=True)
        # Stop when converged; the rule is not stated in [§5]
        if has_converged(z):
            break
    return z

Note that the update does not change the source network. Each batch starts from its own source predictions, so nothing is carried across batches §5 §4.

Hyperparameters

Results

Author-reported. Setup: seven test scenarios (four i.i.d., three non-i.i.d., see Hyperparameters), each averaged over 10 random runs §6. Source models are ImageNet-trained. Batch size 16 for the tables fig. 5.

Training recipes, RN-50 (accuracy, averaged over the seven scenarios) tab. 1:

Architectures tab. 1:

Other reported results (text and figures):

• Per-scenario, RN-50 fig. 4: LAME does not help the i.i.d. likelihood-shift scenarios and is about 0.5 points below the baseline in the worst case §7. In non-i.i.d. prior-shift scenarios its average absolute gain is 6.7 points, rising to 15 points on ImageNet-V2 §7. The network-adapting methods do not beat the baseline on prior shift §7.
• Robustness to training recipe: LAME keeps its gain over the baseline with the transferred settings, with about 4 points on SimCLR §7 tab. 1. The other methods lose ground §7.
• Batch size: the gain is close to 4 points on average across batch sizes app. E fig. 9.
• Cross-shift validation matrix: LAME's worst-case degradation is much smaller than TENT's §7 fig. 3.
• Runtime and memory: the text says LAME avoids a second forward pass and needs about half the GPU memory of the network-adapting methods §7 fig. 6. The introduction says both time and memory are reduced by half §1. The paper gives no runtime numbers in the text; they are only in fig. 6.

Paper vs code

Official repository: https://github.com/fiveai/LAME, linked from the paper's abstract, at commit d2e5f63. Its LICENSE file is CC BY-NC-SA 4.0 (non-commercial) LICENSE. This card does not quote any code. The pseudocode above follows the paper's equations.

Discrepancies, each with both citations:

1. Affinity default. The paper's stated affinity is kNN app. C. The code default is rbf defaults.py. The best-config file sets only LAME_KNN and LAME_SIGMA overall_best.yaml, so RBF is used there. Which affinity produced Table 1 is not stated.
2. kNN direction. The paper defines w_ij = 1 when phi(x_i) is among the neighbours of phi(x_j) app. C. The code sets W[i, j] = 1 when j is among the neighbours of i lame.py. The matrix is transposed relative to the paper's formula, and the two differ when the default non-symmetric matrix is used.
3. Symmetrisation. The code option LAME_FORCE_SYMMETRY defaults to False defaults.py and, when set, averages W with its transpose lame.py. The paper does not mention symmetrisation.
4. RBF bandwidth. The paper's sigma is the average distance to the k-th neighbour app. C. The code takes the k-th smallest distance per row of a distance matrix that includes the zero self-distance, so it uses the (k-1)-th true neighbour, then averages over the batch lame.py. The diagonal is kept at 1 lame.py. The config value LAME_SIGMA is overwritten and has no effect lame.py defaults.py.
5. Feature normalisation. The paper's phi is the pre-trained feature extractor §5 with no normalisation stated. The code L2-normalises features before any affinity lame.py. With the linear affinity this turns it into cosine similarity.
6. Superclass pooling. The paper uses the average of the source probabilities over each superclass, with no renormalisation app. A. The code sums the probabilities of each superclass and then renormalises across superclasses classifier.py. These are different mappings unless the superclass sizes are equal.
7. Start point and stop rule. The paper states neither. The code starts from the source probabilities and stops after at most 100 iterations or when the relative change in the energy is at most 1e-8 lame.py. It also adds 1e-10 inside the logarithm of the source probabilities lame.py.
8. Update rule. Matches eq. 5 with unit Laplacian weight. The code computes a softmax of the log source probabilities plus the affinity term lame.py, which is the same update in the log domain.
9. Feature source for EfficientNet. The code stores the global average of the final feature map as the feature used for the affinity efficient_net.py. The paper does not state which layer the feature is taken from.

No further discrepancies found in the objective or the update beyond the items above.

Pitfalls

• The convergence argument needs a positive semi-definite W §5. The Gaussian (RBF) and linear kernels are PSD by construction. A binary kNN matrix is generally neither symmetric nor PSD. The paper does not discuss this, so implementers using kNN should expect no convergence guarantee, and should check the energy trace.
• Start point, stop rule and iteration count are not in the paper. A fixed low iteration count changes the answer, so record the count used.
• Numerics: the update is a ratio of exponentials of sums over neighbours, which can overflow on large or dense batches. Compute it in the log domain (log-sum-exp or softmax), as the code does lame.py.
• Memory: the code builds pairwise difference tensors of shape (N, N, d) for both kNN and RBF lame.py lame.py. Memory grows with N squared times d, not just N squared. An implementation can compute distances from squared norms and a matrix product to avoid this.
• Batch-only information: each batch is handled alone, so the method cannot learn across batches §5 §4. The gain depends on whether the batch contains enough similar samples. The objective has no class-balance term by design, since the authors consider class balance irrelevant for test-time adaptation §5.
• The kNN code drops the first sorted neighbour, assuming it is the sample itself lame.py. With duplicate features, a different sample may be dropped.
• Self-affinity: with RBF the diagonal stays at 1 lame.py. Whether the paper's diagonal is included is not stated.
• "No hyperparameter" is only partly true: the paper's claim covers the bound optimisation §7 and the kernel choice and k are still tuned app. C.
• Results are author-reported, cover ImageNet-trained classifiers only §6, and are reported as averages over scenarios. The i.i.d. result is a slight loss §7.
• The runtime advantage is shown only in a figure; no numbers appear in the text §7 fig. 6.
• Licence: the code is CC BY-NC-SA 4.0 (non-commercial). Implement from the paper's equations, not from copied code, and do not copy snippets into any commercial product.

Lineage

• Builds on: Laplacian (graph) regularisation used in semi-supervised learning and graph clustering §1 §5, and Laplacian-regularised few-shot learning by Ziko et al. (ICML 2020), whose kNN affinity choice the paper follows app. C. The optimisation uses the concave-convex procedure §5.
• Compared against: the unadapted baseline; TENT (entropy minimisation); SHOT-IM; PseudoLabel; AdaBN §6.
• Related components cited: Taori et al. (2020), whose max-pooling over superclasses the paper compares with average pooling app. A; Burns and Steinhardt (2021), cited for adapting only early layers under prior shift app. C.
• Superseded by: not stated in this wiki.