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

Problem Statement

Large-scale pre-trained vision-language models, such as CLIP, exhibit remarkable zero-shot performance and robustness to distribution shifts. However, fine-tuning these models to improve in-distribution (ID) performance on downstream tasks often leads to a degradation in out-of-distribution (OOD) robustness. This phenomenon, known as "representational drift" or "feature suppression," occurs because standard fine-tuning distorts the pre-trained representation space, overwriting semantically meaningful features with task-specific noise or arbitrary directions.

Existing robust fine-tuning methods attempt to mitigate this trade-off but face significant limitations:

• WiSE-FT: Averages weights between the zero-shot and fine-tuned models. While effective, it often ignores the text encoder and relies on simple interpolation, offering limited control over which features are preserved.
• Text-Guided Approaches (e.g., FLYP, StarFT): These methods fine-tune both vision and text encoders or inject synthetic features and external priors into text prompts. These approaches introduce computational complexity and rely on external data engineering, complicating the fine-tuning pipeline.

The core challenge is to achieve robust fine-tuning that improves ID performance without sacrificing OOD generalization, while doing so in a computationally efficient manner that operates solely on the visual modality and provides interpretability regarding how the model adapts.

Methodology: SAE-FT

The authors propose SAE-FT (Sparse Autoencoder Fine-tuning), a regularization framework that constrains the fine-tuning process to the interpretable feature span of the pre-trained model. The method relies on the Linear Representation Hypothesis, which posits that high-level concepts are encoded as linear directions in activation space.

Core Components

1. Sparse Autoencoder (SAE) Training:

   • A Top-k SAE is trained unsupervised on the frozen visual representations ($r_0$) of the pre-trained zero-shot CLIP model.
   • The SAE decomposes dense, polysemantic representations into a sparse, overcomplete basis of semantically meaningful features ($s$).
   • The SAE weights are frozen after training and serve as a fixed dictionary of interpretable features.

2. Regularization Strategy:
   During fine-tuning of the vision encoder ($f_{ft}$), the model is constrained to ensure that updates to the representations ($r_{ft}$) remain within the semantic span of the pre-trained SAE. The total loss function combines the standard cross-entropy loss ($L_{CE}$) with two specific regularization terms:

   • Residual Alignment Penalty ($L_{resid}$): Ensures that the change in representation ($\Delta r = r_{ft} - r_0$) can be reconstructed by the SAE decoder. This forces updates to be expressible as a linear combination of the learned semantic features, preventing the model from drifting into arbitrary geometric directions orthogonal to the dictionary.
     $$L_{resid} := \|\Delta r - W_d(\Delta s)\|_2^2$$
   • Feature Preservation Penalty ($L_{add}$): Explicitly penalizes the activation of features that were inactive in the zero-shot model. This prevents the "addition" of new, potentially spurious semantic concepts.
     $$L_{add} := \lambda_{add} \sum_{k=1}^{p} (1 - m_k)|s_{ft}^k|$$
     where $m_k$ is an indicator for features active in the zero-shot model.

3. Training and Inference:

   • The SAE is not updated during fine-tuning and is not used during inference. It serves only as a geometric constraint during the training loop.
   • The method operates exclusively on the vision encoder using a linear-head cross-entropy setup, avoiding the complexity of text-side data engineering.

Key Contributions

1. SAE-FT Framework: A novel fine-tuning strategy that constrains representation updates to the interpretable feature span of a pre-trained backbone. It ensures adaptation occurs by re-weighting existing semantic concepts rather than overwriting them, preventing catastrophic forgetting of general knowledge.
2. Performance and Efficiency: Extensive experiments demonstrate that SAE-FT matches or exceeds state-of-the-art (SOTA) robustness on ImageNet and distribution shift benchmarks. Crucially, it achieves this without text-side augmentations or injected priors, making it computationally efficient and easier to deploy than text-guided methods.
3. Mechanistic Insight: The framework provides granular analysis of feature preservation. It explicitly demonstrates that robust fine-tuning involves re-weighting existing features (e.g., amplifying "ladder" and "red paint" for a fire truck) rather than learning entirely new representations, offering a transparent view of the adaptation process.

Experimental Results

The authors evaluated SAE-FT on the OpenAI ViT-B/16 model across multiple benchmarks:

• ImageNet and Distribution Shifts: SAE-FT achieved an ImageNet accuracy of 82.9% (matching SOTA) and the highest average accuracy (64.6%) across distribution shift benchmarks (ImageNet-R, A, S, V2). It outperformed other vision-encoder-only methods like CAR-FT and WiSE-FT.
• Downstream Transfer: When evaluated on CIFAR-10, CIFAR-100, Caltech-101, and STL-10 without further fine-tuning, SAE-FT achieved the highest average transfer accuracy (87.8%), indicating that its representations generalize effectively to diverse tasks.
• iWilds Benchmarks: On real-world OOD datasets (iWildCam and FMoW), SAE-FT achieved the best or near-best performance, particularly in worst-group accuracy and OOD macro F1 scores, outperforming L2 regularization and other robust fine-tuning methods.
• Comparison to Baselines: While generic regularization methods (L1, L2, PCA) also improved robustness by limiting geometric drift, SAE-FT provided superior control over specific feature changes. It achieved slightly higher average accuracy than these baselines but offered the distinct advantage of operating in a learned, semantically meaningful feature basis.

Significance and Claims

The paper claims that SAE-FT represents a significant step forward in making fine-tuning both robust and interpretable.

• Interpretability: Unlike standard norm regularization which treats all directions equally, SAE-FT allows researchers to directly analyze which semantic features are preserved, re-weighted, or suppressed. The authors show that the model corrects misclassifications (e.g., distinguishing a pirate ship from a schooner) by amplifying task-relevant features while retaining the original feature vocabulary.
• Efficiency: By operating solely on the visual modality and avoiding complex text-side manipulations or external LLM data, SAE-FT offers a streamlined pipeline for robust fine-tuning.
• Mechanism of Success: The authors argue that the success of SAE-FT stems from its ability to force the model to focus on a sparser, more relevant subset of pre-existing features. This prevents the "crippling" of foundation model features that typically occurs during standard fine-tuning, thereby preserving the model's general knowledge while sharpening task-specific distinctions.

The work concludes that interpretable regularization can enable efficient fine-tuning without erasing the learned concepts of pre-trained models, offering a viable alternative to complex, text-dependent robust fine-tuning strategies.