Patching open-vocabulary models by interpolating weights

In plain words

CLIP-style models can label images with any list of class names you type, but they are weak on some specific jobs, such as reading digits or traffic signs. The usual fix is to fine-tune the model on that job. That helps the job, but the model often gets worse at everything else it could do before.

This paper's method, called PAINT, is a two-step recipe. First, fine-tune the model on the weak task. Second, mix the old and new weights with a single dial between 0 and 1. A low setting stays close to the original model, a high setting stays close to the fine-tuned one, and something in between can keep most of the new skill while losing little on the old tasks. No new parameters are added, so the model still accepts any class names.

The authors report large gains on nine weak tasks while ImageNet accuracy barely moves. They also find that patching on one task can help a related task with different class names, and that the method works better on larger models.

TL;DR

• Two steps: fine-tune the visual encoder on the patch task, then interpolate between the zero-shot and fine-tuned weights with alpha §2.
• No task-specific parameters: the classification head comes from the frozen CLIP text encoder §3 app. C.
• ViT-L/14 gains 15 to 60 points on nine weak tasks, with ImageNet down less than 1 point (author-reported) §4 §1.
• Broad transfer: patching on one task raised accuracy on a different task with disjoint classes in 7 of 8 pairs §6 tab. 2.
• The authors report no extra cost at fine-tuning or inference time versus standard fine-tuning §1.

Problem

Zero-shot open-vocabulary models are often far from optimal on specific classification tasks §1. Fine-tuning fixes the target task but can cause catastrophic forgetting on tasks where the zero-shot model was already adequate, and a fine-tuned model with a fixed head loses the free-form, text-defined class interface §1. Re-training an open-vocabulary model from scratch with the new data is the alternative, and the authors call it practically infeasible in most settings §1. The goal is to improve accuracy on a patch task without reducing accuracy on a supported task such as ImageNet §1 §3.

The paper assumes labelled data for the patch task, and a held-out validation set for the supported task and the patch task, used to choose the mixing coefficient §2.

Method

Patching on one task §2. Let the zero-shot weights be theta_zs and the patch task be D_patch. A representative supported task D_supp is used for checking, not for training §2.

Step 1 fine-tunes theta_zs on training data from D_patch to get theta_ft. Step 2 interpolates: theta_patch = (1 - alpha) * theta_zs + alpha * theta_ft, with alpha in [0, 1] §2, steps 1 and 2. Alpha is chosen on held-out validation sets of D_supp and D_patch §2. The interpolation applies to the visual encoder only; the text encoder and the head stay fixed §3 app. B app. C.

Frozen head. Instead of a learned classifier, the head is the output of CLIP's text tower for each class name, so no class-specific parameter is added §3 app. C. App. C reports that freezing it has a negligible effect on accuracy for MNIST, EuroSAT and SUN397, over three seeds app. C.

Patching on several tasks §2 §5. The paper compares three strategies:

• Joint: merge all patch datasets into one task, then patch once §5.
• Sequential: patch on each new task in turn, starting each step from the last patched model §5.
• Parallel: fine-tune each task separately from theta_zs, then mix. The authors search one scalar alpha against the average of the fine-tuned models, because searching one alpha per task was impractical §5.

Joint patching needs data for every task at once. Sequential patching fits a stream of tasks. Parallel patching can use distributed hardware §5.

Broad transfer §6. A dataset is split into two disjoint sets of classes, task A and task B. Patching on A is then tested on B §6. Patching on one dataset and testing on a different dataset with related content is also tested §6 tab. 2.

Scale §4.1. The paper measures how far a patched model is from the best interpolation, with accuracy distance to optimal (eq. 1), and finds the gap shrinks as models grow §4.1 fig. 2. Representations and weights of the zero-shot and fine-tuned models become more similar at larger scale §4.1 fig. 2 app. D.

# PAINT: patch one task by interpolating weights  [§2]
# Input:  zeroshot_weights (visual encoder), patch_data, alphas
# Output: one patched encoder per candidate alpha

def patch(zeroshot_weights, patch_data, alphas):
    # Step 1: fine-tune on the patch task, head frozen  [§2, step 1]
    finetuned_weights = finetune(zeroshot_weights, patch_data)
    patched_models = {}
    for alpha in alphas:
        # Step 2: interpolate every tensor  [§2, step 2]
        # [code: src/patch.py#L48-L51]
        patched_weights = {
            name: (1 - alpha) * zeroshot_weights[name]
                  + alpha * finetuned_weights[name]
            for name in zeroshot_weights
        }
        patched_models[alpha] = patched_weights
    # Choose alpha on held-out validation of supported and patch tasks  [§2]
    return patched_models

# Fine-tuning used in step 1  [§3] [app. C] [code: src/finetune.py#L18-L125]
# Input:  encoder (CLIP visual tower), head (frozen, from text encoder), data
# Output: fine-tuned encoder

def finetune(encoder, head, data):
    # Only encoder parameters are trained; the head is frozen  [§3]
    optimizer = adamw(encoder.parameters(), lr=1e-5, weight_decay=0.1)
    # Linear warm-up, then cosine decay to zero
    # [§3] [code: src/utils.py#L19-L31]
    scheduler = cosine_with_warmup(optimizer, warmup_steps=200)
    for step in range(2000):  # "2000 iterations", batch size 128  [§3]
        images, labels = next_batch(data, batch_size=128)
        # Normalised features times frozen class weights
        # [code: src/modeling.py#L51]
        logits = head(encoder(images))
        loss = cross_entropy(logits, labels)
        loss.backward()
        # Clip gradient norm at 1.0 (not in the paper)
        # [code: src/finetune.py#L95]
        clip_grad_norm(encoder.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
    return encoder

Hyperparameters

Fine-tuning and interpolation values. Where the paper and the code differ, the row says so; see Paper vs code.

Results

Author-reported. Setup: CLIP ViT-L/14 unless stated; combined accuracy is the mean of supported and patch-task accuracy §3.

Per-dataset numbers for each model and alpha are in app. O (single-task tables), with multi-task per-dataset numbers in app. J.

Paper vs code

Official repository: https://github.com/mlfoundations/patching, linked in the paper's footnote §7. Code commit 3ca86220b926849ea5539531722ba7143f67a30e. The code covers single-task patching only: src/patch.py takes one --train-dataset patch.py. I found no joint, sequential or parallel multi-task code in src/.

Differences:

• Warm-up. Paper: 200 steps §3. Code default: 500 args.py. The README sample command does not override it README.md.
• Length. Paper: 2000 iterations §3. Code: --epochs, default 10 args.py. The README sample uses 5 epochs on MNIST README.md. The step count therefore depends on the dataset size and is not 2000.
• Learning rate. Paper: 1e-5 §3. Code default: 1e-3 args.py. The README sample uses 1e-5, which matches the paper.
• Gradient clipping. Code clips the global norm at 1.0 finetune.py. The paper does not mention clipping.
• Alpha selection. Paper: alpha is picked on held-out validation sets §2. patch.py evaluates every alpha in its --alpha list and saves each model patch.py. It contains no selection step, so choosing alpha is left to the user.
• Head. Paper: the head is the frozen output of CLIP's text tower §3 app. C. Code builds it from several prompt templates per dataset, averages them, normalises, and multiplies by the exponential of CLIP's logit scale heads.py. The paper does not say which prompts it used for the head.
• Checkpoints. Code saves whole modules with pickle utils.py. Loading them requires the same code and classes.

No other discrepancies found in the parts of the code I read (src/modeling.py, src/finetune.py, src/patch.py, src/utils.py, src/heads.py).

Pitfalls

• Alpha is a per-model choice. The paper picks alpha on validation data. In the Avg column of the single-task tables in app. O, ViT-B/32 peaks around alpha 0.5 to 0.6 and ViT-L/14 around 0.65 to 0.70. Treat these as examples and choose alpha on your own validation data.
• Interpolation covers every tensor of the saved encoder patch.py. Whatever fine-tuning changes, such as norms or biases, is blended the same way.
• Sequential patching still forgets. Interpolation reduces forgetting but does not remove it §5 §9. Forgetting can compound across steps.
• Parallel patching needs a search. Black-box search for the per-task alphas gave little gain over one shared alpha. The authors' exhaustive search over two tasks showed headroom for better search §5 app. J.
• Small models lose more. The paper reports larger drops on supported tasks for smaller models §9; ViT-B/32 sequential patching without interpolation lost 4.6 points §5.
• Datasets without public test sets. The validation set is used as the test set, and the training set is split stratified app. A. Compare numbers only on the same split.
• Head depends on class names. The head is built from text prompts for each class list. A new class list needs a new head, computed with the same frozen text encoder app. C heads.py.
• Cost. The method adds an interpolation step to standard fine-tuning, and the paper reports no extra cost at fine-tuning or inference time §1. The paper's full experiments used about 10 thousand GPU-hours app. N; it does not give a per-patch figure.
• Typographic data. The synthetic attack data is built from SUN397 with text overlays app. K. The real-world set has 110 images §7 and is released in the repository as a zip file typographic_attack_data.zip.

Lineage

• Builds on: linear weight interpolation for robustness to distribution shift, Wortsman et al. (2021), cited in §1 as the basis for mixing fine-tuned and original weights. No card in this wiki.
• Compared against: early stopping, regularisation towards the initialisation, lower learning rates, EMA, EWC, LwF, re-training with patch data, and objective mixing §4.2 app. E.
• Superseded by: not stated.