Flatness Preserves Instruction Following in Vision-Language-Action Models

1. Introduction: The Hidden Fragility of Vision-Language-Action (VLA) Models

A critical failure mode in the deployment of embodied AI is “instruction blindness,” a phenomenon where Vision-Language-Action (VLA) models ignore language directives in favor of visual shortcuts learned during downstream finetuning. In high-stakes environments, this lack of robustness presents a clear safety risk: a robot might proceed to handle a hazardous material—for example, “bleach”—despite a corrective semantic instruction to “pick up the vinegar,” simply because the model has overfit to the visual scene layout of its training data.

Research by Zhang and Bisk (2026) suggests that this fragility is a geometric consequence of the loss landscape. The authors hypothesize that finetuning on sparse, biased robot data causes models to converge into “sharp” minima, wherepretrained semantic grounding is sacrificed for local task performance. The paper demonstrates that applying Sharpness-Aware Minimization (SAM) during the finetuning phase significantly improves instruction following by over 60% across multiple benchmarks. Crucially, this improvement is achieved without additional training data, architectural modifications, or retraining of the base model.

2. The Geometry of Failure: Sharp Minima and Limited Data

The disparity between internet-scale pretraining and robot-specific finetuning is the primary driver of representation loss. While pretraining provides dense coverage of the vision-language manifold, robot datasets are often sparse and imbalanced, applying gradients to only a limited set of points. This over-indexing leads to high-curvature minima and brittle policies.

Zhang and Bisk define the problem as follows:

“Instruction blindness occurs when the learned manifold $M$ is insufficiently sensitive to semantically distinct language perturbations: the manifold has converged such that novel $(o, l)$, which are within support of the internet-scale pretrained VLM, are now poorly supported on $M$.”

As depicted in the analysis of [SOURCE_IMAGE_1], standard VLA finetuning creates narrow “purple” pits in the loss landscape. These pits represent regions where the model has overfit to specific visual cues, making the policy unstable under weight perturbations. In contrast, the authors’ proposed flatness-preserving optimization creates stable “green” basins. These flatter regions indicate a reduction in “Keskar Sharpness,” a metric used by the paper to quantify how sensitive the model is to input variations that should otherwise be resolved by the linguistic instruction.

3. The Proposed Remedy: Sharpness-Aware Minimization (SAM)

To counter the emergence of sharp minima, the paper implements SAM as a bi-level optimization process. Instead of minimizing loss at a single point, SAM seeks parameters in uniformly flatter regions by minimizing the worst-case loss within a neighborhood. This prevents the model from relying on sharp, vision-dominated features that appear predictive in small datasets but do not generalize.

The three-step implementation for VLAs is as follows:

• Step 1: Compute Perturbations. The algorithm identifies the maximizing perturbation $\hat{\epsilon}$ (the “worst-case” direction) at the current parameters $\theta$.
• Step 2: Evaluate Gradients. The gradient $\tilde{g}$ is computed at the perturbed point $\theta + \hat{\epsilon}$ rather than the original point.
• Step 3: Update Parameters. The original parameters are updated using $\tilde{g}$ via a base optimizer such as AdamW.

By enforcing a flatter loss basin, the embedding function remains grounded in the pretrained VLM’s semantic structure, ensuring that the model remains sensitive to language-based corrections.

4. Empirical Evidence: Breaking Benchmarks in Simulation

The authors evaluated the $\pi_{0.5}^{SAM}$ model across three core benchmarks: LIBERO-PRO Task, LangGap, and LIBERO-CF. These benchmarks encompass 138 tasks designed to test “counterfactual” scenarios, where the robot must ignore visual biases to follow a novel instruction.

The following table compares the success rates of the standard $\pi_{0.5}$ model against the proposed $\pi_{0.5}^{SAM}$:

Benchmark	$\pi_{0.5}$ (Standard) Success Rate	$\pi_{0.5}^{SAM}$ (Zhang & Bisk) Success Rate	Absolute Improvement
LIBERO-PRO Task (Avg)	26.6%	42.6%	+16.0%
LangGap (Avg)	24.5%	41.7%	+17.2%
LIBERO-CF (Avg)	13.2%	47.8%	+34.6%

The most significant gains were observed in “Object” subsets, which require grounding specific target objects among similar items. The authors argue that this specifically validates the method’s ability to preserve vision-language coupling. Furthermore, Appendix D.4 of the paper notes that SAM is complementary to existing techniques; a combination of SAM + CFG (Classifier-Free Guidance) resulted in an additional 17.8% average improvement on LIBERO-PRO, establishing a significantly higher baseline for robustness.

5. Visualizing the Shift: Representation Probing and Hessian Analysis

To provide a technical quantification of why SAM preserves semantic grounding, the authors conducted an object-level representation probe using the COCO dataset [SOURCE_IMAGE_2]. Specifically, they extracted the Gemma last-layer hidden states at object tokens for categories such as “cup,” “bottle,” and “chair.”

In the standard $\pi_{0.5}$ model, these semantic clusters were largely disrupted after finetuning. However, the $\pi_{0.5}^{SAM}$ model preserved the coherent clusters found in the pretrained VLM. The authors quantified this using the Silhouette Score and K-NN Accuracy (Appendix B), which measure cluster separability and local neighborhood structure quality.

Additionally, the authors used the Lanczos algorithm to approximate the local geometry of the loss landscape, focusing on the action head’s parameter space.

Metric	$\pi_{0.5}$ (Standard)	$\pi_{0.5}^{SAM}$
Keskar Sharpness ($S$)	0.012	0.005
Max Eigenvalue ($\lambda_{max}$)	0.93	0.52

The significant reduction in $\lambda_{max}$ confirms that SAM effectively smooths the manifold, lowering the local curvature and reducing the risk of catastrophic forgetting of linguistic grounding.

6. Real-World Validation: From Simulation to the DROID Setup

The paper validates these findings in a real-world setting using a 7-DoF Franka Panda and the DROID setup. Five pick-and-place tasks were designed where visual bias was intentionally induced by training the model on only a single task per scene layout.

When evaluated on counterfactual instructions (e.g., picking up a hand cream bottle instead of the training-set salt shaker in the same scene), the standard model achieved only a 13.8% success rate. The $\pi_{0.5}^{SAM}$ model achieved 36.3%, representing a 163% increase in success in counterfactual real-world deployment.

[SOURCE_IMAGE_3] illustrates the distinction between “Grounding Rate” and “Task Success.” Grounding Rate—defined as the binary measure of whether the gripper makes contact with the correct object—remained high for $\pi_{0.5}^{SAM}$ even when the grasp itself failed. This suggests that the method effectively aligns the robot’s “intent” with the instruction, even when physical execution remains challenging.

7. Strategic Insights: Global vs. Selective Application

Component ablation studies (Table 6) were conducted to determine if SAM should be targeted at specific modules. The authors compared applying SAM to the vision, language, and action components individually.

Application Area	Average Success Rate (LIBERO-PRO)
Action Head Only	21.0%
Vision Component Only	22.7%
Language Component Only	23.0%
Global Application	42.6%

[!IMPORTANT] Researcher’s Note: The Necessity of Global Optimization Zhang and Bisk found that the curvature induced by low-data finetuning is not localized; it propagates throughout the entire network. Selective application—such as targeting only the action head—might reduce local sharpness in that component, but high sharpness in the language backbone continues to cause instruction blindness. For robust VLA performance, flatness must be enforced globally to ensure the preservation of original pretrained representations across all bottlenecks.

8. Conclusion: Takeaways for AI Safety and Future Research

While SAM provides a significant boost to robustness, the authors identify key trade-offs. The method requires two forward passes per training step, doubling the compute cost. Furthermore, while the relative improvement is substantial, absolute success rates in complex counterfactual scenarios suggest that loss landscape geometry is only one part of the safety puzzle.

Key Takeaways for the AI Safety Community:

1. Landscape Curvature as a Proxy for Risk: Red-teamers should look for high “sharpness” ($\lambda_{max}$) in the loss landscape as a precursor to catastrophic forgetting and instruction blindness in low-data finetuning environments.
2. Flatness as a Foundational Defense: Enforcing flatness provides a stable baseline that allows inference-time techniques (like Classifier-Free Guidance) to perform more effectively, suggesting that safety should be “baked in” at the optimization level.
3. Optimization over Augmentation: In regimes where high-quality expert robot data is scarce, modifying the optimization geometry (via SAM) is a more efficient path to generalization than the costly collection of counterfactual demonstration data.

Read the full paper on arXiv · PDF