【轻量化复刻大模型】线性注意力(Linear Attention)
0x00 概要
「线性注意力」不是一个单一论文,而是一条技术路线。核心文献如下:
注:
- 项目github.com/RossiLH/Mod… 和系列博客是我的一个尝试,也是我学习的过程。不确定是否满足读者的需求,我自己会反思改进,也希望大家提出宝贵意见。
- 涉及 KIMI K3,DeepSeek V4,PI0.5,PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻,不涉及 GPU,也不涉及大规模训练(这是另外一个庞杂深邃的领域)和硬件。
- 模型的组件尽量是原始大模型的缩小版,如果组件没有开源版本,则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么,以及组件之间如何咬合成一个整体。重点不是记名词,而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本(plot_*.py),可生成静态图。
- 所有文件行号基于本地源码快照;源码后续变化时,行号可能移动,但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文,所以有些内容会在不同文章中都存在,可能在系列层面上存在冗余,还请大家理解。
0x01 演进脉络与定位
1.1 前身与痛点
标准 softmax Attention:Attention(Q,K,V) = softmax(QK^T/√d)V
- 痛点:计算和内存复杂度 O(T²),T=100 万时 T² 矩阵难以存储。
1.2 同期竞品对比
1.3 定位一句话
不同于标准 softmax 注意力通过 O(T²) 显式计算所有 token 对的做法,线性注意力通过核函数 + 递推状态(R=φ(K)^T V)将复杂度降到 O(T),其中 DeltaNet/KDA 进一步用可修正的 Delta Rule 替代纯加性更新,实现了 O(1) 状态下的可覆写记忆。
1.4 演进总结
Linear Attention → 只能加,不能改
↓ 加门控
GLA → 能遗忘,但不能纠错
↓ 加 delta rule
DeltaNet → 能纠错,但遗忘是固定的
↓ 门控 + delta rule
Gated DeltaNet → 数据依赖遗忘 + 纠错
↓ Householder 反射
Parallel DeltaNet → 旋转替代对角衰减(更强表达力)
↓ softplus 门控 + dt_bias + A_log
Kimi Linear (KDA) → 最精细的门控 + 纠错(K3 实际使用)
每一步都是在解决上一步的局限:不能改 → 不能忘 → 遗忘不灵活 → 衰减不够激进。
0x02. 原理讲解
2.1 直觉类比视角
想象一间黑板教室。普通注意力是"老师每讲一句,就把所有旧板书重抄一遍"——随着课越讲越长,重抄的代价越来越大(O(T²))。线性注意力换了个思路:老师只维护一块固定大小的黑板(状态矩阵 S),每讲一句新知识,就在黑板上加/改几笔(O(1)),不再重抄。
难点在于:如果两句讲的是同一件事但答案不同,怎么改(以 Delta Rule 为例)?纯加性只会越写越糊(混叠);Delta Rule 则先"看清黑板上的旧答案"(检索),算出"差多少"(误差),再精准擦掉旧答案写下新的(修正)。
2.2 专业水平视角
线性注意力的基本想法,是不再保存所有历史 Token 的完整 Key-Value,而是把历史信息累积到一个固定大小的状态矩阵中。
普通的 Full Attention 在读取历史时,需要访问历史上的许多 Key 和 Value。上下文越长,KV Cache 通常也越大。线性注意力则把历史压缩成一个固定形状的状态 S,之后的新 Token 只需要读取和更新这个状态,而不需要逐个访问全部历史 Token。
最基本的状态更新可以写成:
这里:
- k_t 是当前 Token 的 Key,用来表示这条信息应该通过什么模式被寻址;
- v_t 是当前 Token 的 Value,表示真正要存储的内容;
- v_t k_t^\top 是一个外积矩阵,把当前的 Key-Value 关联写入状态;
- S_t 是处理完第 t 个 Token 后的记忆状态。
这个更新具有一个重要特征:当前 Token 的 v_t 和 k_t 在完成外积并累积到状态后,就可以立即丢弃。之后的计算只依赖状态矩阵 S_t,不需要保留随上下文长度增长的完整 KV Cache。不过,前面的短卷积仍然需要保留固定长度的、卷积之前的 qkv 输入。也就是说,线性注意力不代表所有中间状态都完全不需要保存;它主要消除了完整历史 Key-Value 的逐 Token 存储。
输入输出:
- 输入:Q, K, V ∈ R^{B, H, T, D}(batch, heads, 序列, 每头维度)
- 输出:O ∈ R^{B, H, T, D}
核心推导:softmax 可以近似为核函数的乘积:softmax(q·k^T) ≈ φ(q)·φ(k)^T,则
Attention(Q,K,V) = φ(Q) · (φ(K)^T · V)
= φ(Q) · S 其中 S = φ(K)^T · V ∈ R^{D×D} 固定大小
写成递推:S_t = S_{t-1} + φ(k_t)^T v_t(纯加法),o_t = φ(q_t) S_t
与标准做法的公式级对比:
2.3 讨论
- 为什么理论上 work:线性注意力把注意力重写为核方法的线性形式,借助矩阵结合律把"先算 QK^T 再乘 V"改为"先算 φ(K)^T V 再乘 φ(Q)",从而利用结合律消除对 T 的二次依赖。这揭示了 Transformer 推理时等价于一个 RNN(带 D×D 隐藏状态)。
- 为什么纯加性不 work:核方法假设单调记忆(新信息不影响旧信息表示),导致状态 S 里新旧信息混叠,无法覆写。
- 论文消融结论:DeltaNet 证明用 Delta Rule 替代加法可显著提升合成检索任务的精确性;K3 论文进一步论证了下界衰减 + 全秩门控在数值稳定性与表达力上的优势。
- 未解决理论问题:D×D 状态的信息容量上限(多少 token 可无损压缩进固定 S)仍是开放问题;如何衡量"检索精度 vs 状态容量"的权衡也未被严格刻画。
0x03. 子知识点拆解
3.1 树状结构
线性注意力
├── [1] 核特征映射 φ(x) = ELU(x)+1 (基础,所有变体依赖)
├── [2] 纯加性 Linear Attention (依赖 [1])
├── [3] Delta Rule 更新 (改进 [2])
│ └── [3a] RETRIEVE-COMPARE-CORRECT 三段式
├── [4] Gated Linear Attention (GLA) (改进 [2] 加门控)
├── [5] Gated DeltaNet (合并 [3]+[4] + L2归一化)
├── [6] Parallel DeltaNet (改进 [3] 用 Householder)
└── [7] KDA (Kimi Delta Attention) (工程化 [5],下界衰减+全秩门控)
├── [7a] 下界衰减 g_min=-5
├── [7b] 全秩输出门控
└── [7c] ShortConv + SiLU + L2 归一化投影链
3.2 代码映射
0x04. 子知识点深度分析
4.0 核心概念
4.0.1 状态矩阵 S
最基本的线性注意力可以写成:
普通线性层的投影参数在推理过程中通常保持不变,而 S_t 会随着每个 Token 的到来持续改变。
S 是所有方法共享的同一个抽象:这是一个 (B, H, Dk, Dv) 的大小固定的"记忆矩阵",替代了标准 Transformer 中随序列长度增长的 KV Cache。每一步的操作都是:读 S → 改 S → 返回输出。
S_t 又像一个普通线性层的权重矩阵:它通过矩阵乘法作用在 Query 上。因此,从 fast-weight 的角度看,S_t 可以理解为:
- 一组随上下文快速更新的权重;
- 一个动态的 Key-Value 关联记忆;
- 一个在推理过程中在线变化的线性模型。
如果初始状态为:
把递推展开可得:
如果 S_0 不为零,还需要在右侧加上初始状态 S_0。
用 Query 读取历史
将状态乘以任意 Query q:
其中:
表示第 i 个历史 Key 与当前 Query 的匹配程度。匹配程度决定对应 Value 被取出的比例。因此:
- v_i k_i^\top 可以看成向状态中写入一条 Key-Value 关联;
- S_tq 可以看成根据 Query 与历史 Key 的匹配程度,对历史 Value 做加权聚合。
用当前 Key 查询旧记忆
既然状态可以接受 Query,也可以把某个 Key 直接输入状态:
它表示当前关联记忆根据 Key k 召回的 Value。
因此,在写入当前关联 (k_t,v_t) 之前,可以先计算:
这个量可以解释为:旧记忆对当前 Value 的重建、预测或近似。
Delta Rule 的本质,就是比较这个旧预测和当前真正目标 v_t 之间的差异,然后只更新没有预测正确的部分。
4.0.2 特征映射函数
elu_plus_one 是特征映射函数,其核心是 elu(x) + 1 → 保证输出 > 0。
为什么需要正数?因为 Linear Attention 把 softmax(QK^T)V 拆成了 φ(Q) · [φ(K)^T V],中间的 φ(K)^T V 就是状态矩阵 S。如果 φ 输出有负数,S 中的正负会互相抵消,导致信息丢失。ELU+1 保证所有特征值为正,S 只做加法累积。
4.0.3 linear_attn_step - 纯加法更新
以下是最简单的线性注意力实现。
# S_new = S + φ(k) ⊗ v ← 只加不删
# out = φ(q) · S_new ← 读出
# ============================================================
# 1. Linear Attention (Katharopoulos et al., ICML 2020)
# Core idea: S_t = S_{t-1} + phi(k_t) * v_t^T
# Pure additive update, cannot overwrite
# ============================================================
def linear_attn_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Single step of Linear Attention (recurrent form).
S shape: (B, H, Dk, Dv) — state matrix
q shape: (B, H, Dk) — query
k shape: (B, H, Dk) — key
v shape: (B, H, Dv) — value
Returns: (output, new_S)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
S_new = S + torch.einsum('bhd,bhv->bhdv', phi_k, v) # pure additive
out = torch.einsum('bhd,bhdv->bhv', phi_q, S_new)
return out, S_new
直觉:一个只能往里塞东西、永远不扔东西的笔记本。每来一个 (k,v),就把 φ(k) × v 叠加到 S 上。查询时用 φ(q) 去匹配。
问题:只能叠加,不能覆盖。如果两个不同的 v 对应相似的 k,旧值永远不会被修正。
4.0.4 linear_attn_chunkwise - 批量训练形式
linear_attn_chunkwise 是 linear_attn_step 的循环展开版:等价地逐 token 调用 step 函数、收集全部输出,数值上与串行递推完全一致。
注意:这份 demo 代码本身并不并行——它只是把"一条 T 长序列的递推"一次性跑完,方便训练时拿到所有位置的输出。名字里的 "chunkwise(分块)"指的是并行化策略:把序列按时间切成若干 chunk,chunk 内部用矩阵运算并行(纯加法用 cumsum 前缀和即可;带 Delta Rule/门控这类递推依赖的用 UT/WY 变换),chunk 之间才做串行递推。真正的 chunk 并行实现见 (kda_chunkwise);
此处用循环展开,是为了数值等价、便于用测试验证 "chunkwise == recurrent"。
def linear_attn_chunkwise(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
"""
Chunkwise form of Linear Attention (parallel training form).
q, k, v shape: (B, H, T, D)
Returns: output (B, H, T, D)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
# S = phi(K)^T @ V (prefix sum for causality)
S = torch.einsum('bhkd,bhvd->bhdv', phi_k, v)
# But we need causal: o_t = phi(q_t) @ S_{:t}
# We use cumulative sum to get each prefix
S_cum = phi_k.transpose(-2, -1) @ v # (B, H, Dk, T) @ (B, H, T, Dv) -> (B, H, Dk, Dv)
# Actually simpler: just do it step by step for demo
B, H, T, Dk = q.shape
Dv = v.shape[-1]
S = torch.zeros(B, H, Dk, Dv, device=q.device)
outputs = []
for t in range(T):
out, S = linear_attn_step(S, q[:, :, t], k[:, :, t], v[:, :, t])
outputs.append(out)
return torch.stack(outputs, dim=2)
4.0.5 讨论
Q:为什么训练时用 chunkwise、推理时用 step 版?
这背后是线性注意力「推理等价 RNN」带来的两种计算形态,本质区别在于「这次前向需要多少个位置的输出」:
- 训练要"所有位置"的输出:训练时手头有完整输入序列(T 个 token 全部已知),而语言建模损失需要对每个位置 t 都算出一个输出 o_t 去预测下一个 token。所以训练走「全序列前向」——把 T 步一次性跑完、收集 o_1..o_T,才够算整条序列的 loss。linear_attn_chunkwise 就是这 "全序列前向" 的实现(这份演示代码为了等价直观直接循环展开;生产实现会用 chunk 并行 + UT/WY 变换把 chunk 内并行掉)。
- 推理只要"这一步"的输出 + 状态递推:推理/解码是自回归的,每次只生成一个 token,且它只能来自上一轮的输出(顺序拿到)。这时只需要当前这一步的输出 o_t 和更新后的状态 S_t,把它传给下一步即可,没有也不需要历史 T 步的输出。所以用 step 版:输入 (S_{t-1}, q_t, k_t, v_t),输出 (o_t, S_t)。
- 一句话:训练要"所有位置的输出"(全序列前向);推理只要"这一步的输出 + 状态递延"(单步前向) 。
这不是线性注意力的特例,而是**所有可递推序列模型(LSTM / RWKV / RetNet / KDA)**的通用范式:训练走全量前向(并行化机会大),推理走单步自循环(状态跨 token 传递)。chunkwise 正是回答「有 Delta Rule 递推依赖的模型,训练时怎么并行」的桥梁——纯加法可用 cumsum 前缀和并行,而带修正/门控的递推要靠分块并行来解决。
4.0.6 朴素线性注意力的缺陷:只会累加
最简单的线性注意力只有累加,没有遗忘,也没有覆盖。
如果同一个 Key 被写入两次,那么从这个 Key 读取时,得到的是两次 Value 的和,而不是第二次写入的 Value。换句话说,系统无法判断:
- 这条信息是不是已经存在;
- 新信息是不是旧信息的修正;
- 旧信息是不是已经过期;
- 当前写入是否应该覆盖以前的内容。
因此,写错的信息和已经过期的信息都会继续留在固定容量的状态矩阵里。随着 Token 不断到来,状态可能逐渐饱和;如果 Key 的范数较大,累加的数值甚至可能发生爆炸。
线性注意力虽然把历史压缩成了固定状态,但“固定大小”也意味着它必须在有限容量里不断处理新的信息。如果没有任何选择性更新机制,它只能盲目地把所有内容叠加进去。
4.1 Delta Rule
DeltaNet 的核心改进是:写入之前,先使用当前 Key 查询旧状态,看看旧状态已经记住了什么(旧状态对当前 Value 的预测);然后只写入“目标 Value 和旧状态预测结果之间的差(预测残差)”,从而避免同一关联被无条件重复累加。
4.1.1 直觉思路
一个会"纠错"的笔记本。来新 (k,v) 时,先用 k 查一下旧记忆里存了什么,如果查出来的和真实 v 不一样,只把差值写回去。即,像记笔记时"用修正液改错"——只涂掉写错的那几笔,不动其他内容。
4.1.2 专业角度
。检索 是记忆对当前 key 的预测值,v_t - 预测 是误差,beta*k 外积误差 只修正与当前 key 相关的行。
这是 Widrow-Hoff (LMS) 算法的引入——把状态 S 看作线性联想记忆,每次看到 (k,v) 对就用 LMS 规则更新,只修错误部分,因此能精准覆写而非常规加法混叠。
公式解读
旧状态对当前 Key 的读取结果是:
它可以理解为:旧记忆根据当前寻址模式 k_t,对当前 Value 做出的预测。
当前真正想写入的是 v_t,因此写入残差为:
这个残差表示旧状态还没有记住的部分。如果旧状态已经能够正确预测当前 Value,那么残差就是零,也就不需要重复写入。
DeltaNet 引入写门:
在本文的描述中,每个 Head 使用一个标量写门,写门由当前 Hidden State 计算得到。它控制本次残差写入的比例。
完整更新式为:
这个式子的含义可以分成三步:
- 用旧状态读取当前 Key,得到旧状态对 Value 的预测;
- 计算目标 Value 与旧预测之间的差;
- 把这部分差值沿当前 Key 的方向写回状态,并由 β_t 控制写入强度。
右乘形式
将上式整理,可以得到:
这个形式更容易看出状态发生了什么变化。
其中:
- β_t v_t k_t^T 是当前 Value 的写入;
- I-\beta_t k_tk_t^T 是对旧状态的变换。
如果 k_t 做了 L2 归一化,那么这个变换会让旧状态在当前 Key 方向上按 (1-β_t) 的比例缩小,而在与当前 Key 正交的方向上保持不变。
因此,DeltaNet 不再像朴素线性注意力那样无条件累加,而是在当前 Key 方向上执行一种平滑修正:
- β_t 较小,表示只修正一部分;
- β_t 较大,表示更积极地覆盖;
- 当 β_t=1 且 Key 已归一化时,可以在当前 Key 方向上实现精确覆写。
由于 Key 已经归一化,状态转移不会因为 Key 范数过大而无条件放大旧状态。因此,它比朴素线性注意力更不容易出现状态爆炸。
损失函数与 Delta Rule
DeltaNet 把状态 S 看成一个从 Key 预测 Value 的线性模型。
给定当前样本 (k_t,v_t),模型的预测是:
目标是让预测尽可能接近真实 Value v_t。因此,当前样本的回归损失为:
对 S 求梯度:
以 β_t 作为当前样本的学习率,在旧状态 S_{t-1} 上执行一步梯度下降:
代入梯度后得到:
这正是 Delta Rule。
因此,Delta Rule 可以理解为一次在线学习:
- 把当前 Key 当作输入;
- 用状态矩阵预测当前 Value;
- 计算预测误差;
- 用写门 β_t 控制学习率;
- 把误差写回状态。
如果:
说明旧记忆已经正确预测当前 Value,此时残差为零,不需要重复写入。
4.1.3 对比基线
4.1.4 讨论(反方视角)
下面是一些相关讨论和思考。
- a) 最强反对:Delta Rule 引入了序列依赖循环,无法像纯加法那样做前缀和并行化,训练变慢。
- b) 若失败,最可能错在:假设误差能线性传播修正,但多头、非线性下该假设不严格成立。
- c) 胜出条件:需要精确 key-value 检索的任务 (合成检索、指代消解);失败条件:极长序列 + 强非线性。
- d) 更稳健做法:GLA (门控加法)——放弃纠错换可并行性。
- e) 受影响方:训练方 (循环难并行)、推理方 (需状态回滚)。
4.1.5 计算复杂度
4.1.6 参数敏感性
alpha 对检索精度的影响如下图。具体是 固定 beta,扫 alpha(衰减因子):
关键推论:误差随 alpha 单调下降,在这个「同 key 重复写入」的合成任务上 alpha 越大越好。 但这不意味着实际应该取 alpha→1——真实长序列下 alpha→1 会导致状态无界累积、不同 key 之间开始混叠,这正是 KDA 引入 g_min = -5.0 下界衰减(等价 alpha 上界 ≈ exp(-ε))来强制一个最小遗忘量的原因。
plot_sensitivity
4.1.7 工程实现
# variants.py delta_net_step
# RETRIEVE → COMPARE → CORRECT (关键创新)
# ① RETRIEVE: retrieved = k · S ← 用 k 查旧记忆
# ② COMPARE: error = v − retrieved ← 算出"差多少"
# ③ CORRECT: S_new = α·S + β·(k ⊗ error) ← 只修正误差部分
# ④ READ: out = q · S_new ← 用 q 读出
# • α (decay):旧记忆按比例遗忘,防止无限累积
# • β (write strength):修正力度,太大可能震荡,太小收敛慢
# 关键区别 vs Linear Attention:能覆盖旧信息,因为 error = v − retrieved 允许"减去旧的、加上新的"
retrieved = einsum('bhd,bhdv->bhv', k, S) # RETRIEVE: S^T @ k
error = v - retrieved # COMPARE: 误差
correction = beta * einsum('bhd,bhv->bhdv', k, error) # CORRECT
S_new = alpha * S + correction # 衰减 + 修正
out = einsum('bhd,bhdv->bhv', q, S_new) # READ
4.1.8 探讨代码(验证反方 a:并行性损失)
可运行脚本:experiments\models\llm\linear_attn\discussion.py(第一段实验)。
思路:纯加性 S = Σ k_i v_i^T 满足结合律,因此「一次 matmul 并行」与「逐 token 串行」必须得到完全相同的结果;而 DeltaNet 的 S_t 依赖 S_{t-1},无法用 cumsum 前缀和展开。
实测输出:
============================================================
[纯加性] 前缀和(并行, matmul) vs 逐token(串行) 最大差: 4.77e-07
-> 若 diff≈0, 证明纯加性可用一次 matmul 并行算出所有前缀
[DeltaNet] S_t 依赖 S_{t-1} (循环依赖)
-> 无法用 cumsum 前缀和, 只能串行递推 (时间复杂度 O(T), 墙钟更长)
============================================================
[反方 a] KDA 固定状态的信息容量瓶颈
============================================================
写入值 | 检索 S[0,0] | 说明
---------------------------------------------
+0.9 | +0.450 | 朝新值方向移动, 旧值被冲淡
+0.5 | +0.284 | 朝新值方向移动, 旧值被冲淡
+0.1 | +0.072 | 朝新值方向移动, 旧值被冲淡
-0.3 | -0.145 | 朝新值方向移动, 旧值被冲淡
+0.7 | +0.339 | 朝新值方向移动, 旧值被冲淡
结论: 固定 D×D 状态只记住最近的写入, 同一 key 的旧 value 无法保留,
验证了'状态容量有限'的反方观点 (这正是 K3 混 3 KDA + 1 MLA 的原因).
结论:讨论 a 成立。这正是 Parallel DeltaNet(arXiv:2406.06484)用 Householder / WY 表示做 chunkwise 并行的动机——不是为了降复杂度,而是为了恢复可并行性。
4.1.9 DeltaNet 仍然存在的局限
DeltaNet 会保留与当前 Key 正交的方向,但这些“方向”不一定等同于单独的特征通道。如果某些方向长期没有被新的 Key 更新,那么其中的旧关联就可能一直保留下来。即使这些内容已经过期,也没有一个独立的机制主动把它们清掉。
因此,DeltaNet 具备残差覆写能力,但它的遗忘仍然依赖当前 Key 的方向。它不能独立地说“这一部分记忆现在应该过期”,也不能对所有记忆维度进行主动、选择性的衰减。
4.2 Gated Linear Attention (GLA)
原理如下:
- 直觉:在 Linear Attention 基础上加了一个"橡皮擦"。gate 是从输入学来的,每个维度独立决定"保留多少旧记忆"。但没有 delta rule,仍然不能纠错——只能选择遗忘或保留,不能定向修正。
- 专家:,σ 把门控压到 (0,1),逐通道缩放旧记忆。
gate 是数据依赖的逐通道遗忘因子,等价给每个通道一个可学习的"寿命";代价是只能整体缩放、无法定向改写具体条目(无纠错能力)。
S_new = σ(gate) ⊗ S + φ(k) ⊗ v ← gate 控制遗忘
out = φ(q) · S_new
具体代码如下:
# ============================================================
# 3. Gated Linear Attention (Yang et al., 2023)
# Core idea: S_t = G_t * S_{t-1} + phi(k_t) * v_t^T
# Data-dependent gating controls forget/remember
# ============================================================
def gla_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
g
本文由 AI图鉴自动抓取并翻译,版权归原作者所有;如需查看原文请点下方链接。
Linear Attention
0x00 Summary
Linear Attention is not a single paper, but a technical route. The core documents are as follows:
Notes:
- Project Github.com/RossiLH/Mod... and series of blogs are an attempt and a learning process. I will myself reflect on improvements and hope that you will make valuable comments, in the absence of certainty as to whether the needs of the reader are met.
- Models such as KIMI K3, DeepSeek V4, PI0.5, PI0.7 etc. New models may be followed up.
- All CPUs, not GPUs, not large-scale training (an additional area of depth) and hardware.
- The components of the model are as narrow as possible from the original LLM and, if there is no open-source version of the components, are replaced by approximate components.
- Explains how each component solves a problem, where it is difficult, why it can only be so solved, what the rationale for the solution is, and how the components combine into a whole. The focus is not on naming, but on understanding a thread.
- Each major component has a direct running Gradio demo - no GPU. There are also several visual scripts (plot *.py) that generate static maps.
- All file line numbers are based on local source snapshots; if the source code changes, the line numbers may be moved, but the function name and call relationship can still serve as the location entry.
- In order to ensure that each article can be written as independently as possible, some elements will be present in different articles and may be redundant at the level of the series.
0x01 integrator and positioning
1.1 Front and pain
Standard softmax Attention: Attention(Q,K,V)=softmax(QK^T/√d)V
- Pain points: Computation and memory complexity O (T2), T = 1 million times T2 matrix is difficult to store.
1.2 Competitions over the same period
1.3 Positioning sentence
Unlike standard softmax attention calculates all token pairs by O(T2) visible, linear attention reduces complexity to O(T) by a nuclear function + (R=(K)^TV), where DeltaNet/KDA further uses the modified Delta Rule instead of pure addition updates, which achieves overwritten memory under O(1) status.
1.4 Summary of evolution
Linear Attention
♪ Garmon control ♪
GLA can forget, but not correct.
delta rule
DeltaNet can correct mistakes, but forget is fixed.
+ delta rule
Gated DeltaNet
Householder Reflection
Parallel DeltaNet Rotation instead of Diagonal Decline (more expressive)
+dt bias+A log
Kimi Linear (KDA) fixes the most sophisticated door control + corrects (K3 actually used)
Each step is to address the limitations of the previous step: not to change, not to forget, not to be flexible, not to be radical.
0x02. Rationale
2.1 Intuitive analogy perspective
Imagine a blackboard classroom. The general focus is "Teacher repeats all old books every word he says" — the cost of copying is increasing as the lesson grows longer (O (T2). Linear attention changed the thinking: the teacher only maintains a fixed-sized blackboard (state matrix S), and every new piece of knowledge is added/changed to the blackboard (O (1)) and is not copied.
The difficulty is: if the two sentences are the same but the answer is different (as in Delta Rule) how? Pure additions will only get more and more obfuscated (mixed); Delta Rule will first "see the old answer on the blackboard" (retribution), calculate "the difference" (in error) and then accurately erase the old answer and write new (amended).
2.2 Professional-level perspective
The basic idea of linear attention is to move away from preserving the integrity of Token, Key-Value, but to accumulate historical information into a fixed-size state matrix.
Normal FullAttention needs to access many of the history Key and Valle when reading history. The longer context, the bigger the KV Cache. Linear attention compresses history into a fixed shape state S, followed by a new Token just to read and update the state, without having to access the whole history one by one.
The most basic state updates can be:
Here:
- k t is the current Token's Key, which indicates the mode in which this information should be located;
- v t is the current Token's Value, which is the real content to store;
- v t k t^top is an external matrix that writes the current Key-Value connection to the state;
- S t is the state of memory after processing the t Token.
This update has an important feature: the current Token 's v t and k t can be discarded as soon as they have been accumulated and accumulated. The subsequent calculations rely only on the state matrix S t and do not require the retention of the complete KV Cache as the length of context increases. However, the short volume in the front still requires the input of a fixed length prior to the volume. In other words, linear attention does not mean that all intermediate states need to be saved at all; it eliminates the whole history of Key-Value's token storage.
Input output:
- Enter: Q, K, V R^B, H, T, D}(batch, heads, sequence, each head dimension)
- Output: O ∈R{B, H, T, D}
Core extrapolation: softmax can approximate the product of a nuclear function: softmax (q.k^T)≈ (q)φT, or
Attention(Q,K,V) = φ(Q) · (φ(K)^T · V)
=(Q) = S where S =(K)^T.V ^R^D} Fixed Size
Written by: S t = S t t t t t ^T v t (pure addition), o t = q t) S t
Formula level comparison with standard practice:
2.3 Discussion
- Why in theory work: Linear attention is rewritten as a linear form of a nuclear approach, using matrix combinations to replace "first count QK^T" with "first count ^TV" with "first count ^TV and then multiply Q" to eliminate T's secondary dependence. This reveals that Transformer's reasoning is priced as a RNN (with DXD hidden).
- Why is it pure and non-working: the nuclear method assumes a monotonous memory (new information does not affect the representation of old information), leading to the mixing of old and new information in state S that cannot be overwritten.
- Discretionary conclusion: DeltaNet proves that the addition of Delta Rule to synthetic search missions significantly enhances accuracy; K3 further demonstrates the downside decay + the advantage of whole door control in numerical stability and expression.
- Unresolved theoretical issues: the information capacity cap of DXD (how much token can be compressed to fix S) is still open; the trade-off on how to measure "retribution accuracy vs state capacity" is not strictly painted.
0x03. Dismantling of sub-knowledge points
3.1 Tree structure
Linear attention.
=(x) = ELU(x)+1 (base, all variants dependent)
Ideas - [2] Pure Linear Attention (dependent [1]
Ideas - [3] Delta Rule Update (Improved [2]
│ - [3a] RETRIEVE-COMPARE-CORRECT Three-part format
Ideas - [4] Gated Linear Attention (GLA) (improvements [2] Gate control)
Ideas -[5] Gated DeltaNet (consolidated [3] + [4] + L2
Ideas - [6] Parallel DeltaNet (improvements [3] with Householder)
└ - [7] KDA (Kimi Delta Attention) (engineered [5], lower boundary decay + full door control)
Ideas - [7a] Decline in lower boundary g min=5
Ideas - [7b] Full output door control
– [7c] ShortConv + SiLU + L2 unified projection chain
3.2 Code Map
0x04. Sub-knowledge point depth analysis
4.0 Core concepts
4.2.1 Status matrix S
The basic linear focus can be:
The projection parameters of the general linear layer are usually unchanged during the reasoning process, and S t will continue to change as each Token arrives.
S is the same abstraction that all methods share: this is a fixed "memory matrix" (B, H, Dk, Dv), which replaces the KV Cache in the standard Transformer that increases with the length of the sequence. Each step of the operation is to read S → change S → back to output.
S t is also like a normal linear layer weight matrix: it multiplications by matrix on Query. Thus, from the point of view of last-right, S t can be understood as:
- A set of rapidly updated weights with context;
- A dynamic Key-Value association memory;
- A linear model of online changes in the reasoning process.
If the initial state is:
Extending the push is possible:
If S 0 is not zero, you need to add an initial state to the right.
Read history with Query
Query q:
Of which:
Indicates the extent to which the i history Key matches the current Query. Match levels determine the proportion of Value to be removed. Therefore:
- v i k i^top can be seen as writing a Key-Value connection to a directional state;
- S tq can be seen as a weighted aggregation of history Value based on the match between Query and history Key.
Query old memories with current Key
If Query is acceptable, you can also enter a Key directly:
It indicates that the current association memory is based on the Value recalled by Keyk.
Therefore, before writing the current association (k t, v t), you can calculate:
This amount can be interpreted as the reconstruction, prediction or approximation of old memories to the current Value.
The essence of Delta Rule is to compare the difference between this old forecast and the current real target v t, and then to update only the parts that do not predict correctly.
4.0.2 Characteristic mapping function
eu plus one is the feature map function, the core of which is eu(x) + 1 guaranteed output > 0.
Why do you need a positive number? Because Linear Attention splits softmax (QK^T)V into φ(Q) [K)^TV], the middle ^TV is state matrix S. If the output has a negative number, the positives and negatives in S offset each other, resulting in the loss of information. ELU+1 ensures all feature values are positive, and S only adds up.
4.3 Linear attn step - pure addition update
The simplest linear focus is achieved.
# S new = S + φ (k) ⊗ v← add only
# out = (q) S new #
# ============================================================
# 1. Linear Attention (Katharopoulos et al., ICML 2020)
# Core idea: S_t = S_{t-1} + phi(k_t) * v_t^T
# Pure additive update, cannot overwrite
# ============================================================
def linear_attn_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Single step of Linear Attention (recurrent form).
S shape: (B, H, Dk, Dv) — state matrix
q shape: (B, H, Dk) — query
k shape: (B, H, Dk) — key
v shape: (B, H, Dv) — value
Returns: (output, new_S)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
S_new = S + torch.einsum('bhd,bhv->bhdv', phi_k, v) # pure additive
out = torch.einsum('bhd,bhdv->bhv', phi_q, S_new)
return out, S_new
Intuitive: a notebook that only stuffs in, never throws. Each (k,v) comes with a superseding of k (k) x v to S. Use φ(q) to match the query.
Problem: It can only be stacked, not covered. If two different vs correspond to similar k, the old value will never be corrected.
4.4.4 Linear attn chunkwise - Mass training forms
Linear attn chunkwise is a circular extension for linear attn step: equal value token calls the step function, collects all outputs, and values are exactly the same as serial increments.
Note: This demo code does not by itself coincide... It's just a one-time run of "a long T-sequence push" to get the output of all positions when it's easy to train. The name "chunkwise" refers to a parallel strategy: Cuts the sequence in time into several chunks, and the chunk internal is run in array parallels (purely plus cumsum prefix is sufficient; UT/WY conversion with Delta Rule/ Door control for this type of push-dependent type) before a serial push is made between chunk. Real chunk parallels (kda chunkwise);
This cycle is being developed for numerical equivalence and for testing purposes.
def linear_attn_chunkwise(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
"""
Chunkwise form of Linear Attention (parallel training form).
q, k, v shape: (B, H, T, D)
Returns: output (B, H, T, D)
"""
phi_q = elu_plus_one(q)
phi_k = elu_plus_one(k)
# S = phi(K)^T @ V (prefix sum for causality)
S = torch.einsum('bhkd,bhvd->bhdv', phi_k, v)
# But we need causal: o_t = phi(q_t) @ S_{:t}
# We use cumulative sum to get each prefix
S_cum = phi_k.transpose(-2, -1) @ v # (B, H, Dk, T) @ (B, H, T, Dv) -> (B, H, Dk, Dv)
# Actually simpler: just do it step by step for demo
B, H, T, Dk = q.shape
Dv = v.shape[-1]
S = torch.zeros(B, H, Dk, Dv, device=q.device)
outputs = []
for t in range(T):
out, S = linear_attn_step(S, q[:, :, t], k[:, :, t], v[:, :, t])
outputs.append(out)
return torch.stack(outputs, dim=2)
4.5 Discussion
Q: Why use chunkwise in training, step in reasoning?
Behind this is the two forms of calculation offered by linear attention to the "pricing equivalents of RNN", which differ essentially from "the output of the number of locations needed this time":
- Training requires "all positions" output: full input sequences (all T tokens are known) are available at the time of training, and language modelling losses require an output for each location t o t to predict the next token. That's why training to go "full-sequence forward" -- to run the T-step one-off and collect the o 1.o T -- is enough to count the entire series ass. Linear attn chunkwise is the achievement of this "full-sequence forward" (this demonstration code is active for a direct cycle of equal value intuitives; production is achieved using chunk parallels + UT/WY transformations to drop in chunk).
- The reasoning is only as long as the output of this step + state is extrapolated: the reasoning/decode is self-returned, only one token is generated each time and it can only come from the previous round of output (in order of receipt). All you need is the output of the current step o t and the updated state S t to pass it to next step, and the output of the historical T step is not needed. So use the step version: Enter (S t-), q t, k t, v t), output (o t, S t).
- One sentence: Training requires "output of all positions" (full sequence forward); reasoning requires "output of this step + status deferred" (single forward).
This is not an exceptional case of linear attention, but rather a generic paradigm for **all recalcitrant sequence models (LSTM / RWKV / RetNet / KDA)**: training goes full forward (with great opportunities for parallelization) and reasoning goes one-step-round (state-crossing token). chunkwise is the answer to the bridge that says, "There's a Delta Rule relying model, how do training go hand in hand"? - Pure additions can be used with cumsum prefixes and parallels, while thrusts with amendment/door control are to be resolved in parallel in segments.
4.2.6 Impairment of simple linear attention: only cumulative
The simplest linear attention is cumulative, unforgettable and uncovered.
If the same Key is written twice, then when read from this Key it gets two times the combination of Value, not the second time it is written. In other words, the system cannot judge:
- Whether this information already exists;
- Whether the new information is a revision of old information;
- Is the old information outdated;
- Whether the current writing should cover previous content.
Therefore, both miswritten and expired information will remain in a fixed capacity status matrix. As Token arrives, the state may become saturated; if Key has a larger model, the cumulative value may even explode.
Linear attention, while compressing history into a fixed state, “fixed size” also means that it has to continuously process new information within limited capacity. If there is no selective updating mechanism, it can only blindly superimpose everything.
4.1 Delta Rule
The core improvement for DeltaNet is that before writing, use the current Key query on the old state to see what the old state remembers (the old state predicts the current Value); then write only "the difference between the target Value and the old state forecast results (forecast differences) ", thereby avoiding an unconditional repetition of the same connection.
4.1.1 Intuitive thinking
A notebook that would be "corrected." When new (k,v) comes, use k to find out what's in the old memory, and if it's not true v, write back the margin. That's, like, "removed with fixes" when taking notes -- only the wrong ones, without the rest.
4.1.2 Professional perspective
I don't know. Retrieve is the memory forecast for the current key, v t - the forecast is error, bita*k external error only fixes the rows associated with the current key.
This is the introduction of the Widrow-Hoff (LMS) algorithm -- to view state S as linear memory, updated with the LMS rule every time you see (k, v), and to fix only the wrong part, so it is accurate to overwrite instead of adding it in a normal way.
Formula Interpretation
The old state reads the current Key result:
It can be understood that the old memory predicts the current Value based on the current search mode k t.
Now, what really wants to write is v t, so the difference is:
This disability represents parts of the old state that have not been remembered. If the old state is able to predict the current Value correctly, then the disability is zero, and there is no need to repeat it.
DeltaNet introduces writing doors:
In the description of this paper, each Head uses a marker writing door, which is currently calculated by Hidden State. It controls the percentage of this disability written.
The complete update is:
This formula can be divided into three steps:
- Read the current Key with the old state and get the old state forecast for Valle;
- The difference between the Value target and the old forecast;
- Write this part of the margin back to status in the direction of the current Key and write strength by beta t.
Right Multiplication
This will be organized so that:
It is easier to see how the situation has changed.
Of which:
- β t v t k t^T is written by the current Value;
- I-\beta t tk t^T is a transformation of the old state.
If k t makes L2 homogenization, this transformation reduces the ratio (1-beta t) of the old state in the current Key direction, and remains unchanged in the direction that Key is currently dealing with.
So DeltaNet no longer adds up like simple linear attention, but instead implements a smooth amendment in the current Key direction:
- β t is smaller and means only one part of the correction;
- Beta t is larger and means more active coverage;
- When beta t=1 and Key has been unified, it can be accurately overwritten in the current Key direction.
As Key has been normalized, the shift does not unconditionally magnify the old state because the Key model is too large. As a result, it is less likely than simple linear attention to a state explosion.
Loss function and Delta Rule
DeltaNet looks at state S as a linear model for forecasting Value from Key.
For the current sample (k t, v t), the model predicts that:
The goal is to get the prediction as close as possible to the real Value v t. Thus, the current return losses of sample are:
For S:
With β t as the current sample study rate, one step down in the old state S t e:
Gradient after replacement:
That's Delta Rule.
So Delta Rule can be understood as an online study:
- Current Key as input;
- Forecast the current Value with a status matrix;
- Calculating prediction errors;
- Control of learning rates with writing doors β t;
- Write back the error.
If:
This indicates that the old memory has correctly predicted the current Value, at a time when the gap is zero and does not need to be repeated.
4.1.3 Comparison of baselines
4.1.4 Discussion (reverse perspective)
Some relevant discussions and reflections follow.
- (a) The strongest objections: Delta Rule introduced a sequence-dependent cycle, which did not allow for prefixing and parallelization as pure additions and slower training.
- (b) If failure occurs, it is most likely that the assumption that the error energy linear transmission correction is not strictly valid under multiple, non-linear scenarios.
- c) Winner condition: mission requiring precise key-value search (composite search, decomposition); failure condition: extremely long sequence + strong non-linear.
(d) More robust practice: GLA (door control plus) - waiver of error correction for parallelity.
(e) Affected parties: Trainers (recycling difficult to parallel), reasoning (requirement to roll back).
4.1.5 Complexity of calculations degrees
4.1.6 Parameter sensitivity
The impact of alpha on retrieval accuracy is illustrated below. In particular, fix Beta, sweep alpha:
Key inference: The error drops with the alpha monotonous, and the alpha is as big as possible on this "repeat with key" synthesis. But this does not mean that it should actually take alpha1 – the alpha1 under the real long sequence would lead to an unbridled accumulation of the state and a mixing of different keyes, which is why KDA introduced g min = -5.0 lower boundary decay (the alpha upper border ≈exp (-ε)) to enforce a minimum level of oblivion.
plot_sensitivity
4.1.7 Project realization
# variants.py delta_net_step
# RETRIEVE # COMPARE # CORRECT (Key innovation)
# 1 RETRIEVE: retrieved #
# 2 COMPARE: error = v-retrieved #
#3 CORRECT: S new = alpha S + beta (k ⊗ error)
#4 READ: out = q. S new ← read in q
• α (decay): Old memory is proportionally forgotten to prevent infinite accumulation
• Beta (write strength): correction strength, too big to shock, too small to slow
# Key distinction vs Linear Attention: can overwrite old information because error = v - retrieved allows "mind old, add new"
retrieved = einsum('bhd,bhdv->bhv', k, S) # RETRIEVE: S^T @ k
# COMPARE: Error
correction = beta * einsum('bhd,bhv->bhdv', k, error) # CORRECT
S new = alpha *S + decay #+ fix
out = einsum('bhd,bhdv->bhv', q, S_new) # READ
4.1.8 Explore the code (validation of a: concurrent loss)
Runable Script: Experiments \models\llm\linear attn\discussion.py (first paragraph experiment).
Thinking: Purely S = Σ k i v i^T meets the amalgamation rule, so "matmul" and "token-by-line" must have exactly the same result; and DeltaNet 's S t relies on S {t m m s, unable to prefix and expand.
Actual output:
============================================================
Prefixes [pure plus] prefixes and (parallel, matmul) vs by token (string) Max deviation: 4.77e-07
- >If diff≈0, prove pure additions can be calculated at the same time as matmul
[DeltaNet] S t depends on S t t s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s s t s s s s t s s s s s s s s s s s s s s s s s s s s s
- > Unable to use the cumsum prefix, only symmetry (time complex O(T), wall clock longer)
============================================================
[Reverse a] KDA fixed-state information capacity bottlenecks
============================================================
Write Value | Search S[0,0] | Description
---------------------------------------------
+0.9 | 0.450 | Moves in the direction of new values, the old values are diluted
+0.5 |0.284 | moving in the direction of new values, old values being diluted
+0.1 | +0.072 | Moves in the direction of new values, the old values are diluted
-0.3 |-0.145 | Move to new values, old values diluted
+0.7 | +0.339 | Move in the direction of new values, old values diluted
Conclusion: Fixed DxD status only remembers the most recent writing, and the old value of the same key cannot be retained.
Checked the "state limited" counterpoint.
Conclusion: Discussion a Established. This is precisely what Parallel DeltaNet (arXiv:24066.0484) used Householder/WY to indicate as a parallel motivation for chunkwise - not to reduce complexity but to restore competitivity.
4.1.9 The remaining limitations of DeltaNet
DeltaNet will retain the current orientation of Key, but these “directions” are not necessarily the same as separate feature channels. If certain directions have not been updated for a long time by the new Key, the old links may remain. Even if these elements have expired, there is no independent mechanism to clear them out.
So DeltaNet has the ability to write poorly, but its oblivion still depends on the direction of Key. It cannot say independently that “this part of the memory should now expire”, nor can it undertake a proactive and selective decline in all memory dimensions.
4.2 Gated Linear Attention (GLA)
The rationale is as follows:
- Intuitive: Add a rubber rub on Linear Attention. It's from input, and each dimension independently determines how much old memory to keep. But without delta rule, the error cannot be corrected — the choice is to forget or retain, and there can be no targeted amendment.
- Specialist: Consider pressing the door control (0,1) and scaling up old memories.
Gate is the channel-by-channel forgotten factor on which the data depend, with an equivalent price of one learning "life" per channel; the cost is that it can only be scaled down as a whole and it cannot be directed to recast a specific entry (without ability to correct errors).
S new =gate
out = φ(q) · S_new
The specific codes are as follows:
# ============================================================
# 3. Gated Linear Attention (Yang et al., 2023)
# Core idea: S_t = G_t * S_{t-1} + phi(k_t) * v_t^T
# Data-dependent gating controls forget/remember
# ============================================================
def gla_step(S: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
g
Auto-fetched and translated by AI Atlas. All rights belong to the original author.