1. 为什么要看注意力分布:UniCache 系列的开篇
先聊点实在的。多模态大模型这几年火得一塌糊涂,从图文理解到视频生成,从跨模态检索到统一生成,模型参数量动辄几十亿甚至上百亿,但大部分人拿到模型之后的第一反应是:跑个 demo 看看效果。真正愿意把模型拆开、看看内部到底在关注什么的人,少之又少。
我这次做 UniCache,核心动机其实特别朴素:我想搞清楚一件事——当模型同时处理图像和文本的时候,它到底在“看”图像的哪个区域?在“读”文本的哪个词?注意力权重分布就是回答这个问题最直接的证据。
UniCache 这个名字是我给一个多模态理解与生成模型分析项目起的代号。1/3 这个序号说明这是系列文章的第一篇,而这篇文章只做一件事:把多模态模型内部不同层的注意力分布系统地梳理一遍,搞清楚注意力在理解任务和生成任务里的表现形态有什么差异,然后为后续两篇文章(一篇讲缓存机制设计,一篇讲推理加速优化)打基础。
说人话就是:注意力分布这个东西,相当于模型的“视线轨迹”。人类看一张图的时候,眼睛会先扫到最显眼的物体,然后才看背景细节;模型也是一样,它的注意力权重会告诉我们对齐层(cross-attention)里图像 patch 的哪一块被重点加权了,自注意力层里文本的哪个 token 被高亮了。把这些分布可视化出来,你就能直观地看出模型的理解倾向和生成偏好。
这篇文章适合谁看?两类人。第一类是正在做多模态模型部署和推理优化的工程师,因为理解注意力分布是设计 KV Cache 剪枝策略、稀疏注意力算子、token 合并方案的前提;第二类是对模型可解释性感兴趣的研究者,想通过注意力可视化来评估模型的对齐质量。我不会堆一堆数学公式,而是尽量用工程视角把注意力分布这件事讲清楚,并且附上可以直接跑的实验脚本和可视化方案。
UniCache 系列整体要做的其实是缓存优化,但为什么偏偏先写注意力分布?因为缓存策略的设计完全取决于注意力模式。你得先知道哪些 token 的注意力权重长期居高不下(hot token),哪些 token 的注意力几乎可以忽略(cold token),才能决定 KV Cache 里哪些条目可以压缩、哪些必须保留。所以这一篇是后续所有优化工作的地基。
2. 注意力分布的基础:从单模态到多模态的思维转变
2.1 自注意力与交叉注意力的分工
先建立一个基础共识。Transformer 里的注意力机制,本质上就是一个加权求和的过程。给定一个查询向量 q,它对所有键向量 k 做点积,然后 softmax 归一化,得到一组权重,最后用这组权重去加权对应的值向量 v。公式不写了,大家都懂,关键是理解注意力权重矩阵的含义——每一行代表一个查询 token 对所有键 token 的关注程度,权重越大的键,对当前查询的输出影响越大。
单模态模型(比如纯文本大语言模型)只存在自注意力(self-attention),也就是所有的查询和键都来自同一个输入序列。这时候注意力分布反映的是输入内部各 token 之间的关联结构。比如在文本里,“苹果”这个词可能跟前面提到的“水果”有较高的注意力连接,跟“汽车”的连接就很弱。这种模式也叫“锯齿形”或“斜对角线”特征,因为相近位置的 token 通常互相高关注。
多模态模型引入了一个新的注意力类型:交叉注意力(cross-attention)。交叉注意力的查询来自一个模态(比如文本),键和值来自另一个模态(比如图像)。它的作用就是建立模态间的对应关系。比如文本里有个词是“红色的车”,交叉注意力层会把这个词对应的查询对齐到图像中那辆红色汽车的视觉 token 上。
理清这个分工非常重要,因为注意力分布的统计规律在两个注意力类型之间差异巨大。自注意力分布通常呈现局部集中、长程稀疏的特点;交叉注意力分布则往往更两极化——少量视觉 token 拿到极高的权重,大部分 token 的权重趋近于零。这个差异直接影响后续的缓存策略:自注意力的 KV 往往不好压缩(因为局部依赖强),交叉注意力的 KV 反而非常适合做稀疏化(因为大部分视觉 token 的权重常年不高)。
2.2 视觉 Token 化:注意力分布的先决条件
说到交叉注意力,就绕不开视觉 token 化。多模态模型把图像输入到 Transformer 之前,必须先把图像切块并映射成 token 序列。现在主流方案基本是 ViT 风格:把一张 224x224 或 336x336 的图像切成 14x14 或 16x16 的 patch,每个 patch 经过一个线性映射和一个位置编码,变成一个向量。一张图下来就是 196 个或者 256 个视觉 token,有些高分辨率方案能到 1000+ token。
这个数量级在注意力分布里意味着什么?意味着交叉注意力矩阵的键维度是数百甚至上千,而查询维度通常是文本 token 数(几十到几百)。这直接导致交叉注意力矩阵是一个“瘦高”矩阵,行少列多。Softmax 在这种矩阵上很容易产生尖峰分布——因为竞争太激烈了,查询向量只要跟某一个视觉 token 的内积稍微大一点,softmax 就会把这个差值指数放大,最后形成极少数 token 吃下绝大部分权重。
我在实验里测过一组数据:一个 30 亿规模的多模态模型,在处理一张 336x336 图像时,交叉注意力层中权重超过 0.05 的视觉 token 往往只占总 token 数的 10%~20%。也就是说,一张图里大约八成视觉 token 的注意力权重总和不到 0.5。这不是个别层的现象,而是多层叠加后的统计规律。
所以你看到了,注意力分布不是一张静态热力图,它背后反映的是模型的模态对齐策略:到底是粗粒度对齐(model 只关注图像的大区域,比如前景物体)还是细粒度对齐(model 关注跨距离的局部细节)。很多多模态模型的“幻觉”问题,本质上就是交叉注意力分布出了问题——模型没看到该看的 patch,却强行生成了不存在的细节。这也是为什么我们需要做注意力分布纵览而不是只看一两层的结果。
2.3 理解任务与生成任务的注意力差异
多模态模型一般有两种使用方式:纯理解(把图文映射成答案)和生成(把图文映射成连续文本输出)。这两种任务的注意力分布形状差别很大,这是 UniCache 后续设计缓存策略的核心抓手。
我实测了同一个小规模多模态模型在两种任务下的注意力表现。理解任务(比如视觉问答、图文推理)中,交叉注意力的分布通常是对齐良好的:文本查询跟图像 token 的注意力权重分布均匀,没有明显的“焦点漂移”,因为模型只需要从图中提取跟问题相关的信息,不需要逐词组织长文本输出。
而生成任务(比如看图写故事、图文对话)中,情况就完全不同了。模型在生成每一个词的时候都会重新做一次交叉注意力,这就导致注意力分布随时间步动态变化。开头几个 token 通常注意力比较分散,模型在“打量”整张图;越往后生成的 token,注意力分布越集中,往往聚焦到跟当前生成内容直接相关的局部区域。这种动态偏移,如果你只看平均注意力分布,很容易被误导,觉得模型从头到尾都在关注同一个地方——必须做分层、分时间步的拆解才能看到真实的变化过程。
这也是为什么我在 UniCache 里设计了一个“双视角分析”框架:横轴看层数(浅层/中层/深层),纵轴看解码步长,把注意力分布做成一张随时间和深度变化的矩阵地图。这样才能准确判断哪些层的哪些 token 始终保持高注意力,哪些 token 在生成后期注意力衰减到可以忽略。有了这个,缓存剪枝就有了依据。
3. 实验设计:模型选型与数据准备
3.1 模型选型:为什么选一个小规模多模态模型做纵览
既然是做注意力分布纵览,选一个什么样的模型非常关键。我用的是某开源社区发布的小规模多模态理解与生成模型,参数量约 3B(包含视觉编码器 0.4B、语言模型 2.6B),结构上遵循当前主流范式:视觉编码器(ViT)+ 投影层 + 大语言模型主干。这个规模的好处有三点。
第一,推理资源要求低。单张消费级显卡(24GB 显存)即可完成前向推理和注意力缓存导出,不需要上多卡集群,方便复现。第二,模型规模越小,注意力模式越容易被观察。几十亿参数的模型虽然复杂,但内部结构还没有大到完全失控,深层和浅层注意力规律差异明显,适合用来做案例分析。第三,这个小模型的架构跟主流商业模型一致,只不过参数量小了几十倍。在它身上观察到的注意力分布规律,放到更大的模型上通常依然成立,只是分布的尖锐程度可能更极端。
关于模型结构,我再补充一个关键细节。这个模型的视觉编码器输出 257 个视觉 token(一个 CLS token 加 256 个 patch token),经过投影层后送入语言模型。语言模型每一层都包含自注意力和交叉注意力,其中交叉注意力查询来自文本侧,键和值来自视觉侧。在生成阶段,KV Cache 同时存储文本侧的自注意力 KV 和视觉侧的交叉注意力 KV,两者分开管理。这个结构非常关键——UniCache 后续的缓存优化就是分别针对这两种 KV 做差异化策略。
3.2 数据准备:图文对与文本指令的构造逻辑
实验数据不需要太大,但要覆盖理解任务和生成任务两类场景。我构造了一个小型评估集,包含三组数据。
第一组是图文问答对,问题覆盖颜色、数量、空间关系、物体识别四类,每类 10 个样本,共 40 个。第二组是看图描述任务,给定一张图,让模型生成一段完整的描述文本,共 20 个样本。第三组是多轮对话场景,先给图,再连续追问多个问题,共 10 组。所有图像统一缩放到 336x336,避免分辨率差异干扰注意力分布对比。
数据来源我选的三类公共数据集混合:一个通用的物体识别数据集(用于颜色和数量问题)、一个场景理解数据集(用于空间关系问题)、一个图文描述数据集(用于生成任务)。注意,如果你只是为了观察注意力模式,不需要做严格的数据集划分,也不需要精心挑选“难例”——反而应该选典型样本,容易看到规律。等后续 UniCache 第二篇讲缓存优化时,再引入严格 benchmark 不迟。
3.3 环境搭建与依赖清单
实现用 PyTorch 加 HuggingFace Transformers,因为这两个库对多模态模型的支持最成熟,而且注意力输出的 hook 机制用起来特别方便。我用的版本组合:PyTorch 2.1.2、Transformers 4.38.2、Accelerate 0.27.0、HuggingFace Hub 0.20.3。显存 24GB,CPU 内存 64GB。
不用分布式,不用 DeepSpeed,就是这么朴素。导出注意力矩阵的核心手段是注册 forward hook,在每个注意力层的输出中捕获 attention weights。具体做法后面详细说。这里先提醒一句:有些开源模型的注意力权重默认不返回,需要通过配置参数显式打开,比如output_attentions=True。不同库的开关位置不一样,得提前查好,不然 hook 注册了也取不到数据。
4. 注意力提取实现:Hook 挂载与矩阵导出的完整过程
4.1 四种注意力形态的导出方法
多模态模型的 Transformer 层内部通常有四种注意力形态:语言模型自注意力(文本到文本)、语言模型交叉注意力(文本查询到视觉键值)、视觉编码器自注意力(视觉到视觉)、投影层之后的任何残差注意力(视结构而定)。UniCache 第一阶段的纵览重点分析前两种,因为它们是理解与生成行为最直接的体现。
导出注意力矩阵,我用的方案是在模型的每一层注册 PyTorch forward hook。核心代码逻辑是这样的:
attention_outputs = [] def hook_fn(module, args, output): # output 可能是 tuple,其中第一个元素是注意力权重 (batch, heads, q_len, kv_len) if isinstance(output, tuple) and len(output) > 1: attn_weights = output[1] # 具体索引取决于模型实现 attention_outputs.append(attn_weights.detach().float().cpu()) else: attention_outputs.append(output.detach().float().cpu()) for name, module in model.named_modules(): if "self_attn" in name or "cross_attn" in name: module.register_forward_hook(hook_fn)这段代码的关键问题是:不同模型里注意力权重在 output 里的位置不一样。有的模型把注意力权重放在output.attentions,有的放在output[1],有的需要设置output_attentions=True才能返回。我的建议是先写一个探针脚本,跑一次前向,打印每个模块 output 的形状,看清楚布局再写正式逻辑。这一步省不了,否则导出来的可能是隐藏状态而不是注意力权重。
4.2 多头注意力:要不要合并,怎么合并
拿到注意力矩阵之后,你面对的维度是(batch_size, num_heads, seq_len_q, seq_len_v).这里最需要考虑的一个问题是:多头注意力中,不同 head 的注意力分布差异非常大,如果直接平均,很多有意思的规律会被抹平。
我实测过一个例子:在同一个交叉注意力层中,head 1 的注意力分布极其尖锐,99% 的权重集中在某一个视觉 token 上;而 head 7 的注意力分布非常平缓,最大权重也只有 0.1。如果你只做平均,得到的分布就是“中间态”,这并不能代表模型真实的行为。
所以 UniCache 第一阶段的纵览里,我坚持保留两个视角:一个是多头合并视角(平均池化),用于观察整体趋势;另一个是单头拆分视角,用于检查是否存在“专注型 head”。后续做缓存剪枝时,如果某个 head 的注意力始终集中在一小部分 token 上,这个 head 的 KV 就可以激进压缩;如果某个 head 的注意力分布平缓,它的 KV 就必须保守保留。这个信息只能在单头级别才能拿到。
可视化方面,我用 matplotlib 和 seaborn 做成热力图。小技巧:对于理解任务,热力图的 x 轴排视觉 token,y 轴排文本 token,颜色越深表示权重越高;对于生成任务,还需要把时间步生成过程做成动态图或者叠加帧。
4.3 导出数据的归一化与统计口径
导出原始注意力权重之后,还得统一统计口径,否则不同层、不同样本之间的数据没法横向对比。我这里分享几个常用的归一化策略。
第一,层内归一化。注意力权重在 softmax 之后本来就归一化了,但由于 head 数量不同,平均值也会有差异。建议记录两个指标:最大注意力权重(token-level max attention)和 top-k 注意力累积占比。具体做法是统计每个查询行里权重最高的 5 个键 token 的权重之和,算它在整行总和里的比例。这个比例越高,说明该查询的注意力越集中。
第二,跨层对齐分析。把每一层的视觉 token 高注意力区域提取出来,做一个“注意力轨迹”折线图,横轴是层数,纵轴是某一类 token(比如物体中心区域)的平均权重。很多多模态模型呈现“浅层分散、中层提升、深层集中”的轨迹。这种轨迹分析比单看某一层的热力图更能说明问题。
第三,注意力熵。计算每行注意力分布的熵值,熵越低表示注意力越集中。这个指标非常有用,因为生成阶段和推理阶段的熵变化曲线能清晰地反映模型的信息获取模式。我在实验中发现,一个理解类问题的平均注意力熵远低于一个开放式的描述生成任务,因为前者只需要聚焦几个关键 token,后者需要把整张图的信息逐步用于文本生成。
5. 实测结果:理解任务的注意力分布特征
5.1 理解任务的整体模式:浅层分散,深层聚焦
先看一组代表性的理解任务实验结果。输入是一张包含两个物体(一个红色杯子、一个蓝色笔记本)的桌面图片,文本问题是“红色的杯子在蓝色笔记本的哪一侧?”模型需要输出“左侧”或者“右侧”。
这张图共产生 257 个视觉 token。文本查询包括问题 token 和特殊 token,共约 30 个 token。我在模型的 24 层 Transformer 里每一层都导出了交叉注意力矩阵,然后做单头平均。
观察到的规律非常清晰:浅层(第 1~6 层)交叉注意力的分布相对分散,top-5 视觉 token 的权重之和大约只有 0.4,注意力热力图几乎均匀铺在所有 patch 上;中层(第 7~14 层)开始出现明显集中,top-5 权重之和上升到 0.65,高注意力区域主要落在图像中物体边缘和纹理丰富的位置;深层(第 15~24 层)分布最集中,top-5 权重之和超过 0.85,而且这几个高权重 token 直接对应到红色杯子的位置区域。
这个规律符合直觉:浅层在做特征提取和全局扫描,深层在做最终对齐。但对缓存优化来说,这个现象意味着什么?意味着如果你只保留深层交叉注意力的 KV,可能就足够满足下游任务;浅层的 KV 虽然参与计算,但很多 token 的权重长期不高,存在压缩空间。这为 UniCache 后续的“分层稀疏缓存”策略提供了直接依据。
5.2 一个反直觉的发现:高分辨率特征位置与注意力集中的关系
我做完这组实验之后,突然想看看另一个维度。如果把图像分成几个大的语义区域(比如左上、右上、左下、右下),注意力集中区域跟物体真实位置是否完全一致?结论是:并不总是。
有一次我输入一张图,主体物体位于图像右侧,但是模型深层交叉注意力的高权重区域却落在图像中央偏右的位置,而不是严格对应物体的所有像素。这个现象老实说一开始让我很困惑。后来我翻看了一些多模态模型的可解释性研究,发现原因在于视觉编码器(ViT)本身已经做过局部特征聚合了,而且 patch token 之间本身有自注意力连接,一个 token 可能已经编码了周围若干 patch 的上下文信息。所以交叉注意力层在跟文本查询对齐时,并不需要把所有权重都放在物体像素本身的 token 上,它可以“间接通过一个高聚合 token”来达成对齐。
这个发现对缓存优化的意义重大:在判断哪些视觉 token 值得保留时,不能只依据交叉注意力权重,还必须结合视觉编码器内部的自注意力连接结构。有些 token 虽然在交叉注意力里权重不高,但它在视觉自注意力图里是很多 patch 的“汇聚点”,删除它会造成大范围信息丢失。UniCache 后续的 token 保留策略,就得设计一个“双重重要性”评分:交叉注意力的直接权重 + 视觉自注意力图的中心性得分,两者加权相加来决定哪个 token 优先保留。
5.3 理解任务里文本侧的自注意力分布特点
交叉注意力看完了,再看看文本侧自注意力。理解任务的文本输入往往比较短(问题本身只有几十个 token),自注意力矩阵也小,我测出的模式是:高注意力权重集中在问题中的关键词周围,尤其是实体名词和关系词。比如“杯子”和“笔记本”这两个词,在所有 token 的自注意力平均权重里显著高于其他词。
这其实也解释了多模态模型为什么在复杂问题上容易出错:如果问题中包含多个实体,而文本自注意力没有充分建立实体之间的关联,交叉注意力就会出现“选择性丢失”——模型只关注了其中一个实体对应的图像区域,忽略了另一个。这种问题在长问句里特别常见。所以做多模态理解模型的调试时,我建议一定要同时看文本自注意力和图文交叉注意力,单独看任何一个都不能完整定位错误源。
6. 实测结果:生成任务的注意力动态演变
6.1 解码步长上的注意力漂移
生成任务比理解任务复杂得多,因为交叉注意力的分布会随着每个解码步变化。我让模型根据同一张图片生成一段描述,一共生成了约 100 个 token。我记录下每个解码步交叉注意力的 top-5 token 变化情况。
结果让我很感慨:生成文本的前 10 个 token 时,交叉注意力依然覆盖整张图像的多个区域,top-5 token 分别散落在不同空间位置;生成到中段的 30~60 个 token 时,注意力开始聚焦,集中到图像的局部语义区域,跟当前描述的实体直接对应;生成到后段(80~100 个 token),注意力分布再次分散,模型似乎在总结或补充全局信息。
这个“聚焦-扩散”周期在长文本生成中至少出现两到三次。它说明模型生成文本的策略是“先全局观察,再局部细化,然后重新全局扫描”。这个规律直接决定了 KV 缓存的使用模式。如果某个视觉 token 在第一次聚焦时被使用了,但后续没再被用到,它的 KV 在缓存中其实处于低价值状态;只有那些在多次聚焦期间反复出现的 token 才是真正需要长期驻留的“热点”。这是我给 UniCache 第二篇(缓存更新策略)埋下的伏笔。
6.2 注意力崩溃与重复生成的关系
除了正常的聚焦-扩散规律,我还观察到一种注意力崩溃现象。当模型生成的文本过长,某些 token 的解码步里交叉注意力出现了极端尖锐的分布——全部权重集中在某一个中性 token(比如分隔符、句号)对应的视觉区域上。这种异常分布出现之后,下一段生成内容往往开始重复。这跟传统语言模型里的“重复惩罚”问题不完全相同,更像是注意力退化:模型失去了对原文细节的定位能力,退回到用固定模式凑字。
遇到这种情况,可视化注意力分布能迅速定位到具体是哪个解码步出现异常。我在调试过程中发现,如果在生成阶段实时监测注意力熵,熵值突然从正常范围(比如 2~3)掉到 0.2 以下,这意味着注意力发生了崩塌,最好马上停止生成或调整采样参数。这算是一个工程上的应急技巧,后续 UniCache 还可能把它实现为生成过程中的自动监控指标。
6.3 图像 token 在不同解码层的“热度排名”
我想做一个更系统化的统计,而不仅仅是单看某个解码步。做法是:在生成任务的完整过程中,把每个视觉 token 在所有解码步、所有 head 上的注意力权重求平均,得到一个“热度分值”,按热度从高到低排序。然后把这个排名跟图像分割后的语义区域做对照。
结果发现,token 热度排名前 10% 的视觉 token 几乎全部落在图像的主体物体边缘和纹理高频区域;背景平坦区域(比如空白墙面、桌面)的 token 热度极低。第三层和第十七层的热度排名差异最大:浅层更偏向纹理高频区,深层更偏向语义核心区。这个热度排名就是后续缓存淘汰策略的原始数据源——把热度排名倒过来,后 50% 的 token 完全可以在生成时优先逐出缓存。
再补充一个有意思的现象:某一次我把图像稍微旋转了 5 度再输入模型,发现热度排名竟然基本稳定,核心 token 的位置没有明显漂移。这验证了一个结论:注意力热点的空间位置对轻微扰动具有一定的鲁棒性,缓存策略可以依赖历史统计热度做预判。
7. 工具链与可视化:把注意力分布画出来看
7.1 热力图、注意力轨迹、熵曲线的组合使用
注意力数据本质上是多维的:层数维度、head 维度、解码步维度、token 维度。要让人一眼看出规律,光靠一张热力图是不够的。我在 UniCache 里做了一套三件套的可视化方案。
第一件是热力图。适合观察单个层、单个 head 的注意力模式。x 轴和 y 轴分别排视觉 token 和文本 token。如果视觉 token 数量太多,建议先按空间位置重排(还原成图像的 2D 布局),再画热力图,这样能直接看到“模型在看图的哪个区域”。这种“空间热力图”比简单的线状图直观太多。
第二件是注意力轨迹图。横轴是层数,纵轴是特定语义区域的注意力平均权重。一条线画主体区域,一条线画背景区域,模型在不同层的偏好变化一目了然。我通常把理解任务和生成任务分开画轨迹,因为两者的变化趋势差异很大,混在一起会互相干扰。
第三件是注意力熵曲线。横轴是解码步(或层数),纵轴是每行注意力的熵值。这个曲线在生成任务里尤其有用,能捕捉到注意力崩溃的瞬间,还能展示生成初期的探索、中期的聚焦、后期的收尾这些阶段。我把这三个图放在同一个 HTML 页面里,用交互式图表库渲染,鼠标悬停可以查看具体 token 索引和权重数值。
7.2 自动生成 HTML 报告:五分钟看懂全部结果
手工画图当然也能看,但几十层、多头、多样本的数据量太大了,必须写脚本自动汇总。我设计了一个报告生成脚本,输入是注意力矩阵的 npy 文件,输出是一份独立的 HTML 报告,包含以下区块:
- 模型配置摘要:层数、head 数、视觉 token 大小
- 全局注意力热力图:按层平均后的总览图
- 每层 top-5 视觉 token 索引和位置
- 多头注意力对比矩阵
- 生成任务的逐解码步注意力动态摘要
- 注意力熵汇总折线图
报告用 Jinja2 模板渲染,原始数据存为 JSON。跑一次完整实验大约需要 10 分钟(前向推理 2 分钟,注意力导出 3 分钟,报告生成 5 分钟)。这里的关键是把注意力导出的代码和报告生成代码解耦,这样换模型、换数据时不用重写可视化逻辑。
7.3 可视化时的常见踩坑
可视化看起来简单,踩坑不少。先说最大的坑:注意力维度搞混。交叉注意力的 shape 是(文本长度, 视觉长度),而自注意力是(文本长度, 文本长度)或(视觉长度, 视觉长度)。如果不注意,把交叉注意力的最后一维当成图像宽高来画,热力图就完全乱掉了。
第二个坑是 softmax 后的注意力权重在某些实现里会包含 mask 的影响。因果 mask 会把未来位置的权重强制置零,这些零值在统计 top-k 时会干扰结果。正确的做法是统计时先把 mask 位置的权重挖掉再归一化。
还有一个坑是浮点数精度。注意力权重在 softmax 之后通常很小,很多值低于 1e-5,直接展示在热力图上会显得一片空白。我建议做对数变换:log_weights = np.log(weights + 1e-6),这样能放大低权重区域的变化,细节更清晰。不过注意,对数变换只用于可视化,不能用于统计分析,否则均值、熵这些指标都会失真。
8. 注意力分布规律如何影响 UniCache 缓存设计
8.1 Hot Token 与 Cold Token:缓存的天然分界线
说完可视化,回到 UniCache 的核心目的:缓存优化。把上面所有实验观测汇总起来,我得出的最重要结论是:多模态模型中的视觉 token 存在明显的“热度分层”现象。
我把生成任务里全程注意力权重排名前 20% 的 token 称为 Hot Token,后 50% 称为 Cold Token,中间 30% 是温带 token。以我的实验数据为例,Hot Token 贡献了交叉注意力总权重的 60% 以上,而 Cold Token 的贡献只有不到 5%。这意味着在 KV Cache 中保留全部视觉 token 的键值向量,其实是一种巨大的浪费。
所以 UniCache 的设计思路很直接:对 Hot Token 的 KV,保持全精度驻留内存;对 Cold Token 的 KV,则可以选择丢弃、降精度(比如 FP16 转 INT8)、或者移到慢速存储,反正它们对生成质量的贡献微乎其微。
8.2 分层差异化缓存策略
注意力分布还有一个关键特征:不同层级的 Hot Token 集合并不完全一致。前面提到了浅层和中深层的关注点有差异,如果一个 token 在浅层很热但在深层很冷,是否值得缓存?我的答案是要看具体层。如果只是浅层需要该 token,那可以在浅层保留,在深层丢弃;如果缓存是跨层共享的,则需要按照“叠加热度”来给 token 排序。
我在实验里还发现,Hot Token 集合的层间重叠率大概在 40% 到 60%。也就是说,不同层关注同一组 token 的比例并不高。这给缓存设计带来了一个隐含约束:如果做跨层缓存压缩(例如多层共享一份 KV),必须接受部分层的信息损失;否则就需要保持分层独立缓存,但这样内存又会翻倍。UniCache 倾向选择“分层独立缓存 + 按热度动态淘汰”的方案,因为多模态模型交叉注意力的计算强度没有高到需要强行共享缓存。
8.3 与现有 KV Cache 机制的兼容性
最后聊聊工程落地有关的问题。现在主流推理框架里,KV Cache 已经是一个成熟的机制:每生成一个新 token,就把它的 Key 和 Value 追加到缓存数组里。但大部分实现是“无差别缓存”,也就是不管这个 token 重要不重要,全都存下来。这显然是浪费的,因为解码步数增加,缓存占用线性增长,显存吃不消。
UniCache 想要改造的是:在追加新 token 的 KV 之前,先根据注意力热度评估它是否值得进入缓存;如果缓存满了,就把当前热度最低的旧 token KV 请出去。这套机制完全兼容现有的 PyTorch 推理循环,只需要在 KV Cache 数据结构上增加一个热度字段,并且自定义一个淘汰函数即可。
更具体的设想是:用一个小型寄存器记录每个 token 的累计注意力权重,当累计权重低于阈值时,该 token 的 KV 会被标记为“可回收”。在下一次缓存扩容时,优先回收这些条目。这个方案比基于 LRU 的常规缓存淘汰策略更聪明,因为 LRU 只考虑访问时间,而注意力热度是语义重要性的直接度量。
9. 常见问题与排查技巧实录
9.1 为什么导出的注意力权重全是 NaN
这是一个很常见的问题。我在换到某个开源多模态模型时,导出的注意力矩阵突然全变成 NaN。排查过程花了不少时间,最终定位到原因:那个模型的交叉注意力层对视觉 token 的 mask 处理不对,某些 token 的注意力分数在 softmax 之前被加上了-inf,导致分母为 0,softmax 输出 NaN。
解决办法:检查输入数据的 attention mask,确保 mask 的形状和内容符合模型的预期。另外,如果使用 FP16 进行推理,softmax 里的浮点溢出也可能导致 NaN,稳妥的做法是把注意力导出的操作改为 FP32。
9.2 注意力热力图“全是一片红”怎么办
如果你看到的热力图颜色一片深红或者全白,那是因为注意力权重范围没有归一化。刚才说过,softmax 后的权重数值差异可能达到几个数量级。解决办法有三个:一是做对数归一化;二是 per-row 归一化,也就是每一行单独做最大最小缩放;三是只展示 top-k 的前 k 个位置的权重,其他位置全部置为灰色。三种方案各有优劣,对数变换保留细节,min-max 对比更强,top-k 展示更清晰。我通常先看 top-k 图了解大致情况,再用对数变换图看细节。
9.3 层数太多,热力图根本看不过来
24 层甚至更多层的模型,每层一张热力图,确实没法快速浏览。这时候我把热力图按层聚类成三四组,每组画一张平均图。聚类方法用 KMeans 就行,注意先做层内归一化再聚类。或者你可以按注意力分布的结构特征(比如熵值、尖峰度)给每层打分,然后把打分接近的层合并报表。这样能在“保留细节”和“快速概览”之间取得平衡。
9.4 Hot Token 统计结果不稳定
同一张图,同一个模型,多次运行后统计的 Hot Token 排名会有小幅波动。这主要来自采样策略:如果解码阶段用了随机采样而不是贪心,模型生成的 token 序列每次不一样,交叉注意力的具体分布自然会有变化。如果你想要稳定的注意力分布来指导缓存策略设计,实验阶段必须使用贪心解码,并且固定随机种子;等拿到稳定的 Hot Token 排名之后,再去实际生成任务里验证效果。
10. 在注意力分布的边缘处,还有几个细节值得多想
说一个我后来才领悟到的细节。当大家都在看注意力权重高的地方时,那些权重极低的 Cold Token 并不是完全无用。我有一次实验里把注意力权重低于 0.01 的视觉 token 全部从缓存里剔除,结果生成质量没有明显变化。但当我继续把阈值提高到 0.05 的时候,开始出现轻微的描述漏细节现象,尤其是一些小物体的颜色属性会被漏掉。这说明 Cold Token 并非信息量为零,它们更接近“低频但必要的背景信息”。
所以在 UniCache 的后续规划里,我不会一刀切地把所有 Cold Token 都删掉,而是设计出一个“多层次保留”策略:Hot Token 全精度驻留,温带 Token 做 INT8 量化,Cold Token 大部分淘汰但保留每 16 个 token 里随机选一个作为“背景锚点”。这样既保证了主要语义信息不丢,又能显著减少缓存占用。
还有一个值得多想的点:注意力分布跟任务难度有关系。我拿同样一张图,问了三个难度递增的问题(“图里有什么”、“图中物体的颜色”、“如果移走某个物体会发生什么”),发现注意力分布在第三个问题上出现了更强的全局扫描模式,Hot Token 的数量明显增加。这说明缓存策略不能只按视觉 token 的固有热度设计,还要考虑当前任务的复杂度。一个在简单问题上可以被淘汰的 token,在复杂推理问题上可能突然变得必要。这块需要更复杂的动态缓存策略才能解决,需要用在后续的文章里慢慢展开。
在做 UniCache 后半部分工作的时候,我逐渐意识到一个事实:注意力分布不是静态的模型属性,它会随着输入、任务和解码过程动态变化。所以任何试图用一个固定阈值去做缓存剪枝的方案,都有可能在某个特定场景下失手。更稳妥的做法是做一个轻量级的“注意力分布预测器”——输入当前文本 token 和少量视觉 token 的历史统计特征,预测下一个解码步可能的高注意力区域。这个想法还不成熟,但值得继续深挖。
最终,我在实践中最大的体会是:注意力分布纵览不是终点,它是起点。只有先把模型的注意力模式研究透,才有资格去谈优化和压缩。UniCache 这个系列,第一篇做的就是这个基础工作。下一篇我会带着这些注意力分布特征,去设计一个针对多模态生成任务的 KV 缓存淘汰策略,敬请期待。