Diverse Data Augmentation with Diffusions for Effective Test-time Prompt Tuning

In plain words

CLIP can label pictures it was never trained on, but accuracy drops when the pictures look unlike its training data. Test-time prompt tuning adjusts the text prompt for each picture on the fly, with no labels, by checking how confident the model is across several versions of that one picture.

The earlier method makes those versions by cropping and flipping, and the versions end up looking much alike. DiffTPT adds versions drawn by an image generator, which vary more in appearance but keep the same object. It then drops generated versions that look too different from the original and versions the model is unsure about. On the averages the paper prints, it beats the earlier method, though by less than its headline number suggests.

TL;DR

DiffTPT is test-time prompt tuning for CLIP. For each unlabeled test image it builds 128 views: the original, 64 crop-and-flip views, and 63 Stable Diffusion image variations. It keeps the 80% of generated views most similar to the original, then the 30% lowest-entropy views, and tunes four context tokens for four optimizer steps at learning rate 5e-3. The paper reports gains of 0.7 to 2.9 points over TPT on the averages in its own tables (author-reported); the 5.13% headline figure cannot be reproduced from them. The code is public but has no LICENSE file.

Problem

Zero-shot CLIP degrades under distribution shift. Prompt tuning fixes this only with labelled data from the target domain, so the paper studies the setting where a model sees one unlabeled test image at a time and must adapt to it §1 §3.1. The earlier test-time method (TPT) makes its views with cheap transforms, so the views vary little, and it filters them only by prediction entropy, which does not guarantee that a kept view still shows the same class §1 §3.2. Generating views with a diffusion model gives more variety, but the generated images can drift to other classes, so the paper pairs diffusion views with a similarity filter §3.2 §3.4.

Method

Setting and TPT baseline §3.1. A CLIP image encoder f and text encoder g score an image x against K class prompts with a temperature-scaled cosine similarity eq. 1. The generic objective in the paper is a loss L over a prompt v and one test image eq. 2. TPT's concrete objective minimizes the entropy of a mean prediction taken over the confident augmented views, where confidence is judged by each view's own entropy eq. 3 eq. 4.

Diffusion views §3.3. The paper uses Stable Diffusion V2. Its usual text-conditioning is replaced by the CLIP image embedding of the test image f(x), because no label is available at test time. Each generated view is D_n(x) = G(f(x), n_n), where n_n is fresh Gaussian noise eq. 5. Equation 6 swaps the standard views for these diffusion views in the TPT sum. The paper then uses both sources together, to combine the standard view's variety with the diffusion views' variety §3.3.

Cosine filter §3.4. Each generated view is compared with the test image by cosine similarity. The paper keeps views whose similarity is above a threshold ε chosen so that a fraction ρ_C of the views survives §3.4. The combined objective keeps the entropy indicator and adds the cosine indicator, normalized by ρ_H ρ_C N eq. 7. The paper does not state which representation the cosine uses; see Paper vs code.

Per-image procedure. The loop below condenses §3.1 to §3.4 and the released code. The paper has no algorithm box, so this is our reading.

# Per-image loop, condensed from §3.1-§3.4 and eq. 4-7.
# Inputs: one unlabeled test image and its 63 pre-generated SD views.

def difftpt_adapt(clip, x_test, sd_views):
    # Original, 64 crop/flip views, 63 SD views: 128 in total
    # [§4.1] [code: tpt_classification.py#L177-L178]
    views = [x_test] + augment(x_test, count=64) + sd_views
    # Context starts from "a photo of a", 4 tokens, reset per image
    # [§4.1] [code: tpt_classification.py#L279-L283]
    ctx = init_context("a photo of a", n_ctx=4)
    # Paper says Adam; released code uses AdamW at lr 5e-3
    # [§4.1] [code: tpt_classification.py#L146]
    optimizer = AdamW([ctx], lr=5e-3)
    keep = None
    # Four optimization steps  [§4.1] [code: scripts/do_tpt_difftpt.sh#L6]
    for step in range(4):
        # Class logits for every view  # (128, K)
        logits = clip.logits(views, ctx)
        if keep is None:
            # Filter once at step 1, then reuse the indices
            # [code: difftpt_utils.py#L133-L138]
            keep = select_views(logits)
        # Entropy of the mean probability over kept views  [eq. 4] [eq. 7]
        loss = avg_entropy(logits[keep])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    # Predict the original image with the tuned context
    return clip.logits(x_test, ctx)


def select_views(logits):
    # Cosine of each view's logits with the original view's logits
    # [§3.4] [code: difftpt_utils.py#L99-L102]
    sims = [cosine(logits[0], logits[i]) for i in range(1, 128)]
    # Keep the top 80% most similar views: floor(127 * 0.8) = 101
    # [§4.1] [code: difftpt_utils.py#L102]
    kept = top_k(sims, floor(127 * 0.8))
    # Original plus the kept views, mapped back to view indices
    # [code: difftpt_utils.py#L107]
    candidates = [0] + [i + 1 for i in kept]
    # Keep the 30% lowest-entropy candidates: floor(102 * 0.3) = 30
    # [§4.1] [code: difftpt_utils.py#L109-L110]
    entropies = [entropy(logits[i]) for i in candidates]
    lowest = lowest_k(entropies, floor(102 * 0.3))
    return [candidates[j] for j in lowest]

Key equations, as written in the paper with its numbering:

\tilde{p}_v = \frac{1}{\rho_H N} \sum_{n=1}^{N} \mathbb{1}[H(p_n) \le \tau]\, p_n(y \mid \mathcal{A}_n(x_{test})) eq. 4

\mathcal{D}_n(x_{test}) = \mathcal{G}(f(x_{test}), \mathbf{n}_n) eq. 5

\tilde{p}_v = \frac{1}{\rho_H \rho_C N} \sum_{n=1}^{N} \mathbb{1}[H(p_n) \le \tau]\, \mathbb{1}[\mathcal{M}_n]\, p_n(y \mid \mathcal{D}_n(x_{test})) eq. 7

Here M_n is the similarity mask from §3.4. Equation 6 is the same sum with D_n in place of A_n. The paper writes H(p_i) in eq. 6 and eq. 7 but sums over n; we read it as H(p_n).

Hyperparameters

Results

Author-reported. Setup: CLIP RN50 and ViT-B/16 backbones; top-1 accuracy (%) on 1000 randomly selected test images §4.1; baselines are zero-shot CLIP with the default prompt and with an 80-prompt ensemble, TPT, and CoOp and CoCoOp trained on ImageNet with 16 shots §4.1. Arrows against each group's baseline are omitted here; they can be recomputed from the values.

Scenario 1, natural distribution shift tab. 1. Columns are ImageNet (IN), ImageNet-A, ImageNet-V2, ImageNet-R, ImageNet-Sketch, the 5-dataset average, and the 4-dataset OOD average.

* Printed as 42.26. The five printed values average to 44.26 (see Pitfalls).

Scenario 2, cross-dataset generalization tab. 2. Each row lists the five datasets of its half of the table, then the 10-dataset average. The averages are over all 10 datasets, which is how the printed values reconcile. Baselines CoOp and CoCoOp are trained on ImageNet with 16 shots.

Differences against TPT, computed from the printed averages above (not printed by the paper): Scenario 1, RN50 +2.06 and ViT-B/16 +0.69; Scenario 2, RN50 +2.87 and ViT-B/16 +1.48. The appendix gives full-dataset averages of 54.42 → 57.52 for RN50 and 62.21 → 63.27 for ViT-B/16 app. A.

Ablations (author-reported; read from text, plots not reproduced) §4.3. ρ_H = 0.3 and ρ_C = 0.8 give the best result, and ρ_H above 0.5 lowers accuracy fig. 5. Accuracy rises with the number of views up to about N = 64 and holds at N = 128 fig. 6. Going from one to four optimization steps lifts Scenario 2 from 56.6 to 58.8, and further steps slightly reduce it fig. 7. Mixing standard and diffusion views is studied on a grid fig. 4.

Inference cost §4.3. The paper gives timings for TPT and for stable diffusion without a shared setting (see Pitfalls).

Paper vs code

Code paths are relative to the repository root at commit 03f11696ccc793888a2cac303a8cbf643ea171d5. The repository is github.com/chunmeifeng/DiffTPT, linked from the paper's author block. The repository owner is the paper's first author.

• Cosine space. The paper computes cos(D_n(x), x) between the generated view and the test image §3.4. The paper does not name the embedding. The code takes the model's per-class logits for each view and compares them with the logits of the original view difftpt_utils.py.
• Selection timing. The paper's objective uses the indicators in every step eq. 7. The code computes the cosine and entropy selection once, at the first step, and reuses the indices afterwards difftpt_utils.py.
• Optimizer. The paper says Adam §4.1. The code uses AdamW and passes no weight decay tpt_classification.py.
• Integer rounding. The paper gives ratios §4.1. The code floors them: 101 of 127 views survive the cosine step and 30 of 102 survive the entropy step difftpt_utils.py difftpt_utils.py.
• Diffusion model. The paper calls SD V2 a text-to-image model and swaps the text embedding for the CLIP image embedding §3.3. The code uses the image-variation checkpoint lambdalabs/sd-image-variations-diffusers, revision v2.0, with guidance 3 and the pipeline's default sampler image_generator_R.py image_generator_R.py.
• Augmented views. The paper says TPT's standard augmentation uses random resized crops §4.2. The code adds AugMix operations only when the dataset name is longer than one character tpt_classification.py tpt_classification.py. For the ImageNet family the augmented views are crop and flip only datautils.py.
• Scripts. Only RN50 runs appear in scripts/do_tpt_difftpt.sh do_tpt_difftpt.sh. The ViT-B/16 numbers in tables 1 and 2 have no command in the repository. A comment in the code still names a ViT-L/14 model tpt_classification.py.
• Few-shot baselines. The TPT&CoOp and TPT&CoCoOp rows need the CoOp or CoCoOp weights, passed with --load tpt_classification.py. The repository does not contain those weights.
• Diffusion views are precomputed. The paper's inference-cost discussion includes generation time §4.3. The released test loop only reads the precomputed _<i> images listed in selected_data_list.txt difftpt_utils.py image_generator_R.py.
• Unused generated views. The generator writes 65 variants per image image_generator_R.py and the loader reads 63 difftpt_utils.py.

Pitfalls

• Printed average does not match its parts. In Table 1 the CLIP RN50 average is printed as 42.26, but its five values (58.10, 22.81, 53.00, 53.90, 33.50) average to 44.26. The OOD average (40.80) does match. Recompute averages from the per-dataset values before comparing.
• Headline gain not reproduced. The abstract and conclusion state an average gain of 5.13% over TPT. §4.2 gives 5.1% (RN50) and 2.3% (ViT-B/16) for Scenario 2. The printed Scenario 2 averages give 2.87 and 1.48 over TPT, and 4.73 and 2.44 over zero-shot CLIP. The paper does not say how the 5.13% figure is computed. Treat it as unverified.
• Backbone attribution in §4.2. The sentence on ViT-B/16 gains gives the pair 61.27 → 64.12 and 61.01 → 61.57 as ResNet-50 results. Those values are the ViT-B/16 rows of Table 1, and the RN50 rows show 47.12 → 50.87 and 46.71 → 49.61 §4.2 tab. 1.
• Subset versus full data. The main tables use 1000 sampled images per dataset §4.1. Appendix A reports full-dataset numbers; its text gives CLIP-RN50 ImageNet-V2 as 51.41, while Table 1 prints 53.00 app. A tab. 1. Do not compare across the two.
• Inference-cost figures conflict. §4.3 gives 6 s for 10 images with TPT, 36 min for standard SD, and 70 s for the original SD to generate 10 images, with the consistency model at 0.5 s §4.3. No hardware or step count is given for these figures, so they cannot be compared directly.
• Selection is not reproducible from code alone. The generator shuffles the candidate images without a seed image_generator_R.py, and its diffusion calls set no seed. The generated images are not in the repository, so the exact 1000-image set and its views must be regenerated.
• Safety checker disabled. The generator replaces the pipeline's safety checker with a pass-through image_generator_R.py. Stated for transparency; the paper does not mention it.
• Model terms not stated. The paper does not state the license or terms of the Stable Diffusion checkpoint used for generation. Check them before redistributing any generated images.
• Entropy term index. Equations 6 and 7 write H(p_i) inside a sum over n. Read it as H(p_n) eq. 4 eq. 7.
• Precision. The code runs under autocast with a gradient scaler at initial scale 1000 tpt_classification.py. The paper does not state a precision.
• Argparse defaults are not the paper's. --tta_steps defaults to 1 and --selection_p to 0.1 tpt_classification.py. Use the script values: 4 steps, ρ_H 0.3, ρ_C 0.8.
• No license file. The repository has no LICENSE file, so code_license is not stated. Do not quote code from it. Read the code and reimplement.
• Other method in the same repository. The IT3A/ folder holds code for a different method, IT3A (its README links arXiv 2412.09706). It is not covered by this card.

Lineage

• Builds on: Test-time prompt tuning (TPT), Shu et al., 2022 ("Test-time prompt tuning for zero-shot generalization"). Its arXiv id is not stated in the paper's source, so no link is given. DiffTPT keeps TPT's augmented views, adds diffusion views next to them, and adds a cosine filter on top of TPT's entropy filter §3.3 §3.4.
• Also uses: CLIP §3.1; CoOp and CoCoOp (Zhou et al., 2022) as baselines and as starting points for TPT&CoOp and TPT&CoCoOp §4.1; Stable Diffusion V2 §3.3.
• Compared against: TPT, CoOp, CoCoOp, zero-shot CLIP with the default prompt and with an 80-prompt ensemble §4.1.
• Superseded by: not stated.