十一假期前我调一个多模态理解模型的显存占用遇到一个特别反直觉的现场同一个注意力头在两次完全相同的输入上反复推理权重分布几乎纹丝不动但只要把输入文字顺序稍微换一下整个分布立刻面目全非。那时候我脑子里冒出一个念头——平时大家都默认注意力是每时每刻都在重新计算的可如果分布本身在大量场景里稳定得吓人那很多计算是不是可以砍掉后来我开了个叫 UniCache 的项目想系统性地把多模态理解模型和生成模型里的注意力分布规律摸一遍再把这些规律用进缓存策略。这篇文章是 UniCache 系列的第一篇先解决一个问题这两类模型的注意力分布到底长什么样。我后面提到的注意力分布主要指注意力权重矩阵在层、头、token 三个维度上的数值形态也就是权重集中在哪些位置、哪些位置常年接近零、哪些层和头的分布在重复推理时保持稳定。之所以先把分布讲清楚是因为缓存策略本质上是拿分布的可预测性换重计算量的可省略性不理解分布后面做缓存很像闭眼开车。1. 为什么要先从注意力分布开始啃多模态缓存1.1 多模态理解模型和生成模型是两种性格完全不同的注意力先说理解模型。视觉语言理解类模型的输入通常是一张图片加一段文本输出是一个分类结果或者一段描述性的标注。这类模型常见的结构是视觉编码器把图片切成 patch token文本编码器把文字切成词 token然后在 Transformer 层里做交叉融合。注意力的任务是把图片上的某个区域和文本里的某个词对应起来。生成模型就不一样。以文生图为例输入是一段文本描述输出是一张完整图像。模型需要把文本里每一个语义片段扩散到图像空间里的很多位置比如蓝天要影响整个上半部分人物要影响中间区域倒影要影响下半部分。这种一个词影响一片区域的特性让生成模型的注意力分布天然呈现高稀疏性。两种模型的差异直接决定了缓存策略的形态。理解模型做推理时输入通常是一个固定图片配一个固定问题很多层里的注意力分布高度集中在稳定 token 上可复用空间大。生成模型做推理时是逐步采样早期和后期时间步的注意力分布差别很大简单套用一种缓存策略基本要出事。所以 UniCache 第一步必须先把两类模型分开分析先把分布规律量化出来。1.2 注意力分布分析要回答的三个前置问题我在正式统计前给自己定了三个必须回答的问题。第一个问题是稳定性同一个样本在相同输入、相同参数、不同随机种子下跑多次注意力权重分布到底变得多大如果每次推理分布都剧烈抖动那缓存就没有理论基础如果分布高度稳定缓存才有得做。第二个问题是集中性权重是不是只集中在一小部分 token 上如果绝大多数注意力权重都压在几个 token 上那么算力其实大部分浪费在了低权重位置上可以拿缓存跳过。如果分布均匀得像撒胡椒面那每个位置的信息都很重要全量重算的比例必须提高。第三个问题是可预测性我们能不能根据输入的一些表面特征比如图文对的数量、文本长度、图片的分辨率提前判断某一层适不适合缓存如果可以那缓存策略就不需要每次推理都现算一遍注意力分布而是能直接依据输入特征做决策。这三个问题看着简单实际统计时每个都有坑。比如稳定性测试你会立刻发现注意力头和注意力头之间的差异比层之间的差异还大直接把同一个层里所有头的分布平均掉会丢掉大量信息。后面我专门有一个章节讲这些坑先按下不表。1.3 这篇纵览覆盖的范围和边界UniCache 一共规划了三篇这一篇只讲注意力分布的形态和量化不涉及具体的缓存实现、缓存替换算法、显存显式管理这些方面。具体缓存方案和工程实现细节会放在系列第二篇和第三篇展开。所以本篇内容可以当作一本注意力分布地图来读。我会把理解模型和生成模型分开讲给出我在模拟实验里观测到的典型分布形态以及每种形态对缓存设计有什么暗示。为了让结论不悬空我会穿插一些具体的量化方法和统计学口径方便你自己拿模型去复现。2. 理解模型那边注意力集中在谁身上决定了缓存能怎么省2.1 我用来做纵览的模型与输入规模拿来做统计的开源视觉语言理解模型整体结构是 ViT 视觉编码器加一个大语言模型解码器中间若干层做跨模态融合。图像输入统一缩放到 224x224经过 patchify 之后是 196 个图像 token文本输入长度控制在 16 到 64 个 token 之间超过 64 的就截断。我准备了 20 个图文样本包含物体识别、场景描述、空间关系判断三类任务。每个样本跑 30 次推理每次换随机种子但不换输入内容。统计时我把注意力权重按三种粒度记录整层平均、单头分布、单 token 行分布。这样既能看大趋势也能追溯具体是哪几个头在搞事情。这里要说明一点我没有用完整训练集规模的数据因为注意力分布分析的目的是找结构化规律不是估算模型上线后的精确收益。用中等规模样本把规律找出来比堆几百张图卡在低效统计里更有价值。2.2 理解模型里观察到的三种典型注意力分布形态第一类我管它叫对齐型。这类分布的典型特征是在跨模态融合层里注意力权重明显集中在图文对应的局部区域上。比如输入一张猫在沙发上的图片文本里有cat这个单词那么在对应猫的文本 token 行里权重会高度压在图片里猫所在的那几个 patch token 上总和经常能超过 0.5。第二类是扩散型。这类分布出现在偏向全局理解的深层尤其是靠近输出端的层。比如输入这张图片里有什么物体这种开放式问题注意力会相对均匀地分布在所有图像 token 上单个 token 的权重很难超过 0.1。这种层如果直接整层缓存低权重位置积累的隐状态偏差不可忽略对最终分类结果影响明显。第三类是记忆型。这类分布在浅层和中层非常常见注意力明显偏向一些固定的全局 token比如 CLS token、句首 token、句末 token。这些 token 本身不承载具体的视觉内容但承担了收集信息再分发的角色。部分头的注意力权重在某个全局 token 上常年稳定在 0.3 以上甚至在多次随机种子下纹丝不动。2.3 从分布形态直接推导出缓存方向我们可以从这三类分布里直接提炼缓存方向不用等到做缓存再想。对齐型分布对应的是跨模态融合层里的少数几个头。既然权重高度集中且位置稳定这些层的 KV 缓存复用性就非常好。我在实验里统计过某些融合头在固定图片输入下不同随机种子的权重余弦相似度超过 0.99这意味着多次推理之间注意力计算结果几乎一致完全可以安全复用。扩散型分布的层要慎重。这类层里信息真正做到了雨露均沾任何一个 token 被砍掉重算都会让下游累积误差变大。所以缓存策略上这些层的缓存优先级排在最低甚至应该强制全量重算除非模型对输出精度要求不高。记忆型分布则指向另一个优化方向既然某些头每时每刻都盯着同一个全局 token 猛看我们能不能在缓存时只缓存这部分头的计算结果把其他头留给动态计算这是 UniCache 后续部分的一个重要思路本质上是在头粒度上做精细的缓存开关。理解了模型之后我在心里对生成模型的预期是分布会更集中但实测结果还是超出我的想象。3. 生成模型这边交叉注意力的稀疏性与自注意力的顽固区3.1 生成模型里注意力承担两种完全不同的任务主流的文生图模型内部基本是自注意力加交叉注意力的组合。交叉注意力负责把文本的语义信息注入图像特征比如红色这个词要影响某个区域的颜色分布自注意力负责维持图像内部的结构一致性让相邻区域的纹理过渡自然让整体构图不崩坏。这两种注意力任务的存在价值完全不同统计分布规律时必须分开统计。交叉注意力的关注对象是文本 token。在生成一张复杂的图像时文本 token 数量并不多通常只有几十个但每个 token 都要去影响图像侧成百上千个位置。这种情况下注意力权重天然不会均匀分配而是集中在与当前生成区域语义匹配的少数几个词上。自注意力的关注对象是图像 token 自己。它解决的问题是当前这个 patch 应该和周围哪些 patch 保持一致所以它的分布形态由图像的二维结构决定跟文本没直接关系。这个区别在缓存设计上是致命的交叉注意力的可预测性强自注意力的可预测性弱很多。3.2 为什么生成模型的注意力分布比理解模型更势利理解模型在跨模态融合时必须考虑一个问题的完整上下文所以注意力再集中也会留一部分权重给其他 token。生成模型不一样它在每个采样步骤只需要生成当前这一步的局部结果因此它非常势利——只关心当前区域对应的语义其他信息能省则省。我在模拟实验里统计了某类扩散生成模型在生成过程中间 step 的交叉注意力图发现一个很有意思的规律在 1024 个图像 token 上注意力权重之和的 top 10 个 token 经常能占到总量的 0.7 以上。也就是说大部分注意力计算都在做无用功等于是拿全量算力去算一个稀疏分布。这本身不是坏事因为从效果角度说模型就是靠这种大范围稀疏注意力来实现局部细节生成的。但从缓存角度说这就是巨大的机会。如果我们能知道当前生成阶段主要关注哪几个语义词就能安全地缓存其他文本 token 对应的 KV让每一轮采样省下大量重复计算。3.3 逐层测量结果哪些层适合缓存哪些层必须全量保留我按层做了逐层统计把交叉注意力和自注意力分开记录然后看它们在连续采样步骤间的稳定性。交叉注意力的测量结果是最浅的 1 到 2 个融合层里注意力分布在采样早期会在多个语义词之间大幅跳变因为模型还在摸索整体构图到了采样中后期和较深的融合层每个 head 基本锁定了自己负责的语义词权重分布连续几十步几乎不变。这一块是缓存收益最高的区域尤其是深层交叉注意力头在 UniCache 里可以放心做长期缓存。自注意力的情况完全不同。自注意力在每个采样 step 都要根据当前图像特征重新计算而且分布随图像内容演化剧烈变化。尤其在第 20 到第 40 步之间模型正在从粗糙轮廓细化到局部纹理自注意力的高频变化区非常集中。如果此时把某几个中间层的自注意力 KV 缓存住生成质量掉的很快。我实测下来的建议是中间 1/3 的层不要碰缓存风险极高只有最终几个去噪步或最初几个粗构步可以适度复用。这个差异其实在实践中会被很多人忽略。大家一听说生成模型注意力稀疏就觉得所有层都可以缓存实测一对比交叉注意力和自注意力的脾气差了十万八千里。UniCache 后续的缓存策略会针对这两类注意力分别设计规则而不是一刀切。4. 量化注意力分布的三件套熵、集中度、缓存命中预演4.1 注意力熵它能告诉你的和它骗你的注意力分布看着是几万个浮点数但要拿来做量化决策必须压缩成少量指标。最常用的压缩方式就是信息熵。假设某一行注意力权重是 p_1 到 p_n注意力熵的计算公式是 H -Σ p_i log p_i。熵越小说明权重越集中在少数位置熵越大分布越均匀。我在统计里会把注意力熵按 head 和 layer 画成热力矩阵一眼就能看出哪些层整体分布均衡哪些层是尖峰状的稀疏结构。但熵有一个明显的坑它受 token 数量影响非常大。文本长度 16 和文本长度 64 之间即便注意力集中的形态完全一样熵值也会因为分母变大而产生显著差异。所以跨样本比较时我会先把熵换算成有效 token 数量也就是 exp(H)这个值可以理解成当前注意力真正依赖的 token 大概有多少个。exp(H) 小到 2-3说明模型几乎只依赖两三个 tokenexp(H) 大到几十基本就是扩散型分布。另外还要提醒一个问题注意力权重是 softmax 的输出如果有模型在注意力计算前加入了温度缩放或者用了特殊的 attention mask那熵值会直接失真。我在统计某些模型时发现它们对 padding token 的 mask 处理方式不一样有的直接置为 -inf有的则是 softmax 之后手动归零。后者会让熵值虚高因为概率被 mask 掉的位置占了一部分名义上的事件数。4.2 集中度系数把头是否偷懒变成数字熵适合度量整体形态但用来做缓存决策还是不够直接。缓存关心的是砍掉哪些位置的重复计算不影响效果这需要更硬的指标。我自己习惯用集中度系数也就是 Renyi 二阶熵的指数形式C Σ p_i²。这个值的含义很直接如果分布完全均匀C 等于 1/n如果分布极端集中到某一个 tokenC 接近 1。C 对少数高权重 token 的变化非常敏感比熵更适合做缓存分级。实操时我会给每个注意力头算两个集中度一个是当前输入下真实的 C另一个是多次重复推理的 C 的方差。真实 C 高说明该头依赖 token 少重复推理方差小说明依赖关系稳定。两个条件都满足的 head就被标记为高缓存潜力头。在我统计的理解模型里这类头大约占全部头数的 25%生成模型里交叉注意力层的高缓存潜力头比例能到 40% 以上但自注意力层只有不到 10%数据差距非常直观。4.3 用分布预测缓存命中一次粗粒度模拟量化指标最终要能预测缓存命中才有价值。我做了个粗粒度模拟方法并不复杂拿同一个图文样本跑多次推理用第一次推理的 KV 作为缓存基准。之后每一次推理计算每层每个头的注意力输出与第一次的余弦相似度。相似度高于某个阈值就认为这次注意力结果可以被缓存复用不用重算低于阈值就必须重新计算。把所有头的可复用比例平均就得到一个模拟命中率。这个模拟对我来说最大的价值不在于精确而在于揭示哪层平均命中率高、哪层波动大。统计结果里理解模型融合层的顶层平均命中率可以达到 0.85 以上而中间某一层会突然掉到 0.4 以下这种断层如果不做分布分析根本想不到。生成模型这边就更夸张。交叉注意力深层在采样中后段的模拟命中率接近 0.95自注意力中间层在细节生成阶段的命中率只有 0.1 左右。所以后续 UniCache 的缓存设计会采用分层分阶段的策略而不是用一个全局命中率去代表整层。我还把集中度 C 和模拟命中率的相关系数算了一下大概是 0.7 左右说明集中度可以作为命中率的一个有效预指标。这意味着在线推理时我们可以不跑完整相似度计算只根据输入特征估算每层头的集中度就大致知道这一层适不适合缓存这个思路会直接进入系列第二篇。5. 复现这套纵览时的实操注意点5.1 搭建注意力分布分析管线的最小方案复现这套纵览不需要特别复杂的平台我用的方案是在目标模型的 Transformer 层里插入 forward hook把每一层、每一个头的注意力权重矩阵钩出来存成 numpy 或者张量格式。重点是采样策略。我不建议直接把整个验证集全跑一遍那样数据量太大而且大部分样本的分布形态是重复的。我的做法是先跑 5 个样本观察每个层的熵分布区间锁定几个代表层比如一个浅层、一个中间层、一个深层。之后再拿 20 个样本在这几个代表层上集中记录既减小存储压力又不失规律性。存储格式上我建议保存 float16 而不是 float32能省一半空间。注意力权重是概率值float16 的精度完全够用。如果你还想做相似度分析那就额外存一个未经 softmax 的 logits 副本因为某些统计口径需要在 softmax 之前计算。5.2 三个统计口径上容易踩的坑第一个坑是 padding mask 处理。很多模型在处理变长文本时会把 padding token 的注意力分数设成一个非常大的负值。如果你直接从 softmax 之后的权重矩阵做统计padding token 的概率是 0这没有问题。但如果你是从 logits 自己算 softmax就必须手动把 padding 位置排除在分母之外否则熵值和集中度都会被显著拉低。第二个坑是 text embedding 输入信息不足不要直接在模型输入层统计注意力分布因为输入层能获取的信息量有限。要把统计点设在跨模态融合层靠近 text 侧的地方或者干脆把所有层都钩出来然后按层位置分组。否则你会在输入层观察到大量均匀分布误以为缓存没法做。第三个坑是平均统计掩盖了 head 的差异。很多现成的分析库会把一个层内所有 head 的注意力权重平均成一个矩阵来可视化。这个做法用于画图没问题用于缓存决策就很不严谨。一个层内可能 8 个 head 里 2 个稳定得离谱、6 个抖动得离谱平均后看起来中等稳定最后缓存方案做出来两个都顾不好。5.3 这些发现给后续两篇留了哪些接口UniCache 系列的第二篇会围绕分层分阶段缓存展开把这一篇得到的分布规律转换成具体的缓存策略哪些层用长期缓存、哪些层用短期缓存、哪些层完全不缓存以及如何在运行时根据输入特征动态切换缓存模式。第三篇则聚焦工程实现缓存存储格式、显存复用、缓存失效策略以及推理延迟和显存占用的实测数据。这一篇里标记出的高缓存潜力头和低命中率层就是后续缓存规则最基本的输入项。你在复现时得到的统计表可以直接当作第二篇实验的配置参数来使用。最后说一点个人体会。做注意力分布纵览最大的收获不是画了几张漂亮的热力图而是改变了我对注意力是动态计算这件事的直觉。很多层的注意力根本不是每次推理现算出来的它更像一个被输入牢牢锁定的函数输入不变输出就几乎不变。多模态模型里这种稳定性到处存在只是平时没有人把它摊开来看。希望这份分布地图能给你一些启发也欢迎在复现时多试试不同的量化口径也许你会发现比我更有意思的规律。