Robust fine-tuning of zero-shot models

1 Introduction

• Fine-tuned robustness varies with small hyperparameter changes, and aggressive fine-tuning can substantially harm shifted accuracy despite improving target accuracy.For linear-classifier fine-tuning, prediction complementarity may explain gains; end-to-end observations suggest linear mode connectivity, although the loss landscape and its relation to shifted-distribution error remain open questions.Source: §1 · pp. 2–3
• WiSE-FT first fine-tunes a zero-shot model, then linearly interpolates its weights with the original model, simplifying hyperparameter choice.The method is weight-space ensembling and incurs no additional fine-tuning or inference cost.Source: §1 · p. 3
• By 4 to 6 percentage points (pp), WiSE-FT improves accuracy under distribution shift over prior work on ImageNet and five natural shifts while maintaining or improving fine-tuned CLIP’s ImageNet accuracy.The result is shown for standard end-to-end fine-tuning.Source: Figure 1 · p. 2, §1 · p. 3
• By 1 to 9 pp, WiSE-FT improves shifted-distribution accuracy over zero-shot models, with no additional computational cost during fine-tuning or inference.It also improves over regularization and evaluating at various points during fine-tuning.Source: §1 · p. 3
• On BASIC-L, WiSE-FT improves ImageNet accuracy by 0.4 pp and average shifted-distribution accuracy by 2 to 11 pp.Similar trends were observed for ALIGN and a ViT pre-trained on JFT.Source: §1 · p. 3
• On six additional shifts, WiSE-FT consistently improves robustness; across seven datasets, it reduces relative error rate by 4-49% versus standard fine-tuning.The shifts include geographic, CIFAR-10 reproduction, and video temporal perturbations; the seven datasets include ImageNet, CIFAR-10, CIFAR-100, Describable Textures, Food-101, SUN397, and Stanford Cars.Source: §1 · p. 3

Experimental setup

• They compare zero-shot, fine-tuned, and WiSE-FT models on reference and shifted test sets, with reference-distribution training data and a goal of high, consistent accuracy.They avoid ID/OOD terminology because zero-shot models are not trained on the reference distribution.Source: §2 · pp. 3–4

Distribution shifts

• The experiments focus on natural ImageNet-derived shifts: ImageNet-V2, ImageNet-R [37], ImageNet Sketch [100], ObjectNet, and ImageNet-A [38].These shifts capture naturally occurring variations rather than synthetic perturbations.Source: §2 · p. 4

Effective robustness and scatter plots

• Effective robustness is shifted accuracy beyond the reference-trained baseline, ρ(f) = Accshift(f) − β(Accref(f)); positive values indicate robustness to distribution shift [97].Scatter plots place reference accuracy on the x-axis and shifted accuracy on the y-axis, separating robustness changes from reference-accuracy changes.Source: §2 · pp. 4–5

Zero-shot models and CLIP

• They primarily study CLIP, also evaluating ALIGN, BASIC, and a JFT-pretrained ViT; unless noted, experiments use CLIP ViT-L/14@336px.CLIP-like models learn image and text encoders from web image-caption pairs and classify by matching images to class-name captions.Source: §2 · p. 5

• WiSE-FT first fine-tunes the zero-shot model on application-specific data, then linearly interpolates the original and fine-tuned weights.The method is motivated by combining the zero-shot model’s strength under distribution shift with standard fine-tuning’s high accuracy on the reference distribution.Source: §3 · pp. 5–6
• Standard fine-tuning can modify all model parameters end-to-end or keep the pretrained encoder fixed and fine-tune only a linear classifier.Source: §3 · p. 6
• For linear-classifier fine-tuning, weight-space interpolation is equivalent to an output-space ensemble; for end-to-end fine-tuning, accuracy remains high along the zero-shot-to-fine-tuned linear weight path.Source: §3 · p. 6
• WiSE-FT improves accuracy under distribution shift while maintaining high reference-distribution performance, using a single set of weights with no additional computational cost.Source: §3 · p. 6

Main results: ImageNet and associated distribution shifts

• WiSE-FT (end-to-end, α=0.5) outperforms numerous strong models in average shifted accuracy and average reference-and-shifted accuracy on ImageNet and five derived shifts.The reported α=0.5 performs close to optimally across experiments and is recommended when no domain knowledge is available.Source: §4 · p. 6, Table 1 · p. 7

Robustness on additional distribution shifts

• Across six additional shifts, WiSE-FT (α=0.5) improves shifted accuracy by 3.5, 6.2, 1.7, 2.1, 9.0, and 23.2 pp, with reference accuracy decreasing by at most 0.3 pp.The shifts cover WILDS-FMoW [49] [13], WILDS-iWildCam [49] [6], CIFAR-10.1, CIFAR-10.2, ImageNet-Vid-Robust, and YTBB-Robust; zero-shot accuracy is below 30% on the WILDS shifts.Source: §4 · pp. 6–7

Hyperparameter variation and alternatives

• A 0.3 pp ImageNet accuracy difference between two 10-epoch learning rates accompanies up to 8 pp variation under distribution shift, while WiSE-FT removes the accuracy trade-off.Raising the learning rate from 10−7 to 3 · 10−5 improves ImageNet accuracy by 5 pp but reduces shifted accuracy by 8 pp; varying α requires no new model training. The best frontier is reached by weight averaging, including EMA [95].Source: §4 · p. 7

Accuracy gains on reference distributions

• On seven datasets, WiSE-FT improves accuracy after fine-tuning on reference distributions, extending its reported benefits beyond robustness to distribution shift.Source: §4 · p. 7

Beyond CLIP

• On BASIC-L [77], WiSE-FT (α=0.5) improves average accuracy on five ImageNet shifts by over 7 pp and ImageNet accuracy by 0.4 pp versus fine-tuning.BASIC-L is fine-tuned with contrastive loss on half the ImageNet training data; WiSE-FT improves both reference and shifted accuracy despite these experimental differences.Source: §4 · p. 8, Figure 4 · p. 9
• For the JFT-300M-pretrained ViT-H/14 [21] [93], WiSE-FT (α=0.8) beats fine-tuning by 2.2 pp on shifts while keeping ImageNet accuracy within 0.2 pp.The zero-shot classifier uses a manually identified class correspondence; the result demonstrates applicability without contrastive image-text pre-training.Source: §4 · p. 8, Figure 4 · p. 9
• WiSE-FT also shows similar trends on ALIGN [45], which resembles CLIP but uses a different pre-training dataset.Source: §4 · p. 9, Figure 4 · p. 9

5.1 Zero-shot and fine-tuned models are complementary

• Zero-shot and fine-tuned models exhibit prediction diversity on reference and shifted distributions despite sharing the same backbone.The paper compares their diversity with that between linear classifiers fine-tuned on different halves of ImageNet.Source: §5.1 · p. 9, Figure 5 · p. 10
• On shifts other than ImageNetV2, zero-shot predictions override fine-tuned predictions more often; zero-shot models are also more confident under shift, while fine-tuned models are more confident on the reference distribution.Confidence is measured by the margin between each classifier’s largest and second-largest output.Source: §5.1 · pp. 9–10, Figure 5 · p. 10

5.2 An error landscape perspective

• On ImageNet and five associated shifts, some end-to-end weight-space ensembles may match or outperform both the zero-shot and fine-tuned models.Figure 6 states that there exists an α for which ensemble accuracy is at least the accuracy of either endpoint.Source: Figure 6 · p. 11, §5.2 · p. 11
• For every α ∈ [0, 1], interpolated weights achieve accuracy at least as high as the linear interpolation of endpoint accuracies on ImageNet and its five associated shifts.This is the accuracy inequality stated as Observation 1.Source: §5.2 · p. 10, Figure 6 · p. 11

Robustness

• Prior work finds synthetic-shift interventions often fail on natural shifts [97], while reference accuracy usually predicts shifted accuracy [106] [97] [70], with exceptions [16] also seen across fine-tuning hyperparameters.The passage connects the exception to D’Amour et al.’s finding that accuracy under certain shifts cannot be inferred reliably from reference accuracy.Source: §6 · p. 11

Pre-training and transfer learning

• Language-supervised vision pretraining enables zero-shot inference [18] [86] [111] [82] [45] [77] [109], but standard fine-tuning can reduce effective robustness at convergence [3].Prior NLP methods for stable fine-tuning incur computational overhead [46] [113]; quadratic regularization was found to work best for fine-tuning, while many continual-learning approaches did not improve robustness to multiple natural shifts [3].Source: §6 · pp. 11–12

Traditional (output-space) ensembles

• WiSE-FT’s two-model weight-space ensemble improves shifted accuracy by 2–15 pp, versus 0.8–1.6 pp for a similar-sized ensemble of models pre-trained on different datasets [72], without extra compute.The comparison covers ImageNetV2, ImageNet-R, ObjectNet, and ImageNet-A.Source: §6 · p. 12

Weight-space ensembles

• Weight-space ensembles linearly interpolate model weights, and this work applies that approach to models trained on different data to study distributional robustness.Prior examples include averaging training checkpoints [43] and averaging fast and slow weights every k steps.Source: §6 · p. 12

Concurrent and subsequent work

• Kumar et al. report an accuracy trade-off for end-to-end versus linear-classifier fine-tuning, while this work starts from a zero-shot classifier; output ensembles can mitigate trade-offs even without calibration.Kumar et al. study calibrated output-space ensembles, while the paper reports uncalibrated mitigation in Appendix Figures 10 and 25.Source: §6 · p. 12
• Hewitt et al. find output ensembles mainly outperform distillation; Gontijo-Lopes et al. [31] find subdomain specialization supports ensemble performance; Matena and Raffel combine weights using Fisher information.Matena and Raffel’s method goes beyond single-coefficient linear interpolation, and their experiments do not examine shifted accuracy; their goal of combining differing expertise aligns with this work.Source: §6 · p. 12

7 Limitations, impact, and conclusion

• The investigation is limited to image classification; object detection and natural language processing are identified as future directions.The authors expect the findings may apply more broadly, but did not investigate those domains here.Source: §7 · p. 12
• Although alpha=0.5 performs well overall, finding the optimal alpha for specific target distributions remains future work.Source: §7 · p. 13
• Because WiSE-FT builds on large zero-shot models, it may perpetuate harms such as model bias and malicious uses including surveillance.The paper attributes discussion of these potential harms to Radford et al. [82] and Brown et al. [12].Source: §7 · p. 13
• WiSE-FT substantially improves performance under distribution shift with minimal or no target-distribution accuracy loss compared with standard fine-tuning.The authors present it as a first step toward more sophisticated fine-tuning schemes that leverage zero-shot model robustness.Source: §7 · p. 13

A Pseudocode for WiSE-FT

• WiSE-FT interpolates corresponding checkpoint weights as (1-alpha) * theta_0[key] + alpha * theta_1[key], then loads the resulting weights into the model.The pseudocode loads zero-shot and fine-tuned state dictionaries before forming the interpolated weights.Source: A Pseudocode for WiSE-FT · p. 21
• The pseudocode initializes the model from zero-shot weights, fine-tunes it on the dataset, then applies weight-space ensembling.The routine calls the interpolation procedure in place after standard fine-tuning.Source: A Pseudocode for WiSE-FT · p. 21

B Mixing coefficient

• The optimal α averages 0 to 0.4 percentage points better than α=0.5 across ImageNet and five derived shifts, but α=0.5 is recommended when domain knowledge is unavailable.Testing different α values requires no retraining, making tuning relatively cheap; finding the optimal value for any distribution is left for future work.Source: B Mixing coefficient · p. 21

C Additional experiments

• WiSE-FT improves effective robustness on six additional distribution shifts without reducing performance on the reference distribution.Section C.2 presents effective robustness scatter plots for these shifts.Source: C Additional experiments · p. 21
• The additional experiments compare WiSE-FT with distillation and CoOp [112].Source: C Additional experiments · p. 21
• Section C.5 examines WiSE-FT accuracy improvements on reference data, focusing on the low-data regime.Source: C Additional experiments · p. 21
• Section C.6 examines whether accuracy improvements under distribution shift persist across scales of pre-training compute.It reports similar trends across scales.Source: C Additional experiments · p. 21
• Section C.7 explores WiSE-FT with ALIGN [45], a ViT-H/14 model pre-trained on JFT [21], and BASIC [77], while C.8 ensembles zero-shot CLIP with an independently trained classifier.Source: C Additional experiments · p. 21
• Section C.1 provides a breakdown of Figure 1 results for each distribution shift.Source: C Additional experiments · p. 21

C.1 Breakdown of CLIP experiments on ImageNet

• Table 3 compares optimal-coefficient WiSE-FT with fixed α=0.5 for two CLIP models, choosing the optimal coefficient separately to maximize each metric.It reports differences in percentage points and includes averages over the five shifts and over ImageNet plus those shifts.Source: Table 3 · p. 22
• The scatter-plot context includes standard-trained models, models trained with more data, and models using robustness interventions, including augmentation and adversarial robustness [97] [70].The cited examples of interventions include special data augmentation [19] [28] [36] and adversarially robust models [65] [15] [85] [87].Source: §C.1 · p. 22
• WiSE-FT improves accuracy on ImageNet and five derived distribution shifts, with results displayed separately for each shift.Figures 7 and 8 present the per-dataset breakdown of the key ImageNet results.Source: §C.1 · p. 21, Figure 7 · p. 22, Figure 8 · p. 23
• Tables 4 and 5 report WiSE-FT accuracy across mixing-coefficient values for CLIP ViT-L/14@336 and ViT-B/16, respectively.α=0.0 is the zero-shot model and α=1.0 is standard fine-tuning; the tables include ImageNet and five derived shifts.Source: §C.1 · p. 22, Table 4 · p. 24, Table 5 · p. 25

C.2 Robustness on additional distribution shifts

• WiSE-FT improves accuracy under distribution shift relative to standard fine-tuning on six additional datasets.The datasets are ImageNet-Vid-Robust, YTBB-Robust, CIFAR-10.1, CIFAR-10.2, WILDS-FMoW, and WILDS-iWildCam.Source: Figure 9 · p. 26
• The six additional shifts cover temporal video perturbations, shifted CIFAR-10 reproductions, geographic and temporal satellite imagery, and geographic wildlife imagery.ImageNet-Vid-Robust and YTBB-Robust involve temporal perturbations; WILDS-FMoW shifts geographically and temporally, while WILDS-iWildCam shifts geographically.Source: §C.2 · p. 26

C.3 Comparison with alternative methods

• The alternative-method comparisons include six shifted datasets, with linear-classifier fine-tuning for ImageNet-Vid-Robust and YTBB-Robust and end-to-end fine-tuning for the others.The comparisons contrast weight-space and output-space ensembling, EMA, and additional linear-classifier methods.Source: §C.3 · p. 26, Table 6 · p. 27

C.3.1 Output-space ensembles

• The output-space ensemble has a favorable trend like weight-space ensembling but requires twice as much compute.The comparison uses end-to-end fine-tuning with learning rate 3·10−5.Source: §C.3.1 · p. 27, Figure 10 · p. 27

C.3.2 Comparison to exponential moving averages

• EMA decay can move solutions along or slightly outside the WiSE-FT curve, and applying WiSE-FT to EMA solutions can yield further robustness.The initialization-biased EMA follows the WiSE-FT curve, while EMA evaluations along the trajectory underperform solutions generated by applying WiSE-FT.Source: Figure 12 · p. 28, Figure 11 · p. 28, §C.3.2 · p. 29

C.3.3 Additional comparisons when fine-tuning a linear classifier

• WiSE-FT matches methods requiring more compute or performs better in effective-robustness comparisons, while many examined approaches show concave trends.The comparisons fine-tune only a linear classifier and are shown in Figure 13.Source: §C.3.3 · p. 29, Figure 13 · p. 29
• Linear-classifier comparisons include random interpolation, softmax-output ensembling, regularized classifiers, distillation, and regularization toward zero-shot weights.Softmax-output ensembling performs comparably to weight-space ensembling but requires slightly more compute; regularization toward zero-shot weights performs slightly worse than distillation in most cases.Source: §C.3.3 · pp. 29–30

C.4 Changes in data augmentation

• Using default ImageNet training augmentation—random cropping, horizontal flipping, and color jitter—produces trends similar to those with minimal augmentation.The paper identifies further investigation of augmentation effects as future work.Source: §C.4 · p. 30, Figure 14 · p. 31

C.3.3 Additional comparisons when fine-tuning a linear classifier

• Table 7 and Figure 14 compare WiSE-FT with CoOp [112]; both methods fine-tune the ViT-B/16 CLIP model using 16 examples per class of ImageNet.Source: Table 7 · p. 30, Figure 14 · p. 31

C.5 Accuracy improvements on reference datasets

• WiSE-FT improves accuracy on ImageNet and six transfer datasets, including CIFAR-10/100, Describable Textures, Food-101, SUN397, and Stanford Cars.In low-data linear-classifier experiments, for k = 1 the zero-shot model outperforms the fine-tuned classifier and smaller α yields high performance; with more data, higher α improves performance.Source: §C.5 · pp. 31, 33, Figure 16 · p. 32, Figure 17 · p. 33

C.6 Robustness across scales of pre-training compute

• Across orders of magnitude of pre-training compute, WiSE-FT improves shifted-distribution accuracy with minimal reference-performance loss, with similar trends for the smaller CLIP ViT-B/16.Figure 20 considers reference-accuracy losses below ε ∈ {0, 0.1, 1} pp; the larger-model experiments otherwise use CLIP ViT-L/14@336px.Source: §C.6 · p. 36, Figure 20 · p. 36, Figure 21 · p. 36

C.7 WiSE-FT and additional models

• Table 8 summarizes WiSE-FT results for CLIP, ALIGN, BASIC, and a ViT model pre-trained on JFT.The table reports ImageNet and derived-shift accuracy for models fine-tuned end-to-end.Source: §C.7 · p. 37, Table 8 · p. 37

C.7.1 ALIGN

• WiSE-FT is effective for ALIGN, with results reported for end-to-end and linear-classifier fine-tuning across mixing coefficients.The end-to-end model is trained with AdamW for 40,000 steps; the linear classifier uses L-BFGS.Source: §C.7.1 · p. 37, Figure 22 · p. 38, Table 9 · p. 39

C.7.2 JFT pre-training

• The JFT-300M experiments construct zero-shot ImageNet weights for ViT-H/14 by mapping 973 of 1,000 ImageNet classes to JFT classes.They also evaluate linear-classifier and end-to-end fine-tuning, with results summarized for ImageNet and five derived shifts.Source: §C.7.2 · p. 40, Table 10 · p. 41

C.7.3 BASIC

• For BASIC-L, WiSE-FT improves accuracy relative to the fine-tuned model on ImageNet and five derived distribution shifts.The zero-shot classifier uses ImageNet class names, and BASIC fine-tuning uses half of the ImageNet training data.Source: Figure 24 · p. 42, §C.7.3 · p. 42, Table 12 · p. 44

C.8 Ensembling zero-shot CLIP with independently trained models

• 13.6 pp and 6.9 pp gains in average accuracy under distribution shift came from ensembling zero-shot CLIP with Noisy Student EfficientNet-B6 and EfficientNet-L2, respectively, without reducing reference accuracy.Because these models can have different architectures, the experiments use output-space rather than weight-space ensembling, increasing inference cost.Source: §C.8 · p. 45, Figure 25 · p. 45

D.1 CLIP zero-shot

• CLIP inference uses ViT-L/14@336px by default, factors its learned temperature into the zero-shot weights, and ensembles 80 prompts except for five datasets with manually engineered prompts.The exceptions are WILDS-FMoW, WILDS-iWildCam, Stanford Cars, Describable Textures, and Food-101.Source: §D.1 · p. 45

D.2 End-to-end fine-tuning

• End-to-end fine-tuning initializes the final layer from CLIP’s zero-shot classifier, scales its weights by the pretrained temperature, omits a fine-tuning temperature, and continues normalizing image-encoder outputs.Source: §D.2 · p. 46
• End-to-end fine-tuning uses AdamW [76], default hyperparameters, a 3 × 10−5 learning rate, 10 epochs, and cosine annealing with 500 warm-up steps unless noted otherwise.It uses the largest batch fitting on 8 GPUs (512 for ViT-B/16), gradient clipping at global norm 1, and the augmentations in [82], with crop dimensions of 336 pixels for ViT-L/14@336px and 224 for other models.Source: §D.2 · p. 46

D.3 Fine-tuning a linear classifier

• Linear-classifier training evaluates 64 configurations for each of four regularization strategies, randomly varying batch size, learning rate, and strategy-specific coefficients; label smoothing follows [71].The classifier selected for ensembling attains the best in-distribution performance, and its hyperparameters are reused in distillation and regularization experiments.Source: §D.3 · pp. 46–47

D.4 ObjectNet

• 2.3% [82] performance improvement on ObjectNet comes from adapting the zero-shot classifier to ObjectNet class names; Figure 26 compares weight-space ensembles with and without this adaptation.The zero-shot models in Table 1 use ImageNet class names; ObjectNet is the only one of the five Section 3 datasets where class-shift adaptation is possible.Source: Figure 26 · p. 47, §D.4 · p. 47

D.3 Fine-tuning a linear classifier

• For low-data linear-classifier training, the epoch cap scales inversely with removed data, is doubled until performance saturates, and each run samples its epoch count uniformly from 1 through that cap.The procedure uses default PyTorch AdamW hyperparameters with learning rate 0.001 and weight decay 0.01.Source: §D.3 · p. 47

E Diversity measures

• Prediction Diversity measures the fraction of samples on which two classifiers disagree while one classifier is correct.Source: E Diversity measures · p. 47
• Cohen’s Kappa Complement uses the complement of Cohen’s agreement coefficient as a diversity measure between classifiers.The coefficient compares empirical agreement with expected agreement.Source: E Diversity measures · p. 48
• KL diversity is the average Kullback-Leibler divergence between the classifiers’ predicted class-probability distributions across samples.Source: E Diversity measures · p. 48
• Centered Kernel Alignment Complement measures representation diversity by taking the complement of CKA, for which larger similarity indicates smaller diversity [51].Source: E Diversity measures · p. 48
• For datasets larger than 10,000 samples, CKAC is computed on a random sample of 10,000 because the measure is computationally expensive.Source: E Diversity measures · p. 48
• The diversity comparisons are extended across combinations of measures, datasets, and CLIP models, using same-encoder linear classifiers trained on separate halves of ImageNet as baselines.Results are shown in Figures 27–30.Source: E Diversity measures · p. 48

F When do weight-space ensembles approximate output-space ensembles?

• When the model is linear along the path between two parameter settings, their weight-space and output-space ensembles are equivalent.The section notes that this approximation can apply in the NTK regime, and that later training—closely resembling fine-tuning—can be better described by a linear expansion [44] [24].Source: F When do weight-space ensembles… · p. 51