Make It Count: Text-to-Image Generation with an Accurate Number of Objects

## Abstract

Abstract Despite the unprecedented success of text-to-image diffusion models, controlling the number of depicted objects using text is surprisingly hard. This is important for various applications from technical documents, to children’s books to illustrating cooking recipes. Generating object-correct counts is fundamentally challenging because the generative model needs to keep a sense of separate identity for every instance of the object, even if several objects look identical or overlap, and then carry out a global computation implicitly during generation. It is still unknown if such representations exist. To address count-correct generation, we first identify features within the diffusion model that can carry the object identity information. We then use them to separate and count instances of objects during the denoising process and detect over-generation and under-generation. We fix the latter by training a model that predicts both the shape and location of a missing object, based on the layout of existing ones, and show how it can be used to guide denoising with
correct object count. Our approach, CountGen , does not depend on external source to determine object layout, but rather uses the prior from the diffusion model itself, creating prompt-dependent and seed-dependent layouts. Evaluated on two benchmark datasets, we find that CountGen strongly outperforms the count-accuracy of existing baselines.

## 1 Introduction

Text-to-image diffusion models provide an accessible way to control the generation of visual content. A major failure mode is their inability to count, that is, they often fail to generate the correct number of items in response to text prompts. For instance, when asked to generate an image of Goldilocks and the three bears, models may generate only two bears (Figure [1](#S0.F1)). Counting failures are particularly frustrating: The accuracy is surprisingly low, and mistakes are often obvious for people to detect.

To illustrate the difficulty of the problem, consider some naive attempts to work around it. First, one can manually design layouts per count, to determine the spatial organization and the number of objects, then provide it as a conditioning signal to a generative model []. This approach would fail to generate prompt-dependent layouts, which is highly desirable.
One could also try asking large vision-language models to propose layouts (e.g., [, ]), but these approaches do not use the visual prior information that text-to-image models already collected, and, as we show below, their performance is rather poor for the counting task.

Why is it so hard for diffusion models to count while they generate? First, counting objects requires that models capture "objectness" – the high-level coherent representation of something being a separate entity, even if surrounded by other similar entities. Capturing objectness is by itself a hard task in image understanding [, ], and long studied in cognitive psychology []. It is currently not known to what extent diffusion models represent objectness of entities they generate.
A second main challenge is that text-to-image diffusion models struggle with controlling spatial layout just from text. Producing a correct number of objects requires obeying a global and complex spatial relation between entities in an image [, ].

To address the problem of accurate count generation we describe several new contributions, which together form our method CountGen. First, we analyze the representations of the self-attention layers in SDXL [], and identify features that capture objectness and instance identity. We then use these features to develop ways to detect instances of objects during the denoising process, find their spatial layout and count them. Specifically, we localize the features that correspond to objects using the cross-attention maps and cluster them to form object instance segmentation.
Importantly, we do not have to wait for an image to be fully generated, and we can accurately count the number of objects already at an intermediate step of the denoising process.

Given this new capability to count the number of objects being generated during the denoising process, we further develop methods to correct generation when the count does not match the prompt. First, we train a layout-modification network we call ReLayout. It takes a spatial layout of $k$ objects and generates a similar spatial layout with one more instance of an object added in a natural location for the input layout. For example, given a row of five kittens sitting on a branch, it learned to add a sixth kitten to the same row. This model is trained using image-pair samples generated by the diffusion model itself. Finally, we show how to use the new layouts in a new test-time-optimization procedure.

We evaluate CountGen on text prompts from the T2I-CompBench [] which includes prompts with numbers. CountGen greatly improves accuracy, as evident by human evaluation experiments, from 29% accuracy for SDXL to 48% by our method. It also improves over all other baseline methods including large commercial models like the recent DALL-E 3 []. To support future work in this field, we design and release a dataset that can be evaluated automatically. Specifically, we release CoCoCount, a set of prompts based on COCO classes[], which can easily be evaluated using COCO-trained object detectors, like YOLO []. CountGen also significantly improves over all baseline methods on CoCoCount, importantly from 26% accuracy for SDXL to 54% by our method.

In summary, this paper makes the following new contributions (1) We identify novel features that represent objectness and instance identity in SDXL []. (2) We design an inference-time optimization to guide SDXL to generate an accurate number of instances for an object. (3) We describe a learning approach to automatically modify layouts to add a new instance of an object while preserving the structure of the scene. (4) We achieve state-of-the-art results in count-accurate generation.

## 2 Related Work

Generating images with accurate object count.
Numerous papers noted that text-to-image diffusion models often fail to produce images that accurately match text prompts, especially when these prompts specify an exact number of objects [, , , , , , , , , , , , ]. Various efforts were made to improve the accuracy of these models.
Most relevant to our work, Kang et al. [] proposed a classifier-guidance approach to improve object count accuracy. The method “counts” instances at each diffusion step using a pretrained counting network and adjusts the denoising process using gradient guidance. However, it requires using an additional U-Net in *every* denoising step.

An important line of work suggests breaking the generation process into two steps: (1) Text-to-layout - setting a spatial location for every object instance; and (2) Layout-to-image - generating an image with the correct object count using the given layout.
Text-to-layout: Several studies used large language models (LLMs) to propose spatial layouts [, , , ]. LayoutGPT [] injects visual commonsense into the LLM prompt which enables it to generate desirable layouts. Gani et al. [] suggest decomposing complex prompts into smaller prompts before injecting them into the LLM.
Layout-to-image: Providing a predefined layout with the exact number of subjects helps ensure that the generated images reflect the intended count [, ]. Bounded Attention [] addresses this challenge by channeling attention to bounding boxes corresponding to object instances. However, this approach requires users to manually provide the bounding boxes for all the instances of each object. In contrast to these separate-step approaches, CountGen, addresses the two steps of count-accurate generation. It first corrects the layout that emerges during generation so it contains the correct number of instances. It then uses a novel test-time optimization method to generate a count-accurate image.

Controlling text-to-image models through attention-based loss.
To address the issue of object neglect—when objects mentioned in a text prompt fail to appear in the generated image—Chefer et al. [] developed a novel loss function that ensures all objects in the prompt are reflected in the cross-attention maps used during image generation. Rassin et al. [] tackled the challenge of incorrect attribute association by designing a loss function that binds the cross-attention maps of subjects and their attributes more effectively. Inspired by these advancements, CountGen includes a novel cross-attention maps loss function designed to ensure the generation adheres closely to the input layout.

## 3 Our Approach: CountGen

Figure: Figure 2: Architecture outline. Given a prompt that includes a quantity, we begin generating a corresponding image using pretrained SDXL until timestep $t=500$. We then perform Instance Localization, where we combine cross-attention maps corresponding with the object, and self-attention features extracted at timestep $t$ to generate object clusters for each generated object. Then we apply ReLayout, which generates an object layout with the correct number of instances, while preserving the composition of the extracted layout. Finally, we perform Layout Guided generation, which applies an inference time optimization based on the layout through cross-attention loss $L_{\text{cross}}$ and self-attention masking.
Refer to caption: /html/2406.10210/assets/x2.png

Our method, CountGen, aims to enhance text-conditioned image generators to accurately produce the intended number of objects for complex input prompts. Our methodology involves a two-step process: initially, we generate a natural layout that specifies where and how objects should appear in the image ([Section 3.2](#S3.SS2)). That layout is based on a layout that emerges naturally from the text-conditioned generation ([Section 3.1](#S3.SS1)). At the second step (LABEL:{sec:layout_guided_generation}), we use this layout as a blueprint to generate the final image.

### 3.1 Discovering Object-Instance Layout during Early Generation

To count object instances during generation, one must first find an internal representation that captures the separate identity of different object instances. It is not known if this representation exists in diffusion models like SDXL. We now discuss this representation and then show how we can detect the layout of object instances during early generation.

Figure: Figure 3: PCA Visualization. to explore the notion of objectness inside SDXL latent space, we visualize dimension-reduced self-attention feature maps from various layers across the network at timestep $t=500$. We notice that although most layers do not exhibit a clear separation between object instances, layer $l^{up}_{52}$ displays a robust separation indicated by different ones having distinct colors. Visualization across different timesteps is shown in the appendix ([Figure 12](#A2.F12)).

##### An emerging instance-identity representation in SDXL.

We begin by exploring the notion of ‘objectness’ in SDXL. While previous work [, , ] utilized the cross-attention mechanism to localize objects of a given class in generated images, little research has been conducted on whether the model encodes information about object instances and how to distinguish between different instances of an object. We tackle this problem by exploring a variety of features across different layers and timesteps of the diffusion process to determine if and where the model encodes instance-level information. [Figure 4](#S3.F4) illustrates this analysis using PCA visualization of self-attention features from various layers across SDXL at timestep $t=500$, which shows the most robust instance representation (See appendix [Figure 12](#A2.F12)). While most layers do not exhibit separability at the instance level, we notice that layer $l^{up}_{52}$ tends to generate different features for different instances of the same object, with each instance having its distinct color.
Based on this finding, we select the self-attention features from layer $l^{up}_{52}$ at timestep $t=500$ to serve as our instance-level features.

##### Identifying object instances.

Building on the findings of Hertz et al. [], which show that cross-attention maps can pinpoint a token’s position in a generated image, we create a foreground mask for each object described in the prompt. By contrasting these foreground masks, derived from the cross-attention, with the self-attention features, we effectively segregate pixels associated with objects from those belonging to the background. Subsequently, we cluster the object-associated pixels from the self-attention map into distinct masks for each object. This approach allows us to refine our object representations and enhance the accuracy of the generated layouts.

Formally, let $A^{self}_{l,t},A^{cross}_{l,t}$ represent the self-attention and cross-attention maps, respectively, for layer $l$ at timestep $t$ within our diffusion network. We aggregate cross-attention maps associated with the tokens corresponding to the objects specified in the input prompt. We then use these cross-attention maps to extract a foreground mask $M$ based on dynamic thresholding $M=\text{Otsu}(A^{cross}_{l,t})$,
where Otsu applies the Otsu thresholding method [, ] to determine foreground (object) pixels.
We define ${p_{k}}\subseteq{A^{self}_{l,t}}$ as the set of features from the self-attention map that are identified as foreground by mask $M$. We then cluster these patches: $Clusters=DBSCAN({p_{k}},\epsilon)$,
where $DBSCAN(\cdot,\epsilon)$ is the DBSCAN [] clustering algorithm with a dynamic parameter $\epsilon$. Finally, the initial layout $L$ is created by grouping the object clusters: $L=\bigcup_{C\in{Clusters}}C$.

At the end of this process, we obtain a set of masks, one for each object being generated.
This is illustrated in [Figure 2](#S3.F2), left gray box.

### 3.2 ReLayout: Correcting the Number of Objects in the Mask

We now introduce our layout-correction component, *ReLayout*, which preserves the overall scene composition while correcting the number of objects. For example, [Figure 2](#S3.F2) depicts an image generated using the prompt “a photo of six cats”, but only four cats were generated. Our ReLayout generates a new layout with the correct number of instances while keeping the overall composition of the kittens sitting in a row. More examples are shown in [Figure 4](#S3.F4).

The input to the ReLayout is an object-layout described in [Section 3.1](#S3.SS1), from which we initially infer the number of generated instances. Next, our ReLayout component takes one of two corrective actions based on the discrepancy between the generated and expected counts. In cases of over-generation, where more instances were generated than requested, ReLayout deterministically removes the smallest instances to achieve the desired cluster count. We find that this simple strategy produces appealing results.
In cases of under-generation, a more intricate challenge arises:
the ReLayout must insert new instances to the scene in a way that preserves the original scene structure.
This process involves a sophisticated understanding of different object layouts—like the stark contrast between linearly arrayed *bottles* and the clustered arrangement of *elephants*—to seamlessly augment the layout.
In [Section 3.2.1](#S3.SS2.SSS1), we detail our approach for handling under-generation.
In cases where the number of instances is correct, the ReLayout maintains the initial layout.

#### 3.2.1 Handling Under-generation

To address under-generation issues, we train a U-Net model to predict a new layout, represented as a multi-channel mask, from an existing layout. In practice, each forward pass of the U-Net generates a mask with an additional instance. This process is applied in iterations until the mask reflects the correct number of instances. In what follows we provide detailed information on the architecture and training of our U-Net model.

##### Creating a training dataset.

To train our ReLayout U-Net, we need a dataset of layout pairs with $k$ and $k+1$ objects, that maintain the same scene composition.
We begin with the empirical observation that slight variations in the object count specified in the prompt—while keeping the starting noise and the rest of the prompt consistent—typically results in images with similar layouts, as shown in [Figure 5](#S3.F5). This consistency is crucial as it allows us to generate a training dataset of layout pairs where each pair has a similar object composition, differing by only one object, thereby preserving the overall scene structure.

Following this observation, we generate a set of ~10K pairs of images of $I_{k}$ and $I_{k+1}$, where each pair consists of images that differ by only one in the number of objects depicted.
Each pair is generated with random fixed seeds and prompts that fit the same template, such as "a photo of two cats" versus "a photo of three cats”.
To confirm that each image pair accurately represents an $k$ and $k+1$ object scenario, we extract object masks $M_{k}$ and $M_{k+1}$ as described in [Section 3.1](#S3.SS1), and verify the object count in one image is exactly one more than in its paired image.
Overall, the final dataset for training consists of pairs of binary masks $(M_{k},M_{k+1})$, representing the U-Net task of learning to generate a mask with $k+1$ objects from a mask with $k$ objects.

Figure: Figure 5: A training set for a ReLayout. We created pairs of images using SDXL, using the same seed and prompts that only differ by object count. We filtered out images that did not conform to the prompt, using the techniques described in Section [3.1](#S3.SS1). The resulting image pairs preserve the scene and layout except adding one object.
Refer to caption: /html/2406.10210/assets/x5.png

##### Matching objects.

To train the U-Net, we need to establish a correspondence between each object $i$ in $M_{k}$ to its new position in $M_{k+1}$. We aim to find a matching that minimizes the shift in objects positions. We use the Hungarian algorithm [] to find the optimal matching. More details in [Section C.1](#A3.SS1).

##### Training the U-Net module.

We trained the U-Net architecture by adapting it to handle 9 input channels – corresponding to the source tensor $M_{k}\in\{0,1\}^{W\times H\times k}$ with $k$ objects, and output 10 channels – for the target tensor with $k+1$ objects, to support counts up to 10. We optimized the U-Net parameters using two loss functions: (1) A Dice loss [] between a predicted masks $\hat{M}_{k+1}$ and the target masks $M_{k+1}$ of that object; and (2) Mask-to-mask overlap loss, designed to reduce the overlap between output masks of different instances.
Specifically, this was computed as $1-L_{Dice}$ between all pairs of predicted masks $M^{i}_{k+1}$, $M^{j}_{k+1}$.

$$ $\mathcal{L}=\mathcal{L_{\text{DICE}}}+\lambda\mathcal{L_{\text{overlap}}}$ (1) $$

with $\lambda$ being a weighing hyperparameter. Detailed definitions are provided in [Section C.2](#A3.SS2).

##### Inference.

At inference time, as a pre-processing step, we first add padding to input masks. After each iteration, we gradually and consistently increase the padding size around the original masks. This operation is beneficial when we need to add a large number of objects, as it creates a “zoom-out” effect, making space for new objects.

### 3.3 CountGen Image: Layout-based Image Generation

Provided with correct object mask layouts ([Section 3.2](#S3.SS2)), our goal is to guide the image generation process to adhere to the input layout. Given a mask for each object in the desired layout, we apply an inference time optimization to match the layout in the generated image.
To optimize object layouts at inference time, we propose a dual approach:
object layout loss to encourage object creation in the foreground, i.e. pixels within the object masks, and self-attention masking to prevent object generation in the background.

##### Object layout loss.

Consider the optimization of object placement within layouts using a weighted binary cross-entropy loss. Given $c$, the aggregated cross-attention scores, and $m$, a binary mask denoting object presence (foreground), the weighted binary cross-entropy loss is computed pixel-wise and is defined as follows:

$$ $L(c,m)=-\sum_{i}w_{i}\left(m_{i}\log c_{i}+(1-m_{i})\log(1-c_{i})\right),$ $$

where $c_{i}$ is the cross-attention score at pixel $i$, $m_{i}$ is the value of the binary mask at pixel $i$, and $w_{i}$ is the weight assigned to each pixel $i$ where $w_{i}=10$ if $m_{i}=1$, otherwise $w_{i}=1$.

##### Self-attention masking.

The object-layout loss encourages objects to be generated in the foreground, but when applied on itself, generated objects may appear outside the object masks ([Figure 8](#S5.F8)). To address this, we mask the self-attention connections between pixels in the background to pixels in the foreground. By disrupting these links, we stop the information flow from the objects to the rest of the image and prevent the model from forming objects in the background.
Formally, at layer $l$ and timestep $t$, the masked self-attention $S^{*(l)}_{t}$ is defined as:

$$ $S^{*(l)}_{t}[i,j]=\begin{cases}0&\text{if }i\in\mathcal{B}^{(l)}\text{ and }j\in\mathcal{F}^{(l)},\\ S^{(l)}_{t}[i,j]&\text{otherwise}.\end{cases}$ $$

where $i$ and $j$ are pixels indices, $\mathcal{B}^{(l)}$ and $\mathcal{F}^{(l)}$ represents the set of pixels belonging to the background and the foreground respectively, and $S^{(l)}_{t}$ is the self-attention map at layer $l$ and timestep $t$.
We discuss implementation details and computational efficiency in [Appendix A](#A1).

## 4 Experiments

##### Compared methods.

We compare CountGen against six baseline methods: (1) SDXL []; (2) Repeated Object: SDXL, with a modified prompt, where an object is repeated in the prompt the number of times it is required to generate, as in replacing “three cats” with “a cat and a cat and a cat”. This is a naive approach that parallels prompts like “A cat and a dog”.
(3) Reason Out Your Layout: [] uses GPT-3.5 [] to generate layouts then trained an adapter to integrate it to SD-1.4 []; (4) DALL-E 3 []; (5) Random masks + BoundedAttn : generate a layout with the correct amount of clusters placed randomly in the image and apply a layout-guidance generation method on top; (6) Counting Guidance [] : boost generation of SD with a counting network;
Full details on how we used these baselines are given in [Section B.2](#A2.SS2).
We also compared our layout-to-image phase, CountGen-Image, described in section [Section 3.3](#S3.SS3) with Bounded Attention [].

##### Datasets.

We evaluate our method and the baselines using two datasets. (1) T2I-Compbench-Count. A subset of T2I-Compbench [], which is a benchmark for open-world compositional text-to-image generation. This subset specifically includes 218 prompts that specify a single object and its number (between 2 to 10).
(2) CoCoCount (ours). We generate a dataset with automatic evaluation in mind. Specifically, we sample classes from COCO, which are more favorable to accurate and automatic detection by methods, like YOLOv9 []. We design simple prompts around these classes, with a number between 2 and 10. In total, there are 200 prompts with various classes, numbers and scenes (See full details in [Section C.3](#A3.SS3)).

##### Count accuracy evaluation.

We evaluate the results of CountGen and the baselines using human and automatic evaluation method, which is standardized and reproducible. In both settings, we seek to identify if the number of instances generated by the object matches the request in the prompt.

Human evaluation. We quantified the count-accuracy of our method and baselines using human raters. Raters were asked for every image: (1) Is the object in the image?; (2) Are its instances well-formed?; (3) How many instances of the object are in the image? If the answer to question (1) or (2) is “no”, then we do not ask question (3).
We provide details on the platform, rater selection and pay, and screenshots of the task in [Section D.1](#A4.SS1).

Automatic evaluation. For automatic evaluation, we use the YOLOv9 model [] with its default settings, as it represents the current state-of-the-art in the YOLO object detection benchmarks. To extract the number of objects in the image, we simply count the number of detected bounding-boxes corresponding to the target object.

##### Image quality evaluation.

Forcing the diffusion model to obey the count in the text prompt is inevitably expected to reduce the naturalness and visual appeal of generated images, simply because more constraints are added.
This effect has been observed in other studies using test-time optimization[, ].
We evaluate the image quality of CountGen by presenting human raters with two images, by CountGen and SDXL, and asking them to select whether one image is more natural and well-formed than the other or to indicate that both images are equally good.

**Table 1: Generated count accuracy. Values are the percent of generated images that have the correct number of objects, for CoCoCount and T2I-Compbench-Count.**
|  | CoCoCount | T2I-Compbench [] |  |
| --- | --- | --- | --- |
|  | YOLOv9 | Human | Human |
| Model | Accuracy | Accuracy | Accuracy |
| SDXL | 28 | 26 | 29 |
| Repeated Object | 17 | 18 | 14 |
| Reason Out Your Layout | 21 | 26 | 15 |
| DALL-E 3 | 25 | 38 | 36 |
| Random masks + BoundedAttn | 29 | 30 | 35 |
| Counting Guidance | 21 | 22 | 22 |
| CountGen (ours) | 50 | 54 | 48 |

## 5 Results

##### Quantitative results.

[Table 1](#S4.T1) compares CountGen with competing baselines, showing its significant improvement over baselines in both CoCoCount and T2I-compbench-Count. [Figure 8](#S5.F8),
and [Figure 10](#A2.F10) show CountGen outperforms all baselines for all values, except for two and three instances, where DALL-E 3 slightly outperforms. We hypothesize that DALL-E 3 is larger and was trained on higher-quality data than SDXL (our base model). In terms of image quality, out of 200 comparisons, in only 23 cases the majority of the raters preferred SDXL over our model. This indicates there is no significant loss of quality. We also include the confusion matrix figure of CountGen based on human evaluation in [Section D.1](#A4.SS1).

Figure: Figure 6: Qualitative comparisons. We evaluated CountGen against DALLE 3, Reason Out Your Layout, SDXL, Repeated Object SDXL and Counten Layout + Bounded Attention. Our method successfully generates the correct number of objects, while other methods struggle in some or all of the examples. Additional results are shown in the supplemental material.
Refer to caption: /html/2406.10210/assets/x6.png

##### Qualitative results.

Figure [6](#S5.F6) shows examples of prompts and the images generated by various methods. In contrast to other methods, CountGen consistently generates the correct number of object instances.

Figure: Figure 7: Component ablation. We ablate over two components of the layout-guided generation model: the optimization loss and Self-Attention Masking. Disabling the loss causes the generated image to deviate from the required layout. Removing the Self-Attention masking typically causes objects to appear outside of the layout foreground.
Refer to caption: /html/2406.10210/assets/x7.png

## 6 Ablation Study

##### Contributions of CountGen-Layout and CountGen-Image.

Table [3](#S6.T3) quantifies the contribution of each of these components to the overall accuracy, by replacing it with a baseline alternative. Compared with a baseline (Random Masks + Bounded Attention) our first phase CountGen-Layout improves accuracy measured by people by 14% (from 30 to 44), and our second phase CountGen-Image by 12%. Together, the two components add up to improve accuracy by 21 points.

##### Layout guided generation ablation study.

The second phase of our method, CountGen-Image, consists of two components: self-attention masking and object layout loss, as described at [Section 3.3](#S3.SS3). To evaluate the contribution of each component, we deactivate it and compare the results.
In [Figure 8](#S5.F8), we qualitatively observe that removing the layout loss leads to the objects scattering in the image, not constrained by the required mask. When removing the self-attention masking the objects tend to obey the mask unwanted object instances occur in the background.

We confirm these observations quantitatively in [Table 3](#S6.T3), where we evaluate the adherence of the generated image to the input mask. We use YOLOv9 to detect the bounding boxes of generated objects and compare them to the input mask using three metrics: Precision is the percentage of bounding boxes that highly overlap (IOU>0.6) the mask (union of all object masks), Recall is the percentage of mask pixels that are covered by bounding boxes, and IOU is measured between the boxes and the mask. Our findings align with the qualitative observation: removing the self-attention masking leads to a worse precision score, meaning objects are generated in the background. Removing the layout loss leads to low recall and IOU, meaning poor adherence to the mask. CountGen-Image, employing both components, achieves balanced results by generating objects in accordance with the mask. Overall, these results emphasize the critical roles that both components in ensuring accurate adherence to the input mask.

**Table 2: Model components Accuracy(%)**
|  |  | CoCoCount | Compbench [] |  |
| --- | --- | --- | --- | --- |
| Text $\rightarrow$ | Layout $\rightarrow$ | YOLOv9 | Human | Human |
| Layout | Image | Acc. | Acc. | Acc. |
| CountGen | CountGen | 50 | 54 | 48 |
| CountGen | B-Attn | 40 | 42 | 40 |
| Random | CountGen | 37 | 44 | 42 |
| Random | B-Attn | 29 | 30 | 35 |

## 7 Limitations

Occasionally, our optimization ([Section 3.3](#S3.SS3)) results in multiple instances of an object in an area intended for just one by the layout. In other cases CountGen generates plain backgrounds compared to SDXL ([Figure 11](#A2.F11)).
Finally, the scope of our experiments may seem narrow, since we focus on generating scenes with up to 10 instances and a single object per prompt. Nevertheless, we have shown in [Section 5](#S5) that even this setup is highly challenging to contemporary models, especially as the number of instances required to generate grows, as evident by the massive drop in performance, even for DALL-E 3 (see [Figure 8](#S5.F8)).

## 8 Conclusions

The task of generating images that depict the number of requested objects correctly is a hard task. It requires models to capture “objectness", and obey global spatial constraints, at the same time they generate a well-formed natural image. Current text-to-image diffusion models perform poorly in this task (Table [1](#S4.T1)), especially when asked to generate more than three objects (Figure [8](#S5.F8)).

Our CountGen approach took three steps to address this task. First, we identified a notion of objectness from the self-attention layers of the diffusion model. Then, we trained a U-Net model that learned to correct the number of instances of an object in a given layout, whether it is removing or adding instances of an object such that the structure of the layout is preserved. Third, we developed a layout-guidance optimization method method to generate images from the corrected layout. Together, this approach almost doubled the counting accuracy from 26% in standard SDXL to 54% using our method applied to SDXL. We expect the lessons learned from this method, specifically the features that represent objectness and the process of learning to automatically fix a layout, to become useful in other problems of structured generation like spatial constraints in text-to-image models or spatio-temporal constraints in video generation.

## Acknowledgments and Disclosure of Funding

We thank Yuval Atzmon for useful discussions and for providing feedback on an earlier version of this manuscript. This study was supported by an equipment grant from the Israel Science Foundation (ISF 2332/18), and by a grant from the National committee of budgeting (VATAT).

## Appendix A Appendix / Supplemental Material

##### Efficiency.

CountGen takes ~36 seconds on average to generate an image on a single A100 80GB. We arrive at this number by iterating over CoCoCount. To put in context, Bounded-Attention[] takes ~55 seconds and requires bounding boxes as input, while our solution is not input-dependent. SDXL takes ~8 seconds.

##### Compute.

All experiments were conducted over a period of a week on a single A100 80GB.

## Appendix B Implementation Details & Reproducibility

### B.1 CountGen

##### Layout guided generation.

In our implementation, the self-attention masking is applied at timesteps $t\in[1000,900]$, in the decoder layers of the U-Net. The object layout loss is applied at timesteps $t\in[1000,500]$, in all layers of the U-Net. Our pipeline used the Attend-and-Excite [] code base as a starting point.

##### ReLayout

The ReLayout U-Net was built upon the U-Net Implementation of []. We trained the U-Net with a learning-rate of *8e-6*, a batch-size of size *32* and the Adam optimizer. The intersection penalty is set to *0.25* and the Dice penalty is set to *1.* During training we apply a horizontal flip augmentation across all masks, and shuffle augmentation where we randomly re-arrange the input channels.

##### Instance identification.

In the DBSCAN clustering algorithm, we used a dynamic epsilon value in the range of $[0.1,0.2]$ and used cosine similarity as the distance metric.

### B.2 Compared Methods

Each prompt in CoCoCount and T2I-CompBench-Count was assigned a unique random seed and was used by all baselines and CountGen.

We compared CountGen with the following baselines:

SDXL []. We used the stable-diffusion-xl-base-1.0 model.

Repeated Object In this baseline, we used the same model and seeds as in SDXL but modified the prompts. We repeated the object in the prompt as many times as the target count. For example, “a photo of three cats” was changed to “a photo of a cat and a cat and a cat”.

Reason Out Your Layout []. This baseline has two main steps. First, it leverages GPT-3.5-turbo to generate spatially reasonable coordinates to be used as a bounding box for each instance of an object (i.e., “a photo of three cats” results in three bounding boxes, one for each cat). Second, it uses the generated layout to guide the generation process. We followed the prompt used by the authors, however, it seems that the responses by GPT-3.5-turbo and the author’s parser are not completely cohesive, which at times leads to zero bounding boxes. We count such cases as failures. For the CoCoCount experiment, it successfully generated 134/200 images, and for T2I-CompBench-Count, just 89/200. Failures were counted as errors in the reported results. We did not need to make changes to the code to run it.

DALL-E 3 []. We used the OpenAI API interface for the DALL-E 3 model with “standard” image quality. We did not use seeds in this baseline.

Random masks + BoundedAttn [].
Given a prompt with a required number of object instances, we create a corresponding layout with the correct number of objects randomly placed in the image plane in a way they do not intersect one another. Then we used Bounded Attention to generate an image condinitioned on that layout.

Counting Guidance []. The authors provided us with their code. We did not need to change it to run our experiments.

Figure: Figure 9: Accuracy, as a function of the number of generated objects. Accuracy evaluated by human raters, over the set of 200 evaluation images. CountGen (blue) outperforms all methods for $n>3$, and is on par with DALL-E 3 for 2 and 3 objects.
Refer to caption: /html/2406.10210/assets/x9.png

Figure: Figure 11: Limitations. Failure modes of CountGen.
Refer to caption: /html/2406.10210/assets/x11.png

Figure: Figure 12: PCA visualization across timestamps. to explore the notion of objectness inside SDXL latent space, we visualize a dimension-reduced self-attention feature maps across different timestamps range from $t=900$ to $t=100$. Initially, up to timestamp $t=500$, clear separation is not observed in some objects (e.g., some eggs appear in similar colors). However, starting from $t=500$, a distinct separation emerges, with each object clearly distinguished by different shades.
Refer to caption: /html/2406.10210/assets/x12.png

## Appendix C Extended Details on CountGen

### C.1 ReLayout: Matching Objects

We aim to understand how $M_{k}^{i}$ transitions to $M_{k+1}^{i}$.
Specifically, for each object $i\in{1,\ldots,k}$ in the original $M_{k}$ layout, our ReLayout objective is designed to predict how the corresponding mask $M_{k}^{i}$ changes in the new image $M_{k+1}^{i}$, and additionally where to insert the added object $k+1$. This design encourages the model to slightly modify existing objects while preserving spatial and shape consistency across the images.

To this end, we first have to establish a correspondence between the object masks $(M_{k},M_{k+1})$. We employ the Hungarian algorithm [] to find the optimal one-to-one matching between masks in the two images based on the overlap and similarity of the masks.
This algorithm effectively pairs each object in $M_{k}$ with a corresponding object in $M_{k+1}$. The object in $M_{k+1}$ that remains unmatched represents the additional object introduced in the new image, providing a clear identifier for the increment in object count.

### C.2 Losses for Training the ReLayout

We use two training losses:

Dice Loss: measures the overlap between the predicted mask and target mask across all channels containing objects:

$$ $L_{\text{Dice}}^{i}=1-\frac{2\sum_{p\in P}M_{k+1}^{i}(p)\cdot M_{k+1}^{*i}(p)}{\sum_{p\in P}(M_{k+1}^{i}(p)+M_{k+1}^{*i}(p))}$ (2) $$

Here, $p$ iterates over all pixels $P$ in the masks, and $i$ ranges over all possible object channels. For all $k+1$ channels, the total dice loss is:

$$ $L_{\text{Dice}}=\sum_{i=1}^{k+1}L_{\text{Dice}}^{i}$ (3) $$

Intersection Loss:
To ensure distinctiveness among the predicted masks and to minimize overlap between different object masks, the intersection loss for all possible pairs of different masks in the output mask containing objects is given by:

$$ $L_{\text{Overlap}}=\sum_{i=1}^{k+1}\sum_{j\neq i}^{k+1}\frac{2\sum_{p\in P}M_{k+1}^{i}(p)\cdot M_{k+1}^{j}(p)}{\sum_{p\in P}(M_{k+1}^{i}(p)+M_{k+1}^{j}(p))}$ (4) $$

### C.3 Datasets

##### CoCoCount.

To create this set, we first select at random 20 classes from MSCOCO []. We then sample from six counting categories: 2,3,4,5,7, and 10. The two and three categories contain 34 samples, while the rest contain 33. Our prompts consist of the pattern “a photo of {number} {object}” with an optional variation of scenes: “on the grass”, “on the road”, or “on the ground”, which we incorporate for 50% of the prompts, also randomly. In total, we have 200 prompts.

Below are the complete lists from which elements were chosen:

- •
Objects: ’car’, ’airplane’, ’bird’, ’cat’, ’dog’, ’horse’, ’sheep’, ’cow’, ’elephant’, ’bear’, ’backpack’, ’tie’, ’sports ball’, ’baseball glove’, ’cup’, ’bowl’, ’apple’, ’donut’, ’cell phone’, ’clock’
- •
Counting Categories: ’two’, ’three’, ’four’, ’five’, ’seven’, ’ten’
- •
Scenes: ’on the grass’, ’on the road’, ’on the ground’

## Appendix D Evaluation

##### Automatic evaluation.

We use the implementation by [Ultralytics YOLO](https://github.com/ultralytics/ultralytics) of YOLOv9e (large).

### D.1 Human Evaluation

We use the Amazon Mechanical Turk platform and ensure the evaluation is of high quality by hiring raters with a minimum of 5,000 approved HITs and an approval rate exceeding 98%. Each example was shown to three raters and the majority selection was taken. The compensation was $15 per hour. Screenshots of the count precision task can be viewed in [Figure 13](#A4.F13), [Figure 14](#A4.F14), [Figure 15](#A4.F15) and the image fidelity task in [Figure 16](#A4.F16).

Figure: Figure 13: Instructions for the Image Evaluation Task - Part 1.
Refer to caption: /html/2406.10210/assets/figures/amt_figures/amt_instructions_1.png

Figure: Figure 14: Instructions for the Image Evaluation Task - Part 2.
Refer to caption: /html/2406.10210/assets/figures/amt_figures/amt_instructions_2.png

Figure: Figure 15: Example task to count the number of objects and assess their well-formedness.
Refer to caption: /html/2406.10210/assets/figures/amt_figures/count_precision_task.png

Figure: Figure 16: Example task to compare image fidelity based on prompt matching and naturalness.
Refer to caption: /html/2406.10210/assets/figures/amt_figures/image_fidelity_instructions.png

Figure: Figure 17: Confusion matrix of human evaluation ([Section 4](#S4)) of the count accuracy experiment for CountGen.
Refer to caption: /html/2406.10210/assets/figures/amt_figures/confusion_mat_human_eval.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: * Does CountGen increase inference time compared to baselines? If so, how does its runtime compare to the one of previous methods?
* How does CountGen compare with related works in terms of performance mentioned in the section above?
* Is the discovery that instance feature exists in text-to-image diffusion models generalizable (i.e., not just for SDXL)?
soundness: 2
strengths: * The work identifies features that represent object instances rather than semantics in SDXL, offering a new understanding of the internal representations of diffusion models.
* The work proposes CountGen, a technique to guide SDXL for accurate number of object instances, with a novel object layout loss and self-attention masking strategy.
* Detailed benchmarking demonstrates that CountGen improves accuracy in counting without degrading image quality.
summary: Existing text-to-image models struggle with generating a certain number of objects in images. The paper attributes this inability to keep a sense of separate identity for every instance of the object. The paper resolves this problem by first identifying features within a diffusion model, SDXL, that carry object identity information. Then the paper proposes CountGen, which counts instances of objects during denoising, detects over- and under-generation, and fixes the latter by predicting missing objects to guide denoising in order to obtain a correct object count. Results on two benchmarks show that CountGen outperforms existing baselines in the accuracy of generating a specified number of objects.
weaknesses: * The analysis in the introduction is unclear. Speficically, the work identifies the cause of inability to generate a specified number of objects to be the lack of instance-level features (L61-65). However, the work later discovers the presence of such features (L199-210), which contradicts the claim. The fact that previous works have not found such features is not an explanation on why current diffusion models do not perform well in counting tasks.
* Lack of comparisons with existing baselines. The ReLayout proposed in this work is very similar to the closed-loop generation proposed in Self-correcting Diffusion Models [1], which aims at solving problems in a larger scope that includes count-specific generation. The authors are recommended to discuss the relevance and compare with [1] in terms of performance.
* The proposed method requires finding features that are related to instanceness for each model. Furthermore, the authors did not provide explanations on the process of finding such features, which are the core of the method.
* The authors only experimented with SDXL, which does not show generalization capabilities of the proposed method.
* The proposed method introduces an additional module to generate the layouts. This might increase the run-time, impacting downstream use cases. The authors are recommended to show the wall-clock runtime and compare with previous baselines.

[1] Self-correcting LLM-controlled Diffusion Models. T. Wu, et al. CVPR 2024. https://arxiv.org/abs/2311.16090.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Please refer to the weakness. 

I would like to say this paper is solving a good problem existing in T2I diffusion models. However, this method needs to create a dataset and train the ReLayout module for each kind of T2I diffusion model, which limits its application scenario. As a comparison, the paper “Counting Guidance for High Fidelity Text-to-Image Synthesis” doesn't need the training process.
soundness: 2
strengths: 1. The authors identify a specific layer $L_{52}^{up}$ within the SDXL model that exhibits strong object instance separability in its self-attention features. 
This discovery is significant because it reveals an inherent capability of diffusion models to represent individual objects as distinct entities, which is crucial for accurate counting. 
The paper provides evidence for this through PCA visualization of self-attention feature maps.

2. The paper proposes a novel ReLayout module to correct the number of object instances in the generated layout. This module tackles both over-generation and under-generation. This paper also introduces a layout-guided generation process that combines object layout loss and self-attention masking. The object layout loss encourages object creation within the defined masks, while self-attention masking prevents objects from being generated in the background. All these techniques ensures that the final image closely adheres to the corrected layout.

3. The paper introduces CoCoCount, a new dataset designed specifically for automatic evaluation of count accuracy. Based on this new benchmark, it presents a comprehensive evaluation of CountGen on the CoCoCount and T2I-Compbench-Count benchmarks.
summary: This paper presents a novel method, CountGen, which is a two-stage method for improving the accuracy of object counts in text-to-image generation using diffusion models. 
The first stage identifies an “objectness” representation in the SDXL model by analyzing self-attention features at layer $L_{52}^{up}$ at time t=500. It then uses cross-attention maps to locate and cluster object instances, creating an initial layout. 
The second stage, ReLayout, corrects this layout for over-generation or under-generation of objects. 
For under-generation, ReLayout employs a U-Net trained on layout pairs with varying object counts, adding objects while preserving scene composition. 
Finally, layout guidance, using object layout loss and self-attention masking, directs the generation process to adhere to the corrected layout. This approach significantly improves count accuracy, especially for larger numbers of objects, without compromising image quality.
weaknesses: 1. CountGen relies on the specific architecture of the SDXL model, particularly the self-attention features extracted from layer $L_{52}^{up}$. This dependence raises concerns about the method's transferability to other diffusion models. Further investigation is needed to assess whether the concept of objectness, as captured by CountGen, can be generalized to different architectures or whether model-specific adaptations would be required. Since the PCA visualization in Fig.3 should be easily reproduced using other models, it could be verified in time.

2. The ReLayout module is trained on a dataset of layout pairs generated by SDXL using the same random seed but slightly different object counts. This approach, while clever in its utilization of consistent layout generation by SDXL, could potentially lead to overfitting. ReLayout might learn to exploit specific artifacts or biases in SDXL's output, limiting its effectiveness when applied to layouts generated by other models or even by SDXL with different seeds. For example, how to make sure that SDXL is able to always generate correctly with "two apples" and "three apples" with the similar layout. Over here, I'm not sure if any VLM model (like visual question answering) could help.

3. While the paper briefly mentions the time taken by CountGen to generate an image (~36 seconds), it lacks a detailed discussion of the computational cost compared to other methods, especially in relation to the achieved accuracy gains.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Asked above.
soundness: 2
strengths: The problem that this paper is addressing is critical, as all SOTA image generative models fail in solving counting tasks. The authors attempt to solve this problem by adopting a pre-trained text-to-image diffusion model during inference time to correct the number of objects expected in the scene. Hence, no costly pretraining is needed for unlocking these new abilities. The paper addresses the task of generating images using a specific prompt format, such as "A photo of [number of objects] [object type]." It conducts numerous evaluations and ablation studies to demonstrate the effectiveness of each component. In my view, all these parts are highly effective in enabling the diffusion model to accurately count objects.
summary: This paper presents a pipeline designed to facilitate counting tasks using pre-trained text-to-image diffusion models. The pipeline consists of three stages: instance localization, ReLayout, and layout-guided image generation. The results demonstrate that the proposed approach significantly outperforms all state-of-the-art baselines in counting tasks.
weaknesses: I think the paper introduces many exciting components that address the counting problem. However, I believe the paper suffer from a serious shortcoming which should be discussed/addressed in the next version of the paper:

**Overfitting to one specific prompt format**: In all visualizations and explanations of the paper, I notice that it consistently uses a prompt template structured as "A photo of [number of objects] [object name]." For example, "A photo of five apples" or "A photo of five backpacks." Why only use this prompt template in all ablations and qualitative visualization? What about exploring additional prompt templates, such as those provided in the T2I-CompBench dataset? I would suggest authors to discuss about following points during rebuttal or future submissions:

1. Explicitly state the limitations of their current approach in handling more complex prompts.
2. Provide results on a wider range of prompt templates from T2I-CompBench, even if performance is lower.
3. Discuss potential ways their method could be extended to handle multiple object types and attributes in future work.

More specifically, given following prompts:

```
Prompt 1: one apple and two oranges
Prompt 2: four balloons, two rabbits, three cars, four butterflies, and a dog created a festive atmosphere
Prompt 3: three printers, four people, a bicycle, three microwaves, and two refrigerators cluttered the office
```

Can authors:
1. show results of their instance localization, ReLayout, and layout-guided generation steps on 2-3 example prompts with multiple object types
2. discuss which components of their pipeline would need to be modified to handle these more complex prompts
3. quantify performance on a subset of T2I-CompBench prompts with multiple object types, if possible

***Note***: The above issue is concerning because I believe this solution cannot scale to other prompt templates, which require more object types, each associated with an attribute, such as "A photo of 6 green apples." If the authors believe their proposed method can effectively generate different objects from various categories within a scene, I would ask them to provide this information during the rebuttal. Otherwise, I believe this paper has limited contributions/applications.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1.I hope the authors can achieve fair comparisons by comparing different baselines under the same base model (e.g., sd1.5 or sd2). For example, you can try to migrate CountGen to sd1.5, or migrate the baselines to be compared to sdxl.

2.I hope the authors can explain why they did not include experiments with the YOLO metric in T2IBench. And I hope they can include this experiment in the new version.

3.I hope the authors can explain the advantages and disadvantages of ReLayout and GPT-Based Layout, and clarify the novelty of CountGen.

4.I hope the authors can explain the reason for selecting l_{52}^{up} in UNet and whether it has generality.

If these confusions are clarified, I will consider raising the score rating.
soundness: 2
strengths: This article focuses on a subproblem within diffusion inference process——how to generate images with a more precise number of objects. The authors provide a detailed discussion of related work, including Counting Guidance [1] and RPG [2]. They also introduce a new benchmark, which contributes to advancing research within the whole community. Lastly, the paper covers a wide range of techniques, including improving guidance, enhancing instance segmentation, and assembling these components into a pipeline. Particularly, their discussion on "Self-attention masking" is interesting.

[1] Kang, Wonjun, Kevin Galim, and Hyung Il Koo. "Counting guidance for high fidelity text-to-image synthesis." arXiv preprint arXiv:2306.17567 (2023).

[2] Yang, Ling, et al. "Mastering text-to-image diffusion: Recaptioning, planning, and generating with multimodal llms." Forty-first International Conference on Machine Learning. 2024.
summary: This paper addresses the problem of generating a precise number of objects in Diffusion models.  From my understanding, it introduces: 1) an instance segmentation framework based on Attention; 2) a network for modifying the layout (especially in cases where the count of objects needs to be increased); and 3) a well-designed guidance method to direct the image generation towards the modified layout. The authors conducted main experiments on two benchmarks (one of which was introduced in this paper) as well as ablation studies, demonstrating the effectiveness of the proposed method.
weaknesses: 1.The comparison of baselines is unfair. 

In Table 1, the authors compare CountGen with other baselines, but CountGen is based on the SDXL model, while some of these baselines are not. For example, the base model for Reason Out Your Layout [1] is SD1.4. And Counting Guidance [2], as a further improvement of Universal Guidance, uses SD1.5 or SD2. If you haven’t modified its source code (as mentioned in Line 797), it means the base models are different, making this comparison unfair. Therefore, migrating CountGen to other SD models for a fair comparison should be part of your work. Otherwise, this kind of unfair comparison is detrimental to the progress of the entire field. I think you should conduct experiments on other Diffusion Base Models (e.g, sd1.4, sd1.5, sd2) to achieve fair comparisons. It's worth noting that this is not an unnecessary increase in experimental workload.

2.No quantitative experiments in T2IBenchmark.

The authors use both human metrics and YOLO metrics. They apply these metrics in the CoCoCount benchmark, but what confuses me is why they do not use objective quantitative indicators (YOLO) in the T2ICompBench, and only report human evaluation scores in the paper. Given that CoCoCount is one of the contributions of this paper and has not yet been publicly released, the quantitative evaluation only on the self-built dataset is not convincing. I hope the authors can first explain why they only report human evaluation results in T2IBench? What is the reason for this? If possible, further provide the evaluation results of T2IBench.

3.The importance of the "Corrected Layout" is not sufficiently emphasized.

The core of CountGen lies in first providing the original layout and then correcting it. What is the advantage of CountGen over models like RPG[3]?  CountGen first detects the incorrect layout and then corrects it, finally guiding the generation. Why not directly plan the layout from the beginning and then directly guide the generation using this layout?  I guess the author's idea is similar to the work in OMG [4], where they first use the original capabilities of Diffusion to provide a layout and then modify it on this basis, possibly improving the quality. However, this paper does not provide sufficient quantitative analysis to compare the advantages and disadvantages of the "GPT-Based-Layout" and "Corrected Layout" methods, but only conducts ablation experiments on "Random Layout".  I think this paper should further explain the necessity of using the "Corrected Layout" and provide qualitative or quantitative ablation experiments for comparison.

Additionally, the "ReLayout" module can only handle the case where the number of extra objects is "1" (this is due to the setup of the training dataset, see line 266). However, in Diffusion generation, there are often more than one extra objects, for example, generating 2 cats but the actual generated image has 4 cats. In such cases, ReLayout cannot be applied, so the universality of this method is not strong. The authors might try to construct more diverse datasets and add conditional quantity signals in ReLayout to output the corrected Layout.

4.The selection of UNet features lacks generality.

During the first stage of object segmentation, the authors use a specific layer in UNet and empirically believe that this layer works the best. Similarly, like weak 1, is this layer the best for sd, sd2, and sdxl? If not, then for each different Diffusion Model, the optimal layer would have to be manually and empirically selected again, increasing the difficulty of deploying this method. I think the authors should further research this part to enhance the generality of feature selection.


[1] Chen, Xiaohui, et al. "Reason out your layout: Evoking the layout master from large language models for text-to-image synthesis." arXiv preprint arXiv:2311.17126 (2023).

[2] Kang, Wonjun, Kevin Galim, and Hyung Il Koo. "Counting guidance for high fidelity text-to-image synthesis." arXiv preprint arXiv:2306.17567 (2023).

[3] Yang, Ling, et al. "Mastering text-to-image diffusion: Recaptioning, planning, and generating with multimodal llms." Forty-first International Conference on Machine Learning. 2024.

[4] Kong, Zhe, et al. "OMG: Occlusion-friendly Personalized Multi-concept Generation in Diffusion Models." arXiv e-prints (2024): arXiv-2403.