Mini Diffuser: Fast Multi-task Diffusion Policy Training Using Two-level Mini-batches
This work addresses efficiency bottlenecks for researchers and practitioners in robotics and AI, enabling faster and more resource-friendly training of diffusion policies, though it is incremental as it builds on existing diffusion methods.
The paper tackles the problem of high time and memory costs in training multi-task vision-language robotic diffusion policies by introducing Mini Diffuser, which reduces training time by 95% and memory usage by 93% while achieving 95% of state-of-the-art performance.
We present a method that reduces, by an order of magnitude, the time and memory needed to train multi-task vision-language robotic diffusion policies. This improvement arises from a previously underexplored distinction between action diffusion and the image diffusion techniques that inspired it: In image generation, the target is high-dimensional. By contrast, in action generation, the dimensionality of the target is comparatively small, and only the image condition is high-dimensional. Our approach, \emph{Mini Diffuser}, exploits this asymmetry by introducing \emph{two-level minibatching}, which pairs multiple noised action samples with each vision-language condition, instead of the conventional one-to-one sampling strategy. To support this batching scheme, we introduce architectural adaptations to the diffusion transformer that prevent information leakage across samples while maintaining full conditioning access. In RLBench simulations, Mini-Diffuser achieves 95\% of the performance of state-of-the-art multi-task diffusion policies, while using only 5\% of the training time and 7\% of the memory. Real-world experiments further validate that Mini-Diffuser preserves the key strengths of diffusion-based policies, including the ability to model multimodal action distributions and produce behavior conditioned on diverse perceptual inputs. Code available at mini-diffuse-actor.github.io