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

TL;DR

This work tackles fidelity and semantic alignment in conditional image generation by addressing inaccuracies in reward-model feedback. It introduces Ctrl-U, which performs two forward passes with different diffusion timesteps to estimate reward uncertainty and then applies uncertainty-aware regularization to adaptively weight reward signals during diffusion training. The method yields substantial gains in controllability and image quality across segmentation, edge, and depth controls on ADE20K, COCO-Stuff, and MultiGen-20M, with robust ablations validating the approach. By explicitly handling uncertainty, Ctrl-U enhances the reliability and scalability of conditional diffusion-based generation with minimal inference overhead.

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. Codes are publicly available at https://grenoble-zhang.github.io/Ctrl-U-Page/.

Figures (10)

• Figure 1: Given a test image and the layout condition, we employ a diffusion model to generate new images by adding noise and then recovering from the noisy input. (a) Ground-truth segmentation results with the category illustration. (b) Here we show the reward changes, i.e., mIoU error, on newly generated images at different timesteps. The horizontal axis represents the current timestep $t$ and the vertical axis shows the error, i.e., 1-mIoU. As shown, even at $t=0$, there are non-zero mIoU errors. As $t$ increases, although the visual layout aligns with the condition, the reward model tends to increase the error, leading to the backpropagation of incorrect gradients.
• 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 KL-divergence prediction discrepancy between $\hat{c}_1, \hat{c}_2$ as the uncertainty indicator $U_1, U_2$ (see Eq. \ref{['eq2']}). Based on $U_1, U_2$, we then calculate the rectified reward loss between the predicted label $\hat{c}_1, \hat{c}_2$ and the ground-truth label $c$, as Eq. \ref{['eq3']}.
• Figure 3: Qualitative comparisons with different conditional controls on unseen test images. We observe that our generated image preserves condition alignment with good visual quality. Some models do not have open-source weights for Hed or Lineart condition, and thus we skip them.
• Figure 4: Ablation study on the uncertainty-aware reward modeling. Here, we show the recovered test image with different denoising timesteps. Since we mitigate the negative impact of noisy rewards, our output maintains consistent semantic conditions with fewer blurred areas.
• Figure 5: Text prompts generated by MiniGPT-4