1. 从一次显存告急说起:为什么要盯住注意力分布
做多模态推理优化的人,大概都经历过这种时刻:一张图加一段长 prompt 丢进去,模型还没开始吐字,显存先红了。你去看显存占用,权重只占了一部分,剩下的大头全在 KV cache 上。尤其是多模态场景,图像经过视觉编码器之后动辄产生几百上千个 visual token,再叠上文本 token,序列长度轻松翻倍。这时候如果还按文本模型那套"每层都缓存、每个 token 都保留"的思路去做,显存和带宽都会被吃干抹净。
UniCache 这个系列要解决的就是这件事。它的核心命题很直接:多模态理解与生成模型里,注意力到底是怎么分布的,哪些 KV 是真正被反复读到的,哪些其实可以丢。只有先把注意力分布摸清楚,后面的缓存淘汰、复用、压缩策略才有依据。这篇是系列的第一篇,定位是"纵览",也就是把理解模型和生成模型这两大类多模态架构的注意力行为先摊开来看一遍,建立一个统一的观察框架。
我写这篇的出发点很实际。过去一年我在本地和集群上跑过不少多模态模型,从早期的 CLIP 类双塔,到后来的 LLaVA 式拼接结构,再到 DiT 这类生成架构,一个共同的感受是:理解模型和生成模型的注意力分布差异极大,但很多缓存优化方案却把它们当成一回事。结果就是理解任务上效果还行,一换到生成任务就崩,或者反过来。所以这篇不急着讲具体怎么压缩,而是先把"分布长什么样"这件事讲透。
适合谁看?如果你正在做多模态推理加速、KV cache 管理、长上下文优化,或者只是单纯好奇为什么多模态模型这么吃显存,这篇都能给你一个可复现的观察视角。我会尽量用能跑起来的代码和实测数据说话,而不是停留在概念层面。
2. 多模态注意力分布的基本盘:先搞清楚在看什么
2.1 理解模型与生成模型的注意力结构差异
要谈分布,先得把结构对齐。多模态理解模型的典型代表是 LLaVA 系列:视觉编码器(通常是 ViT)把图像切成 patch,输出一批 visual token,然后通过一个投影层(MLP 或 Q-Former)映射到语言模型的 embedding 空间,和文本 token 拼在一起送进 decoder。这里的注意力是因果的,每个 token 只能看到自己之前的 token,visual token 一般放在序列前部,文本 token 在后。
生成模型就完全是另一套逻辑。以 DiT(Diffusion Transformer)和各类多模态生成架构为例,图像 token 和文本 token 之间往往是双向注意力,或者至少在同一模态内部是双向的。扩散过程还要反复迭代去噪,同一个序列会被前向计算几十次。这意味着生成模型的注意力分布不是"一次成型",而是随去噪步数动态变化的。
这个结构差异直接决定了 KV cache 的行为。理解模型里,KV 是自回归逐 token 增长的,缓存是"追加式"的;生成模型里,如果做缓存,更多是"跨步复用",同一位置的 KV 在不同去噪步之间被反复读取。两者对"哪些 KV 重要"的答案完全不同。
2.2 注意力分布到底在观察什么指标
很多人一说注意力分布,就只想到 softmax 之后的注意力权重矩阵。但真正做缓存优化,光看权重是不够的。我一般会同时观察四个维度:
- 注意力熵:某一层某个 head 的注意力分布有多集中。熵低说明这个 head 只盯着少数几个 token,这些 token 就是缓存的重点。
- KV 命中频次:在自回归解码过程中,某个位置的 KV 被后续 token 读取了多少次。频次高的位置,淘汰代价大。
- 层间差异:浅层和深层的注意力模式往往截然不同。浅层偏向局部和位置,深层偏向语义聚合。
- 模态间注意力流向:文本 token 对 visual token 的注意力占比,以及反向的流向。这决定了跨模态缓存能不能分开管理。
这四个指标里,KV 命中频次是最贴近缓存决策的,因为它直接对应"这个 KV 被读了几次"。注意力熵则更像是一个先验,帮你判断哪些 head 值得重点保护。实际做的时候,我建议先把命中频次统计出来,再回头看熵,两者交叉验证。
2.3 为什么不能直接套用文本模型的结论
文本模型上有一套很成熟的观察:attention sink(前几个 token 吸收大量注意力)、局部性(近期 token 更重要)、层间稀疏性差异。这些结论搬到多模态上,一部分成立,一部分会翻车。
最典型的翻车点是 attention sink。文本模型里,序列开头的 BOS token 经常成为 sink,几乎所有 head 都会分一部分注意力给它。但在多模态理解模型里,序列开头是 visual token,情况就复杂了:有些 head 确实会盯着某几个"锚点"visual token,但另一些 head 的注意力会均匀铺在整片图像区域上。如果你按文本模型的经验,粗暴地只保留前几个 token 的 KV,图像信息会丢得很惨。
另一个差异是位置编码的干扰。多模态模型里 visual token 的位置编码方式和文本不同,有的是二维位置编码,有的是插值后的一维编码。这会导致注意力分布在 visual 区域内部呈现出和文本完全不同的几何结构。所以我的态度很明确:文本模型的缓存经验可以参考,但必须重新验证,不能直接照搬。
3. 理解模型的注意力分布:visual token 到底被谁在读
3.1 LLaVA 式拼接结构的实测观察
我拿一个标准的 LLaVA 结构做了实测,图像分辨率 336,ViT 输出 576 个 visual token,加上一段约 200 token 的文本 prompt,总序列长度接近 800。用一张信息量中等的图片(比如一张有多个物体的街景),跑一次完整前向,把每层的注意力权重和 KV 读取情况都记录下来。
第一个发现是:visual token 内部的注意力高度不均匀。576 个 visual token 里,真正被文本 token 大量读取的,集中在少数几十个位置上。这些位置往往对应图像里的显著物体或高对比度区域。换句话说,图像被"看"的时候,模型并不是均匀地扫描,而是有选择地聚焦。
第二个发现是层间分化非常明显。浅层(前 1/3)的注意力基本是局部的,visual token 主要看自己附近的邻居,文本 token 也主要看邻近的 visual token。到了中层,跨模态注意力开始活跃,文本 token 大量读取视觉信息。深层则反过来,visual token 开始聚合全局语义,注意力铺得更开。
这个分层特征对缓存策略的启示很直接:浅层和深层的 KV 重要性判断标准应该不一样。浅层可以更激进地按局部性淘汰,深层则要保守一些,因为深层 KV 承载的是聚合后的语义。
3.2 视觉 token 的注意力锚点现象
在实测里我注意到一个很有意思的现象,我把它叫做视觉锚点。在每一层里,总有那么几个 visual token,它们的 KV 被读取的频次远高于其他 token,而且这些锚点在不同层之间有一定的延续性——浅层的锚点位置,往往在中层还是锚点。
这跟文本模型的 attention sink 有点像,但成因不同。文本 sink 更多是位置编码和 softmax 归一化导致的"无处可去"的注意力倾倒;视觉锚点则更像是模型真的在用它作为图像信息的"索引"。我猜测这和 ViT 的 patch 聚合方式有关,某些 patch 在编码阶段就吸收了更多信息,到了语言模型里自然成为被反复查询的对象。
实操上的价值是:如果能稳定地识别出这些锚点,就可以给它们开小灶,优先保留。识别方法也不复杂,跑一批样本,统计每个 visual token 位置在各层的平均读取频次,取 top-k 作为锚点候选。我在几个不同图片上验证过,锚点位置有一定稳定性,但不是完全固定,所以最好按任务类型分别统计。
3.3 文本到视觉的注意力流向分析
跨模态注意力流向是理解模型缓存优化的关键。我统计了文本 token 对 visual token 的注意力占比,发现一个规律:不是所有文本 token 都同等程度地"看"图像。
具体来说,实体词、名词短语、形容词这些承载视觉语义的 token,对 visual token 的注意力占比很高;而功能词、标点、连接词,注意力基本留在文本内部。这意味着如果你要做跨模态的缓存压缩,可以按 token 类型区别对待:视觉相关的文本 token 需要保留对视觉 KV 的访问能力,纯功能性 token 则可以更激进地压缩。
还有一个细节值得注意:文本对视觉的注意力在生成回复的不同阶段是变化的。当模型开始描述图像内容时,对 visual token 的读取会突然增强;当它在做纯语言推理时,读取会减弱。这个动态特征说明,静态的缓存策略可能不是最优的,理想情况下应该根据解码阶段动态调整。
4. 生成模型的注意力分布:跨步复用才是重点
4.1 扩散架构下的注意力动态
生成模型这边,我主要观察了 DiT 类架构。和自回归理解模型最大的不同是,扩散模型的注意力是跨去噪步的。同一个序列,在 50 步去噪里会被前向计算 50 次,每次的注意力分布都不一样。
实测下来,去噪早期(高噪声阶段)的注意力比较弥散,模型还在探索整体结构;去噪后期(低噪声阶段)注意力变得集中,模型在打磨细节。这个变化对缓存的意义是:早期步的 KV 复用价值低,因为分布还在剧烈变化;后期步的 KV 复用价值高,因为分布趋于稳定。
我试过一个简单的策略:只在去噪后 1/3 的步数里启用 KV 缓存复用,前面的步数每步重算。结果在几乎不损失生成质量的前提下,省下了可观的计算量。这个思路和很多缓存方案"全程复用"的默认假设是相反的,但实测更稳。
4.2 图像 token 与文本 token 的双向注意力特征
生成模型里,图像 token 和文本 token 之间通常是双向注意力。我观察到的一个特征是:文本 token 对图像 token 的注意力,在去噪过程中相对稳定;而图像 token 对文本 token 的注意力,波动较大。
这个不对称性很有意思。文本作为条件,它的语义是固定的,所以图像去查询文本的模式比较稳定;而图像本身在去噪中不断变化,它查询文本的方式也就跟着变。对缓存的启示是:文本侧的 KV 可以更放心地跨步复用,图像侧的 KV 要更谨慎。
另外,生成模型里图像 token 数量通常远大于文本 token(比如 1024 个图像 token 对 77 个文本 token),所以缓存压力主要在图像侧。这就更需要精细地判断哪些图像 token 的 KV 值得留。
4.3 跨步 KV 复用的可行性边界
跨步复用听起来很美,但边界在哪?我做了几组对比实验,结论是:复用能不能成立,取决于相邻去噪步之间注意力分布的相似度。
我用余弦相似度衡量相邻步的注意力分布,发现在高噪声阶段,相邻步相似度可能只有 0.3 到 0.5,这时候复用基本等于瞎猜;到了低噪声阶段,相似度能到 0.8 以上,复用就很安全。所以一个可操作的判据是:设定一个相似度阈值(比如 0.7),只在超过阈值的步之间复用 KV。
这个阈值不是拍脑袋定的,需要按模型和任务调。我一般会先跑一批样本,画出相似度随步数变化的曲线,找到那个"拐点",把阈值设在拐点之后。这样既保证了质量,又最大化了复用收益。
5. 统一观察框架:把两类模型的分布放在一起看
5.1 理解与生成的分布对比表
把前面两部分的观察整理成一张表,能更清楚地看到差异:
| 观察维度 | 多模态理解模型 | 多模态生成模型 |
|---|---|---|
| 注意力方向 | 因果(单向) | 双向或模态内双向 |
| KV 增长方式 | 自回归追加 | 跨去噪步复用 |
| 关键 token | 视觉锚点 + 视觉相关文本 | 稳定文本条件 + 后期图像 token |
| 分布稳定性 | 随解码位置变化 | 随去噪步数变化 |
| 缓存优化重点 | 淘汰低频 visual token | 判断跨步复用边界 |
| 浅层特征 | 局部、位置主导 | 结构探索、弥散 |
| 深层特征 | 全局语义聚合 | 细节打磨、集中 |
这张表是我做后续缓存策略设计的基础。它说明一件事:不存在一套通用的多模态缓存方案,必须按模型类型分而治之。UniCache 这个系列后面要讲的统一缓存框架,也是建立在这个"分而治之再统一抽象"的思路上的。
5.2 一个可复用的注意力统计脚本
光看结论不够,得能自己跑。我整理了一个最小可用的统计脚本,基于 PyTorch hook 抓取每层注意力,统计 KV 读取频次和注意力熵。核心逻辑如下:
import torch from collections import defaultdict class AttentionStats: def __init__(self): self.read_count = defaultdict(lambda: None) self.entropy = defaultdict(list) def hook_fn(self, layer_idx): def fn(module, inputs, outputs): # outputs 里通常包含 attention weights attn = outputs[1] if isinstance(outputs, tuple) else None if attn is None: return # attn: [batch, heads, q_len, kv_len] with torch.no_grad(): # 统计每个 kv 位置被读取的总次数 read = attn.sum(dim=(0, 1, 2)) # [kv_len] if self.read_count[layer_idx] is None: self.read_count[layer_idx] = read.cpu() else: self.read_count[layer_idx] += read.cpu() # 计算每个 head 的注意力熵 p = attn.clamp(min=1e-9) ent = -(p * p.log()).sum(dim=-1).mean().item() self.entropy[layer_idx].append(ent) return fn stats = AttentionStats() handles = [] for idx, layer in enumerate(model.layers): handles.append(layer.self_attn.register_forward_hook(stats.hook_fn(idx)))跑完之后,stats.read_count里就是每层的 KV 读取频次,stats.entropy是每层的平均注意力熵。这两个量结合起来,就能画出分布图,判断哪些位置值得保留。
注意:不同框架的 attention 模块输出格式不一样,有的返回 attention weights,有的默认不返回(需要开启
output_attentions=True)。抓之前先确认一下,否则 hook 拿到的可能是 None。
5.3 统计结果怎么读
拿到数据之后,怎么读很关键。我一般会做三件事:
第一,画 KV 读取频次的热力图,横轴是 token 位置,纵轴是层。这样一眼就能看出哪些位置在哪些层被重点读取。视觉锚点会表现为几条竖直的高亮带。
第二,画注意力熵的层间曲线。熵低的层说明注意力集中,这些层的 KV 更"专一",淘汰风险高;熵高的层注意力分散,KV 冗余度大,可以更激进地压缩。
第三,做模态分组统计。把 visual token 和文本 token 分开算平均读取频次,看跨模态注意力的整体强度。如果文本对视觉的读取占比很低,说明这个任务其实不太依赖视觉,缓存可以大幅压缩视觉侧。
这三步做完,你对一个模型的注意力分布就有了量化认识,后面做任何缓存决策都有据可依,而不是凭感觉。
6. 实操中的坑与排查经验
6.1 常见问题速查表
| 问题现象 | 可能原因 | 排查方向 |
|---|---|---|
| hook 抓不到 attention | 框架默认不返回权重 | 开启 output_attentions |
| 读取频次全为 0 | hook 挂错模块 | 确认挂在 self_attn 而非整个 block |
| 熵值异常低 | softmax 前未缩放 | 检查是否用了原始 logits |
| 视觉锚点不稳定 | 样本差异大 | 分任务类型分别统计 |
| 跨步复用后质量崩 | 阈值设太低 | 提高相似度阈值 |
| 显存没降 | 缓存没真正释放 | 检查是否只标记未释放 |
这张表是我踩坑之后总结的,基本覆盖了统计阶段最常见的几类问题。尤其是 hook 挂错模块这一条,我一开始挂在 transformer block 上,拿到的输出里根本没有 attention weights,白白折腾了半天。
6.2 几个容易忽略的细节
第一个细节是batch 维度。统计的时候如果 batch size 大于 1,不同样本的注意力分布会混在一起。做锚点分析时最好 batch size 设为 1,或者按样本分别统计再聚合,否则锚点位置会被平均掉。
第二个细节是精度。用 fp16 跑统计,注意力权重可能会有数值误差,熵的计算会偏。建议统计阶段用 fp32,虽然慢一点,但数据可靠。真要省时间,至少把 softmax 那一步放到 fp32 里算。
第三个细节是序列长度对齐。不同样本的 visual token 数量可能不同(如果做了动态分辨率),统计时要么固定分辨率,要么按比例归一化位置。我一般固定分辨率,省得后面解释麻烦。
6.3 我个人的几条经验
做了这么多轮统计,有几条经验我觉得比技术细节更重要。
别一上来就追求全模型统计。先挑几层代表性的(比如第 1、中间层、最后一层)看,建立直觉,再决定要不要全量跑。全量统计又慢又占内存,很多时候没必要。
分布是会随 prompt 变的。同一个模型,换一种提问方式,注意力分布可能就不一样。所以统计要覆盖多样化的 prompt,不能只用一两个样本就下结论。我一般至少准备 20 到 30 个不同类型的样本。
把统计和实际缓存实验结合起来。光看分布不够,得真的按分布去淘汰 KV,看效果掉多少。分布告诉你"哪些可能重要",实验告诉你"淘汰了到底行不行"。两者结合,结论才靠谱。
7. 从分布到策略:下一步该往哪走
把理解模型和生成模型的注意力分布摊开看之后,一个自然的下一步就是:怎么把这些观察转化成可执行的缓存策略。这篇作为纵览,先把地基打好,具体的淘汰算法、复用调度、跨模态缓存分区,会在后续展开。
但有一个方向我现在就可以说:统一缓存框架的关键,不在于找一套万能规则,而在于建立一个能自适应判断"当前该用哪种规则"的机制。理解模型和生成模型的分布差异这么大,硬套一套规则必然顾此失彼。更现实的做法是,先识别当前处于什么模式(理解还是生成、高噪声还是低噪声、浅层还是深层),再切换到对应的缓存策略。
我在实际项目里试过一个简化版的自适应逻辑:用一个轻量的相似度检测判断当前步的注意力分布是否稳定,稳定就复用,不稳定就重算。这个逻辑在理解模型和生成模型上都能跑,虽然还不够精细,但至少证明了"自适应"这条路是通的。
最后分享一个我常用的调试技巧:把注意力分布可视化出来,存成图片,按层和步数排成网格。很多时候数值看不出的模式,一眼图就看出来了。尤其是跨模态注意力的流向,画出来比看表格直观得多。这个习惯帮我省了很多瞎猜的时间。