PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Adaptation via Conjugate Pseudo-labels

Sachin Goyal, Mingjie Sun, Aditi Raghunathan, Zico Kolter

arXiv:2207.09640v2cs.LG

TL;DR

The paper addresses how to choose an effective unsupervised loss for test-time adaptation under distribution shift. It derives a conjugate adaptation loss from the classifier’s supervised training loss, explaining why meta-learning recovers entropy for cross-entropy models and different losses for others. Across datasets and training losses, the method consistently outperforms prior TTA losses, while its meta-learning exploration is limited to functions of individual-input logits and practical adaptation still requires design heuristics.

  • Problem

    The paper asks how to choose a TTA loss when prior unsupervised objectives, including entropy minimization, have largely been selected heuristically.

  • Method

    The paper derives an unsupervised TTA loss from the convex conjugate of the supervised training loss and interprets it as self-training with conjugate pseudo-labels.

  • Results

    Across datasets and training losses, conjugate pseudo-labeling consistently outperforms prior TTA losses; for PolyLoss classifiers, it achieves 36.83% error versus 38.23% for entropy minimization on CIFAR-100-C.

  • Takeaways & Limitations

    The appropriate TTA objective depends on the classifier’s training loss, with entropy emerging for cross-entropy models but not necessarily for classifiers trained with other losses.

  • Takeaways & Limitations

    The meta-learning framework is restricted to functions of individual-input logits, while effective TTA also involves design choices such as parameter updates and temperature scaling.

Abstract

from arXiv · show

Test-time adaptation (TTA) refers to adapting neural networks to distribution shifts, with access to only the unlabeled test samples from the new domain at test-time. Prior TTA methods optimize over unsupervised objectives such as the entropy of model predictions in TENT [Wang et al., 2021], but it is unclear what exactly makes a good TTA loss. In this paper, we start by presenting a surprising phenomenon: if we attempt to meta-learn the best possible TTA loss over a wide class of functions, then we recover a function that is remarkably similar to (a temperature-scaled version of) the softmax-entropy employed by TENT. This only holds, however, if the classifier we are adapting is trained via cross-entropy; if trained via squared loss, a different best TTA loss emerges. To explain this phenomenon, we analyze TTA through the lens of the training losses's convex conjugate. We show that under natural conditions, this (unsupervised) conjugate function can be viewed as a good local approximation to the original supervised loss and indeed, it recovers the best losses found by meta-learning. This leads to a generic recipe that can be used to find a good TTA loss for any given supervised training loss function of a general class. Empirically, our approach consistently dominates other baselines over a wide range of benchmarks. Our approach is particularly of interest when applied to classifiers trained with novel loss functions, e.g., the recently-proposed PolyLoss, where it differs substantially from (and outperforms) an entropy-based loss. Further, we show that our approach can also be interpreted as a kind of self-training using a very specific soft label, which we refer to as the conjugate pseudolabel. Overall, our method provides a broad framework for better understanding and improving test-time adaptation. Code is available at https://github.com/locuslab/tta_conjugate.

1 Introduction

The paper studies how to choose unsupervised test-time adaptation losses under distribution shift. It finds that meta-learned losses reflect the source classifier’s training loss and proposes a convex-conjugate framework that consistently improves adaptation.

  • Distribution shifts can sharply reduce deep-network performance, while existing robustness methods are often specialized to predefined perturbations or domains.
  • TTA adapts a source model to new test inputs by optimizing an unsupervised objective, but selecting that objective has remained largely heuristic.
  • Meta-learning a broad class of TTA losses surprisingly recovers a temperature-scaled softmax-entropy objective for cross-entropy-trained classifiers.
  • The convex-conjugate formulation explains these meta-learning results and provides a recipe for deriving TTA losses from general supervised training losses.
  • Across datasets and training losses, conjugate pseudo-labeling consistently outperforms prior TTA losses and improves state-of-the-art TTA performance.

2 Background and preliminaries.

The paper defines TTA as updating a trained source model on unlabeled test inputs using an unsupervised objective. It uses meta-learning to explore which such objectives produce effective adaptation, rather than proposing meta-learning itself as the deployment method.

  • TTA updates a trained source model on test inputs before prediction using an unsupervised objective, such as TENT’s entropy of softmax-normalized predictions.
  • The TTA loss is parameterized by a neural network that maps model predictions or logits to a scalar loss, with parameters learned by differentiating through adaptation.
  • The meta-learning procedure uses corresponding test labels through a supervised surrogate objective, but only to investigate what effective TTA losses look like.

3 Test-time Adaptation via Meta-Learnt Losses

Meta-learning reveals structured TTA losses: scaled softmax-entropy for cross-entropy-trained models and a negative squared-loss-like function for squared-loss-trained models. Directly fitting these forms reproduces or slightly improves the learned adaptation performance.

  • 12.32% error from fitted scaled softmax-entropy essentially matches the meta-TTA loss, which achieves 12.35% versus TENT’s 13.14%.The fitted objective is α · H(softmax(hθ(x)/T)).
  • For squared-loss-trained classifiers, the learned meta-loss differs from entropy, resembles negative squared error, and slightly outperforms the meta-TTA loss when fitted directly with a quadratic.The squared-loss meta-TTA loss obtains 13.48% error versus 14.57% for entropy.
  • For cross-entropy-trained classifiers, the learned meta-loss consistently takes the form of temperature-scaled softmax-entropy across architectures and meta-learning variants.
  • The experiments show that the best TTA loss depends cleanly on the loss used to train the source classifier, despite the meta-loss network being able to represent complex functions.

4 Conjugate Pseudo Labels

The paper derives a generic test-time adaptation loss from the convex conjugate of the supervised training loss, yielding loss-specific conjugate pseudo-labels. This framework recovers entropy minimization for cross-entropy classifiers while providing distinct objectives for squared-loss and PolyLoss classifiers.

  • 4.1 Losses and the convex conjugate: The convex-conjugate framework defines an unsupervised adaptation loss as the negative conjugate evaluated at the gradient of the model output.Under over-parameterization and near-optimal source parameters, this expression locally approximates the original supervised empirical loss without requiring target labels.
  • Cross-entropy: For cross-entropy classifiers, the conjugate adaptation loss is exactly softmax-entropy, recovering the TENT objective and the loss found by meta-learning.The cross-entropy choice f(h)=log sum_i exp(h_i) yields softmax model predictions through the gradient of f.
  • Squared loss: Squared-loss classifiers instead receive a quadratic conjugate adaptation loss rather than softmax-entropy.For squared loss, f(h)=1/2||h||_2^2, so the optimality condition is y=h and the conjugate is quadratic.
  • Self-training with conjugate pseudo-labels: Conjugate adaptation is equivalent to self-training with soft pseudo-labels given by the gradient of f at the current hypothesis.For cross-entropy these pseudo-labels are the current softmax outputs, whereas other training losses can produce different pseudo-labels.
  • PolyLoss: The method extends to PolyLoss through an expanded conjugate form that yields loss-specific pseudo-labels based on the Jacobian of its transformed hypothesis.The paper notes that operating directly on this expanded form is typically easier for PolyLoss.
  • Test-time adaptation: Across models, datasets, and training losses, conjugate pseudo-labeling consistently outperforms prior TTA losses and improves test-time adaptation over the state of the art.The method is applied to shifted-domain data by self-training with conjugate pseudo-labels as a proxy for fine-tuning on the unavailable out-of-distribution loss.

5 Experiments

The experiments evaluate conjugate pseudo-labeling across corruption and domain-adaptation benchmarks for classifiers trained with cross-entropy, PolyLoss, and squared loss. Conjugate pseudo-labeling consistently improves adaptation, matching temperature-scaled entropy for cross-entropy while outperforming entropy-based and other baselines for alternative training losses.

  • 5.2 Results on classifiers trained with cross-entropy: For cross-entropy-trained classifiers, conjugate pseudo-labeling exactly reduces to entropy minimization with additional temperature scaling.The experiments omit Softmax PL and ENT because they coincide with conjugate PL in this setting.
  • 5.2 Results on classifiers trained with cross-entropy: 36.10% error for temperature-scaled entropy on CIFAR-100-C surpasses vanilla ENT, robust PL, and MEMO.The passage reports 41.15% for vanilla ENT, 39.80% for robust PL, and 38.52% for MEMO.
  • 5.3 Results on classifiers trained with polyloss and squared loss: Across PolyLoss corruption and domain-adaptation benchmarks, conjugate PL outperforms all other TTA losses.The reported comparison spans Tables 2 and 3 and includes CIFAR-10-C, CIFAR-100-C, ImageNet-C, ImageNet-R, VisDA-C, and digit adaptation tasks.
  • 5.3 Results on classifiers trained with polyloss and squared loss: 36.83% error versus 38.23% for entropy minimization on CIFAR-100-C demonstrates PolyLoss-specific gains from conjugate PL.With PolyLoss classifiers, entropy minimization is no longer the best adaptation objective.
  • 5.3 Results on classifiers trained with polyloss and squared loss: 44.01% versus 45.34% error on ImageNet-C shows PolyLoss with conjugate PL outperforming every tested cross-entropy adaptation method.Using softmax-entropy instead would yield 45.5% for the PolyLoss classifier, worse than the cross-entropy result.
  • 5.3 Results on classifiers trained with polyloss and squared loss: 12.87% error on CIFAR-10-C for squared-loss conjugate PL beats ENT at 13.24% and Softmax PL at 31.81%.The passage reports consistent gains over baselines on CIFAR-10 and CIFAR-100 experiments.

6 Related Works

Related work studies test-time adaptation through unsupervised objectives, self-supervised auxiliary tasks, and domain-adaptation settings with different data access. This paper connects TTA to self-training by selecting soft pseudo-labels from the source training objective.

  • Test-time adaptation methods: TENT minimizes prediction entropy, while later TTA methods propose alternative unsupervised objectives.These approaches update the source classifier on target-distribution samples without labels.
  • Test-time adaptation methods: Self-supervised TTA methods require modifying source training with an auxiliary task, unlike methods for standard supervised classifiers.Examples include rotation prediction and masked autoencoders.
  • Source-free domain adaptation: Source-free domain adaptation accesses the entire unlabeled target dataset at once, whereas TTA adapts over test inputs.Source-free methods optimize the trained source hypothesis using objectives including entropy, diversity, and self-training.
  • Unsupervised domain adaptation: Unsupervised domain adaptation typically uses labeled source and unlabeled target data during training to align representations.The cited approaches match statistics, use discriminators, or apply self-supervised learning.
  • Pseudo-labels and self-training: The paper generalizes self-training by replacing labels with soft pseudo-labels derived from the conjugate analysis of the source training objective.TENT is characterized as self-training with instantaneous softmax pseudo-labels.

7 Conclusion, Limitations and Future Directions

The paper proposes conjugate pseudo-labeling as a general TTA loss and argues that it approximates the supervised loss for a broad class of training objectives. It identifies unresolved questions about optimal TTA losses, regularization heuristics, real-world shift settings, and extensions beyond TTA.

  • Conclusion: The proposed conjugate loss is motivated by meta-learning experiments and is shown to approximate the oracle supervised loss for a broad set of functions.The authors note that this suggests, but does not fully establish, optimality of the loss.
  • Limitations: The meta-learning framework is limited to functions over individual-input logits, excluding intermediate representations and interactions across batches.The paper identifies these richer function classes as possible extensions.
  • Limitations: Good TTA also depends on heuristics such as updating only batch-normalization parameters and choosing temperature scaling, whose roles remain insufficiently understood.The authors connect these choices to regularizing adaptation so updates do not move the model too far.
  • Future Directions: The paper leaves open which real-world distribution shifts are most suitable for self-training-based adaptation and suggests applying conjugate pseudo-labeling to semi-supervised learning.These are presented as future research directions rather than established capabilities.

A.1 Conjugate Derivations

This appendix derives the conjugate formulation for the cross-entropy setting. It introduces the function whose convex conjugate yields the corresponding adaptation loss.

  • Cross-entropy derivation: For cross-entropy, the derivation defines f through the log-sum-exp expression associated with class logits and probabilities.The supplied derivation passages show the definition and subsequent conjugate calculation, although the displayed formulas are fragmented.
  • Conjugate loss: The appendix then evaluates the convex conjugate f* and uses it to obtain the conjugate loss for cross-entropy under the stated constraint.The sequence proceeds from defining f to computing f* and identifying the resulting loss.

A.2 Experiments on Binary Classification with Exponential Loss

The appendix evaluates conjugate adaptation for binary classification with exponential training loss on synthetic Gaussian clusters under three shift severities. Conjugate adaptation is compared with hard pseudo-labeling and sigmoid-entropy baselines.

  • Experimental setup: The experiment uses a 100-dimensional synthetic binary-classification dataset with Gaussian clusters and evaluates easy, moderate, and hard shifts controlled by cluster drift.The shift settings are λ = 0.6, 0.65, and 0.7, respectively.
  • Results: Conjugate adaptation generalizes faster under easy and moderate shifts, while under hard shift it is the only method that improves shifted-data accuracy.Entropy deteriorates performance under the hard shift, according to the reported Figure 3 comparison.
  • Evaluation: Figure 3 plots accuracy against adaptation iterations across the three shift levels to assess generalization to shifted Gaussian clusters.The comparison covers the conjugate method and both baseline approaches.

A.3 Meta Learning Experiment Details

The meta-learning experiment parameterizes the TTA loss with a transformer and MLP operating on model logits, then differentiates through adaptation to optimize the loss. The procedure uses supervised task loss only to update the meta-loss during exploration.

  • Architecture: A transformer followed by an MLP maps the model’s logits to a scalar meta-loss value used to update the source model.The logits are organized per class, transformed into a learned representation, and passed through the MLP.
  • Input representation: The transformer input is constructed from per-class logits, with each class represented by a concatenated feature vector.The supplied implementation description specifies the class-wise input structure and transformer output.
  • Optimization procedure: The meta-learning loop updates the source model with the learned meta-loss, then updates the meta-loss using a supervised task loss.The supervised loss is used in the exploratory meta-learning procedure rather than proposed as a TTA objective.
  • Implementation details: The architecture and MLP depths are selected from small discrete search spaces, and Adam with learning rate 1e−3 trains the meta-loss for 100 epochs with batch size 200.These are implementation settings for the meta-learning experiments.

A.4 Effect of Task Loss in Meta Learning

The meta-learning task loss barely changes the learned TTA loss; instead, the source classifier and its training loss determine its form. Cross-entropy-trained classifiers yield temperature-scaled softmax-entropy-like losses, whereas squared-loss-trained classifiers yield quadratic losses.

  • Effect of Task Loss: The learned meta loss is determined primarily by the source model and its training loss rather than the task loss used during meta-learning.The paper reports that meta losses differ substantially across source classifiers trained with different source losses, while remaining insensitive to the meta task loss.
  • Effect of Task Loss: The meta task loss barely affects the learned TTA loss, including for classifiers trained with cross-entropy or squared loss.For a cross-entropy-trained ResNet-26, losses learned using cross-entropy and squared-loss meta objectives are nearly unchanged; the same observation holds for squared-loss-trained classifiers.
  • Cross Entropy: For cross-entropy-trained classifiers, the learned meta loss is well approximated by temperature-scaled softmax-entropy across digit, dataset, and architecture settings.The approximation is observed for SVHN-to-MNIST and SVHN-to-USPS adaptation and across ResNet-26/ResNet-50 and CIFAR-10/CIFAR-100 settings.
  • Squared Loss: For classifiers trained with squared loss, the learned meta loss follows a quadratic form, consistent with the convex-conjugate formulation.This behavior is reported across the squared-loss experiments and contrasts with the temperature-scaled softmax-entropy behavior found for cross-entropy-trained classifiers.
Loading 2207.09640v2…