深入理解Transformer通过SmallThinker-3B-Preview剖析模型内部注意力机制当我们谈论大语言模型时“注意力机制”和“Transformer架构”是绕不开的核心。但你是否曾好奇当模型在回答你的问题时它的“注意力”究竟聚焦在哪里那些抽象的“多头注意力”和“层”是如何协同工作最终生成流畅文本的今天我们就借助一个轻量级的开源模型——SmallThinker-3B-Preview来一场“模型大脑”的探秘之旅。我们将不再停留在理论层面而是通过实际的代码和可视化工具深入模型内部亲眼看看注意力机制是如何运作的。你会发现理解Transformer的内部工作远比想象中更有趣和直观。1. 为什么选择SmallThinker-3B-Preview来“解剖”在深入技术细节之前我们先聊聊为什么选它。SmallThinker-3B-Preview是一个参数量为30亿的预训练语言模型。对于“解剖”注意力机制这个任务来说它有几个得天独厚的优势。首先它“个头”适中。动辄数百亿、上千亿参数的大模型内部结构极其复杂注意力头可能多达上百个层数也深不见底。对于初学者甚至中级研究者来说直接分析它们就像用显微镜观察整个森林容易迷失在细节中。SmallThinker-3B-Preview的规模则刚刚好结构清晰便于我们追踪和理解信息流动的完整路径。其次它是完全开源的。这意味着我们可以获得其完整的模型架构定义、权重文件并且能够毫无障碍地使用像transformers这样的标准库来加载它并利用丰富的工具来提取中间层的激活值包括我们最关心的注意力权重。最后它的能力足够“示范”。虽然只有30亿参数但它已经具备了不错的语言理解和生成能力足以展示注意力机制在完成具体任务如回答问题、续写文本时的典型模式。通过分析它在具体例子上的表现我们能获得对Transformer工作原理的普适性理解。简单来说用它来学习就像用一台结构清晰、零件可见的经典发动机来学习内燃机原理比直接研究最新款的混动超跑要直观得多。2. 搭建我们的“注意力观测站”理论准备就绪接下来我们搭建实验环境。整个过程非常简单只需要几个常见的Python库。# 环境搭建安装必要的库 # 在终端或Jupyter Notebook中执行以下命令 # pip install torch transformers numpy matplotlib seaborn import torch from transformers import AutoModelForCausalLM, AutoTokenizer import numpy as np import matplotlib.pyplot as plt import seaborn as sns # 设置中文字体可选用于图表标签 plt.rcParams[font.sans-serif] [SimHei, DejaVu Sans] plt.rcParams[axes.unicode_minus] False print(环境准备完成主要库版本) print(fPyTorch: {torch.__version__}) print(fTransformers: {transformers.__version__})接下来我们加载SmallThinker-3B-Preview模型和对应的分词器。这里的关键是我们需要以能够输出中间注意力权重的模式来加载模型。# 加载模型与分词器并启用输出注意力权重的选项 model_name SmallThinker-3B-Preview # 请根据实际的Hugging Face模型ID或本地路径替换 tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) # 关键在加载模型时设置 output_attentionsTrue model AutoModelForCausalLM.from_pretrained( model_name, trust_remote_codeTrue, torch_dtypetorch.float16, # 使用半精度节省显存 device_mapauto, # 自动分配模型层到可用设备GPU/CPU output_attentionsTrue # 告诉模型前向传播时需要返回注意力权重 ) print(f模型 {model_name} 加载成功) print(f模型设备{model.device})代码中的output_attentionsTrue是本次探索的“钥匙”。它指示transformers库在模型进行前向计算时不仅要返回最终的预测结果还要把每一层、每一个注意力头的注意力权重矩阵也一并返回给我们。3. 第一次观测注意力如何聚焦在一个简单句子上让我们从一个简单的例子开始。假设我们给模型输入一句话“人工智能正在改变世界。”然后让它生成下一个词。我们来看看在生成过程中模型的注意力都放在了哪里。# 准备输入文本 input_text 人工智能正在改变世界。 inputs tokenizer(input_text, return_tensorspt).to(model.device) # 进行模型推理获取输出包括注意力权重 with torch.no_grad(): outputs model(**inputs, output_attentionsTrue) # outputs 现在包含 # - logits: 模型对下一个词的预测分数 # - attentions: 一个元组包含每一层输出的注意力权重 # attentions 的 shape: (层数, 批大小, 注意力头数, 目标序列长度, 源序列长度) attentions outputs.attentions # 这就是我们需要的所有注意力权重 num_layers len(attentions) num_heads attentions[0].size(2) # 假设所有层的头数相同 seq_len attentions[0].size(-1) # 序列长度 print(f模型总层数{num_layers}) print(f每层注意力头数{num_heads}) print(f输入序列长度含特殊符号{seq_len}) print(f注意力权重元组长度层数{len(attentions)})现在我们有了一个四维张量attentions[layer][batch][head][target, source]。它记录了在生成序列中每个“目标”位置target时对输入序列中每个“源”位置source的关注程度。对于解码器式的语言模型这通常是“因果注意力”即位置i只能关注位置j (j i)。让我们可视化最后一层通常被认为整合了最丰富语义信息的某个注意力头看看它关注什么。def plot_attention(attention_matrix, layer_idx, head_idx, tokens): 绘制指定层和头的注意力热力图。 attention_matrix: 从 outputs.attentions 中提取的权重 layer_idx: 层索引从0开始 head_idx: 头索引从0开始 tokens: 对应的token列表用于坐标轴标签 # 提取指定层、指定头、第一个批次的注意力权重 # 我们通常取最后一个解码位置的注意力即生成下一个词时的注意力分布 attn_data attention_matrix[layer_idx][0, head_idx].cpu().numpy() # 对于因果语言模型我们通常看最后一个token即要预测的下一个位置关注了哪些前面的token # 这里我们取最后一行目标位置是序列末尾 last_token_attention attn_data[-1, :] # 形状: (seq_len,) # 或者我们可以看整个注意力矩阵上三角部分因为是因果的 plt.figure(figsize(10, 8)) # 绘制热力图 sns.heatmap(attn_data, xticklabelstokens, yticklabelstokens, cmapYlOrRd, squareTrue, cbar_kws{shrink: 0.8}) plt.title(f注意力热力图 - 第{layer_idx1}层第{head_idx1}个头, fontsize14) plt.xlabel(源序列 (Source Tokens)) plt.ylabel(目标序列 (Target Tokens)) plt.tight_layout() plt.show() # 也可以绘制最后一个token的注意力条形图更清晰 plt.figure(figsize(12, 4)) bars plt.bar(range(len(tokens)), last_token_attention) plt.xticks(range(len(tokens)), tokens, rotation45, haright) plt.title(f序列末位Token的注意力分布 - 第{layer_idx1}层第{head_idx1}个头) plt.ylabel(注意力权重) # 高亮最高权重的bar max_idx np.argmax(last_token_attention) bars[max_idx].set_color(red) plt.tight_layout() plt.show() # 获取输入文本的tokens用于标签 tokens tokenizer.convert_ids_to_tokens(inputs[input_ids][0]) print(输入Tokens:, tokens) # 绘制最后一层假设是第24层索引23的第一个头的注意力 plot_attention(attentions, layer_idxnum_layers-1, head_idx0, tokenstokens)运行这段代码你会得到两张图。热力图展示了整个注意力矩阵你可以看到一条明显的对角线这是“自回归”或“因果”注意力的特征——每个位置主要关注自身及之前的位置。条形图则清晰地告诉我们当模型试图预测“人工智能正在改变世界。”之后的下一个词时它最关注的是哪个或哪些词。你可能会发现它高度关注“世界”或“改变”这样的关键词这为它的预测提供了上下文。4. 深入案例注意力如何解决指代与逻辑单个句子可能太简单。我们来看一个更复杂的例子其中包含指代和逻辑关系。比如“小明去了图书馆因为他想借一本书。他找到了一本关于历史的书。” 我们让模型续写“他感到...”在这个例子中“他”指代“小明”而“感到”的情绪可能与“历史书”或“借书”这个行为有关。注意力机制需要将这些信息关联起来。# 复杂文本示例指代与逻辑 complex_text 小明去了图书馆因为他想借一本书。他找到了一本关于历史的书。他感到 inputs_complex tokenizer(complex_text, return_tensorspt).to(model.device) with torch.no_grad(): outputs_complex model(**inputs_complex, output_attentionsTrue) attentions_complex outputs_complex.attentions tokens_complex tokenizer.convert_ids_to_tokens(inputs_complex[input_ids][0]) print(复杂文本Tokens:, tokens_complex) # 我们关注生成“感到”之后情绪词时的注意力。 # 假设“感到”是倒数第二个token我们看看模型在生成最后一个位置预测情绪词时注意力如何分布。 # 我们需要提取对应“感到”之后那个预测位置的注意力。 # 在transformers的输出中attentions的最后一个维度对应的是“源”序列即我们输入的整个序列。 # 对于因果模型在生成位置t时其注意力权重向量的长度也是t只能看到前面的t个token。 # 但为了方便我们通常看最后一个已生成token“感到”对前面所有token的注意力。 # 选择中间层如第12层和最后一层第24层的某个头进行对比 layer_to_observe [11, num_layers-1] # 第12层和最后一层的索引 head_to_observe 4 # 观察第5个头可以换不同的头看看 for l_idx in layer_to_observe: # 获取该层该头的注意力矩阵 attn_mat attentions_complex[l_idx][0, head_to_observe].cpu().numpy() # 我们看目标位置为“感到”所在位置假设是序列的倒数第二个token的注意力分布 target_pos -2 # “感到”的位置 source_attention attn_mat[target_pos, :target_pos1] # 因果注意力只看到当前位置 # 由于source_attention长度短于总tokens数我们需要对齐 relevant_tokens tokens_complex[:target_pos1] plt.figure(figsize(14, 5)) bars plt.bar(range(len(relevant_tokens)), source_attention) plt.xticks(range(len(relevant_tokens)), relevant_tokens, rotation60, haright, fontsize10) plt.title(f生成‘感到’时对前文的注意力 - 第{l_idx1}层第{head_to_observe1}个头) plt.ylabel(注意力权重) # 高亮前三个最关注的token top_k 3 top_indices np.argsort(source_attention)[-top_k:][::-1] colors [red, orange, green] for idx, color in zip(top_indices, colors): bars[idx].set_color(color) # 在柱子上方标注权重值 plt.text(idx, source_attention[idx]0.01, f{source_attention[idx]:.3f}, hacenter, vabottom, fontsize9, colorcolor) plt.tight_layout() plt.show() print(f第{l_idx1}层第{head_to_observe1}个头最关注的Token及其权重) for i in top_indices: print(f Token: ‘{relevant_tokens[i]}’ 权重: {source_attention[i]:.4f})通过对比中间层和最后一层的注意力图你可能会发现有趣的现象中间层可能更关注局部的语法和短语结构比如“他”与“小明”的指代关系共指消解或者“书”与“历史”的修饰关系。最后一层可能更关注全局的语义和逻辑比如将“感到”与“图书馆”地点、“借书”目的、“历史”书的内容等更广泛的上下文联系起来为预测“高兴”、“无聊”、“满意”等情绪词做准备。这就是多头注意力的魅力所在不同的头、不同的层各司其职共同协作完成复杂的语言理解任务。5. 注意力模式的归纳从个案到规律通过观察大量例子研究者们总结出Transformer注意力头的一些常见模式。虽然SmallThinker-3B-Preview的具体模式可能有所不同但大体上可以归为以下几类我们可以尝试在我们的观测中寻找它们语法头关注特定的语法结构。例如一个头可能专门关注动词和其宾语之间的关系另一个头关注名词和其修饰词如形容词之间的关系。在我们之前的例子中关注“一本”和“书”之间关系的头可能就属于此类。语义头关注语义上相关的词即使它们相隔很远。例如在“苹果公司发布了新款手机它的设计很惊艳”中一个语义头需要将“它”与远处的“手机”或“苹果公司”关联起来。位置头关注固定偏移的位置。例如总是关注前一个词或后一个词这有助于捕捉局部依赖和序列顺序。罕见词/实体头特别关注输入中的命名实体、数字或罕见词这些词通常携带关键信息。全局头注意力分布相对均匀可能起到“广播”全局信息或“残差”连接的作用。我们可以写一个简单的函数批量分析几个不同输入样例统计某个头在不同情况下最关注什么类型的词来猜测它的“职能”。def analyze_head_function(head_idx, layer_idx, test_sentences, model, tokenizer): 简单分析一个特定注意力头的功能倾向。 patterns [] for sent in test_sentences: inputs tokenizer(sent, return_tensorspt).to(model.device) with torch.no_grad(): outputs model(**inputs, output_attentionsTrue) attn outputs.attentions[layer_idx][0, head_idx].cpu().numpy() # 取最后一个生成位置的注意力分布 last_attn attn[-1, :] tokens tokenizer.convert_ids_to_tokens(inputs[input_ids][0]) # 找出最关注的token top_idx np.argmax(last_attn) top_token tokens[top_idx] patterns.append((sent, top_token, last_attn[top_idx])) print(f 对第{layer_idx1}层第{head_idx1}个头的初步分析 ) for sent, top_tok, weight in patterns: print(f句子: \{sent}\) print(f 生成句末时最关注的词: ‘{top_tok}’ (权重: {weight:.3f})) print(*50) # 准备一些测试句子 test_sentences [ 猫在沙发上睡觉。, # 简单主谓宾 穿着红色裙子的女孩在公园里唱歌。, # 长修饰结构 尽管下雨了他们还是决定去野餐。, # 转折关系 北京是中国的首都它是一座历史悠久的城市。 # 指代关系 ] # 分析最后一层的第0个头和第5个头 analyze_head_function(head_idx0, layer_idxnum_layers-1, test_sentencestest_sentences, modelmodel, tokenizertokenizer) analyze_head_function(head_idx5, layer_idxnum_layers-1, test_sentencestest_sentences, modelmodel, tokenizertokenizer)通过这样的分析你可能会发现某个头总是关注句号、动词或主语而另一个头则可能对“尽管”、“但是”这样的转折词或者“它”、“他们”这样的代词特别敏感。这为我们理解模型内部的“分工”提供了直观的线索。6. 总结与展望通过这次对SmallThinker-3B-Preview模型注意力机制的可视化探索我们亲手揭开了Transformer“黑箱”的一角。我们看到注意力并非均匀分布而是像一束智能的聚光灯根据任务需求动态地、有选择地照亮输入序列的不同部分。从关注相邻词的“语法头”到串联远距离信息的“语义头”这些机制共同构成了模型理解语言的基石。这种可视化分析的价值远不止于教学和好奇。对于研究人员它是诊断模型偏见、理解错误来源、甚至设计更高效模型架构如剪枝不重要的注意力头的重要工具。对于开发者深入理解注意力有助于设计更好的提示词Prompt引导模型更关注输入中的关键信息。当然我们今天看到的只是冰山一角。完整的Transformer还包括前馈神经网络、层归一化、残差连接等组件它们与注意力机制协同工作。注意力权重本身也只是模型决策过程的一部分。要全面理解模型还需要结合其他中间激活值、梯度等信息进行分析。不过万事开头难。掌握了提取和解读注意力权重这项基本技能你就拥有了探索更大、更复杂模型内部世界的钥匙。下次当你使用任何一个基于Transformer的AI模型时不妨在脑海中想象一下那些无形的“注意力”正在文本的海洋中如何跳跃、关联最终汇聚成你眼前那句流畅而智能的回答。这或许能让你对这项技术的精妙之处有更深一层的体会。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。