SnapFlow:通过渐进式 Self-蒸馏 为流量匹配 VLA 生成一步操作
SnapFlow: One-Step Action Generation for Flow-Matching VLAs via Progressive Self-Distillation
摘要
基于流匹配的视觉语言动作 (VLA) 模型(例如 pi0、pi0.5 和 SmolVLA)实现了最先进的通用机器人操作,但它们的迭代去噪(通常为 10 个 ODE 步骤)会引入大量延迟:在现代 GPU 上,仅去噪一项就占了端到端推理时间的 80%。单纯地减少步数是不可靠的,由于单步跳跃的速度场未校准,因此会降低大多数任务的成功率。我们提出了 SnapFlow,一种即插即用的自 蒸馏 方法,可将多步去噪压缩为用于流匹配 VLA 的单个前向传递 (1-NFE)。 SnapFlow 将标准流匹配样本与一致性样本混合在一起,一致性样本的目标是根据模型自身的边际速度预测计算出的两步欧拉捷径速度,避免了条件速度引起的轨迹漂移,正如我们理论上分析的那样。零初始化的目标时间嵌入允许网络在单一架构内的局部速度估计和全局一步生成之间切换。 SnapFlow 不需要外部教师,无需更改架构,并且在单个 GPU 上训练大约 12 小时。我们在跨越 6 倍参数范围、具有相同超参数的两个 VLA 架构上进行验证:在 pi0.5 (3B) 上,跨越四个 LIBERO 套件(40 个任务,400 个片段),SnapFlow 实现了 98.75% 的平均成功率,与 10 步教师的 97.75% 匹配,并稍稍超过了它,降噪加速提高了 9.6 倍,端到端延迟从 274 毫秒减少至 83 毫秒;在 SmolVLA (500M) 上,它使 MSE 降低了 8.3%,端到端加速提高了 3.56 倍。对长期任务的行动步骤扫描表明,SnapFlow 在执行范围内保持了优势,在 n_act=5 时实现了 93%,而基线仅达到 90%。 SnapFlow 与层 蒸馏 和词元修剪方法正交,从而实现组合加速。