Precise Parameter Localization for Textual Generation in Diffusion Models

Authors: Łukasz Staniszewski, Bartosz Cywiński, Franziska Boenisch, Kamil Deja, Adam Dziedzic
Primary area: generative models
Keywords: diffusion models, text edition, LoRA, localization, SD-XL, SD3, DeepFloyd IF

Abstract:
Novel diffusion models can synthesize photo-realistic images with integrated high-quality text. Surprisingly, we demonstrate through 
attention activation patching that only less than $1$\% of diffusion models' parameters, all contained in attention layers, influence the generation of textual content within the images.  Building on this observation, we improve textual generation efficiency and performance by targeting cross and joint attention layers of diffusion models.  We introduce several applications that benefit from localizing the layers responsible for textual content generation. We first show that a LoRA-based fine-tuning solely of the localized layers enhances, even more, the general text-generation capabilities of large diffusion models while preserving the quality and diversity of the diffusion models' generations. Then, we demonstrate how we can use the localized layers to edit textual content in generated images. Finally, we extend this idea to the practical use case of preventing the generation of toxic text in a cost-free manner. In contrast to prior work, our localization approach is broadly applicable across various diffusion model architectures, including U-Net (e.g., SDXL and DeepFloyd IF) and transformer-based (e.g., Stable Diffusion 3), utilizing diverse text encoders (e.g., from CLIP to the large language models like T5). Project page available at https://t2i-text-loc.github.io/.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730143957654

thread_type:
official_review

<thread_object>
creator: reviewer_mh7c

<content>
summary:
The paper introduces a method for precisely localizing the parameters within diffusion models (DMs) responsible for generating textual content in images. The key findings reveal that less than 1% of parameters, specifically in cross-attention and joint-attention layers, influence text generation. This localization approach, applicable across various DM architectures, including U-Net and transformer-based models, enables several enhancements:
1. Efficient finetuning: By targeting only the localized attention layers, fine-tuning improves text generation capabilities while preserving image quality and diversity.
2. Text editing: The method allows for precise text editing within generated images without altering other visual attributes.
3. Toxic content prevention: The localized layer approach provides a cost-free way to prevent generating harmful text in images.
The authors validate their method across different models, demonstrating improvements in text alignment and visual consistency while showcasing its adaptability and efficiency across architectures.

strengths:
The paper is innovative in isolating specific cross and joint attention layers that directly influence text generation in diffusion models. This novel approach to localizing model parameters is broadly applicable across various architectures, including U-Net and transformer-based models, making it an impactful step forward in the interpretability and control of DMs. Furthermore, applying this localization technique to prevent toxic text generation is a unique and practical application, expanding the scope of ethical safeguards in generative models.
This research has notable significance for both practical applications and theoretical advancements in generative AI. By pinpointing and leveraging specific layers responsible for text content, the paper introduces a method that improves the efficiency and quality of DMs in text generation tasks. The application to toxic text prevention adds a meaningful dimension, contributing to the broader conversation on safe and ethical AI. The findings open up new avenues for fine-tuned, targeted adjustments in diffusion models, making it easier for practitioners to achieve high-quality text rendering and control content generation more responsibly.

weaknesses:
The method for identifying key cross and joint attention layers is central to the paper’s contribution, yet the description of this process could benefit from greater specificity. For instance, while the authors mention leveraging activation patching techniques, the exact steps involved in isolating specific layers and evaluating their effect on text generation are not fully detailed. Providing pseudocode or a workflow diagram would make the process clearer and more reproducible.
The approach is tested on models using CLIP and T5 encoders, yet diffusion models incorporating different or proprietary text encoders could present challenges to the proposed localization technique. Extending the evaluations to other popular encoders (or discussing limitations for certain encoder types) would improve the generalizability of the method.
While LoRA fine-tuning on localized layers is shown to enhance text generation quality, the paper does not deeply analyze potential trade-offs involved. Specifically, fine-tuning localized layers only may limit generalization to unseen prompts or introduce subtle biases in text rendering quality across different styles or languages. Including a comparative analysis of localized vs. whole-model fine-tuning would be helpful, especially with examples illustrating where one approach might be preferred over the other.
While the method is shown to be effective on certain models, scaling the approach for larger models or more computationally intensive tasks may present challenges. A more detailed discussion on these potential limitations would provide a balanced view and help guide future research in optimizing model efficiency.

questions:
Could the authors provide more details on the method for identifying the critical cross and joint attention layers? Specifically, is the layer selection based solely on empirical observations, or are there theoretical underpinnings guiding this choice?

Have the authors tested whether fine-tuning only the localized layers with LoRA introduces any long-term risks, such as overfitting or loss of generalization? Are there any benchmarks or datasets where this approach might underperform relative to fine-tuning all cross-attention layers?

How does the method handle diffusion models that utilize text encoders other than CLIP and T5, or models with proprietary modifications? Does the localization process adapt effectively to these scenarios, or are there specific challenges?

Why were SimpleBench and CreativeBench chosen as the primary benchmarks? Would the method be equally effective on more complex, real-world benchmarks that include diverse text styles, such as cursive or stylized fonts?

Could the approach for localizing parameters influencing text generation extend to other applications? The paper would benefit from a discussion of possible applications beyond text generation and toxicity prevention, potentially sparking ideas for future work or new applications of the method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730443879670

thread_type:
official_review

<thread_object>
creator: reviewer_9sea

<content>
summary:
This paper explores to find a small subset of attention layers in DMs which determines text generated in images. It finds only a few layers related to this. Based on this observation, this paper further introduces a new fine-tuning strategy and a text edit technique. It can be used to prevent the generation of harmful or toxic text in images. Extensive experiments and analysis are conducted to demonstrate the observation and effectiveness of the proposed method.

strengths:
The paper is well-structured, with a logical progression from the introduction to the conclusions. The writing is clear and concise, allowing complex exploration to be conveyed effectively. Key terms are introduced and defined appropriately, enhancing readability and ensuring that readers can follow the arguments easily. This observation is interesting and can effectively help us understand the inherent properties of the DMs in textual generation.

weaknesses:
1、	The generalizability of the observations proposed in the paper needs further exploration. For instance, it remains to be determined whether this property holds consistently across texts of varying lengths and images with different aspect ratios.

2、	Different attention layers tend to focus on varying aspects, such as content or style. This has already been explored in previous works. This paper applies this feature to textual generation, yet more clarification is needed to highlight the novelty of this contribution.

3、	From some visualizations, it appears that the proposed approach not only alters the content of the text but also changes its style, such as the font. Are both style and content localized within the same layer?

questions:
Hope the author can answer the questions in the weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730612485140

thread_type:
official_review

<thread_object>
creator: reviewer_axtp

<content>
summary:
This paper made an interesting observation on how visual text specified for text-to-image generation is realized through the cross-attention layers - there are only a few cross-attention layers responsible for processing the information about the text information. This is identified by swapping the key and values to be generated across every layer. With such observation, the authors perform several experiments showing the effectiveness of using these layers for different visual text applications.

strengths:
[Originality]

The finding about the function of cross-attention layers in generating visual text is quite novel and interesting.

[Quality]

The authors did rigorous experiments to show that only a few layers are responsible for generating visual texts with multiple experiment setups. Additionally, the authors show that these layers only focus on these visual texts instead of generating other content described by the text prompts. The applications of improving text generation, editing, and preventing the generation of toxic text also demonstrate the usefulness of the findings. 


[Clarity]

The paper did a very good job of presenting the analysis and applications.

[Significance]

I personally find that this is an intriguing finding about the interpretation of the text-to-image diffusion model. The paper paves a way towards understanding how text embeddings are used through attention layers to control the text-to-image model. The finding also shows promise for several downstream tasks.

weaknesses:
I did not find significant weaknesses in the paper. However, the authors could consider the following feedback to improve the readiness and clarity of the paper:

1. Although I'm fairly familiar with the architectures of SD3 and SDXL, I still need to guess how localizing by patching is different from injection. So, for SD3, are only the keys and values corresponding to text embeddings swapped by the target prompt caching? 

2. There are many papers in text-to-image generation/editing working on cross-attention layers and manipulating the denoising steps through keys and values. The authors should consider doing more comprehensive survey and here are some of them:

[1] Kumari, Nupur, et al. "Multi-concept customization of text-to-image diffusion." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[2] Geyer, Michal, et al. "TokenFlow: Consistent Diffusion Features for Consistent Video Editing." The Twelfth International Conference on Learning Representations.

[3] Patashnik, Or, et al. "Localizing object-level shape variations with text-to-image diffusion models." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[4] Phung, Quynr, et al. "Grounded text-to-image synthesis with attention refocusing." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[5] Tumanyan, Narek, et al. "Plug-and-play diffusion features for text-driven image-to-image translation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

questions:
Minor questions and suggestions:

1. Is there any specific reason to classify the target text into different buckets and use the most frequent texts as the validation set? It seems other buckets are not used.

2. Lines 144 and 145 look like new paragraph names.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730623899253

thread_type:
official_review

<thread_object>
creator: reviewer_tspc

<content>
summary:
This paper proposes to localize the attention layers which influence the text-rendering ability of tex-to-image generation models. With the localized layers, the authors propose to utilize them in many tasks including generation and editing.

strengths:
Distribution of image samples with text may be quite different from image samples without text. As a result, we may face a trade-off between image quality and text-rendering ability in practice when training text-to-image generation model.
Thus the idea of localizing corresponding layers for text rendering ability in text-to-image generation models is interesting. 
If we can localize such layers, then we can use carefully designed fine-tuning so that the resulting model performs well in both image quality and text accuracy.



From the experiments, we can see that the localized parameters indeed influence the text-rendering ability.

weaknesses:
The fine-tuning experiment is conducted on a small subset of MARIO-10M dataset. So it is expected that fine-tuning the whole model may lead to overfitting. The experiment results can show fine-tuning localized layers indeed works, but it can not show that fine-tuning only localized layers is better than fine-tuning the whole model.
To illustrate the effects of the localized layers, it is suggested to conduct the experiment on large-scale dataset, i.e. fine-tuning the localized layers or whole model on all samples from MARIO-10M dataset.

Another case which is more related to the fine-tuning experiment with smaller dataset is fine-tuning the model on complicated samples of high-quality. Different from MARIO-10M which may contain some simple and low-quality samples, high-quality samples with long sentences rendered might be difficult to collect. It can strength the contribution of the proposed method if it can improve model performance under this setting. 


The author mentioned "We focus on LoRA for SDXL since this model has a significantly lower text generation
quality compared to other studied DMs". Can fine-tuned SDXL outperform SOTA methods on generation task with commonly used benchmark? If it can not, why not fine-tuning one of the SOTA models to show that the proposed method also works and can be used to improve not only SDXL? Only fine-tuning a model with lower performance is not convincing.

More qualitative results (especially with more words to be generated) and human evaluation are suggested.

questions:
Since MARIO-10M dataset is utilized in training, why is MARIO-Eval Benchmark not directly used to evaluate the generation performance in the paper? 

What is the reason that the authors only use a small subset of Mario-10M dataset?

How is the editing performance compared to other methods like textDiffuser?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730636536558

thread_type:
official_review

<thread_object>
creator: reviewer_rlch

<content>
summary:
This paper discovered that only a small amount of parameters within modern diffusion models are in charge of textual content generation. This phenomenon is widely observed among U-Net-based diffusion models like SDXL and transformer-based ones like SD3. Based on this observation, this paper developed three possible application scenarios, including LoRA-based fine-tuning for enhanced text generation, textual content editing for images generated by diffusion models, and the prevention of toxic text generation.

strengths:
- By careful experimental analysis, this paper localized a small subset of cross and joint attention layers in diffusion models that are responsible for textual content generation.
- Based on this observation, this paper developed an effective fine-tuning strategy for enhanced text generation while maintaining the model’s overall generation performance.

weaknesses:
For the application of preventing the generation of toxic text within images, this paper claimed that their approach is able to remove the toxic text in the text prompt from the generated images. They achieved this target by first detecting the toxic text in the text prompt using state-of-the-art large language models, and then replacing the toxic text with non-harmful text for image generation using their approach. However, it seems that directly replacing the toxic text with non-harmful text and feeding the modified text prompt to the diffusion models can achieve the same target. Since the image generated by the source prompt won’t be provided to the users, there seems no need to prevent other visual content from being influenced when changing the source prompt (with toxic text) to the target prompt (without toxic text).

questions:
Request the authors clarify the necessity of using their approach to prevent toxic textual content generation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732722622510

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response by Authors

comment:
We would like to express our sincere gratitude to all the reviewers for their valuable feedback and insightful suggestions, which have enabled us to improve our work further. We are pleased to see that most reviewers found the findings in our submission novel and interesting and are inclined to accept it.

Reviewers found our *"idea of localizing corresponding layers for text rendering ability in text-to-image generation models"* interesting (Reviewers TSpC and 9sEa), innovative (Reviewer mH7c), novel, and intriguing (Reviewer aXtp). This unique work *"paves a way towards understanding how text embeddings are used through attention layers to control the text-to-image model and shows promise for several downstream tasks*" (Reviewer aXtp). *"The method is broadly applicable across various architectures, making it an impactful step forward in the interpretability and control of DMs"* (Reviewer mH7c). The experiments were carried out carefully (Reviewer RLCH), rigorously (Reviewer aXtp), and extensively (Reviewer 9sEa) to show that only a few layers are responsible for generating the visual text. The paper was clearly written and well-structured (Reviewer 9sEa), and *"did a very good job of presenting the analysis and applications"* (Reviewer aXtp).


Based on the reviewers' comments, we conducted the following experiments to improve our submission:

- **Toxic text removal.** We conducted an experiment demonstrating that our text-editing method effectively removes toxic content from images while preserving user-specified emotions in the prompt. This makes our approach more practical than simple baselines, such as replacing toxic words in the prompt and passing them to all the layers. The results are detailed in **Sec. A.6** of the Appendix.
  
- **Fine-tuning experiments.** We extended our fine-tuning experiments to varying training set sizes. We show that fine-tuning all cross-attention layers underperforms compared to fine-tuning only localized layers, even with a significant increase in training set size. Furthermore, fine-tuning localized layers achieves consistent performance across different training set sizes, demonstrating robustness. These results are presented in **Sec. A.8** of the Appendix.
  
- **Multi-word text editing:** We evaluated our method on images with texts containing multiple words. Although text alignment slightly decreases as the number of words increases, our method successfully edits text while preserving visual attributes. Results and examples are provided in **Sec. A.11** of the Appendix.
  
- **Localization of text style.** We conducted a study showing that control of text style, contrary to its content, is distributed across multiple cross-attention layers in diffusion models. The results of this study are presented in **Sec. A.9** of the Appendix.

On top of that, we improved the clarity of the paper, addressing some reviewers' concerns. In Section A.9 of the Appendix, we included **pseudocode for localizing the subset of cross-attention layers that control visual text**. In Section A.1  of the Appendix, we reviewed **related work** on manipulating diffusion models with cross-attention layers, indicating how those works differ from our approach. Additionally, we improved Section 4.1 to clarify details of the **patching technique** and how it differs from the injection.

We strongly believe that the revised version of our work presents a compelling contribution, offering valuable insights for the community. It advances the understanding of large-scale neural architectures and supports the development of novel methods for fine-tuning, image editing, and mitigating toxic content in generative models.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper identifies a small subset of attention layers in diffusion models that drives textual generation. Based on this finding, the authors explore efficient fine-tuning to enhance text generation, text editing, and toxic content prevention. Experiments demonstrate improved text alignment and visual consistency.

After the discussion between the authors and reviewers, 4 out of 5 reviewers leaned toward acceptance. One reviewer expressed concerns about the performance of this work on a downstream task, specifically preventing toxic content generation. However, the strengths of this submission outweigh its limitations, making it worthy of publication at ICLR. The authors are encouraged to provide additional experimental results on preventing toxic content generation.