DLPO: Diffusion Model Loss-Guided Reinforcement Learning for Fine-Tuning Text-to-Speech Diffusion Models

## Abstract

Abstract Recent advancements in generative models have sparked significant interest within the machine learning community. Particularly, diffusion models have demonstrated remarkable capabilities in synthesizing images and speech. Studies such as those by Lee et al. [ 19 ] , Black et al. [ 4 ] , Wang et al. [ 37 ] , and Fan et al. [ 8 ] illustrate that Reinforcement Learning with Human Feedback (RLHF) can enhance diffusion models for image synthesis. However, due to architectural differences between these models and those employed in speech synthesis, it remains uncertain whether RLHF could similarly benefit speech synthesis models. In this paper, we explore the practical application of RLHF to diffusion-based text-to-speech synthesis, leveraging the mean opinion score (MOS) as predicted by UTokyo-SaruLab MOS prediction system [ 29 ] as a proxy loss. We introduce diffusion model loss-guided RL policy optimization (DLPO) and compare it against other RLHF approaches, employing the NISQA speech quality and naturalness assessment model [ 21 ] and human preference experiments for further evaluation. Our results show that RLHF can enhance diffusion-based text-to-speech synthesis models, and, moreover, DLPO can better improve diffusion models in generating natural and high quality speech audios.

## 1 Introduction

Diffusion probabilistic models, initially introduced by Sohl-Dickstein et al. [], have rapidly become the predominant method for generative modeling in continuous domains. Valued for their capability to model complex, high-dimensional distributions, these models are widely used in fields like image and video synthesis. Notable applications include Stable Diffusion [] and DALL-E 3 [], known for generating high-quality images from text descriptions. Despite these advancements, significant challenges persist where large-scale text-to-image models struggle to produce images that accurately correspond to text prompts. As highlighted by Lee et al. [] and Fan et al. [], current text-to-image models face challenges in composing multiple objects [, , ] and in generating objects with specific colors and quantities [, ].
Black et al. [], Fan et al. [], and Lee et al. [] propose reinforcement learning (RL) approaches for fine-tuning diffusion models which can enhance their ability to align generated images with input texts.

In contrast, text-to-speech (TTS) models face different limitations. TTS models must handle a sequence of sounds over time, which introduces complexity in terms of processing time-domain data. Additionally, the size of audio files varies based on text length, quality, and compression—for instance, an audio file encoding a few seconds of speech could range from a few hundred kilobytes to several megabytes. These models must also capture subtle aspects of speech such as intonation, pace, emotion, and consistency to produce natural and intelligible audio outputs [, , ]. These challenges are distinct from those faced by text-to-image models. It thus remains uncertain whether similar techniques can deliver improvements comparable to those seen in text-to-image models. We are keen to explore the potential of using RL approaches to fine-tune TTS diffusion models.

Figure: Figure 1: The steps of online RL for fine-tuning text-to-speech diffusion models: (1) Audios are sampled from the pretrained TTS model, (2) The reward model predicts a reward for the each sampled audio from step, (3) The text-to-image model is updated to maximize the reward.
Refer to caption: extracted/2405.14632v1/image/rl.jpg

In this study, we explore the application of online RL to fine-tune TTS diffusion models (Figure 1). Our evaluation focuses on the performance of TTS model fine-tuning using reward-weighted regression (RWR) [], denoising diffusion policy optimization (DDPO) [], diffusion policy optimization with KL regularization (DPOK)[] and diffusion policy optimization with a KL-shaped reward [], which we refer to as KLinR in this paper. We also introduce and assess diffusion loss-guided policy optimization (DLPO) where the reward is shaped by the diffusion model’s loss. In our experiments, we fine-tune WaveGrad 2, a non-autoregressive generative model for TTS synthesis [], using reward predictions from the UTokyo-SaruLab mean opinion score (MOS) prediction system []. We find that RWR, DDPO, DPOK, and KLinR do not improve TTS models as they do for text-to-image models; because they cannot successfully control the magnitude of deviations from the original model. However, DLPO does, improving the sound quality and naturalness of the generated speech.

We can summarize our main contributions as follows:

- $\bullet$
We are the first to apply reinforcement learning (RL) to fine-tune TTS diffusion models.
- $\bullet$
We evaluate current RL methods for fine-tuning diffusion models in the TTS setting: RWR, DDPO, DPOK, and KLinR.
- $\bullet$
We introduce diffusion loss-guided policy optimization (DLPO). Unlike other RL methods, DLPO aligning with the training procedure of TTS diffusion models by incorporating the original diffusion model loss as a penalty in the reward function to effectively prevent model deviation and fine-tune TTS models.
- $\bullet$
We further investigate the impact of inference steps on TTS diffusion model fine-tuning by comparing a Monte Carlo approximation approach to fixed inference step sampling. We conduct a human experiment to evaluate the speech quality of DLPO-generated audio.

## 2 Related Work

#### Text-to-speech diffusion model.

Diffusion models [, , ] have shown impressive capabilities in generative tasks like image [, , ] and audio production [, , ]. TTS systems using diffusion models generate high-fidelity speech comparable to state-of-the-art systems [, , , ]. These models use a Markov chain to transform noise into structured speech waveforms through a series of probabilistic steps, which can be optimized with RL techniques.

However, unlike text-to-image diffusion models, TTS diffusion models face the challenge of high temporal resolution. Audio input for TTS is a one-dimensional signal with a high sample rate; for example, one second of audio at 24,600 Hz consists of 24,600 samples. This high dimensionality requires managing thousands of data points, necessary to capture the nuances of human speech for natural-sounding audio. In contrast, a typical 256x256 image has 65,536 pixels, fewer than the samples in one second of high-fidelity audio. To handle and process the vast number of audio samples without compromising speed or accuracy, diffusion/denoising steps of TTS models are usually larger than that for text-to-image diffusion models, requiring memory optimization strategies during training. To address this, some TTS models only include part of the input sequence in training, which introduces a small mismatch between training and inference.

#### Reinforcement learning from human feedback.

Human assessments have been crucial in guiding model learning across various tasks, from behavior [] and language modeling [, ] to text-to-image alignment [, , ]. This process typically involves collecting feedback at scale by having annotators compare outputs, and subsequently training a reward model to generate rewards aligned with human feedback. These reward models enhance the quality of text-to-image models through fine-tuning pre-trained generative models. Beyond these applications, human feedback has also optimized models in diverse settings including robotic control [], game-playing [], and more.

#### RL fine-tuning of diffusion models.

Recent studies have focused on fine-tuning diffusion text-to-image models using alternative reward-weighted regression methods and reinforcement learning (RL) to enhance their performance. Lee et al. [] demonstrate that developing a reward function based on human feedback and using supervised learning techniques can improve specific attributes like color, count, and background alignment in text-to-image models. While simple supervised fine-tuning (SFT) based on reward-weighted regression loss improves reward scores and image-text alignment, it often results in decreased image quality (e.g., over-saturation or non-photorealistic images).

Fan et al. [] suggest this issue likely arises from fine-tuning on a fixed dataset generated by a pre-trained model. Black et al. [] argue that the reward-weighted regression lacks theoretical grounding and only roughly approximates optimizing denoising diffusion with RL and they propose denoising diffusion policy optimization (DDPO), a policy gradient algorithm that outperforms reward-weighted regression methods. DDPO improves text-to-image diffusion models by targeting objectives like image compressibility and aesthetic appeal, and it enhances prompt-image alignment using feedback from a vision-language model, eliminating the need for additional data or human annotations. Fan et al. [] also show that RL fine-tuning can surpass reward-weighted regression in optimizing rewards for text-to-image diffusion models. Additionally, they demonstrate that using KL regularization in RL methods helps address issues like image quality deterioration in fine-tuned models.

For fine-tuning TTS diffusion model, we develop an diffusion policy optimization shaped by diffusion model objective. This method outperforms RWR, DDPO, and DPOK and KLinR. We find that the latter methods, although reported to enhance text-to-image diffusion models, do not yield improvements in TTS model quality.

## 3 Models

### 3.1 Text-to-Speech Diffusion Model

Figure: Figure 2: WaveGrad2 follows a Markov process where forward diffusion $q(x_{t+1}|x_{t})$ iteratively adds Gaussian noise to the signal starting from the waveform $y_{0}$. $q(x_{t+1}|x_{0})$ is the noise distribution used for training. The inference denoising process progressively removes noise, starting from Gaussian noise $x_{T}$. This figure is adapted from Chen et al. [], Ho et al. [].
Refer to caption: extracted/2405.14632v1/image/wg2.jpg

Diffusion models gradually transform random noise into structured outputs, like speech, through a series of learned transformations. This involves two phases: the forward phase, where noise is incrementally added until only noise remains, and the reverse phase, where a model learns to reverse this noise addition to generate the desired output from pure noise.In this study, we use WaveGrad2 [] as the pretrained TTS diffusion model for its ability to balance inference speed and sample quality by adjusting refinement steps and produce high-fidelity audio. Due to the lack of official code, we use an unofficial version from MINDs Lab ([https://github.com/maum-ai/wavegrad2](https://github.com/maum-ai/wavegrad2)). WaveGrad2 is a non-autoregressive TTS model adapting the diffusion denoising probabilistic model (DDPM) from Ho et al. []. As shown in Figure 2, WaveGrad2 models the conditional distribution $p_{\theta}(x_{0}|c)$ where $x_{0}$ represents the waveform and $c$ the associated context. The distribution follows the reverse of a Markovian forward process $q(x_{t}|x_{t-1})$, which iteratively introduces noise to the data. Reversing the forward process can be accomplished by training a neural network $\mu_{\theta}(x_{t},c,t)$ with the following objective:

$$ $\mathcal{L}_{DDPM}(\theta)=\mathbb{E}_{c\sim p(c)}\mathbb{E}_{t\sim\mathcal{U} \{0,T\}}\mathbb{E}_{p_{\theta}(x_{0:T}|c)}\left[\|\tilde{\mu}(x_{t},t)-\mu_{ \theta}(x_{t},c,t)\|_{2}\right]$ (1) $$

where $\tilde{\mu}$ is the posterior mean of the forward process and $x_{t}$ is the prediction at timestep $t$ in the denoising process. This objective is justified as maximizing a variational lower bound on the log-likelihood of the data [], which is trained to predict the scaled derivative by minimizing the distance between ground truth added noise $\epsilon$ and the model prediction:

$$ $\mathbb{E}_{c\sim p(c)}\mathbb{E}_{t\sim\mathcal{U}\{0,T\}}\mathbb{E}_{p_{ \theta}(x_{0:T}|c)}\left[\|\tilde{\epsilon}(x_{t},t)-\epsilon_{\theta}(x_{t},c ,t)\|_{2}\right]$ (2) $$

where $\tilde{\epsilon}(x_{t},t)$ is the ground truth added noise for step $x_{t}$ and $\epsilon_{\theta}(x_{t},c,t)$ is the predicted noise for step $x_{t}$. As discussed in section 2, given the high waveform resolution in WaveGrad2 (24,000 samples per second), computing the loss across all waveform samples in an utterance is impractical due to significant computational and memory demands. Instead, after learning the representations on the whole input sequence, a small segment is sampled to synthesize the waveform. Random segments are sampled individually in each minibatch and the corresponding waveform segment is extracted based on the upsampling rate (300 in WaveGrad2 setup). WaveGrad2 employs a strategy to randomly select a specific time step $t$ and its associated noise, then calculates the loss by predicting the noise at this step. During inference, the entire input sequence is utilized, which introduces a small mismatch between training and inference. Additionally, to generate and process the vast number of audio sampled at audio at 24 kHz without compromising speed or accuracy, the number of diffusion/denoising steps is set to 1000 in WaveGrad2. This contrasts with text-to-image diffusion models from previous studies, which utilize the entire input sequence during both training and inference phases and involve 50 diffusion/denoising steps throughout their inference processes.

Sampling from a diffusion model begins with drawing a random $x_{T}\sim\mathcal{N}(0,I)$ and following the reverse process $p_{\theta}(x_{t-1}|x_{t},c)$ to produce a trajectory {$x_{T},x_{T-1},...,x_{0}$} ending with a sample $x_{0}$.
As discussed in Chen et al. [, ], the linear variance schedule used in Ho et al. [] is adapted in the sampling process of WaveGrad2, where the variance of the noise added at each step increases linearly over a fixed number of timesteps. The linear variance schedule is a predefined function that dictates the variance of the noise added at each step of the forward diffusion process. This schedule is crucial because it determines how much noise is added to the data at each timestep, which in turn affects the quality of the generated samples during the reverse denoising process.

Additionally, Chen et al. [, ] found that replacing the original $L^{2}$ distance metric with $L^{1}$ improves training stability. However, the unofficial WaveGrad2 implementation by MINDs Lab uses the original $L^{2}$ metric.

### 3.2 Reward Model

We use the UTokyo-SaruLab mean opinion score (UTMOS) prediction system [] to predict the speech quality of generated audios from Wavegrad2. Mean opinion score (MOS) is a subjective scoring system that allows human evaluators to rate the perceived quality of synthesized speech on a scale from 1 to 5, which is one of the most commonly employed evaluation methods for TTS system []. UTMOS is trained to capture nuanced audio features that are indicative of human judgments of speech quality, providing an accurate MOS prediction without the need for extensive labeled data. Thus, we use UTMOS as the reward model for fine-tuning Wavegrad2.

## 4 RL for Fine-tuning TTS Diffusion Models

In this section, we present a Markov decision process (MDP) formulation for WaveGrad2’s denoising phase, evaluate four RL algorithms for training diffusion models, and introduce a modified fine-tuning method incorporating diffusion model loss, comparing it with other RL approaches.

### 4.1 Denoising as a multi-step MDP

We model denoising as a $T$-step finite horizon Markov decision process (MDP). Defined by the tuple $(S,A,\rho_{0},P,R)$, an MDP consists of a state space $S$, an action space $A$, an initial state distribution $\rho_{0}$, a transition kernel $P$, and a reward function $R$. At each timestep $t$, an agent observes a state $s_{t}$ from $S$, selects an action $a_{t}$ from $A$, receives a reward $R(s_{t},a_{t})$, and transitions to a new state $s_{t+1}\sim P(s_{t+1}|s_{t},a_{t})$. The agent follows a policy $\pi_{\theta}(a|s)$, parameterized by $\theta$, to make decisions. As the agent operates within the MDP, it generates trajectories, sequences of states and actions $(s_{0},a_{0},s_{1},a_{1},...,s_{T},a_{T})$. The goal of reinforcement learning (RL) is to maximize $\mathcal{J}_{RL}(\theta)$, which is the expected total reward across trajectories produced under its policy:

$$ $\mathcal{J}_{\mathrm{RL}}(\theta)=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{T }R(\mathrm{s_{t}},\mathrm{a_{t}})\right]$ (3) $$

Therefore, we can define a Markov decision process (MDP) formulation for the denoising phase of WaveGrad2 as follows:

$$ $\begin{split}s_{t}\triangleq(c,x_{T-t})\quad P(s_{t+1}|s_{t},a_{t})\triangleq( \delta_{c},\delta_{a_{t}})\quad a_{t}\triangleq x_{T-t-1}\\ \quad\rho(s_{0})\triangleq(p(c),\mathcal{N}(0,I))\quad R(s_{t},a_{t}) \triangleq\left\{\begin{array}[]{lr}r(s_{t+1})=r(x_{0},c)&\textrm{if}\quad t=T -1\\ 0&\textrm{otherwise,}\\ \end{array}\right.\end{split}$ (4) $$

where $\delta$ is the Dirac delta distribution, $c$ is the text prompt sampled from $p(c)$, and $r(x_{0},c)$ is the reward model UTMOS introduced in subsection 3.2. $s_{t}$ and $a_{t}$ are the state and action at timestep $t$, $\rho(s_{0})$ and $P$ are the initial state distribution and the dynamics, and $R$ is the reward function. We let $\pi_{\theta}(a_{t}|s_{t})\triangleq p_{\theta}(x_{T-t-1}|x_{T-t},c)$ be the initial parameterized policy, where $p_{\theta}(x_{T-t-1}|x_{T-t},c)$ is the WaveGrad2 model discussed in subsection 3.1. Trajectories consist of $T$ timesteps, after which $P$ leads to a terminating state. The cumulative reward of each trajectory is equal to $r(x_{0},c)$. Maximizing $r(x_{0},c)$ to optimize policy $\pi_{\theta}$ in Equation 4 is equivalent to fine-tuning WaveGrad2 ($\mathcal{L}_{DDPM}$ Equation 1), the denoising diffusion RL objective is presented as follows:

$$ $\mathcal{J}_{\mathrm{DDRL}}(\theta)=\mathbb{E}_{c\sim p(c)}\mathbb{E}_{x_{0} \sim p_{\theta}(x_{0}|c)}\left[r(\mathrm{{x_{0}},\mathrm{c}})\right]$ (5) $$

### 4.2 Reward-Weighted Regression

As discussed by Black et al. [], using the denoising loss $\mathcal{L}_{DDPM}$ Equation 1 with training data $x_{0}\sim p_{\theta}(x_{0}|c)$ and an added weighting of reward $r(x_{0},c)$, we can optimize $\mathcal{J}_{DDRL}$ with minimal changes to standard diffusion model training. This approach can be referred as reward-weighted regression (RWR) []. Lee et al. [] use this approach to update the diffusion models, their objective is presented as follow:

$$ $\mathcal{J}_{\mathrm{DDRL}}(\theta)=\mathbb{E}_{c\sim p(c)}\mathbb{E}_{p_{ \mathrm{pre}}(x_{0}|c)}\left[-r(x_{0},c)\log p_{\theta}(\mathrm{x_{0}}|\mathrm {c})\right]$ (6) $$

However, Lee et al. [] note that fine-tuning the text-to-image diffusion model with reward-weighted regression can lead to reduced image quality, such as over-saturation or non-photorealistic images. Fan et al. [] suggest that this deterioration might be due to the model being fine-tuned on a static dataset produced by a pre-trained model, meanwhile, Black et al. [] argue that reward-weighted regression aims to approximately maximize $\mathcal{J}_{RL}(\pi)$ subject to a KL divergence constraint on $\pi$ []. Yet, the denoising loss $\mathcal{L}_{DDPM}$ Equation 1 does not compute an exact log-likelihood; it is instead a variational bound on log $p_{\theta}(x_{0}|c)$. As such, the RWR procedure approach to training diffusion models lacks theoretical justification and only approximates optimization of $\mathcal{J}_{DDRL}$ (Equation 5).

### 4.3 Denoising Diffusion Policy Optimization

RWR relies on an approximate log-likelihood by disregarding the sequential aspect of the denoising process and only utilizing the final samples $x_{0}$. Black et al. [] propose the denoising diffusion policy optimization to directly optimize $\mathcal{J}_{DDRL}$ using the score function policy gradient estimator, also known as REINFORCE [, ]. DDPO alternately collects denoising trajectories ${x_{T},x_{T-1},...,x_{0}}$ via sampling and updates parameters via gradient descent:

$$ $\nabla_{\theta}\mathcal{J}_{\mathrm{DDRL}}(\theta)=\mathbb{E}_{c\sim p(c)} \mathbb{E}_{p_{\theta}(x_{0:T}|c)}\left[\sum_{t=1}^{T}\nabla_{\theta}\log p_{ \theta}(x_{t-1}|x_{t},c)r(x_{0},c)\right]$ (7) $$

where the expectation is calculated across denoising trajectories generated by the current parameters $\theta$. This estimator only allows for one step of optimization for each data collection round, since the gradient needs to be calculated using data derived from the current parameters. In Black et al. []’s study, DDPO is shown to achieve better performance in fine-tuning the text-to-image diffusion model compared to reward-weighted regression methods.

In the training process of WaveGrad2, as described in subsection 3.1, it is computationally and memory-intensive to compute the loss for every waveform sample in a speech segment. Therefore, WaveGrad2 employs a strategy where it randomly selects a specific time step $t$ along with its associated noise. The loss is then calculated by predicting the noise at this selected time step. We can adapt DDPO to the WaveGrad2 setting as follows by leveraging a Monte Carlo approximation of the $\mathcal{J}_{\mathrm{DDRL}}$ objective:

$$ $\nabla_{\theta}\mathcal{J}_{\mathrm{DDRL}}(\theta)=\mathbb{E}_{c\sim p(c)} \mathbb{E}_{p_{\theta}(x_{0:T}|c)}\mathbb{E}_{t\sim\mathcal{U}\{1,T\}}\left[ \nabla_{\theta}\log p_{\theta}(x_{t-1}|x_{t},c)r(x_{0},c)\right]$ (8) $$

### 4.4 Diffusion Policy Optimization with a KL-shaped Reward

Fan et al. [] demonstrate that adding KL between the fine-tuned and pre-trained models for the final image as a regularizer $KL(p_{\theta}(x_{0}|z)\|p_{\mathrm{pre}}(x_{0}|z)$ to the objective function helps to mitigate overfitting of the diffusion models to the reward and prevents excessively diminishing the "skill" of the original diffusion model. As discussed in subsection 3.1, $p_{\theta}(x_{0}|z)$ is calculated as a variational bound, Fan et al. [] propose to add an upper-bound of this KL-term to the objective function $\mathcal{J}_{DDRL}(\theta)$, which they call DPOK. They find that DPOK outperforms reward weighted regression in fine-tuning text-to-image diffusion models:

$$ $\mathbb{E}_{c\sim p(c)}\left[\alpha\mathbb{E}_{p_{\theta}(x_{t-1}|x_{t},c)} \left[-r(x_{0},c)\right]+\beta\sum_{t=1}^{T}\mathbb{E}_{p_{\theta}(x_{t}|c)} \left[\mathrm{KL}(p_{\theta}(x_{t-1}|x_{t},c)\|p_{\mathrm{pre}}(x_{t-1}|x_{t}, c))\right]\right]$ (9) $$

where $\alpha,\beta$ are the reward and KL weights, respectively. They use the following gradient to optimize the objective:

$$ $\operatorname*{\mathbb{E}}_{\begin{subarray}{c}c\sim p(c)\\ p_{\theta}(x_{0:T}|c)\end{subarray}}\Bigg{[}-\alpha r(x_{0},c)\sum_{t=1}^{T} \nabla_{\theta}\log p_{\theta}(x_{t-1}|x_{t},c)+\\ \beta\sum_{t=1}^{T}\nabla_{\theta}\mathrm{KL}(p_{\theta}(x_{t-1}|x_{t},c)\|p_{ \mathrm{pre}}(x_{t-1}|x_{t},c))\Bigg{]}$ $$

As DDPO, DPOK cannot be applied directly to WaveGrad2 for memory reasons, but we can adapt it using a Monte Carlo approximation:

$$ $\mathbb{E}_{c\sim p(c)}\mathbb{E}_{p_{\theta}(x_{0:T}|c)}\mathbb{E}_{t\sim \mathcal{U}\{1,T\}}\left[-\alpha r(x_{0},c)\nabla_{\theta}\log p_{\theta}(x_{t -1}|x_{t},c)+\beta\nabla_{\theta}\log\frac{p_{\theta}(x_{t-1}|x_{t},c)}{p_{ \mathrm{pre}}(x_{t-1}|x_{t},c)}\right]$ (10) $$

Another RL objective presented by Ahmadian et al. [] also involves a KL-penalty to prevent degradation in the coherence of the model (KLinR). In contrast to DPOK, this objective function $\mathcal{J}_{\mathrm{DDRL}}(\theta)$ includes the KL penalty within the reward function:

$$ $\mathcal{J}_{DDRL}(\theta)=\mathbb{E}_{c\sim p(c)}\mathbb{E}_{p_{\theta}(x_{0: T}|c)}\left[-\left(\alpha r(x_{0},c)-\beta\log\frac{p_{\theta}(x_{0}|c)}{p_{ \mathrm{pre}}(x_{0}|c)}\right)\sum_{t=1}^{T}\log p_{\theta}(x_{t-1}|x_{t},c)\right]$ (11) $$

where $\beta$ is the KL weight. We can adapt it to WaveGrad2 similarly:

$$ $\nabla_{\theta}\mathcal{J}_{DDRL}(\theta)=\operatorname*{\mathbb{E}}_{\begin{ subarray}{c}c\sim p(c)\\ p_{\theta}(x_{0:T}|c)\\ t\sim\mathcal{U}\{1,T\}\end{subarray}}\Bigg{[}-\Bigg{(}\alpha r(x_{0},c)-\\ \beta\nabla_{\theta}\log\frac{p_{\theta}(x_{t-1}|x_{t},c)}{p_{\mathrm{pre}}(x_ {t-1}|x_{t},c)}\Bigg{)}\nabla_{\theta}\log p_{\theta}(x_{t-1}|x_{t},c)\Bigg{]}$ $$

We evaluate both approaches with KL for fine-tuning WaveGrad2; for KLinR approach, we follow Fan et al. [] and calculate the KL upper-bound. The algorithms are shown in Appendix A and Appendix B.

### 4.5 Diffusion Model Loss-Guided Policy Optimization

Ouyang et al. [] indicate that mixing the pretraining gradients into the RL gradient shows improved performance on certain public NLP datasets compared to a reward-only approach. Therefore, adding the diffusion model loss to the objective function can be another way to improve performance and prevent degradation in the coherence of the model. We propose the following objective:

$$ $\mathbb{E}_{c\sim p(c)}\mathbb{E}_{t\sim\mathcal{U}\{1,T\}}\mathbb{E}_{p_{ \theta}(x_{0:T}|c)}\left[-\alpha r(x_{0},c)-\beta\|\tilde{\epsilon}(x_{t},t)- \epsilon_{\theta}(x_{t},c,t)\|_{2}\right]$ (12) $$

where $\alpha,\beta$ are the reward and weights for diffusion model loss, respectively. We use the following gradient to update the objective:

$$ $\mathbb{E}_{c\sim p(c)}\mathbb{E}_{t\sim\mathcal{U}\{0,T\}}\mathbb{E}_{p_{ \theta}(x_{1:T}|c)}\left[-\left(\alpha r(x_{0},c)-\beta\nabla_{\theta}\|\tilde {\epsilon}(x_{t},t)-\epsilon_{\theta}(x_{t},c,t)\|_{2}\right)\nabla_{\theta} \log p_{\theta}(x_{t-1}|x_{t},c)\right]$ (13) $$

where we follow Ahmadian et al. [] and add the diffusion model objective to the reward function as a penalty. The pseudo-code of our algorithm, which we refer to as DLPO, is summarized in Algorithm Algorithm 1. This algorithm aligning with the training procedure of TTS diffusion models by incorporating the original diffusion model objective $\beta\|\tilde{\epsilon}(x_{t},t)-\epsilon_{\theta}(x_{t},c,t)\|_{2}$ as a penalty in the reward function effectively prevent model deviation.

Figure: Algorithm 1 DLPO: Diffusion model loss-guided policy optimization

## 5 Experiments

We now present a series of experiments designed to evaluate the efficacy of different RL fine-tuning methods on TTS diffusion model.

### 5.1 Experimental Design

Figure: Figure 3: (fine-tuning effectiveness) The relative effectiveness of seven RL algorithms. (Left) shows the change of UTMOS during training with 550 reward samples (reward samples = episodes * batch size) of each RL approach while (right) shows the change of UTMOS during the initial 200 training episodes of the diffusion model loss-guided policy optimization (DLPO), which illustrates an increase of UTMOS from 3.41 to 3.59. $\textrm{DPOK}_{\textrm{tail}}$ and $\textrm{KLinR}_{\textrm{tail}}$ are only trained with samples of step $t=996,997,998,999$.
Refer to caption: extracted/2405.14632v1/image/klpenalty_allanddlpo.png

We use the unofficial implementation of WaveGrad2 as our baseline model, pre-trained on the LJSpeech dataset [] which consists of 13,100 short audio clips of a female speaker and the corresponding texts, totaling approximately 24 hours. For the reward model, we use UTMOS [], trained on datasets from the VoiceMOS Challenge 2022 [], including 4,974 audio clips from male and female speakers with MOS ratings. For automatic evaluation of the fine-tuned models, we use another speech quality and naturalness assessment model (NISQA) [], trained on the NISQA Corpus including more than 14,000 speech samples produced by male and female speakers along with samples’ MOS ratings. (This use of a separate MOS network is intended to guard against overfitting the reward model.) We also conduct a human experiment to have people evaluate the speech quality of the generated audios from fine-tuned WaveGrad2.

We follow Fan et al. [] and use online RL training for fine-tuning WaveGrad2 to evaluate the performance of RWR, DDPO, DLPO, KLinR, $\textrm{KLinR}_{\textrm{tail}}$, DPOK, and $\textrm{DPOK}_{\textrm{tail}}$. In each episode (amount of batch processed) during the training, we sample a new trajectory based on the current model distribution $\pi_{\theta}$ and calculate the rewards of the new trajectory. Online RL training is claimed to be better at maximizing the reward than the supervised approach which only uses the supervised dataset based on the pre-trained distribution. Due to significant computational and memory demands of online RL training for fine-tuning WaveGrad2, we train WaveGrad2 with 8 A100-SXM-80GB GPU for 5.5 hours and 550 samples are fed to WaveGrad2.

During online RL fine-tuning, the model is updated using new samples from the previously trained model. In every training episode, the UTMOS score is computed using the final state $x_{0}$ of the trajectory generated given sampled text $c$ from the training dataset, where $c\sim p(c)$. We plot the UTMOS score obtained in each training episode in Figure 3 to illustrate the fine-tuning progress of Wavegrad2. In Figure 3 (left) shows the UTMOS score obtained during 550 training episodes for each RL approach. For Wavegrad2 baseline model trained with RWR, DDPO, DLPO, KLinR, and DPOK, we follow Wavegrad2 by randomly sampling step $t$ from 1000 diffusion/denoising steps during training, while for Wavegrad2 baseline model trained with $\textrm{DPOK}_{\textrm{tail}}$ and $\textrm{KLinR}_{\textrm{tail}}$, we only sample step $t$ from steps 996, 997, 998, 999, as we find that the KL divergence of current distribution and the distribution of previously trained model differs greatly in the last few steps. As presented in Figure 3 (left)
DLPO increases UTMOS in the first 200 training episodes, shown in Figure 3 (right), and then maintains the UTMOS score around 3.4 for the rest of the samples during the training, while other methods have UTMOS score drop to below 1.5 in the first 200 training episode and then maintain the UTMOS score. Moreover, the $x_{0}$ generated by methods other than DLPO gradually becomes acoustically noisy as the number of seen samples grows. Due to randomness in sampled text, different models receive different initial samples, resulting in varied UTMOS at the beginning, as shown in Figure 3.

Figure: Figure 4: (a) shows the mean UTMOS and NISQA scores for audio based on 200 unseen texts generated by models fine-tuned using seven RL algorithms, the error bar shows the standard deviation of each result. (b) shows the proportion of raters who prefer the audios generated from the DLPO fine-tuned model or baseline model and the proportion of raters who think audios generated by DLPO fine-tuned model and baseline model are about the same (Tie).
Refer to caption: extracted/2405.14632v1/image/infere.png

### 5.2 Algorithm Comparisons

We save the top three checkpoints for each model and use them to generate audios for 200 unseen texts. We then use UTMOS and NISQA to predict MOS score of these generated audios. The results are shown in Figure 4(a).

We see that only audios generated by DLPO fine-tuned Wavegrad2 gain UTMOS and NISQA scores higher than the audios generated by the baseline pretrained WaveGrad2 model. Two two-sample t-tests confirmed significant differences ($p<10^{-20}$) between the baseline and DLPO for both UTMOS and NISQA scores. For models trained with random sampling steps, RWR outperforms DDPO, KLinR and DPOK in fine-tuning TTS diffusion models, though its UTMOS and NISQA scores are below the baseline. Nonetheless, it is observed that audios generated through methods other than DLPO exhibit significant fricative noise, indicating lower speech quality. We observe that the models trained with tail steps sampling, both $\textrm{DPOK}_{\textrm{tail}}$ steps and $\textrm{KLinR}_{\textrm{tail}}$ steps, generate audio with higher UTMOS and NISQA scores compared to the model trained with DPOK and KLinR using random sampling. The detailed results of UTMOS and NISQA scores for each checkpoints and links to sample audios are shown in Appendix C.

We further conduct an experiment recruiting human participants to evaluate the speech quality of the audios generated by DLPO. We randomly select 20 audios pair from the generated audios (among audios of 200 unseen texts). Each pair includes one generated audio from DLPO and one from the baseline pretrained WaveGrad2 model, both audios based on the same text. 11 listeners are recruited for the experiment. They are asked to assess which audio is better regarding to speech naturalness and quality. We show the results in Figure 4. In 67% of comparisons, audios generated from DLPO fine-tuned Wavegrad2 is rated as better than audios generated by the baseline Wavegrad2 model, while 14% of comparisons have audios generated by the baseline Wavegrad2 model rated as better. 19% of comparisons are rated as about the same (Tie). The experiment question is shown in Appendix E.

## 6 Discussion

Our research indicates that fine-tuning TTS diffusion models using DLPO results in better outcomes compared to other reinforcement learning methods. This is likely due to DLPO aligning with the training procedure of TTS diffusion models by incorporating the original diffusion model loss as a penalty in the reward function to effectively prevent model deviation in fine-tuned TTS models.
As demonstrated by Ouyang et al. [], RL fine-tuning can be improved by constraining the model to avoid arbitrary deviations from the original model. However, DDPO, DPOK and KLinR do not improve Wavegrad2 in our experiments despite the latter two incorporating KL divergence terms intended to provide this constraint.
The formulas used to calculate loss in these methods incorporate a KL upper-bound estimate [] which we must further approximate in training due to memory constraints, by sampling both trajectories and then inference steps $t$. These approximations are very unstable in practice due to the large number of steps in each inference trajectory (in Appendix D).

Moreover, as demonstrated in our experiments, models trained with KLinR_tail and DPOK_tail generate audios with higher UTMOS and NISQA scores compared to KLinR and DPOK. This suggests that certain inference steps are more crucial for the performance of the fine-tuned model, but approximating KL by sampling steps may miss these. Previous studies did not encounter this issue because the text-to-image diffusion model they used had fewer inference steps and DDPO, DPOK, and KLinR for images could calculate KL based on all inference steps.
It would likely be possible to trade time for performance by using more samples to obtain better KL bounds which would allow stable fine-tuning with these methods. However, we show that the original loss term (DLPO) is the only one which remains stable enough for fine-tuning with single samples. Future work should explore more efficient methods for approximating KL in models with long inference chains.

## Appendix A Algorithm for fine-tuning Wavegrad2 with DPOK

Figure: Algorithm 2 WaveGrad2 with DPOK: Diffusion Policy Optimization with a KL-shaped Reward

## Appendix B Algorithm for fine-tuning Wavegrad2 with KLinR

Figure: Algorithm 3 WaveGrad2 with Diffusion Policy Optimization with KLinR

## Appendix C Fine-tuned Results: UTMOS and NISQA scores

**Table 1: Mean NISQA mos and naturalness scores for generated audios**
| RL Algorithms | $\alpha$ | $\beta$ | UTMOS | NISQA |
| --- | --- | --- | --- | --- |
| baseline | - | - | 2.90 | 3.74 |
| RWR | 1 | - | 2.18 | 3.00 |
| RWR | 5 | - | 1.67 | 2.95 |
| RWR | 0.5 | - | 2.50 | 3.42 |
| RWR | 0.1 | - | 1.46 | 2.84 |
| DDPO | 1 | - | 1.30 | 2.30 |
| DDPO | 5 | - | 1.20 | 1.09 |
| DLPO | 1 | 1 | 3.52 | 3.70 |
| DLPO | 5 | 1 | 3.47 | 3.92 |
| DLPO | 10 | 1 | 3.31 | 3.63 |
| DLPO | 1 | 0.8 | 3.71 | 3.96 |
| DLPO | 1 | 0.6 | 3.53 | 3.70 |
| DLPO | 1 | 0.5 | 3.61 | 3.83 |
| KLinR | 1 | 1 | 1.45 | 2.33 |
| KLinR | 1 | 10 | 1.43 | 0.82 |
| KLinR | 10 | 1 | 1.37 | 1.46 |
| KLinR | 1 | 0.1 | 1.51 | 1.47 |
| KLinR | 1 | 0.01 | 1.42 | 1.29 |
| $\textrm{KLinR}_{\textrm{tail}}$ | 1 | 1 | 1.51 | 2.84 |
| $\textrm{KLinR}_{\textrm{tail}}$ | 1 | 0.1 | 1.51 | 2.68 |
| $\textrm{KLinR}_{\textrm{tail}}$ | 1 | 0.01 | 2.60 | 3.62 |
| $\textrm{KLinR}_{\textrm{tail}}$ | 1 | 0.001 | 1.89 | 2.78 |
| DPOK | 1 | 1 | 1.41 | 2.83 |
| DPOK | 1 | 10 | 1.34 | 2.55 |
| DPOK | 1 | 100 | 1.29 | 1.17 |
| DPOK | 10 | 1 | 1.37 | 1.46 |
| $\textrm{DPOK}_{\textrm{tail}}$ | 1 | 1 | 1.51 | 2.84 |
| $\textrm{DPOK}_{\textrm{tail}}$ | 1 | 0.1 | 1.51 | 2.68 |
| $\textrm{DPOK}_{\textrm{tail}}$ | 1 | 0.01 | 1.70 | 3.62 |
| $\textrm{DPOK}_{\textrm{tail}}$ | 1 | 0.001 | 1.98 | 2.78 |

Sample audios are presented in this website [https://anonymous.4open.science/r/audios-C7F7/](https://anonymous.4open.science/r/audios-C7F7/)

## Appendix D Experiment of KL divergence stability

We run two experiments to evaluate the stability of KL divergence estimation. We first use two fixed checkpoints (baseine checkpoint vs. checkpoint finetuned with DPOK) to calculate the KL divergence of samples at selected inference step between two fixed checkpoints, we refer to as Monte Carlo KL. Then we calculate the KL divergence of the entire audio between two fixed checkpoints, which we refer to as KL upperbound. The results are shown in the following figure. We can see that both Monte Carlo KL and KL upperbound estimation are very unstable in practice, and the KL upperbound has more outliers may due to the large number of steps in estimation.

Figure: Figure 5: shows KL divergence of samples at selected inference steps from fixed policy trained with DPOK and fixed baseline policy, labeled as MonteCarlo_KL; exact KL upperbound of samples from fixed policy trained with DPOK and fixed baseline policy, labeled as KL_upperbound.
Refer to caption: extracted/2405.14632v1/image/klpenalty_ap.png

## Appendix E Human Experiment Example

The following figure shows one example question of our human experiment. We ask participants to listen to two audios and choose the audios that sounds better, regarding the quality of the speech (clear or noisy, intelligibility)

[Uncaptioned image]: extracted/2405.14632v1/image/human_exp.jpg

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Would DLPO affect the diversity of the samples? The rewards (i.e., MOS predictions) appear not capturing any variability of the generator. The authors may also report the KL divergences as another reasonable metric.
soundness: 3
strengths: **Originality**
- This work partially bases on several recent works: (1) DDPO, (2) DPOK, and (3) "Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs" to derive the DLPO. Different from (3), DLPO adds the diffusion model objective to the reward function as a penalty. In this sense, the proposed algorithm is new in the field of diffusion RL.

**Quality**
- The authors introduce the diffusion model loss-guided reinforcement learning policy optimization (DLPO) and provide thorough comparisons with existing RL approaches. The results also indicate that DLPO could enhance the quality of generated speech, achieving higher MOS compared to both baseline models and other RL methods. 
- As long as I checked, the derivation of DLPO is technically sound. 

**Clarity**
- The authors provide clear RL training procedures in Algorithm 1, which helps in understanding and reproducing the proposed method.

**Significance**
- This work is notable for being the first to apply RL techniques to improve the speech quality of TTS diffusion models. This suggests a new direction in enhancing TTS technology.
- The introduction of DLPO is significant as it aligns the training process of TTS diffusion models with the original diffusion model loss, effectively preventing model deviation. This innovative approach leads to better sound quality and naturalness in generated speech, outperforming other RL methods.
summary: The paper introduces DLPO (Diffusion Model Loss-Guided Reinforcement Learning), a method for fine-tuning text-to-speech (TTS) diffusion models using reinforcement learning with human feedback (RLHF). It explores the application of RLHF to TTS models trained on waveform data, leveraging the mean opinion score (MOS) as a proxy loss. The authors compare DLPO against existing RLHF methods and demonstrate that it significantly improves the naturalness and quality of generated speech audio. Through experiments, they find that traditional RL methods like reward-weighted regression (RWR) do not enhance TTS models, while DLPO effectively aligns with TTS training procedures, leading to better performance.
weaknesses: My main concern would be the generalization ability of the proposed DLPO method, while the validation of its effectiveness is only determined with a specific model (i.e., WaveGrad 2) in a specific task (TTS) trained on a limited set of data. This argument seems to be consistent if we look at the performance of DDPO in TTS (i.e., being worse than the base model), which is contradictory to its results in text-to-image generation (as the authors said, DDPO can achieve better performance in text-to-image diffusion model). To enlarge the impact of this work, I suggest the authors additionally apply the DLPO either to an alternative speech synthesis model, or to another task, e.g., the typical text-to-image task using CIFAR-10 as to validate the consistency of improvements.


Minor typos:
- Eq. (3): The s and a are inconsistently formatted (should be italic).
- Page 8: "We plot both UTMOS and NISQA MOS scores obtained in each training episode in Figure 1 to illustrate the fine-tuning progress of Wavegrad2." - "Wavegrad2" should be capitalized as "WaveGrad2" for consistency.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: No more questions.
soundness: 3
strengths: 1.	The first to apply reinforcement learning (RL) to improve the speech quality of TTS diffusion models.
2.	Introduce a new diffusion loss-guided policy optimization (DLPO) which aligns with the training procedure of TTS diffusion models.
summary: They explore the practical application of RLHF to diffusion-based text-to-speech with long-chain diffusion directly on waveform data, leveraging the mean opinion score (MOS) as predicted by UTokyo-SaruLab MOS prediction system as a proxy loss. They introduce diffusion model loss-guided RL policy optimization (DLPO) and compare it against other RLHF approaches, employing the NISQA speech quality and naturalness assessment model and human preference experiments for further evaluation. Their results show that RLHF can enhance diffusion-based text-to-speech synthesis models, and, moreover, DLPO can better improve diffusion models in generating natural and high-quality speech audio.
weaknesses: 1. It is hoped that the human experiment will include not only comparisons between DLPO and baseline, but also comparisons between DLPO and other RL methods. And it is hoped to increase the number of audio pair in human experiment.
2. UTMOS is used in both training and testing, which may lead to model overfitting to UTMOS, and other auto-scoring MOS models can be added for verification.
3. The method is innovative in general, just put the combination of losses previously explored in the image generation diffusion model into the TTS field.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: N.A.
soundness: 1
strengths: The motivation behind this paper is reasonable, as subjective human metrics are often regarded as the golden standard in TTS systems. However, mainstream TTS system training does not typically consider this factor. Therefore, modeling human evaluations or preferences into speech generation is indeed a meaningful pursuit.
summary: This paper explores the practical application of RLHF to diffusion-based text-to-speech with long-chain diffusion directly on waveform data, leveraging the  MOS as predicted by MOS prediction system as a proxy loss. The results shows that DLPO can improve diffusion models in generating natural and high-quality speech audio.
weaknesses: This paper has several issues that, in my view, make it unsuitable for acceptance at a conference of ICLR’s caliber：

1) The current academic focus is primarily on zero-shot TTS, while single-speaker TTS has already reached a relatively mature stage. More importantly, the training data (24 hours) or the pre-trained model used in this paper is too limited. Optimizing based on such a toy model is of limited value, as the performance gains achieved are less than what could be obtained with a modest increase in training data.

2) The majority of training examples in the NISQA dataset are not generated by TTS systems; instead, they are simulated using clean speech with certain techniques, which creates a noticeable domain shift from TTS-synthesized speech. A clear indication of this is that MOS predictors trained on NISQA perform poorly on datasets like SOMOS. Therefore, this metric is not appropriate.

3) Using only a single small model, WaveGrad 2, makes it difficult to demonstrate the generalizability of the proposed method, especially in the absence of any groundbreaking theoretical contributions in diffusion optimization.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Can DLPO generalized to other TTS models and multi-speaker TTS?

2. What are results of speaker similarity and word error rate of the methods?

3. Which cases of TTS generation can DLPO improve?
soundness: 2
strengths: The proposed method can achieve better performance in the task of TTS than previous RLHF methods, like RWR, DDPO, DPOK and KLinR.
summary: The paper proposed a reinforcement learning policy optimization method based on diffusion models. The method could enhance the TTS models, resulting in better naturalness and quality.
weaknesses: The main weaknesses of this paper is that experiments and analysis are insufficient. 

1. The paper used a baseline model (WaveGrad 2) which was proposed in 2021. However, there are many new models that could achieve much better TTS performance. The authors should use more modern baseline models to compare like [1,2,3].

2. The paper was trained on LJSpeech dataset, which only contains a single speaker. It will be more interesing to see the performance of multi-speaker TTS.

3. The method used MOS predicted by a model as the main human feedbacks and evaluations. However, it is known that MOS cannot totally reflect the performance of a TTS model. Speaker similarity and word error rate should also be considered.

4. The model is lack of comparison with other RL-based TTS methods like [4,5,6]. The authors should also compare with them.

5. It would be better to provide some audio samples to listen to. In addition, the authors should analysis more about what RL really benefit the TTS, like analyzing some examples.

[1] Yang, Dongchao, et al. "SimpleSpeech: Towards Simple and Efficient Text-to-Speech with Scalar Latent Transformer Diffusion Models." arXiv preprint arXiv:2406.02328 (2024).

[2] Li, Yinghao Aaron, et al. "Styletts 2: Towards human-level text-to-speech through style diffusion and adversarial training with large speech language models." Advances in Neural Information Processing Systems 36 (2024).

[3] Chen, Yushen, et al. "F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching." arXiv preprint arXiv:2410.06885 (2024).

[4] Chen, Chen, et al. "Enhancing Zero-shot Text-to-Speech Synthesis with Human Feedback." arXiv preprint arXiv:2406.00654 (2024).

[5] Tian, Jinchuan, et al. "Preference Alignment Improves Language Model-Based TTS." arXiv preprint arXiv:2409.12403 (2024).

[6] Hu, Yuchen, et al. "Robust Zero-Shot Text-to-Speech Synthesis with Reverse Inference Optimization." arXiv preprint arXiv:2407.02243 (2024).