PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Autoencoders enable Robust and Interpretable Fine-tuning of CLIP models

Fabian Morelli, Arnas Uselis, Ankit Sonthalia, Seong Joon Oh

arXiv:2605.15961v1cs.CV

TL;DR

Fine-tuning CLIP can improve downstream performance while degrading robustness under distribution shifts, and existing regularizers offer limited semantic control. SAE-FT regularizes vision representations within a pretrained SAE feature span, preserving and re-weighting interpretable concepts; it achieves strong robustness and downstream generalization, while remaining limited to vision-side regularization.

  • Problem

    CLIP fine-tuning can improve in-distribution performance while reducing robustness under distribution shifts, and standard regularization does not distinguish semantic from incidental features.

  • Method

    SAE-FT constrains vision-encoder updates to the interpretable feature span of a pretrained Sparse Autoencoder and preserves existing semantic concepts rather than overwriting them.

  • Results

    SAE-FT achieves state-of-the-art robustness on distribution shifts and superior downstream-task generalization compared with existing methods, especially vision-encoder-only methods.

  • Takeaways & Limitations

    SAE-FT provides fine-grained, interpretable control over feature changes while re-weighting pretrained concepts and avoiding their erasure.

  • Takeaways & Limitations

    The method regularizes only the vision modality; future work could extend SAE-based regularization to both vision and text encoders.

Abstract

from arXiv · show

Large-scale pre-trained vision-language models like CLIP demonstrate remarkable zero-shot performance across diverse tasks. However, fine-tuning these models to improve downstream performance often degrades robustness against distribution shifts. Recent approaches have attempted to mitigate this trade-off, but often rely on computationally expensive text-guidance. We propose a novel method for robust fine-tuning, SAE-FT, which operates only on the model's visual representations. SAE-FT regularizes changes to these representations by penalizing the addition and removal of semantically meaningful features identified by a Sparse Autoencoder trained on the pre-trained model. This constraint prevents catastrophic forgetting and makes the fine-tuning process interpretable, enabling direct analysis of semantic changes. SAE-FT is both mechanistically transparent and computationally efficient, matching or exceeding state-of-the-art performance on ImageNet and its associated distribution shift benchmarks. Code is publicly available at: https://github.com/Fabian-Mor/sae-ft.

1 Introduction

CLIP fine-tuning can improve downstream in-distribution performance while reducing robustness under distribution shifts. SAE-FT addresses this trade-off by constraining visual updates to interpretable pretrained features and achieves strong robustness and transfer performance.

  • Fine-tuning improves in-distribution performance but often decreases out-of-distribution performance under distribution shifts.
  • SAE-FT adapts classification by amplifying discriminative features such as “ladder” and reducing shared features such as “red paint.”
  • SAEs decompose dense CLIP representations into sparse, semantically meaningful features, enabling finer control over which concepts change.
  • SAE-FT constrains vision-encoder updates to the pretrained SAE feature span, preserving and reusing existing semantic concepts.
  • Experiments on ImageNet and distribution-shift benchmarks show state-of-the-art robustness without text-side augmentations or injected priors, with improved downstream transfer on CIFAR-10 and CIFAR-100.
  • SAE-FT provides mechanistic analysis showing that zero-shot features are explicitly retained and re-weighted.

2 Related Work

Prior work identifies a robustness–adaptation trade-off and links it to feature suppression and representational drift. SAEs provide an interpretable basis for analyzing these changes in CLIP representations.

  • WiSE-FT interpolates fine-tuned and zero-shot weights, while FLYP continues contrastive fine-tuning across vision and text modalities.
  • Feature suppression describes supervised fine-tuning diminishing pretrained features that are not aligned with the downstream objective.
  • The Linear Representation Hypothesis motivates SAEs as decompositions of dense activations into sparse, interpretable feature directions.

3 Preliminaries

CLIP maps images and text into a shared representation space for contrastive learning and zero-shot classification. SAE-FT operates in the vision-only linear-head setting and uses SAEs to expose sparse semantic features.

  • CLIP uses image and text encoders that map both modalities into a shared d-dimensional space, trained by contrasting matched and mismatched image-text pairs.
  • Zero-shot classification embeds prompted class descriptions and compares image representations with the resulting class embeddings.
  • Standard CLIP fine-tuning updates the vision encoder and linear classification head using cross-entropy loss.
  • SAE-FT remains in this vision-only linear-head setting while adding regularization to the visual representations.
  • SAEs decompose dense, polysemantic CLIP representations into sparse, human-understandable directions.
  • A Top-k SAE maps image representations into a sparse higher-dimensional latent space using TopK activation.
  • The SAE encoder maps representations into latent features, while decoder weights define a dictionary mapping sparse activations back to CLIP-space directions.

4 Representational Drift in CLIP Fine-Tuning

Fine-tuning substantially changes CLIP’s representation geometry and can collapse the pretrained SAE dictionary. Weight averaging partially restores geometry, motivating SAE-FT’s explicit span and feature-change constraints.

  • Representational-drift analysis motivates geometric constraints because fine-tuning limits interpretability and robustness.
  • 0.40 CKA similarity shows that standard fine-tuning produces major changes relative to zero-shot representations.
  • WiSE-FT reaches 0.83 similarity to zero-shot geometry, compared with 0.67 for direct representation averaging.
  • Figure 2 reports FVU and preserved zero-shot SAE-feature percentages across fine-tuning epochs and the weight-averaged model.
  • Standard fine-tuning yields FVU > 1.0 for the pretrained SAE, meaning the original dictionary performs worse than a zero-vector baseline.
  • SAE-FT constrains changes to the fixed SAE span and penalizes additions of new SAE features.
  • L2 regularization keeps CKA at 1.0 and FVU low, but its neuron-based penalty restricts changes along potentially non-semantic axes.

5 SAE-FT

SAE-FT constrains fine-tuning updates to an interpretable feature span while preserving representation geometry and selectively limiting semantic feature changes. It regularizes residual and sparse feature changes using a frozen Sparse Autoencoder, without using the SAE during inference.

  • Residual regularization: The residual alignment penalty forces representation updates to be expressible as linear combinations of the decoder’s interpretable feature span.It minimizes the component of the representation change orthogonal to that span.
  • Sparse feature regularization: Sparse feature regularization encourages the model to change as few SAE features as possible by penalizing feature differences with an L1 norm.The method contrasts this semantic feature-space constraint with standard regularization in raw representation space.
  • Feature preservation: Feature preservation penalizes newly activated features while allowing the magnitudes of existing features to change.With a Top-K SAE, this acts as a support-set constraint that encourages re-weighting the original active features.
  • Training and inference: The SAE is trained on frozen zero-shot representations, kept fixed during fine-tuning, and omitted during inference to limit computational overhead.During training, it computes regularization terms for zero-shot and fine-tuned representations.
  • Semantic feature regularization: SAE-FT constrains representation changes to directions in a learned, semantically meaningful feature basis rather than arbitrary neuron axes.This allows substantial raw activation changes when they correspond to limited updates in SAE feature space.

6 Experiments

Experiments evaluate SAE-FT on ImageNet distribution shifts, downstream transfer datasets, and real-world OOD benchmarks, alongside robust and generic representation-regularization baselines. SAE-FT achieves strong robustness, transfer performance, and competitive in-distribution accuracy while operating on the vision encoder only.

  • Method and efficiency: SAE-FT operates on the vision encoder with a cross-entropy linear-head framework and avoids text-side augmentations or injected priors.Its regularization constrains updates to the semantic span of a pretrained Sparse Autoencoder.
  • ImageNet and distribution shifts: SAE-FT reaches 82.9% ImageNet accuracy and the highest average accuracy across distribution-shift benchmarks for OpenAI ViT-B/16.It improves ImageNet accuracy by 1.0 percentage point over CAR-FT and average OOD performance by 0.2 percentage points over WiSE-FT among vision-encoder-only methods.
  • Downstream transfer: SAE-FT achieves the highest average transfer accuracy across CIFAR-10, CIFAR-100, Caltech-101, and STL-10.The transfer protocol evaluates ImageNet-fine-tuned visual representations with the original zero-shot CLIP text classifier.
  • Real-world distribution shifts: Across iWildCam and FMoW, SAE-FT achieves the best reported performance among the evaluated methods.These benchmarks measure robustness under real-world distribution shifts using macro-F1, ID accuracy, and worst-group OOD accuracy.
  • Representation regularization baselines: L1, L2, and PCA regularization substantially reduce robustness degradation, while SAE-FT achieves slightly higher average accuracy.The alternatives have broadly similar performance, making SAE-FT’s semantic constraint its primary distinction rather than a large performance gain.

7 Analysis of SAE-FT

Analysis finds that SAE-FT preserves the zero-shot model’s feature support while re-weighting features for downstream decisions. A pirate-ship example illustrates how this mechanism can correct a zero-shot error by increasing a task-relevant feature’s priority.

  • Feature statistics: Both L2 regularization and SAE-FT keep representations close to zero-shot representations, with CKA approximately 1.0 and FVU at most 0.25.SAE-FT nevertheless differs from L2 in feature overlap and entropy because it specifically penalizes feature addition and removal.
  • Feature preservation: SAE-FT preserves more zero-shot features than L2 regularization while re-weighting them rather than replacing them.The analysis compares zero-shot, L2-regularized, and SAE-regularized models using SAE feature statistics.
  • Feature re-weighting: In a pirate-ship example, SAE-FT corrects a schooner misclassification by amplifying the pirate-ship feature while retaining the schooner feature with reduced importance.The comparison examines activation patterns for an image misclassified by the zero-shot model but correctly classified by SAE-FT.
  • Mechanistic interpretation: The analysis suggests SAE-FT focuses fine-tuning on a sparse subset of pretrained features aligned with target classes rather than learning representations from scratch.This provides a feature-level account of its behavior on datasets that are difficult for zero-shot and L2-regularized models.

8 Conclusion

SAE-FT regularizes CLIP fine-tuning in an interpretable feature basis, achieving robust downstream adaptation while preserving control over semantic changes. Its gains over direct geometric-drift regularization are marginal, but it provides finer-grained interpretability; future work may extend regularization to both modalities.

  • 8 Conclusion: SAE-FT achieves state-of-the-art robustness under distribution shifts and superior downstream generalization compared with existing methods.The paper emphasizes especially strong performance relative to other vision-encoder-only methods.
  • 8 Conclusion: SAE-FT regularizes representation learning in an interpretable feature basis rather than directly constraining only geometric drift.This gives the method finer control over changes to the pretrained representation space.
  • 8 Conclusion: Compared with direct geometric-drift regularization, SAE-FT offers only marginal performance gains but substantially finer-grained control and interpretability.The conclusion distinguishes performance improvement from mechanistic transparency.
  • 8 Conclusion: SAE-FT selectively re-weights pre-existing interpretable concepts, amplifying task-relevant features and dampening irrelevant variations.This preserves learned concepts while adapting representations to downstream tasks.
  • 8 Conclusion: Future work could regularize both vision and text modalities to preserve features in both encoders and potentially improve robustness and generalization.The proposed extension would expand the FLYP fine-tuning protocol.

C Feature importance – task relevance correlation

The feature-task alignment analysis finds that SAE-FT adapts CLIP by re-weighting existing features toward the target class rather than introducing new ones. An optimal-transport regularizer provides an additional geometrically informed way to compare feature representations.

  • Feature-task alignment: SAE-FT yields the highest average feature-task alignment between active features and the correct class embedding.The table caption identifies SAE-FT as the leading method on the ImageNet test set.
  • Feature-task alignment: Fine-tuning increases active-feature alignment with the target class, and SAE-FT exceeds both the zero-shot model and L2 regularization.The analysis attributes this adaptation to re-weighting existing features toward class-embedding directions.
  • Feature-task alignment: High feature-task alignment corresponds to feature-task similarity concentrated in the alignment zone, unlike uniform similarity across features.The visualization contrasts concentrated alignment with uniform similarity.
  • Additional feature regularization: An optimal-transport regularizer measures distances between feature representations using a cost function based on feature cosine similarity.The resulting regularization loss is the Wasserstein distance between probability measures over activated features.
  • Additional feature regularization: The PCA baseline restricts representation changes to the first K truncated-SVD directions and adds residual and sparsity penalties.The baseline uses K=16 components and an L1 penalty on latent changes.

F.2 SAE-FT training times and computational overhead

SAE-FT adds limited computational cost: fine-tuning time is comparable to standard fine-tuning, with roughly 5% total ImageNet training overhead. Its per-step and memory overhead over L2 regularization is negligible, although SAE training requires preprocessing storage.

  • Training times: 1:47:23 hours per epoch for SAE-FT versus 1:49:34 hours for normal fine-tuning indicates no added fine-tuning overhead in testing.The reported difference is attributed to standard system variance.
  • Training times: Approximately 5% additional compute time is required for 10 ImageNet epochs, but this cost disappears when another CLIP model reuses the same SAE.The one-time SAE cost is not repeated for the reused model.
  • Computational overhead: Compared with L2 regularization, SAE-FT increases per-step time by 0.4% and peak GPU memory by only 19.7 MB.Both regularized methods require approximately 586 MB additional GPU memory over standard fine-tuning.
  • Computational overhead: Storing zero-shot ImageNet representations before SAE training requires 2.5 GB of disk space.This storage requirement is part of SAE-FT preprocessing rather than per-step fine-tuning.

G.2 SAE architecture ablation

SAE-FT is sensitive to the Sparse Autoencoder architecture, especially on distribution-shift benchmarks, while in-distribution accuracy changes only mildly. The default configuration consistently performs best across metrics, and feature-addition regularization outperforms alternative SAE regularizers.

  • Sensitivity to SAE architecture: 81.6–82.9 ImageNet accuracy remains relatively stable across SAE architectures, whereas IN-A ranges from 46.3 to 52.6 and IN-R from 72.0 to 78.5.The reported variation is substantially larger on out-of-distribution benchmarks than in-distribution evaluation.
  • Dictionary size and active features: Both smaller and larger dictionaries degrade out-of-distribution performance relative to the default mult = 4 configuration.The authors attribute this pattern to insufficient semantic capacity for smaller dictionaries and redundant or poorly learned features for larger ones.
  • Dictionary size and active features: Fewer active features provide a coarse representation, while more active features dilute regularization across directions and substantially degrade out-of-distribution robustness.The K = 32 configuration may improve with re-tuned hyperparameters because changing active-feature count changes the effective feature-addition penalty.
  • Default configuration: The default configuration, K = d/32 and mult = 4, consistently achieves the best results across all metrics.The ablation varies active features and dictionary size while keeping λres and λadd fixed at default-configuration values.
  • Regularization comparisons: Feature-addition regularization outperforms Wasserstein and sparsity regularization on average distribution-shift accuracy, despite similar in-distribution performance.The comparison evaluates alternative SAE regularization methods on ImageNet and its distribution shifts.
  • Generalization beyond ImageNet: L2 regularization matches SAE-FT on downstream transfer benchmarks, while SAE-FT slightly outperforms L2 on the Describable Textures Dataset.Both SAE-FT and L2 outperform standard fine-tuning on DTD.

H Further qualitative results

The qualitative analysis examines ImageNet samples where zero-shot and SAE-FT predictions or features differ, including rare cases where SAE-FT misclassifies an image correctly classified by the zero-shot model. Examples include shifts such as Barbershop to Barber Chair and Safety Pin to Border Terrier.

  • Analysis scope: The qualitative analysis focuses on ImageNet samples where zero-shot and fine-tuned predictions or features differ.It also includes rare examples that SAE-FT misclassifies while the zero-shot model is correct.
  • Prediction changes: Another prediction changes from Safety Pin under zero-shot CLIP to Border Terrier under SAE-FT.The listed qualitative examples include additional zero-shot predictions such as Cairn Terrier, Brambling, Magnetic compass, Carved Pumpkin, and Music Speaker.
  • Unchanged prediction: For another sample, both zero-shot CLIP and SAE-FT predict Carved Pumpkin.This example illustrates that the qualitative set includes samples with differing features or predictions, not only changed labels.
  • Prediction changes: One prediction changes from Barbershop under zero-shot CLIP to Barber Chair under SAE-FT.This is an example of a prediction difference examined in the qualitative analysis.
Loading 2605.15961v1…