Kaleido Diffusion: Improving Conditional Diffusion Models with Autoregressive Latent Modeling

Authors: jiatao  gu
ying  shen
shuangfei  zhai
yizhe  zhang
navdeep  jaitly
josh  susskind
†  apple
virginia  tech
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:16:01

Abstract: Diffusion models have emerged as a powerful tool for generating high-quality images from textual descriptions. Despite their successes, these models often exhibit limited diversity in the sampled images, particularly when sampling with a high classifier-free guidance weight. To address this issue, we present Kaleido, a novel approach that enhances the diversity of samples by incorporating autoregressive latent priors. Kaleido integrates an autoregressive language model that encodes the original caption and generates latent variables, serving as abstract and intermediary representations for guiding and facilitating the image generation process.
In this paper, we explore a variety of discrete latent representations, including textual descriptions, detection bounding boxes, object blobs, and visual tokens. These representations diversify and enrich the input conditions to the diffusion models, enabling more diverse outputs.
Our experimental results demonstrate that Kaleido effectively broadens the diversity of the generated image samples from a given textual description while maintaining high image quality. Furthermore, we show that Kaleido adheres closely to the guidance provided by the generated latent variables, demonstrating its capability to effectively control and direct the image generation process.

Main paper: 
 introduction: 
 Diffusion models have become pervasive in many text-to-image generation tasks for their ability to generate high-quality images based on textual descriptions. A pivotal mechanism in these models is classifier-free guidance (CFG) [Ho and Salimans, 2021] , which effectively steers the sampling process towards better alignment to textual prompts and improved sampling quality at the same time. CFG can be interpreted as tuning the temperature of the conditional distribution, whereas increasing the guidance scale sharpens the conditional distribution. This guides the generation to focus on regions of high conditional probability, effectively reducing sampling noise which is typically of lower density. However, while high CFG improves sampling quality, it simultaneously narrows the diversity in the generated samples. This manifests in the models' inability to produce diverse images from the same caption, even when there are variations in the initial noise that seeds the generation process. For instance, given a fixed textual description, "a cat sits on a mat", existing text-to-image diffusion models predominantly produce image samples depicting cats with similar colors and patterns, as illustrated in Figure 1 . Such limited visual diversity hinders the practical application of diffusion models in scenarios where a wide range of creative and diverse visual interpretations are desired from identical textual inputs. It also poses challenges in scenarios demanding the representation of underrepresented data or accommodating a wide range of user preferences. Therefore, enhancing diversity in diffusion models without compromising the quality remains a critical research problem.To tackle this, we introduce Kaleido, a general framework that improves diffusion models with autoregressive priors. Kaleido first defines a discrete encoding of images (eg, detailed captioning, bounding boxes), which captures desirable abstractions of images that's not included in the default text prompts. Next, Kaleido integrates an encoder-decoder language model that encodes the original text caption and autoregressively predicts the discrete latent tokens. Lastly, the diffusion model is conditioned on both the original text prompt and the autoregressively generated discrete latents and generates an image. This enriched conditioning allows Kaleido to produce a more diverse array of high-quality images, even at high guidance scales. We explore various forms of latents, including textual descriptions, detection bounding boxes, object blobs, and abstract visual tokens -all designed to refine and guide the conditional image generation process.We experiment on both class and text conditioned image generation benchmarks1 . We show that Kaleido not only outperforms standard diffusion models in terms of diversity but also maintains the high quality of the generated image. Additionally, the generated latents effectively control the characteristics of the generated images, ensuring that the image samples closely align with the intended latent variables. This modeling of latent tokens not only increases the diversity of image outputs but also provides a degree of interpretability and control over the image generation process.To summarize, Kaleido exhibits the following advantages:1. Kaleido promotes the diversity in generated image samples even with high CFG, allowing the image generation of both high quality and diversity.2. The generated latent variables are interpretable, offering an explainable mechanism behind the image generation process, and facilitating an understanding of how different latents affect the outputs.3. Kaleido provides a fine-grained, editable interface that allows users to adjust the discrete latent codes before final image production, granting greater flexibility and control over the output. 
preliminaries: 
 Autoregressive Image Generation The success of large language models (LLMs) in NLP has demonstrated their scalability and universality of modeling any complex data, motivating the development of using autoregressive models for image generation. Typically, autoregressive image generation operates on discrete image tokens obtained from vector-quantization (VQ) [Van Den Oord et al., 2017] . More precisely, given an image x ∈ R 3×H×W , we first obtain a sequence of discrete tokens z 1:N = E(x) which approximately reconstructs the input with a learned decoder D(z 1:N ) ≈ x.Then, an autoregressive model is learned to predict the discrete tokens one after another, mirroring the sequential language modeling:EQUATIONwhere c is the condition (e.g., class, text prompt, etc.), and z 0 is a special start token. At inference time, we first sample from the learned distribution, and then pass the sampled latents to the decoder (D) to get the final output. Such VQ-based paradigm has been the foundation for various text-toimage [Esser et al., 2021 , Yu et al., 2021 , Zheng et al., 2022 , Yu et al., 2022] and multi-modal generation [Team et al., 2023 , Team, 2024] .However, these methods share a common limitation: they primarily rely on discretization, which struggles to capture all the nuances of an image when using a limited length of discrete image token sequence. To generate higher-resolution images, a longer sequence of image tokens is necessary.Yet, this inherently leads to increased capacity demands. For instance, Yu et al. [2022] requires 20B parameters to work properly. Additionally, the left-to-right properties of these autoregressive models prevent the rewriting of previously generated image tokens, resulting in suboptimal image quality.Diffusion-based Image Generation Diffusion models [Sohl-Dickstein et al., 2015 , Ho et al., 2020] are latent variable models with a pre-determined posterior distribution and are trained using a denoising objective, which has quickly become the new de-facto approach for image generation.Unlike autoregressive models which predict images as a sequence, diffusion-based models iteratively generate the whole image in a non-autoregressive fashion. Specifically, given an image x ∈ R 3×H×W and a signal-noise schedule {α t , σ t } where the signal-to-noise ratio (SNR) (α 2 t /σ 2 t ) decreases monotonically with t, we define a series of latent variables x t , t = 0, . . . , T that adhere to:EQUATIONwhere x 0 = x, α t|s = α t /α s , and σ 2 t|s = σ 2 t -α 2 t|s σ 2 s for s < t. The model then learns to reverse this process using a backward model p θ (x s |x t , c), which reformulates a denoising objective:EQUATIONwhere x θ (x t , c) is a neural network (typically a UNet [Ronneberger et al., 2015] or Transformer [Peebles and Xie, 2022] ) that maps the noisy input x t to its clean version x, based on the time step t and conditional input c; ω t ∈ R + is a loss weighting factor. In practice, x θ can be re-parameterized with noise-or v-prediction [Salimans and Ho, 2022] for enhanced performance, and can be applied on raw pixel space [Saharia et al., 2022 , Gu et al., 2023] or latent space [Rombach et al., 2022] .Classifier-free Guidance An intriguing property of conditional diffusion models is that we can easily guide the iterative sampling process for better sampling quality. For instance, Ho and Salimans [2021] introduced Classifier-free Guidance (CFG), which utilizes the diffusion model itself to perform guidance at test time. More specifically, we perform sampling using the following linear combination:EQUATION) where γ is the guidance weight, and x θ (x t ) = x θ (x t , c = ∅) is the unconditional denoising output.During training, we drop the condition c with certain probability p uncond to facilitate unconditional prediction. When γ > 1, CFG takes effect and amplifies the difference between conditional and unconditional generation, leading to a global control of high-quality generation.Compared to autoregressive models, diffusion models are more flexible in adjusting sample steps, allowing for the utilization of noise schedules to learn different frequencies. Additionally, with the use of CFG, diffusion models can achieve higher quality images with much fewer parameters than autoregressive models. However, it's notable that CFG can significantly impact the diversity of the diffusion output, which motivates us to revisit the basics and combine the strengths of both.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: Due to page limits, we include discussion of limitations in appendix section. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
kaleido diffusion: 
 We propose Kaleido, a general framework that integrate an autoregressive prior with diffusion model to enhance image generation. As illustrated in Fig. 2 , Kaleido comprises two major components: an AR model that generates latent tokens as abstract representations, and a latent-augmented diffusion model that iteratively synthesizes images based on these latents together with the original condition. For following sections, we first describe the importance to introduce additional latents in standard diffusion models ( § 3.1), and show how we can model them with AR models ( § 3.2). The training and inference procedure are described in § 3.3 and § 3.4.As demonstrated in Ho and Salimans [2021] , diffusion with CFG (Eq. ( 4)) is equivalent to follow which can be interpreted as sampling from a "temperature-adjusted" distribution:EQUATIONEQUATION)Here γ can be seen as inverse temperature, which sharpens the conditional distribution p θ (c|x) when γ > 1. That is to say, CFG is crucial as it guides the generation to only focus on high-probability regions, avoiding sampling noise (which tends to have low density). However, sharpening the distribution also reduces the diversity, causing undesirable phenomena like "mode collapse". This is because c (e.g., class label, text prompt, etc.) normally does not contain all the information that describes x. Suppose we introduce a hypothetical variable z to represent the "modes" of x which we care most -p θ (z|c), and leave p θ (x|z, c) to model other variations including local noise. In this case, CFG will simultaneously sharpen both distributions, considering:EQUATIONwhere standard diffusion models implicitly learn mode selection step together with generation.Therefore, a natural solution is to explicitly model "mode selection" before applying diffusion steps so that the mode distribution will not be distorted by guidance. In this way, the sampling procedure (Eq. ( 6)) is modified as two steps: z ∼ p θ (z|c), x ∼ pθ (x|z, c), where CFG can be applied after z is sampled. From the perspective of score function, we rewrite pθ (x|c) as pθ (x|c, z) in Eq. ( 5):∇ x log pθ (x|c, z) = γ ∇ x log p θ (x|c) + log p θ (z|x, c) -log p θ (x) + ∇ x log p θ (x). (8)Compared to standard diffusion process, the highlighted term above pushes the updating direction towards the sampled modes at each step. This ensures diverse generation as long as p θ (z|c) is diverse.A Toy Example We visualize the effect of explicitly introducing latent priors using a toy dataset with two main classes, each containing two modes. We compare two models: a standard diffusion model conditioned on the major class ID, and a latent-augmented model incorporating subclass ID as priors. Fig. 3 shows that while the standard diffusion model tends to converge to one mode (subclass) with increased guidance, the latent-augmented model captures all modes, showing the benefit of latent priors for improving diversity under high guidance. In practice, given the challenge of identifying all "modes" in real-world data distribution, we next propose to employ an autoregressive model to universally model various latent modes.To capture the complex distribution of real images, it is clearly impossible to assign classes for each mode. However, it is non-trivial to determine (1) the best representations for modes z;(2) the suitable generative model that can model p θ (z|c). Fortunately, the modes that humans can perceive from an image are largely abstract, and such abstract semantics are easily represented in discrete symbols.For example, we can easily describe content differences through natural language, create composite images based on spatial locations, and imagine novel visual concepts from experience. Therefore, it is logical to use such abstract discrete tokens, i.e., z = [z 1 , . . . z N ]. Naturally, the most convenient way to model such distribution is using an autoregressive model p θ (z|c). Note that this is distinct from the conventional autoregressive image generation ( § 2), as we only learn such models as "latent modes", where the sampled z 1:N are not supposed to reconstruct an image directly. As a result, it eases the modeling difficulty and improves the sampling performance.In this work, we explore four types of abstract latents: textual descriptions (text), detection bounding boxes (bbox), object blobs (blob), and visual tokens (voken), all of which can be predicted from multimodal large language models (MLLMs) [Bai et al., 2023 , Liu et al., 2024 , Ge et al., 2023a] given the condition-image pair (c, x). Each type aims to enrich the mode-to-image correspondence, covering different aspects of image formation. These extracted tokens can either be predicted separately or modeled together with a single autoregressive model. Fig. 4 shows the examples of these four generated latent variables. The methodology for constructing the training dataset for these variables is detailed in Appendix A.Similar to other latent variable models like VAEs [Kingma and Welling, 2013] , Kaleido can be trained to maximize the evidence lower bound (ELBO) as follows:max θ log p θ (x|c) ≥ E z∼q(z|x,c) [log p θ (z|c)L AR Eq. ( 1)EQUATIONwhere q is the inference model, and H(q) is the entropy. In this paper, we always assume a fixed inference process (as explained § 3.2). Therefore, the entropy term can be omitted, and we can efficiently sample and store z for the entire dataset before training starts. We illustrate the training pipeline in Fig. 2 . Compared to standard diffusion models which typically involves a context encoder and a denoising network, Kaleido integrates the additional autoregressive decoder for modeling the discrete latents. Such decoder uses cross-attention to gather the encoder states at every step, and the final decoder layer states are concatenated with the encoder as the inputs for diffusion.Following common practices, we freeze the context encoder during training, and jointly optimize the autoregressive decoder together with the denoising model. The training objectives (Eq. ( 9)) is equivalent to the combination of both models, denoted as L = L DM + η • L AR , with η as a hyperparameter for balancing the contributions of the autoregressive and diffusion models in practice.During the inference stage, given the provided textual description, the autoregressive model will first predict the discrete latents before image generation. These latents, being predominantly humanreadable, add a layer of interpretability to the image-generation process, allowing humans to observe its internal "thought" process. This transparency also provides users with the flexibility to modify the latents as desired. To incorporate user modification, the altered latents are re-input into the autoregressive decoder to obtain the modified final hidden states. The latent-augmented diffusion model then synthesizes the final image conditioned on the updated representation. Dataset We validate our approach on both class-and text-conditioned image generation benchmarks. For the former, we use ImageNet [Deng et al., 2009] , and we learn the textto-image models on CC12M [Changpinyo et al., 2021] , a large image-text pair dataset where each image is accompanied by a descriptive alt-text. All models are trained to synthesize at 256 × 256. We generate all four types of latents as discussed in Appendix A for both datasets.Evaluation Metrics To assess the performance of our models, we employ Fréchet Inception Distance (FID) [Heusel et al., 2017] to capture the overall performance (considering both quality and diversity) of the generated images, and use Recall [Kynkäänniemi et al., 2019] to specifically measure the diversity of the generated images. Furthermore, we employ two additional quantitative assessments of diversity: Mean Similarity Score (MSS) and Vendi scores [Friedman and Dieng] . We use SSCD [Pizzi et al., 2022] as the pretrained feature extractor for calculating both MSS (SSCD) and Vendi (SSCD). Additionally, we utilize DiNOv2 [Oquab et al.] as the feature extractor for Vendi (DiNOv2), based on evidence from Stein et al. [2024] that suggests DiNOv2 provides a richer evaluation of generative models.We implement Kaleido with Matryoshka Diffusion Models (MDM) [Gu et al., 2023] , a recently proposed approach that generates images directly in the raw pixel space with efficient training. The default MDM consists of a frozen T5-XL [Radford et al., 2021] context encoder and a nested UNet-based denoiser. We initialize the additional autoregressive decoder with the decoder of T5-XL, and make the parameters trainable. The vocabulary is resized to adapt special visual tokens. For fair comparison, we use MDM with the same hyper-parameters as our baseline model, and train both types in almost identical settings on 64 A100 GPUs. Additionally, we compare Kaleido with the Condition Annealed Diffusion Sampler (CADS) [Sadat et al.] , a general sampling strategy that enhances the diversity of diffusion models by annealing the conditioning signal during inference. Given that CADS is applicable to different model architectures, we also evaluate CADS integrated with both baseline model MDM (MDM + CADS) and our model (ours + CADS).Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA]Answer: [No] Justification: Code will be released after acceptance and internal review. Guidelines:• The answer NA means that paper does not include experiments requiring code. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment Statistical Significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [No] Justification: Due to the costly nature of each run, we did not include error bars Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Due to page limits, we include societal impacts in the appendix Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to 
quantitative results: 
 Fig. 5 quantitatively compares Kaleido with the baseline diffusion models (MDM) with various guidance scales on ImageNet. Both metrics are evaluated with 50K samples against the full training set, where both our models and the baseline use DDPM sampling with 250 steps. Our findings reveal that Kaleido consistently enhances the diversity of samples without compromising their quality across different CFG, evidenced by the general improvement in both FID and Recall. Moreover, while the baseline's FID increases and Recall decreases significantly with higher CFG, Kaleido demonstrates a steadier performance profile. To further investigate, we examine image quality and diversity between Kaleido and baseline models.As shown in Table 1 , Kaleido outperforms the MDM + CADS combination in terms of FID-50K and precision, demonstrating that our method more effectively maintains high image quality while generating diverse samples. Furthermore, integrating CADS with our model yields the best FID-50K results. Note that precision cannot accurately evaluate models with diverse outputs since a model producing high-quality but non-diverse samples could artificially achieve high precision [Sadat et al.] .Moreover, we assess the diversity of the generated images using 10K samples. Following CADS, we select 1, 000 random classes from ImageNet and generate 10 samples per class. Table 1 shows that both Kaleido and CADS significantly enhance sample diversity. While CADS achieves better performance in diversity, our model maintains superior image quality. Additionally, the methodologies used in CADS are complementary to ours, suggesting potential benefits from integrating CADS with our Kaleido. In fact, incorporating CADS into our model not only further improves image quality but also improves diversity, achieving the best scores in FID-50K, MSS (SSCD), and Vendi (DiNOv2).Lastly, we provide visual comparisons for class-and text-conditioned image generation in Fig. 11 . Notably, we observe that MDM + CADS fails to generate cats of diverse breeds from the prompt "a cat sleeping on the bed." In contrast, Kaleido can produce images of cats from various breeds with more diverse surrounding environments, showcasing its superior diversity capabilities. This observation contrasts with the trend of diversity scores in Table 1 , suggesting that these diversity metrics may not fully capture certain aspects of diversity.We present a comparative analysis of the images generated by Kaleido against baseline models (MDM). Fig. 7 demonstrates the comparison between baseline models and Kaleido on two conditional generation tasks: the class-conditioned image generation and the text-to-image generation. In both tasks, Kaleido consistently produces more diverse images from identical condition (class or textual description) across varying CFG scales. For instance, in the task of class-to-image generation, the baseline diffusion models generate predominantly frontal views of a "husky" at high CFG, while Kaleido produces diverse images depicting huskies in various poses and numbers. A similar improvement in diversity is observed in the text-to-image generation as well, highlighting the robustness of Kaleido in generating diverse images under identical conditions.Control from Latent Tokens We show the efficacy of latent variables in guiding the image generation process in Fig. 6 . Fig. 6 demonstrates images generated with different types of latent variables: (a) textual descriptions, (b) object blobs, (c) detection bounding boxes, (d) visual tokens, and (e) combined latents, which integrate textual descriptions, detection bounding boxes and visual tokens. We visualize the generated latents tokens alongside the resulting images, showing how closely the images generated by Kaleido align with the latent tokens. Such alignment is evident in fine-grained visual information -such as object appearance, background, and atmosphere -, spatial location and orientation of different objects, and the stylistic elements of generated images. This alignment confirms that Kaleido can effectively interpret and utilize generated latent variables to guide and refine the image generation process.Latent Editing Fig. 8 showcases the impact of latent editing in image generation. The first row displays images generated using autoregressively produced latent tokens. In the second row, we demonstrate the effect of manual modifications to the textual descriptions: changing "log" to "cobblestones" and "a body of water" to "forest". These changes result in a modified image where a frog is now positioned on cobblestones with a forest background. Additionally, by further augmenting the bounding box of a cup to a different position, we observe that the cup's position in the image changes accordingly, while most other visual elements remain unchanged. The precise control of image characteristics via latent editing underscores Kaleido's flexibility and controllability, offering a powerful interactive interface for users to customize the generated images. Furthermore, the high fidelity of the re-generated images to their original versions indicates Kaleido's potential for applications requiring personalization or customizations. 
related work: 
 Augmenting Diffusion Models Various enhancements have been proposed to improve the versatility and controllability of diffusion models with augmented latents. Innovations such as Diffusion AE [Preechakul et al., 2022] integrates diffusion models with a learnable encoder that extracts high-level semantics and enables the diffusion model to add details directly in image space. Further efforts have focused on incorporating specific control signals, such as bounding boxes, layout, and segmentation masks to guide and control the image generation process. [Balaji et al., 2022 , Li et al., 2023 , Zheng et al., 2023a , Hu et al., 2023] . Recently, BlobGen [Nie et al., 2024] proposes to ground existing text-to-image diffusion models on object blobs -tilted ellipses that capture spatial details of the objects -for compositional generation. While these approaches improve the models' capacity to adhere to specified spatial layouts, they often necessitate modifications to the attention mechanism, potentially limiting their generality. In contrast, our method enhances the generative capabilities of diffusion models without altering the model architecture.Connecting Diffusion Models with LLMs The remarkable success of Large Language Models (LLMs) and diffusion models has spurred interest in connecting these models, aiming to leverage the capabilities of LLMs in understanding and generating complex data and combine it with the powerful image synthesis capabilities of diffusion models [Ge et al., 2023b , Zheng et al., 2023b , Sun et al., 2023] . Ge et al. [2023b] , Zheng et al. [2023b] propose image tokenizers that encodes images into visual tokens, enabling multimodal language modeling. This line of work focuses on empowering LLM with image generation ability by aligning its output embedding space with the pre-trained diffusion models. Our work leverages the LLMs' robust capabilities in textural understanding and generation to model the generation of abstract latents from the original text. These latents are then integrated with latent-augmented diffusion model, enabling a more interpretable and diverse image generation process.Our approach also distinguishes itself from the re-captioning method introduced in DALL-E 3 [Betker et al., 2023] . Unlike re-captioning, which typically replaces the original captions with more descriptive captions, our method retains the original condition and supplements it with latent variables of various forms (beyond textual captions like bbox, blob and "vokens"). The sampled latents serves as a unifying interface for various types of inputs, and introduce diversity compared to recaptioning where no sampling is involved at inference time. 
conclusion: 
 In this work, we address the challenge of improving sample diversity under high CFG in diffusion models. We introduce Kaleido Diffusion, which combines an autoregressive prior with a latentaugmented diffusion model. Results show Kaleido increases diversity without compromising quality, even at high CFG. With human interpretable latent tokens, Kaleido offers an explainable mechanism behind the image generation process and provides a fine-grained editable interface, enabling precise user control over the generated images. 
 
 Appendix: In this work, we explore four types of abstract latents, including textual descriptions (text), detection bounding boxes (bbox), object blobs (blob), and visual tokens (voken). Each type is designed to enrich the mode-to-image correspondence, covering different aspects of image formation. Examples of these abstract latents are illustrated in Fig. 4 . In the following paragraphs, we detail the methodology employed in constructing the training dataset for these abstract latents. Additionally, Fig. 9 outlines the pipeline for the step-by-step generation of these abstract latents.Textual descriptions Typical text-image datasets often provide captions that fail to fully capture the details of the image. For instance, as shown in Fig. 4 , the original caption "Dog laying on a human's lap" omits crucial details such as the presence of "laptop", which is essential for accurate imageOriginal caption: {} Using the information provided in the caption above, Please provide a detailed description of the image in 50-80 words, incorporating relevant information from the caption and expanding on the visual elements:-Include names, objects, events, and locations mentioned in the caption -Do not include placeholders like <PERSON> in the caption -Describe people, characters, animals, and notable entities -Mention the setting, background, and overall environment -Note colors, lighting, composition, and style aspects -Refer to any text, symbols, or logos in the image Combine caption details with your observations to create a comprehensive description of the key elements and overall scene, focusing on the most salient aspects of the image.Object labels: {} Using the provided object labels, generate a detailed description of the image in 50-80 words, incorporating the relevant information about prominent objects identified and expanding on the visual elements:-Describe each labeled object, including its size, shape, and placement in the scene -Depict relations between objects, such as proximity or arrangement -Highlight interactions or functions implied by the objects -Describe the surrounding environment or context that complements the labeled objects -Any notable features or characteristics of the objects, like color, texture, or design elements -Describe people, characters, animals, and notable entities -Mention the setting, background, and overall environment -Note colors, lighting, composition, and style aspects -Refer to any text, symbols, or logos in the image Craft a comprehensive depiction of the scene based on the identified objects, utilizing both the labels and contextual observations to enrich the description.Generate the caption in English with grounding: generation. To address this, we employ detailed textual descriptions as latent variables. These textual descriptions supplement the original captions by providing additional information that might be missing from the original captions. Specifically, we leverage Qwen-VL-Chat [Bai et al., 2023] , a large visual language model, designed for effective instruction-following across a variety of multimodal tasks. We instruct Qwen-VL-Chat to produce a detailed textual description given the original caption and corresponding image. The specific instructions used for generating the textual descriptions are detailed in Table 2 under the section Textual descriptions (caption). Fig. 4 shows an example of the generated detailed textual description that provides a more comprehensive depiction of the scenes than the original captions, thus allowing for a richer image generation.Additionally, for the ImageNet dataset, which consists of label-image pairs for class-to-image generation, we instruct Qwen-VL-Chat to generate detailed descriptions based on the class label and corresponding image. The instruction for this procedure is similarly documented in Table 2 under the section Textual descriptions (label).The spatial location of objects within an image is also crucial information for accurate representation of the image, yet such information is typically absent in textual descriptions. To incorporate this spatial information into the image generation process, we use detection bounding boxes as one type of abstract latents. Specifically, we use Qwen-VL [Bai et al., 2023] to prompt the model to "Generate the caption in English with grounding:". This approach results in captions where the spatial locations of objects are explicitly annotated within the text. For instance, as shown in Fig. 4 , the caption with grounding for this example is: "Dog (1, 33, 995, 995) resting head on owner's lap (1, 630, 785, 998) while they work on a laptop (39, 336, 999, 972) ." Each bounding box is described in a string format "x 1 , y 1 , x 2 , y 2 ", where x 1 , y 1 and x 2 , y 2 are the coordinates of the top-left and bottom-right corner, respectively. All the coordinates are normalized to a [0, 1000] range. The coordinates string is treated as part of the text, obviating the need for an additional positional vocabulary.Object Blobs Inspired by Nie et al. [2024] , we utilize object blobs as the abstract latents that contain more advanced spatial information. An object blob is defined as a tilted ellipse that specifies the position, size, and orientation of an object within an image. Specifically, a blob is represented as "(x c , y c , r major , r minor , θ)" where (x c , y c ) denotes the center point of the ellipse, r major and r minor are the radii of its semi-major and semi-minor axes, respectively, and θ ∈ [0, 180) denotes the orientation angle of the ellipse. To extract the blobs for meaningful objects, we leverage the results from bounding box detection and employ SAM [Kirillov et al., 2023] to generate the segmentation maps using the bounding boxes as prompts. Subsequently, an ellipse fitting algorithm is applied to these segmentation maps to determine the blob parameters for each identified object. This method allows for a more precise representation of objects' spatial characteristics, thus improving the integration of spatial and structural information within the image generation process.Visual Tokens Representing images via discrete visual tokens, especially using technologies like Vector Quantized Variational Autoencoder (VQ-VAE) [Van Den Oord et al., 2017] , has become a prevalent technique in generative modeling due to its ability to encode high-dimensional image data into a more manageable, discrete space. In this work, we utilize SEED [Ge et al., 2023b] , a VQ-based image tokenizer, to encode an image into a sequence of abstract discrete image tokens. These tokens encapsulate high-level semantic information of the visual elements in the image, serving as potent latent variables for guiding the diffusion model. The visual tokens are concatenated with the delimiter "#", forming a sequence of visual tokens represented as "I 1 #I 2 #...#I 32 ", where each "I i " denotes the image token id.In this paper, we use the following NestedUNet architecture proposed in Gu et al. [2023] to implement the denoising model. The total number of parameters is about 500M. For the autoregressive prior, we employ T5-XL [Raffel et al., 2020] for all experiments regardless of the input latent types. Both the denoiser and T5 decoder receive gradients and are trained end-to-end. For all experiments, we share all the following training parameters for both the baseline model and the proposed Kaleido Diffusion.default training config : batch_size =512 num_updates =400 _000 optimizer = ' adam ' adam_beta1 =0.9 adam_beta2 =0.99 adam_eps =1. e -8 learning_rate =1 e -4 l e a r n i n g _ r a t e _ w a r m u p _ s t e p s =10 _000 weight_decay =0.0 gradient_clip_norm =2.0 ema_decay =0.9999 m ix e d _p r e ci s i on _ t ra i n i ng = bp16All experiments are performed on 64 A100 GPUs which takes roughly 2 weeks for training 400k steps for both ImageNet and CC12M datasets. For text-to-image models, we perform an additional 400k steps progressive training at 64 × 64 resolution, while we train the entire model from scratch directly at 256 × 256 for ImageNet. Due to the memory cost of the T5-decoder, we can only fit 4 ∼ 8 images per GPU, causing at least ×3 slower training compared to the original MDM models.To demonstrate the effectiveness of various latents, we train our model with 5 types including text, bbox, blob, voken, and combined for text-to-image generation. For combined setting, we use the autoregressive model to predictcombined = text | bbox | vokenin a sequential way such that the latter latents will be controlled by earlier latents. We also trained models on ImageNet using combined latents for quantitative comparison.Training Complexity: The enhanced diffusion model may require more complex and extended training processes compared to standard models. This could lead to increased computational costs and longer development times, potentially limiting accessibility for smaller organizations or individual researchers.Difficulty in Finding Optimal Latents: Identifying the most effective latent variables to achieve the desired output diversity can be challenging. This process might involve extensive experimentation and fine-tuning, which can be time-consuming and resource-intensive. Additionally, covering a broader range of modes, such as depth and semantic maps, adds another layer of complexity to the model development, requiring sophisticated techniques to integrate these diverse forms of data effectively.Memory Usage: The improved diffusion model, with its increased output diversity, might demand higher memory usage due to the integration of the heavy language models. However, potential strategies such as partial training or joint training with LLMs could be explored to mitigate this issue. These methods could help distribute the computational load more effectively and reduce the memory footprint during the training process.The proposed method to enhance diffusion models and increase output diversity has significant social implications. By advancing the diversity and accuracy of generated outputs, this technology can be leveraged in various fields such as art, media, and content creation, providing more inclusive and representative outputs that reflect a broader spectrum of human experiences and creativity. Moreover, in areas like healthcare and education, diverse and precise models can lead to more personalized and effective solutions, addressing the unique needs of individuals and communities. This innovation also promotes ethical AI practices by reducing biases in model outputs, fostering a more equitable digital landscape. Ultimately, the enhanced diffusion models will contribute to the democratization of AI, making sophisticated tools accessible to a wider range of users and applications, thereby driving societal progress and innovation.We have included an alternative approach that constructs latent tokens without relying on additional knowledge. Specifically, we train a model using color clusters as latent tokens. For each color channel (R, G, and B) within the range of 0 -255, we equally segment it into eight clusters, resulting in a total of 8 × 8 × 8 = 512 color clusters. Given an image, we resize it to 4x4 pixels and assign a color cluster ID to each pixel based on its RGB value. The image is then encoded into a sequence of color cluster IDs (e.g., "C 1 #C 2 #...#C 512 "), with each C i representing a color cluster ID. This sequence serves as the condition for training our Kaleido diffusion.In Fig. 10 , we showcase images generated using color clusters as latent tokens on ImageNet. Our results demonstrate that, compared to the baseline MDM, our Kaleido diffusion can generate more diverse images with latent tokens derived purely from color clustering. This highlights that Kaleido diffusion's capability to generate diverse images is independent of distilled external knowledge, confirming that our approach can produce varied images without the aid of any other pre-trained models. We show additional results randomly sampled from our models. For all results including the baseline model, we use DDPM sampling with 250 steps. Segment-Anything (SAM) (1, 33, 995, 995) Fit Ellipse Dog (1, 33, 995, 995) resting head on owner's lap (1, 630, 785, 998) while they work on a laptop (39, 336, 999, 972) (1, 630, 785, 998) (39, 336, 999, 972) Segment-Anything (SAM) Segment-Anything (SAM)Fit Ellipse Fit Ellipse Dog (120, 164, 112, 130, 84) resting head on owner's lap (384, 881, 138, 486, 92) while they work on a laptop (562, 632, 238, 468, 78) Visual Tokens 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We made the main claim of "improving diffusion models with autoregressive latent modeling" in both the abstract and introduction. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.the class conditioning setting can be considered as a special case of text conditioning. 


Decision: This paper has received consistent feedback from four reviewers, all of whom are in favor of acceptance. The reviewers engaged in thorough discussion and rebuttal, ultimately reaching a consensus. This paper proposes Kaleido Diffusion for text-to-image generation. The method enhances the diversity of generated images, and the resulting latent variables are interpretable. Therefore,  the AC  has decided to accept this paper.

Reviews: review: paper_summary: paper_summary: In this paper, the authors propose a principled pipeline called Kaleido Diffusion for text-to-image generation with better mode coverage and diversity. The main intuition is that, conventional DM requires large CFG to make samples locate at high-likelihood modes, which would somehow constrain the diversity by only generating those limited modes. To solve this problem, the authors propose to first use AR to capture the distribution of such latent tokens, then use diffusion model to take them as extra conditions for generation. Experiments on both quantitative numeraical results and qualitative visualizations are used to evaluate the effectiveness of this work.
strength_weakness: strength_weakness: 1) I still don't well comprehend why choosing AR as the first-stage distribution learner. One of the reason I guess here is that, the CFG in AR model is only modulating the predicted logits, so that we can still sample across multiple reasonable modes for the next-stage DM through temperature parameter or top-k/top-p sampling. However, I'm still a little bit confused by this part: what if we use a diffusion model w/o CFG sampling to learn such an intermediate distribution? Is the model architecture design choice mainly for the fact that the authors use detailed textual descriptions as the latent token, so that it's natural to choose AR? If so, then this is a little bit ad-hoc to me.

2) Choosing these four specific types of latent tokens is also kind of ad-hoc to me. For example, somebody may also say the excessive CFG could result in similar artistic styles in generation. In this way, we have to additionally extract the style token as ground truth, then laboriously train the whole pipeline again to enable it. Another drawback of such design is that: the training cost is large and it's hard to scale up, every time we want to solve the mode collapse of new categories, we have to first extract/tokenize that new aspect, then jointly train the diffusion model and AR part, which is too ad-hoc and hard to scale up to me.

3) I hope the authors could also discuss about the diffusion decoder idea and the difference between that and this work. Specifically, AR + diffusion decoder would similarly take the latent tokens predicted by AR and feed into DM decoder as conditions. The major difference is that their setting doesn't need the ground truth for those latent tokens, which means the latent tokens are implicitly learned via end2end training. My question is what's the benefit of Kaleido compared to their pipeline, and won't their pipeline better at scaling up?
questions: questions: They are mainly elaborated in the weakness section. Overall I really like the intuition of this work. I'd like to see the authors' responses on my questions.
limitations: limitations: The limitations are discussed in the appendix.
paper_summary: paper_summary: This paper improves the diversity of diffusion generation by incorporating autoregressive latent priors. It leverages the autoregressive model to generate specific discrete latent features, and then concat them with the original extracted text features to serve as the condition of diffusion model. Experiments show that the method can achieve high quality and diversity even with high CFG.
strength_weakness: strength_weakness: (1) The writing lacks training and inference details. During training, whether all the modes are used in each step or they are randomly selected? For the visual tokens, whether the latent features from SEED are directly used or the sequence IDs are re-embedded? During inference, the autoregressive model is responsible for generating the latent modes if I understand correctly. Then how different latent modes are generated in control especially the combination of them?

(2) Evaluation is poor. In the experiment section, only a figure is provided without providing numerical numbers. The main experiment setting is confused (class conditioned or text conditioned). The construction of the toy example is also not clearly stated.

(3) Comparison with related works is not sufficient, making the position of this work unclear. There are related works also trying to improve the generation diversity of diffusion models such as [a].

[a] CADS: Unleashing the diversity of diffusion models through condition-annealed sampling.
questions: questions: I hope the authors can answer my questions in Weakness section point by point. I am willing to raise my score if my concerns are well addressed.
limitations: limitations: See Weakness.
paper_summary: paper_summary: This paper introduced Kaleido Diffusion which leverages an autoregressive model to first model the latent mode and then generate latents based on the sampled mode. The proposed method is reasonable. The authors explain the insight from a classifier-free guidance perspective. Several experimental results can support the authors' claim.

In addition, the paper with the same contents is published as a workshop paper in ICML24.
strength_weakness: strength_weakness: - The quantitative results are very limited (only Figure 5). 
- What is the context extractor (MLLM) used in the experiments? Although the authors claimed the mode selection is the major contribution of this work, the additionally enrolled pseudo labels perhaps also contribute to the performance improvements. 
- I am not convinced that the proposed approach is closely connected to the CFG explanation. Several previous studies, eg, ControlNet (ICCV23) and MaskComp (ICML24) have proven that dense controls will improve image quality. The proposed method is more like distilling knowledge from the pre-trained MLLM to obtain more detailed semantic information about the original dataset.
- It would be better to discuss more implementation details to enhance reproducibility.
questions: questions: - The dataset section claimed the usage of ImageNet and CC12M. However, the results of ImageNet are missing in the entire paper.
- Which dataset was used for the results reported in Fig 5?
- Can you quantitatively measure the diversity?
limitations: limitations: Yes
paper_summary: paper_summary: Diffusion models generate high-quality images from text but often lack diversity, especially with high classifier-free guidance. Kaleido addresses this by using autoregressive latent priors, which generate diverse latent variables from captions. It enriches input conditions, resulting in more diverse outputs while maintaining quality. Experiments confirm Kaleido's effectiveness in increasing image diversity and adherence to guidance.
strength_weakness: strength_weakness: 1. The formulation in Section 3.1 seems somewhat unintuitive. From my perspective, both text descriptions and the additional autoregressive priors you construct are forms of condition signals for diffusion models. Therefore, I do not see a compelling reason to modify or complicate the original classifier-free guidance formulation. Why not simply regard your method as an extension of the condition signal for classifier-free guidance, while still adhering to the existing CFG formulation?
2. The quantitative results are somewhat limited, with MDM being the only baseline. More comparisons with state-of-the-art diffusion models are encouraged.
questions: questions: 1. What is the motivation behind introducing visual tokens in the autoregressive prior? From the overall narrative of your paper, it seems that the autoregressive prior is intended to provide more customized control for users. However, visual tokens are essentially uninterpretable, making them unusable for user control.
2. The generation process involves an autoregressive process at the beginning. How does this process affect the efficiency of your method?
limitations: limitations: More results on the efficiency and quantitative performance of your method are encouraged.

Comments: title: Discussion Period
Thank you!
Discussion Periods
Discussion Periods
Post-rebuttal response
Dual submission policy
Request for clarification for W3 "the diffusion decoder idea"
comment: Thank you for your valuable feedback! We hope our responses have clarified your concerns!
Thank you for your valuable feedback! We hope our responses have clarified your concerns, and it would be nice if you consider raising your score to reflect this. However, thank you so much again!
Thank you for your valuable feedback and for increasing your ratings. We completely agree that further simplifying the pipeline to enhance scalability is crucial. This will be a key focus in our future work as well. Stay tuned!
I thank authors for the detailed reply. Overall I agree with the choice of AR to model the first-stage mode distribution. The reason for choosing those four specific types of latent tokens is mainly for paper demonstration. So this rebuttal has addressed most of my concerns. The only point is that I still think training separate tokenizer + accompanied AR model + diffusion makes the pipeline too complex to me. Even though I understand the model size is reasonable w/o need for too much compute, we still have to train the whole set of these three components every time we have a new requirement. But indeed this paper is clearly above the bar of this venue. I hence increase my score and advocate for acceptance of this work.
Thanks for your clarification. I would like to keep my positive rating.
Thanks for the clarification and new results. I would like to keep my positive rating.
Dear Reviewer,

As the discussion period deadline nears, we would greatly appreciate it if you could review our rebuttal and share any further feedback. If there are still concerns, we would greatly appreciate a list of specific changes you would need to reconsider your rating.

Thank you for your time and consideration. 

Best regards,
Thank you and we are glad our rebuttal addressed your concerns. Let us know anything you still have in mind or unclear.

Best regards,
Dear Reviewer,

As the discussion period deadline nears, we would greatly appreciate it if you could review our rebuttal and share any further feedback. If there are still concerns, we would greatly appreciate a list of specific changes you would need to reconsider your rating.

Thank you for your time and consideration. 

Best regards,
Dear Reviewer,

As the discussion period deadline nears, we would greatly appreciate it if you could review our rebuttal and share any further feedback. 
If there are still concerns, we would greatly appreciate a list of specific changes you would need to reconsider your rating.

Thank you for your time and consideration.
Best regards,
Thank the authors for providing their rebuttal. Most of my concerns are addressed and I will increase my rating to 6.
We would like to direct the reviewer’s attention to the dual submission policy outlined for NeurIPS [1]. According to this policy, “Papers previously presented at workshops are permitted, so long as they did not appear in a conference proceedings (e.g., CVPRW proceedings), a journal or a book.” Presentations at ICML workshops are not considered as part of the archival proceedings of the ICML conference.

[1] https://neurips.cc/Conferences/2024/CallForPapers
Thank you for your prompt response! We will carefully compare our approaches with theirs in our rebuttal.
Dear Authors:

Thanks for letting me know about the un-clarified point. Yes, this is exactly what I mean about the "diffusion decoder idea", which would in general uses the output from first-stage model (e.g., AR in your setting) as input to a diffusion model for image/video generation. Hope this can make my concerns clear.

Best
Dear Reviewer,

We are currently in the process of drafting our rebuttal response and would greatly appreciate your clarification on a point mentioned in Weakness 3. Specifically, we are seeking clarification on the "diffusion decoder idea" referenced in your feedback. At present, we believe you might be referring to approaches like multimodal large language models (MLLM) that employ diffusion models as decoders for image generation, similar to the approach used in EMU [1] [2].

Could you please confirm if our understanding is correct? If not, we would be grateful if you could provide a reference to the specific paper or work related to the "diffusion decoder idea" so that we can address this point more accurately in our response.

Thank you very much for your assistance!

[1] EMU: Generative Pretraining in Multimodality

[2] Generative Multimodal Models are In-Context Learners