Paper: Kosmos-G: Generating Images in Context with Multimodal Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
I don't see much novelty from AlignerNet. Compared with GlueNet, AlignerNet merely replaces the MLP with encoder-decoder Transformers but they have the same loss and the same domains (both aligning text embedding). AlignerNet is useful from the experiments, but not novel IMO.

Author response:
**A1:** Thanks for your comments that "The idea of bridging MLLM and SD for versatile image generation is interesting. MLLM naturally can accept both image and text input, which can provide a more diverse signal to the image generation module and therefore enable new applications." We acknowledge that AlignerNet may not be novel, it is adopted from GlueNet [1] for alignment and modified to deal with variable-length outputs. However, we would like to emphasize that for Kosmos-G, the key idea is not about how to align MLLM with SD. As shown in Figure 6, the rough semantic alignment is far from achieving our objective of "Image as a foreign language in image generation". Our main contribution is the idea of leveraging the advanced multimodal perception of MLLM for subject-driven generation. Through instruction tuning the MLLM, we can approach images as a foreign language when prompting Stable Diffusion. We also include additional examples with 3 to 4 image inputs and diverse text interleaving in our revised Appendix D and [here (https://i.ibb.co/BstzDqL/challenge.jpg)](https://i.ibb.co/BstzDqL/challenge.jpg), which shows more about this capability.

Reviewer comment (Experiments):
Since the training data includes the image editing dataset from InstructPix2Pix, a comparison between previous works on image editing benchmarks should also be conducted. Similarly, how is kosmos-g compared with GILL in visual storytelling?

Author response:
**A2:** We utilize the data from InstructPix2Pix to enable attribute and accessory editing, yet it wouldn't be appropriate to compare with InstructPix2Pix on image editing task. Because we are focusing on subject-driven generation, which does not require preserving the layout of the original input, unlike typical image editing task that requires layout consistency. They are completely two tasks for different purposes. For visual storytelling, we argue that Kosmos-G and GILL [2] are instruction fine-tuned for different purposes. Kosmos-G is fine-tuned on a compositional generation dataset for subject-driven generation, while GILL is fine-tuned on Conceptual Captions (CC3M). Considering the composition nature of our data, our model will not be suitable for visual storytelling, while GILL also struggles with subject-driven generation and composition. However, given the aligned multimodal perception, our method possesses the potential to deal with visual storytelling when trained with interleaved images and captions following GILL.

Reviewer comment (Experiments):
In AlignerNet, both MSE and REC losses are used. However, no ablation is done about those two losses.

Author response:
**A3:** As shown below or in the revised Appendix C, we present a comprehensive set of ablation studies on the COCO-30k, focusing on the model design and supervision methods used for AlignerNet. We observed that using MSE loss alone leads to slightly worse results, which is consistent with the observation from GlueGen [1].

### Table 1: FID Score $\downarrow$ on COCO-30k with a guidance scale of 3.0: measuring the fidelity of generated images

| Methods                 | MSE + REC | MSE    | Diff   |
| ----------------------- | --------- | ------ | ------ |
| Linear                  | Failed    | Failed | Failed |
| MLP                     | Failed    | Failed | Failed |
| Perceiver Resampler [2] | Failed    | Failed | Failed |
| 12-Layers Decoder       | Failed    | -      | -      |
| 12-Layers AlignerNet    | 9.89      | 10.23  | 11.30  |
| 24-Layers AlignerNet    | 9.55      | -      | -      |

### Table 2: CLIP Score $\uparrow$ on COCO-30k with a guidance scale of 3.0: measuring the alignment between text input and image output

| Methods                 | MSE + REC | MSE    | Diff   |
| ----------------------- | --------- | ------ | ------ |
| Linear                  | Failed    | Failed | Failed |
| MLP                     | Failed    | Failed | Failed |
| Perceiver Resampler [3] | Failed    | Failed | Failed |
| 12-Layers Decoder       | Failed    | -      | -      |
| 12-Layers AlignerNet    | 25.48     | 25.31  | 24.22  |
| 24-Layers AlignerNet    | 25.57     | -      | -      |

Reviewer comment (Evaluation):
In Tab2, it seems the E2E fine-tuning fails. However, recent works such as BLIP-Diffusion, EMU, and MGIE can successfully connect MLLM with SD via E2E fine-tuning without any specific alignment. Why Kosmos-G's behavior is different from others and relies on additional alignment?

Author response:
**A4:** Thanks a lot for this important inquiry, we will clarify it here. E2E Fine-tuning in the paper means we directly connect MLLM with SD and e2e fine-tuning the MLLM. As shown in the above table, direct e2e fine-tuning is feasible with the help of AlignerNet. While it is more costly because we need to pass it through U-Net, leading to a worse performance under the same GPU days. So in practice, we choose to add a specific alignment to accelerate training.

Reviewer comment (Experiments):
When constructing the compositional generation dataset, what if multiple objects of the same class exist in the same image? Would the corresponding segmentation mask cover multiple instances in the same mask?

Author response:
**A5:** Thank you for this insightful question. There do exist cases that multiple instances of the same object class appear in a single image. We simply choose the largest connected sub-region. This approach is intuitive as it picks the clearest, unobstructed object for training. It is also consistent with the inference usage, as we represent each object in the prompt using a single image.

Reviewer comment (Experiments):
The ablation study seems not very comprehensive; for example, if the goal is to align the two representation spaces, there should be other options to achieve the alignment design.

Author response:
Dear Reviewer VSX4,

Thanks a lot for your valuable feedback and advice.

**Q:** More justification and ablation study for AlignerNet

**A:** In the revised Appendix C, we have expanded the ablation study for AlignerNet on the COCO-30k dataset, focusing on its model design and supervision. We investigated various configurations, including the use of MSE loss, reconstruction loss, and diffusion loss in supervising the aligned network, while keeping other modules such as the MLLM frozen. 

The results demonstrate that simple networks such as Linear, MLP, and Resampler fail to achieve satisfactory alignments. Moreover, due to MLLM's variable-length output, MLP-based GlueNet [1] and encoder-only Transformer models are unsuitable for learning alignment. Decoder-only architecture also fails due to its inefficiency in modeling dimension transformations.

For loss functions, using only MSE loss results in worse results, which is consistent with findings from GlueGen [1]. Direct end-to-end training is more costly because we need to pass it through U-Net, leading to worse performance under the same GPU days. 

Notably, AlignerNet surpasses the performance of GlueNet+T5. For instance, GlueNet + T5 attained CLIP scores of 20.67, 23.24, and 23.74 on the COCO-5k dataset at guidance scales of 1.5, 5, and 7.5 respectively, and a FID score of 14.32 at a guidance scale of 3.0 on COCO-30k. These results highlight the effectiveness of AlignerNet in achieving superior alignment, validating the soundness of its design and supervision.

### Table 1: FID Score $\downarrow$ on COCO-30k with a guidance scale of 3.0: measuring the fidelity of generated images

| Methods               | MSE + REC | MSE   | Diff   |
| --------------------- | --------- | ------ | ------ |
| Linear                | Failed    | Failed | Failed |
| MLP                   | Failed    | Failed | Failed |
| Perceiver Resampler [2] | Failed    | Failed | Failed |
| 12-Layers Decoder     | Failed | -      | -      |
| 12-Layers AlignerNet  | 9.89     | 10.23 | 11.30 |
| 24-Layers AlignerNet  | 9.55     | -      | -      |

### Table 2: CLIP Score $\uparrow$ on COCO-30k with a guidance scale of 3.0: measuring the alignment between text input and image output

| Methods               | MSE + REC | MSE   | Diff |
| --------------------- | --------- | ------ | ------ |
| Linear                | Failed    | Failed | Failed |
| MLP                   | Failed    | Failed | Failed |
| Perceiver Resampler [2] | Failed    | Failed | Failed |
| 12-Layers Decoder     | Failed | -      | -     |
| 12-Layers AlignerNet  | 25.48   | 25.31 | 24.22 |
| 24-Layers AlignerNet  | 25.57   | -      | -     |

## Reference

[1] Qin, C., Yu, N., Xing, C., Zhang, S., Chen, Z., Ermon, S., ... & Xu, R. (2023). GlueGen: Plug and Play Multi-modal Encoders for X-to-image Generation. *arXiv preprint arXiv:2303.10056*.

[2] Alayrac, J. B., Donahue, J., Luc, P., Miech, A., Barr, I., Hasson, Y., ... & Simonyan, K. (2022). Flamingo: a visual language model for few-shot learning. *Advances in Neural Information Processing Systems*, *35*, 23716-23736.

Reviewer comment (Experiments):
It is unclear why the current AlignerNet design is the best; maybe more justification and ablation study are needed here.

Author response:
Dear Reviewer VSX4,

Thanks a lot for your valuable feedback and advice.

**Q:** More justification and ablation study for AlignerNet

**A:** In the revised Appendix C, we have expanded the ablation study for AlignerNet on the COCO-30k dataset, focusing on its model design and supervision. We investigated various configurations, including the use of MSE loss, reconstruction loss, and diffusion loss in supervising the aligned network, while keeping other modules such as the MLLM frozen. 

The results demonstrate that simple networks such as Linear, MLP, and Resampler fail to achieve satisfactory alignments. Moreover, due to MLLM's variable-length output, MLP-based GlueNet [1] and encoder-only Transformer models are unsuitable for learning alignment. Decoder-only architecture also fails due to its inefficiency in modeling dimension transformations.

For loss functions, using only MSE loss results in worse results, which is consistent with findings from GlueGen [1]. Direct end-to-end training is more costly because we need to pass it through U-Net, leading to worse performance under the same GPU days. 

Notably, AlignerNet surpasses the performance of GlueNet+T5. For instance, GlueNet + T5 attained CLIP scores of 20.67, 23.24, and 23.74 on the COCO-5k dataset at guidance scales of 1.5, 5, and 7.5 respectively, and a FID score of 14.32 at a guidance scale of 3.0 on COCO-30k. These results highlight the effectiveness of AlignerNet in achieving superior alignment, validating the soundness of its design and supervision.

### Table 1: FID Score $\downarrow$ on COCO-30k with a guidance scale of 3.0: measuring the fidelity of generated images

| Methods               | MSE + REC | MSE   | Diff   |
| --------------------- | --------- | ------ | ------ |
| Linear                | Failed    | Failed | Failed |
| MLP                   | Failed    | Failed | Failed |
| Perceiver Resampler [2] | Failed    | Failed | Failed |
| 12-Layers Decoder     | Failed | -      | -      |
| 12-Layers AlignerNet  | 9.89     | 10.23 | 11.30 |
| 24-Layers AlignerNet  | 9.55     | -      | -      |

### Table 2: CLIP Score $\uparrow$ on COCO-30k with a guidance scale of 3.0: measuring the alignment between text input and image output

| Methods               | MSE + REC | MSE   | Diff |
| --------------------- | --------- | ------ | ------ |
| Linear                | Failed    | Failed | Failed |
| MLP                   | Failed    | Failed | Failed |
| Perceiver Resampler [2] | Failed    | Failed | Failed |
| 12-Layers Decoder     | Failed | -      | -     |
| 12-Layers AlignerNet  | 25.48   | 25.31 | 24.22 |
| 24-Layers AlignerNet  | 25.57   | -      | -     |

## Reference

[1] Qin, C., Yu, N., Xing, C., Zhang, S., Chen, Z., Ermon, S., ... & Xu, R. (2023). GlueGen: Plug and Play Multi-modal Encoders for X-to-image Generation. *arXiv preprint arXiv:2303.10056*.

[2] Alayrac, J. B., Donahue, J., Luc, P., Miech, A., Barr, I., Hasson, Y., ... & Simonyan, K. (2022). Flamingo: a visual language model for few-shot learning. *Advances in Neural Information Processing Systems*, *35*, 23716-23736.

Reviewer comment (Experiments):
The paper only demonstrates cases with a maximum of two images, failing to showcase examples with more than two images as inputs.

Author response:
Thank you for your valuable suggestion. We acknowledge the importance of showcasing more complex scenarios in demonstrating the capabilities of Kosmos-G. It's worth noting that subject-driven generation across multiple images, especially more than two, is challenging. This is true even for state-of-the-art fine-tuning methods like Break-a-Scene [1], which supports up to three subjects. Zero-shot methods like FastComposer [2] and Subject-Diffusion [3] can handle up to two images but struggle with attribute and accessory editing and also require attention map manipulation. Owning to our collected training data, it is possible to treat the image input as a text token and prompt Kosmos-G like Stable Diffusion. This unique feature allows for more diverse and complex multi-image and text interleaving scenarios. We have included additional examples with 3 to 4 image inputs and diverse text interleaving in our revised Appendix D and [here (https://i.ibb.co/BstzDqL/challenge.jpg)](https://i.ibb.co/BstzDqL/challenge.jpg). These examples demonstrate Kosmos-G's robust performance in handling these challenging cases.

Reviewer comment (Experiments):
The paper predominantly showcases and evaluates image-text-image input scenarios, leaving more diverse multi-image and text interleaving cases unexplored.

Author response:
Thank you for your valuable suggestion. We acknowledge the importance of showcasing more complex scenarios in demonstrating the capabilities of Kosmos-G. It's worth noting that subject-driven generation across multiple images, especially more than two, is challenging. This is true even for state-of-the-art fine-tuning methods like Break-a-Scene [1], which supports up to three subjects. Zero-shot methods like FastComposer [2] and Subject-Diffusion [3] can handle up to two images but struggle with attribute and accessory editing and also require attention map manipulation. Owning to our collected training data, it is possible to treat the image input as a text token and prompt Kosmos-G like Stable Diffusion. This unique feature allows for more diverse and complex multi-image and text interleaving scenarios. We have included additional examples with 3 to 4 image inputs and diverse text interleaving in our revised Appendix D and [here (https://i.ibb.co/BstzDqL/challenge.jpg)](https://i.ibb.co/BstzDqL/challenge.jpg). These examples demonstrate Kosmos-G's robust performance in handling these challenging cases.

Reviewer comment (Theory):
Section 2.3 discusses the "Score distillation instruction tuning" technique, but the description lacks clarity. The paper should provide a more precise definition of the entities involved in calculating the KL divergence, along with any specific mathematical formulas or equations for better understanding.

Author response:
Thank you for your question regarding the clarity of "Score Distillation Instruction Tuning". We would like to emphasize that during the instruction tuning stage, we still utilize the diffusion loss for model training. However, this loss can also be regarded as score distillation, and it is effectively equivalent to optimizing the diffusion loss as the SDS loss. We will elaborate on these in the following text.

Our objective is to distill the learned score function from the Stable Diffusion U-Net into Kosmos-G. This process enables Kosmos-G to encode image features into embeddings that the Stable Diffusion model can understand for subject-driven generation. Essentially, this approach is like pre-training a generalized textual inversion [4] model, with all conditions learnable by model-seeking.

Consider the Kosmos-G model, denoted as $\phi$, which takes an input $\mathbf{x}$ and produces an output $\mathcal{C}=\phi(\mathbf{x})$. Alongside this, we have the Diffusion U-Net, represented as $\theta$. In our process, we optimize the Kosmos-G model $\phi$ using the diffusion loss, while keeping the parameters of the Diffusion U-Net $\theta$ frozen. The diffusion loss is expressed as follows:

$$
\mathcal{L} _{diff}(\phi) = \mathbb{E} _{\mathbf{z} _0, \boldsymbol{\epsilon} \sim \mathcal{N}(0, 1), t} \Big[w(t)\|\boldsymbol{\epsilon} _\theta(\mathbf{z} _t; \mathcal{C}, t) - \boldsymbol{\epsilon}\|^2 \Big]
$$

consider the gradient of $\mathcal{L} _{diff}$:
$$
\nabla _\phi \mathcal{L} _{diff}(\phi) = \mathbb{E} _{\mathbf{z} _0, \boldsymbol{\epsilon} \sim \mathcal{N}(0, 1), t} \Bigg[w(t)
\underbrace{\left(\boldsymbol{\epsilon} _\theta(\mathbf{z} _t; \mathcal{C}, t)  - \boldsymbol{\epsilon} \right) \vphantom{\partial \mathcal{C} \over \partial \phi}} _{\text{Noise Residual}} 
\underbrace{{\partial \boldsymbol{\epsilon} _\phi(\mathbf{z} _t; \mathcal{C}, t) \over \mathcal{C}} \vphantom{\partial \mathcal{C} \over \partial \phi}} _{\text{U-Net Jacobian}}
\underbrace{{\partial \mathcal{C} \over \partial \phi}} _{\text{Kosmos-G Jacobian}}\Bigg]
$$

Following the approach in Dreamfusion [6], we can simplify this equation by omitting certain terms, leading to the SDS loss for Kosmos-G:
$$
\nabla _\phi \mathcal{L} _{SDS}(\phi) \triangleq \mathbb{E} _{\mathbf{z} _0, \boldsymbol{\epsilon} \sim \mathcal{N}(0, 1), t} \Bigg[w(t)\left(\boldsymbol{\epsilon} _\theta(\mathbf{z} _t; \mathcal{C}, t) - \boldsymbol{\epsilon} \right) {\partial \mathcal{C} \over \partial \phi}\Bigg]
$$

As established in Dreamfusion [6], optimizing the SDS loss is effectively equivalent to optimizing the diffusion loss when the U-Net $\theta$ is frozen. From the perspective of score distillation, when using diffusion loss, the KL divergence defined by conditions and the pre-learned score function is equivalently minimized for distilling learned probability density in conditional image synthesis [5, 7]:
$$
\mathop{\min}\limits _{\phi} \mathcal{L} _{Diff}(\phi) = \mathbb{E} _{\mathbf{z} _0, t, \mathcal{C}} \Big[D _{\rm{KL}}\big(q(\mathbf{z} _{t-1}|\mathbf{z} _t, \mathbf{z} _{0})~\| ~p _\theta(\mathbf{z} _{t-1}|\mathbf{z} _t; \mathcal{C})\big)\Big]
$$

Reviewer comment (Experiments):
The paper highlights the exceptional subject-driven generation capabilities of KOSMOS-G, particularly when not training the diffusion model. It is worth asking if the authors have explored the possibility of further enhancing subject-driven generation by training the diffusion model.

Author response:
Thank you for your valuable suggestion regarding the enhancement of subject-driven generation through training the diffusion model. We designed Kosmos-G to function as a minimal, plug-and-play module within Stable Diffusion, offering a more community-friendly approach compared to a fully modularized system. This design choice allows Kosmos-G to seamlessly replace CLIP, facilitating easier integration with other U-Net techniques, such as ControlNet and LoRA, as demonstrated in our paper. While we acknowledge that some research, like BLIP-Diffusion [8], involves training the diffusion model, our approach as outlined in Table 1 achieves superior performance without the need to train the U-Net. We agree the subject-driven generation capability can be improved by tuning more parameters. However, our current resources, specifically V100 GPUs with 32GB VRAM, constrain our ability to fully explore the outcomes of training Kosmos-G with an unlocked U-Net.