Finetune like you pretrain: Improved finetuning of zero-shot vision models

In plain words

Image-and-text models such as CLIP learn to match pictures with captions. You can use one to classify pictures by writing a short text for each class and picking the class whose text best matches the picture. To get better accuracy on your own labelled pictures, the usual route is to retrain part or all of the model with a classification loss. That can make the model better on your data but worse on pictures that look a bit different, and small changes to the recipe can shift the result a lot.

This paper's idea is to keep the training recipe the same as the original pretraining. Each labelled picture is paired with a sentence built from its class name, and both the picture side and the sentence side are trained with the same matching loss that the model learned from. Afterwards the model is used exactly like a zero-shot classifier: write the class sentences, embed them, and pick the closest. The authors report that this beats several standard finetuning alternatives across many benchmarks, and that it also works when only a handful of labelled examples per class are available.

TL;DR

FLYP (Finetune Like You Pretrain) continues CLIP-style contrastive training on labelled data: each label becomes a text prompt, and both encoders are updated with the same symmetric InfoNCE-style loss as in pretraining. Inference is zero-shot classification with the finetuned encoders. The authors report gains over linear probing, full finetuning, L2-SP and LP-FT on ImageNet distribution shifts, WILDS-iWILDCam and FMoW, transfer tasks, and 4/16/32-shot binary tasks. Results are author-reported, and most runs use one OpenAI CLIP ViT-B/16 setting.

Problem

Zero-shot classifiers of CLIP type need no training and are robust to many kinds of distribution shift. Supervised finetuning usually raises in-distribution accuracy but can lower out-of-distribution accuracy, and small choices in the finetuning procedure change the outcome a lot. Prior fixes (two-stage LP-FT, and weight ensembling in WiSE-FT) modify the usual cross-entropy recipe, and there is no agreed "correct" modification §1. The assumptions: labelled data for the target task, class names available as text, and a pretrained image-text model §1 §2.

Method

The problem is set up as classification of an image I into a label y. A pretrained image encoder and text encoder embed images and text into a shared space, with l2-normalised features §2.

The contrastive pretraining loss matches each image to its own caption and away from the other captions in the batch, in both directions. CLIP uses this objective §2 eq. 1.

Step 1: pair labels with text. For each labelled sample (I, y), draw a text description T from a uniform distribution over the text descriptions of class y. The descriptions are the CLIP prompt templates (for example "a photo of a small {class}") §3 alg. 1. The experiments use the same template sets as CLIP and WiSE-FT §4. Each batch then has the same size as the labelled batch alg. 1.

Step 2: update both encoders. Apply the same pretraining loss (eq. 1) to the image-text batch, and update the image and text encoder parameters by a gradient step, with learning rate alpha §3 alg. 1. The text encoder is trained; standard finetuning updates only the image side §3.

Inference. Build a zero-shot head from the finetuned text encoder: embed each class's prompts, average the normalised embeddings over the templates, and normalise again. Predict the class with the highest cosine similarity to the normalised image embedding §3 zeroshot.py.

Weight ensembling (an optional add-on used for the main tables). Linearly interpolate the zero-shot weights and the finetuned weights, theta_we = alpha * theta_ft + (1 - alpha) * theta_zs, with alpha in [0, 1] eq. 2. The paper uses 10 mixing coefficients and selects the one with the best in-distribution validation accuracy §4.

# FLYP: one training step on a labelled batch  [alg. 1]
# Input:  images (B, 3, H, W); labels (B,); templates[y] per class
# Output: updated image and text encoder parameters

def flyp_step(images, labels, templates, model, optimizer):
    # 1. Turn each label into a caption, drawn uniformly from its templates
    #    [§3] [alg. 1, step 1]
    captions = [sample_template(templates[y]) for y in labels]

    # 2. Embed both sides and l2-normalise  [§2]
    image_emb = normalise(model.encode_image(images))  # (B, d)
    text_emb = normalise(model.encode_text(captions))  # (B, d)

    # 3. Learned scale on the similarities. Eq. 1 has none; the code uses
    #    the model's trained logit_scale  [code: clip/model.py#L378]
    scores = model.logit_scale * image_emb @ text_emb.T  # (B, B)

    # 4. Matching pairs sit on the diagonal. Cross-entropy over rows
    #    (image to text) and over columns (text to image)  [§2] [eq. 1]
    targets = arange(B)
    loss_image = cross_entropy(scores, targets)
    loss_text = cross_entropy(scores.T, targets)
    # The code averages the two terms; eq. 1 sums them
    #    [code: clip/loss.py#L207-L208]
    loss = (loss_image + loss_text) / 2

    # 5. Gradient step on both encoders  [alg. 1, step 2]
    optimizer.zero_grad()
    loss.backward()
    # The code uses AdamW here, not plain SGD  [code: src/models/flyp_loss.py#L69]
    optimizer.step()
    return loss

# Zero-shot-style inference with the finetuned encoders  [§3]

def build_head(class_names, templates, text_encoder):
    # Average the normalised prompt embeddings per class, then renormalise
    #    [code: src/models/zeroshot.py#L47-L66]
    head = []
    for name in class_names:
        prompts = [t(name) for t in templates]
        embs = normalise(text_encoder(prompts))  # (T, d)
        head.append(normalise(embs.mean(dim=0)))  # (d,)
    return stack(head)  # (K, d)

def predict(image, head, image_encoder):
    # Highest cosine similarity wins  [§3]
    feature = normalise(image_encoder(image))  # (d,)
    return argmax(head @ feature)  # (K,) scores

Equation 1, the pretraining loss, in the paper's form eq. 1:

\mathcal{L}(D,\theta)=\sum_{i=1}^{B}-\log\frac{\exp(\hat f(I_i)\cdot\hat g(T_i))}{\sum_{j=1}^{B}\exp(\hat f(I_i)\cdot\hat g(T_j))}+\sum_{i=1}^{B}-\log\frac{\exp(\hat f(I_i)\cdot\hat g(T_i))}{\sum_{j=1}^{B}\exp(\hat f(I_j)\cdot\hat g(T_i))} eq. 1

Here f and g are the image and text encoders, and hats mark l2 normalisation §2. The paper gives no separate temperature term in this equation.

Few-shot binary tasks use the same loss. With k examples per class and two classes, the training batch is all 2k examples, taken as one step per epoch §4 flyp_loss_few_shot.py.

Hyperparameters

Scripts in the repository set these values per dataset; the paper lists the sweep ranges but not the final chosen values app. B README.md.

Results

Author-reported. Setup: CLIP ViT-B/16 from OpenAI unless stated; finetuned on the named ID dataset; OOD numbers average over shifts as stated; few-shot accuracy is on the test set §4. The tables below omit the paper's 95% confidence intervals, which the paper shows in parentheses §4 tab. 1.

Main results without weight ensembling tab. 1:

With weight ensembling (alpha chosen on ID validation) tab. 1:

Large model, WILDS-iWILDCam with OpenAI ViT-L/14 at 336px, with ensembling tab. 2:

Few-shot binary classification, OpenAI ViT-B/16, test accuracy tab. 3:

Transfer learning, full finetuning, ID accuracy tab. 4:

Binary few-shot with OpenAI ViT-L/14 tab. 6:

Few-shot ImageNet shifts (average OOD over 5 shifts, with weight ensembling): FLYP reports gains of 1.5 points over LP-FT at 4-shot and 0.8 at 16-shot fig. 3 §4.2.2. ImageNet per-shift tables are in tab. 7.

Ablations (author-reported, ViT-B/16) §5 app. A:

• Replacing the contrastive loss with cross-entropy and a text-embedding head (the paper's "CE" variant) is worse on ImageNet and iWILDCam §5 fig. 2.
• Freezing the text encoder while keeping the contrastive loss is worse on ImageNet and on iWILDCam (drop of 2.5 ID and 1.3 OOD) §5 fig. 2.
• Masking same-class pairs in the batch lowers PCam accuracy by 1.3 points §5.
• A single template instead of 80 leaves ImageNet ID and OOD unchanged §5 fig. 4.
• Adding cross-entropy to the contrastive loss lowers ID and OOD accuracy on iWILDCam, and lowers OOD on FMoW app. A tab. 5 fig. 5.
• Half batch size: ImageNet ID similar, OOD about 1 point lower; iWILDCam similar; FMoW OOD about 0.5 lower app. A.

Paper vs code

Searched: the paper links https://github.com/locuslab/FLYP, which is the official repository (owner locuslab; Zico Kolter is an author). Cloned at code_commit 215d5bb6feeda6675f60e5818abcb4f6465c83af. The repository carries an MIT license LICENSE.

• Loss scale. Eq. 1 sums the two cross-entropy terms. The code averages them eq. 1 loss.py.
• Temperature. Eq. 1 has no scale on the dot products. The code multiplies by the model's trained logit scale, initialised to log(1/0.07) and not clamped, and AdamW updates it with the rest of the parameters model.py model.py flyp_loss.py.
• Optimizer. Alg. 1 writes a plain gradient step alg. 1. The experiments and code use AdamW with weight decay and cosine decay §4 flyp_loss.py.
• Epoch definition. The epoch length is the number of batches in the labelled training loader. Batches for the update are drawn from a separate shuffled image-caption loader flyp_loss.py flyp_loss.py.
• ImageNet captions. The caption file pairs every image with all 80 templates, so each image appears 80 times. Sampling is over those rows imagenet_csv_creator.py. The paper describes sampling from a distribution over templates §3. The two agree in expectation for ImageNet, but an epoch of batches covers only part of the expanded set. Caption files for the other datasets were not reviewed.
• Classifier at evaluation. Each epoch rebuilds the zero-shot head from the current text encoder, using the dataset's templates, and evaluates it flyp_loss.py. This matches the paper's inference description §3.
• Model selection. The paper says all methods are early-stopped on in-distribution validation accuracy §4 app. B. The main script evaluates each epoch, writes stats and saves a checkpoint per epoch, and does not pick one flyp_loss.py flyp_loss.py. The few-shot script picks the epoch with the lowest validation cross-entropy, not the highest accuracy, and it includes the zero-shot model (epoch -1) as a candidate flyp_loss_few_shot.py flyp_loss_few_shot.py.
• Weight ensembling. The repository has a --alpha argument but no interpolation code, so the WiSE-FT results cannot be reproduced from it args.py. The text search over src/ and datacreation_scripts/ found no other ensembling code.
• Hyperparameters. The paper gives sweep grids app. B. The repository gives one fixed setting per dataset in the README README.md. The paper does not say which sweep value it reports.
• Warmup and min LR. The code uses 500 warmup steps and a minimum LR of 0 by default args.py args.py. The paper does not state warmup.
• Model variants. --model=ViT-B/16 loads OpenAI CLIP. The LAION-400M ViT-B/16 needs ViT-B-16, which loads open_clip weights modeling.py. The ViT-L/14@336px setting for iWILDCam is not in the README.
• Training transform. The OpenAI train transform uses a random resized crop (scale 0.9 to 1.0) rather than the eval resize-and-crop clip.py clip.py. The paper does not mention augmentation.
• Loss helpers. ClipLoss has masking and supervised variants loss.py loss.py. The FLYP path calls it without labels, so those branches are unused flyp_loss.py.
• iWILDCam metric. The paper reports macro F1 §4. The code returns the first value from the dataset's eval() call iwildcam.py. The WILDS function it calls is not in this repository, so the metric was not verified here.

Pitfalls

• Pairing choice. Each batch can hold several samples of the same class. The loss treats them as negatives. The paper tests masking those pairs and finds it worse, not better §5. Do not "fix" it without rerunning the comparison.
• Text side is trained. Standard finetuning keeps the text encoder fixed, so FLYP changes two networks. Freezing the text encoder is one of the paper's ablations and loses accuracy §5.
• Learning rate and weight decay. The paper gives sweep ranges but not the chosen values for each main result. The repository's README values apply to its own scripts only app. B README.md. Low LR (1e-5 or 1e-6) is used throughout the scripts.
• Learned logit scale and weight decay. The logit scale is an ordinary parameter with weight decay in the code, and it is not clamped. The paper does not state this, so an implementation that fixes the temperature differs from the released code.
• Batch size sensitivity. The paper reports that halving the batch changes OOD by about 1 point on ImageNet and about 0.5 on FMoW app. A. Contrastive losses depend on batch size in general, as the paper notes §5.
• Few-shot selection. Few-shot runs pick epochs by validation cross-entropy on a small validation set flyp_loss_few_shot.py. Variance is large: the paper reports 50 repeated draws §4.
• Expensive stacks. The paper does not report compute, memory or wall-clock cost. Training updates both encoders at each step, which is more compute per step than training the image encoder alone.
• Numbers that disagree inside the paper. Several figures in the text, captions and tables do not match. Use the table values, and check any number you copy.
  • The abstract says gains over LP-FT are "more than 1%" ID and OOD averaged over 7 shifts. Table 1 gives ImageNet ID gains of 0.8 points over LP-FT in both settings, and an OOD gain of 1.4 points over LP-FT with ensembling. The 7-dataset averages are not in any table §1 tab. 1.
  • Section 4 says FLYP beats the most competitive baseline on CalTech101 by 1.2%. Table 4 gives 97.6 against 97.2 for full finetuning, the best baseline there, which is 0.4 points §4.3 tab. 4.
  • ImageNet with ensembling: Table 1 gives FLYP OOD 63.2. The appendix per-shift table gives an average of 63.1 for the same setting tab. 1 tab. 7.
  • ImageNet without ensembling: Table 1 gives LP-FT OOD 60.5. The appendix gives 60.3 tab. 1 tab. 7.
  • The appendix ablation table lists FLYP on iWILDCam at 53.03 ID and FMoW at 68.56 ID. The main table gives 52.2 and 68.6 for the same method tab. 5 tab. 1.
  • The Figure 2 caption gives the LP-FT OOD gain as 1.3 points. The text in the same section gives 1.4 fig. 2 §4.1.
• Run-to-run variance. Most intervals are 0.1 to 0.6 points, and some few-shot intervals are above 4 points tab. 1 tab. 3. Small gaps (for example FLYP against LP-FT on ImageNet OOD without ensembling) are within noise tab. 1.
• No win everywhere. Without ensembling, FLYP's ImageNet OOD (60.2) is below LP-FT's (60.5), and the paper calls the two similar tab. 1 §4.1.

Lineage

• Builds on: CLIP's contrastive pretraining objective §2 (the CLIP paper is not yet a card). Standard finetuning and linear probing as baselines §4.
• Related prior work: LP-FT, which the paper treats as the current state of the art §1 (see LP-FT); WiSE-FT, whose weight interpolation is used as an optional add-on eq. 2 §1 (no arXiv id in the paper's sources).
• Compared against: L2-SP, full finetuning, linear probing, LP-FT, WiSE-FT and Model Soups on the WILDS leaderboard §4 tab. 1 tab. 2.
• Superseded by: not stated in the content.