Shallow-π：面向流式VLA模型的知识蒸馏框架

通过知识蒸馏将VLA模型从18层压缩到6层，实现2倍推理加速且精度损失<1%

前置知识

Vision-Language-Action (VLA) 模型

VLA是一种端到端的机器人基础模型架构，它将视觉-语言模型(VLM)作为骨干网络处理图像和语言指令，然后将提取的多模态特征注入到动作生成模块中作为条件信号。这种架构使机器人能够通过单一模型执行多样化的操作任务，实现从感知到动作的端到端映射。典型的VLA模型如π0和π0.5采用了这种设计，其中VLM骨干和动作头紧密耦合，动作头在每一层都从VLM骨干接收特征注入。

本文的核心目标就是压缩这类模型的transformer深度，理解VLA的架构设计是理解本文方法的前提

Flow-based 扩散模型

Flow-based模型（如flow matching）通过学习一个条件向量场来将噪声分布传输到目标动作轨迹分布。给定观测o、语言指令l和真实动作a，模型构建带噪插值 a_τ = τ·a + (1-τ)·ε，其中ε服从标准正态分布，τ在0到1之间。模型通过flow matching目标训练来预测目标速度u = a - ε，损失函数为预测速度与目标速度差的L2范数期望。推理时从τ=0逐步积分到τ=1完成去噪。

本文方法专门针对flow-based VLA设计，理解flow matching的训练和推理过程对于理解蒸馏目标函数至关重要

知识蒸馏 (Knowledge Distillation)

知识蒸馏是一种模型压缩技术，通过训练一个较小的学生模型来模仿较大的教师模型的行为。在传统蒸馏中，学生模型不仅学习真实标签（硬标签），还学习教师模型的输出分布（软标签）。这使得学生模型能够继承教师模型的知识，通常比从头训练同等大小的模型表现更好。本文将蒸馏扩展到VLA领域，同时压缩VLM骨干和动作头。

这是本文的核心方法论基础，理解蒸馏的基本原理有助于理解本文设计的三种损失函数

Diffusion Transformer (DiT)

Diffusion Transformer是一种将Transformer架构用于扩散模型的架构，由Peebles和Xie在2023年提出。在VLA模型中，DiT被用作动作头来利用flow-matching的强生成能力。它由多层transformer层组成，每层包含自注意力机制和前馈网络。在π0.5等模型中，动作头通常包含18层transformer，推理时需要迭代执行多次去噪步骤，使得计算成本成为实时部署的主要瓶颈。

理解DiT架构有助于理解为什么层压缩对推理延迟有如此显著的影响——层是顺序执行的，而token计算可以并行化

层跳过 (Layer Skipping)

层跳过是一种在推理时动态跳过冗余transformer层的方法。现有方法如DeeR-VLA和EfficientVLA基于层间特征相似度（如余弦相似度）来决定哪些层可以跳过，而MoLE-VLA则使用学习的路由器机制进行条件跳过。然而，这些方法要求完整模型常驻GPU内存，因为层是动态跳过而非结构化移除。此外，相似度指标无法捕捉层的语义角色，且现有评估主要针对模拟环境。

本文的核心论点之一就是层跳过方法不足以解决flow-based VLA的效率问题，理解其局限性有助于理解本文动机

LIBERO基准测试

LIBERO是一个用于评估机器人操作策略的基准测试套件，包含多个任务类别：Spatial（空间推理）、Object（物体操作）、Goal（目标导向）和Long（长序列任务，包含10个子任务）。它用于评估VLA模型在模拟环境中的操作成功率。本文使用LIBERO作为主要的模拟评估平台，教师模型π0和π0.5在此基准上训练和评估。

这是本文的核心评估基准，理解其任务结构有助于解读实验结果

研究动机

随着实时机器人部署需求的增长，VLA模型的推理效率成为关键瓶颈。Flow-based VLA模型（如π0、π0.5、GR00T、CogACT）结合了大型VLM骨干和基于扩散transformer的动作头，需要数十层transformer层，并且在推理时需要多次迭代去噪步骤。在Jetson Orin等边缘设备上，这种计算成本使得实时部署极具挑战性。本文的基准测试显示，在H100 GPU上，仅减少视觉token（从144减至64）只能带来有限的延迟改善，因为现代GPU能高效并行化token级计算；但减少transformer层（从18层减至4层）能带来显著的延迟降低，因为层是顺序执行的。具体而言，在H100上，π0.5模型的完整版本需要25.5ms CUDA时间，而6层版本仅需11.3ms；在Jetson Orin上，这一差距更为显著。这表明transformer层深度是实时推理的关键瓶颈。

本文的目标是本文提出Shallow-π，一个专为π-like flow-based VLA架构设计的知识蒸馏框架，目标是通过联合压缩VLM骨干和动作头的transformer深度（从18层压缩到6层，实现70%的层减少），同时保持任务性能。具体目标包括：（1）实现超过2倍的推理加速；（2）保持与教师模型相近的成功率（差距<1%）；（3）在边缘设备（Jetson Orin和Jetson Thor）上验证实际部署效果；（4）在复杂动态操作场景和人形机器人平台上验证有效性。

与已有工作不同的是，本文的独特切入角度体现在三个层面。首先，与现有层跳过方法不同，本文采用结构化的层移除而非动态跳过，这意味着压缩后的模型不再需要完整模型常驻内存，且支持高效批处理推理和提前编译优化。其次，与现有小骨干方法（如SmolVLA）不同，本文通过蒸馏大容量教师模型的知识而非从头训练小模型，获得了更好的性能。最后，本文首次为π-like架构设计了专门的蒸馏策略——π-like架构的特点是动作头在每一层都接收VLM骨干的特征注入，这要求蒸馏时同时压缩骨干和动作头，而现有方法要么只压缩骨干，要么只在骨干上做层跳过。此外，本文首次在真实机器人部署中验证了层压缩方法的有效性。

核心方法

Shallow-π的整体思路是：给定一个预训练的flow-based VLA教师模型，通过知识蒸馏训练一个具有更少transformer层的学生模型。直觉上，虽然直接移除大量层会导致性能崩溃（实验显示移除超过3层后成功率急剧下降），但通过精心设计的蒸馏损失函数，可以让浅层学生模型学习到教师模型的去噪行为和跨模态特征对齐。技术路线上，首先使用TinyBERT风格的均匀子采样策略初始化浅层学生（直接复制教师模型的底部层），然后使用三种互补的损失函数进行训练：任务损失L_task监督学生预测真实动作，知识蒸馏损失L_kd让学生匹配教师的预测输出，注意力蒸馏损失L_attn在中间层对齐动作token对视觉-语言token的注意力分布。

本文的核心创新在于为π-like flow-based VLA架构设计了专门的蒸馏策略，这与已有方法有本质区别。首先，与层跳过方法不同，本文进行结构化层移除，压缩后的模型是独立的小模型，不需要完整模型常驻内存。其次，注意力蒸馏损失L_attn只针对动作token（而非所有token），这是因为flow-based VLA中只有动作token参与去噪生成，而视觉-语言token仅作为条件上下文。如果对所有token进行注意力匹配，会导致训练不稳定和接近0的成功率（如表1c所示）。第三，本文发现注意力蒸馏在中间层效果最佳——早期层由于学生直接复制教师底部层，特征已经对齐；最终层则已有L_task和L_kd的输出匹配约束。这种设计与传统NLP蒸馏中对全token集做注意力匹配的做法有本质不同。

方法步骤详情

Shallow-π的完整方法包含以下步骤：（1）学生模型初始化：给定预训练教师模型（如18层的π0.5），使用均匀子采样策略选择层子集来初始化浅层学生（如6层）。具体做法是每隔k层选取一层，复制其参数作为学生模型的对应层。实验表明，基于敏感度分析选择初始化层并无额外收益，只要训练步数足够。（2）联合训练损失：学生模型使用三种损失的加权组合训练。任务损失L_task监督学生预测真实速度u = a - ε。知识蒸馏损失L_kd让学生匹配教师的预测速度。注意力蒸馏损失L_attn在中间transformer层处对齐动作token对视觉-语言token的注意力分布。（3）训练配置：教师和学生同时加载，使用batch size 64训练30K步（LIBERO）或batch size 128训练100K步（真实世界）。（4）推理部署：压缩后的学生模型可独立部署在边缘设备上，无需教师模型。

技术新颖性

Shallow-π的技术新颖性体现在多个方面。首先，这是首个为π-like flow-based VLA架构设计的知识蒸馏框架，该架构的独特性在于动作头在每一层都从VLM骨干接收特征注入，这要求蒸馏时同时压缩两个组件。其次，注意力蒸馏的设计极具创新性——只蒸馏动作token对视觉-语言token的注意力，而非传统的全token注意力匹配，这一设计是基于对flow-based VLA架构特性的深刻理解。第三，实验揭示了层跳过方法在flow-based VLA上的根本局限：相似度profile随噪声水平τ变化剧烈，相似度不能预测功能重要性（层1-2的相似度高于层16-17，但跳过前者导致更大性能下降），且使用敏感度分析作为oracle的路由机制在移除超过3层后就会崩溃。第四，本文首次在工业级真实机器人平台上验证了层压缩方法，在Jetson Orin上实现了近10Hz的端到端推理。

**Figure 2**: CUDA inference time as a function of transformer depth and visual token count

**Figure 3**: Feature similarity trend along noise level τ and layer sensitivity analysis

**Figure 4**: Success rate (%) on LIBERO as a function of the number of skipped transformer layers

**Figure 5**: Shallow-π reduces the transformer depth via knowledge distillation

实验结果

本文的实验结果全面验证了Shallow-π的有效性。在LIBERO模拟基准上（表2），蒸馏后的π0.5-L6（6层）模型在四个任务上分别达到99%、98%、97%、93%的成功率，平均97%，仅比教师模型π0.5的96%低1个百分点（实际上在某些子任务上甚至略高）。同时，FLOPs从3.39T降至1.82T（减少46%），CUDA推理时间从25.5ms降至14.8ms（减少42%）。更激进的4层版本π0.5-L4保持94.2%的成功率。与小骨干方法SmolVLA相比，Shallow-π显著优于其87%的平均成功率，同时推理延迟更低（11.3ms vs 26.0ms）。在真实世界实验中（表3），Shallow-π0在ALOHA平台上的动态任务中表现优异：Peg in hole任务达到10/10成功率（教师π0为7/10），Insert foam为7/10（教师5/10），Scoop apple为9/10（教师6/10），Pour beans为8/10（教师5/10）。端到端推理时间从364ms降至110ms，实现3.3倍加速。在RB-Y1人形机器人平台上，Shallow-π0.5在Recycle任务上达到17/20（教师12/20），Open lid & peg (A)为7/10（教师5/10），Open lid & peg (B)为5/5（教师1/5），E2E推理从130ms降至78ms。在未见环境泛化测试中，Shallow-π在Peg in hole（π0）上达到3/5（教师0/5），Recycle（π0.5）为15/20（教师8/20），展示了更强的鲁棒性和泛化能力。

**Table 2**: Success rates on LIBERO benchmark and computation measured on H100

**Table 3**: Task performance across dynamic, complex, and unseen settings

**Figure 6**: Experiment task suites for Shallow-π

**Figure 7**: Degrees of freedom (DoFs) and camera configurations for the robot platforms

**Figure 8**: Skill sequences required by the tasks of the hand-typed robot (RB-Y1)

**Figure 9**: Flow matching loss vs training step and samples of chunk to compare motion quality

**Figure 10**: Histogram of the per-frame average translation of the right-arm end effector

**Figure 11**: Experimental snapshots illustrating open-loop failures of the teacher model

查看结构化数据

任务	指标	本文	基线	提升
LIBERO Spatial	Success Rate (%)	π0.5-L6: 98	π0.5: 98	持平，同时FLOPs减少46%
LIBERO Object	Success Rate (%)	π0.5-L6: 96	π0.5: 96	持平
LIBERO Goal	Success Rate (%)	π0.5-L6: 94	π0.5: 97	下降3%
LIBERO Long (10)	Success Rate (%)	π0.5-L6: 90	π0.5: 93	下降3%
LIBERO Average	Success Rate (%)	π0.5-L6: 95	π0.5: 96, SmolVLA: 87	比教师低1%，比SmolVLA高8%
Peg in hole (动态)	Success Rate	Shallow-π0: 10/10	π0: 7/10, SmolVLA: 0/10	比教师高30%，SmolVLA完全失败
Recycle (RB-Y1)	Success Rate	Shallow-π0.5: 17/20	π0.5: 12/20	提升25%
ALOHA E2E推理	Latency (ms)	110	π0: 364, SmolVLA: 230	比教师快3.3倍，比SmolVLA快2.1倍
RB-Y1 E2E推理	Latency (ms)	78	π0.5: 130	快1.7倍

局限与改进

本文的局限性包括：（1）训练成本高：知识蒸馏需要同时加载教师和学生模型，相比层跳过方法增加了训练时的计算和显存开销。作者提出未来可通过选择性冻结模型组件来减少VRAM消耗。（2）蒸馏数据依赖：蒸馏效果可能依赖于训练数据的质量和多样性，作者建议未来研究如何筛选最具信息量的训练样本。（3）层减少有上限：实验显示，当层减少超过一定程度（如从18层减至4层）时，性能下降开始变得明显，特别是在长序列任务上（π0.5-L4在Long任务上从93%降至约88%）。（4）仅评估了均匀子采样：本文使用均匀间隔选择层进行初始化，未探索更复杂的层选择策略。（5）单一压缩维度：本文仅关注transformer深度压缩，未结合视觉token剪枝或扩散步数减少等其他效率优化方向。（6）模拟与真实世界评估的差距：LIBERO模拟结果中Shallow-π在某些子任务上略低于教师，但真实世界实验中反而优于教师，这可能与推理延迟改善有关。

独立分析的弱点

本文的独立弱点分析如下：（1）损失函数权重未详细说明：论文中三种损失L_task、L_kd、L_attn的权重比例未在正文中明确给出，这影响了方法的可复现性。建议未来工作提供详细的超参数敏感度分析。（2）注意力蒸馏仅在单一中间层：本文选择在中间层进行注意力蒸馏，并声称这是最优位置，但实验仅对比了初始层、中间层和较后层三种选择。更细致的层选择策略（如多层蒸馏或自适应选择）可能进一步提升性能。（3）与层跳过方法的对比不完全公平：本文批评层跳过方法需要完整模型常驻内存，但未讨论层跳过方法可以与量化等技术结合使用的场景。（4）真实世界任务多样性有限：虽然本文在两个机器人平台上进行了评估，但任务主要集中在桌面操作（抓取、插入、倾倒），未涉及移动操作、接触丰富的任务或长时间任务。（5）缺乏与更多蒸馏基线的对比：本文主要与从头训练的小模型（SmolVLA）对比，但未与其他蒸馏方法进行系统对比。

未来方向

基于本文成果，未来研究方向包括：（1）多维度效率优化组合：将层压缩与视觉token剪枝（如CogVLA、LightVLA）和扩散步数减少（如一致性蒸馏）结合，实现更大的综合加速。作者在结论中明确提到了这一方向。（2）选择性组件冻结：为减少蒸馏训练的VRAM消耗，研究如何选择性冻结VLM骨干或动作头的某些层，仅蒸馏关键组件。（3）自适应层压缩：探索非均匀的层压缩策略，根据层的重要性分析保留更多关键层，移除更多冗余层。（4）跨具身泛化：将Shallow-π应用于更多样化的机器人平台和任务，包括双足人形机器人的全身控制、移动操作等。（5）与图级优化结合：本文未使用torch.compile等图级优化技术，将层压缩与编译器优化结合可能实现更大的加速。（6）在线蒸馏：探索在真实机器人交互中进行在线蒸馏，使学生模型能够持续从教师模型学习。

复现评估

本文的复现评估如下：（1）开源情况：论文提供了项目页面https://icsl-jeon.github.io/shallow-pi/，但正文未明确说明代码是否开源。（2）数据可用性：LIBERO是公开基准，可直接获取。真实世界数据由作者自行采集，可能不公开。训练数据使用2个摄像头输入（第三人称和腕部摄像头），移除了原始实现中的第三个摄像头输入。（3）算力需求：教师模型训练需要较大算力（使用H100 GPU），但学生模型训练相对轻量。在Jetson Orin（边缘设备）上实现了近10Hz推理，表明部署门槛较低。（4）复现难度：中等。核心蒸馏损失函数设计清晰，但需要预训练的VLA教师模型（如π0或π0.5）。作者基于官方代码进行了修改，使用batch size 128训练100K步（真实世界）或64训练30K步（LIBERO）。（5）关键超参数：action chunk size为50，控制循环30Hz，每步执行7个动作后请求下一次推理，使用时间集成确保平滑执行。

Shallow-π：面向流式VLA模型的知识蒸馏框架 Shallow-π: Knowledge Distillation for Flow-based VLAs