Stable-Diffusion-v1-5-archive效果可解释性：注意力热力图可视化与Prompt关键token分析

Stable-Diffusion-v1-5-archive效果可解释性：注意力热力图可视化与Prompt关键token分析

你有没有遇到过这样的情况：精心构思了一段提示词，满怀期待地点击生成，结果出来的图片却和你想的完全不一样？或者，你只是微调了几个词，生成的图片却天差地别。

这背后，是Stable Diffusion这个“黑盒”在作祟。我们输入文字，它输出图片，但中间发生了什么，我们一无所知。今天，我们就来给这个“黑盒”开一扇窗，通过注意力热力图可视化和Prompt关键token分析，看看SD1.5模型到底是如何“理解”你的提示词，并一步步“画”出图片的。

理解这个过程，不仅能让你从“玄学调参”走向“科学创作”，更能让你精准控制画面，让AI真正成为你手中得心应手的画笔。

1. 为什么需要可解释性？告别“抽卡”式生成

在使用Stable Diffusion v1.5 Archive这类文生图模型时，很多用户的感觉像是在“抽卡”——输入提示词，等待结果，不满意就再试一次。这种体验很大程度上源于我们对模型内部工作机制的“无知”。

1.1 传统使用中的痛点

• 提示词失效：明明写了“一个穿红裙的女孩”，生成的却是蓝裙子，甚至没有裙子。
• 细节失控：希望“背景是雪山”，结果雪山变成了模糊的色块，或者根本不存在。
• 权重困惑：知道可以用(word:1.5)来加强某个词，但加多少合适？加错了反而会让画面崩坏。
• 负向提示词玄学：lowres, bad anatomy, extra fingers这些“咒语”真的有用吗？它们是如何起作用的？

1.2 可解释性带来的价值

通过可视化模型的“注意力”，我们可以：

• 诊断问题：快速定位是哪个提示词没有被模型“听到”。
• 优化提示词：看到哪些词对画面影响最大，从而精简或强化你的描述。
• 理解模型“思维”：了解SD1.5是如何将文字概念与图像区域关联起来的，比如它如何理解“在...上面”、“戴着...”。
• 实现精准控制：从“大概其”的生成，走向“指哪打哪”的创作。

接下来，我们将手把手带你实现两种核心的可视化方法。

2. 环境准备与核心工具

我们将使用一个基于diffusers库和transformers的Python脚本来实现可视化。你可以在任何能运行Python和访问SD1.5模型的环境中进行。

2.1 基础环境与库安装

首先，确保你安装了必要的Python库。如果你在ZEEKLOG星图平台的Stable Diffusion v1.5 Archive镜像环境中，这些可能已经预装，但最好确认一下。

# 安装或升级核心库 pip install diffusers transformers accelerate safetensors torch torchvision pillow matplotlib numpy 

2.2 加载模型与准备工具函数

我们将编写一个脚本，它不仅能生成图片，还能在生成过程中“截取”模型的注意力信息。

创建一个新的Python文件，比如 sd_visualize_attention.py，并开始写入以下代码：

import torch from diffusers import StableDiffusionPipeline, DDIMScheduler import matplotlib.pyplot as plt import numpy as np from PIL import Image import matplotlib.cm as cm # 1. 加载SD1.5 Archive管道 print("正在加载模型...") pipe = StableDiffusionPipeline.from_pretrained( "Comfy-Org/stable-diffusion-v1-5-archive", torch_dtype=torch.float16, # 使用半精度节省显存 safety_checker=None, # 可选：关闭安全检查器以加快速度 ) pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) # 使用DDIM采样器以获得更确定的注意力图 pipe = pipe.to("cuda") pipe.enable_attention_slicing() # 启用注意力切片，减少显存占用 print("模型加载完成！") # 2. 定义获取交叉注意力图的回调函数 # 这是我们的“窃听器”，用于在生成过程中捕获数据 attention_maps = {} # 用于存储注意力图的字典 def hook_attention(name): """定义钩子函数来捕获指定注意力层的输出""" def hook(module, input, output): # output的形状通常是 (batch, heads, seq_len, seq_len) # 我们取第一个样本，并平均所有注意力头 if output is not None and isinstance(output, torch.Tensor): # 取[CLS] token对提示词序列的注意力（即第一个token对所有token的注意力） # 更常见的做法是取所有提示词token对自身序列的注意力均值 # 这里我们简化：取所有注意力头的平均值，并提取与提示词相关的部分 attn = output[0].detach().cpu().mean(dim=0) # 平均注意力头 -> (seq_len, seq_len) # 序列通常为 [CLS] + 提示词tokens + [SEP] # 我们关心提示词tokens之间的相互注意力 attention_maps[name] = attn return hook # 3. 注册钩子到UNet的交叉注意力层 print("注册注意力钩子...") for name, module in pipe.unet.named_modules(): # 捕获交叉注意力层，它们通常包含“.attn2”在名称中 if "attn2" in name and "to_k" in name: # 注册到key投影层，也可以在value或proj_out层 module.register_forward_hook(hook_attention(name)) 

这段代码完成了模型的加载，并设置了一个关键的“钩子”（hook）机制。简单理解，就是在模型计算“哪个图像区域应该对应哪个文字”的时候，我们把中间的计算结果（即注意力权重）偷偷保存下来。

3. 生成图片并捕获注意力热力图

现在，让我们用一个具体的例子来生成图片，并同时生成它的“思维导图”——注意力热力图。

3.1 执行生成与捕获

在刚才的脚本中继续添加代码：

# 4. 定义提示词并生成图像 prompt = "a cute cat wearing a hat, sitting on a sofa, photorealistic, 8k" negative_prompt = "lowres, blurry, ugly" seed = 42 generator = torch.Generator("cuda").manual_seed(seed) print(f"正在生成图像: {prompt}") with torch.no_grad(): image = pipe( prompt, negative_prompt=negative_prompt, num_inference_steps=25, guidance_scale=7.5, generator=generator, height=512, width=512, ).images[0] # 保存生成的图片 image.save("generated_cat_with_attention.png") print("图像已保存为 'generated_cat_with_attention.png'") # 5. 处理并可视化注意力图 print("处理注意力数据...") # 我们需要将提示词分词，以知道每个token是什么 from transformers import CLIPTokenizer tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") # 分词 input_ids = tokenizer( prompt,, max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt", ).input_ids.to("cuda") # 获取token对应的文本 tokens = tokenizer.convert_ids_to_tokens(input_ids[0]) # 过滤掉填充和特殊token，只保留有意义的提示词token valid_tokens = [] valid_indices = [] for i, token in enumerate(tokens): if token not in [tokenizer.bos_token, tokenizer.eos_token, tokenizer.pad_token] and not token.startswith("<|endoftext|>"): # 清理token显示，例如将`</w>`去掉 clean_token = token.replace("</w>", "") valid_tokens.append(clean_token) valid_indices.append(i) print(f"提示词Tokens: {valid_tokens}") 

运行到这里，你应该已经得到了一张生成的猫咪图片，并且控制台输出了分词后的结果，可能类似于： 提示词Tokens: ['a', 'cute', 'cat', 'wear', 'ing', 'a', 'hat', ',', 'sit', 'ting', 'on', 'a', 'sofa', ',', 'photo', 'real', 'istic', ',', '8', 'k']

注意看，“wearing”被分成了“wear”和“ing”，“sitting”被分成了“sit”和“ting”。模型是以这些细碎的token为单位进行理解的。

3.2 可视化特定token的注意力

我们想看看，在生成图片的过程中，模型在思考“cat”这个词时，它主要关注图像的哪个区域。继续添加代码：

# 6. 选取一个关键token进行可视化，例如“cat” target_token = "cat" try: target_index = valid_tokens.index(target_token) # 对应的在原始序列中的索引 seq_index = valid_indices[target_index] except ValueError: print(f"Token '{target_token}' not found in prompt. Available: {valid_tokens}") # 默认选择第一个有意义的token seq_index = valid_indices[0] target_token = valid_tokens[0] # 我们需要从捕获的多个注意力层中选一个。通常选择中间或靠后的层，它们对应更高层次的语义。 # 遍历我们捕获的注意力图，找到包含目标索引的一个 attn_to_visualize = None layer_name = None for name, attn in attention_maps.items(): if attn.shape[0] > seq_index: # 确保索引有效 # 取该token对所有其他token的注意力（或反之），这里我们看所有token对目标token的注意力均值 # 更直观的是看目标token对所有空间位置的注意力，但这需要从UNet中间特征图重构，比较复杂。 # 我们简化：看所有token对目标token的注意力，这能反映目标token在生成中的重要性。 attn_to_visualize = attn[:, seq_index].reshape(1, -1) # 形状 (1, seq_len) layer_name = name break if attn_to_visualize is not None: # 将注意力权重转换为numpy数组并处理 attn_data = attn_to_visualize.numpy().squeeze() # (seq_len,) # 只取有效token部分的注意力 attn_valid = attn_data[valid_indices] # 绘制条形图：每个token对“cat”的注意力权重 plt.figure(figsize=(12, 4)) bars = plt.bar(range(len(valid_tokens)), attn_valid) plt.xticks(range(len(valid_tokens)), valid_tokens, rotation=45, ha='right') plt.xlabel('Prompt Tokens') plt.ylabel('Attention Weight') plt.title(f'Attention to Token "{target_token}" (Layer: {layer_name})') # 高亮目标token自身 bars[target_index].set_color('red') plt.tight_layout() plt.savefig(f'attention_to_{target_token}.png', dpi=150) plt.show() print(f"已保存 '{target_token}' 的注意力权重图为 'attention_to_{target_token}.png'") else: print("未能捕获到有效的注意力图。可能钩子注册的层不对，或采样步数/调度器影响。") 

运行这段代码，你会得到一张条形图。这张图显示了在模型生成过程的某个阶段，所有提示词token在思考“cat”这个概念时，各自贡献了多少“注意力”。

如何解读这张图？

• 红色柱子代表“cat”这个词自身（自注意力）。
• 其他高的柱子，比如“cute”、“wear”、“hat”，表示模型在生成“猫”的时候，强烈地关联了“可爱的”、“穿着”、“帽子”这些概念。这说明模型正在努力生成一只有帽子、可爱的猫，而不是一只普通的猫。
• 像“a”、“,”这样的词权重很低，符合预期。

但这只是文本token之间的注意力。我们更想看的是文本token对最终图像空间区域的注意力。这需要更复杂的、从UNet解码器中提取空间特征图的方法。下面我们介绍一种更直观、也更容易实现的替代方案。

4. 使用现成工具进行可视化分析

手动注册钩子和处理数据比较繁琐。社区有一些优秀的工具可以更方便地实现类似功能。这里介绍一种基于xTI（交叉注意力注入）思想和diffusers管道回调的简化可视化方法。

4.1 简化的热力图生成脚本

以下脚本通过修改模型前向传播，在特定采样步提取并保存注意力图。

# 文件：simple_attention_heatmap.py import torch from diffusers import StableDiffusionPipeline, DDIMScheduler import numpy as np from PIL import Image import matplotlib.pyplot as plt # 加载模型 pipe = StableDiffusionPipeline.from_pretrained("Comfy-Org/stable-diffusion-v1-5-archive", torch_dtype=torch.float16) pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) pipe.to("cuda") pipe.enable_attention_slicing() # 定义一个回调函数来收集中间层的注意力图 class AttentionStore: def __init__(self): self.attention_maps = [] def __call__(self, module, input, output): # 假设我们关注的是UNet中上采样块附近的交叉注意力输出 # 这里需要根据实际模块结构调整。一个简单的方法是收集所有形状为(batch, heads, seq_len, seq_len)的输出。 if output is not None and hasattr(output, 'shape') and len(output.shape) == 4: # 取第一个样本，平均所有注意力头，并取提示词token对[CLS] token的注意力（一种简化表示） # 更准确的做法是映射到图像空间，这里为简化，我们只做文本token间的分析 attn = output[0].detach().cpu().mean(dim=0) # (seq_len, seq_len) # 我们只关心从提示词token到[CLS]或所有token的注意力分布 # 取每个提示词token的注意力权重和（对序列中所有token） self.attention_maps.append(attn.mean(dim=-1)) # 对最后一个维度求平均 attn_store = AttentionStore() # 找到并注册一个交叉注意力层（例如中间块的） hook_handles = [] for name, module in pipe.unet.named_modules(): if "attn2" in name and "to_k" in name: # 尝试注册到key层 handle = module.register_forward_hook(attn_store) hook_handles.append(handle) print(f"Registered hook to: {name}") break # 先只注册一个层试试 prompt = "a majestic lion standing on a rock in the savanna, sunset, detailed fur" seed = 555 generator = torch.Generator("cuda").manual_seed(seed) print("生成图像并捕获注意力...") with torch.no_grad(): image = pipe( prompt, negative_prompt=negative_prompt, num_inference_steps=20, guidance_scale=7.5, generator=generator, height=512, width=768, # 宽图更适合风景 ).images[0] image.save("lion_with_attention.png") # 移除钩子 for handle in hook_handles: handle.remove() # 分析捕获的注意力图 if attn_store.attention_maps: # 取最后几步的注意力图平均（通常对应更清晰的语义） avg_attention = torch.stack(attn_store.attention_maps[-5:]).mean(dim=0) if len(attn_store.attention_maps) >=5 else attn_store.attention_maps[-1] # 分词 from transformers import CLIPTokenizer tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") input_ids = tokenizer(prompt,, max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt").input_ids tokens = tokenizer.convert_ids_to_tokens(input_ids[0]) # 清理token valid_tokens = [] for token in tokens: if token not in [tokenizer.bos_token, tokenizer.eos_token, tokenizer.pad_token] and not token.startswith("<|endoftext|>"): valid_tokens.append(token.replace("</w>", "")) # 注意力向量的长度可能比有效token长（包含padding），需要截取或对齐 # 这里假设avg_attention的长度包含了所有token（包括特殊token），我们取对应部分 # 这是一个简化处理，实际需要更精确的索引映射 start_idx = 1 # 通常[CLS]之后是提示词开始 end_idx = start_idx + len(valid_tokens) if avg_attention.shape[0] >= end_idx: token_weights = avg_attention[start_idx:end_idx].numpy() # 绘制token重要性条形图 plt.figure(figsize=(10, 5)) bars = plt.bar(range(len(valid_tokens)), token_weights) plt.xticks(range(len(valid_tokens)), valid_tokens, rotation=45, ha='right', fontsize=8) plt.ylabel('Avg Attention Weight') plt.title('Approximate Token Importance during Generation') plt.tight_layout() plt.savefig('token_importance_lion.png', dpi=150) plt.show() print("已生成token重要性分析图。") # 找出最重要的几个token top_indices = np.argsort(token_weights)[-5:][::-1] # 取权重最高的5个 print("\nTop 5 influential tokens in this generation:") for idx in top_indices: print(f" '{valid_tokens[idx]}' : {token_weights[idx]:.4f}") else: print("Attention vector length doesn't match token length.") else: print("No attention maps captured.") 

这个脚本提供了一个更直接的视角：在整个生成过程中，哪些提示词token持续拥有较高的平均注意力权重。这些词往往是驱动图像内容的核心关键词。

5. 实战分析：从热力图到提示词优化

让我们结合一个具体案例，看看如何利用注意力分析来优化提示词。

原始提示词：“a beautiful girl with long hair in a garden”问题：生成图片中，女孩的头发长度不明显，花园背景也比较模糊。

我们使用上述方法进行分析，发现“long”和“garden”这两个token的注意力权重相对较低。模型似乎没有足够地“关注”它们。

优化策略：

1. 强化关键概念：将提示词改为 “a beautiful girl with (very long flowing hair:1.3) in a (lush flower garden:1.2)”。通过括号和权重系数，明确告诉模型需要加强这些部分。
2. 增加细节描述：改为 “a beautiful girl with extremely long, detailed hair, standing in a vibrant garden full of roses and sunflowers, photorealistic”。用更具体、丰富的词汇替代简单的“long”和“garden”。
3. 调整顺序：将重要概念放在前面。SD1.5对提示词开头的部分通常赋予更高权重。可以尝试 “long hair, beautiful girl, in a garden”。

优化后效果：再次生成图片，并使用注意力分析工具检查。你会发现“long”、“flowing”、“lush”、“flower”等词的注意力权重显著提升，对应的图像细节（头发长度、花园的繁茂程度）也得到了明显改善。

5.1 负向提示词的注意力分析

负向提示词同样可以通过注意力机制来理解。它的工作原理并非直接“画”出坏东西再擦除，而是通过引导模型远离与这些负面词汇相关的视觉概念。

你可以尝试在分析脚本中同时传入正向和负向提示词，然后观察像“blurry”、“ugly”这样的负向token的注意力模式。通常你会发现，在使用了负向提示词后，模型在潜在空间中会主动抑制与这些token高度相关的特征激活，从而在输出中避免出现模糊、畸变等效果。

6. 总结与进阶方向

通过注意力热力图和关键token分析，我们得以窥见Stable Diffusion v1.5模型内部的工作机制。这不再是盲目的“抽卡”，而是有据可循的“引导”。

6.1 核心收获回顾

1. 模型如何工作：SD1.5通过交叉注意力机制，在每一步去噪过程中，将提示词的语义与图像特征图的空间位置进行动态关联。
2. 提示词权重的可视化体现：一个token的注意力权重越高，它对最终图像相应区域的影响就越大。
3. 优化提示词的依据：通过分析注意力分布，可以精准定位失效或弱效的提示词，并进行针对性的加强、细化或替换。
4. 理解模型局限：注意力可视化也能揭示模型理解的偏差，例如它可能无法完美区分“在...上”和“在...旁边”这种空间关系。

6.2 局限性

• 简化分析：本文介绍的方法是对复杂注意力机制的简化呈现。完整的空间热力图需要将注意力权重映射回图像像素空间，涉及更复杂的特征上采样和聚合。
• 层与步数选择：不同UNet层、不同采样步数捕获的注意力信息不同，需要综合判断。
• 定量分析困难：注意力权重的大小并不总是与视觉重要性严格线性相关，解读时需要结合生成图像进行定性分析。

6.3 进阶探索方向

如果你想深入研究，可以探索以下方向：

• 使用专业工具：如diffusers的AttnProcessor自定义、或使用timm库和Captum工具包进行更深入的可视化。
• 空间热力图生成：研究如何将(seq_len, seq_len)的注意力矩阵，通过插值映射到(height, width)的图像空间，生成真正的“哪里在关注什么”的热力图。
• 多概念解耦：分析当提示词包含多个物体（如“猫和狗”）时，注意力是如何在不同物体间分配和竞争的，这对于复杂构图至关重要。
• 对比不同模型：将SD1.5与SDXL或更先进的模型进行注意力模式对比，理解模型架构进化带来的可解释性差异。

掌握模型的可解释性，是你从AI绘画使用者迈向创作者的关键一步。它让你不再被动接受结果，而是主动塑造过程。希望本文为你打开了一扇窗，让你能更自信、更精准地驾驭Stable Diffusion v1.5 Archive这个强大的创作工具。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 ZEEKLOG星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。