Test-time Sparsity for Extreme Fast Action Diffusion

Abstract: 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 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.

• The paper demonstrates that test-time sparsity prunes up to 95% of computations while preserving high-fidelity action generation.
• It introduces a parallelized inference pipeline with a 3D lattice cache that reuses features across multiple axes to eliminate redundant computation.
• The approach achieves a 5× speedup with 47.5 Hz inference, validating robust performance on complex robotic manipulation tasks.

Test-Time Sparsity for Extreme Fast Action Diffusion: An Expert Summary

Introduction and Motivation

Diffusion-based action policies, particularly diffusion transformers, achieve high-fidelity, multimodal action generation for complex robotic and vision-language-action (VLA) tasks. Despite their effectiveness, a critical limitation is the severe computational overhead entailed by their iterative denoising, resulting in low control frequencies (frequently $<$10 Hz), which is insufficient for real-time robotic manipulation.

Fundamental bottlenecks preventing speedup are: (1) the latency imposed by redundant conditional encoding and recurrent pruning stages, and (2) the inability of feature caches from adjacent denoising steps to adequately constrain approximation error under aggressive computation pruning. This work introduces a comprehensive sparsification framework—test-time sparsity—for action diffusion, integrating a parallelized inference pipeline with an omnidirectional reusing strategy, supported by 3D lattice cache modeling and trajectory-level training. This results in substantial acceleration, enabling 47.5 Hz inference and a 5$\times$ empirical speedup at up to 95% sparsity, with no loss in generation fidelity.

System Architecture

Parallelized Inference Pipeline

Test-time sparsity requires adaptive, data-conditioned skipping of residual computations in each forward pass. A lightweight pruner—sharing its encoder with the diffusion model—generates a binary pruning mask predicting which transformer blocks can be omitted. To minimize the inefficiency from sequential pruner invocations in iterative denoising, all timesteps are encoded in parallel, collapsing the temporal axis into the batch dimension. This yields a large, flat batch amenable to a single GPU call; the pipeline asynchronously overlaps pruning and decoding, relegating pruner overhead to well below 1 ms.

Figure 1: The paradigm integrates a parallelized pipeline, omnidirectional caching, 3D lattice management of features, and trajectory-level optimization for dynamic sparsity adaptation.

Figure 2: The parallelized pipeline minimizes redundant encoder and pruner calls, in contrast to naively sequential baselines.

Figure 3: Empirical inference latency is dominated by decoder steps after aggressive pruning; the remainder is nearly eliminated by parallelization.

This approach allows the caching and reuse of condition encodings and pruning masks, obviating sequential bottlenecks and making residual block skipping a negligible portion of wall-clock latency.

Omnidirectional Reusing Strategy and 3D Lattice Modeling

Prior approaches cache features only along the denoising-timestep axis, yielding suboptimal performance at high sparsity. This work reveals that cached features from three orthogonal axes—current forward computation, previous denoising timesteps, and prior rollout iterations—exhibit substantial alignment in latent space, as demonstrated via similarity metrics and PCA visualization.

Figure 4: Anchored feature alignment and complementary utility of cached features across denoising steps and rollouts, visualized via PCA.

The cache is organized as a 3D lattice $(b, k, r)$, where $b$ denotes block index, $k$ is denoising step, and $r$ is rollout iteration. Each feature thus obtains three reusable candidates, and the pruner's output is designed as a four-dimensional gating vector, indicating whether to compute or reuse from each axis. Supervision is provided across full action trajectories, using a fidelity objective for action reconstruction and a sparsity penalty to target a desired pruning rate.

Training Protocol and Mask Dynamics

The pruner is optimized jointly with the action diffusion transformer via trajectory-level, stepwise backpropagation, employing the Straight-Through Estimator for discrete mask decisions. During inference, mask dynamics are highly nonstationary, adapting to evolving visuomotor policy structure and changing environmental states.

Figure 5: Example rollout iteration—dynamic mask patterns underscore the test-time adaptation to trajectory context.

Experimental Evaluation

Quantitative evaluation covers a broad suite of manipulation and VLA tasks, including Lift, Can, Square, Transport, Tool Hang, and multi-phase Kitchen scenarios, under diverse human-demonstration datasets and with both transformer-based diffusion policies and large-scale RDT-1B VLA models.

Key empirical results:

• Up to 95% of decoder computations pruned, with 5$\times$ real-world speedup and preserved performative success rates even on challenging multi-stage policies.
• On standard manipulation tasks, the accelerated model achieves 47.5 Hz inference on an NVIDIA 4090, surpassing the $30$ Hz real-world control threshold.
• In ablation, omnidirectional caching decisively outperforms single-axis strategies, which collapse on at least one task under high sparsity.

  Figure 6: Success rates on the Kitchen benchmark as a function of the number of training trajectories; high data efficiency is maintained under extreme sparsity.

Mask and Cache Analysis

Qualitative analysis shows that the mask usage pattern adapts between tasks and over rollout iterations, with significant proportions of cached features sourced from the rollout direction in complex environments—highlighting the necessity of aggregating multiple axes for high-fidelity, aggressive sparsity.

Figure 7: Mask dynamics on the Lift task—cache selection evolves throughout extended rollouts.

Practical and Theoretical Implications

The architecture demonstrates that real-time action diffusion is compatible with extreme sparsity if dynamic, context-aligned reuse of historical representations is enabled. This approach addresses both theoretical and practical limitations of static or single-axis cache schedules, opening pathways to:

• deployment of high-fidelity diffusive action policies in low-latency, real-time robotics,
• training-efficient sparsity adaptation at trajectory-level granularity,
• extension to other generative domains such as video and event-based sequence modeling,
• further exploration of adaptive, lattice-based caching for efficient large-scale models.

Conclusion

This work establishes test-time sparsity with omnidirectional feature reuse and parallelized inference as a scalable paradigm for accelerating diffusion-based action generation. The methodology achieves extreme computation pruning (up to 95%) and 5$\times$ speedup, without sacrificing policy performance in long-horizon, high-frequency robotic scenarios. These findings not only enhance the practical accessibility of action diffusion but suggest new research directions in dynamic model sparsification and cache orchestration for both policy execution and broader generative modeling (2605.13316).