Transformer 架构演进 可视化实验室
用交互式 SVG 呈现注意力复杂度、线性注意力、Mamba/SSM、RWKV、Flash Attention 等下一代架构。八个模块,纯原生实现。
注意力复杂度与 O(n²) 瓶颈
标准注意力的计算量与内存随序列长度二次增长。拖动滑块改变序列长度 n,观察 Attention FLOPs、Attention 显存、线性注意力 FLOPs 与 KV cache 的相对规模,并在对数刻度上理解长上下文的代价。
当前序列长度 n = 8192
标准注意力的注意力矩阵规模为 n×n,因此显存与计算量随 n 平方增长。当 n 从 4K 增长到 128K,注意力显存将膨胀 1024 倍,这正是长上下文建模的核心瓶颈。
线性注意力与 Kernel Trick
将 softmax(QK^T) 替换为可分解的 φ(Q)φ(K)^T,先计算 φ(K)^T V(与 n 无关的 d×d 矩阵),再左乘 φ(Q),把复杂度从 O(n²d) 降到 O(nd²)。点击切换核函数并观察分解流向。
核函数 φ 映射
线性注意力的关键在于 kernel trick:把不可分解的 softmax 拆成可分解的特征映射 φ。标准做法是 φ(x)=elu(x)+1,使 attn = (φ(Q)(φ(K)^T V)) / (φ(Q)∑φ(K)),分母分子均可递推,推理时成为线性 RNN。
Attn(Q,K,V) = φ(Q) · (φ(K)ᵀ V) / φ(Q) · ∑φ(K)
Mamba 与状态空间模型
SSM 通过连续微分方程 h'(t)=Ah(t)+Bx(t) 建模,经 ZOH 离散化为 h_t=Āh_{t-1}+B̄x_t,输出 y_t=Ch_t。Mamba 的关键创新是让 A、B、C 依赖于输入(选择性),从而可沿序列扫描。点击播放观察状态传递。
选择性扫描
传统 SSM 的 A/B/C 是固定参数(LTI 系统);Mamba 让 B、C 依输入变化,Ā 也通过门控改变,使模型能“遗忘”无关 token、“记住”关键 token。这把 RNN 的归纳偏置与 Transformer 的上下文选择性结合,推理为线性时间。
h_t = Ā_t h_{t-1} + B̄_t x_t , y_t = C_t h_t
RWKV:线性 RNN 与注意力混合
RWKV 用 WKV 机制把注意力改写为可递推形式:每个 token 的“键值累加”随序列增长。time-mixing 类似注意力,channel-mixing 类似 FFN。点击播放观察 WKV 的指数衰减累加。
WKV 机制
WKV(t) = ∑_i e^{-(t-1-i)w+k_i} · v_i / ∑_i e^{-(t-1-i)w+k_i}。指数衰减项 w 使远处的 key/value 贡献随距离指数衰减,因此可在线性时间递推累加。time-mixing 处理时序混合,channel-mixing 处理通道间非线性。
wkv_t = (Σ_{i
稀疏注意力模式对比
通过限制每个 query 只关注部分 key,把注意力矩阵稀疏化。Longformer 滑窗 + 全局 token;BigBird 加入随机块;Sparse Transformer 用条纹模式。切换模式查看不同的稀疏掩码与有效密度。
模式:滑窗 (Sliding Window)
每个 token 只关注前后 w 个 token,复杂度 O(nw)。Longformer 在此基础上加入少量全局 token(如 [CLS])以聚合整体信息。密度 = 关注数 / n²,n 越大稀疏度越高。
滑动窗口注意力 (Mistral SWA)
Mistral 采用滑动窗口注意力,每个 token 只关注前 w 个 token。配合层堆叠,第 L 层的有效感受野为 L×w,从而以局部计算达到全局覆盖。调节窗口与层数观察感受野扩展。
有效感受野 = 层数 × 窗口大小
每一层的输出都聚合了上一层窗口内的信息;堆叠 L 层后,顶层输出可触达 L×w 个 token。因此用较小窗口配合较多层数即可达到长程依赖,而单层复杂度仅 O(nw)。
Flash Attention 与 IO 复杂度
Flash Attention 不改变数学结果,但通过分块 (tiling) 把注意力计算限制在快速 SRAM 中,避免反复读写 HBM。标准注意力 IO 为 O(N²d),Flash 降至 O(N²d²/M)(M 为 SRAM 大小)。切换模式对比访问模式。
IO 复杂度对比
标准注意力需把 N×N 的中间矩阵写回 HBM 再读出,IO 量 O(N²d)。Flash 按 Q 行分块,每块加载对应 K、V 行块到 SRAM,计算局部 softmax 并在线累积输出,全程不物化 N×N 矩阵。IO 降为 O(N²d²/M)。
IO_naive = O(N²d) · IO_flash = O(N²d² / M_SRAM)
架构综合对比雷达图
从推理速度、长序列支持、参数效率、训练成本、效果、并行性六个维度对比 Transformer、Linear Attention、Mamba、RWKV、Flash Attention。勾选架构叠加显示,查看各自的优势区。
综合对比说明
Transformer 在效果与并行性上领先,但长序列成本高;Mamba/RWKV 推理为线性、长序列友好,但训练并行性弱;Flash Attention 是 Transformer 的 IO 优化,效果不变而训练更快。各架构在不同任务上各有所长。