File size: 14,791 Bytes

# ZeroPoint：分层记忆槽网络

> 为 zeroOS 设计的序列模型架构。用它搭的语言模型叫 **ZeroX**，训练框架是 **ZeroCG**。（原名 ZeroMem。）文档结构参照《Attention Is All You Need》（Vaswani et al., 2017），
> 每节开头注明对应论文的哪一节，方便对照阅读。
>
> 代码：[`models/zeropoint.py`](../models/zeropoint.py) ·
> 自检：[`tests/test_zeropoint.py`](../tests/test_zeropoint.py) ·
> 对照实验：[`experiments/compare.py`](../experiments/compare.py) → [`compare.md`](../experiments/compare.md) ·
> 消融实验：[`experiments/ablation.py`](../experiments/ablation.py) → [`ablation.md`](../experiments/ablation.md) ·
> 精确回忆：[`experiments/recall.py`](../experiments/recall.py) → [`recall.md`](../experiments/recall.md) ·
> 历次结果存档：[`experiments/results/`](../experiments/results/)
>
> 状态：**v0（研究原型）**。

---

## 摘要

Transformer 用自注意力取代了循环，换来了训练时的完全并行；代价是推理时每生成一个字都要回看之前所有的字，
KV 缓存随文本长度线性增长。这对 GPU 服务器不是问题，对**没有 GPU、内存按页精打细算的操作系统内核**却是。

我们提出 **ZeroPoint**：用一组**带不同遗忘速度的记忆槽**取代自注意力。每读一个字，模型决定把它**写进**哪个槽、
从哪些槽**读出**；槽的遗忘速度按"寄存器 → 缓存 → 内存"的存储层级初始化，并在训练中学习。
ZeroPoint 有两种数学上等价的算法：训练时用**分块并行**（块长 32，每块一次矩阵乘，块间串行传递状态），
推理时**逐字递推**，状态大小固定，与文本长度无关，也不需要位置编码。

在 **115,240** 参数、1,558,395 字符的字符级语言建模上，ZeroPoint 与同规模 Transformer 的验证 loss
在 2× 种子标准差内无显著差异（1.5214 ± 0.0134 vs 1.5452 ± 0.0100，差 0.0238，判定线 0.0268；**不声称更好**）。
推理状态是后者在第 1024 个字时的 **1/128**（2,048 vs 262,144 个浮点数），
4 倍训练长度的文本上 loss 升高 0.037 nats/字符（64 倍时为 0.219）。

## 1 引言　〔论文 §1〕

论文 §1 的出发点是：循环网络"天生是串行的"，第 t 步必须等第 t−1 步算完，没法在一条样本内并行。
自注意力把任意两个位置之间的路径缩短到 O(1)，训练可以完全并行。

但论文没有讨论的另一面是**推理**：生成第 t 个字时，自注意力要和前面 t−1 个字的 Key/Value 逐一比较。
这些 Key/Value 必须一直存着（KV 缓存）。对 zeroOS 来说：

| | GPU 服务器 | zeroOS 内核 |
|---|---|---|
| 内存 | 几十 GB 显存 | 内核堆，按 4KB 页分配 |
| 算力 | 上万个并行单元 | 一个 CPU 核，x87 浮点 |
| 运行方式 | 一次处理一大批 | 用户敲一个字处理一个字 |

**我们想要的是两者的好处：训练几乎像 Transformer 一样并行（分块并行，串行步数是长度的 1/32），推理像 RNN 一样只带一个固定大小的状态。**

## 2 背景　〔论文 §2〕

论文 §2 提到 End-to-End Memory Networks [34] 用"循环的注意力"读外部记忆。ZeroPoint 和这条线以及
后来的"线性注意力 / 带衰减的循环模型"家族是近亲，不是凭空出现的：

| 已有工作 | 相同点 | ZeroPoint 的不同 |
|---|---|---|
| 记忆网络 / 神经图灵机 | 软寻址读写记忆槽 | 不需要逐步串行训练；每个槽有自己的遗忘速度 |
| 线性注意力（Katharopoulos 2020） | 并行 / 递推双形式 | 状态是 M 个可寻址的槽，不是 D×D 的外积矩阵 |
| RetNet（Sun 2023） | 衰减 λ^(t−u) 实现双形式、不需要位置编码 | λ 可学且按存储层级初始化；写入时 softmax 竞争选槽 |
| Mamba / 状态空间模型 | 递推推理、状态固定 | 不需要专门的扫描算子，一次矩阵乘就能并行 |

真正属于 ZeroPoint 的是这个组合：**多时间尺度的可学习遗忘 + 写时竞争选槽 + 写强度门 + 一次矩阵乘的并行形式**，
以及"为一个没有 GPU 的内核设计"这个出发点。

## 3 模型结构　〔论文 §3〕

### 3.1 层的堆叠　〔论文 §3.1〕

和论文的解码器一样：N 个相同的层堆叠，每层两个子层——**记忆混合层**（代替多头自注意力）和**逐位置前馈网络**，
每个子层外面套残差连接和 LayerNorm。

两处和论文不同：

- **LayerNorm 放在子层前面（Pre-LN）**：`x + Sublayer(LayerNorm(x))`。论文用的是 Post-LN
  `LayerNorm(x + Sublayer(x))`，后来的 GPT-2 等工作发现 Pre-LN 训练更稳，不太依赖学习率预热。
  对照组 Transformer 用的也是 Pre-LN，两边一致。
- **只有解码器（自回归）**。论文是做翻译的编码器-解码器；zeroOS 的任务（续写、意图识别）用不上编码器。

### 3.2 记忆混合层　〔论文 §3.2 注意力〕

#### 3.2.1 读与写　〔对应 §3.2.1 缩放点积注意力〕

论文的注意力是"拿查询去和所有的键比"。ZeroPoint 反过来：**每个字在写入时就决定自己存到哪里**，
读的时候只和 M 个槽打交道，不和之前的字直接打交道。

设第 t 个字的向量为 `x_t`，记忆槽为 `S[1..M]`（每个是 D 维向量），保留率为 `λ[1..M]`：

```
写地址   w_t = softmax(x_t · Ww) × sigmoid(x_t · Wg)    写到哪些槽 × 写多重
写入值   v_t = x_t · Wv
更新     S[m] ← λ[m] · S[m]  +  w_t[m] · v_t           所有槽先按自己的速度遗忘，再写入
读地址   a_t = softmax(x_t · Wr)                        从哪些槽读
读出     r_t = Σ_m a_t[m] · S[m]
输出     LayerNorm(r_t) · Wo
```

- **写地址用 softmax（槽之间竞争）**：一个字要么写进短期槽，要么写进长期槽，模型得决定"这个字值不值得长期记住"。
- **写强度门**：空格、换行这类没信息量的字可以几乎不写，不冲掉已有的记忆。
- **读出后做 LayerNorm**：长期槽可能累积了几十个字，短期槽只有一两个，量级差别很大。
  这和论文 §3.2.1 除以 √d_k 的动机相同——别让数值落到 softmax / 后续层不好处理的区间。
- **槽的初始状态 S0 可学习**（不是全 0）：否则序列开头几个字读出来的向量接近 0，
  后面的 LayerNorm 会把噪声放大。

#### 3.2.2 多头　〔对应 §3.2.2 多头注意力〕

和论文一样分 H 个头，每头 D = d/H 维，各自有一套读写地址和 M 个槽，最后拼起来乘 Wo。
论文的理由同样适用：单头会把不同的"记什么"平均到一起。

#### 3.2.3 保留率与存储层级　〔ZeroPoint 独有〕

```
λ = 0.5 + 0.5 · sigmoid(θ)       λ ∈ (0.5, 1)，θ 可学
```

初始化时，每头的 M 个槽的**半衰期**（旧内容衰减到一半要经过几个字）按等比数列从 1.5 铺到 64：

```
槽 0      λ ≈ 0.63   半衰期 1.5 个字    ← 像寄存器：记"刚才那个字"
槽 1~14   λ 逐渐变大                    ← 像缓存：记"这个词""这一句"
槽 15     λ ≈ 0.989  半衰期 64 个字     ← 像内存：记"这一段在讲什么"
```

论文 §3.5 的正弦位置编码"波长从 2π 到 10000·2π 成等比数列"，ZeroPoint 的半衰期等比铺开是同一个思路：
用一组不同尺度的分量覆盖从很近到很远的距离。区别是论文把尺度编码进**输入**，ZeroPoint 把尺度编码进**记忆本身**。

### 3.3 逐位置前馈网络　〔论文 §3.3〕

与论文相同的结构 `FFN(x) = act(x·W1 + b1)·W2 + b2`，内层宽度 4d（论文 512 → 2048 也是 4 倍）。
激活函数用 GELU（论文是 ReLU；GPT 系列之后普遍用 GELU），对照组也用 GELU。

### 3.4 嵌入与输出　〔论文 §3.4〕

可学习的字嵌入，输出前 LayerNorm + 线性层 + softmax。没有采用论文的"输入/输出嵌入共享权重"，对照组也没有。

### 3.5 位置信息　〔论文 §3.5〕

论文说："模型里没有循环也没有卷积，所以必须注入位置信息。"

**ZeroPoint 不需要位置编码。** 它本质上是一个（可并行的）循环结构：第 t 个字读到的内容里，
u 个字之前写入的信息被乘上了 λ^(t−u)。不同的槽衰减速度不同，所以"某条信息在哪些槽里还剩多少"
本身就告诉了模型它是多久以前的。消融实验 (G) 专门检验了这一点。

## 4 为什么用记忆槽　〔论文 §4，表 1〕

论文表 1 比较了三个指标：每层计算量、最少串行步数、任意两位置之间的最长路径。
我们加上对内核最重要的两列：推理时每个字的计算量、推理时要保存的状态。

n = 序列长度，d = 宽度，M = 每头槽数（实验里 M = 16，d = 64）。

| 层类型 | 每层计算量（训练） | 串行步数 | 最长路径 | 推理：每字计算 | 推理：状态大小 |
|---|---|---|---|---|---|
| 自注意力 | O(n²·d) | O(1) | O(1) | O(n·d)，越往后越慢 | O(n·d)，随长度增长 |
| 循环网络 | O(n·d²) | O(n) | O(n) | O(d²) | O(d) |
| **ZeroPoint** | O(n·c·(M+d))（分块并行，c = 32） | O(n/c) | **O(1)**，但信号按 λ^(t−u) 衰减 | **O(d²)** | **O(M·d)，固定** |

关键在于 ZeroPoint 有两种等价算法：

**并行形式（训练用）**。把递推展开：

```
r_t = Σ_{u ≤ t} A[t,u] · v_u,    A[t,u] = Σ_m a_t[m] · w_u[m] · λ[m]^(t−u)
```

利用 `λ^(t−u) = λ^t · λ^(−u)`：

```
A = ( a ⊙ λ^t ) @ ( w ⊙ λ^(−u) )ᵀ    再把 u > t 的部分清零（因果）
```

只要一次矩阵乘，和注意力的 `Q·Kᵀ` 同样能整段并行。但 `λ^(−u)` 最大是 2^u，float32 到 u ≈ 120 就溢出。

所以实际用的是**分块并行**：序列切成每块 c = 32 个字，进入一块时槽的状态是 `S_in`：

```
块内    r_i  = Σ_{j ≤ i} A[i,j] · v_j            A 同上，但 i、j 是块内下标，λ^(−j) 最大 2^31
跨块    r_i += (a_i ⊙ λ^(i+1)) @ S_in
块尾    S_out = λ^c ⊙ S_in + (w ⊙ λ^(c−1−j))ᵀ @ v
```

块内是矩阵乘，块和块之间只传一个 M × d 的状态。训练时每层计算量 O(n·c·(M+d))，
**序列多长都行**（自检里 float32 跑 1000 个字不溢出，和递推形式一致）。

**递推形式（推理用）**。逐字更新 `S`，每层只存 M·d 个浮点数，跟文本多长无关。

两种形式结果一致：自检在**所有 64 种开关组合**下、300 个字（跨多个块）上逐位比对，float64 误差 < 1e-8。
另有一份固定的参考数据 `tests/golden/zeropoint_ref.json`，以后任何实现（包括内核里的 C 版本）都拿它对答案。
**训练好的模型可以直接用递推形式放进内核**：不用 KV 缓存，每层固定一块 M × d 的状态，不需要动态分配。

关于"最长路径 O(1)"要说实话：u 处写入的信息确实一步就能被 t 处读到，但已经乘上了 λ^(t−u)。
注意力能"原样找回"任意远处的某个字，ZeroPoint 只能找回**还没忘干净、而且没被其他写入冲淡**的信息。

## 5 训练　〔论文 §5〕

| | 论文 | ZeroPoint 实验 |
|---|---|---|
| 数据 | WMT 2014 英德，450 万句对，BPE 37000 | Python 标准库源码 150 万字符，字符级，词表 96 |
| 批次 | 约 25000 个词元 | 32 × 64 = 2048 个字符 |
| 硬件 | 8 块 P100 | 一台 Mac 的 CPU（numpy，自研框架 zerocg） |
| 优化器 | Adam β=(0.9, 0.98)，公式 3（预热 4000 步后按 1/√步数衰减） | AdamW β=(0.9, 0.999)，权重衰减 0.01；预热 100 步后余弦退火到 10% |
| 正则化 | 残差 dropout 0.1、标签平滑 0.1 | 无（模型小、训练步数少，没观察到过拟合） |
| 步数 | 10 万步（12 小时） | 2000 步（每个模型约 2 分钟） |

ZeroPoint 和对照组用**完全相同**的训练设置、相同的随机种子（决定初始化和数据顺序）。

## 6 结果　〔论文 §6〕

### 6.1 对照实验　〔论文 §6.1，表 2〕

> 结果见 [`experiments/compare.md`](../experiments/compare.md)（`compare.py` 自动生成）。

### 6.2 模型变体（消融）　〔论文 §6.2，表 3〕

> 结果见 [`experiments/ablation.md`](../experiments/ablation.md)（`ablation.py` 自动生成）。

### 6.3 长度外推　〔论文 §3.5 的猜想〕

论文选正弦位置编码的理由是它"**可能**可以外推到比训练时更长的序列"，但没有给出实验。
我们在 64 字上训练 ZeroPoint 和正弦位置编码的 Transformer，然后在 4,096 字的文本上分段测试
（32 条验证序列 × 3 个种子，两种模型用同一种分块并行算法）。

- ZeroPoint：64~255 字（训练长度的 4 倍）loss 为 1.4742 ± 0.0341，比训练长度内（0~63：1.4372）高 **+0.037**；
  1024~4095 字（64 倍）为 1.6932 ± 0.0305，高 **+0.219 ± 0.063**。**外推会变差，但没有崩掉。**
- 正弦位置编码的 Transformer：出了训练长度就升到约 4.64（≈ ln 96 = 4.56，即随机水平）。
  这是 8 条序列 × 3 个种子的测量，基线权重未保留，无法用 32 条序列复测。

> 早期 8 条序列的结果（1.5550 / 1.6164 / 1.6410 / 1.7040）抽样噪声大，已由上面的 32 条序列结果取代。
> 原始数据见 [`experiments/results/`](../experiments/results/)。

### 6.4 精确回忆　〔无结论，不能用来比较架构〕

> 结果见 [`experiments/recall.md`](../experiments/recall.md)。

`recall.py` 的训练配置（恒定 lr 3e-3）把两个架构都困在"两个候选值各猜一半"的局部解（loss ≈ ln 2）。
诊断显示换成 1e-3 余弦退火后，只训 N=2 的 Transformer 在 8,000 步内到 100%。在按修正后的学习率重跑之前，
这张 20,000 步的表只能说明训练配置有问题，**不能说明 ZeroPoint 或 Transformer 谁更擅长精确回忆**。

## 7 结论与局限　〔论文 §7〕

**局限**（都是实话）：

- **规模很小**：只在约 11.5 万参数、150 万字符、单一语料上测过，结论不能外推到大模型。
- **精确回忆**：只有 M 个槽，信息会被压缩、混合。**推测**（未验证）它不擅长原样复制很久以前的内容；现有 recall 实验因学习率配置问题无结论。
- **没有编码器-解码器版本**：论文的核心应用是翻译；ZeroPoint 目前只做了解码器。

**下一步**：

1. 如果精确回忆确实弱：加一层"滑动窗口注意力"（只看最近 32 个字，状态仍然固定）做混合结构；
2. 在 zeroOS 的意图识别任务上训练；
3. 用 C 实现递推形式放进内核（x87 浮点）——这是 ZeroPoint 存在的原因。