Decoupled Early Exits for Task-Dependent Compute Allocation in Flow-Matching VLAs

2026-09-24

AI Summary

This study addresses the prohibitive computational overhead of flow-matching Vision-Language-Action (VLA) models, which hinders real-time robotic control. We propose a decoupled early-exit framework that pioneers a triaxial joint configuration mechanism encompassing backbone depth, action expert depth, and denoising steps. By integrating a lightweight Exit Transformer with KV cache synthesis techniques, our approach enables task-adaptive computation allocation and optimizes inference paths without retraining the original policy. Evaluated on the LIBERO and Meta-World benchmarks, the proposed method reduces inference latency by 79.2% and FLOPs by 31.8%, while improving average success rate by 5.6%.

Abstract

Flow-matching Vision-Language-Action (VLA) models have emerged as a potential solution for generalist robot control, designed by combining a pretrained Vision-Language Model (VLM) backbone with an action expert that generates continuous robot actions. While these models exhibit impressive capabilities, due to their very high number of parameters, their computational requirements are often prohibitive for robotics control. To mitigate these inefficiencies, existing methods predominantly skip VLM backbone layers with early exits or reduce denoising steps, while leaving action expert depth untouched. We propose a framework that exposes backbone depth $V$, action expert depth $A$, and denoising steps $D$ as three jointly configurable compute axes in a VLA. Starting from a pretrained VLA, we attach lightweight Exit Transformers (ET) at intermediate depths in both the backbone and the action expert, trained to distil the last layer of the policy into each exit. Furthermore, we introduce a KV Cache synthesis mechanism that manages the missing keys and values of the skipped backbone layers, allowing the action expert to exit deeper than the backbone. Finally, we show that the optimal compute budget is task-dependent, with different tasks benefiting from different axes and depths. Notably, our method does not require training the original policy from scratch, and for each exit, it increases the number of parameters by only $2.1\%$ for SmolVLA and $4.1\%$ for $π_{0.5}$. We validate our approach across two flow-matching VLAs (SmolVLA, $π_{0.5}$) and two benchmarks (LIBERO, Meta-World), revealing complementary effects: $V$ and $A$ respectively reduce FLOPs and latency, while $D$ improves both. Our joint configurations $(V,A,D)$ reduce latency by $79.2\%$ and computation (FLOPs) by $31.8\%$, while improving mean success rate by $5.6\%$.

Problem

• Vision-Language-Action models
• computational efficiency
• robotics control
• compute allocation
• flow matching

Innovation

• Decoupled Early Exits
• Flow-Matching VLA
• Exit Transformers
• KV Cache Synthesis
• Compute Allocation