Enregistré dans:
Détails bibliographiques
Auteurs principaux: Qiu, Zhiying, Lin, Tao
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:https://arxiv.org/abs/2506.16688
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866909654087368704
author Qiu, Zhiying
Lin, Tao
author_facet Qiu, Zhiying
Lin, Tao
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2506_16688
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Fast and Stable Diffusion Planning through Variational Adaptive Weighting
Qiu, Zhiying
Lin, Tao
Machine Learning
Artificial Intelligence
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.
title Fast and Stable Diffusion Planning through Variational Adaptive Weighting
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2506.16688