面向极速动作扩散的测试时稀疏性
Test-time Sparsity for Extreme Fast Action Diffusion
把「测试时稀疏化 + 跨去噪步/跨 rollout 的特征复用」做成可插拔的加速范式,对视觉里做扩散 Transformer 推理加速(如视频/图像扩散的 few-step 与缓存复用)可直接迁移;其并行化去噪时间步、剪枝器与解码器异步重叠的流水线设计,也是降低非主干延迟的通用工程思路。
原文依据 We propose test-time sparsity to tackle this challenge, which aims to accelerate action diffusion by dynamically predicting prunable residual computations
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论动作扩散在高保真动作生成方面表现出色,但其迭代去噪的特性带来了沉重的计算开销。尽管现有技术通过复用缓存特征在加速扩散 Transformer 上展现出潜力,但它们难以适应开放环境中由多样感知与多轮 rollout 迭代所引发的策略动态变化。我们提出测试时稀疏性来应对这一挑战,其目标是在测试时对每次模型前向动态预测可剪枝的残差计算,从而加速动作扩散。然而该范式仍存在两个瓶颈:1)重复的条件编码与剪枝抵消了大部分潜在的速度收益;2)在激进稀疏度下,从先前去噪时间步缓存的特性无法约束较大的剪枝误差。针对第一个瓶颈,我们设计了一条高度并行的推理流水线,将非解码器延迟降至毫秒级。具体而言,我们首先设计了一个与扩散 Transformer 共享编码器的轻量剪枝器;随后通过并行处理所有去噪时间步,将编码与剪枝从自回归去噪循环中解耦,并借助异步机制使剪枝器与解码器前向推理相互重叠。为克服第二个瓶颈,我们引入全向复用策略,通过有选择地复用当前前向、先前去噪时间步以及更早 rollout 迭代所缓存的特性,实现 95% 的稀疏性。为学习 rollout 级别的复用策略,我们采样少量动作轨迹,逐步对稀疏化扩散过程进行监督。大量实验表明,我们的方法将浮点运算量降低 92%,动作生成加速 5 倍,在 47.5 Hz 的推理频率下实现无损性能。代码已开源至 https://github.com/ky-ji/Test-time-Sparsity。
英文原文摘要 1924 字符
Action diffusion excels at high-fidelity action generation but incurs heavy computational costs owing to its iterative denoising nature. Despite current technologies showing promise in accelerating diffusion transformers by reusing the cached features, they struggle to adapt to policy dynamics arising from diverse perceptions and multi-round rollout iterations in open environments. We propose test-time sparsity to tackle this challenge, which aims to accelerate action diffusion by dynamically predicting prunable residual computations for each model forward at test time. However, two bottlenecks remain in this paradigm: 1) repetitive conditional encoding and pruning offset most potential speed gains, and 2) the features cached from previous denoising timesteps cannot constrain large pruning errors under aggressive sparsity. To address the first bottleneck, we design a highly parallelized inference pipeline that minimizes the non-decoder delay to milliseconds. Specifically, we first design a lightweight pruner that shares the encoder with the diffusion transformer. Then, we decouple the encoding and pruning from the autoregressive denoising loop by processing all denoising timesteps in parallel, and overlap the pruner with the decoder forward inference through asynchronism. To overcome the second bottleneck, we introduce an omnidirectional reusing strategy, which achieves 95% sparsity by selectively reusing the features cached from the current forward, previous denoising timesteps, and earlier rollout iterations. To learn the rollout-level reusing strategies, we sample a few action trajectories to supervise the sparsified diffusion step by step. Extensive experiments demonstrate that our method reduces FLOPs by 92% and accelerates action generation by 5x, achieving lossless performance with an inference frequency of 47.5 Hz. Our code is available at https://github.com/ky-ji/Test-time-Sparsity.