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

Authors: Guiyu Zhang, Huan-ang Gao, Zijian Jiang, Hao Zhao, Zhedong Zheng
Primary area: generative models
Keywords: Conditional generation, semantic alignment, uncertainty learning

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730463340987

thread_type:
official_review

<thread_object>
creator: reviewer_xjaq

<content>
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.

strengths:
I appreciate the motivation behind this paper, and I believe the method presented is quite sound.

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.

questions:
Please refer to Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730561545718

thread_type:
official_review

<thread_object>
creator: reviewer_4hgt

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730861763227

thread_type:
official_review

<thread_object>
creator: reviewer_gccp

<content>
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.

strengths:
The method is straightforward and effective, and the paper is well-organized and easy to understand.

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.

questions:
see Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1731206122634

thread_type:
official_review

<thread_object>
creator: reviewer_qrcq

<content>
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.

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)

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.

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)
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.