Fast and Stable Diffusion Planning through Variational Adaptive Weighting

• URL: http://arxiv.org/abs/2506.16688v1
• Date: Fri, 20 Jun 2025 02:12:04 GMT
• Title: Fast and Stable Diffusion Planning through Variational Adaptive Weighting
• Authors: Zhiying Qiu, Tao Lin,
• Abstract summary: Diffusion models have recently shown promise in offline RL.<n>These methods often suffer from high training costs and slow convergence.<n>We introduce a closed-form approximation method for its online estimation under the flow-based generative modeling framework.<n> Experimental results on Maze2D and Kitchen tasks show that our method achieves competitive performance with up to 10 times fewer training steps.
• Score: 3.745003761050674
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• AdaLRS: Loss-Guided Adaptive Learning Rate Search for Efficient Foundation Model Pretraining [12.630306478872043]
  We propose textbfAdaLRS, a plug-in-and-play adaptive learning rate search algorithm that conducts online optimal learning rate search.<n>Experiments show that AdaLRS adjusts suboptimal learning rates to the neighborhood of optimum with marked efficiency and effectiveness.
  arXiv  Detail & Related papers  (2025-06-16T09:14:01Z)
• Leveraging Stochastic Depth Training for Adaptive Inference [1.996143466020199]
  We propose a simpler yet effective alternative for adaptive inference with a zero-overhead, single-model, and time-predictable inference.<n>Compared to original ResNets, our method shows improvements of up to 2X in power efficiency at accuracy drops as low as 0.71%.
  arXiv  Detail & Related papers  (2025-05-23T08:36:56Z)
• VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL [28.95582264086289]
  VAlue-based Reinforced Diffusion (VARD) is a novel approach that first learns a value function predicting expection of rewards from intermediate states.<n>Our method maintains proximity to the pretrained model while enabling effective and stable training via backpropagation.
  arXiv  Detail & Related papers  (2025-05-21T17:44:37Z)
• Fast Adaptation with Behavioral Foundation Models [82.34700481726951]
  Unsupervised zero-shot reinforcement learning has emerged as a powerful paradigm for pretraining behavioral foundation models.<n>Despite promising results, zero-shot policies are often suboptimal due to errors induced by the unsupervised training process.<n>We propose fast adaptation strategies that search in the low-dimensional task-embedding space of the pre-trained BFM to rapidly improve the performance of its zero-shot policies.
  arXiv  Detail & Related papers  (2025-04-10T16:14:17Z)
• Prompt Tuning with Diffusion for Few-Shot Pre-trained Policy Generalization [55.14484317645865]
  We develop a conditional diffusion model to produce exceptional quality prompts for offline reinforcement learning tasks. We show that the Prompt diffuser is a robust and effective tool for the prompt-tuning process, demonstrating strong performance in the meta-RL tasks.
  arXiv  Detail & Related papers  (2024-11-02T07:38:02Z)
• AdaRankGrad: Adaptive Gradient-Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning [9.51289606759621]
  Training and fine-tuning large language models (LLMs) come with challenges related to memory and computational requirements.<n>Various techniques have been developed to tackle these challenges, such as low-rank adaptation (LoRA)<n>We introduce a new method inspired by a phenomenon we formally prove: as training progresses, the rank of the estimated gradient gradually decreases.
  arXiv  Detail & Related papers  (2024-10-23T13:53:26Z)
• Efficient Diffusion as Low Light Enhancer [63.789138528062225]
  Reflectance-Aware Trajectory Refinement (RATR) is a simple yet effective module to refine the teacher trajectory using the reflectance component of images. textbfReflectance-aware textbfDiffusion with textbfDistilled textbfTrajectory (textbfReDDiT) is an efficient and flexible distillation framework tailored for Low-Light Image Enhancement (LLIE)
  arXiv  Detail & Related papers  (2024-10-16T08:07:18Z)
• Enhancing Robustness of Vision-Language Models through Orthogonality Learning and Self-Regularization [77.62516752323207]
  We introduce an orthogonal fine-tuning method for efficiently fine-tuning pretrained weights and enabling enhanced robustness and generalization. A self-regularization strategy is further exploited to maintain the stability in terms of zero-shot generalization of VLMs, dubbed OrthSR. For the first time, we revisit the CLIP and CoOp with our method to effectively improve the model on few-shot image classficiation scenario.
  arXiv  Detail & Related papers  (2024-07-11T10:35:53Z)
• Unmasking Bias in Diffusion Model Training [40.90066994983719]
  Denoising diffusion models have emerged as a dominant approach for image generation. They still suffer from slow convergence in training and color shift issues in sampling. In this paper, we identify that these obstacles can be largely attributed to bias and suboptimality inherent in the default training paradigm.
  arXiv  Detail & Related papers  (2023-10-12T16:04:41Z)
• Training Generative Adversarial Networks by Solving Ordinary Differential Equations [54.23691425062034]
  We study the continuous-time dynamics induced by GAN training. From this perspective, we hypothesise that instabilities in training GANs arise from the integration error. We experimentally verify that well-known ODE solvers (such as Runge-Kutta) can stabilise training.
  arXiv  Detail & Related papers  (2020-10-28T15:23:49Z)
• Extrapolation for Large-batch Training in Deep Learning [72.61259487233214]
  We show that a host of variations can be covered in a unified framework that we propose. We prove the convergence of this novel scheme and rigorously evaluate its empirical performance on ResNet, LSTM, and Transformer.
  arXiv  Detail & Related papers  (2020-06-10T08:22:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.