Enhancing Image Layout Control with Loss-Guided Diffusion Models

Generative models learn to estimate a data distribution with the goal of generating samples from this distribution. Recently , a new family of generative models, known as diffusion models Sohl-Dickstein et al. (2015); Ho et al. (2020); Song and Ermon (2019), have achieved superior results on image synthesis compared to the previous state of the art Dhariwal and Nichol (2021), with samples exhibiting incredible diversity and image quality. While early diffusion models were formulated as Markov chains and relied on a large number of transitions to generate samples, Song et al. Song et al. (2021) showed that such a sampling procedure can be viewed as a discretization of a certain stochastic differential equation (SDE). In particular, Song et al. Song et al. (2021) showed that there exists a family of SDEs, whose solutions are sampling trajectories from the diffusion model. One such SDE, called the probability flow ODE, is completely deterministic and contains no noise. This enables the use of various ordinary differential equation (ODE) solvers for efficient sampling Karras et al. (2022); Lu et al. (2022).

Diffusion models can use a wide range of techniques for controllable generation. SDEdit Meng et al. (2022) allows the user to specify a layout by using paint strokes, which are noised to time $t<T$ to provide an initialization for solving the reverse-time SDE. The realism of the final image is, however, sensitive to the initial noise level $\sigma_{t}$, and guiding the generation of new images requires all high-level features to be specified in the stroke image. A more user-friendly method by Voynov et al. Voynov et al. (2023) requires only a simple sketch to guide the denoising process. The user-provided sketch is compared with the edges extracted by a latent edge predictor in order to compute a loss, which is used to iteratively refine the latent. In this case, the latent edge predictor must be trained, and sketches of more complicated scenes may be tedious. Zhang et al. Zhang et al. (2023) propose a more general method in which a separate encoder network takes as input a conditioning control image, such as a sketch, depth map, or scribble, to guide the generation process. Unfortunately, this requires finetuning a large pretrained encoder. When the inputs are specified as bounding boxes, smaller trainable modules can be used between the layers of the denoising UNet to encode layout information Zheng et al. (2023). Cheng et al. Cheng et al. (2023) also use an additional module which takes in bounding box inputs, injecting it directly after the self-attention layers. Once again, however, neither method can be adapted immediately, as the additional parameters necessitate further training.

A number of other methods have been proposed that are training-free. Bansal et al. Bansal et al. (2024) define a generic loss on the noiseless latent $\hat{\mathbf{z}}_{0}$ predicted from $\mathbf{z}_{t}$ by Tweedie’s formula, and subsequently perform loss guidance on the latent $\mathbf{z}_{t}$ at each time step. One downside of this method is that, while the loss expects clean images, the predicted latent $\hat{\mathbf{z}}_{0}$ is an approximation to only an average of possible generated images, and so can be blurry for large times $t$. An alternative approach called MultiDiffusion was proposed by Bar-Tal et al. Bar-Tal et al. (2023), who used separate score functions on various regions in a latent diffusion model, with an optimization step at each iteration designed to fuse the separate diffusion paths. While this method can be effective in many cases, it can nonetheless exhibit patchwork artifacts, where the final image appears to be composed of several images rather than depicting a single scene.

Alternatively, several works have explored the use of cross-attention to achieve training-free layout control. Hertz et al. Hertz et al. (2023) demonstrate how attention maps can be injected from one diffusion process to another, and reweighted to control the influence of specific tokens. Subsequent work by Balaji et al. Balaji et al. (2023) builds upon this idea by directly manipulating the values in the attention map to obtain the desired layout, although it is difficult to precisely localize objects appearing in an image with this method alone. Singh et al. Singh et al. (2023) also use this technique to improve semantic control in stroke-guided image synthesis, which they combine with loss guidance based on the stroke image to improve the realism of generated images. Instead of using strokes, Chen et al. Chen et al. (2024) show that controlling layout with bounding boxes is possible by using loss guidance, where the loss is defined on the attention maps, although this method requires searching for a suitable noise initialization. Concurrent work by Xie et al. Xie et al. (2023) and Couairon et al. Couairon et al. (2023) also use attention-based loss guidance; the former adds spatial constraints to control the scale of the generated content, while the latter uses segmentation maps instead of bounding boxes. Epstein et al. Epstein et al. (2023) show that it is even possible to control properties of objects in an image, such as their shape, size, and appearance, through their attention maps, and subsequently manipulate these properties through loss guidance.

These works demonstrate the utility of injection, loss guidance, and the general role of cross-attention in layout control. We take a joint approach where we use cross-attention injection to assist loss guidance in producing the desired layout from simple bounding box inputs. In analyzing the role of each technique in layout control, we offer justification for their complementary use. The result is a powerful and intuitive method for layout control which maintains the quality of the generated images.

To perform conditional image synthesis with text, Stable Diffusion leverages a cross-attention mechanism Vaswani et al. (2017). Cross-attention enables the modelling of complex dependencies between two sequences $\mathbf{X}^{T}=(\mathbf{x}_{1},\mathbf{x}_{2},\ldots,\mathbf{x}_{n})$ and $\mathbf{Y}^{T}=(\mathbf{y}_{1},\mathbf{y}_{2},\ldots,\mathbf{y}_{k})$, whose elements are projected to query, key and value vectors using projection matrices

where $d_{k}$ denotes the dimension of the query and key vectors, and $d_{v}$ denotes the dimension of the value vectors. Subsequently, the attention weights are computed as

The new representation for the sequence $\mathbf{X}$ is

In diffusion models, the sequence $\mathbf{X}$ represents the image, where each $\mathbf{x}_{i}$ represents a pixel, and $\mathbf{Y}$ is a sequence of token embeddings. The attention weights $A$, also called the attention or cross-attention map, follow the same spatial arrangement as the image, and a unique map $A_{j}$ is produced for each token $\mathbf{y}_{j}$ in $\mathbf{Y}$. Each entry $A_{ij}$ describes how strongly related a spatial location $\mathbf{x}_{i}$ is to the token $\mathbf{y}_{j}$. We leverage this feature of cross-attention to guide the image generation process.

The forward and reverse processes can be modelled by solutions of stochastic differential equations (SDEs) Song et al. (2021). While determining the coefficients of the forward process SDE is straightforward, the reverse process corresponds to a solution of the reverse-time SDE, which requires learning the score of the intractable marginal distribution $q(\mathbf{x}_{t})$. Instead, Song et al. Song et al. (2021) use the score-matching objective

for some positive function $\lambda\colon[0,T]\to\mathbbm{R}$.

While Eq. (1) does not directly enforce learning the score of $q(\mathbf{x}_{t})$, it is nonetheless minimized when $\mathbf{s}_{\theta}(\mathbf{x}_{t},t)=\nabla_{\mathbf{x}_{t}}\log q(\mathbf{x_{t}})$ Vincent (2011). Conditioning on $\mathbf{x}_{0}$ provides a tractable way to obtain a neural network $\mathbf{s}_{\theta}(\mathbf{x}_{t},t)$ which, given enough parameters, matches $\nabla_{\mathbf{x}_{t}}\log q(\mathbf{x}_{t})$ almost everywhere. Because the forward process $q(\mathbf{x}_{t}|\mathbf{x}_{0})$ is available in closed form, it can be shown that the neural network which minimizes this loss is

where $\sigma_{t}$ is the standard deviation of the forward process at time $t$, and $\bm{\epsilon}_{\theta}(\mathbf{x}_{t},t)$ predicts the scaled noise $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ in $\mathbf{x}_{t}$.

When the predicted score function is available, any one of a family of reverse-time SDEs, all with the same marginal distributions $p_{\theta}(\mathbf{x}_{t})\approx q(\mathbf{x}_{t})$, can be solved to sample from $p_{\theta}(\mathbf{x}_{0})\approx q(\mathbf{x}_{0})$. One of these SDEs is noise-free, and is known as the probability flow ODE. A highly efficient way to sample from $p_{\theta}(\mathbf{x}_{0})$ is to solve the probability flow ODE using a small number of large timesteps Song et al. (2021). An important benefit of sampling using an ODE is that the sampling process is deterministic, in the sense that it associates each noisy image $\mathbf{x}_{T}$ with a unique noise-free sample $\mathbf{x}_{0}$ Song et al. (2021).

In practice, the denoising process usually operates on a latent variable $\mathbf{z}_{t}$, where an autoencoder is used to project to and from image space $\mathbf{x}_{t}$. In order to generate images which correspond to a user-supplied text prompt, the noise predictor $\bm{\epsilon}_{\theta}(\mathbf{z}_{t},t,\mathbf{y})$ is conditioned on text prompts $\mathbf{y}$, and samples are computed using this noise predictor by classifier-free guidance (CFG). In this paper, we denote the CFG noise prediction by $\bm{\epsilon}_{\theta}(\mathbf{z}_{t},t,\mathbf{y})$, or sometimes just $\bm{\epsilon}_{\theta}(\mathbf{z}_{t},t)$, when the dependence on $\mathbf{y}$ is clear. We provide more details about diffusion models in Appendix A.

We use BoxDiff Xie et al. (2023), the method of Chen et al. Chen et al. (2024), and MultiDiffusion Bar-Tal et al. (2023) as our primary points of comparison. In Xie et al. (2023) and Chen et al. (2024), the authors apply spatial constraints on the attention maps of a latent diffusion model to derive a loss, and directly update the latent at time step $t$ by replacing it with

where $\mathcal{L}_{\mathbf{y}}$ is a loss function which depends on the neural network, the prompt $\mathbf{y}$, and a set of bounding boxes. The parameter $\alpha_{t}$ controls the strength of the loss guidance in each iteration. In BoxDiff, $\alpha_{t}$ decays linearly from $t=T$ to $t=0$, while in Chen et al., $\alpha_{t}=\eta\cdot\sigma_{t}^{2}$ for some scale factor $\eta>0$. In MultiDiffusion, letting $\Phi_{\theta}(\mathbf{z}_{t},t,\mathbf{y})$ denote the function which computes $\mathbf{z}_{t-1}$ from $\mathbf{z}_{t}$ using the predicted scaled noise ${\bm{\epsilon}}_{\theta}(\mathbf{z}_{t},t,\mathbf{y})$, a new latent $\mathbf{z}^{\prime}_{t-1}$ is produced by solving the optimization problem

where each $\mathbf{m}_{i}\in\{0,1\}^{n}$ is a mask corresponding to the bounding box of the prompt $\mathbf{y}_{i}$.

Hertz et al. Hertz et al. (2023) observed that, by extracting the attention maps from a latent diffusion process and applying them to another one with a modified token sequence, it is possible to transfer the composition of an image. This technique is very effective, but requires an original set of attention maps which produce the desired layout, and it is not feasible to generate images until this layout is obtained.

Instead, we rely on the observation that the attention maps early in the diffusion process are strong indicators of the generated image’s composition. For early timesteps, both the latents and the attention maps are relatively diffuse, and don’t suggest any fine details about objects in the image. Motivated by this, we manipulate the attention maps by artificially enhancing the signal in certain regions. Given a list of target tokens $S\subset\{1,2,\ldots,k\}$, we define $\mathbf{m}=\{\mathbf{m}_{1},\mathbf{m}_{2},\ldots,\mathbf{m}_{k}\}\in\mathbbm{R}^{n\times k}$ by setting each mask $\mathbf{m}_{j}$, $j\in S$, equal to $1$ over the region that the text token $\mathbf{y}_{j}$ should correspond to, and zero otherwise, and perform injection by replacing the cross-attention map $A_{t}$ at time $t$ with

where $\nu_{t}>0$. We follow Balaji et al. Balaji et al. (2023) and use the scaling

Scaling by $\sigma_{t}$ ensures the injection strength is appropriate for a given timestep, and $\nu^{\prime}$ is a constant which controls the overall strength of injection.

This way, we directly bias the model’s predicted score $\mathbf{s}_{\theta}(\mathbf{x}_{t},t)\approx\nabla_{\mathbf{z}_{t}}\log{q_{t}(\mathbf{z}_{t}|\mathbf{y})}$ so that each latent $\mathbf{z}_{t-1}$ more closely corresponds to the desired layout. We denote the corresponding modified noise prediction by ${\bm{\epsilon}}_{\theta}(\mathbf{z}_{t},t,A_{t}\xrightarrow{\mbox{\raisebox{-1.72218pt}[0.0pt][0.0pt]{\tiny$\nu^{\prime}$,$\mathbf{m}$}}}A_{t}^{\prime})$.

Conditional latent diffusion models predict the time-dependent conditional score $\nabla_{\mathbf{z}_{t}}\log{q(\mathbf{z}_{t}|\mathbf{y})}$, so that the resulting latent at the end of the denoising process is sampled from $q(\mathbf{z}_{0}|\mathbf{y})$. We can modify the conditional score at time $t$ by introducing a loss term $\ell_{\mathbf{y}}(\mathbf{z}_{t})$,

which corresponds to the marginal distribution

The scaling constant $\eta$ controls the relative strength of loss guidance. By using annealed Langevin dynamics Song and Ermon (2019), it is possible to use the predicted score function together with the loss term to sample from an approximation to $\hat{q}(\mathbf{z}_{0}|\mathbf{y})$. While this provides a clear interpretation for the effect of loss guidance, the cost of annealed Langevin dynamics can be fairly large. Instead, we solve the probability flow ODE using the score function

This no longer corresponds to sampling from an approximation to $\hat{q}(\mathbf{z}_{0}|\mathbf{y})$ (see Song et al. (2023)), however, so long as the latents $\mathbf{z}_{t}$ are not too out-of-distribution with respect to the marginals $q(\mathbf{z}_{t}|\mathbf{y})$, then this process influences the trajectory to favor samples from $p_{\theta}(\mathbf{z}_{0}|\mathbf{y})$ for which the loss term is small.

For layout control, given a list of target tokens $S\subset\{1,2,\ldots,k\}$, we choose the simple loss function

where $\mathbf{m}_{j}$ is a mask whose value is $1$ over the region where token $\mathbf{y}_{j}$ should appear, and $0$ otherwise, and $\mathbf{\bar{m}}_{j}=1-\mathbf{m}_{j}$. Intuitively, this simple loss encourages the sampling of latents whose attention maps for each target token take their largest values within the masked regions.

When high levels of loss-guidance are used, the specific choice of the loss function heavily influences the behavior of the denoising process. In BoxDiff, Xie et al. Xie et al. (2023) use sums over the $P\ll n$ most attended-to pixels in the masked attention maps, while in Chen et al. Chen et al. (2024) the authors use a normalized loss which depends on the ratio $\text{sum}(\mathbf{m}_{j}\odot(A_{t})_{j})/\text{sum}((A_{t})_{j})$. These choices seem to be essential for enabling those methods to maintain good image quality at very high levels of loss guidance. Since we will use much lower levels of loss guidance, we find that our method is significantly less sensitive to the choice of loss function, with the result that our simple loss function works well to contain the attentions within the regions defined by $\mathbf{m}$.

In practice, diffusion models are trained to predict the scaled noise $\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ in $\mathbf{z}_{t}$. Revisiting Eq. (2), we observe that we can define the corresponding modified noise prediction by scaling the loss-guidance term appropriately Dhariwal and Nichol (2021):

Since the loss function $\ell_{\mathbf{y}}(\mathbf{z}_{t})$ is a scalar function, computing its gradient with respect to $\mathbf{z}_{t}$ by reverse accumulation requires only a single sweep through the computational graph. As a result, the runtime and memory requirements of sampling using loss guidance are approximately twice those of sampling without loss guidance. This can be seen in, for example, Table 1 of Chen et al. Chen et al. (2024).

One weakness of attention injection is that, by directly modifying the attention maps, we disrupt the agreement between the predicted score function $\mathbf{s}_{\theta}(\mathbf{z}_{t},t)$ and the true conditional score $\nabla_{\mathbf{z}_{t}}\log q(\mathbf{z}_{t}|\mathbf{y})$. The discrepancy between the actual and predicted score functions becomes especially pronounced at smaller noise levels, where fine details in the cross attention maps are destroyed by the injection process. This results in a cartoon-like appearance in the final sampled images, as seen in Figure 17 of Balaji et al. (2023) and in the bottom row of Figure 1. The sensitivity of the cross-attention maps at low noise levels makes it difficult to use injection to precisely control the final layout, without negatively influencing image quality. One big advantage of attention injection, however, is that at high noise levels it is possible to strongly influence the cross-attention maps while still producing latents that are in-distribution. Interestingly, even the cartoon-like images resulting from excessive attention injection at low noise levels are not actually out-of-distribution, but are only in an unwanted style.

A weakness of loss-guidance is that the ad-hoc choice of the loss function may not compete well with the predicted score $\mathbf{s}_{\theta}(\mathbf{z}_{t},t)$. In each denoising step, the latent $\mathbf{z}_{t-1}$ must fall in a high probability region of of $q(\mathbf{z}_{t-1}|\mathbf{z}_{t})$, otherwise the predicted latent $\mathbf{z}_{t-1}$ will be out-of-distribution, which results in degraded image quality in the final sampled latent $\mathbf{z}_{0}$. This means that the guidance term’s influence in Eq. (6) should be small enough to avoid moving the latents into low probability regions. On the other hand, small loss guidance strengths may exert too little influence on the sampling trajectory. In this case, the model produces in-distribution samples, but ones which do not fully agree with the desired layout. This tradeoff is illustrated in the first row of Figure 1, which shows images with unnaturally high contrast and saturation appearing at high loss-guidance strengths . One advantage of loss-guidance is that, even at low strengths, it is able to exert some influence over sampling trajectories without biasing the style of the images.

We take the point of view that injection is better suited as a course control over the predicted latents, and cannot fully replace loss guidance. Instead, we propose using injection together with loss guidance in a complimentary fashion, in such a way that they compensate for each other’s weaknesses. We call this approach to layout control injection loss guidance (iLGD). Instead of delegating the layout generation task entirely to loss guidance, we rely on injection to first bias the latent, as illustrated in the second row of Figure 2. The first row of this figure shows that, when using Stable Diffusion alone, the ball appears in random locations near the center of the image, while the second row shows that injection encourages it to appear more frequently inside of the bounding box, towards the top right. When we also perform loss guidance, the third row of the figure shows that the averaged attention map is significantly more concentrated inside of the bounding box. Since the original score function is modified additively, the amount of loss guidance can be made sufficiently small so that the latents stay in-distribution, while still being large enough to influence the sampling trajectory. This is true even at small noise levels, when fine details are present in the images.

We also observe that, when using both injection and loss guidance, the details of the objects better reflect the context of the scene due to higher levels of attention on those objects. In the third row, many of the balls appear in the style of a soccer ball, which is likely the most common type of ball to appear together with grass.

We outline our algorithm in Algorithm 1 and depict it visually in Figure 3. We perform attention injection from timestep $T$ to $t_{\text{inject}}$ in order to obtain the modified predicted noise $\bm{\epsilon}^{\prime}_{\theta}(\mathbf{z}_{t},t)$. In each timestep, we further refine this latent by performing loss-guidance simultaneously, from timestep $T$ to $t_{\text{loss}}$. In practice, we find it useful to perform loss-guidance for several more steps after we stop injection, $t_{\text{loss}}>t_{\text{inject}}$.

We present a comparison between our proposed method (iLGD), BoxDiff Xie et al. (2023), Chen et al. Chen et al. (2024), MultiDiffusion Bar-Tal et al. (2023), and Stable Diffusion Rombach et al. (2022) in Figure 4. Qualitatively, we observe that BoxDiff and Chen et al. both produce images with higher contrast, and BoxDiff produces especially saturated images. For instance, for the prompt “a donut and a carrot,” both objects in BoxDiff’s image appear unnaturally bright or dark in certain regions. Chen et al. produces high contrast images for the prompts “a dog on a bench” and “a banana and broccoli,” though with visibly less saturation than BoxDiff. High contrast and saturation are even visible in some images when just Stable Diffusion is used, e.g., in the prompts “a balloon and a cake and a frame…” and “a suitcase and a handbag” in Figure 5.

These observations are reflected quantitatively in Table 1, which reports the average contrast and saturation across all generated images. The same high-contrast phenomenon observed in BoxDiff and Chen et al. is seen to occur in Stable Diffusion when increasing the classifier-free guidance (CFG) scale. Furthermore, high staturation can also be seen in both BoxDiff and Stable Diffusion at high CFG scales. Both loss guidance and CFG appear to move the predicted latent into similar low-probability regions if the strength is too large, which may nonetheless be necessary to obtain either the appropriate layout or the appropriate agreement with the semantics of the text prompt. Neither attention injection nor MultiDiffusion appear to have this biasing effect, and the contrast and saturation for iLGD are similar to those of attention injection.

We note that both attention injection and iLGD produce images of lower contrast than even Stable Diffusion without classifier-free guidance, although visual inspection suggests that this does not manifest as any obvious abnormality in the images. We hypothesize that injection favours scenes that naturally contain lower contrast. When the objects appearing in the images receive high levels of attention as a result of injection, they tend to be clear, in-focus, and easily discernible, and do not have either very bright reflections or very dark shadows.

Biasing the latents using injection also clearly helps to preserve image quality and achieve stronger layout control, particularly for layouts which are difficult for the model to generate. For the prompt “a red book and a clock,” BoxDiff struggles to move the clock into the correct position, instead generating it as a lower quality, shadow-like figure, while Chen et al. omits the clock entirely, and MultiDiffusion generates a chaotic scene with many artifacts. Using iLGD, the clock and book appear in the correct positions, and the clock maintains its fine details, e.g., the numbering on the face, as well as its texture. For the prompt “a bowl and a spoon,” the image produced by iLGD is more faithful to the bounding boxes than image produced by BoxDiff, in which it is also difficult to distinguish the spoon from its shadow as dark colours are exaggerated, as well as the image of Chen et al., which struggles to reposition the objects. The layout produced by MultiDiffusion is more accurate, but the image itself is not sensible. In general, it appears that BoxDiff performs very well, but produces images of unnaturally high contrast and saturation. The method of Chen et al. produces images of reasonable quality, but struggles to move the objects to the desired bounding boxes. MultiDiffusion appears to consistently make objects appear in the bounding boxes, but often produces chaotic scenes or scenes with cut-and-paste artifacts, in which the foreground images seem glued to an incompatible background (see the prompt “a banana and broccoli” for one such example).

In Table 2, we report various metrics related to image quality (CLIP-IQA), bounding box precision (YOLOv4) and text-to-image similarity (T2I-Sim). While both BoxDiff and iLGD achieve similar T2I-Sim and [email protected] scores, the latter achieves better performance on all CLIP-IQA metrics. While the method of Chen et al. achieves T2I-Sim and CLIP-IQA scores that are almost as good as those of iLGD, it achieves a much lower [email protected] score. MultiDiffusion has the highest [email protected] score of all methods considered, but it also has the lowest T2I-Sim and CLIP-IQA metrics (all except for Clear), showing that MultiDiffusion sacrifices an unacceptable amount of image quality in exchange for layout control. The ablation study (SD + I) using just injection shows good image quality, but substantially worse layout control than iLGD, while the ablation study (SD + L) using just loss guidance shows slightly worse image quality and much worse layout control than iLGD. These results indicate that iLGD is able to achieve superior image quality without sacrificing layout control, qualifying it as a meaningful improvement over BoxDiff, Chen et al., and MultiDiffusion.

We provide additional comparisons in Appendix C.

A visual comparison between iLGD and injection alone is presented in Figure 5. We observe that, while injection typically does generate an object in each bounding box, the object itself may be incorrect. To illustrate, for the prompt “a balloon, a cake, and a frame, …,” using injection leads to a frame appearing where the cake should be; for “a donut and a carrot,” it generates a hand where the donut should be; for “a suitcase and a handbag,” a second suitcase is generated instead of a handbag; and for “a cat sitting outside,” something akin to a tree stump appears instead of a cat. In the example for “a castle in the middle of a marsh,” the castle does not fill up the bounding box, and for “a cat with a tie” the resulting image has the correct layout, but begins to appear like a cartoon cutout. When loss guidance is added using iLGD, all of these images appear correctly. These observations are reflected in Table 2, where injection achieves a noticeably lower [email protected] score compared to iLGD.

We note that the results in Figure 5 provide some additional evidence that injection is able to successfully bias the image according to the desired layout. For instance, revisiting the example of the prompt “a cat sitting outside,” injection causes a tree stump to appear in the region where the attention was enhanced. When loss guidance is also applied in each step, this stump is instead denoised into a cat. In this case, a small amount of loss guidance suffices to generate the appropriate layout, as it is augmented by the biasing effect of injection.