Source-linked AI summary
Shallow-π: Knowledge Distillation for Flow-based VLAs
Boseong Jeon, Yunho Choi, Taehan Kim
TL;DR
Real-time deployment of flow-based VLAs is limited by the cost of large backbones, deep action heads, and iterative diffusion. Shallow-π distills both transformer components into a shallower model, preserving performance while achieving faster inference and real-world edge deployment. The paper reports improved efficiency and generalization under practical manipulation settings, while noting higher training-time costs than layer skipping.
Problem
Flow-based VLAs face challenging edge deployment because large VLM backbones, deep action heads, and iterative diffusion increase inference cost.
Method
Shallow-π uses knowledge distillation to jointly reduce transformer depth in the VLM backbone and flow-based action head.
Results
Shallow-π reduces FLOPs and CUDA inference time by more than two times while keeping success rates within 1% of teacher models.
Takeaways & Limitations
The distilled model preserves teacher performance and generalization while reducing computation for practical edge-device deployment.
Takeaways & Limitations
Knowledge distillation requires higher training-time computation because teacher and student models must be loaded simultaneously.
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
Flow-based VLAs face substantial inference costs because they combine large VLM backbones with iterative diffusion action heads. Shallow-π addresses this by jointly distilling both components into shallower transformer models.
- Flow-based VLAs combine large VLM backbones, diffusion-based action heads, and iterative inference steps, challenging real-time edge deployment.These costs motivate reducing transformer depth in π-like architectures.
- Prior efficiency methods reduce visual tokens, diffusion steps, transformer depth, attention computation, or model precision.
- Existing layer-skipping methods dynamically bypass computation but generally retain the full model in GPU memory.
- Smaller-backbone approaches can require training from scratch and leave the action-head depth unchanged.This is particularly costly for flow-based models with repeated denoising steps.
- Shallow-π jointly distills the VLM backbone and action head, achieving up to 70% layer reduction while preserving layer-wise feature transfer.
- A 6-layer Shallow-π model reaches almost 10 Hz end-to-end inference on Jetson Orin without graph-level optimization or runtime conversion.
2. RELATED WORKS
VLA efficiency research spans token, diffusion-step, depth, attention, quantization, and optimization techniques. Depth reduction is attractive for flow-based models, but existing dynamic or backbone-only approaches retain important deployment and coverage limitations.
- VLA efficiency methods target visual tokens, diffusion steps, transformer depth, attention computation, quantization, and graph optimization.
- Token pruning can reduce floating-point operations, but wall-clock gains depend on workload because accelerators parallelize token computation.
- Sequential transformer layers accumulate directly in wall-clock time, while repeated action-head execution makes depth critical for real-time flow-based inference.
- Similarity-based layer skipping can miss semantic layer roles and has mainly been evaluated on decoder-only models or VLM backbones.
- Router-based skipping targets only the VLM backbone, retains the full model in memory, and can hinder batching and embedded graph optimization.
- Knowledge distillation is established for compact language and vision-language models, where distilled students can outperform similarly sized models trained from scratch.
3. Preliminaries
Flow-based VLAs learn conditional vector fields that transport noise toward action trajectories using multimodal observations and language. Their architecture combines a VLM backbone with an action transformer, while layer reduction offers larger latency gains than token reduction.
- Flow-based VLAs construct noisy action interpolations and train a velocity field to predict the transport direction toward ground-truth actions.
- The model contains a VLM transformer for multimodal tokens and an action transformer processing noisy actions with state information.
- Vision-language key–value pairs can be cached across diffusion steps because action tokens attend to fixed multimodal representations.
- CUDA inference measurements compare transformer depth and visual token count for π0.5 on H100 and Jetson Orin.
- The latency study uses torch.compile and evaluates both high- and low-capability hardware.
- Reducing transformer layers produces a substantially larger inference-time decrease than reducing visual tokens.Token computation is highly parallelized, whereas transformer layers execute sequentially.
4. Why Layer Skipping Is Insufficient
Layer skipping is unreliable for π-like flow-based VLAs because similarity varies with denoising noise and does not reliably indicate functional importance. Removing layers progressively can therefore collapse success rates.
- Figure 3 evaluates layer-wise cosine similarity across denoising timesteps and measures success-rate decreases when individual layers are skipped.
- π-like flow-based VLAs inject visual-language features at every transformer layer, motivating tests of whether layer-skipping methods transfer to this architecture.
- Similarity profiles vary substantially with noise level, preventing a simple monotonic thresholding rule across layer depths.
- Higher inter-layer similarity does not reliably indicate lower functional importance, since skipping one highly similar layer can reduce success more than skipping another.
- Removing layers in lowest-sensitivity order causes success rates to collapse once more than three layers are removed.
- LIBERO success rate is evaluated across Spatial, Object, Goal, and 10 task groups as skipped-layer count increases.
5. Knowledge Distillation for Shallow Layers
Shallow-π trains a reduced-depth student by combining task supervision with teacher imitation and intermediate cross-attention transfer. Ablations identify the composite objective and middle-layer attention distillation as the strongest configuration.
- Student Initialization: The student uniformly subsamples layers from both the vision–language backbone and action head before distillation.This TinyBERT-style initialization targets substantially fewer transformer layers while preserving action-generation performance.
- Distillation Objectives: Three losses supervise ground-truth velocity, teacher-predicted velocity, and intermediate cross-attention distributions.Together, they align the student with both the teacher’s outputs and internal multimodal representations.
- Attention Distillation: Attention distillation aligns action-query attention with vision–language key–value pairs at an intermediate transformer layer.The KL divergence is evaluated across vision–language tokens for each action token.
- Attention Distillation: The method distills attention only for generated action tokens rather than the full token set.Visual and language tokens provide conditioning context, so matching their attention is described as over-constraining the student.
- Ablations: The composite loss with middle-layer attention distillation yields the best ablation performance.The ablations use the same teacher and train each student for 30K steps with batch size 64.
6. Experiment
Shallow-π is evaluated on simulation and real-world manipulation tasks spanning dynamic scenes and articulated robot platforms. The experiments report near-teacher success rates with substantially lower computation and latency, alongside robustness to unseen spatial perturbations.
- Experimental setup: The evaluation covers simulation and real-world task suites involving dynamic manipulation, articulated robots, and coordinated hand–torso movements.Real-world evaluations use onboard sensing and edge devices including Jetson Orin and Jetson Thor.
- Simulation benchmark: The LIBERO simulation benchmark trains π0 and π0.5 teachers and distills students under the same two-image input setup.The benchmark uses third-person and wrist-camera images, with a third image omitted rather than replaced by a black image.
- Simulation benchmark: Success rates remain within 1% of the teachers while FLOPs and CUDA inference time decrease by more than 2×.The comparison is reported for the distilled models on LIBERO, with computation measured on an NVIDIA H100.
- Real-world experiments: The real-world experiments use 6-layer students deployed on Jetson Orin and Jetson Thor across ALOHA and RB-Y1 platforms.The deployment uses action chunks of 50 in a 30 Hz control loop, with seven actions executed before requesting the next inference.
- Real-world experiments: Shallow-π consistently achieves better real-world performance across the evaluated dynamic, complex, and unseen tasks.The task suite includes peg insertion, foam placement, apple scooping, bean pouring, and two lid-and-cylinder variants.
- Latency and control: More than 200 ms of ALOHA inference-time reduction corresponds to approximately 6 frames and over 2 cm of additional teacher open-loop end-effector translation.The reported 0.8 cm/s turntable-edge speed links prolonged open-loop motion to degraded placement precision.
- Robustness and generalization: Over 80% success is achieved on the recycle task, which requires grasping objects across diverse poses.The distilled model also shows improved robustness under shifted cylinder and trash-bin positions by incorporating updated visual observations more frequently.
7. CONCLUSIONS
Shallow-π is presented as a knowledge-distillation framework for reducing computation in flow-based VLA models while preserving teacher performance and generalization. The authors also identify higher training-time cost as a limitation and propose further efficiency improvements.
- Conclusion: Shallow-π injects conditioning information at all intermediate layers while aggressively reducing transformer depth in flow-based VLA models.The framework jointly targets the VLM backbone and action head.
- Conclusion: The distilled model preserves the teacher’s performance and generalization capability while significantly reducing computation.The conclusion also reports better denoising and precision than models with a small backbone.
- Conclusion: Real-world deployment on edge devices demonstrates reliable performance under practical latency constraints.This conclusion follows the reported validation across robot platforms and edge hardware.
- Limitations: Knowledge distillation requires higher training-time computation because teacher and student models must be loaded simultaneously.The authors suggest selectively freezing components and filtering informative training samples to reduce VRAM consumption during distillation.
- Future work: Future work will combine the approach with visual-token reduction and diffusion-step reduction to improve inference throughput.These efficiency axes are identified as complementary directions rather than as results established by the current work.