DiMaS: Distribution Matching for Steering Vision-Language-Action Models¶

🕒 Published (v1): 2026-07-15 18:39 UTC · Source: Arxiv · link

Why this paper was selected

Fine-grained behavioral steering for flow-matching VLA; enables controllable robot policy

Ask a follow-up

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

TL;DR¶

DiMaS introduces a distribution-matching steering strategy for flow-matching Vision-Language-Action (VLA) models that transports internal representations via optimal transport rather than shifting along a fixed linear direction. It addresses the failure of classical linear steering in VLAs, where behavioral features are linearly decodable but not linearly steerable. The method is validated on SmolVLA and $\pi_{0.5}$ on the LIBERO benchmark, achieving reliable fine-grained control of end-effector speed and vertical displacement.

Problem¶

Representation steering—post-hoc inference-time control of model outputs by intervening on internal activations—is well-established for LLMs and VLMs under the linear representation hypothesis (LRH), but does not translate to flow-matching VLAs. Classical methods (mean-difference shift, regression-based shift) fail inconsistently in visuomotor settings: they often shift the target behavioral feature in the same direction regardless of intended direction (H→L vs. L→H). The authors show this is because behavioral features in the VLA action expert are linearly decodable (a linear probe can predict them) but not linearly steerable (a linear additive shift cannot reliably control them), violating the implicit assumption of LRH-based steering.

Method¶

DiMaS (Distribution Matching Steering) operates on residual-stream representations $h_{\ell,p}^{(m)}(x_t) \in \mathbb{R}^d$ extracted from the flow-matching action expert $f_A$ at a chosen layer $\ell^*$ and denoising step $m$.

Source/target construction. A continuous scalar behavioral feature $\phi(a_i)$ (e.g., speed $= \sqrt{\Delta x^2 + \Delta y^2 + \Delta z^2}$) is computed from predicted actions. Representations are partitioned by empirical quantile thresholds: $$\mathcal{D}^{-} = \{h_i : \phi_i \leq q_\tau\}, \quad \mathcal{D}^{+} = \{h_i : \phi_i > q_{1-\tau}\}$$ using tail quantiles for clean separation.

Transport map. A map $\mathcal{T}_\ell^{(m)}$ from $\mathcal{D}^-$ to $\mathcal{D}^+$ is learned by minimizing the Kantorovich $\mathcal{W}_2^2$ objective: $$\mathcal{W}_2^2(\mathcal{D}^{-}, \mathcal{D}^{+}) = \min_{\gamma \in \Pi(\mathcal{D}^{-}, \mathcal{D}^{+})} \int_{\mathcal{Z} \times \mathcal{Z}} \|z^- - z^+\|^2 \, d\gamma(z^-, z^+)$$ solved efficiently via low-rank Sinkhorn on discrete empirical samples.

Inference. A binary gate classifier $g_\ell^{(m)}$ determines if the feature is absent. If active, the representation is updated as: $$h \leftarrow (1-\alpha)\,h + \alpha\,\bigl(\mathcal{T}_\ell^{(m)} \circ P(h)\bigr)$$ where $P(h) = \arg\min_{z \in \mathcal{D}^-} \|z - h\|$ is a nearest-neighbor projection and $\alpha \in [0,1]$ interpolates between unsteered and fully transported representations to preserve task success.

Key Contributions¶

DiMaS: novel OT-based steering mechanism for flow-matching VLAs; operates on action expert representations and scales across model sizes (SmolVLA and $\pi_{0.5}$).
Generalization characterization: four-setting protocol (held-out initializations, cross-task within suite, cross-suite) that progressively stresses transfer; shows DiMaS remains effective even under cross-suite distribution shift.
Representation analysis: empirical demonstration that behavioral features in VLA action experts are linearly decodable but not linearly steerable, explaining classical steering failure and motivating the OT design.
Interpolation mechanism: $\alpha$ knob enabling continuous trade-off between feature control strength and task success rate without pausing the intervention (unlike prior work requiring intervention pauses).

Results¶

Speed steering (both models): DiMaS produces statistically significant shifts in both H→L and L→H directions across LIBERO-Object and LIBERO-Spatial suites ($p < 0.01$); linear baselines (mean-difference, VLM regression, FM regression) and prompt injection fail to reliably modulate speed in opposing directions.
Z-displacement steering ($\pi_{0.5}$): DiMaS achieves significant shifts; accompanying success-rate drops are expected since z-displacement is tied to task completion (e.g., keeping end-effector too low prevents lifting).
Longer-horizon tasks (LIBERO-10, $\pi_{0.5}$): Δmean speed = $-0.024$ at $\alpha=0.5$, $p < 0.01$; success rate 96% → 90%.
Cross-suite transfer ($\pi_{0.5}$, Setting 4): Speed shift persists in the intended direction ($p = 0.032$) even when learned on one LIBERO suite and evaluated on a disjoint one.
Aggregated training (Setting 2 vs. 1): Learning from multiple tasks jointly yields more stable and consistent shifts than single-task steering.
LIBERO-Goal (hardest suite): Weakest results; some settings do not reach $p < 0.01$ (SmolVLA Setting 1: $p=0.09$; $\pi_{0.5}$ Setting 4: $p=0.032$), attributed to high task diversity.

Limitations¶

Z-displacement steering incurs meaningful success-rate drops because the feature is coupled to task completion (e.g., vertical motion required for grasping/placing).
SmolVLA is more sensitive to action-expert intervention than $\pi_{0.5}$, attributed to its per-timestep action-chunking strategy producing less-smooth trajectories.
Steering effectiveness degrades on behaviorally diverse task suites (LIBERO-Goal); cross-suite transfer to Goal does not reach $p < 0.01$ for $\pi_{0.5}$.
The steering method is calibrated on continuous scalar features derived from the action expert itself; the feature extraction depends on having predicted action access, which may limit applicability to settings where post-hoc action annotation is unavailable.
Paper is evaluated only in simulation (LIBERO); real-robot generalization is unverified.

Relevance to Vision-Language Models¶

DiMaS directly implicates VLMs as the backbone of state-of-the-art robot policies and highlights a fundamental limitation in transferring VLM representation-steering recipes to the visuomotor domain: the linear representation hypothesis, which underlies most VLM steering techniques, holds for semantic features in language/vision but breaks for behavioral features in the action expert of flow-matching VLAs. The finding that features are linearly decodable but not linearly steerable is a nuanced and transferable insight for the broader VLM interpretability literature. By adapting OT-based steering to VLAs, the work extends the toolkit for controlling foundation-model-based policies and opens questions about whether similar distribution-structure arguments apply to other multimodal generation settings beyond robotics.