GPT-2在2019年还是个大家伙,1.24亿个参数。到了2026年,Kimi K3塞进了2.8万亿个参数。两者一比,差了22580倍。七年时间,数字翻了这么多,中间到底发生了什么?不是简单地把模型做大。每一代架构都在解决同一个头疼的问题:记忆体塞满了怎么办。
我们从Transformer的起点出发,沿着KV缓存、线性注意力、DeltaNet、Gated DeltaNet、KDA,一直走到Kimi K3的混合设计。这条路上,每个新结构都在对付缓存爆炸、遗忘旧知识、选择性读写这些麻烦事。看完你会明白,做大模型,难的不是堆参数,是怎么让记忆体不炸。
记忆体里塞进全部历史,读写速度就垮了
GPT-2的代码跑起来很直白。输入一句话,先查表得到每个词的向量,再叠加上位置信息,然后丢进12个一模一样的模块里。每个模块里,注意力机制先算一遍,接着过一个全连接层,最后输出。这中间最要命的一步,是计算注意力分数。每一个词都要跟前面所有词算一遍关系。句子越长,计算量按平方往上翻。
代码里能看到这个操作。拿到查询、键、值三个矩阵后,用查询去点乘所有键,得到一个方阵。方阵的每个格子,代表一个词对另一个词的关注度。然后掩码遮住未来的词,只让模型看到过去。最后用这个分数去加权对应的值,得到输出。每一步都老老实实算全部历史,一点便宜都不占。
但生成下一个词的时候,前面所有词的计算结果其实没变。模型却每次都重新算一遍。这就像你每次走到路口,都要从头回忆一遍整条路怎么走,而不是直接掏手机看导航。有人注意到这个浪费,就搞出了KV缓存。把已经算好的键和值存下来,下次直接用。不用重复劳动。
代价是内存越占越多。句子长到几千词,缓存就大得吓人。读写缓存的耗时,开始超过计算本身。模型卡在等数据从内存搬进计算单元的路上,算力再强也使不上劲。1.24亿参数的GPT-2还能撑住。到了千亿万亿的规模,这条路根本走不通。
把指数换成加法,速度上来了但脑子变笨了
线性注意力换了个打法。原来算注意力分数用的是指数函数,指数完再归一化。每个查询都得跟全部键做这个运算。线性注意力先在键和查询上各自套一个简单的映射,叫ELU加一。然后交换乘法顺序。
原来要算Q乘K的转置,得到一个N乘N的大矩阵,再乘V。现在先算K的转置乘V,得到一个固定大小的方阵。查询直接去点乘这个方阵。方阵的尺寸是固定的,跟句子长度没关系。这样就绕开了那个平方级的暴涨。
代价也很明显。原来的指数函数,能把注意力集中到少数几个最相关的词上。ELU加一做不到那么精准。所有词的地位被拉平了,模型分不清谁更重要。短句子上看不出差别,长文章里就开始犯糊涂。需要记住很久以前的某个细节时,线性注意力常常抓瞎。速度换精度,这个买卖在长文本上越来越不划算。
只记新东西不丢旧东西,记忆体迟早爆掉
线性注意力的状态更新是加法。新来的键值对,直接加到已有的方阵上。方阵大小固定,新东西一直往里塞,旧东西从来不扔。等塞满了,新来的就把旧的挤变形。信息开始互相干扰。你想查A,结果A跟B、C、D混在一起,根本分不清。
DeltaNet专门治这个毛病。每次写入新信息之前,先拿当前的键去查一下记忆体,看这个位置已经存了什么。算一个差值,只把真正新的那部分写进去。旧信息被自动抵消掉。这样就不会无限制地堆叠。
代码里能看到这个操作。用键去点乘状态矩阵S,读出旧值。用新值减去旧值,乘上一个写入强度系数,得到增量。再把键的转置乘上这个增量,加回S。读旧值、算增量、覆盖写入,三步走完,一个精确的替换就完成了。记忆体始终装着最新的信息,不会堆垃圾。
但这套机制有个盲区。它只替换那些能被精确查询到的位置。如果模型要清空一整段记忆,或者想主动忘掉某类信息,DeltaNet就没办法了。它没有一个全局的遗忘按钮。方阵里的每个格子,都只能被精确命中才能修改。乱糟糟的信息混在一起时,命中率直线下降。
给记忆体加个遗忘开关,又能删又能改
Mamba先加了这个开关。每次更新状态时,先乘上一个衰减系数,再加新东西。系数介于零和一之间。系数接近一,记忆保留得久。系数接近零,旧信息快速消散。这样就防止了状态无限膨胀。
但Mamba的衰减是对所有信息一视同仁。重要的、不重要的,一起衰减。模型没办法只忘掉过时的,留下有用的。Gated DeltaNet把DeltaNet的精确替换和Mamba的全局衰减揉在一起。既能在特定位置精确改写,又能整体控制记忆的寿命。
公式里多了一个alpha。alpha等于一的时候,退化成纯DeltaNet,精确替换。alpha等于零的时候,旧状态全部清零,重新开始。alpha在中间取值时,新旧按比例混合。这样模型就有了两种遗忘手段。一种是针对某个具体事实的定点清除,另一种是对整个记忆体的老化管理。
实现上用了跟DeltaNet一样的并行分块技巧。每一块里,先处理块内注意力,再更新全局状态。新增的衰减系数,在跨块传递时累乘进去。早写入的信息,经过多个时间步后,影响力自然衰减。晚写入的,权重更高。记忆体的内容始终在流动,旧的沉下去,新的浮上来。
每个通道单独控制遗忘,脑子更灵光了
Gated DeltaNet的alpha是一个标量。对所有通道都一样。KDA把这个标量换成了一个向量。每个通道都有自己的衰减系数。模型可以决定,哪些特征需要快速忘记,哪些特征要长久保留。
这个改动看起来小,效果却很大。一个通道可能负责句法结构,需要稳定一点,衰减慢。另一个通道负责具体实体,可能几句话之后就过时了,衰减快。同一个状态矩阵里,不同维度的信息寿命不同。标量衰减做不到这种精细控制,向量衰减可以。
代码里能看到这个变化。alpha的形状从(batch, heads)变成了(batch, heads, channels)。每个头每个通道都有自己的值。更新状态时,乘的是逐元素乘,不是标量乘。这样记忆体的每个格子都有自己的遗忘节奏。
KDA还引入了混合架构。不是所有层都用这种线性注意力。每隔几层,插入一个传统的多头注意力层。传统注意力负责全局检索,线性注意力负责流式更新。两者配合,短上下文用精确检索,长上下文用状态记忆。各自的短板被对方补齐。
掏出三个大招,把脑子武装到牙齿
Kimi K3在KDA的基础上加了三个硬货:多头潜在注意力、混合专家模型、注意力残差。每一样都解决一个具体痛点。
多头潜在注意力解决KV缓存太大的问题:传统注意力里,键和值的维度跟词向量一样大。存起来占地方。潜在注意力先压缩到一个低维空间,再展开。缓存尺寸大幅缩减。同时保留检索能力。查询还是高维的,但键值被压缩了。读写变快,占用的显存变少。
混合专家模型解决计算量太集中的问题:每次前向传播,不是所有参数都激活。模型先算一个路由分数,从八百九十六个专家里选出十六个。只有选中的专家参与计算。这样参数总量很大,但每次用的只是一小部分。计算量不随参数总量线性增长。规模扩大时,成本增加得慢。
注意力残差解决层数太深信号衰减的问题:模型层数堆到几十层,早期的信息传到后面已经模糊了。每十二层设一个检查点,把这一段的输出存下来。后面的层可以用注意力机制去查这些检查点。需要早期信息时,直接调取。不用等信号一层一层传过来。训练更稳,推理更快。
三层结构组合起来,Kimi K3能在固定尺寸的记忆体里,存更多信息,算更快速度,调更早的数据。每层各有分工,互不干扰。
核心变化不是参数多了两万倍,是每一代都在改记忆体的读写规则。 加法变替换,标量衰减变通道衰减,纯状态记忆混合定期检索,深度残差提供跨越时间的通道。每一步都在对付同一个物理限制:记忆体有限,信息无限。只能精打细算,该存存,该扔扔。
两万八千亿的参数,就是这么抠出来的。