Characterizing the Roles of Attention Heads in Generative Models: Text-to-Image Diffusion Models and Large Language Models = 생성 모델의 어텐션 헤드 역할 규명: 텍스트-이미지 확산 모델과 거대 언어 모델

국문 초록 (Abstract)

최근 딥러닝의 발전은 모델의 동작 이면에 존재하는 내부 메커니즘을 이해하고자 하는 작동 원리 기반 해석 가능성(Mechanistic Interpretability) 분야에 대한 관심을 높이고 있다. 이 분야에서 어텐션 헤드 기여도 분석(Attention Head Attribution), 즉 딥러닝 모델 내 개별 어텐션 헤드의 기능적 역할을 분석하는 연구는 활발히 이루어지고 있는 주제 중 하나이다. 그러나 기존 연구의 대부분은 거대 언어 모델(LLM)에 집중되어 있으며, 다른 유형의 모델에 대한 탐색은 제한적이다. 더 나아가 거대 언어 모델 내에서도 어텐션 헤드의 여러 기능적 역할은 아직 충분히 밝혀지지 않았다. 본 논문에서는 두 가지 유형의 생성 모델인 텍스트-이미지(T2I) 확산 모델과 거대 언어 모델에 대해 각각 새로운 어텐션 헤드 기여도 분석 기법을 제안한다. 먼저 텍스트-이미지 모델에 대해서는, Stable Diffusion을 대상으로 색상, 스타일, 동물과 같은 객체 범주 등 특정 시각적 개념과 관련된 크로스 어텐션 헤드를 식별하는 기여도 분석 방법을 제안한다. 식별된 개념과 관련된 헤드들은 세 가지 이미지 생성 응용에서 해당 개념을 강화하거나 억제하는 방식으로 제어에 활용된다. 실험 결과는 텍스트-이미지 확산 모델의 어텐션 헤드가 거대 언어 모델에서 관찰되는 것과 유사한 기능적 특수화(Functional Specialization)를 보임을 시사하며, 어텐션 헤드 기여도 분석 연구의 범위를 보다 넓은 생성 모델 계열로 확장시킨다. 거대 언어 모델의 경우, 본 논문은 어텐션 헤드 기여도 분석 문제를 연속된 공간에서의 최적화 문제(Continuous Optimization Problem)로 재구성한 새로운 기법을 제안한다. 이 방법은 소프트하고 미분 가능한 헤드 선택 파라미터(Soft, Differentiable Head-Selection Parameters)를 통해 태스크 임베딩(Task Embeddings)을 어텐션 헤드에 주입함으로써, 기울기 하강법(Gradient Descent)을 사용한 효율적인 태스크 관련 헤드(Task-Relevant Heads) 식별을 가능하게 한다. 이 방법은 기존 few-shot 인컨텍스트 학습(ICL)보다 유의미한 성능 향상을 보여주며, 57개 ICL 태스크에서 평균적으로 12개의 다양한 거대 언어 모델을 대상으로 10-shot ICL 성능을 10.2%–14.3% 상회한다. 성능 향상뿐 아니라, 본 방법은 해석 가능성을 위한 도구로서의 역할도 수행한다. 이는 태스크 관련 헤드를 식별할 뿐만 아니라, 해당 헤드들이 태스크-비특이적(Task-Agnostic) 관점이 아닌 태스크-특이적(Task-Specific) 관점에서 더 잘 설명된다는 사실을 보여준다--즉, 하나의 태스크에서 식별된 헤드는 유사한 특성을 가진 태스크로는 잘 전이되지만, 특성이 상이한 태스크에는 그렇지 않다. 전반적으로, 본 연구는 두 유형의 생성 모델에서 어텐션 헤드에 대한 작동 원리 이해를 심화시키며, 모델 동작의 정밀한 제어를 위한 통찰을 제공한다.

다국어 초록 (Multilingual Abstract)

Recent progress in deep learning has led to increased interest in mechanistic interpretability, which aims to understand the internal mechanisms under model behavior. Within this field, attention head attribution, the study of functional roles of individual attention heads in deep learning models, has become an active area of research. However, most existing studies have focused on large language models (LLMs), with limited exploration of other model classes. Even within LLMs, several functional roles of attention heads remain underexplored. This dissertation introduces attribution techniques for two distinct types of generative models: text-to-image (T2I) diffusion models and LLMs. For T2I models, we propose a head attribution method for Stable Diffusion that identifies cross-attention heads associated with specific visual concepts, such as color, style, and object categories like animals. These concept-relevant heads are then used to control image generation in three visual generative applications by modulating their activation. The results suggest that attention heads in T2I diffusion models exhibit functional specialization analogous to that observed in LLMs, extending the scope of attribution research to a broader class of generative models. For LLMs, we develop a method that reformulates head attribution as a continuous optimization problem. By injecting task embeddings into attention heads using soft, differentiable head-selection parameters, the method enables efficient identification of task-relevant heads via gradient descent. It achieves significantly higher task performance than few-shot in-context learning (ICL); averaged over 57 ICL tasks, the method outperforms the 10-shot ICL baseline by 10.2%-14.3% across 12 diverse LLMs. Beyond performance, our method also serves as a tool for interpretability. It not only identifies task-relevant heads but also reveal that these heads are characterized from a task-specific rather than task-agnostic perspective--attention heads identified for one task tend to transfer well to tasks with similar characteristics, but not to dissimilar ones. Overall, our studies provide a deeper mechanistic understanding of attention heads in two classes of generative models and offer insights for targeted control of model behavior.

목차 (Table of Contents)

• Chapter 1. Introduction 1
• 1.1 Dissertation outline 4
• 1.2 Related publications 5
• Chapter 2. Background 7
• 2.1 Two classes of generative models 7
• 2.1.1 Text-to-image generative models 7
• 2.1.2 Large language models 8
• 2.2 Multi-head attention layer 9
• 2.3 Mechanistic interpretability and targeted model control 10
• 2.4 Visual concepts and linguistic tasks 12
• Chapter 3. Cross-Attention Head Position Patterns Can Align with Human Visual Concepts in Text-to-Image Generative Models 14
• 3.1 Chapter overview 14
• 3.2 Related work 16
• 3.3 Method for constructing head relevance vectors 18
• 3.4 Ordered weakening analysis of head relevance vectors 26
• 3.5 Steering visual concepts in three visual generative tasks 29
• 3.5.1 Image generation - reducing misinterpretation of polysemous words 31
• 3.5.2 Image editing - successful editing for five challenging visual concepts 32
• 3.5.3 Multi-concept generation - reducing catastrophic neglect 37
• 3.6 Additional experiments 39
• 3.6.1 Extensions to a larger architecture 39
• 3.6.2 Ordered weakening analysis with more complex images 40
• 3.6.3 Extending human visual concepts 45
• Chapter 4. Soft Head Selection for In-Context Learning Driven Task Embedding Injection 47
• 4.1 Chapter overview 47
• 4.2 Related work 49
• 4.2.1 In-context learning in large language models 49
• 4.2.2 Task embedding injection in large language models 49
• 4.2.3 Attention head attribution in large language models 50
• 4.3 Method 51
• 4.4 Experiments 56
• 4.4.1 Experimental setup 56
• 4.4.2 Experimental results 72
• 4.5 Analysis of task-relevant attention heads through SITE 73
• 4.5.1 Identification of task-relevant attention heads 75
• 4.5.2 Task-specific vs. task-agnostic: which better explains head roles? 76
• 4.6 Empirical findings and efficiency analysis 88
• 4.7 Additional experiments 94
• 4.7.1 Additional evaluation on MMLU-Pro 94
• 4.7.2 Comparison with LoRA 95
• 4.7.3 Task-wise performance of hard injection 97
• 4.7.4 Comparison of hard injection and its complementary variant 100
• Chapter 5. Discussion 102
• 5.1 Do generation timesteps affect how heads relate to visual concept? 102
• 5.2 How does visual concept granularity affect HRV effectiveness? 104
• 5.3 Role of the argmax operation in HRV construction 111
• 5.4 Task adaptation in large language models 113
• 5.5 Implication and ethical considerations 114
• 5.6 Limitation 115
• 5.6.1 Shared limitations across methods 115
• 5.6.2 Limitation of HRV 116
• 5.6.3 Limitation of SITE 117
• 5.7 Future research direction 119
• Chapter 6. Conclusion 123
• Bibliography 124
• Appendices 139
• A Details and additional results on ordered weakening analysis 140
• A.1 Prompts used for ordered weakening analysis 140
• A.2 Additional results on ordered weakening analysis 140
• A.3 Comparison with random order weakening 142
• A.4 Ordered rescaling with varied rescaling factors 142
• B Details on reducing misinterpretation of polysemous words 151
• B.1 Prompts and selected concepts for reducing misinterpretation 151
• B.2 Human evaluation 155
• B.3 Comparison of concept strengthening and concept adjusting 155
• C Details and additional results on image editing 157
• C.1 Prompts for image editing 157
• C.2 Metrics for image editing and human evaluation on two image attributes 157
• C.3 Trade-off effect of self-attention replacement in P2P and P2P-HRV 160
• C.4 Additional results on image editing 164
• D Details and additional experiments on multi-concept generation 175
• D.1 Comparison with attribute-binding method 175
• D.2 Additional results on multi-concept generation 177
• E Additional results using SDXL 179
• F Task-wise performance across 11 additional LLMs 184
• G Extended analysis of task-relevant head identification 196
• G.1 Optimized soft head-selection values for all 57 tasks 196
• G.2 Optimized soft head-selection values for larger language models 200
• H Extended results on head-selection training dynamics 202
• H.1 Full results for all 57 tasks 202
• H.2 Results for larger language models 206
• Abstract (In Korean) 208
• Acknowledgement 210