SnapFlow: 프로그레시브 셀프-디스틸레이션을 통한 플로우-매칭 VLA의 단일 단계 액션 생성

SnapFlow: One-Step Action Generation for Flow-Matching VLAs via Progressive Self-Distillation

pi0, pi0.5, SmolVLA와 같은 플로우-매칭 기반의 비전-언어-액션(VLA) 모델은 최첨단 수준의 일반적인 로봇 조작 성능을 보여주지만, 일반적으로 10단계의 ODE(상미분방정식) 기반 디노이징 과정은 상당한 지연 시간을 유발합니다. 최신 GPU에서 디노이징 과정만 전체 추론 시간의 80%를 차지합니다. 단순히 디노이징 단계를 줄이는 것은 대부분의 작업에서 성공률 저하를 야기하며, 이는 단일 단계 점프 시 속도장이 제대로 보정되지 않기 때문입니다. 본 논문에서는 플로우-매칭 VLA 모델을 위한 플러그 앤 플레이 방식의 셀프-디스틸레이션 방법인 SnapFlow를 제안합니다. SnapFlow는 다단계 디노이징을 단일 패스(1-NFE)로 압축하여 처리하며, 모델 자체의 주변 속도 예측으로부터 계산된 2단계 오일러 단축 속도를 사용한 일관성 샘플과 표준 플로우-매칭 샘플을 혼합하여, 조건부 속도로 인한 궤적 드리프트를 방지합니다. 또한, 초기화된 타겟-타임 임베딩을 통해 네트워크는 단일 아키텍처 내에서 로컬 속도 추정과 글로벌 단일 단계 생성을 전환할 수 있습니다. SnapFlow는 외부 교사 모델이 필요 없으며, 아키텍처 변경 없이 단일 GPU에서 약 12시간 내에 학습이 완료됩니다. 본 연구에서는 6배의 파라미터 범위를 가진 두 가지 VLA 아키텍처에서 동일한 하이퍼파라미터를 사용하여 SnapFlow의 성능을 검증했습니다. pi0.5 모델(30억 파라미터)을 사용하여 LIBERO 스위트(40개 작업, 400개 에피소드)에서 SnapFlow는 98.75%의 평균 성공률을 달성하여, 10단계 디노이징을 사용하는 모델의 97.75% 성능과 거의 동일하거나 약간 더 높은 성능을 보였으며, 디노이징 속도는 9.6배 향상되고 전체 지연 시간은 274ms에서 83ms로 감소했습니다. SmolVLA 모델(5억 파라미터)에서는 MSE(평균 제곱 오차)를 8.3% 줄이고 전체 처리 속도를 3.56배 향상시켰습니다. 장기 작업에서 액션-스텝을 변경한 실험 결과, SnapFlow는 실행 범위에 걸쳐 우수한 성능을 유지하며, 액션 스텝이 5개일 때 93%의 성공률을 보이는 반면, 기존 방식은 90%에 불과했습니다. SnapFlow는 레이어-디스틸레이션 및 토큰-프루닝 접근 방식과 독립적으로 작동하여 추가적인 속도 향상을 제공합니다.

Vision-Language-Action (VLA) models based on flow matching -- such as pi0, pi0.5, and SmolVLA -- achieve state-of-the-art generalist robotic manipulation, yet their iterative denoising, typically 10 ODE steps, introduces substantial latency: on a modern GPU, denoising alone accounts for 80% of end-to-end inference time. Naively reducing the step count is unreliable, degrading success on most tasks due to the velocity field being uncalibrated for single-step jumps. We present SnapFlow, a plug-and-play self-distillation method that compresses multi-step denoising into a single forward pass (1-NFE) for flow-matching VLAs. SnapFlow mixes standard flow-matching samples with consistency samples whose targets are two-step Euler shortcut velocities computed from the model's own marginal velocity predictions, avoiding the trajectory drift caused by conditional velocities, as we analyze theoretically. A zero-initialized target-time embedding lets the network switch between local velocity estimation and global one-step generation within a single architecture. SnapFlow requires no external teacher, no architecture changes, and trains in ~12h on a single GPU. We validate on two VLA architectures spanning a 6x parameter range, with identical hyperparameters: on pi0.5 (3B) across four LIBERO suites (40 tasks, 400 episodes), SnapFlow achieves 98.75% average success -- matching the 10-step teacher at 97.75% and slightly exceeding it -- with 9.6x denoising speedup and end-to-end latency reduced from 274ms to 83ms; on SmolVLA (500M), it reduces MSE by 8.3% with 3.56x end-to-end acceleration. An action-step sweep on long-horizon tasks reveals that SnapFlow maintains its advantage across execution horizons, achieving 93% at n_act=5 where the baseline reaches only 90%. SnapFlow is orthogonal to layer-distillation and token-pruning approaches, enabling compositional speedups.