Flatness Guided Test-Time Adaptation for Vision-Language Models

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

ICLR 2026; flatness-guided TTA links loss landscape geometry to VLM distribution shift

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

FGA (Flatness-Guided Adaptation) is a test-time adaptation framework for VLMs that exploits loss landscape geometry to unify training and test-time procedures. During training, Sharpness-Aware Prompt Tuning (SAPT) locates a flat minimum; at inference, Sharpness-based Test Sample Selection (STSS) picks augmented test views whose loss landscapes align with that flat minimum—without any gradient updates or parameter changes at test time. On ViT-B/16, FGA surpasses TPT+CoOp by an average of 4.88% across ImageNet out-of-domain variants.

Problem

Existing TTA methods for VLMs (e.g., TPT and its variants) treat test-time adaptation as an isolated optimization problem decoupled from training history. They rely on expensive backpropagation through prompt parameters at inference, ignore the geometric structure of the pre-trained model's loss landscape, and consequently deliver suboptimal generalization under distribution shifts.

Method

FGA unifies training and inference around a single geometric principle: alignment of flat minima between the training and test loss landscapes.

Stage 1 – Sharpness-Aware Prompt Tuning (SAPT). Prompts \(p\) on the downstream training set are optimized with a sharpness-regularized objective: $\(\ell_{\text{SAPT}}(p) = \ell_{\text{CE}}(p) + \lambda \max_{\|\epsilon\|\leq\rho} \left[\ell_{\text{CE}}(p+\epsilon) - \ell_{\text{CE}}(p)\right],\)$ where the worst-case perturbation is approximated via a single gradient step: \(\epsilon^\star = \rho\,\nabla_p\ell_{\text{CE}}(p)/\|\nabla_p\ell_{\text{CE}}(p)\|\). This yields prompts positioned at a flat minimum that acts as a geometric anchor.

Stage 2 – Sharpness-based Test Sample Selection (STSS). At inference the prompts are frozen. For a test image, \(N\) augmented views are generated. Each view's sharpness score is computed via \(R\) random prompt perturbations drawn from \(\mathcal{N}\)—no backpropagation required: $\(\ell_{\text{STSS}}(p) = \ell_{\text{SRG}}(p) + \lambda \max_{r=1,\ldots,R} \left[\ell_{\text{SRG}}\!\left(p + \rho'\tfrac{\epsilon_r}{\|\epsilon_r\|}\right) - \ell_{\text{SRG}}(p)\right],\)$ with \(\ell_{\text{SRG}}\) being a surrogate (entropy) when labels are absent. The top-\(s\) views with the lowest sharpness scores are selected; their predictions are aggregated by majority vote. A theoretical analysis (Theorem 4) shows that distributions closer to training tend to yield lower worst-case entropy, justifying the selection criterion.

Key Contributions

• Novel FGA framework that links sharpness-aware training to test-time adaptation via flat-minima alignment, eliminating test-time gradient descent entirely.
• SAPT: a SAM-style prompt tuning objective that provides both generalization-improved prompts and a flatness geometric anchor for subsequent inference.
• STSS: a backpropagation-free augmentation selection strategy using random prompt perturbations, with formal generalization bounds (Theorems 1 and 4) supporting the approach.
• Demonstrated computation savings over TPT-class methods while achieving superior OOD accuracy.

Results

All results on CLIP ViT-B/16 backbone; "OOD Avg" excludes ImageNet source domain.

• FGA vs. zero-shot CLIP: +6.84% average, +7.13% OOD average across five ImageNet-variant datasets.
• FGA vs. best TPT-based baseline (DPE): FGA Avg 68.04% vs. DPE 64.66%; FGA OOD Avg 66.55% vs. DPE 64.43%.
• FGA vs. TPT+CoOp: +4.88% average across all four OOD variants (stated explicitly in abstract).
• Per-dataset (FGA): IN 74.01%, IN-A 65.90%, IN-V2 67.23%, IN-R 81.24%, IN-Sketch 51.81%.
• Ablation: SAPT+CoOp alone achieves 62.64% Avg; STSS+CoOp alone achieves 66.48%; combining both (FGA) yields 68.04%, demonstrating complementary contributions.
• Cross-dataset generalization results (Table 2) are reported but the provided text is truncated.

Limitations

• Requires a labeled downstream training set for the SAPT stage; purely zero-shot settings (no training data) are not addressed.
• STSS adds inference-time overhead proportional to \(N \times R\) (augmented views × perturbation samples), though no backprop is needed.
• Theoretical guarantees (Theorem 4) rely on a \(\gamma\)-separability assumption between test distributions, which may not hold for subtle or gradual domain shifts.
• Designed for single-sample TTA; online methods (TDA, DPE) that leverage accumulated test statistics are a different regime and not directly comparable.
• The ViT-B/16 backbone is the primary evaluation; results for larger encoders (e.g., ViT-L/14) are not reported in the provided text.

Relevance to Vision-Language Models

FGA introduces a principled, geometry-aware perspective on CLIP adaptation that avoids the now-standard but expensive backpropagation-through-prompts paradigm of TPT, making it practically more deployable. The insight that loss-landscape flatness from training (SAM-style) can serve as a test-time signal bridges two previously siloed research threads—prompt tuning and sharpness-aware optimization—directly within the VLM context. For researchers tracking VLMs, this work highlights that the choice of training objective (not just the test procedure) has first-order consequences for adaptation quality, suggesting that future prompt tuning methods may benefit from explicit flatness objectives. The theoretical framing (Theorems 1 and 4) also provides a rare formal foundation for why augmentation-based sample selection improves OOD performance in contrastive VLMs.