【轻量化复刻大模型】Linear Attention 与 Sparse Attention 梳理

0 阅读22分钟

【轻量化复刻大模型】Linear Attention 与 Sparse Attention 梳理


0x00 概要

0.1 新系列

准备开一个新系列:在 CPU(PC)上复刻(微缩)大模型。

为何要开这个系列?是因为:

  • Transformer 书上市之后,很多读者问我 ,是否有可以运行的代码来配套书籍进行阅读。
  • 有同学问我如何在没有GPU的情况下学习。
  • 很多同学都没有GPU或者机器人可以实验(无论是在学校还是在企业)
  • 绝大多数同学没有机会训练超大模型,这些只能在屈指可数的几个大公司才可以。

因此我们会对大模型(LLM & 机器人)进行学习,不但有博客,还有github项目 github.com/RossiLH/Mod… 作为配合使用。

对于项目的一些说明:

  • 涉及 KIMI K3,DeepSeek V4,PI0.5,PI0.7 等模型。后续可能会跟踪新模型。

  • 全 CPU 复刻,不涉及 GPU,也不涉及大规模训练(这是另外一个庞杂深邃的领域)和硬件。

  • 本项目模型的组件尽量是原始大模型的缩小版,如果组件没有开源版本,则用近似组件替代。

  • 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么,以及组件之间如何咬合成一个整体。

  • 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本(plot_*.py),可生成静态图。

  • 所有文件行号基于本地源码快照;源码后续变化时,行号可能移动,但函数名和调用关系仍可作为定位入口。

  • 重点不是记名词,而是理解一条主线,比如:Prefix Cache 不再只是 "找到相同 token,然后搬 KV";它变成 "找到共同前缀边界,并恢复一组语义一致、生命周期不同的模型状态"。

  • 为了保证每个文章可以尽量独立成文,所以有些内容会在不同文章中都存在,可能在系列层面上存在冗余,还请大家理解。

  • 参考项目:在开始这个项目之前,我去github找了一圈,发现一些类似的项目(有用于生产的,也有用于学习的),都非常出色。因此已经有类似项目,我就尽量借鉴。因此,本项目也可以认为是对这些项目的学习和二创。

这个项目和系列博客是我的一个尝试,也是我学习的过程。不确定是否满足读者的需求,也肯定有疏漏和谬误(尤其是demo和测试部分),我自己会反思改进,也希望大家提出宝贵意见。

0.2 注意力

因为会以 KIMI K3, DeepSeek V4 为例来进行,所以我们先来快速浏览其核心。

Linear Attention 的核心

  • 用固定状态替代不断增长的历史 KV,把历史读取成本从 "按 token 扫描" 变成 "读取工作集"。
  • 它的优势是效率、流式和长期状态;它的代价是容量、可寻址性和 Query 无关压缩。

Sparse Attention 的核心

  • 保留可寻址历史,只让当前 Query 读取少量相关 token/block。
  • 它的优势是精确检索、prefix reuse 和 Query 条件化;它的代价(举例)是 indexer、gather 和 KV Cache。

Kimi K3 的核心

  • 用 KDA 处理大多数长序列状态混合,用 Gated MLA 提供周期性的精确全局交互,再用 MoE、AttnRes、低精度和原生多模态把它们组织成大规模 Agent 系统。

DeepSeek‑V4 的核心

  • 用 local、CSA、HCA 组成多分辨率记忆系统:近处精确,中处选择,远处压缩。

0x01 总览:两条路线到底在压缩什么

1.1 Full Attention 的特色

标准 Full Attention 保留完整的 Key‑Value 历史。每来一个 Query,都可以重新决定历史中的每个 token 应该分多少权重:

它最强的地方是读取时仍然拥有选择权,代价是:

  • prefill 需要大量 Query‑Key 交互;
  • decode 要不断读取越来越长的 KV Cache;
  • KV Cache 体积随上下文增长。

1.2 差别

Linear 和 Sparse 的差别可以浓缩成一句话:Linear 压缩"历史怎么存";Sparse 压缩"这次读哪些历史"。

注:此处也将 CSA/HCA 纳入比对。

路线压缩轴历史存储Query 的选择权主要风险
Full Attention不压缩每个 token 都保留最大计算和带宽昂贵
Linear Attention状态/通道轴固定大小状态只能在状态中选方向碰撞、遗忘、不可寻址
Sparse Attentiontoken 轴大部分 token 仍保留可以选 token/block漏检、索引成本、随机访存
CSA/HCAtoken 轴先压缩,再选择压缩 token/block只能选压缩槽块内信息不可逆混合

更深一层看,这其实是在设计一个 Memory Subsystem:

  • Linear 的 state 像芯片里的 register file: 小、快、反复读写,但容量有限;
  • Full/Sparse 的 KV 像 cache: 容量更大、地址可寻,但访问成本更高;
  • K3 和 DeepSeek‑V4 做的不是单纯换一个公式,而是让模型记忆结构和硬件存储结构协同设计。

Linear 和 Sparse 的真正分歧,不是哪个公式更漂亮,而是要不要在读取时保留 "选择对象" 的权利。

1.3 深度思考:三条压缩轴

1.3.1 token 序列轴

压缩 token 数量:

  • block;
  • top‑k;
  • CSA/HCA;
  • 更长跨度的 Linear state。
1.3.2 hidden‑dim 通道轴

压缩每个 token 的表示宽度:

  • MLA;
  • 低秩 KV;
  • 子空间映射。
1.3.3 模型深度轴

用更多层传播、重组和补偿每层丢失的信息。Linear 单层省下了 KV, 但可能把代价转移为:

  • 更多层;
  • 更多激活参数;
  • 更长串行深度;
  • 更高 TPOT。

所以应比较 "每层参数效率", 而不是只比较单层 FLOPs。


0x02 Full Softmax Attention:为什么它难以被完全替代

2.1 数学逻辑

对每个 Query,先用点积给 Key 打分:

再用 Softmax 变成一份总和为 1 的预算:

最后用这份预算加权平均 Value:

通俗理解:

  • 点积:谁和当前问题相关;
  • 缩放:避免维度变大后分数过于极端;
  • Softmax: 所有 token 竞争一笔总预算;
  • Value 加权:把相关内容合成输出。

2.2 为什么 Full Attention 成为基线

如果任务要求:

  • 找到很久以前的某个精确 token;
  • 区分多个相似实体;
  • 根据不同 Query 对同一历史重新排序;
  • 组合多个远程证据;

那么"保存所有 token,读取时再决定"天然可靠。

Linear 或压缩 Sparse 的问题都是:它们在 Query 到来前,已经提前做了一部分信息合并。如果这一步把需要的信息丢掉,后面再聪明的读取函数也无法恢复。

2.3 深度思考

深度思考:为什么固定精度下,任意未来 Query 需要较大的存储?

如果要求对所有未来 Query 都保持通用正确性,就不能提前知道未来会访问哪一个 token。最坏情况下,几乎每个历史 token 都可能成为 needle, 需要保留足够的独立信息和地址。

因此:

  • Full/Sparse 需要随 T 增长的可寻址存储;
  • Linear 把它压成 state, 必然依赖任务分布假设;
  • CSA/HCA 处于二者之间,用 压缩槽换取折中。

0x03 Linear Attention:把历史提前汇总

3.1 Linear 路线

Full Softmax
    -> 核化 Linear Attention
    -> 卷积 + SiLU, 补局部上下文和非线性
    -> Mamba‑2, 加入输入相关记忆寿命
    -> DeltaNet, 沿当前 Key 定向改写
    -> GDN, 全局遗忘 + 定向改写
    -> KDA, 逐通道衰减 + GPU 友好分块
    -> KDA + MLA, 固定状态 + 可寻址检索

每一步的目的:

方法主要新增解决前任什么问题
Vanilla固定状态和核结合律Full 的历史扫描成本
卷积 / SiLU增加局部混合和非线性token 独立、表达过弱
Mamba‑2引入输入相关记忆寿命无法按内容遗忘
DeltaNet按当前 Key 做定向纠错和覆盖只累加、无法覆盖
GDN加全局衰减, 慢慢清理旧状态Delta 缺少快速清空旧状态
KDA逐通道衰减。即将 head‑level decay 细化为 channel‑wise decayGDN 只有一个 head‑level decay
KDA 混合用 MLA 补回固定状态缺乏 token 地址的问题固定状态无法精确检索

每次改进都没有消灭固定状态,而是在 "如何写、如何擦、如何忘、何时精确检索" 上做补偿。

3.2 数学本质

Linear Attention 将相似度近似为可分解形式:

于是:

令:

历史就变成一个固定大小的状态,Query 到来时只需读取 。

左边看起来像:

每个 Query 都扫描历史

右边变成:

历史先求和 → 得到状态 → Query 读取状态

因此,左式能做到每步近似 读取的唯一原因是:求和被移到了 Query 之前。

这同时意味着一个关键牺牲:你把 "先看 Query 再决定保留什么" 的权利卖掉了,换来固定大小状态。

**深度思考:为什么 Linear 的 O(1) 读取等于放弃 Query 选择权?**因为历史已经在 Query 出现前被求和成状态。之后 Query 只能在状态矩阵的方向中读取,无法回到原始 token 列表重新挑选某个对象。

Sparse 则相反:它保留 token 条目,Query 可以选择对象,只是选择过程要付出 indexer 成本。

3.3 Vanilla Linear Attention

3.3.1 它解决了什么

相对于 Full Attention:

  • 历史不再为每个 Query 重复扫描;
  • decode 的持久状态不随 T 增长;
  • 适合流式输入;
  • 可以用固定状态处理超长序列。
3.3.2 它的问题

Vanilla 通常是不断累加:

问题也很直观:

  • 不会主动删除旧信息;
  • 相互冲突的 Key‑Value 会叠在一起;
  • 历史越长,状态越像不断往同一个登记簿上写字;
  • 一阶近似可能产生负权重,失去 Softmax 的概率语义;
  • 固定状态无法保存无限多个独立事实。
3.3.3 "压通道轴" 和 "压 token 轴"

Linear 主要压的是状态通道:

每个 token 最后还能以多少维信息留下来。

”

Sparse 主要压的是 token 数量:

这次到底保留多少个 token 参加注意力。

”

所以二者并非同一种压缩:

  • 压通道轴:改变 "每个东西看得多清"。
  • 压 token 轴:改变 "有几个东西可挑";

深度思考:token 轴和通道轴压缩的区别

  • hidden‑dim 轴:每个 token 保留多宽;
  • token 轴:留下多少个 token/block;
  • 模型深度轴:用多少层传播和补偿信息。

token 轴压缩直接影响 "找不得到谁"; 通道轴压缩更直接影响 "找到后看得多清"。

3.4 DeltaNet:从累加改为纠错

DeltaNet 把状态看作一个在线学习模型。先用当前 Key 查询旧状态:

再用真实 Value 和预测的差值更新:

可以把它理解为:

  1. 先看旧账本对当前 Key 的预测;
  2. 把当前方向上的错误擦掉;
  3. 写入新的 Value。
3.4.1 相对于 Vanilla 的改进
  • 从 "只写不擦" 变成 "先纠错再写";
  • 新值可以覆盖旧关联;
  • 状态污染减轻;
  • 具有在线回归的解释。
3.4.2 局限
  • 擦除通常沿当前 Key 方向;
  • 两个 Key 很相似时,擦一个可能误伤另一个;
  • 写入顺序会影响结果;
  • 没有解决固定状态容量;
  • 没有让未来 Query 参与历史压缩。

3.5 GDN:加入全局遗忘

Gated DeltaNet 在 DeltaNet 上增加衰减:

它包含三步:

旧状态整体衰减 → 当前 Key 方向擦除 → 写入新的残差
3.5.1 相对于 DeltaNet 的改进

DeltaNet 擅长 "快速改写当前方向",但不会主动清空全局旧状态。GDN 增加了输入相关的记忆寿命:

  • 不再需要等某个新 Key 恰好撞上旧信息;
  • 很久不用的信息可以慢慢消失;
  • 状态不会无限累积。
3.5.2 局限

GDN 的一个 head 通常只有一个标量 decay:

这相当于假设状态空间各个方向共享同一个遗忘速度。可是实际信息的寿命不同:

  • 语法局部信息需要快速更新;
  • 主题信息需要长期保留;
  • 实体绑定可能要中等时间保留。

一个标量无法同时满足这些需求。

3.6 KDA:逐通道遗忘

KDA 将 GDN 的标量衰减改成逐通道衰减:

简化写法:

每个状态通道有自己的寿命:

  • 有些通道接近长期记忆;
  • 有些通道快速清空;
  • 有些通道负责中短期上下文。
3.6.1 相对于 GDN 的改进

GDN 是 "一整个房间共用一个温控器";KDA 是 "每个房间有自己的温控器"。

”

数学上, 出现在 Delta 梯度更新之前,意味着后面的纠错不是在同一个统一尺度上发生,而是在不同坐标上使用不同的有效步长。因此:

  • GDN:逐 head 标量门,状态各向同性;
  • KDA:逐 channel 门,状态可以各向异性;
  • KDA 能让不同状态方向有不同学习速度和记忆寿命。

但 KDA 仍然没有解决:

  • 固定状态的容量上限;
  • Query 无关压缩;
  • token 的独立地址;
  • 相近 Key 的覆盖;
  • 单步只能沿有限方向擦除。

0x04 Sparse Attention:保留 token 地址,减少读取对象

4.1 数学本质

Sparse Attention 先为 Query 选择支持集 ,再只在这个集合中做 Softmax:

它没有改变 Softmax 的竞争方式,只是删掉一部分参赛者。

因此 Sparse 的误差核心不是 “top‑k 这个数字多大”,而是:

被删掉的 token 原本拿走了多少注意力预算。

如果被删掉的概率质量很小,Sparse 误差可能很小;如果被删掉的 token 恰好是唯一的 needle,误差会很大。

”

4.2 Sparse 的发展脉络

Full Attention
     -> 静态窗口/全局点
     -> DSA 动态 indexer
     -> HiLS 可学习块摘要
     -> IndexCache 跨层复用
     -> CSA 压缩后选择
     -> HCA 强压缩后全读

每一步的核心矛盾:

静态模式: 便宜, 但不知道要找谁
动态模式: 知道要找谁, 但 indexer 变贵
块路由: 更规则, 但块内会带噪
压缩路由: 更省, 但信息先被混合

具体而言,每一步都解决前一代暴露出的新瓶颈:

阶段主要解决新问题
静态窗口 / 图访问规则、无需索引与内容无关,可能漏远程信息
DSA根据内容动态选 tokenindexer 变贵,gather 不规则
HiLS让块摘要和选择可学习训练和块级表示更复杂
IndexCache降低跨层重复索引索引复用错误可能传播
CSA缩小索引对象并规则化访存块内信息先被混合
HCA更强压缩和规则全局读取精确检索能力进一步降低

深度思考:Sparse Attention 的完整发展脉络

Full Softmax: 保留全部 token, 精确但昂贵

静态 Sparse: 固定窗口/全局点, 访问便宜但与内容无关

DSA: 低成本 indexer 动态选择, 精度提高但索引变贵

HiLS: 让块摘要和选择直接接受任务损失训练

IndexCache: 跨层复用 indexer 结果, 降低重复选择成本

CSA: 先压缩 block, 再做稀疏选择, 降低对象数和访存随机性

HCA: 进一步压缩后全读, 取消 indexer, 但牺牲块内精度

4.3 DSA

DSA 通过低维、低精度 indexer 先筛选候选,再用较高精度注意力读取候选。

它的价值:

  • 保留 token‑level 的动态选择;
  • 不必对全部候选做完整高维点积;
  • top‑k 让 Softmax 的注意力分布更尖锐。

它的问题:

  • 如果 indexer 仍扫描全部历史,二次项只是变便宜,不是消失;
  • top‑k 是硬截断,选错后无法补救;
  • token gather 可能比 FLOPs 账面显示的更慢。

4.4 HiLS

简单 mean pooling 的块摘要会稀释块内的单个关键 token。HiLS 的方向是:

  • 用更合理的可学习摘要;
  • 让块质量和熵信息进入摘要;
  • 让检索分数直接参与前向;
  • 让最终语言模型 loss 训练选择器。

核心变化:

原来: 索引器负责选,主模型负责算
现在: 索引结果也参与主模型的任务损失

因此 “选谁” 不再只是模仿 Full Attention,而是直接由最终任务效果塑形。

4.5 IndexCache

相邻层的 top‑k 支持集通常有较高重叠(只在稳定层 / 对角区成立)。IndexCache 让:

  • 少数层计算 indexer;
  • 多数层复用已有索引。

它体现了 Amdahl 定律:当主注意力已经很便宜时,原本次要的 indexer 会成为新的瓶颈。

IndexCache 的限制是:局部索引重叠高,不代表最终输出误差一定小。因此索引复用策略应以端到端 loss / 任务指标校准,而不是只看相似度。


0x05 CSA 与 HCA:压缩后再做 Sparse

5.1 CSA 的本质

CSA 的流程是:

token KV
    -> 块内可学习压缩
    -> 压缩块 indexer
    -> top‑k 压缩块
    -> 精确注意力 + exact local window
5.1.1 相对于 DSA 的改进
  • indexer 不再面对全部 token;
  • KV Cache 也变小;
  • 选择对象变成连续 block,访存更规则;
  • 局部窗口保留近期 token 细节。
5.1.2 CSA 的核心代价

块内压缩发生在 Query 到来之前:

一个块里的多个 token -> 一个压缩槽

之后 Query 只能决定 “读不读这个块”,无法重新决定块内哪个 token 更重要。

所以 CSA 的误差有两部分:

  1. 选错压缩块;
  2. 压缩时已经把块内重要 token 混掉。

深度思考:Sparse Attention 的压缩为什么说是学习出来的池化?

CSA/HCA 首先把多个 token 聚合为一个压缩槽。这个槽不是简单固定平均,而是由内容、门控和位置共同决定的加权池化。因此 Sparse 的 "压缩" 不只是删 token, 也是在学习:

哪些 token 应该共同形成一个摘要

5.2 HCA 的本质

HCA 更激进:

大块 token -> 更少的压缩槽 -> 对所有压缩槽做稠密注意力 + exact local window

HCA 不再需要 top‑k indexer,因为压缩后的槽已经很少,全部读取仍然可接受。

5.2.1 相对于 CSA 的改进
  • 去掉动态 indexer;
  • 读取规则;
  • 全局分支更便宜;
  • 更适合作为低分辨率全局上下文。
5.2.2 HCA 的核心缺点

HCA 的支持集是全部压缩槽,所以读取阶段基本没有漏选问题;但所有损失都集中在池化阶段:

Query 到来前,L 个 token 已经混成一个槽;如果目标 token 被同窗噪声污染,读取端只剩一个不可再分的对象。

”

因此:

  • HCA 的读取端相对完整;
  • HCA 的写入端不可逆;
  • 更换 Query 不能把被池化抹掉的 token 找回来。

深度思考:为什么 HCA 的损失全在池化阶段?

HCA 会读取全部压缩槽,因此没有 "selected block 漏选" 这一层损失。但 token 在压缩时已经混合。如果关键 token 和噪声进入同一槽,读取端只能得到混合结果,换一个 Query 也不能重新拆分。

深度思考:HCA 为什么容易发生块内污染?

因为 HCA 在 Query 到来前把同一窗口内的多个 token 混成一个槽。假设:

关键 token + 大量无关 token -> 一个 compressed slot

如果关键 token 的池化权重不够大,它的信息就被稀释。后续 Query 只能读取这个混合槽,无法回到原始窗口重新挑出关键 token。

这正是 HCA 用极低成本换取全局概览时付出的主要代价。

5.3 CSA/HCA 的分辨率分工

local window = 高分辨率近期记忆
CSA         = 中分辨率选择性全局记忆
HCA         = 低分辨率全局概览

DeepSeek‑V4 的精华不是某一个模块单独胜出,而是:

让不同注意力分支保存不同分辨率的记忆。

”

深度思考:Sublinear 压缩为什么要数据感知?

如果不同位置的信息重要性不同,固定比例压缩会浪费容量:

  • 背景 token 可以压得很厉害;
  • 关键实体和工具调用 token 需要高保真;
  • 任务切换附近可能需要更高分辨率。

所以未来的 sublinear compression 更可能由数据和任务驱动,动态决定不同片段保留多少信息。


0x06 Linear 与 Sparse 的系统级比较

设上下文长度为 T,Linear 状态维度为 ,Sparse 选择 个 token/block,压缩比例为 ,局部窗口为 。

维度FullLinear/KDAToken SparseCSA/HCA
历史存储完整 KV固定 state通常完整 KV压缩 KV
decode 读取 KV stateindexer + KV压缩索引 + 或
KV Cache约
token 地址完整不完整完整块级
Query 条件化完整受限强受块内压缩限制
prefill 并行强需 chunk 化强但有索引压缩 / 索引流水线
decode 时间依赖逐步逐步推逐步逐步
prefix reusetoken KV 共享状态快照token KV 共享压缩块共享
speculative decoding截断 KV 页即可要保存 / 回滚状态截断 KV 页即可需处理压缩块边界
主要硬件风险大带宽状态读写和串行链随机 gather压缩器和块边界

6.1 Prefill

6.1.1 Full/Sparse

prefill 有大量 Query,可以通过 tile 复用 KV,FlashAttention 等 kernel 能提高算术强度。

6.1.2 Linear

朴素递推有时间依赖。chunkwise 的作用是:

  • 块内把多个递推步骤重新组织成矩阵计算;
  • 块间只传递一个状态;
  • 用结合律恢复一部分并行度。
6.1.3 CSA/HCA

多了一条压缩流水线:

  • 先形成压缩槽;
  • 再做索引或压缩槽注意力;
  • 未完成的尾块需要临时缓存;
  • block 边界必须和缓存布局协调。

6.2 Decode

Decode 常常是 memory‑bound:

  • Full: 读取长 KV 流;
  • Sparse: 读取 indexer + selected KV;
  • Linear: 读取固定状态并写回;
  • CSA/HCA: 读取更短、更规整的压缩 KV。

所以 Sparse 的 “少读” 只有在索引扫描和 gather 也足够便宜时才会兑现;Linear 的 “固定状态” 只有在状态能驻留较近缓存层级、并且递推串行不成为瓶颈时才会兑现。

6.3 Prefix Cache

Sparse/Full 的 KV 是 append‑only:

  • token 可以按前缀树共享;
  • 任意 token 边界都可复用;
  • 分支请求只需引用只读 KV。

KDA 更适合保存 prefix 末尾的完整状态快照:

  • 状态与前缀长度无关;
  • 但快照通常需要复制到新请求;
  • 中间任意分支点可能需要重放;
  • speculative decoding 需要状态回滚。

所以:Sparse 的 prefix reuse 更像共享文件;Linear 的 prefix reuse 更像复制一个程序快照。

深度思考:Sparse 没有主动遗忘,Linear 是否因此更强?

Sparse 的 KV 是 append‑only, 通常不擦除历史,而是让 indexer 在读取阶段不选择某些 token。

Linear 的确有更强的状态更新和遗忘归纳偏置,可能更适合:

  • 流式序列;
  • 长期状态;
  • 近因任务;
  • 持续在线更新。

但 Linear 付出的代价是:

  • 固定容量;
  • 不可寻址;
  • Query 无关压缩;
  • prefix 和 speculative state 管理复杂。

因此不是谁绝对更强,而是记忆任务不同。

6.4 "Linear + Sparse" 值得研究

Linear 擅长隐式、长期、低带宽记忆;Sparse 擅长显式、可寻址、按 Query 检索。

一个自然的组合是:

Linear state: 长期趋势、主题、语法和状态
Sparse/MLA: 精确 token 检索
local window: 短程细节

真正需要优化的不是 “某个 attention 公式最优”,而是:

压缩应该插在哪个轴、哪个时间点、哪个层级,才能尽量不剥夺读取端的选择权。

”

深度思考:"把压缩插在哪一段" 为什么重要?

压缩放在不同阶段,损失不同:

  • 写入前压缩:影响所有未来 Query;
  • 读取前压缩:当前 Query 仍可能先选择;
  • token 轴压缩:减少可选择对象;
  • 通道轴压缩:降低每个对象的分辨率;
  • 层间压缩:可能用更多深度补偿。

因此真正的问题是压缩位置和信息选择权,而不是 "压缩率越高越好"。


TransFormer-封面

TransFormer-封面


0xFF 参考

48 小时啃完 Kimi K3:一个模型装得下 22580 个 GPT-2

详解 DeepSeek V4:Infra 巨鲸 “四连击”,百万上下文走进现实

详解 Kimi K3:强到冲击 Anthropic 估值的模型什么样?丨晚点播客

本文使用 markdown.com.cn 排版