PaperlayerView exact revision on arXiv

Source-linked AI summary

From Inference to Adaptation: A Unified Optimal Transport View of Vision Language Model

Qi Yu, Zhichen Zeng, Katherine Tieu, Xiyuan Yang, Ruizhong Qiu, Yuchen Yan, Lihui Liu, Yanjun Zhao, Lingjie Chen, Jingrui He, Hanghang Tong

arXiv:2608.18339v1cs.CVcs.AIcs.CLcs.LG

TL;DR

VLM test-time adaptation is hindered by unreliable pseudo-labels and coarse objectives that miss sample-level cross-modal relationships under distribution shift. ORIGIN unifies Wasserstein optimal-transport inference with soft-label InfoNCE adaptation, consistently achieving state-of-the-art results across three benchmarks, including a 7.4% mean-accuracy gain on CIFAR10-C without added latency.

  • Problem

    Under distribution shift, raw embedding similarities produce noisy pseudo-labels, while class-level surrogate objectives miss individual image-text relationships needed for sample-level alignment.

  • Method

    ORIGIN uses Wasserstein optimal transport for robust sample-level pseudo-labels, soft-label InfoNCE for fine-grained adaptation, and theoretically unifies both objectives.

  • Results

    7.4% mean-accuracy improvement on CIFAR10-C over the fastest baseline without additional latency, while consistently achieving state-of-the-art performance across three benchmarks.

  • Takeaways & Limitations

    A shared OT perspective aligns inference and adaptation at sample level, supporting robust inference and fine-grained adaptation for VLMs under distribution shifts.

  • Takeaways & Limitations

    Existing test-time adaptation methods remain vulnerable to noisy pseudo-labels from raw similarities and coarse class-level supervision under distribution shifts.

Abstract

from arXiv · show

Vision-language models (VLMs) have demonstrated remarkable zero-shot capabilities yet remain sensitive to real-world distribution shifts during inference. Although significant efforts are devoted to adapting VLMs at test time, they rely heavily on noisy pseudo-labels predicted directly from raw embedding similarities during inference, which are unreliable under distribution shift and mislead the adaptation. To avoid noise amplification, existing works craft coarse-grained surrogate objectives during adaptation, which fail to explicitly model sample-level relationships across different modalities, creating objective mismatch with inference, thus leading to marginal performance improvement. In this work, we aim to bridge the detached objectives of inference and adaptation for VLMs, and propose a principled VLM TTA method called \algname. For VLM inference, we formulate the zero-shot image classification task as a cross-modal alignment problem encoded via a Wasserstein OT formulation, providing robust pseudo-labels at the sample-level to effectively adapt VLMs. For VLM adaptation, we adopt a soft-label InfoNCE loss to adapt VLMs based on the OT-induced pseudo-labels, leveraging fine-grained supervisions to explicitly model relationships of individual image-text pairs via contrastive learning, which empowers accurate inference at the same granularity. Moreover, we theoretically reveal that the InfoNCE loss can be neatly reformulated as a Wasserstein OT formulation, thereby unifying the objectives of the inference and adaptation of VLMs to achieve their mutual benefits. Extensive experiments demonstrate the effectiveness and efficiency of our methods, outperforming the best-performing methods by up to 7% with state-of-the-art efficiency.

1 Introduction

The introduction identifies noisy raw-similarity pseudo-labels and coarse class-level adaptation objectives as key limitations of VLM test-time adaptation. It presents ORIGIN, which uses Wasserstein optimal transport for robust inference and soft-label InfoNCE for fine-grained adaptation, theoretically unifying both objectives and achieving up to 7% mean accuracy improvement.

  • Motivation: VLM embedding spaces enable strong zero-shot multimodal capabilities but remain vulnerable to real-world distribution shifts.VLMs project image and text embeddings into a shared space and support tasks including classification, captioning, and segmentation.
  • Limitations: Raw image-text embedding similarities under distribution shifts produce noisy pseudo-labels that can limit or misguide adaptation.The introduction reports that such pseudo-labels have poor accuracy and may cause suboptimal or degraded adaptation performance.
  • Limitations: Existing methods use coarse class-level prototype objectives to suppress noise, but these objectives fail to model individual cross-modal image-text relationships.Treating samples uniformly creates an objective mismatch with sample-level inference.
  • ORIGIN: ORIGIN formulates zero-shot classification as Wasserstein optimal transport and adapts VLMs with soft-label InfoNCE using OT-induced pseudo-labels.The method uses the OT transport plan for inference and fine-grained contrastive supervision for adaptation, yielding consistent accuracy improvement.
  • Unified framework and results: ORIGIN theoretically reformulates InfoNCE as entropic Wasserstein OT, unifying inference and adaptation through alternating transport-plan and VLM-parameter optimization.Experiments on three standard benchmarks show up to 7% mean accuracy improvement over state-of-the-art test-time adaptation methods.

2 Preliminaries

This section introduces VLM test-time adaptation and optimal transport preliminaries. VLMs map images and class descriptions into a shared embedding space for zero-shot classification, while OT provides transport plans that encode soft correspondences between distributions.

  • VLM TTA: A pre-trained VLM maps images and text descriptions into a shared d-dimensional embedding space using separate image and text encoders.For n classes, class descriptions produce embeddings {t_j} and an image produces embedding v.
  • VLM TTA: Zero-shot predictions compare test-image and class-text embedding similarities, while online TTA processes shifted test images sequentially in batches.The prediction rule is arg max_j v^⊤t_j.
  • Optimal Transport: Discrete OT finds a transport plan minimizing total cost between two probability distributions, with the optimal value defining their Wasserstein distance.The transport plan encodes soft correspondence between points from the two distributions.
  • Optimal Transport: Entropic regularization makes the OT problem strongly convex and efficiently solvable with the Sinkhorn algorithm at quadratic complexity.The optimal plan has scaling form S*=diag(a) K diag(b), where K := exp(−C/ϵ), and Sinkhorn iteratively enforces marginal constraints.

3 Methodology

ORIGIN unifies VLM inference and test-time adaptation through Wasserstein optimal transport: OT produces robust sample-level image-text pseudo-labels, while soft-label InfoNCE adapts image representations with matching fine-grained relationships. Theoretically, these objectives are equivalent, making inference and adaptation mutually beneficial rather than decoupled.

  • 3.1 OT-Based Inference: ORIGIN formulates image-text alignment as an entropic Wasserstein OT problem to reduce noise from shifted embedding similarities and generate reliable pseudo-labels.The constrained OT formulation provides globally informed alignment and predicts each image’s label from the solved transport map.
  • 3.1 OT-Based Inference: OT assigns equal importance to test images and uses a uniform initial text distribution that is dynamically updated with previous-batch predictions through EMA.The text-level marginal distribution is adjusted according to prior test batches, while the image-level marginal remains uniform.
  • 3.2 Soft-Label InfoNCE Adaptation: ORIGIN adapts VLMs with a soft-label InfoNCE objective that explicitly models image-text relationships at the same sample level as inference.The OT transport map S serves as soft pseudo-labels, and KL divergence with temperature τ defines the adaptation objective.
  • 3.2 Soft-Label InfoNCE Adaptation: During image-side distribution shifts, ORIGIN replaces unavailable hard labels with OT-induced soft labels and adapts only the image encoder for image-to-text alignment.The framework can be extended with text-to-image alignment for text-level distribution shifts.
  • 3.3 Theoretical Unification: The soft-label InfoNCE loss is mathematically equivalent to the entropic Wasserstein OT formulation, with temperature τ corresponding to entropic regularization weight ϵ.Thus, inference optimizes the transport plan S while adaptation optimizes image-encoder parameters θ under the same objective family; the unification also extends to text-side adaptation.

4 Experiments

Experiments show that ORIGIN achieves strong performance under distribution shifts while jointly improving accuracy and runtime efficiency. Ablations and sensitivity analyses support the benefits of OT-based pseudo-labels, InfoNCE adaptation, and robust hyperparameter choices.

  • Effectiveness: ORIGIN consistently achieves state-of-the-art performance across three VLM TTA benchmarks, improving mean accuracy by 7.4% on CIFAR10-C and 4.8% on CIFAR100-C.
  • Effectiveness-efficiency trade-off: ORIGIN improves mean accuracy by up to 7.4% over the fastest baselines without additional latency and delivers up to 45× runtime speed-up over other training-based methods.
  • Unified OT view: OT-based inference consistently outperforms embedding-based inference under the same adaptation loss, with average performance improvements of 1.78%, 1.68%, and 15.84%.The comparison evaluates combinations of embedding- or OT-based inference with MINT, BATCLIP, and InfoNCE adaptation losses.
  • Hyperparameter sensitivity: Overly small or large temperature τ or batch size b degrades performance by causing overconfident or noisy supervision, limited global OT information, or fewer adaptation steps.Small τ can overconfidently introduce errors, whereas large τ produces noisy supervision; batch size creates a corresponding adaptation–inference trade-off.
  • Hyperparameter sensitivity: ORIGIN remains insensitive to the EMA parameter α, with less than 1% performance variation across choices.The parameter trades off historical and current text-level distributions.

5 Related Works

Prior VLM test-time adaptation methods address distribution shifts through prompt optimization, memory management, or data augmentation and selection. Optimal transport has also been applied to prompt learning, model selection, and improving vision-language models across downstream tasks.

  • VLM Test-time Adaptation: VLM test-time adaptation operates without source training data or test labels, using prompt optimization, historical-sample memory management, or data augmentation and selection.These methods adapt pre-trained VLMs from a source domain to a shifted target domain at test time.
  • Optimal Transport for VLMs: Optimal transport supports VLM prompt learning, zero-shot model selection across datasets, and improvements to vision-language models on downstream tasks.Representative methods include PLOT, Prompt-OT, SWAB, and OT-CLIP.

6 Conclusion

The paper unifies VLM inference and test-time adaptation through optimal transport, using Wasserstein OT for robust sample-level pseudo-labels and soft-label InfoNCE for fine-grained adaptation.

  • Conclusion: The method theoretically bridges VLM inference and adaptation at test time through optimal transport.This unification is intended to achieve their mutual benefits.
  • Conclusion: Wasserstein OT reformulates VLM inference to generate robust sample-level pseudo-labels for effective adaptation.The formulation targets improved pseudo-label robustness during inference.
  • Conclusion: Soft-label InfoNCE leverages fine-grained supervision signals during VLM adaptation.The conclusion identifies this loss as the adaptation component of the proposed framework.

A Proof · A.1 Proof of Lemma 3.1

The appendix proves Lemma 3.1 by expanding KL(S∥˜Pθ) under the transport plan’s probability-matrix and uniform image-side marginal constraints.

  • A.1 Proof of Lemma 3.1: Under these transport-plan constraints, the proof formulates KL(S∥˜Pθ) from Eq. (7).The passage explicitly identifies KL(S∥˜Pθ) in Eq. (7) as the term being formulated.
  • A.1 Proof of Lemma 3.1: The expansion separates a term identified as constant with respect to S.The passage marks this component as “constant w.r.t. S.”
  • A.1 Proof of Lemma 3.1: The proof begins by assuming that S is a probability matrix with a uniform marginal distribution at the image side.This constraint is the stated starting condition for the derivation.
  • A.1 Proof of Lemma 3.1: The derivation expands KL(S∥˜Pθ) into terms involving the transport entries S_i,j.The displayed expansion includes the summation term involving S_i,j.
  • A.1 Proof of Lemma 3.1: One resulting component is the sum over i,j of S_i,j(log S_i,j −1).This expression is explicitly shown in the proof passage.
  • A.1 Proof of Lemma 3.1: Combining the formulation and expansion establishes the claimed result of Lemma 3.1.The appendix explicitly concludes that Lemma 3.1 has been proven.

B Detailed Experimental Setup

Experiments evaluate VLM test-time adaptation on three corruption benchmarks under the standard highest-severity protocol, using CLIP backbones and accuracy. Results are averaged across randomized runs with matched batch sizes and official baseline settings.

  • Datasets & Metrics: Experiments use CIFAR10-C, CIFAR100-C, and ImageNet-C, each containing 15 corruption types, with Level 5 severity and classification accuracy as the metric.These benchmarks follow the standard test-time adaptation protocol.
  • VLM Backbones: CLIP with ViT-B-16 and ViT-B-32 visual encoders serves as the VLM backbone, with ViT-B-16 used by default.The default changes only when otherwise noted.
  • VLM TTA Baselines: ORIGIN is compared with state-of-the-art VLM TTA methods spanning training-free approaches and training-based methods that optimize prompts or internal tunable parameters.Named training-free baselines include TDA, DMN, VTE, ZERO, ECALP, and TENT; the supplied passage truncates the training-based list.
  • Reproducibility: Reported results average 5 randomized data-loader runs, while all baselines use matched batch sizes of 200 for CIFAR10-C/CIFAR100-C and 64 for ImageNet-C.Baseline hyperparameters use official implementation defaults, and ORIGIN settings are provided in Table 3.
  • Machine: Experiments run on a server with dual Intel Xeon Gold 6240R CPUs and four NVIDIA Tesla V100-SXM2 32GB GPUs.This describes the reported experimental hardware.

C Additional Experimental Results

On ViT-B-32, ORIGIN consistently achieves the best performance across all benchmarks, with up to 10% improvement in mean accuracy.

  • C Additional Experimental Results: ORIGIN achieves the best performance across all ViT-B-32 benchmarks, improving mean accuracy by up to 10%.Table 4 marks the first- and second-best results with bold and underline, respectively.

D Limitations & Future Works · E Boarder Impact

ORIGIN shows strong performance but is limited by its dependence on online adaptation and model gradients, motivating episodic and black-box extensions. The paper targets cross-modal learning, vision-language models, and optimal transport, while identifying no societal consequences requiring specific emphasis.

  • D Limitations & Future Works: ORIGIN is an online VLM TTA method that is not directly applicable to strictly episodic TTA without explicit memory for historical samples.An episodic setting would require maintaining historical samples explicitly.
  • D Limitations & Future Works: ORIGIN requires access to model gradients for test-time adaptation, as do most training-based VLM TTA methods.This gradient requirement is a limitation of the adaptation procedure.
  • D Limitations & Future Works: Gradient dependence may limit ORIGIN’s applicability to closed-source proprietary models.The limitation follows from restricted access to model gradients in such models.
  • D Limitations & Future Works: Future work could develop an episodic extension of ORIGIN using efficient OT variants for incremental transport-plan updates.The paper specifically suggests sliced OT as an example of an efficient OT variant.
  • D Limitations & Future Works: Future work could extend the OT framework to black-box adaptation settings.This direction would address settings where model gradients are unavailable.
  • E Boarder Impact: The paper aims to advance cross-modal learning, vision-language models, and optimal transport.These are the research areas identified by the paper’s broader-impact statement.
  • E Boarder Impact: The authors identify many potential societal consequences but do not consider any to require specific highlighting.This statement summarizes the paper’s assessment of societal consequences.

NeurIPS Paper Checklist

The NeurIPS checklist requires standardized answers with brief justifications and remains part of the submission and published paper. This checklist records claims, limitations, reproducibility, experimental details, statistical significance, and societal-impact considerations.

  • Checklist instructions: The checklist is integral to submission, visible to reviewers and ethics reviewers, and published with the final paper.It follows the references and optional supplemental material and does not count toward the page limit.
  • Checklist instructions: Checklist questions require [Yes], [No], or [N/A] answers, each followed by a 1–2 sentence justification.[N/A] indicates that a question is not applicable or relevant information is unavailable.
  • Claims and limitations: The authors affirm that the abstract and introduction accurately state the paper’s contributions and scope, supported by theoretical analysis and extensive experiments.The checklist guidelines require claims to match theoretical and experimental results and reflect assumptions, limitations, and expected generalization.
  • Claims and limitations: The paper reports limitations and future work in Appendix D, while the guidelines require discussion of strong assumptions and robustness to their violations.The checklist encourages a separate limitations section and distinguishes [N/A] from [No] for limitation reporting.
  • Experiments and reproducibility: The checklist guidelines address reproducibility, open access to data and code, and the need to explain societal-impact considerations, including potential harms and fairness issues.Code release is encouraged but [No] can be acceptable unless code is central to the contribution; societal-impact answers of [N/A] or [No] require explanation.
  • Experiments and reproducibility: The paper answers [Yes] for experimental details and statistical significance, citing Appendix B and results averaged over 5 randomized data-loader runs.The guidelines emphasize sufficient training and test details and clear identification of variability captured by statistical reporting.
Loading 2608.18339v1…