Shallow-π: Knowledge Distillation for Flow-based VLAs

TL;DR

Shallow-π tackles the real-time deployment challenge of flow-based Vision-Language-Action models by aggressively reducing transformer depth in both the vision–language backbone and the diffusion-based action head. It introduces a knowledge-distillation framework with three losses—task, teacher-student velocity matching, and intermediate attention distillation—designed for π-like architectures where the action head densely relies on multi-layer conditioning. Empirically, Shallow-π achieves up to a 70% layer reduction, over 2× speedups, and less than 1% absolute drop in success rate on standard benchmarks, with robust real-world performance on Jetson Orin/Thor across multiple robot platforms. Realistic deployment results show significant latency reductions and improved generalization to unseen perturbations, indicating practical viability for edge robotics without resorting to token pruning or complex routing schemes. The work also outlines limitations (training-time costs) and future directions to combine additional efficiency axes, such as visual token and diffusion-step reductions.

Abstract

The growing demand for real-time robotic deployment necessitates fast and on-device inference for vision-language-action (VLA) models. Within the VLA literature, efficiency has been extensively studied at the token level, such as visual token pruning. In contrast, systematic transformer layer reduction has received limited attention and, to the best of our knowledge, has not been explored for flow-based VLA models under knowledge distillation. In this work, we propose Shallow-pi, a principled knowledge distillation framework that aggressively reduces the transformer depth of both the VLM backbone and the flow-based action head, compressing the model from 18 to 6 layers. Shallow-pi achieves over two times faster inference with less than one percent absolute drop in success rate on standard manipulation benchmarks, establishing state-of-the-art performance among reduced VLA models. Crucially, we validate our approach through industrial-scale real-world experiments on Jetson Orin and Jetson Thor across multiple robot platforms, including humanoid systems, in complex and dynamic manipulation scenarios.

Figures (11)

• Figure 1: Layer reduction strategies and targeted structures in the VLA domain. Previous methods primarily reduce backbone depth only or dynamically skip layers at inference time. In contrast, we propose a systematic knowledge distillation framework that jointly reduces the transformer depth of both the VLM backbone and the action head, which is especially effective for $\pi$-like architectures where the action head mirrors the VLM backbone to receive conditioning information from all layers.
• Figure 2: CUDA inference time as a function of transformer depth and visual token count. Measurements are obtained using the $\pi_{0.5}$ model trained on LIBERO, evaluated on an H100 GPU (left) and Jetson Orin (right).
• Figure 3: (Top) Feature similarity trend along noise level $\tau$. (Bottom) Layer sensitivity analysis of $\pi_{0.5}$ on the LIBERO benchmark. The bar chart shows the decrease in average success rate caused by skipping individual layers.
• Figure 4: Success rate (%) on LIBERO (Spatial, Object, Goal, and 10) as a function of the number of skipped transformer layers.
• Figure 5: Shallow-$\pi$ reduces the transformer depth of the VLM backbone and action head via knowledge distillation, using three loss terms to match ground-truth actions, teacher outputs, and intermediate cross-attention between the backbone and action head.