← 全部论文2402.18668v2 / r1
版本与下载

BASED:简单线性注意力如何平衡召回能力与生成吞吐

论文:Simple linear attention language models balance the recall-throughput tradeoff
作者:Simran Arora、Sabri Eyuboglu、Michael Zhang、Aman Timalsina、Silas Alberti、Dylan Zinsley、James Zou、Atri Rudra、Christopher Ré。
版本:arXiv:2402.18668v2,2025-03-07;首次提交于 2024-02-28。
研读与获取日期:2026-10-10。
材料:arXiv v2 · 本地原文 PDF · 英文汇报 PDF · 英文 TeX。下文页码均为从 1 开始的 PDF 页码,与本版印刷页码一致。

论文总览

BASED 研究的是一个具体取舍:语言模型在生成时,能否保留较小的历史状态,又能准确取回上下文中的信息?Transformer 把历史 key/value 留在 KV cache 中,召回能力强,但缓存随序列增长;固定状态的递归模型省内存,却需要把越来越多的信息压入同一状态。

作者把这个问题拆成三层。首先,通过多查询关联召回(Multi-Query Associative Recall,MQAR)实验和通信复杂度下界,说明状态容量确实约束召回能力。其次,把全局二阶 Taylor 线性注意力与小窗口精确 softmax 注意力组合,并在实际语言模型中加入短门控卷积,让容量与局部计算精度可以分别调节。最后,设计融合特征映射、状态更新与读出的 CUDA 内核,把理论上的线性复杂度转化为真实吞吐收益。

最有代表性的实证结果是:在约 1.3B 参数、50B 训练 token 下,BASED、Mamba、Transformer++ 的 Pile 困惑度分别为 6.30、6.28、6.28,但 BASED 在 SWDE、FDA 和改写后的 SQuAD 上均优于 Mamba。它仍明显落后于 Transformer++ 的部分信息抽取结果。论文另报告特定 H100 解码设置下约 24 倍于 FlashAttention-2 的吞吐,这个数字不包含预填充,也不是所有批量下的统一加速。(Table 1,p. 8;Sec. 6.2,p. 9)

这篇工作的主要价值是给出可操作的设计原则:为远程关联配置足够大的记忆,为局部对齐保留精确计算,再围绕实际状态形状设计内核。

结构地图

原文章节论证作用PDF 页码与关键证据
1. Introduction从整体困惑度转向召回与效率的联合评估pp. 1–2,Fig. 1
2. Preliminaries and Related Worksoftmax、线性注意力、稀疏注意力与递归模型p. 3,Eq. (1)
3. No Free Lunch: Memory-Recall TradeoffMQAR 实验、状态下界、BaseConv 深度下界pp. 3–5,Fig. 2,Theorems 3.1–3.4
4. The BASED ArchitectureTaylor 特征映射与小窗口注意力pp. 5–6,Eqs. (2)–(4),Fig. 3
5. Efficient Implementation分块、融合与寄存器驻留pp. 6–7
6. Results语言建模、召回、消融与速度pp. 7–9,Table 1,Fig. 4
7. Conclusion召回—吞吐前沿的总结p. 9
A. Extended Related Work与已有稀疏/线性混合、SSM 和内核工作的关系pp. 16–19
B. IO Aware Implementations预填充、递归更新、窗口内核及微基准pp. 20–30,Algorithms 1–3,Figs. 5–7
C. Extended Architecture Details短卷积与衰减变体p. 31
D. Extended Results扩展 MQAR、常识、few-shot、DNA、消融pp. 32–35,Fig. 8,Tables 2–6
E. Experimental Details预训练、状态口径、召回数据构造pp. 36–39
F. Theoretical Results模拟关系、容量/深度下界、受限设置的上界pp. 40–74;训练配置表位于 pp. 75–78

一:动机——整体困惑度接近,为什么召回仍可能差很多?

1.1 召回需要保留输入中的临时事实

考虑序列 A 4 B 3 C 6 … A ?。模型需要从当前上下文找回 A → 4,而非依靠训练中记住的固定答案。在网页字段抽取、文档问答中,这类能力对应“在前文找到这次输入特有的事实”。

整体困惑度会平均大量容易预测的 token。即使一个模型在常见语法、词频和局部搭配上很好,它仍可能在稀有名称、重复 bigram 的值或较远字段上失误。因此本文同时测量整体 Pile PPL、关联召回 token 的 PPL,以及要求从输入中提取事实的任务。

1.2 参数量与历史状态量是两个预算

模型参数在不同输入之间共享;递归状态随当前输入变化。增加参数不等于为某段输入增加可保存的信息。MQAR 实验把横轴设为生成时的状态大小,而非仅按参数量比较。

图 2:状态容量与关联召回

原文 Fig. 2,p. 4。横轴为状态字节数,纵轴为 MQAR 准确率;不同点来自不同模型配置。同类架构通常随状态增加而提高召回,但同样状态预算下,不同架构的利用效率不同。BASED 的橙色点扩展了作者所测配置的经验前沿。此图横轴是内存,真实运行速度在 Fig. 4 中另外测量。

全注意力保留逐 token 的 key/value;线性注意力保存累积的关联矩阵;滑动窗口注意力(Sliding Window Attention,SWA)只保留最近的一段 KV。三者对历史信息的表示方式不同。对固定状态方法,“单步复杂度不随序列长度增长”仍不足以决定召回能力和真实速度。

二:方法——全局关联记忆、局部精确计算与硬件实现

2.1 整体架构及信息流

图 1:BASED 的局部与全局组合

原文 Fig. 1,p. 2。左图为不同矩阵/窗口规模下的 GEMM 延迟;中图比较单独线性注意力、单独 SWA 和组合后的召回;右图示意全局覆盖与局部精确交互如何互补。右图表达依赖关系,实际预训练模型通过不同类型的层混合实现。

输入 token 经嵌入后依次经过混合层:线性注意力层读写全局压缩记忆,SWA 层精确访问近邻,短门控卷积层提供低成本的局部混合;这些层配合模型的投影、残差和前馈计算,最终预测下一 token。它不是把一个训练好的 softmax 模型直接替换内核,实验中的模型是从头预训练。

作者提出的机制解释是:局部算子负责形成正确的 key/value 对齐与精细比较,全局线性状态负责跨距离传播关联。以 A 4 … A 为教学例子,模型先需要让 A 与相邻的 4 发生局部交互,再把这种关联保留到较远的查询位置。组合消融支持这种分工,但没有直接测出每一层具体执行了哪一步。

2.2 从注意力求和变成固定大小状态

为把原文混用的行列向量写法统一,下面采用列向量约定,并以 $d_v$ 单独表示 value/head 维度。对应原文 Eqs. (1)–(3),pp. 3、5。

符号含义
$N$、$i$序列长度与当前 token 位置
$d'$投影后的 query/key 维度,常用 16
$d_v$每个头的 value 维度
$D$应用特征映射后的维度,对应原文 $\tilde d$
$q_i,k_i\in\mathbb R^{d'}$当前 query/key
$v_i\in\mathbb R^{d_v}$当前 value
$s_i\in\mathbb R^{D\times d_v}$累积 key–value 关联状态
$z_i\in\mathbb R^D$归一化所需的 key 状态

softmax 注意力的昂贵部分是对历史逐项评分、归一化和加权。选取特征映射 $\phi$ 后,定义核函数 $K(q,k)=\phi(q)^\top\phi(k)$,有

$$ y_i=\frac{\sum_{j\le i}K(q_i,k_j)v_j}{\sum_{j\le i}K(q_i,k_j)} =\frac{s_i^\top\phi(q_i)}{z_i^\top\phi(q_i)}, $$
$$ s_i=s_{i-1}+\phi(k_i)v_i^\top,\qquad z_i=z_{i-1}+\phi(k_i),\qquad s_0=0,\ z_0=0. $$

关键变化是求和与矩阵乘法交换次序:历史被收进 $s_i,z_i$,生成时不必重新访问每个历史 token。分母仍然存在,因此“线性注意力”不能简单理解成去掉所有归一化。

单头保存 $D(d_v+1)$ 个标量,单步核心计算为 $O(Dd_v)$,这些量不随当前序列长度增长。整段计算为 $O(NDd_v)$;称其对 $N$ 线性时,前提是头数和特征维度固定。训练时还需考虑投影、激活和反向传播,不能把这个状态大小当成训练总显存。

2.3 二阶 Taylor 映射:扩容不必依赖额外可学习参数

原文 Eq. (4) 写为

$$ K(q,k)=1+q^\top k+\frac{(q^\top k)^2}{2}. $$

它来自 $e^t\approx 1+t+t^2/2$。正文省略了缩放,附录 B.1 的 Listing 1(pp. 20–22)包含维度归一化。将原始投影向量记为 $q,k$,令 $\bar q=q/(d')^{1/4}$、$\bar k=k/(d')^{1/4}$,可以明确写出

$$ \phi(\bar q)=\left[1,\ \bar q,\ \frac{\operatorname{vec}(\bar q\bar q^\top)}{\sqrt2}\right], $$
$$ \phi(\bar q)^\top\phi(\bar k) =1+\frac{q^\top k}{\sqrt{d'}}+\frac{(q^\top k)^2}{2d'}. $$

常数项保存历史 value 的总和,一阶项保存通常的 key–value 关联,二阶项增加不同 key 分量组合与 value 的关联。二阶展开本身没有新的可学习权重,因此可以在给定投影维度下增加状态容量;改变投影维度 $d'$ 时,投影参数和计算成本仍会变化。

三个维度口径需要区分。

表示方式特征维度$d'=16$ 时
完整外积展开$1+d'+d'^2$273
利用对称性的紧凑展开$1+d'+d'(d'+1)/2$153
Algorithm 2 的硬件填充将 273 填充到 64 的倍数320

紧凑展开保留每个交叉项一次:平方项使用 $\bar q_a^2/\sqrt2$,交叉项使用 $\bar q_a\bar q_b$($a表中的 153 不能直接当作该 CUDA 实现每头实际存储的宽度。(pp. 27、35–37)

从公式还可看出 $1+t+t^2/2=((t+1)^2+1)/2>0$,归一化分母在非空前缀下为正;实现另加 $\epsilon$。但二阶多项式只在适当分数范围近似指数:它对很大的负数仍会增长,不能保留指数函数全域的单调性。这是根据公式得到的性质,也解释了为什么不能只凭“近似 softmax”就认定其检索行为与 softmax 相同。

2.4 小窗口与短卷积各自提供什么?

SWA 仅计算 $j\in[\max(1,i-w+1),i]$ 的标准 softmax,保留最近 $w$ 个 token 的 KV。单头状态规模为 $2\min(N,w)d_v$,序列混合计算为 $O(Nwd_v)$。小窗口能做精确近邻比较,但单独堆叠时,感受野依赖窗口和深度。全局线性层让远端信息能够直接影响当前输出。(Sec. 4.2,p. 6)

论文根据 Tensor Core 的 tile 粒度选择小窗口,强调 64–128 这一区域可增加局部表达能力而只带来有限额外开销。实际配置需看 Table 7,不能统一假定都是 64 或 128:

BASED 配置小模型列大模型列
层数 / 隐藏维度27 / 102436 / 1792
线性注意力层57
SWA 层57
BaseConv 层1722
SWA 窗口6416
线性注意力头数 / $d'$16 / 1616 / 16

Table 7,p. 75;该表列名为 355M/1.4B,Table 1 实测参数数列为 363M/1.35B,因此这里以“小/大模型配置列”表述。线性层不使用位置编码,窗口层使用 Rotary;BaseConv 的卷积宽度为 3、扩展倍率为 4、激活为 SiLU。

短门控卷积可写成“投影分支与局部卷积分支逐元素相乘,再输出投影”的形式,其局部缓冲较小。作者用约 20% 线性注意力、20% SWA、60% 短门控卷积组成实际语言模型;“两种注意力结合”是核心构想,实际实验模型还包含第三类 mixer。(Appendix C,p. 31)

衰减设置也随实验变化:10B-token 实验使用附录描述的固定逐头衰减率及输入相关的跨头组合;30B/50B-token 实验不使用该衰减。后者支持“可以不依赖输入相关衰减达到这些结果”,但不应把全部实验都描述成无衰减。

2.5 IO-aware 实现:避免把巨大中间状态反复搬运

朴素实现先在全局显存生成 $\phi(Q),\phi(K)$,再计算因果点积。Taylor 展开扩大了特征宽度;如果逐位置物化外积和累计状态,临时张量与访存开销就会压过理论复杂度优势。

预填充内核将序列分为 16-token 的 tile。下面是对原文块计算的维度统一写法,$\Phi_{Q_b},\Phi_{K_b}$ 为第 $b$ 块的特征矩阵,$M_b$ 为块内因果掩码:

$$ U_b=(M_b\odot\Phi_{Q_b}\Phi_{K_b}^\top)V_b +\Phi_{Q_b}S_{<b},\qquad S_{<b}=\sum_{a<b}\Phi_{K_a}^\top V_a. $$

第一项在块内显式处理因果关系,第二项通过历史累计状态一次性读取所有旧块。$U_b$ 是分子;完整输出还需要对应的归一化分母。

实现进一步利用 Taylor 的结构:零阶项是 value 前缀和;一阶项使用块内矩阵乘和跨块累计;最大的二阶状态分片保存在多个 warp 的寄存器中,必要的归并通过共享内存完成。融合特征映射与读写,避免把展开特征和每步大状态反复写回慢内存。(Sec. 5.2,p. 7;Appendix B.1,pp. 20–24,Algorithm 1)

解码阶段融合外积写入和 query 读出,并为小窗口融合点积、softmax 与加权求和。Algorithm 2 仍会读取并写回持久 KV 状态;减少临时张量不意味着状态本身没有访存成本。微基准的线性注意力测试只针对分子,SWA 测试排除了 KV cache 更新与 Rotary 处理,因而不能把它们的加速直接当作整个模型的加速。(Appendix B.2,pp. 25–30)

三:理论分析与实验设计

3.1 状态下界:用一次通信来理解为什么不能无限压缩

Theorem 3.1(p. 4;正式版本 Theorem F.12,pp. 43–44)给出:能够完成任意关联召回的因果递归模型,最坏情况下需要 $\Omega(N)$ 比特的状态。这里按有限比特计量状态,不允许一个无限精度实数免费编码任意长序列。

证明用单向通信中的 Index 问题:Alice 有 $N$ 个 bit $x_1,\ldots,x_N$,Bob 只有索引 $i$,想知道 $x_i$。Alice 让模型读完 (1,x₁), …, (N,xₙ),只把模型状态发给 Bob;Bob 输入查询 $i$,就应读出 $x_i$。如果状态是 $o(N)$ bit,就得到了低于 Index 通信下界的协议。论文使用的随机通信下界也覆盖成功概率至少 $2/3$ 的协议。

这个下界同时约束 BASED。BASED 所做的是在有限任务与长度上更有效地利用状态;当需要可靠保存的独立关联数量持续增加时,固定大小状态最终会成为瓶颈。

3.2 深度下界:固定容量以外,计算形式也有代价

理论中的 BaseConv 由线性投影、固定卷积和逐元素乘法组成,单层输出关于输入是次数不超过 2 的多项式。堆叠 $L$ 层,次数上限为 $2^L$。对于二进制编码的 token,$d=\log_2 c$,比较两个 token 的每一位并取回值需要高阶交互,因此得到 $\Omega(\log d)=\Omega(\log\log c)$ 的层数下界。(Theorem 3.2,p. 4;Appendix F.4,pp. 46–49)

可用单个匹配对理解次数来源:

$$ \operatorname{EQ}(q,k)=\prod_{a=1}^{d}\left[q_ak_a+(1-q_a)(1-k_a)\right]. $$

每位相等判定为二次,全部相等需要相乘,再乘 value 位产生 $2d+1$ 次项。一个每层最多使次数翻倍的模型,需要足够深度才能表达这种精确条件。这里的结论针对所定义的 data-independent BaseConv;含任意非线性或任意输入相关滤波器的模型不能直接代入此多项式次数论证。

论文还有两个补充方向:

  • 当 $c\le N$、$\log c\le d\le 2^{(\log N)^{1-\epsilon}}$,参数用 $O(\log N)$ bit 表示时,data-independent BaseConv 需要 $\Omega(\epsilon\log\log N)$ 层。证明让 Alice 发送关于 Bob 查询的多项式系数,再由通信下界约束系数数目和深度。(Theorem 3.4,p. 5;Theorem F.14,pp. 44–46)
  • 对 $p$-hot 编码,$d=p c^{1/p}$,附录 Theorem F.29 给出至少 $\lfloor\log_2(2p)\rfloor$ 层的下界。正文 Remark 3.3 写成 input-dependent,附录定理与证明使用 data-independent,本笔记采用后者的明确假设。(pp. 4、49–51)

这些结论解释某些受限门控卷积为什么表达精确匹配需要更多资源,不意味着所有 SSM 的表现都能仅靠此定理排序。论文中的 attention 常数深度构造也允许特定投影、偏置、ReLU 和 MLP,应按其构造条件理解。

3.3 表达能力联系与匹配上界

Proposition F.8 通过算术电路模拟,将特定 Taylor 多项式 LinAtt 运算表示为 $O(\log^2(Nd))$ 层的 BaseConv,并扩大内部表示。该附录定义的是未归一化的多项式运算,不能由此宣称完整 BASED 语言模型与 BaseConv 在相同资源下等价。(pp. 40–43)

Theorem F.30 给出另一方向的构造:在 one-hot 编码、每个 query 至多匹配一个 key、$\bar N=\max(N,d)\ge8$ 的设置下,BaseConv 可以用 $O(\log\log\bar N)$ 层解决 MQAR,同时使用增大的内部长度与宽度,量级为 $O(\bar N\log\bar N)$。构造先累计 key 对应的位置二进制编码,query 选出匹配位置,再通过乘积树转成 one-hot 选择器,最终取出 value。这说明某些设置下双对数深度量级可达到,代价包含内部表示扩张。(pp. 51–74)

3.4 每组实验究竟检验什么?

实验设计与对照指标及要回答的问题
合成 MQAR训练长度 256、4–64 个 KV 对;评测长度 1024、4–256 个 KV 对;改变状态相关超参数准确率越高越好;相同容量下谁更善于召回,并包含长度外推
语言预训练Pile、GPT-2 BPE、相同数据与顺序;约 360M/1.3B;10B 至 30B/50B token整体与 AR 切片 PPL 越低越好
事实抽取/问答SWDE、FDA、改写后的 SQuAD,零样本生成SWDE/FDA 准确率、SQuAD F1;是否利用输入事实
常识与 few-shotLM Eval Harness 六任务;SuperGLUE 0/1/5-shot平均任务分数;通用质量和增加示例后的变化
效率单 H100,区分 prefill 与 decode;BASED 朴素/定制内核、FlashAttention-2、Mamba延迟越低、吞吐越高越好;数学结构能否兑现为速度
DNAHG38 预训练及 GenomicsBenchmarks 分类PPL、top-1 accuracy;能否扩展到另一种序列

Transformer++ 是采用 Rotary、RMSNorm、SwiGLU 的现代 Transformer 对照;Mamba 是主要的强递归对照。较小规模还报告 GLA、RWKV v5、H3 等。A100 80GB 用于预训练,H100 用于效率测试,两个硬件口径不能混用。(Appendix E,pp. 36–39)

3.5 召回任务的特殊口径

Pile AR 切片选择某个已在上下文出现的 bigram 的末尾 token,同时该 bigram 在训练集出现不超过 1250 次;切片在测试集前 1600 万 token 上构造。它强调临时上下文中可检索、训练记忆相对不足的关联。(Appendix E.3,p. 37)

FDA 文档被分成 1920-token 块。提示为“文档块 + 字段名”,生成文本只要包含目标值(忽略大小写)就计正确,所以该准确率是包含答案的比例,不是完整输出严格匹配率。(p. 38)

SQuAD 由 GPT-4 将问答改写为答案位于句尾的陈述,再移除答案作为续写提示。5000 个验证问题筛选后保留 2984 个,约 40% 因不满足句尾条件被丢弃。Table 1 报告 F1,因此这是一套适配基础语言模型的续写式问答评测,不能直接等同标准原始 SQuAD 测试。(pp. 38–39)

四:结果——收益在哪里,代价又在哪里?

4.1 MQAR 与特征映射:容量有效,参数扩张并非唯一手段

Fig. 2 显示增大状态通常提高召回,BASED 在所测配置中比 Mamba 更有效地利用状态;H3、Hyena 的点离经验前沿更远。Fig. 8 增加 NystromFormer、BigBird、Scatterbrain,延续了用容量而非方法标签比较的思路。(pp. 4、32)

图 3:不同特征映射的容量与参数取舍

原文 Fig. 3,p. 6;使用 256 个 KV 对的较难任务。上图横轴为状态字节数,下图为参数量,纵轴均为召回准确率。Taylor、PosELU、ReLU 都接近状态—召回前沿;Taylor 的特点在下图更明显:多项式展开允许较少参数对应较大关联状态。因而结果支持“Taylor 是一种有效的容量配置方式”,没有证明它是唯一有效的核。

4.2 整体语言质量相近,AR 切片更能区分模型

参数规模 / 训练 token模型Pile PPL ↓AR PPL ↓常识六任务均分 ↑
约 360M / 10BTransformer++8.391.8744.08
同上BASED8.652.0743.03
同上Mamba8.642.2143.51
约 360M / 30BTransformer++7.681.8044.75
同上BASED7.771.9345.36
同上Mamba7.732.0245.62
约 1.3B / 50BTransformer++6.281.6553.33
同上BASED6.301.7153.81
同上Mamba6.281.7453.50

来源:Table 1,p. 8;Table 2,p. 32。均分采用原表报告值;作者定义的常识均分对六个任务的分数取均值,不包含 LAMBADA PPL;HellaSwag 和 ARC-C 使用长度归一化选择得分对应的准确率。

在 1.3B/50B 上,BASED 的整体 PPL 与 Mamba 仅差 0.02,AR 切片则从 1.74 降到 1.71;非 AR 切片为 6.82 对 6.78,略逊于 Mamba。收益具有明确的子任务方向,而非所有 token 一致改善。

4.3 下游召回:优于 Mamba,仍有较大的抽取差距

约 1.3B / 50BSWDE Acc. ↑FDA Acc. ↑SQuAD F1 ↑
Transformer++76.5080.4743.47
BASED64.4530.4041.62
Mamba52.7518.5135.92
BASED − Mamba+11.70+11.89+5.70

来源:Table 1,p. 8;最后一行为据表相减,准确率列单位为百分点,F1 列单位为分数点。BASED 在三项上都优于 Mamba,但 FDA 仍比 Transformer++ 低 50.07 个百分点。整体困惑度几乎相同,并不意味着上下文抽取能力相同。

摘要中的提升数字存在不一致。 PDF 摘要与 Sec. 6 使用 10.36,在线摘要使用 6.22;若按 Table 1 对上述三项分数差取未加权均值,得到:

配置三项分数差的均值(本次据 Table 1 计算)
约 360M / 10B3.89
约 360M / 30B10.36
约 1.3B / 10B9.03
约 1.3B / 50B9.76

因此 PDF 所写“1.3B/50B 提高 10.36”不能由该组三列表值直接复算;10.36 对应表中的 360M/30B。由于还混合了 accuracy 与 F1,报告各任务差值最清楚。这些不一致不改变 BASED 在表内三项召回任务均优于 Mamba 的方向。(在线摘要;PDF pp. 1、7–9)

4.4 消融:局部模块、状态宽度和核函数都重要

所有下表消融来自约 362M 参数、10B Pile token 的 Table 6(p. 35),除注明项外使用 $d'=16$、窗口 64、卷积与衰减。

变化Pile PPL ↓AR PPL ↓主要含义
Taylor 默认组合8.652.07对照
Performer,维度 169.088.53此特征预算下召回退化明显
CosFormer,维度 169.032.42较弱
CosFormer,维度 648.822.18增大容量可缩小差距,同时增大投影参数
去掉 SWA8.912.11仅保留短卷积也损失质量
去掉短卷积8.742.09两种注意力已接近最好 PPL,卷积仍有增益
两者都去掉,$d'=24$9.492.29较大特征宽度下仍不及组合;这行不是完全单因素对照
不使用衰减8.652.04整体 PPL 不变、AR 略好,任务收益不统一
SWA 窗口增至 1288.612.06PPL 小幅改善,部分下游指标改善更明显

Taylor 的 $d'=8,16,24,32$ 对应紧凑维度 $45,153,325,561$,整体 PPL 依次为 8.77、8.65、8.58、8.56。从 24 到 32 的 PPL 改善很小,状态却仍显著增大。下游也不是单调关系,例如 FDA 在 $d'=24$ 时为 20.87,在 32 时为 12.89。

将窗口从 64 增至 128 时,SWDE 从 29.16 增至 32.13,FDA 从 11.71 增至 14.39,SQuAD 列从 25.07 增至 31.84;仅看 0.04 的 PPL 变化会低估其对特定召回任务的作用。Table 6 把 SQuAD 列写成 Acc.,Table 1 则写 F1;因此讨论该消融时保留为“SQuAD 分数”,不据此另行断定它更换了评分方法。

4.5 效率:架构与内核必须一起比较

图 4:预填充与递归生成延迟

原文 Fig. 4,p. 8。纵轴实际是延迟毫秒,越低越好。左侧按输入长度比较 prefill,批量为 2;右侧按批量比较固定生成 1024 token 的 decode。蓝色为 FlashAttention,黄色为 Mamba,橙色虚线/实线为 BASED 朴素/定制内核。缺失曲线后段表示 OOM。

在 1.3B、4K 输入的 prefill 设置下,正文报告 BASED 相对 FlashAttention-2 吞吐高约 56%、相对 Mamba 高约 44%。在单 H100、批量 128、生成 1024 token、不含 prefill 并使用 CUDA graphs 的设置下,正文报告约 24 倍 FlashAttention-2 吞吐,同时约为 Mamba 吞吐的 95%。(Sec. 6.2,p. 9)

Table 1 列出的 1.3B prefill 吞吐为 161.71、103.50、112.22 token/ms(BASED、Transformer++、Mamba),与 56%/44% 的相对变化一致。但该表 caption 将 generation 写为 2048 token,Fig. 4 与正文的 24 倍结果写为 1024 token;不能把表中所有速度值强行归为同一次测量。360M/30B 行还重复了 1.3B 行的效率数值,本笔记不据这些行推导小模型速度结论。

橙色虚线和实线之间的差距尤其关键:朴素 BASED 在某些 prefill 点比优化后的注意力更慢。论文提供的是混合架构加定制实现的收益;只实现数学递推式不足以得到论文中的吞吐。

4.6 扩展结果:few-shot 与 DNA

SuperGLUE 在约 360M、10B-token 模型上,先在任务内部平均多个指标,再对任务平均。报告结果为:

模型0-shot1-shot5-shot
BASED45.747.248.0
Transformer++46.347.848.9
Mamba46.646.945.2

来源:Table 3,p. 33。BASED 和 Transformer++ 随示例增加而提高,Mamba 在 5-shot 时下降。这个现象支持召回与上下文利用能力有关的解释;它是这些模型与提示协议下的观察。

DNA 使用约 46–49M 参数模型,在 HG38 上从头训练 10B token,并分别测试 1K/4K/8K 上下文。BASED 的 PPL 为 2.51/2.50/2.49,Mamba 为 2.51/2.49/2.49,Transformer++ 为 2.52/2.50/2.51。在五项 DNA 分类任务中,BASED 的准确率为 74.6、89.5、89.5、96.8、79.0,呈现竞争力但不全面领先。与引用的既有结果比较时,作者指出 tokenizer 存在差异。(Tables 4–5,pp. 33–34)

五:结论与研究启发

BASED 最有说服力的结论是:固定状态模型的召回能力既依赖容量,也依赖如何组织读写和局部计算。二阶特征展开提供较大的关联空间,小窗口与短卷积帮助形成精确局部关系;融合内核使这些较大状态仍可高效运行。

实际选择时应同时看三项:任务需要保存多少临时事实、模型为这些事实配置多少状态、硬件能以多大代价读写该状态。 Pile PPL、参数量和渐近复杂度中的任何一个都不能单独回答这三个问题。

本工作的实证范围主要是 360M–1.3B、最高 50B 语言预训练 token,以及特定 GPU 和批量下的效率测试。最明显的残余问题是 FDA 抽取与 Transformer++ 仍有大差距;更大的状态虽可能提升召回,也会增加带宽与存储成本。理论下界则提醒我们,局部精度与更好的内核不会消除无限增长的记忆需求。

由此自然引出两个后续问题:能否按输入的有效关联数量动态扩展记忆,而不是按原始序列长度统一缓存;能否在相同实际字节数、精度、层数和内核条件下比较不同写入/遗忘规则,更清楚地区分容量收益与更新规则收益。这些是本文启发的研究方向。