Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Uehara, Masatoshi, Zhao, Yulai, Black, Kevin, Hajiramezanali, Ehsan, Scalia, Gabriele, Diamant, Nathaniel Lee, Tseng, Alex M, Biancalani, Tommaso, Levine, Sergey
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866914693356978176
author Uehara, Masatoshi
Zhao, Yulai
Black, Kevin
Hajiramezanali, Ehsan
Scalia, Gabriele
Diamant, Nathaniel Lee
Tseng, Alex M
Biancalani, Tommaso
Levine, Sergey
author_facet Uehara, Masatoshi
Zhao, Yulai
Black, Kevin
Hajiramezanali, Ehsan
Scalia, Gabriele
Diamant, Nathaniel Lee
Tseng, Alex M
Biancalani, Tommaso
Levine, Sergey
contents Diffusion models excel at capturing complex data distributions, such as those of natural images and proteins. While diffusion models are trained to represent the distribution in the training dataset, we often are more concerned with other properties, such as the aesthetic quality of the generated images or the functional properties of generated proteins. Diffusion models can be finetuned in a goal-directed way by maximizing the value of some reward function (e.g., the aesthetic quality of an image). However, these approaches may lead to reduced sample diversity, significant deviations from the training data distribution, and even poor sample quality due to the exploitation of an imperfect reward function. The last issue often occurs when the reward function is a learned model meant to approximate a ground-truth "genuine" reward, as is the case in many practical applications. These challenges, collectively termed "reward collapse," pose a substantial obstacle. To address this reward collapse, we frame the finetuning problem as entropy-regularized control against the pretrained diffusion model, i.e., directly optimizing entropy-enhanced rewards with neural SDEs. We present theoretical and empirical evidence that demonstrates our framework is capable of efficiently generating diverse samples with high genuine rewards, mitigating the overoptimization of imperfect reward models.
format Preprint
id arxiv_https___arxiv_org_abs_2402_15194
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control
Uehara, Masatoshi
Zhao, Yulai
Black, Kevin
Hajiramezanali, Ehsan
Scalia, Gabriele
Diamant, Nathaniel Lee
Tseng, Alex M
Biancalani, Tommaso
Levine, Sergey
Machine Learning
Artificial Intelligence
Diffusion models excel at capturing complex data distributions, such as those of natural images and proteins. While diffusion models are trained to represent the distribution in the training dataset, we often are more concerned with other properties, such as the aesthetic quality of the generated images or the functional properties of generated proteins. Diffusion models can be finetuned in a goal-directed way by maximizing the value of some reward function (e.g., the aesthetic quality of an image). However, these approaches may lead to reduced sample diversity, significant deviations from the training data distribution, and even poor sample quality due to the exploitation of an imperfect reward function. The last issue often occurs when the reward function is a learned model meant to approximate a ground-truth "genuine" reward, as is the case in many practical applications. These challenges, collectively termed "reward collapse," pose a substantial obstacle. To address this reward collapse, we frame the finetuning problem as entropy-regularized control against the pretrained diffusion model, i.e., directly optimizing entropy-enhanced rewards with neural SDEs. We present theoretical and empirical evidence that demonstrates our framework is capable of efficiently generating diverse samples with high genuine rewards, mitigating the overoptimization of imperfect reward models.
title Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2402.15194