# Fast and Stable Diffusion Planning through Variational Adaptive Weighting

arXiv: https://arxiv.org/abs/2506.16688 · HTML: https://mlsanity.com/paper/2506.16688/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 34 of its 38 references that are in the map.
- First posted on arXiv 20 Jun 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)
- Offline RL and imitation learning from data / [Diffusion and flow policies for offline RL](https://mlsanity.com/offlinerl/diffusion-and-flow-policies-for-offline-rl/)

## Builds on (most cited first, 30 of 34)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Learning Universal Policies via Text-Guided Video Generation](https://mlsanity.com/paper/2302.00111/) (2023)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [IDQL: Implicit Q-Learning as an Actor-Critic Method with Diffusion Policies](https://mlsanity.com/paper/2304.10573/) (2023)
- [Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation](https://mlsanity.com/paper/2303.00848/) (2023)
- [Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models](https://mlsanity.com/paper/2410.11081/) (2024)
- [Efficient Diffusion Training via Min-SNR Weighting Strategy](https://mlsanity.com/paper/2303.09556/) (2023)
- [Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling](https://mlsanity.com/paper/2209.14548/) (2022)
- [AdaptDiffuser: Diffusion Models as Adaptive Self-evolving Planners](https://mlsanity.com/paper/2302.01877/) (2023)
- [Contrastive Energy Prediction for Exact Energy-Guided Diffusion Sampling in Offline Reinforcement Learning](https://mlsanity.com/paper/2304.12824/) (2023)
- [Simple Hierarchical Planning with Diffusion](https://mlsanity.com/paper/2401.02644/) (2024)
- [FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification](https://mlsanity.com/paper/2410.10356/) (2024)

## Abstract (arXiv)

Diffusion models have recently shown promise in offline RL. However, these methods often suffer from high training costs and slow convergence, particularly when using transformer-based denoising backbones. While several optimization strategies have been proposed -- such as modified noise schedules, auxiliary prediction targets, and adaptive loss weighting -- challenges remain in achieving stable and efficient training. In particular, existing loss weighting functions typically rely on neural network approximators, which can be ineffective in early training phases due to limited generalization capacity of MLPs when exposed to sparse feedback in the early training stages. In this work, we derive a variationally optimal uncertainty-aware weighting function and introduce a closed-form polynomial approximation method for its online estimation under the flow-based generative modeling framework. We integrate our method into a diffusion planning pipeline and evaluate it on standard offline RL benchmarks. Experimental results on Maze2D and Kitchen tasks show that our method achieves competitive performance with up to 10 times fewer training steps, highlighting its practical effectiveness.