论文笔记：Shallow-π：Knowledge Distillation for Flow-based VLAs

背景

π0/π0.5 这类流匹配 VLA 用「VLM 骨干 + DiT 动作头」的组合生成连续动作，效果领先，但计算结构有个特殊之处：动作头的层数通常与 VLM 骨干对齐（π0.5 中都是多层堆叠），且每一层都要通过交叉注意力消费 VLM 注入的视觉-语言特征。推理时还要在去噪步之间反复调用动作头。于是部署瓶颈有两个来源：VLM 骨干太大、动作头太深且被反复执行。

VLA 效率研究此前主要沿三条线走：视觉 token 剪枝/缓存、扩散步数压缩、transformer 层数削减。其中层数削减这条线上的已有工作大多是测试期层跳过（DeeR-VLA、EfficientVLA 用相邻层特征相似度决定跳层；MoLE-VLA 用路由机制条件性跳层），有三个硬伤：层是动态跳过的，整个模型仍得驻留显存；不同输入走不同执行路径，阻碍批处理与 ahead-of-time 编译；路由的训练要靠代理梯度估计器，优化不稳定。另一条线是换小骨干（TinyVLA、SmolVLA），但往往需要从头训练 VLM 骨干，与大规模预训练模型的兼容性差，而且不减少动作头深度——对推理时要反复执行动作头的流匹配模型来说尤其致命。此外这些方法大多只在仿真或简单单臂任务上验证。

一个容易被忽略的物理事实是：在现代 GPU 上，token 级计算高度并行（减 token 的加速比有限），而 transformer 层是串行执行的，其开销直接累加到墙钟时间上；流匹配 VLA 的动作头又在每个去噪步被调用一次，深度问题被进一步放大（论文 Figure 2 在 H100 与 Jetson Orin 上实测了深度 vs token 数的延迟对比）。

本文 Shallow-π（Samsung Research，arXiv:2601.20262，项目主页，开源代码 icsl-Jeon/openpi）就是在「为什么不能简单跳层」这个问题上展开的：先用分析说明测试期跳层对流匹配 VLA 从根本上不可靠，再提出一个同时压缩 VLM 骨干与动作头的蒸馏框架，把 18 层结构压到 6 层，并搬到 Jetson Orin/Thor 上做了真机验证。

中心思想

用知识蒸馏做「结构性」的层数削减，同时压 VLM 骨干与动作头，而不是测试期动态跳层——因为 π 类架构中动作头每一层都从骨干接收条件信息，必须整条通路一起变浅。

先看作者对「测试期跳层为什么不行」的两个实证：

1. 特征相似度随去噪时间 $\tau$ 变化，固定阈值很脆：相邻层余弦相似度剖面在不同 $\tau$ 下差异很大（$\tau =0$ 与 $\tau =1$ 时早晚层的趋势相反），简单的单调阈值策略在所有深度上不成立。
2. 相似度与功能重要性脱钩：尽管层 1→2 的相似度高于层 16→17，但跳过层 1–2 导致的成功率下降远大于跳过 16→17。若把层敏感性当作 oracle 按最低敏感度顺序渐进删层，删到三 层以上成功率就崩塌（论文 Figure 4）。

结论是：流匹配 VLA 的层功能与去噪动力学深度耦合，测试期启发式（相似度阈值或路由）不可靠，必须训练期蒸馏出真正的浅结构。初始化用 TinyBERT 式均匀隔层采样教师权重，作者也试过基于敏感性分析挑层，发现在足够训练步数下没有额外收益——均匀隔层就够了。

蒸馏侧的核心是三个互补损失：流匹配任务损失 + 教师速度场蒸馏 + 交叉注意力蒸馏，其中注意力蒸馏只对「动作 token 对视觉-语言 KV」的交叉注意力做对齐，且只在一层中间层做（详见方法）。

主要方法

1. 隔层初始化

学生从教师均匀采样层初始化：VLM 骨干与动作头同时减层，例如 18 层压到 9 层（1/2）或 6 层（1/3），对应教师层 { 1 , 3 , 5 , …   } \{1,3,5,\dots\} {1,3,5,…} 或 { 1 , 4 , 7 , …   } \{1,4,7,\dots\} {1,4,7,…} 这类均匀网格。这一步保证学生起点就在教师函数空间的"低分辨率版本"上，而不是随机初始化。

2. 三元损失

（a）流匹配任务损失（ground-truth 监督）：标准 FM 目标，让学生预测真值速度场

$${\mathcal{L}}_{task}=\mathbb{E}[\Vert v_{\theta }(\cdot )-u{\Vert }^2]$$

（b）教师速度场蒸馏（teacher trajectory imitation）：学生的速度预测对齐教师的预测，提供教师生成的引导信息

$${\mathcal{L}}_{kd}=\mathbb{E}[\Vert v_{\theta }(\cdot )-v_{\varphi }(\cdot ){\Vert }_2^2]$$

（c）交叉注意力蒸馏：在单一中间层上，对齐教师与学生的「动作 token 查询 × 视觉-语言 KV」交叉注意力分布

$${\mathcal{L}}_{attn}=\mathbb{E}\left[KL\right({\mathrm{Attn}}_{\varphi }^{a\to vl}\Vert {\mathrm{Attn}}_{\theta }^{a\to vl}\left)\right],{\mathrm{Attn}}^{a\to vl}=\mathrm{softmax}\left(Q_a{K}_{vl}^{\top }\right)$$

两个关键设计区别于 LLM/VLM 里的经典注意力蒸馏：

• 只对动作 token 的注意力做匹配，不对视觉-语言 token 之间的注意力（${\mathrm{Attn}}^{vl\to vl}$）做匹配。理由是流匹配 VLA 中只有动作 token 定义生成策略，VLM 的视觉-语言 token 只是条件上下文；对非生成的骨干 token 蒸馏注意力会过度约束学生、干扰预训练表示，却不提升控制保真度。实验验证：在全部 token 上做注意力蒸馏时，L9/L6/L4 三个深度的成功率分别为 0.0 / 61.5 / 0.0——训练崩掉。
• 只在一层（中间层）做，不做逐层匹配（借鉴 Align-KD）。理由：学生底层直接复制自教师，早期层已对齐、监督冗余；末层的输出匹配已由 ${\mathcal{L}}_{task}$ 和 ${\mathcal{L}}_{kd}$ 承担。消融确认中间层最优（L9/L6/L4：96.8/94.6/94.2 vs 初始层 94.4/93.9/91.0、末层 95.6/94.1/92.8）。

3. 实验结论

仿真（LIBERO，四套件），教师与学生均为 30k 步训练：

模型	Spatial	Object	Goal	Long-10	平均	FLOPs (T)	CUDA 时延 (ms)
π0（教师，10 步）	97	97	93	92	95	2.93	22.6
π0.5（教师，10 步）	98	96	97	93	96	3.39	25.5
SmolVLA（小骨干路线）	90	96	92	71	87	0.50	26.0
π0.5-L9（1/2 深度）	99	98	97	93	97	1.82	14.8
π0.5-L6（1/3 深度）	98	96	94	90	95	1.30	11.3
π0-L9	98	98	95	87	95	1.62	13.5
π0-L6	98	95	93	85	94	1.18	10.5

• 同时压骨干+动作头的效果显著优于单独压骨干：π0.5-L6 相比 π0.5 教师减少 62% 层数（18→6 后跨 VLM+动作头整体），FLOPs 降 62%、CUDA 时延降 56%，成功率仅从 96 降至 95（−1 pp）。
• 蒸馏大模型优于从头训小骨干：同样推理预算下，Shallow-π 各深度全面超过 SmolVLA（87），且时延更低——SmolVLA 的动作头仍有 16 层，CUDA 时延 26 ms 反而高于 π0.5-L6 的 11.3 ms。训练曲线还显示 SmolVLA 的 FM 损失约为 Shallow-π 的两倍，对应更噪、更不精确的动作。
• 真机（Jetson Orin/Thor，ALOHA 与 RB-Y1 人形双平台）：6 层学生全面超过各自教师，如 Peg-in-hole 10/10 vs 7/10、Recycle 17/20 vs 12/20；端到端时延 110 ms vs 364 ms（ALOHA 上快 200 ms 以上，转台边缘线速约 0.8 cm/s 下相当于少走 2 cm 的开环漂移）。慢推理→开环执行陈旧观测→不可恢复失败，这是学生真机反超的机制解释。
• 未见扰动：圆柱初始位置偏移 3 cm、垃圾箱位移 10 cm+ 的场景下学生依旧更稳，因为更频繁地吸收新观测。

特点（优缺点）

优点

• 结构性的深度削减：与跳层不同，层数是真的删掉了，模型常驻显存随层数同步缩小，也没有动态控制流，方便 AOT 编译与批处理。
• 「同时压骨干+动作头」抓住了 π 类架构的要害：动作头每层都要从骨干拿条件信息，只压其中一侧都会留下短板；表 2 显示联合压缩在精度几乎不降的前提下把 FLOPs/时延都砍半以上。
• 注意力蒸馏的两个设计有清晰的架构动机且有数据支撑（只对动作 token、只在中间层），全 token 匹配直接训练崩溃这一对照很有说服力。
• 真机验证的稀缺价值：在 Jetson Orin/Thor、双平台、动态场景（移动孔位、移动箱子）下做了工业级规模的评测（Peg-in-hole 92 回合、Recycle 2600 回合），且给出「时延→开环漂移→失败」的机制分析，这在 VLA 加速论文里很少见。
• 蒸馏比从头训小模型更好：同预算下精度与时延双赢 SmolVLA。

缺点

• 训练期成本高：教师和学生要同时驻留显存，论文自己把这条列为首要 limitation；蒸馏 100k 步 × batch 128 的算力开销不小，只提了未来方向（冻结部分组件、样本筛选）没有给出解法。
• 一次训练绑定一个深度：不像跳层/路由方法可以推理期动态权衡，L9/L6/L4 是三个独立蒸馏出来的模型，要按部署预算分别训练。
• 隔层初始化的层映射是均匀的：虽然论文试过基于敏感性的非均匀选层且未见收益，但均匀采样对层数非整倍数（如 18→5）或层功能高度不均匀的模型是否仍合适未讨论。
• 对 π 类架构的适配性强，通用性依赖该结构：整套设计（动作头逐层接收骨干特征、交叉注意力蒸馏只对动作 token）都建立在 π 类流匹配 VLA 的结构假设上，迁移到其他 VLA 范式（自回归 token 解码、纯扩散 U-Net 头）需要重新设计。
• 标注为工程观察：官方开源实现（GitHub icsl-Jeon/openpi 中 DistilledPI0Pytorch.forward()）目前只实现了 ${\mathcal{L}}_{task}$ 和 ${\mathcal{L}}_{kd}$ 两个损失，缺少论文 5.2 节的交叉注意力蒸馏 ${\mathcal{L}}_{attn}$——教师前向只返回最终速度预测、不导出中间注意力图，也没有任何注意力对齐损失项。结合消融表（L9 96.2→96.8、L6 93.9→94.6），注意力蒸馏的边际贡献约 0.5–0.7 pp，不是主导项，但这仍意味着开源代码与论文方法描述存在偏差，复现时需要注意（这是对开源代码的观察，非论文原文内容）。
• 评测覆盖仍以仿真为主：真机虽好但教师基线本身较弱（π0/π0.5 未微调到饱和），且未报告 OOD 泛化与多 seed 方差。