DiG-Flow: Discrepancy-Guided Flow Matching for Robust VLA Models

TL;DR

DiG-Flow introduces a geometric regularization framework for VLA models that uses distributional discrepancy to improve robustness and performance, especially on complex tasks and under limited data conditions.

Contribution

It proposes a novel discrepancy-guided regularization method that enhances VLA model robustness without altering the core flow matching process.

Findings

Improves VLA performance on complex multi-step tasks.

Enhances robustness under distribution shifts.

Achieves these gains with negligible computational overhead.

Abstract

Vision-Language-Action (VLA) models trained with flow matching have demonstrated impressive capabilities on robotic manipulation tasks. However, their performance often degrades under distribution shift and on complex multi-step tasks, suggesting that the learned representations may not robustly capture task-relevant semantics. We introduce DiG-Flow, a principled framework that enhances VLA robustness through geometric regularization. Our key insight is that the distributional discrepancy between observation and action embeddings provides a meaningful geometric signal: lower transport cost indicates compatible representations, while higher cost suggests potential misalignment. DiG-Flow computes a discrepancy measure between empirical distributions of observation and action embeddings, maps it to a modulation weight via a monotone function, and applies residual updates to the observation embeddings before flow matching. Crucially, this intervention operates at the representation level without modifying the flow matching path or target vector field. We provide theoretical guarantees showing that discrepancy-guided training provably decreases the training objective, and that guided inference refinement converges with contraction. Empirically, DiG-Flow integrates into existing VLA architectures with negligible overhead and consistently improves performance, with particularly pronounced gains on complex multi-step tasks and under limited training data.