Ctrl-U: Robust Conditional Image Generation via Uncertainty-aware Reward Modeling

## Abstract

Abstract In this paper, we focus on the task of conditional image generation, where an image is synthesized according to user instructions.
The critical challenge underpinning this task is ensuring both the fidelity of the generated images and their semantic alignment with the provided conditions.
To tackle this issue, previous studies have employed supervised perceptual losses derived from pre-trained models, i.e. , reward models, to enforce alignment between the condition and the generated result.
However, we observe one inherent shortcoming: considering the diversity of synthesized images, the reward model usually provides inaccurate feedback when encountering newly generated data, which can undermine the training process.
To address this limitation, we propose an uncertainty-aware reward modeling, called Ctrl-U , including uncertainty estimation and uncertainty-aware regularization, designed to reduce the adverse effects of imprecise feedback from the reward model.
Given the inherent cognitive uncertainty within reward models, even images generated under identical conditions often result in a relatively large discrepancy in reward loss. Inspired by the observation, we explicitly leverage such prediction variance as an uncertainty indicator.
Based on the uncertainty estimation, we regularize the model training by adaptively rectifying the reward.
In particular, rewards with lower uncertainty receive higher loss weights, while those with higher uncertainty are given reduced weights to allow for larger variability.
The proposed uncertainty regularization facilitates reward fine-tuning through consistency construction.
Extensive experiments validate the effectiveness of our methodology in improving the controllability and generation quality, as well as its scalability across diverse conditional scenarios, including segmentation mask, edge, and depth conditions. Code, models, and demo will soon be available at this https URL .

## 1 Introduction

Driven by the emergence of large-scale image-text datasets (Schuhmann et al., ; Changpinyo et al., ; Schuhmann et al., ) and the development of diffusion models (Dhariwal & Nichol, ; Nichol et al., ; Rombach et al., ), text-to-image (T2I) diffusion models (Ramesh et al., ; Saharia et al., ) have become fundamental in the field of controllable visual generation. These models excel at producing realistic, high-quality images that accurately reflect the descriptions provided in natural language. However, due to their inherent properties, text-based conditions often struggle to convey the detailed controls required for final generation results efficiently. This limitation becomes especially evident in certain scenarios where text prompts alone can not capture all necessary details, such as when depicting unique artistic styles or accurately rendering complex scenes. To this end, alongside text descriptions, a substantial body of research focuses on integrating novel conditional controls, such as user-drawn sketches and semantic masks, into T2I diffusion models (Qin et al., ; Ye et al., ; Zhang et al., ; Mou et al., ).
Despite investigations into the controllability of T2I diffusion models and their expanded applications (Gao et al., ; Li et al., ; Chen et al., ),
achieving precise and fine-grained control remains challenging. The primary issue lies in ensuring both the fidelity of the generated images and their semantic alignment with the provided conditions. To tackle this issue, some efforts, *e.g.*, ControlNet++ (Li et al., ) attempt to employ a pre-trained reward model to extract the corresponding condition from the generated images and enforce alignment between the specified condition and the output. However, we observe that the reward model inevitably produces inaccurate feedback (see Fig. 1). As the timestep $t$ increases,
and the generation diversity also increases. Some newly generated data reflects the correct condition but has not been "seen" by the reward model during training, resulting in inaccurate feedback.
For instance, although the denoised image at $t=0$ aligns well with the input condition, there are still some mIoU errors according to the reward model.
If we do not rectify such misleading rewards, enforcing alignment between the provided conditions and the inaccurate
predictions on newly generated data will consistently
compromise the training process.

To mitigate the adverse effects of inaccurate rewards, we introduce a robust, controllable image generation approach via uncertainty-aware reward modeling (Ctrl-U).
Motivated by our observation, we explicitly model the reward uncertainty to facilitate the reward back-propagation.
The proposed uncertainty-based method consists of two phases, *i.e.*, uncertainty estimation and uncertainty regularization.
(1) In the uncertainty estimation phase, we forward the identical input condition twice with different noise timesteps, yielding two similar generation results.
We explicitly leverage the reward variance between these two generations as an uncertainty indicator.
Notably, we do not introduce extra parameters and thus do not impact the inference efficiency of the diffusion model.
(2) Following the uncertainty estimation, we adaptively adjust the loss weights of different reward feedback. Specially,
the alignment reward for each pixel
are rectified by its corresponding pixel-wise uncertainty.
Rewards with lower uncertainty, indicating greater stability, should be given higher weights to encourage the model to learn from these reliable signals.
Conversely, rewards with higher uncertainty, which are less stable, should be assigned reduced weights to minimize the negative impact of potentially inaccurate feedback.
Quantitative and qualitative experiments verify the efficacy of the proposed method on controllability and image quality across various conditions.

- •
We observe an inherent drawback in enforcing alignment for conditional image generation using a pre-trained reward model, as the reward model often fails to generalize to newly generated data. To address the adverse effects of inaccurate reward feedback on conditional image generation, we introduce an uncertainty-aware reward modeling approach, termed Ctrl-U, which adaptively regularizes the reward learning process.
- •
Through extensive experiments on five benchmarks across three datasets, *i.e.*, ADE20k, COCO-Stuff, and MultiGen-20M, we validate the effectiveness of our methodology in improving controllability and generation quality. Our approach also shows scalability across various conditional scenarios, including segmentation masks, edges, and depth conditions.

Figure: (a)
Refer to caption: x1.png

## 2 Related Works

Conditional Generation.
Recently, diffusion probabilistic models (Ho et al., ; Song et al., ) have
become an important cornerstone in general image generation.
As the quality of generated images improves, a critical challenge
remains, *i.e.*, achieving precise control over these generative models to adequately meet the intricate and varied demands of real-world applications.
Motivated by the development of guidance mechanism (Song et al., ; Ho & Salimans, ), T2I diffusion models such as GLIDE (Nichol et al., ), Imagen (Saharia et al., ), DALL·E 2 (Ramesh et al., ) excel at modeling fine-grained structures and texture details. To facilitate the training of diffusion models on limited resources,
Latent Diffusion Model (LDM) (Rombach et al., ) maps the diffusion process from pixel space to the latent space. Building upon the foundation of LDM framework, Stable Diffusion (SD) exhibits exceptional capabilities in T2I generation.
Despite the astonishing capabilities of T2I diffusion models, language, with its sparse and high-level semantic nature, is unsuitable for processing intricate and low-level control images, *e.g.*, depth maps. To achieve conditional control in T2I diffusion models, researchers are exploring the integration of various control signals with text descriptions. ControlNet (Zhang et al., ) integrates image-based conditions by incorporating an additional encoder copy into frozen T2I diffusion models via zero convolutions.
Similarly, GLIGEN (Li et al., ) and T2I-Adapter (Mou et al., ) propose utilizing independent adapters (or extra modules) to synchronize internal knowledge within T2I models with external conditions. Furthermore, Cocktail (Hu et al., ) integrates multi-modal control information within a single diffusion model to yield more high-quality outputs. With the rapid development of large language models, some recent works have explored prompt engineering for regulated generation. For instance, ReCo (Zhang et al., ) specifies text descriptions of different regions to model the fine-grained structure and aesthetic features. Taking one step further, Control-GPT (Zhang et al., ) applies GPT-4 (OpenAI, ) as the sketch generator to incorporate control signals.
However, one drawback is that their generated images usually deviate from input conditions. In this work, we take a closer look at denoised results and explicitly introduce the uncertainty to adaptively impose consistency constraints.

Uncertainty-aware learning. As data-driven technologies rapidly advance, the imperative for model reliability grows. Measuring the "confidence" of predictions has been a long-standing problem.
As a result, researchers are progressively focusing on uncertainty as a viable solution.  Kendall & Gal () categorize uncertainty into two primary types: epistemic uncertainty and aleatoric uncertainty. The former epistemic uncertainty denotes model uncertainty, reflecting variations in model weights even when trained on the same dataset. The representative approach in this direction is Bayesian networks (Jensen & Nielsen, ; Neal, ), which focuses on learning the distribution of weights and estimates uncertainty by assessing the distribution variance.
Similarly, Monte Carlo Dropout (Gal & Ghahramani, ) is proposed to effectively simulate variational Bayesian approximation by randomly using varying dropout masks.
Another line of research on epistemic uncertainty (Raghu et al., ; Nandy et al., ; Lee & AlRegib, ) adapts the original model to measure prediction uncertainty directly. However, incorporating this additional uncertainty estimation typically compromises both the training efficiency and predictive accuracy of the model. In addition, ensemble methods (Lakshminarayanan et al., ; Malinin et al., ; Wenzel et al., ) combine various deterministic models in the prediction process to improve prediction accuracy, but is constrained by the computational burden associated with operating multiple independent networks and the requisite diversity across ensemble models. Another family of research tackles the inherent uncertainty in data, arising from noise interference or ambiguous annotations. This methodology for addressing uncertainty has been applied across various fields: such as image segmentation (Zheng & Yang, ), image retrieval (Chen et al., ), image classification (Litrico et al., ), person re-identification (Zhang et al., ), and 3D object detection (Zhang et al., ). In particular,  Dou et al. () effectively simulates uncertainty by utilizing model prediction variance to inject noise into the latent space. In a manner akin to (Kendall & Gal, ), a dynamic uncertainty-based loss is introduced to enhance training stability. Unlike existing uncertainty-based works, we explicitly introduce noise at varying intensities into the diffusion model to enable the quantification of uncertainty. It is worth noting that we explore the utilization of uncertainty in the untapped conditional image generation task, preserving the original diffusion learning and mitigating the adverse effects of inaccurate feedback from the reward model.

## 3 Method

### 3.1 Uncertainty Estimation

We provide a brief overview of the pipeline in Fig. 2. To simplify the explanation, here we take the segmentation task as an example, if not specified.
Given the triplet input, *i.e.*, one source image input $x_{0}$, one text prompt, and the conditional control $c$, the basic encoders extract the visual feature $z_{0}$ of $x_{0}$, textual feature $f_{t}$ and the control feature $f_{c}$ of $c$.
We fix the weight of the off-the-shelf image encoder and the text encoder, while finetuning the control encoder. During the conditional diffusion training, we first
add Gaussian noise $\epsilon$ to the feature map $z_{0}$ as the noisy latent. In particular, we conduct two generation forward with identical condition $c$ but different $t_{1}$ and $t_{2}$, and resampled Gaussian noise $\epsilon$.
Therefore, the two noisy latents $z_{1}$ and $z_{2}$ can be formulated as:

$$ $z_{1}=\sqrt{\bar{\alpha}_{t_{1}}}z_{0}+\sqrt{1-\bar{\alpha}_{t_{1}}}\epsilon, \quad z_{2}=\sqrt{\bar{\alpha}_{t_{2}}}z_{0}+\sqrt{1-\bar{\alpha}_{t_{2}}} \epsilon,\quad\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I}),$ (1) $$

Following the ControlNet pipeline (Zhang et al., ), we further fuse text condition $f_{t}$ and image condition $f_{c}$ to predict the added noise.
After removing the predicted noise, we obtain the recovered latent, *i.e.*, $\hat{z}_{0}^{1}$ and $\hat{z}_{0}^{2}$.
Then, given the latent $\hat{z}_{0}^{1}$ and $\hat{z}_{0}^{2}$, the pretrained decoder is to reconstruct the input image as $\hat{x}_{0}^{1}$ and $\hat{x}_{0}^{2}$, respectively.

Figure: Figure 2: A brief overview of our pipeline. Here, we take the segmentation mask as a conditional generation example. (a) Conditional Generation. Given text, source image $x_{0}$, and the conditional control $c$, we extract feature $z_{0}$, $f_{t}$, $f_{c}$, respectively. Then, we fine-tune the Diffusion model to generate two intermediate features for the image decoder. (b) Uncertainty Learning. Given the two features, we decode the two images, *i.e.*, $\hat{x}_{0}^{1}$ and $\hat{x}_{0}^{2}$. Then we apply the reward model to obtain the two layout predictions $\hat{c}_{1}$ and $\hat{c}_{2}$. We leverage the prediction discrepancy as the uncertainty indicator to rectify the original reward loss.
Refer to caption: x3.png

To align the condition within the generated images, we follow the existing works and apply an off-the-shelf reward model $D$ to quantify the consistency between the input condition $c$ and the corresponding output condition of the generated images.
However, as shown in Fig. 1(b), the reward model usually contains inaccurate feedback, even if the visual layout has already aligned with the condition.
If we consistently apply strong supervision on the consistency between inaccurate rewards and conditions predicted on the newly generated data, the backpropagation of incorrect gradients will significantly compromise the model.
To estimate inaccrate rewards, we explicitly leverage two diffusion forwards for the identical input conditions.
We compare the reward discrepancy between extracted conditions $\hat{c}_{1},\hat{c}_{2}$ from generated images, which can be considered as a reward indicator at the current timestep. For the segmentation map condition, considering the probability output, we quantify uncertainty by calculating the KL-divergence between the extracted conditions of two generations at the pixel level as:

$$ $\displaystyle U_{1}$ $\displaystyle=\mathbb{E}\left[\hat{c}_{1}\log\left(\frac{\hat{c}_{1}}{\hat{c}_ {2}}\right)\right]=\mathbb{E}\left[D(\hat{x}_{0}^{1})\log\left(\frac{D(\hat{x} _{0}^{1})}{D(\hat{x}_{0}^{2})}\right)\right],$ (2) $\displaystyle U_{2}$ $\displaystyle=\mathbb{E}\left[\hat{c}_{2}\log\left(\frac{\hat{c}_{2}}{\hat{c}_ {1}}\right)\right]=\mathbb{E}\left[D(\hat{x}_{0}^{2})\log\left(\frac{D(\hat{x} _{0}^{2})}{D(\hat{x}_{0}^{1})}\right)\right].$ $$

For other non-probability conditions, *e.g.*, edge and depth, we adopt the $\ell_{1}$ distance $U_{1}=U_{2}=|\hat{c}_{1}-\hat{c}_{2}|$ as the uncertainty indicator.

Discussion. 1). Why not use an auxiliary network to directly regress uncertainty? Some existing methods (Zheng & Yang, ; Jain et al., ) introduce extra modules to model the uncertainty. However, directly regressing uncertainty often leads to overfitting, where the model predicts zero uncertainty for all samples or a large, uniform value. This line of approaches is challenging to optimize effectively. Therefore, instead of directly regressing uncertainty, we perform predictions twice to estimate the prediction variance, which serves as a cognitive uncertainty indicator for the reward model. Additionally, this method has the side benefit of not introducing the extra training parameters.
2). How about choosing the same timestep for the two-time generations?
Using the same timestep for the two-time generations is possible, and the samples will still differ due to the inherent randomness in
$\epsilon$. However, this approach limits the diversity of the generated images, which hinders the full utilization of the reward model. In our experiments, we observe that using different timesteps for the two generated images helps to learn a better and more robust generation model after rectifying the reward.

### 3.2 Uncertainty Regularization

The existing controllability modeling methods (Li et al., ) usually adopt consistency loss between the input
condition and the extracted condition, which can
be regarded as a pixel-wise supervision. For example, when using segmentation mask as input condition, $\mathcal{L}^{c}$ can be defined as a per-pixel cross-entropy loss as:

$$ $\displaystyle\mathcal{L}_{1}^{c}=-c\log(\hat{c}_{1}),\quad\mathcal{L}_{2}^{c}= -c\log(\hat{c}_{2})$ (3) $$

where $c$ denotes the input condition, $\hat{c}_{1},\hat{c}_{2}$ represent extracted conditions from generated images.
Our objective is to adaptively rectify the inaccurate reward feedback. To achieve this, we adapt the original consistency loss by incorporating our estimated uncertainty:

$$ $\mathcal{L}_{1}^{u}=\frac{\mathcal{L}_{1}^{c}}{\exp\left(U_{1}\right)}+\lambda \cdot U_{1},\quad\mathcal{L}_{2}^{u}=\frac{\mathcal{L}_{2}^{c}}{\exp\left(U_{2 }\right)}+\lambda\cdot U_{2},$ (4) $$

where $\lambda$ denotes the regularization factor. The second term $\lambda\cdot U$ is to prevent the model from consistently predicting high uncertainty across all samples. The first term diminishes the impact of reward feedback when significant disagreement is evident, yet retains its influence when predictions are consistent. Notably, when the uncertainty value is constant, the back-propagation gradient is identical to the original consistency loss.
Finally, to optimize robust conditional image generation, we adopt a combination of the diffusion loss $\mathcal{L}_{\text{diffusion}}$ and the proposed uncertainty-regularization loss $\mathcal{L}_{1}^{u},\mathcal{L}_{2}^{u}$. The original diffusion loss at timestep $t_{1}$ and $t_{2}$ can be reformulated as:

$$ $\mathcal{L}_{\text{diffusion}}=\mathbb{E}\left[\|\epsilon_{\theta}(z_{1},t_{1} ,f_{t},f_{c})-\epsilon\|_{2}^{2}+\|\epsilon_{\theta}(z_{2},t_{2},f_{t},f_{c})- \epsilon\|_{2}^{2}\right].$ (5) $$

Therefore, the total loss is as follows:

$$ $\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{diffusion}}+\mu\cdot(\mathcal{L} _{1}^{u}+\mathcal{L}_{2}^{u}),$ (6) $$

During training, we balance the ratio $\mu$ of diffusion training and reward feedback by:

$$ $\mu=\begin{cases}\mu_{0}&\text{if }t\leq t_{\text{thre}}\\ 0&\text{if }t\textgreater t_{\text{thre}}\end{cases}$ (7) $$

where $\mu_{0}$ is the consistency weight. If $t\textgreater t_{\text{thre}}$, we only consider the diffusion optimization as existing works (Zhao et al., ). if $t\leq t_{\text{thre}}$, we incorporate uncertainty-aware reward modeling.

Discussion. 1). Why set timestep threshold $t_{\text{thre}}$ for $\mu$? Considering that when $t$ is large, $z_{t}$ approaches the random noise $\epsilon$, leading to more diverse recovered $\hat{x}_{0}$ (see Fig. 1).
In such cases, it becomes difficult to ensure that two outputs with identical conditions remain close, which is necessary for our uncertainty estimation, as it requires visually similar inputs to evaluate the reward model.
In practise, we, therefore, empirically set a threshold to avoid the large timesteps for the reward modeling.
We have conducted the experiment on the choice of $t_{\text{thre}}$ (see Section 4.3).
2). What is the advantage of uncertainty regularization in the reward feedback? The advantage of uncertainty regularization in the reward feedback is that it helps to mitigate the adverse effects of imprecise feedback from the reward model. Specifically, by incorporating uncertainty estimation, the proposed method, Ctrl-U, can adaptively adjust the weights of the rewards during training. Rewards with lower uncertainty are given higher loss weights, ensuring that more reliable feedback has a stronger influence on the training process. Conversely, rewards with higher uncertainty are given reduced weights, allowing for greater variability and preventing the model from being overly influenced by potentially inaccurate feedback. This adaptive rectification of rewards based on their uncertainty improves the consistency and reliability of the training process. As a result, the model becomes more robust to the diverse and sometimes unpredictable nature of newly generated data, leading to better controllability and generation quality.

## 4 Experiment

### 4.1 Settings

Datasets. Our experiments are conducted using three datasets: ADE20K (Zhou et al., ; ), COCOStuff (Caesar et al., ) and MultiGen-20M dataset (Qin et al., ), adhering to the dataset construction principles of Controlnet++ (Li et al., ). More specifically, we utilize the ADE20K and COCOStuff for the segmentation task. For Hed and Lineart conditions, we utilize pretrained models, which is the same as ControlNet++, to generate annotations for the MultiGen-20M dataset. Considering that existing datasets usually include masks for pixels with unknown depth values, we adapt the MultiGen20M depth dataset in a manner similar to the dataset construction method used by ControlNet (Zhang et al., ). For datasets without image captions, we use the captions generated by MiniGPT-4 (Zhu et al., ). The training and inference are conducted at a resolution of 512x512 for all datasets and methods.

Evaluation and Metrics.
To evaluate semantic segmentation and depth map controls, we adopt mIoU and RMSE as the metrics, respectively. For the Hed edge and Lineart tasks, we calculate SSIM to compare the difference between the extracted control and the ground truth.
For Sampling, we employ the UniPC (Zhao et al., ) sampler, implementing 20 denoising steps to generate images using the original text prompts in accordance with ControlNet v1.1 (Zhang et al., ), without incorporating any negative prompts. For comparative methods, we utilized their publicly available code to generate images, ensuring fairness by using the same validation dataset and adhering to their original inference settings.

Implementation Details.
In our experiments, we first fine-tune the pre-trained ControlNet model to convergence, using Adam as the optimizer with a learning rate of 1e-5, weight decay of 1e-2, and momentum of 0.9. Then, we use the same optimization settings to perform 10k iterations of uncertainty-aware reward fine-tuning. It is worth noting that we employ an one-step efficient reward strategy (Li et al., ) to enhance training efficiency. Following the settings of previous work, we choose slightly weaker models as the reward model and stronger models for evaluation to guarantee fairness in assessment. We utilize 8 A100 (80G) GPUs for all our experiments. See also the supplementary material for detailed reward model settings.

Baselines. In our evaluation, we mainly compare with competitive methods, including ControlNet++ (Li et al., ), T2I-Adapter (Mou et al., ), ControlNet v1.1 (Zhang et al., ), GLIGEN Li et al. (), Uni-ControlNet (Zhao et al., ), and UniControl (Qin et al., ).
These models perform well in controllable text-to-image generation, providing open-source model weights for re-implementation.
For fair comparison, all models are tested under the same image conditions and text prompts. Although these methods utilize the user-friendly SD1.5 for controllable text-to-image generation, recent advancements have led to the adoption of SDXL (Podell et al., ) by several methods. Accordingly, we also present controllability results for ControlNet-SDXL and T2I-Adapter-SDXL. Since currently there are no officially released models, we deploy the third-party ControlNet-SDXL for the following experiments.

### 4.2 Experimental Results

Comparison of Controllability.
We present the results on five benchmarks in Table 1.
Our uncertainty-aware method Ctrl-U significantly outperforms the previous state-of-the-art method ControlNet++ (Li et al., ) by 6.53% in ADE20K and 8.65% in MultiGen20M depth, respectively. Regarding Hed and Lineart edge, the model with uncertainty regularization has obtained +3.76% and +1.06% increase on SSIM. The most significant improvement is observed in COCO-Stuff for segmentation masks, with an increase of 44.42%. This performance enhancement confirms that reducing the adverse effects of imprecise feedback from the reward model can significantly improve the model controllability. We use the same hyper-parameter settings as those in the ControlNet++ experiments, verifying the effectiveness of our uncertainty-aware reward modeling.

**Table 1: Controllability comparison under various conditional controls and datasets. $\uparrow$ denotes higher result is better, while $\downarrow$ indicates lower is better. ’-’ signifies the absence of a publicly available model for testing. The best result in each column is marked bold and the second is underlined. We generate four groups of png images and report their average result to reduce random errors.**
| Condition<br>(Metric) | T2I<br>Model | Seg. Mask<br>(mIoU $\uparrow$) | Hed Edge<br>(SSIM $\uparrow$) | LineArt Edge<br>(SSIM $\uparrow$) | Depth Map<br>(RMSE $\downarrow$) |  |
| --- | --- | --- | --- | --- | --- | --- |
| Dataset | ADE20K | COCO-Stuff | MultiGen-20M | MultiGen-20M | MultiGen-20M |  |
| ControlNet | SDXL | - | - | - | - | 40.00 |
| T2I-Adapter | SDXL | - | - | - | 0.6394 | 39.75 |
| T2I-Adapter | SD1.5 | 12.61 | - | - | - | 48.40 |
| Gligen | SD1.4 | 23.78 | - | 0.5634 | - | 38.83 |
| Uni-ControlNet | SD1.5 | 19.39 | - | 0.6910 | - | 40.65 |
| UniControl | SD1.5 | 25.44 | - | 0.7969 | - | 39.18 |
| ControlNet | SD1.5 | 32.55 | 27.46 | 0.7621 | 0.7054 | 35.90 |
| ControlNet++ | SD1.5 | 43.64 | 34.56 | 0.8097 | 0.8399 | 28.32 |
| Ctrl-U (Ours) | SD1.5 | 46.49 | 49.91 | 0.8401 | 0.8488 | 25.86 |

**Table 2: FID ($\downarrow$) comparison under various conditional controls and datasets. ’-’ signifies the absence of a publicly available model for testing. The best result in each column is marked bold and the second is underlined. We generate four groups of png images and report their average result to reduce random errors.**
| Method | T2I<br>Model | Seg. Mask | Hed Edge | LineArt Edge | Depth Map |  |
| --- | --- | --- | --- | --- | --- | --- |
| ADE20K | COCO-Stuff | MultiGen-20M | MultiGen-20M | MultiGen-20M |  |  |
| Gligen | SD1.4 | 33.02 | - | - | - | 18.36 |
| T2I-Adapter | SD1.5 | 39.15 | - | - | - | 22.52 |
| UniControlNet | SD1.5 | 39.70 | - | 17.08 | - | 20.27 |
| UniControl | SD1.5 | 46.34 | - | 15.99 | - | 18.66 |
| ControlNet | SD1.5 | 33.28 | 21.33 | 15.41 | 17.44 | 17.76 |
| ControlNet++ | SD1.5 | 29.49 | 19.29 | 15.01 | 13.88 | 16.66 |
| Ctrl-U (Ours) | SD1.5 | 28.61 | 15.79 | 11.59 | 11.99 | 15.48 |

Comparison of Image Quality.
We employed the Fréchet Inception Distance (FID) (Heusel et al., ) to measure the distribution distance between generated and real images in Table 2. We could observe that, compared with existing methods, Ctrl-U exhibits superior FID values in various conditional generation tasks. Notably, in the COCO-Stuff for segmentation masks and MultiGen20M for Hed edge, Ctrl-U achieves impressive improvements of 18.14% and 22.74%, respectively. This significant boost indicates that our method not only enhances the controllability but also improves the quality of image generation, underscoring the efficacy of our proposed uncertainty-aware reward modeling in adaptively adjusting the weights of the rewards during training.

Comparison of CLIP Score.
To assess the potential impact on text controllability, we reported the CLIP-Score metrics across various conditional generation tasks in Table 3. We discovered that Ctrl-U generally exhibits comparable or superior CLIP-Score results in most cases, suggesting that our approach not only enhances the controllability of conditional controls but also maintains the original proficiency of text-to-image generation. After communicating with the authors, we re-implemented the CLIP-Score for ControlNet++ (Li et al., ) and marked it with ^∗ in Table 3.

Qualitative Analysis. We show the visual comparisons between our Ctrl-U and previous state-of-the-art methods across various conditional scenarios in Fig. 3. Given the same text prompts and image-based conditions, we observe that existing methods usually generate images with areas that do not align with the image conditions. Other methods, taking the segmentation mask generation task as an example, produce content on the building that is irrelevant to the provided condition, reflecting relatively poor controllability. Similarly, under edge and depth conditions, other methods fail to accurately represent the nuances of various image controls. In contrast, images generated by Ctrl-U exhibit better consistency with the input conditions.

Human Evaluation.
Following ControlNet++ (Li et al., ), we choose four types of conditional controls (segmentation masks, Hed edge, Lineart edge, and depth conditions) for human evaluation. We invite 20 participants to
rank each result based on three distinct
criteria. According to the average rankings shown in Table 4, users prefer the results produced by our method over those generated by the comparative methods.

**Table 3: CLIP-score ($\uparrow$) comparison under various conditional controls and datasets. ’-’ signifies the absence of a publicly available model for testing. The best result in each column is marked bold and the second is underlined. ^∗ and gray denote the CLIP-score we re-implemented using the checkpoints provided by Controlnet++. We generate four groups of png images and report their average result to reduce random errors.**
| Method | T2I<br>Model | Seg. Mask | Hed Edge | LineArt Edge | Depth Map |  |
| --- | --- | --- | --- | --- | --- | --- |
| ADE20K | COCO-Stuff | MultiGen-20M | MultiGen-20M | MultiGen-20M |  |  |
| Gligen | SD1.4 | 31.12 | - | - | - | 31.48 |
| T2I-Adapter | SD1.5 | 30.65 | - | - | - | 31.46 |
| UniControlNet | SD1.5 | 30.59 | - | 31.94 | - | 31.66 |
| UniControl | SD1.5 | 30.92 | - | 32.02 | - | 31.68 |
| ControlNet | SD1.5 | 30.16 | 31.18 | 31.46 | 31.26 | 32.11 |
| ControlNet++ | SD1.5 | 31.96 | 13.13 | 32.05 | 31.95 | 32.09 |
| ControlNet++^∗ | SD1.5 | 31.24 | 30.93 | 30.49 | 30.34 | 30.02 |
| Ctrl-U (Ours) | SD1.5 | 31.26 | 31.23 | 32.05 | 31.87 | 31.72 |

**Table 4: Results of Human Evaluation**
| Evaluation Category | Ours | ControlNet++ | ControlNet | UniControl | UniControlNet |
| --- | --- | --- | --- | --- | --- |
| Image-Condition Alignment | 72.5% | 7.5% | 6.2% | 7.5% | 12.5% |
| Image Quality | 56.2% | 10.0% | 15.0% | 15.0% | 13.7% |
| Image-Text Alignment | 50% | 27.5% | 22.5% | 7.5% | 12.5% |

### 4.3 Ablation Studies and Further Discussion

Design of Uncertainty Estimation. We present an ablation study on the design of the two-time generation in Table 5(b). To mitigate the adverse effects of inaccurate rewards, we forward the identical input condition twice with different noise timestep to estimate uncertainty. Specifically, we adjust the disparity between the two timesteps, *i.e.*, $|t_{1}-t_{2}|$. A short interval, such as $t_{1}=t_{2}$, where the only randomness stems from resampled noise $\epsilon$, limits the diversity of the generated images. Considering the generated images are too similar, the reward discrepancy is small, and could not serve as uncertainty indicator. Conversely, a long interval indicates a significant gap between the two noisy latents, which, in turn, increase the generation discrepancy. Too large discrepancy also impacts the accurate uncertainty estimation, and thus compromises reward modeling. When $|t_{1}-t_{2}|=1$, the model achieves the optimal FID and relatively strong mIoU and CLIP-score.

Impact of Different Timestep Threshold $t_{\text{thre}}$. We investigate the impact of varying timestep threshold $t_{\text{thre}}$ (see Eq. 7) in Table 5(b). We observe that, as the $t_{\text{thre}}$ increases, the scope of our uncertainty-aware reward rectification broadens, leading to improved mIoU. However, while increasing the rectification range can improve alignment between generated images and input conditions, setting the threshold too high will negatively affect image quality, resulting in higher FID scores. We find that the $t_{\text{thre}}=1$ setting strikes an optimal balance, facilitating uncertainty estimation and regularization to adaptively correct the reward learning process.

Impact of Regularization Weight $\lambda$. We conducted a study to analyze the impact of varying uncertainty regularization weight $\lambda$ in Table 5(d). When using a low value for $\lambda$, which imposes a minimal penalty for high uncertainty, it leads to excessively high uncertainty values across all samples. Conversely, a high value of $\lambda$ imposes a strong penalty on high uncertainty, suppressing the impact of uncertainty during reward fine-tuning. The setting with $\lambda=1$ achieves an optimal balance. This configuration provides enough model capacity to accurately fit reward feedback while preventing overfitting to inaccurate rewards. As a result, the adaptive rectification of rewards based on uncertainty regularization improves the consistency and reliability of the training process.

Impact of the Consistency Weight $\mu_{0}$. We conduct an ablation study on the consistency weight $\mu_{0}$ as shown in Table 5(d). We notice that the model shows insensitivity to the $\mu_{0}$ change in terms of CLIP-score. Setting $\mu_{0}$ to 1 yields the highest mIoU but at the cost of FID performance. This is due to the high weight of the consistency loss, which biases model training towards reward learning, thereby enhancing controllability but compromising generation quality. When $\mu_{0}=0.1$, FID performance reaches its optimum, and mIoU remains strong, effectively balancing diffusion training and uncertainty-aware reward learning without causing the uncertainty to either vanish or explode.

Efficacy of Uncertainty-aware Reward Modeling. We validate the superiority of our learnable uncertainty-aware reward method by comparing it with the vanilla reward learning. As depicted in Fig. 4, under the identical input conditions and timesteps, although the generated images align with the given conditions, reward learning without uncertainty produces inaccurate predictions. On the contrary, our uncertainty-aware reward fine-tuning generates higher-quality images. Additionally, the reward model can accurately extract conditions from these images. This improvement is attributed to the reward uncertainty, which adaptively rectifies inaccurate rewards and improves the reliability of the training process. Specifically, rewards with lower uncertainty are given higher loss weights, enhancing the influence of precise feedback in the training process. Conversely, rewards with higher uncertainty are assigned lower weights, leading to greater variability and ensuring the model is not excessively influenced by potentially inaccurate feedback. Consequently, the model becomes more robust to the diversity of the newly generated images, enhancing both controllability and the quality of generation.

**Table 5: Ablation studies on the ADE20K dataset. We report mIoU, FID and CLIP-score to evaluate controllability and image quality respectively. (a) We show the impact of the interval between $t_{1}$ and $t_{2}$. We find that a short interval limits the diversity of the generated images, making it challenging to estimate uncertainty. In contrast, a long interval impacts accurate uncertainty estimation and thus compromises reward modeling. (b) We study the timestep threshold $t_{\text{thre}}$. We find that the $t_{\text{thre}}=1$ setting strikes an optimal balance, facilitating uncertainty estimation and regularization to adaptively correct the reward learning process. (c) The impact of the regularization weight $\lambda$. With $\lambda=1$, we obtain best results. This setting provides enough model capacity to accurately fit reward feedback while preventing overfitting to inaccurate rewards. (d) The impact of the consistency weight $\mu_{0}$. Considering the balance between diffusion training and uncertainty-aware reward learning, we set $\mu_{0}=0.1$ to ensure enhanced controllability while improving generation quality.**
| $|t_{1}-t_{2}|$ | 0 | 1 | 3 | 5 | 7 | 9 |
| --- | --- | --- | --- | --- | --- | --- |
| mIoU | 45.33 | 46.49 | 46.72 | 45.94 | 45.11 | 44.88 |
| FID | 29.01 | 28.61 | 29.01 | 29.81 | 29.10 | 28.94 |
| CLIP-score | 30.92 | 30.89 | 31.01 | 31.06 | 30.93 | 31.09 |

Figure: Figure 3: Qualitative comparisons with different conditional controls on unseen test images.
Refer to caption: x4.png

Figure: Figure 4: Ablation study on the uncertainty-aware reward modeling.
Refer to caption: x5.png

## 5 Conclusion

In this work, we focus on the challenge of ensuring both high fidelity and semantic alignment in conditional image generation. We highlight a significant limitation in existing methods that use pre-trained reward models for enforcing alignment. Yet these models often provide inaccurate feedback when encountering diverse, newly generated data, which can negatively impact the training process.
To alleviate this issue, we propose an uncertainty-aware reward modeling approach, termed Ctrl-U, which incorporates uncertainty estimation and adaptive regularization.
This method assigns higher loss weights to rewards with lower uncertainty and reduces the weights for highly uncertain rewards, thereby enhancing the consistency and reliability of the training process.
Our extensive experiments across five benchmarks on three datasets, *i.e.*, ADE20k, COCO-Stuff, and MultiGen-20M, validate the effectiveness of our methodology in improving controllability and generation quality. Additionally, Ctrl-U shows robust scalability, performing well across various conditional scenarios, including segmentation masks, edges, and depth conditions. These results indicate the potential of our method to contribute to advancements in conditional image generation.

## Appendix A Appendix

### A.1 More Implementation Details.

Dataset Details.
We use the ADE20K dataset for segmentation masks, which includes 20,210 images in the training set and 2,000 images in the validation set. The dataset contains 150 types of segmentation annotations, covering most objects found in daily life scenes. Since ADE20K (Zhou et al., ; ) lacks text prompts, we follow ControlNet++ (Li et al., ) and use MiniGPT-4 (Zhu et al., ) to generate image captions by instructing: "Please briefly describe this image in one sentence". Similarly, COCO-Stuff provides segmentation annotations, with 118,287 images in the training set and 5,000 in the validation set. It is challenging to find datasets with accurate annotations for Hed and LineArt edge. To tackle this issue, we adhere to the dataset construction principles of ControlNet (Zhang et al., ) and train our model using the MultiGen20M (Qin et al., ) dataset, which includes annotations generated by the model.
The images and their corresponding captions are displayed in Fig. 5.

Reward Model Details. We provide details of the models and evaluation metrics in Table 6. Considering that the Hed and Lineart edge extraction models are neural networks without non-differentiable operations, we achieve differentiability by modifying the forward code. We have improved the reward fine-tuning process by the proposed uncertainty-aware reward modeling, resulting in significant enhancement in controllability and generation quality across four control conditions. Our research will expand to encompass additional control scenarios, such as human pose and scribbles.

**Table 6: Details of the reward model, evaluation model, and training loss under different conditional controls. ControlNet* indicates that we utilize the same model for condition extraction as employed in ControlNet (Zhang et al., ).**
|  | Seg. Mask | Depth Edge | Hed Edge | LineArt Edge |
| --- | --- | --- | --- | --- |
| Reward Model (RM) | UperNet-R50 | DPT-Hybrid | ControlNet* | ControlNet* |
| RM Performance | ADE20K (mIoU): 42.05 | NYU (AbsRel): 8.69 | - | - |
| Evaluation Model (EM) | Mask2Former | DPT-Large | ControlNet* | ControlNet* |
| EM Performance | ADE20K (mIoU): 56.01 | NYU (AbsRel): 8.32 | - | - |
| Consistency Loss | CrossEntropy Loss | MSE Loss | MSE Loss | MSE Loss |

### A.2 More Visualization

More visualization results across different conditional controls for our image generation on the validation set are shown in Figures 6, 7, 8, 9. We observe that our uncertainty-aware framework is capable of generating diverse images that align well with the given inputs across semantic masks, edge, and depth conditions. Attributed to our proposed uncertainty estimation and regularization, the adaptive rectification of inaccurate rewards enhances the diversity of the generated data and its alignment with the given conditions.

Figure: Figure 5: Text prompts generated by MiniGPT-4
Refer to caption: x6.png

Figure: Figure 6: More visualization results of ours on unseen test images (Segmentation Mask)
Refer to caption: x7.png

Figure: Figure 7: More visualization results of ours on unseen test images (Depth Map)
Refer to caption: x8.png

Figure: Figure 8: More visualization results of ours on unseen test images (Hed Edge)
Refer to caption: x9.png

Figure: Figure 9: More visualization results of ours on unseen test images (LineArt Edge)
Refer to caption: x10.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The annotation in Fig. 2 and Eq. 2 needs to be clarified. Specifically, is the KL divergence computed between the input control (c) and c1, c2 separately? Or is it computed only between c1 and c2? While Fig. 2 denotes the former, Eq. 2 denotes the latter.
2. What is the range of t1 and t2? Also, do the optimal t1 and t2 vary on different datasets/tasks?
3. Visualization of inaccurate prediction from the reward model and the corresponding uncertainty prediction. Are they aligned? In other words, does the uncertainty reduce the weight loss of positions where the predicted condition matches the input but somehow has a significant loss (the reward model predicts the wrong segmentation mask)?
4. Line 360, can the authors elaborate more on "After communicating with the authors, we re-implemented the CLIP-Score for ControlNet++ (Li et al., 2024b) and marked it with ∗ in Table 3"?
5. Does the reward model impact the performance if a different one is used? Also, does the consistency prediction generalize well to different distributions? (e.g., different segmentation datasets)
soundness: 3
strengths: 1. The proposed method is intuitive and easy to incorporate.
2. The quantitative results outperform the baselines by a large margin (although some metrics used seem to be questionable, e.g., mIoU from another segmentation model; see Weaknesses)
summary: The authors investigate the reward model used in conditional generation diffusion models and find that it tends to output wrong conditions (e.g., segmentation masks) on the de-noised images. They propose leveraging the output uncertainty on two different generated results from different denoising timestamps. They reduce the loss weight of pixels with high uncertainty while increasing ones with low uncertainty.
weaknesses: 1. If I understand correctly, the authors leverage another more powerful segmentation model to evaluate mIoU. However, this seems strange to me. For instance, in Fig. 3, while the authors circled out the additional green masks by ControlNet++, I think this shouldn't be regarded as an error. Because the green boxes seem to be windows, and the powerful segmentation model predicts these masks, it could be that ControlNet++ did a great job generating the windows based on the concept of the pinkish building mask.
2. Please see Questions.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: see Weaknesses
soundness: 3
strengths: The method is straightforward and effective, and the paper is well-organized and easy to understand.
summary: This paper addresses the issue of conditional image generation. The authors identify a significant flaw in the perceptual loss used for conditional image generation, where it may provide incorrect supervisory signals for some newly generated images, thereby affecting the results. To address this, the authors propose a simple yet effective approach to assess the reliability of supervisory loss by estimating the uncertainty in the generated images. Specifically,  during the training process they measure the variability between two different images generated under the same condition as an indicator of the uncertainty in the supervisory loss, suggesting that a high variability implies an unreliable signal. This method automatically corrects the supervisory signals, thereby enhancing the quality of the generated models. The conclusions are supported by detailed experimental evidence.
weaknesses: 1. There is a lack of important comparisons, particularly the absence of direct comparison with methods that use auxiliary networks to measure uncertainty (e.g.  Learning model uncertainty as
variance-minimizing instance weights(ICLR 24), It can be used to conduct an ablation experiment to compare who is the better way to measure uncertainty for Conditional Image Generation), discussed on line 214, which is a significant oversight.

2. The motivation section in the introduction is somewhat abrupt and some statements, such as those mentioned on line 74 regarding "the denoised image at t = 0 aligns well with the input condition," are not intuitive and require clearer explanation.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. As specified in L215-217, authors stated that directly regressing uncertainty can incur indistinguishable predictions over samples. However, In L239-240, the proposed methods also require additional terms to prevent indistinguishable predictions. Considering that the proposed method can introduce double training time, what’s the major advantage of the proposed method?
2. I'm confused by the value of t in equ. 7. Specifically, how does t in equ.7 relate to t1 and t2 introduced in equ.5?
soundness: 2
strengths: 1. This paper targets the uncertainty problem existing in training with reward models, which is critical and meaningful.

2. Despite increasing training time for each training step, the proposed method avoids incorporating additional modules to assess reward uncertainty.
summary: Targeting  conditional image generation, the authors proposed an uncertainty-aware reward modeling method, named Ctrl-U. This method includes uncertainty estimation and uncertainty-aware regularization, aimed at mitigating the negative effects of inaccurate feedback from the reward model. The authors use the prediction variance as an uncertainty indicator, and based on this, they adaptively adjust the reward during model training. Rewards with lower uncertainty are given higher loss weights, while those with higher uncertainty are given reduced weights, allowing for more variability.
weaknesses: 1. Lack of novelty. The major contribution of this work is the methodology that generates two images, calculates their uncertainty, and reweights the reward loss. Such a solution seems to be trivial, double the time of each training step, and bring neglective quantitative improvement. 
Maybe the author can report the computational cost required before and after using the proposed method. 
In addition, why do not the authors collect a dataset composed of paired generated images and their uncertainty scores and then train a simple predictor to predict uncertainty for each image to remove the necessity to generate two images in one step? In my opinion, such an uncertainty predictor can also be adaptive to different generative models and can significantly improve the training efficiency compared to the proposed method. Since the predictor can be more lightweight than generative models and can be trained on image latents.

2. Unfair comparison. The method proposed in this paper is mainly used for post-training. In other words, authors need to pretrain a conditional model such as ControlNet before employing the proposed training losses (Line 292-295). However, during quantitative comparison, authors compare their methods with pretrained models. Such a comparison seems to be inappropriate. It may be better to report the performance improvement when applying the proposed method to different base models, showing the generalization ability of the proposed method.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Please refer to Weaknesses.
soundness: 3
strengths: I appreciate the motivation behind this paper, and I believe the method presented is quite sound.
summary: This paper made improvements to the controllability of diffusion-based conditional generative models by introducing the concept of uncertainty.
The authors designed loss functions for probabilistic and non-probabilistic conditions by assessing the differences between images in two inference processes.
The authors claim that they have achieved SOTA performance on both objective and subjective metrics.
weaknesses: 1. We found that the loss function proposed in this paper can lead to artifacts, such as the linerart in Fig 3. The proposed method can achieve alignment between the outputs and the conditions, but the authors do not discuss how to maintain the visual plausibility of the generated results while ensuring this alignment.

2. In Tab.5a, there seems to be no clear evidence of what the optimal value of delta t is. We hope the authors can further discuss the choice of delta t. I am curious about why a relatively small delta t (1% of total 1000 timesteps) can encourage the diversity of generated samples.

3. From Fig.2, it can be seen that the gradient obtained from 2 different timesteps is back-propagated, and the parameters of the UNet will be also optimized. We hope the authors can provide some discussion on the training computational cost. Have the authors tried to only optimize the parameters of the control net?

4. This point does not affect our rating, but we want to know why the ControlNet method on SDXL in Tab.1 performs significantly worse than on SD1.5.

Meta-review: Summary: 

This paper identifies an issue in the perceptual loss used in diffusion-based conditional image generation, where incorrect supervisory signals can degrade quality. Correspondingly, the authors propose an uncertainty-aware reward modeling method (Ctrl-U) that uses prediction variability between two generated images from different denoising timestamps as an uncertainty indicator to adaptively adjust loss weights – rewards with lower uncertainty are given higher weights, while those with higher uncertainty are down-weighted. This approach demonstrates improved controllability, reliability, and state-of-the-art conditional image generation performance.

Strengths: 

Based on the reviewers' comments, the proposed method is intuitive, sound, and effectively addresses the critical issue of uncertainty in reward models for conditional image generation. It outperforms baselines by a significant margin in quantitative results and avoids the need for additional modules to assess reward uncertainty, making it a straightforward and effective solution. The paper is well-organized and easy to understand.

Weaknesses:

Based on the reviewers' comments, there was a perceived lack of technical novelty, as the method simply generates two images, calculates uncertainty, and reweights loss, which some reviewers (Reviewer 4HGt) felt was trivial and computationally expensive, suggesting that an uncertainty predictor could provide a more efficient solution and be used for an ablation study (Reviewers 4HGt and gccp). Concerns were also raised about the trade-off between controllability/alignment and image quality. Additionally, the fairness of the comparisons was questioned, along with the increased computational costs. Some aspects of the evaluation metrics and methodology configuration, as well as their impact, were also unclear.

Discussion: 

During the author-reviewer discussion phase, the authors provided detailed clarifications and additional experimental results. Reviewers QrcQ and 4HGt acknowledged that the rebuttal effectively addressed their concerns. Although Reviewers gccp and xjaQ did not respond to the rebuttal, their initial ratings were borderline accept. By the end of the discussion phase, all reviewers rated the paper positively.

AC Recommendation: 

AC reads the authors' rebuttal and found that it effectively addressed the concerns raised by Reviewers gccp and xjaQ. There is no basis to overturn the reviews, and AC agrees with the recommendation. In the camera-ready version, the authors should include the additional experiments and explanations provided in the discussion.

Decision: Accept (Poster)