Readout Guidance for Diffusion Control — Poster

论文上下文

Paper title: Readout Guidance: Learning Control from Diffusion Features Abstract: A conference poster detailing 'Readout Guidance' (RG), a method using small readout networks to extract properties from frozen diffusion features for controlling image generation, including architecture, qualitative results, and low-data training comparisons. Paper body (method & results): Readout Guidance: Learning Control from Diffusion Features Grace Luo1,2 Trevor Darrell2 Oliver Wang1 Dan B Goldman1 Aleksander Holynski1,2 1Google Research 2UC Berkeley Readout Readout Readout Readout (a) Readout Guidance With Readout Guidance Control Input Pose Appearance Correspondence (b) Applications Text Guidance Only ▽dMSE ▽dHinge ▽dCosine Sim Figure 1. Given a frozen pre-trained text-to-image diffusion model [50], we learn parameter-efficient readout heads to interpret relevant signals, or readouts, from the intermediate network features. These readouts can be single-image concepts such as pose and depth, or relative concepts between two images, such as appearance similarity and correspondence. We use the readouts for sampling-time guidance to enable controlled image generation. Abstract We present Readout Guidance, a method for controlling text-to-image diffusion models with learned signals. Readout Guidance uses readout heads, lightweight networks trained to extract signals from the features of a pre-trained, frozen diffusion model at every timestep. These readouts can en- code single-image properties, such as pose, depth, and edges; or higher-order properties that relate multiple images, such as correspondence and appearance similarity. Furthermore, by comparing the readout estimates to a user-defined tar- get, and back-propagating the gradient through the readout head, these estimates can be used to guide the sampling pro- cess. Compared to prior methods for conditional generation, Readout Guidance requires significantly fewer added pa- rameters and training samples, and offers a convenient and simple recipe for reproducing different forms of conditional control under a single framework, with a single architecture and sampling procedure. We showcase these benefits in the applications of drag-based manipulation, identity-consistent generation, and spatially aligned control. Project page: https://readout-guidance.github.io. 1. Introduction Diffusion models have shown enormous potential in accu- rately modeling the space of natural images. However, one remaining open challenge that is critical to many applications is enabling arbitrary user control over their outputs. Existing solutions for enabling custom user control typically involve substantial model training on large annotated datasets, a pro- cess that is cumbersome and often infeasible for the average user. In this work, we provide an alternative solution for enabling user control that combines two ideas; first, that dif- fusion models contain rich internal representations that are useful for extracting relevant image properties, and second, that these extracted image properties can be used to guide the generation process towards desired user constraints. We call the combined approach Readout Guidance, because it makes use of small auxiliary readout heads that can be easily trained on top of a frozen diffusion model. Taking as input the set of intermediate diffusion features, these readout heads can be trained to extract arbitrary proper- ties about the image being generated (Figure 1, left). These can include image-space properties, such as human pose, arXiv:2312.02150v2 [cs.CV] 2 Apr 2024 depth maps, and edges; but can also be higher-order prop- erties that relate two or more images, such as appearance similarity, correspondence, or shared identity. These readout heads consist of very few parameters, meaning that they can be trained on a single consumer GPU in a matter of hours— and since they bootstrap the already-rich diffusion features, they only need as few as 100 training examples. Beyond their utility in efficiently extracting properties about the gen- erated images, the readout heads offer a useful mechanism for guiding the sampling process (Figure 1, right). At each step of the sampling process, properties can be extracted from the readout heads and compared to user-defined targets. In a similar fashion to classifier guidance [13], this compar- ison can be used as a guidance signal that encourages the generated image to match the target constraints. We show that Readout Guidance can implement a num- ber of popular forms of user control—all within the same simple guidance framework. In particular, we demonstrate state-of-the-art performance on the task of drag-based image manipulation, an application that has previously required be- spoke architectural modifications and additional per-example fine-tuning. We also showcase our method on the task of identity-consistent image generation, in which outputs can be guided to contain the same person as a reference image. Finally, our method can also be used for spatially-aligned controls, such as depth-guided or pose-guided generation, as popularized by ControlNet [70] and T2IAdapter [40]. Notably, when compared to ControlNet [70], our method requires significantly less training data (as few as 100 super- vised pairs vs. 200k), much less training time (a few hours vs. more than a week), and fewer added parameters (49MB vs. 1.4GB) [69]. 2. Related Work Conditional Diffusion Models. It has become increas- ingly popular to fine-tune text-to-image diffusion models to condition generation on signals beyond text, including camera pose [35], a reference identity [51], a reference im- age [6, 32, 38], a depth map [3], and more [23]. Due to the high cost of training diffusion models, many methods propose to keep the base model frozen and instead train an additional network that takes in the control signal and modu- lates the intermediate diffusion features accordingly. These types of models, including ControlNets [70], Adapters [40], and LoRAs [22], require less compute to train, fewer training samples, and can be combined in various ways because they build on the same frozen base model. Our method is orthog- onal and complementary to this kind of adapter tuning, since our method aims to guide the sampling process based on the diffusion model’s features, rather than modulate it, and can therefore be applied to any base model with any additional adapter. Our approach is more similar to classifier guidance, which is shown in Dhariwal & Nichol [13] to not only en- able conditional generation from unconditional models, but also reinforce the capabilities of a conditional model. In our experiments, we similarly demonstrate that our method can be applied to these adapter-based models [40, 70] to further improve their control capabilities. Sampling-Time Guidance. Because diffusion models syn- thesize images through an iterative sampling process rather than a single forward pass, one can guide the sampling pro- cess in a particular direction without modifying the base model. Such guidance was first achieved with gradients from an ImageNet [12] classifier trained on noisy images from the diffusion forward process [13]. Follow-up works explored alternative guidance approaches including the difference of score estimates between a conditional and unconditional model [21], off-the-shelf models that operate on the pre- dicted clean image [4, 60], and hand-designed training-free functions that operate on diffusion features [10, 15, 33, 43]. In our work, we focus on learning this guidance function from diffusion features. Depth-Aware Guidance [29] first explored a similar idea for depth refinement, Sketch-Guided Diffusion [59] for edge guidance, and Mid-U Guidance [61] for aesthetic guidance. We expand this line of work into a more general form that can be used for non spatially-aligned applications such as drag-based manipulation. Diffusion Model Representations. There exist works that demonstrate the rich and expressive nature of diffusion fea- tures beyond image synthesis, and how they can be readily applied to tasks like segmentation [5, 63, 66, 71], depth esti- mation [63, 71], human pose prediction [63], and semantic correspondence [20, 36, 55, 68]. These methods typically train additional decoders to extract an application-specific representation from the features of a particular, single diffu- sion timestep. To use these signals for guidance, we build on this work to instead extract an evolving prediction of these image properties at each step of the entire diffusion process. Drag-based Manipulation. The concurrent works DragDif- fusion [53] and DragonDiffusion [39] explore drag-based editing [42] with diffusion models. In this task, the user provides an image and a sparse set of point correspondences and the method produces an edited image that respects the specified appearance and deformation. These works en- force their appearance constraint using custom sampling- time processes which include single-image finetuning [53] or shared attention features [39] (following work in video stylization [7, 18, 28, 62]). The deformation constraints are enforced by encouraging high similarity between the corre- sponding points in the feature maps extracted from the diffu- sion model, as proposed in diffusion-based semantic corre- spondence methods [36, 55]. In contrast, our method learns the appearance similarity and correspondence constraints automatically from video data and applies both through the same sampling-time guidance procedure without any hand- designed operators or per-example fine-tuning. Figure 2. Drag Based Manipulation (Generated Images): We show generated images with a single user correspondence constraint (with overlay) followed by the Readout Guidance generated result. Please see the Supplemental for the associated text prompts. 3. Readout Guidance In this section, we describe the process of Readout Guidance: first describing the use of readout heads in the sampling pro- cess, then providing a general recipe for training a readout head on a custom dataset. In Sec. 4, we describe a num- ber of example readout heads, showing that they can model spatially aligned properties like depth and human pose, as well as pairwise relative properties, such as correspondence and appearance similarity. We later showcase these readout heads in a number of conditional control applications, such as drag-based manipulation (Figure 2, 5), appearance preser- vation (Figure 6, 7), and spatially aligned control (Figure 4). 3.1. Background Classifier guidance is a sampling procedure that enables conditional synthesis in an unconditional diffusion model. For the deterministic DDIM sampler [54], Dhariwal & Nichol [13] derive the update rule: ˆϵt ←ϵθ(xt) − √ 1 −¯αt∇xtlog pϕ(y|xt) where ϵθ(·) is the diffusion model, pϕ(y|·) is the classifier, and √1 −¯αt is a timestep-conditional scaling factor. One can think of this process as the classifier “nudging” the generation at each step such that it remains on the natural image manifold but moves towards to the classifier’s criteria. In our method, we bootstrap pre-trained diffusion features from the frozen U-Net decoder to train a small readout head on top of the diffusion model. This allows us to design simple heads that are the same across entire classes of ta