FIND: Fine-tuning Initial Noise Distribution with Policy Optimization for Diffusion Models

Venue: acmmm_org_ACMMM_2024_Conference
Published: 2024-04-05 11:50:46

Abstract: In recent years, large-scale pre-trained diffusion models have demonstrated their outstanding capabilities in image and video generation tasks. However, existing models tend to produce visual objects commonly found in the training dataset, which diverges from user input prompts. The underlying reason behind the inaccurate generated results lies in the model's difficulty in sampling from specific intervals of the initial noise distribution corresponding to the prompt. Moreover, it is challenging to directly optimize the initial distribution, given that the diffusion process involves multiple denoising steps. In this paper, we introduce a Fine-tuning Initial Noise Distribution (FIND) framework with policy optimization, which unleashes the powerful potential of pre-trained diffusion networks by directly optimizing the initial distribution to align the generated contents with user-input prompts.  To this end, we first reformulate the diffusion denoising procedure as a one-step Markov decision process and employ policy optimization to directly optimize the initial distribution. In addition, a dynamic reward calibration module is proposed to ensure training stability during optimization. Furthermore, we introduce a ratio clipping algorithm to utilize historical data for network training and prevent the optimized distribution from deviating too far from the original policy to restrain excessive optimization magnitudes. Extensive experiments demonstrate the effectiveness of our method in both text-to-image and text-to-video tasks, surpassing SOTA methods in achieving consistency between prompts and the generated content. Our method achieves 10 times faster than the SOTA approach.

Main paper: 
 introduction: 
 Nowadays, the generative capabilities of diffusion models [14, 28, 35] have gained widespread recognition in the domain of text-toimages [27, 28, 30 ], text-to-videos [11, 38] , and text-to-3D [22, 37] .Although diffusion-based generative models excel at creating highquality content, the semantics of the content generated frequently fail to align closely with the input text prompts. The current model tends to generate highly correlated objects and concepts even with explicit and clear prompts. For example, as shown in Fig. 1 , the generated image is still about an astronaut on the Moon with the prompt 'a dog on the moon'. It remains a significant challenge to ensure consistency between the generated content and the prompt's description.To address this problem, numerous works [24, 25, 49, 50] are proposed to ensure consistency through the use of additional control signals, such as depth maps, edge maps, etc. Although achieving promising results in alignment with users' intentions, these methods require substantial computational resources for large-scale training with auxiliary models. Besides, the provision of additional control mediums is a burden for users. On the other hand, several pioneer methods [3, 10] leverage reinforcement learning to align the generated images with the prompt by fine-tuning network parameters iteratively. The advantage of these approaches is that the output content is approaching to align with the input prompt without needing additional training data. However, due to the sampling nature of reinforcement learning and the extensive optimization of LoRA-like [10] networks, these methods result in longer training times for individual prompts.Inspired by recent works [4, 8, 43] , we observe the essential capability of large-scale pre-trained models to generate diverse and high-quality controllable visual content with a zero-shot fashion. The initial noise distribution significantly impacts the final generated outcomes, influencing aspects such as layout, color, and semantics of the generated content [35] . The reason for the misalignment of the baseline diffusion model partially comes from the sampling bias between the standard normal distribution and the unusual complex prompts provided by users. The training of diffusion utilizes the standard normal distribution, making it easier for the network to generate samples similar to those in the training set when sampling from a normal distribution in testing, as illustrated by the red boxes in the first row of Fig. 1 . Based on these findings, our motivation is to directly adjust the initial noise to align the generated content with user prompts without any training of the baseline model and extra network structures. After adjusting the initial noise, the baseline model can generate highly aligned images and videos with unconventional prompts, as shown in the second row of Fig. 1 . However, since diffusion processes require multiple denoising steps, it is challenging to calculate the loss on the final generated result to backpropagate gradients to the initial noise distribution.In this paper, we propose a novel framework Fine-tuning Initial Noise Distribution (FIND) of diffusion model to align the content generated more closely with the input text prompt by adjusting the initial noise distribution with policy optimization. To this end, we formulate the entire optimization process as a one-step Markov decision process and employ policy gradients to optimize the initial distribution. The proposed approach allows the baseline model to bypass the multiple intermediate denoising steps and directly optimize the initial distribution based on the reward derived from the final generated outcome. To ensure the accuracy of the optimization direction for the parameters of the initial distribution, we introduce a dynamic reward calibration module to predict the expected reward of the current initial distribution. As shown on the right side of Fig. 1 , we need to optimize our initial noise distribution to increase the expected value of the reward without compromising generative performance. The optimization process is then guided by the difference between the reward of sampled data and the expected reward. To further stabilize fine-tuning, a ratio clipping algorithm is proposed to reuse the historical data to minimize the discrepancy between new and old policies by directly constraining the difference in output action probabilities. Extensive experiments demonstrate the effectiveness and efficiency of our proposed framework in both text-to-image and text-to-video tasks, surpassing SOTA methods in consistency and speed.Our main innovations are as follows:• To the best of our knowledge, we are the first to propose an initial distribution optimization framework based on policy gradients. The proposed framework is a general approach for diffusion-based generative models to produce content that is semantically closer to its input prompt. • We formulate the optimization as a one-step MDP to efficiently adjust the initial distribution. Dynamic reward calibration module and ratio clipping algorithm are proposed to ensure the accuracy and stability of optimization.• Extensive experiments demonstrate that our proposed work can be applied to both image and video diffusion models. Our proposed method is about an order of magnitude faster compared with SOTA. The source code will be released. 
related works 2.1 diffusion-based generation models: 
 Diffusion model [14, 35] is a novel type of generative model that progressively denoise a Gaussian noise into a sample conforming to a learned data distribution by predicting the noise. Generating samples at high resolutions leads to significant computational costs for the denoising model. Latent Diffusion Model (LDM) [28] addresses this issue by utilizing a Variational Autoencoder (VAE) [16] to shift the denoising process from the pixel level to the latent space, significantly reducing computational overhead. Diffusion models have been applied across various generative tasks in different domains, achieving impressive results. These applications include image generation [7, 27, 28, 30] , audio generation [9, 12, 21, 31] , 3D object generation [6, 23, 29, 44] , and robotics-related generation [5, 36, 42, 51] tasks. Despite their ability to generate high-quality content, these models exhibit limited control over the generated outcomes, which affects their applicability in practical scenarios.To address the issue of limited control, researchers have proposed a variety of solutions. Some works leverage fine-tuning techniques to enhance models with extra conditioning layers, building upon the foundation of pretrained models. ReCo [47] and GLIGEN [20] use bounding boxes as conditional controls. SceneComposer [48] and SpaText [1] generate images by segmentation maps. ControlNet [49] is a significant contribution to this field, introducing a parallel network alongside the U-Net architecture. Beyond segmentation maps, ControlNet [49] is capable of processing various types of input, including depth maps, normal maps, canny maps, and so on. Several other works, such as Uni-ControlNet [50] , UniControl [25] , and T2I-Adapter [24] , similarly integrate various conditional inputs to control the generation. To save on computational costs, some approaches [4, 8, 43] directly control the generation of objects during the inference phase through attention maps. These control methods modulate the generated content with additional inputs, but they offer no assistance in enhancing the control of the generated content through more precise prompts. Recent works [18, 41] also try to improve the alignment of textto-image models by a reward model. The reward model is trained from a pre-trained vision-language model such as CLIP [26] or BLIP [19] by asking annotators to compare generations (learn from human feedback). Several studies [3, 10] frame fine-tuning as a multi-step decision-making process, showing that RL fine-tuning exceeds the performance of supervised fine-tuning with rewardweighted loss in reward optimization. These approaches enhance the alignment between generated content and prompts. However, the need to optimize the entire network results in high training costs and the significant fluctuations between new and old policies lead to instability during the training process. Our method only requires optimizing the initial noise distribution, significantly reducing computational overhead and our dynamic ratio clipping algorithm smoothens the policy updates, making the optimization process more stable. 
preliminaries: 
 In this section, we briefly revisit the fundamental concepts of diffusion models and the optimization objectives based on rewards.Diffusion models are designed to produce high-quality, diverse content controlled by text prompts. To reduce computational costs, Rombach et al. [28] proposed a Latent Diffusion Model (LDM) that conducts the denoising process in a latent space. This model features a Variational Autoencoder (VAE) with an encoder E to condense the original image from pixel to latent space, and a decoder D to revert from latent to pixel space. The U-Net, denoted as 𝜖 𝜑 , is involved and its structure comprises alternating down-sampling and upsampling blocks, connected by middle blocks, each equipped with convolutional layers and spatial transformers to streamline image creation. The training of the U-Net hinges on a noise prediction loss function:EQUATIONwhere 𝑧 0 is the latent code of the training sample, c is the text prompt condition, 𝜖 is the Gaussian noise, and 𝑡 is the time step. The noised latent code z 𝑡 is determined as:EQUATIONwhere 𝑎 𝑡 is a hyper-parameter used for controlling the noise strength based on time 𝑡.Sampling from a diffusion model initiates by selecting a random vector z T ∼ N (0, 𝐼 ), which then undergoes the reverse diffusion process 𝑝 𝜃 (z 𝑡 -1 |z 𝑡 , c). This procedure generates a sequence {z 𝑇 , z 𝑇 -1 , . . . , z 0 }, culminating in the final sample z 0 . When employing DDIM [35] as the sampling method, the reverse process is described as follows:EQUATIONwhere 𝜎 2 𝑡 is fixed timestep-dependent variance.The optimization problem of policy gradients [34] is framed within the context of a Markov Decision Process (MDP), which is defined by the tuple (𝑆, 𝐴, 𝜌 0 , 𝑃, 𝑅), with 𝑆 as the state space, 𝐴 as the action space, 𝜌 0 indicating the distribution of initial states, 𝑃 as the transition kernel, and 𝑅 representing the reward function. At each timestep 𝑡, an agent observes a state s 𝑡 ∈ 𝑆, chooses an action a 𝑡 ∈ 𝐴, earns a reward 𝑅(s 𝑡 , a 𝑡 ), and transitions to the next state s 𝑡 +1 following 𝑃 (s 𝑡 +1 |s 𝑡 , a 𝑡 ). Actions are determined by adhering to a policy 𝜋 (a|s).In the context of MDP, the agent's interactions generate trajectories, defined as sequences of states and actions 𝜏 = (s 0 , a 0 , . . . , s 𝑇 , a 𝑇 ). The objective of policy optimization is to maximize the agent's expected cumulative reward over these trajectories, denoted as J 𝜋 , which are sampled according to its policy: EQUATION 
methods 4.1 overview: 
 In this section, we provide a detailed presentation of the proposed FIND framework. Firstly, we introduce FIND formulation which is utilized along with policy gradients to optimize our initial distribution. To ensure the accuracy of our optimization direction, we introduced DRCM. Moreover, RCA is proposed to leverage historical data and enhance the stability of our training process. The pipeline is shown in Fig. 2 .Based on the findings from DDIM [35] , it becomes evident that the initial noise plays a crucial role in our final generated content. Given that the diffusion model requires a multi-step denoising process, optimizing our initial distribution through direct value-based methods is not feasible. Instead, policy gradient techniques optimize the distribution of initial noise directly based on the reward, which is determined by the consistency between the generated content and the prompt. We model the process of optimizing our initial distribution using policy gradient as a one-step MDP as follows:EQUATIONin which s and a represent the state and action, respectively. The value of s is a constant, specifically the input prompt c. a is the initial noise z 𝑡 . 𝑓 𝑟 denotes the reward function, which is used to calculate the similarity between the generated contents and the input prompt. 𝑝 𝜃 refers to the probability of sampling 𝑧 𝑇 given 𝜃 . We formulate the policy 𝜋 𝜃 as follows:EQUATIONwhere 𝜃 = {𝜇, 𝜎 2 }. 𝜇 and 𝜎 2 are two tensors, and their sizes are the same as that of z 𝑇 . Each element in them represents the mean and variance of an individual Gaussian distribution, corresponding to the element in z 𝑇 , respectively. We formulate it as z 𝑇 ( 𝑗) ∼ N (𝜇 ( 𝑗), 𝜎 ( 𝑗) 2 ), where 𝑗 is the index of the element in z 𝑇 . We target 𝜋 𝜃 as our optimization goal and employ the method of policy gradients to refine it. The term 𝜋 𝜃 (a) refers to the probability of sampling action a under the policy 𝜋 𝜃 . DDIM [35] is utilized as our denoising strategy, ensuring that once the initial noise is specified, the resultant denoised image remains constant. Consequently, the entire T-step denoising process is treated as our environment. As shown in the blue part of Fig. 2 , the U-Net is frozen and does not participate in the backward process. Specifically, the (i+1)-th optimization iteration unfolds as follows: Firstly, action z t is sampled from the policy N (𝜇 𝑖 , 𝜎 2 𝑖 ) which is optimized i times. Then, z 𝑡 is denoised to z 0 by the environment. The reward function 𝑓 𝑟 assigns a reward based on the generated content z 0 and state c. We use this reward to optimize the policy 𝜋 𝜃 via policy gradients [34] . This optimization process is repeated multiple times until the reward is maximized. Our objective is to maximize the expected reward: To optimize this objective function, we express it in the form of gradients:EQUATIONEQUATIONWe present the proof of Eq.8 in Appendix A.1.According to the theory of policy gradient [34] , the optimization direction is determined by the reward. Theoretically, if the generated content matches the prompt, the reward should be positive; otherwise, it should be negative. Since the 𝑓 𝑟 is a pre-trained model, zero is not the dividing line for the quality of the reward. As illustrated in right part of Fig. 3 , we observe that although the reward value of our sampled yellow point is greater than 0 (indicated by the red plane), it is less than the expected reward of the current initial distribution (indicated by the blue plane). Considering this sampling point as a positive reward for optimizing the initial distribution is incorrect.The distance to the expected reward 𝑟 of 𝜋 𝜃 is what we required. A straightforward approach is to estimate 𝑟 through a large amount of samplings, which results in significant time consumption. We propose a Dynamic Reward Calibration Module (DRCM) to predict 𝑟 of 𝜋 𝜃 by a simple 3-layer MLP network 𝑔 defined as 𝑟 = 𝑔(𝜃 ). The loss function of 𝑔 is formulated as:EQUATIONHowever, due to the lack of a ground truth dataset for the distribution and the expected reward, it is difficult to pre-train 𝑔. Considering the 𝑁 times multiple optimization steps involved in the entire process, we optimize 𝑔 using the rewards 𝑟 corresponding to the sampled z 𝑇 at each optimization step, as well as the initial distribution, as shown in the left part of Fig. 3 . We reformulate the loss function of 𝑔 as follows:EQUATION)where 𝑚 is the number of samples in the current optimization iteration. Considering the efficiency of optimization, here 𝑚 is set to 1. We define the optimized reward for our current sample as 𝑟 * = 𝑟 -𝑟 , as the difference between the reward obtained from sampling and the reward predicted by the network. As shown in the right part of Fig. 3 , the yellow sampling point is treated as a negative reward for optimization by DRCM. 𝑟 * is then used in Eq.7 to optimize our initial distribution. 𝑟 ← 𝑓 𝑟 (z 0 , c)7:𝑟 ← 𝑔(𝜃 ) optimize 𝑔 by Eq.10 9:𝑟 * ← 𝑟 -𝑟 10:Optimize 𝜃 by Eq.13 11: end while 12: output: Optimized initial distribution 𝜋 𝜃When using Eq.7 to optimize the initial distribution, the update process is limited to the data samples that are currently sampled. The requirement for multiple denoising steps significantly slows down the diffusion model's inference time, resulting in suboptimal optimization efficiency when repeated sampling is necessary. Further, optimizing the initial distribution solely based on the feedback from the reward function without any constraints compromises the generative performance of the original diffusion model. This issue arises because the diffusion model is trained with initial noise sampled from a standard normal distribution. If our optimized distribution deviates too far from the initial distribution, it creates a gap between the training and generation processes. We propose the Ratio Clipping Algorithm (RCA) to limit the extent of each optimization step by the historical data. Inspired by TRPO [34] , we employ importance sampling, which enables the network to incorporate historical data into its updates, thereby enhancing the overall efficiency of the optimization process. We reformulate Eq.7 to a loss function as follows:EQUATIONwhere 𝜋 𝜃 old is the policy of the previous step. We formulate Eq.11 in a gradient form:EQUATIONWe present the proof of Eq.12 in Appendix A.2. After establishing Eq.11, we optimize the diffusion network's initial noise distribution by leveraging historical data. Distinct from DPOK [10] , which utilizes the KL divergence from the initial model to moderate the extent of parameter updates to avoid too much deviation from the original model. Our RCA, inspired by the findings of sDPO [15] , adopts 𝜋 𝜃 old as the reference model. This is based on the insight that comparing parameters with those from the previous step provides a more effective upper bound for updating parameters. Specifically, we define the ratio of new policy and old policy as 𝜂 = 𝜋 𝜃 𝜋 𝜃 old . When the new policy is equal to the old policy, 𝜂 is equal to 1. To limit the magnitude of updates to the new policy, we set a margin 𝜆, ensuring that 𝜂 falls within the range of [1 -𝜆, 1 + 𝜆].We reformulate Equation 10 as follows:L * 𝑝 = L 𝑝 , if 𝜂 ∈ [1 -𝜆, 1 + 𝜆] 0. else (13)Our entire optimization process is outlined in Algo.1. 
experiments 5.1 experimental setups: 
 We utilize Stable Diffusion v1.5 [28] and ModelScpoe [38] as our image and video generation base model, whose parameters remain frozen throughout our optimization process. Our optimization targets are the mean 𝜇 and variance 𝜎 2 of the initial distribution, whose sizes are 4x64x64 for image generation and 4x16x64x64 for video generation. Compared to optimizing the entire model, this approach significantly reduces the computational cost of optimization. For the reward model, we employ ImageReward [44] trained on a large dataset with human judgments for image generation and ViCLIP [39] for video generation. 𝜆 is set to 0.02. The number of total optimization steps 𝑁 is 150. The learning rate is set to 0.001 and the optimizer is AdamW. All our experiments are conducted on a single 3090 GPU and the VRAM consumption is less than 10GB for image generation and 15GB for video generation.To verify the effectiveness of our method, we conduct both qualitative and quantitative experiments. We compare the proposed method with the standard Stable Diffusion v1.5 [28] . Additionally, we also compare our approach with state-of-the-art approaches DPOK [10] that optimize the entire U-Net using Reinforcement Learning to highlight the efficiency and effectiveness of optimizing the initial noise. Quality Results. Following the similar setting from DPOK [10] , we select four prompts, A green dog is running on the grass, A dog and a cat, Four pandas, A dog on the moon, for fair comparison. As shown in Fig. 4 , in the color aspect, we can see that the baseline model generates clear images of dogs but struggles with unusual colors. DPOK generates green dogs, but in the first column, the appearance of the dog is blurred, possibly due to finetuning the network, which weakened its generative performance. In the second column, the dog generated by DPOK is not entirely green. Our method generates not only green dogs but also dogs with a very complete appearance.In terms of composition, the baseline struggles to generate multiple subjects in one image. Using DPOK, the generated images of cats and dogs exhibit some overlap, impacting the quality of the generation. Our method is capable of combinations of multiple subjects.In the counting aspect, the baseline method and DPOK struggle to generate multiple subjects of the same type in one image while our method achieves higher completeness. In terms of location, for unusual positions, the baseline method struggles and only generates common ones, such as an astronaut on the moon. The DPOK method generates dogs on the moon, but the clarity is not high. Our method is capable of generating dogs with brighter colors and higher completeness. Quantity Results. In this section, we validate our proposed method from two perspectives: generative capability and computational cost. For assessing the quality of generation, we employ two metrics: ImageReward [45] and Aesthetic Score [32] . ImageReward evaluates the alignment between the generated images and the prompts. The Aesthetic Score assesses the aesthetic quality of the generated images. To validate our results, we conduct tests using prompts in four different aspects as quality evaluation. For each prompt, we select 100 images for evaluation. As shown in Tab.1, we observe that our method has advantages in terms of ImageReward, which assesses text-image alignment, particularly in the aspects of Color, Composition, and Location. In terms of Count, our method demonstrates only a slight discrepancy compared to DPOK. This highlights our method's significant advantage over both the baseline and previous SOTA methods in aligning text and images, demonstrating that we unleash the potential of the pre-trained model. In terms of Aesthetics, our method shows advantages in Composition and Count, but overall, the difference between our method, the baseline, and DPOK is minimal. This indicates that our approach, while optimizing for control effects as dictated by the prompts, does not negatively impact generative performance; instead, it may even enhance it. Versatility. Followed by the experiment setting of DPOK, we evaluate our method on prompts from Drawbench [30] (10 images per each prompt). As shown in the bottom of Tab.1, our method outperforms the baseline and DPOK as well in a larger set of prompts evaluation settings. This demonstrates the comprehensive generative capability of our proposed method. Time Consumption. As shown in Tab.2, our method achieves approximately 13 times faster over DPOK, completing optimization in around 14 minutes for a single prompt. This demonstrates the practicality of our approach.We conduct ablation studies by utilizing two complex prompts: A red book and a yellow vase. and oil portrait of Batman holding a picture of Spiderman, intricate, elegant, highly detailed, lighting, painting, art station, smooth, illustration, art by Greg Rutkowski and Alphonse Mucha. We generate 100 samples for each scenario to serve as our test dataset. Impact of DRCM. The DRCM primarily predicts the expected reward value under the current initial distribution, aiming to prevent the network from optimizing in incorrect directions. As illustrated on the left side of the second row in Fig. 5 , removing the DRCM leads to generated vases and books similar to the baseline, but there's a noticeable discrepancy between the colors of the vases and books and the user-input prompts. The first column shows multiple books, the vase in the second column is not yellow, the third column produces multiple vases, and in the fourth column, books turn into a table. On the right side of the second row, we observe that the generated Batman has some features of Spiderman, and the Spiderman image appears somewhat blurred. As shown in Tab.5, removing the DRCM results in a decline in both ImageReward and Aesthetic Score metrics. This is attributed to the absence of calculated expected rewards, relying solely on the sign of the reward to determine the direction of optimization leads to suboptimal solutions.Impact of RCA. As demonstrated on the left side of the third row in Fig. 5 , the first column transforms a red book into a red bowl, the second column morphs the concept of a red book into a red vase and a yellow book, the third column generates only a red vase, and the fourth column changes the vase's color to brown. On the right half of the third row, although Batman is well generated, the spider picture he holds is poorly generated, and in the third column, although the content of the painting is well generated, the shape of the painting has turned into a trapezoid. From Tab.3, we observe a significant decrease in the ImageReward metric after We recruited 33 voters from social media to assess the advantages of our method. We compared the Baseline model, DPOK, and our method. Each model generated two images from prompts in four categories same as in Sec. 5.2, Color, Composition, Count, and Location, resulting in a total of eight images. Voters were presented with 24 images and their corresponding prompts. Each image was evaluated on two dimensions: Appearance and Alignment, with scores ranging from 1 to 5, from low to high. As indicated in Tab.4, our method outperforms both the Baseline model and DPOK in all aspects. Furthermore, it's evident that users significantly prefer our method for its text-image consistency. It's also noteworthy that our method is nearly 13 times faster than DPOK.Our approach is theoretically applicable to any diffusion-based method, whether it be text-to-image, text-to-video, text-to-3D, and so forth. To demonstrate the versatility of our method, we use text-to-video as a case study, analyzing its performance both qualitatively and quantitatively. Specifically, we employ ModelScope [38] as our baseline model, which is a large-scale text-to-video diffusion model trained on large-scale datasets [2, 33, 46] . ViCLIP [39] is a pre-trained model used to evaluate the similarity between text and video, which is utilized as our reward function. Quality Results. To verify the effectiveness of our method, we selected four sets of prompts that the baseline models struggle to generate directly: A green dog is running on the grass., A dog is running on the moon., A panda is walking on the grass, from left to right. and A monkey is playing guitar. These cover unusual colors, displacement control, anomalous positions, and abnormal behaviors. As illustrated in the top left of Fig. 6 , generating dogs with colors that do not exist in real life proves to be challenging and the dogs generated remain yellow-brown. After optimization with our method, the color of the generated dogs matches the green specified in the prompt. As illustrated in the bottom left, it is challenging for baseline models to control the objects' motion trajectories directly through prompts, leading to objects moving randomly. The panda generated by the baseline model merely turns to the right. In contrast, our method allows the panda to smoothly move from the left to the right side of the screen as requested by the prompt. In the top right corner, the quality of generated objects in anomalous positions is compromised, making the dogs appear blurry. After our optimization, the generated dogs are much clearer, and their movement process becomes smoother. As depicted in the bottom right corner, despite the ability to generate objects engaged in abnormal motions, such as monkeys and guitars, there is no interaction between them. Following our optimization, the model accurately generates behaviors such as monkeys playing guitars. Quantity Results. We select four prompts same as the quality evaluation, producing 100 videos for each prompt to serve as our test dataset. ViCLIP [39] is selected to evaluate the text-video consistency of the generated videos. Following the methodology of the LOVEU-TGVE competition [40] , we employ the CLIP score [13] to assess the consistency between frames. Additionally, PickScore [17] is used to predict user preferences for our model. As shown in Tab.5, our method surpasses the baseline across all three metrics, demonstrating that the videos generated after optimization with our approach have improved in terms of text-video consistency, inter-frame consistency, and predicted user preferences. 
conclusion: 
 In this paper, we introduced FIND, a novel Fine-tuning Initial Noise Distribution with policy optimization framework, to align the content generated by diffusion models with user-input prompts. Unlike previous methods that required extensive training or additional controls, our approach was capable of optimizing for any prompt within just 13 minutes. We observed that the initial noise significantly influences the final output of diffusion models, leading us to optimize the initial noise. However, optimizing the initial distribution from the generated images was challenging due to the multi-step denoising required by diffusion models. We utilized policy gradients to circumvent the multi-step denoising, optimizing the initial distribution directly through the reward function. To ensure that the optimization direction was not solely determined by the sign of the reward, we proposed the DRCM to predict the expected value of the reward under the current distribution. Additionally, we developed the RCA module to leverage past samples and ensured optimization stability. Both quantitative experiments and qualitative tests have proven the effectiveness of our proposed method. Moreover, our approach can be applied to kinds of diffusion-based generative models, demonstrating its high generalizability and versatility.  
find: fine-tuning initial noise distribution with policy optimization for diffusion models: 
 Reward SpaceUpdate distributionThe expected reward ( , ) Figure 1 : Our FIND framework optimizes the initial distribution of any diffusion-based model to enhance the consistency between generated content and the prompts provided by users. Before optimization, the semantics of generated images and videos could diverge from the prompt for complex scenes, as indicated by the red boxes. By optimizing the overall expected reward of consistency, the content within the green box becomes consistent with the prompt.In recent years, large-scale pre-trained diffusion models have demonstrated their outstanding capabilities in image and video generation tasks. However, existing models tend to produce visual objects commonly found in the training dataset, which diverges from user input prompts. The underlying reason behind the inaccurate generated results lies in the model's difficulty in sampling from specific intervals of the initial noise distribution corresponding to the prompt. Moreover, it is challenging to directly optimize the initial distribution, given that the diffusion process involves multiple denoising steps. In this paper, we introduce a Fine-tuning Initial Noise Distribution (FIND) framework with policy optimization, which unleashes the powerful potential of pre-trained diffusion networks by directly optimizing the initial distribution to align the generated contents with user-input prompts. To this end, we first reformulate the diffusion denoising procedure as a one-step Markov decision process and employ policy optimization to directly optimize the initial distribution. In addition, a dynamic reward calibration module is proposed to ensure training stability during optimization. Furthermore, we introduce a ratio clipping algorithm to utilize historical data for network training and prevent the optimized distribution from deviating too far from the original policy to restrain excessive optimization magnitudes. Extensive experiments demonstrate the 
 
 Appendix:  


Decision: The reviewers agreed that this paper proposes novel and interesting research. LtJH calls the work “highly relevant and timely”. This reviewers extensive list of strengths includes that the paper is detailed in terms of overall examination of the topic, that it is well-referenced, that the proposed method is “noteworthy for it’s innovative approach” and “outperforms state of the art methods”.

2/3 reviewers suggest that this work is “definitely suitable” and I agree that it is very appropriate work for ACM MM.

In terms of limitations, LtJH suggests further exploration of the motivation and extension of experiments to other media, e.g., 3D models.

Tykj notes several questions related to the approximation of intial distribution, and the potential for generating a good range of outputs from one prompt.

All reviewers noted issues with the time required to train the noise distribution for each prompt and this was still a significant issue following the rebuttal.

tYkJ puts it well: “According to Table 3, it takes approximately 14 minutes to fine-tune a single prompt, which seems excessive given that generating a single image with a diffusion model typically only takes a few seconds. In contrast, DPOK[1] learns the information into the model parameters and can be fine-tuned on multiple prompts simultaneously (section 5.5 of [1]).”

This issue is certainly discussed in the rebuttal and I felt that the response was reasonable given that this is an unconventional approach.

My feeling is that, although this length of time is unusually long compared to present diffusion models, this unique idea has shown effectiveness and may inspire other modifications to the diffusion process. Although 14 minutes per prompt is impractical today, this good idea may last for a long time in the literature. I feel that we should support interesting, appropriate albeit unusual ideas to support intellectual diversity at MM.
TYKJ gave me permission to not worry about this concern and so given the mix of ratings, I will lean towards acceptance, particularly given that this is a very suitable topic for MM.

Reviews: review: paper_summary: paper_summary: This paper addresses the challenges of aligning generated content with user prompts in diffusion-based generative models. The authors introduce the FIND framework, which optimizes the initial noise distribution using policy optimization to enhance consistency between generated content and user prompts. The paper thoroughly examines the limitations of existing methods, highlighting the difficulties in achieving prompt-consistent generation and the challenges in directly optimizing the initial noise distribution.
strength_weakness: strength_weakness: 1.The paper addresses a highly relevant and timely topic by introducing the FIND framework, which focuses on optimizing the initial noise distribution of diffusion models to enhance the consistency between generated content and user prompts. This topic is highly relevant in current generative model research.
2.The paper provides a comprehensive examination of this topic, including the analysis of reasons for inconsistencies between generated content and prompts in existing models, the design of the framework, and experimental validation. The content is detailed and rich.
3.The introduction of policy optimization into the initial noise distribution optimization framework and the proposal of dynamic reward calibration module and ratio clipping algorithm to ensure stability and accuracy during the optimization process is particularly noteworthy for its innovative approach.
4.The manuscript is well-referenced, indicating that the authors have conducted an in-depth review and analysis of existing research in the fields of diffusion models, controllable generative models, and policy optimization.
5.The authors' application of the one-step Markov decision process to optimize the initial noise distribution and the extensive experiments demonstrating the effectiveness and superiority of the method in both image and video generation tasks are commendable. The proposed method significantly outperforms existing state-of-the-art methods in achieving consistency between generated content and prompts, with a training speed improvement of an order of magnitude.
limitations: limitations: 1.Although the paper proposes the idea of improving generative performance by optimizing the initial noise distribution, the explanation of this motivation and its significance is not sufficiently in-depth. Readers might need more background information to understand why existing methods do not adequately address this issue and what the unique advantages of this approach are.
2. Although the paper conducts extensive experiments, they are primarily focused on text-to-image and text-to-video generation tasks. It could be further extended to other generative tasks, such as 3D object generation, to demonstrate the broader applicability of the method.
3. The proposed method in the paper includes multiple modules (such as the dynamic reward calibration module and the ratio clipping algorithm). While these modules help with optimization, they also add complexity to the method. The paper does not adequately discuss the impact of these modules on the overall system complexity and computational cost.
4. The paper does not provide sufficiently detailed descriptions of the optimization process and parameter settings, especially the implementation details of the policy gradient optimization. This may lead to difficulties for readers attempting to reproduce the work.
paper_summary: paper_summary: This paper proposes a method to fine-tune the initial noise of the conditional diffusion model to improve the consistency between prompts and the generated image. This method outperforms SOTA methods and reduces the generation time.
strength_weakness: strength_weakness: This paper reveals that the initial noise of diffusion models can be controlled to tune the generated image. The initial noise is optimized using a reinforcement learning method, which is novel and interesting.
limitations: limitations: 1) It seems that the network $g$ needs to be trained for each prompt. I am a bit concerned about the time required for generating large-size images. In Table 2, the size of the images is not mentioned.

2) It would be better to have a comparison of the memory usage. For generating high-resolution images, it seems like the network $g$ requires a huge memory. Assume the size of images as 256$\times$256, The input size of network $g$ will be $393216=2\times3\times256\times256$. The network $g$ is a 3-layer MLP. If the first layer outputs features with a size of 393, the weight of the first layer will be $w\in\mathbb{R}^{393216\times 393}$ which is huge. Can authors provide more detail about network $g$?
paper_summary: paper_summary: This paper proposes an algorithm, FIND, to enhance the alignment between input prompts and generated samples in diffusion-based text-to-image/video generation models. FIND refines the initial noise distribution of a given input prompt using a policy gradient algorithm. Additionally, the authors apply the dynamic reward calibration module and ratio clipping to improve the optimization performance.
strength_weakness: strength_weakness: 1. FIND algorithm optimizes the initial noise distribution for each input prompt instead of the diffusion model parameter. This changes the multi-step MDP setting into a single-step MDP setting, which might reduce the optimizing difficulties.

2. The qualitative samples and quantitative scores look better the baselines.
limitations: limitations: I have several concerns:
1. The FIND algorithm relies on per-prompt optimization of the initialization distribution, which raises concerns about sampling efficiency. According to Table 3, it takes approximately 14 minutes to fine-tune a single prompt, which seems excessive given that generating a single image with a diffusion model typically only takes a few seconds. In contrast, DPOK[1] learns the information into the model parameters and can be fine-tuned on multiple prompts simultaneously (section 5.5 of [1]). The finetuned model from [2] can be also applied to different prompts. Therefore, it might not be fair to directly compare the optimization times of FIND and DPOK, as DPOK can work on various prompts after fine-tuning, whereas FIND is only suitable for one specific prompt. The lengthy optimization time for a single prompt could limit FIND's applicability in many real-world applications.

2. FIND approximates the initial distribution for a given prompt as a simple Gaussian distribution. I am uncertain whether this approximation is valid. If we consider this distribution as a policy, then this straightforward policy form needs to capture the landscape formed by the full diffusion process and the reward model, which seems to be very complex. Specifically, I have the following questions:

(a) After the distribution is fitted, can different samples from the same initial distribution generate different valid images/videos corresponding to the same prompt, or do they collapse into a single mode of nearly identical samples? Figure 4 shows different samples for the same prompt, but I'm not sure whether these samples come from different instances of the same fitted initial distribution or from different fitted initial distributions for the same prompt.

(b) Similar to question (a), under different random seeds, will the algorithm converge to the same initial distribution, or will it capture different modes?

(c) Failure case study: Given different random seeds and prompts, can FIND always find a valid initialization? Given an initialization distribution, are all samples from this distribution valid, or will there be cases where some samples are valid while others are not?

3. Besides building the sampling policy upon the initialization distribution, FIND essentially follows the same reinforcement learning framework as [2]. The ratio clipping algorithm is also very similar to the formulations of TRPO [3] and PPO [4].

[1] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models.
[2] Training diffusion models with reinforcement learning.
[3] Trust region policy optimization.
[4] Proximal policy optimization algorithm.