The Geometry of Robustness: Optimizing Loss Landscape Curvature and Feature Manifold Alignment for Robust Finetuning of Vision-Language Models

1. Problem Statement

Vision-Language Models (VLMs) like CLIP possess strong zero-shot capabilities but face a critical three-way trade-off when fine-tuned for downstream tasks:

1. In-Distribution (ID) Accuracy: Performance on the training distribution.
2. Out-of-Distribution (OOD) Generalization: Performance on natural distribution shifts (e.g., sketches, renditions).
3. Adversarial Robustness: Resistance to gradient-based attacks (e.g., PGD).

Existing robust fine-tuning strategies typically resolve only two of these axes at the expense of the third:

• Generalization-focused methods (e.g., WiSE-FT, TPGM) preserve ID/OOD performance but remain vulnerable to adversarial attacks.
• Adversarial training methods (e.g., TeCoA, FARE) improve robustness but often degrade ID and OOD accuracy due to overfitting to specific attack patterns or drifting too far from the pre-trained manifold.

The authors identify that this trade-off stems from two geometric failures in the model's optimization landscape:

1. Sharp, Anisotropic Minima: The model converges to sharp regions in parameter space, making it sensitive to weight perturbations.
2. Feature Manifold Instability: Feature representations deform significantly under distribution shifts or attacks, breaking alignment between clean, OOD, and adversarial inputs.

2. Methodology: GRACE Framework

The authors propose GRACE (Gram-aligned Robustness via Adaptive Curvature Estimation), a unified fine-tuning framework grounded in Robust PAC-Bayes theory. GRACE jointly regularizes parameter-space curvature and feature-space invariance using three core components:

A. Theoretical Foundation

Based on a PAC-Bayesian bound, the robust risk is decomposed into three terms that must be minimized simultaneously:

• (A) Proximity to Prior: Keeping fine-tuned weights close to pre-trained weights ($\|\theta - \theta_0\|^2$) to preserve zero-shot knowledge.
• (B) Parameter-Space Sharpness: Minimizing the trace of the Hessian ($\text{Tr}(\nabla^2 R)$) to find flatter minima.
• (C) Feature-Space Stability: Minimizing domain discrepancy ($d_{H\Delta H}$) to ensure feature consistency across ID, OOD, and adversarial domains.

B. Core Components

1. LoRA Fine-Tuning (Proximity Control):

   • GRACE uses Low-Rank Adaptation (LoRA) to update only a low-rank subspace of the frozen pre-trained weights. This implicitly controls the KL divergence term, ensuring the model stays close to the pre-trained manifold and prevents catastrophic forgetting.

2. Layerwise Adaptive Low-Rank Adversarial Weight Perturbation (LAR-AWP):

   • Mechanism: Instead of uniform perturbations, GRACE injects adversarial weight perturbations in the LoRA subspace.
   • Adaptivity: Crucially, the rank of the perturbation is adapted per layer based on local curvature estimates (using a Gauss-Newton proxy derived from mini-batch gradients).
   • Goal: Layers with high curvature (sharp regions) receive higher perturbation ranks to force the optimizer toward flatter minima, while flat layers receive minimal perturbation. This addresses the anisotropic nature of VLM loss landscapes.

3. Gram-Volume Alignment Loss:

   • Mechanism: For each input, the model computes a $3 \times 3$ Gram matrix using feature embeddings from three states: Clean, Adversarial (PGD), and LAR-AWP Perturbed (weight perturbation).
   • Objective: The loss minimizes the volume of the parallelepiped spanned by these three vectors ($L_{GV} = \sqrt{|\det(G)|}$).
   • Goal: This forces the feature representations to remain consistent (low volume) across clean, adversarial, and weight-perturbed inputs, ensuring the feature manifold does not collapse or scatter under shifts.

Training Objective:
$$\mathcal{L}_{GRACE} = \mathcal{L}_{task} + \lambda_{LAR} \mathcal{L}_{LAR-AWP} + \lambda_{GV} \mathcal{L}_{GV}$$

3. Key Contributions

1. Geometric Analysis: The paper provides a theoretical and empirical analysis linking VLM robustness failures to sharp minima and feature manifold distortion. It categorizes existing methods by which geometric terms they fail to optimize.
2. Unified Framework (GRACE): Introduces the first framework to simultaneously optimize ID accuracy, OOD generalization, and adversarial robustness by jointly controlling parameter curvature and feature invariance.
3. Layerwise Adaptivity: Proposes a novel LAR-AWP mechanism that dynamically adjusts perturbation ranks based on layer-wise curvature, offering a more efficient and targeted approach to flattening the loss landscape than uniform methods.
4. Gram-Volume Regularization: Introduces a differentiable surrogate for feature manifold stability that aligns clean, adversarial, and weight-perturbed embeddings.

4. Experimental Results

Experiments were conducted on CLIP ViT-B/32 and ViT-B/16 fine-tuned on ImageNet-1K, evaluated across ID, OOD (ImageNet-V2, -R, -S), and Adversarial (ImageNet-A, A-Plus, AutoAttack) benchmarks.

• Performance Gains:

  • ID Accuracy: GRACE achieves 74.21% (vs. 74.86% for vanilla FT), maintaining near-state-of-the-art clean accuracy.
  • Adversarial Robustness: Achieves 25.44% adversarial accuracy on ImageNet (vs. 0% for vanilla FT and ~22% for specialized AT methods).
  • OOD Generalization: Maintains 57.01% average OOD accuracy, comparable to the zero-shot baseline (57.4%) and superior to adversarial training methods which often drop below 40%.
  • Harmonic Mean: GRACE achieves a harmonic mean of 39.69%, significantly outperforming the next best method (PMG-AFT at 32.85%) and bridging the gap between specialized robustness strategies.

• Comparative Analysis:

  • Unlike WiSE-FT (good OOD, bad Adversarial) or TeCoA (good Adversarial, bad OOD), GRACE successfully balances all three axes.
  • Ablation Studies confirm that both LAR-AWP (curvature control) and Gram-Volume (feature alignment) are necessary; removing either leads to significant drops in specific robustness metrics.

• Efficiency: GRACE is 1.4x faster than prior adversarial fine-tuning methods while achieving higher overall performance, as shown in the Pareto frontier analysis.

5. Significance

This work establishes geometry-driven fine-tuning as a general paradigm for robust multimodal adaptation. By demonstrating that the trade-off between accuracy, generalization, and robustness is solvable through the joint optimization of loss landscape curvature and feature manifold alignment, GRACE offers a principled solution for deploying VLMs in real-world, high-stakes environments where both natural distribution shifts and adversarial attacks are concerns. The method effectively resolves the "trilemma" that has previously limited the deployment of robust foundation models.