ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """ZeroPoint：分层记忆槽网络。 | |
| 设计文档见 docs/ZeroPoint.md。核心是 MemoryMixer（代替 Transformer 里的自注意力）： | |
| 每层有 H 个头，每个头有 M 个记忆槽 S[m]（D 维向量），每个槽有自己的保留率 λ[m]。 | |
| 读到第 t 个字时： | |
| w_t = softmax(x_t Ww) * sigmoid(x_t Wg) 写地址（槽之间竞争）× 写强度 | |
| S[m] ← λ[m] · S[m] + w_t[m] · v_t 所有槽先按自己的速度遗忘，再写入 | |
| a_t = softmax(x_t Wr) 读地址 | |
| r_t = Σ_m a_t[m] · S[m] 从槽里读出 | |
| 两种等价算法： | |
| forward() 分块并行，训练用。序列切成每块 CHUNK 个字： | |
| 块内 r = (A ⊙ 因果掩码) @ v，A[i,j] = Σ_m a_i[m] w_j[m] λ[m]^(i-j) | |
| 利用 λ^(i-j) = λ^i · λ^(-j)，一次矩阵乘 | |
| 跨块 r += (a ⊙ λ^(i+1)) @ S_in S_in = 进入这一块时的槽状态 | |
| 块尾 S_out = λ^c ⊙ S_in + (w ⊙ λ^(c-1-j))ᵀ @ v | |
| λ^(-j) 最大 2^(CHUNK-1)，不会溢出，所以序列多长都行 | |
| step() 递推形式，推理用：每个字 O(H·M·D)，状态大小固定，没有 KV 缓存 | |
| 槽的初始状态 S0 是可学习参数（不是全 0），这样序列开头读出来的向量不会接近 0、被 LayerNorm 放大噪声。 | |
| ema=True（指数滑动平均写入）：S[m] ← λ[m]·S[m] + (1-λ[m])·w_t[m]·v_t | |
| 槽的值永远落在"写进去过的值"的范围里，不会随序列变长而越积越大。 | |
| 系统直觉模型要连续跑几天，用这个；实现上只是把写地址 w 乘上 (1-λ)，并行和递推两种形式都不用改。 | |
| """ | |
| import numpy as np | |
| from zerocg import Tensor, concat | |
| from zerocg import functional as F | |
| from zerocg import nn | |
| from .transformer import sinusoidal_encoding | |
| LAMBDA_MIN = 0.5 # λ = 0.5 + 0.5·sigmoid(θ) ∈ (0.5, 1) | |
| CHUNK = 32 # 分块并行的块长：块内 λ^(-j) 最大 2^31，float32 绰绰有余 | |
| def _layernorm_np(x, ln): | |
| mu = x.mean(-1, keepdims=True) | |
| var = ((x - mu) ** 2).mean(-1, keepdims=True) | |
| return (x - mu) / np.sqrt(var + ln.eps) * ln.weight.data + ln.bias.data | |
| def _softmax_np(x): | |
| e = np.exp(x - x.max(-1, keepdims=True)) | |
| return e / e.sum(-1, keepdims=True) | |
| def _gelu_np(a): | |
| return 0.5 * a * (1 + np.tanh(np.sqrt(2 / np.pi) * (a + 0.044715 * a ** 3))) | |
| class MemoryMixer(nn.Module): | |
| """参数里的开关只给消融实验用（experiments/ablation.py），默认值就是 ZeroPoint 本身： | |
| gate 写强度门 | |
| write "softmax"：写入时槽之间竞争；"sigmoid"：每个槽独立决定写不写 | |
| init "hier"：半衰期按存储层级铺开；"uniform"：所有槽一样 | |
| learn_lambda 保留率 λ 是否可学 | |
| ema 写入按 (1-λ) 缩放（指数滑动平均），长时间运行不漂移 | |
| """ | |
| def __init__(self, n_embd, n_head, n_slot, max_half_life=64.0, | |
| gate=True, write="softmax", init="hier", learn_lambda=True, ema=False): | |
| assert n_embd % n_head == 0 and write in ("softmax", "sigmoid") and init in ("hier", "uniform") | |
| self.H, self.M, self.D = n_head, n_slot, n_embd // n_head | |
| self.write_mode = write | |
| self.ema = ema | |
| self.read = nn.Linear(n_embd, n_head * n_slot) | |
| self.write = nn.Linear(n_embd, n_head * n_slot) | |
| self.gate = nn.Linear(n_embd, n_head) if gate else None | |
| self.value = nn.Linear(n_embd, n_embd) | |
| self.out = nn.Linear(n_embd, n_embd) | |
| self.norm = nn.LayerNorm(self.D) # 每个头读出来的向量各自归一化（不同槽的量级差别很大） | |
| # 初始化成"存储层级"：槽 0 的半衰期约 1.5 个字（像寄存器），最后一个约 max_half_life 个字（像内存） | |
| if init == "hier": | |
| half_life = np.geomspace(1.5, max_half_life, n_slot) | |
| else: # 对照：所有槽都取层级的几何中点 | |
| half_life = np.full(n_slot, np.sqrt(1.5 * max_half_life)) | |
| lam = 0.5 ** (1 / half_life) | |
| theta = np.tile(np.log((2 * lam - 1) / (2 - 2 * lam)), (n_head, 1)) # sigmoid 的反函数，(H, M) | |
| self.theta = nn.Parameter(theta) if learn_lambda else Tensor(theta) | |
| self.s0 = nn.Parameter(np.random.randn(n_head, n_slot, self.D) * 0.1) # 可学习的初始状态 | |
| self.usage = None # 最近一次 forward 的 (读地址, 写地址) 平均值，各 (H, M)，用来查槽有没有塌缩 | |
| def retention(self): | |
| """每个槽的保留率 λ，形状 (H, M)。""" | |
| return F.sigmoid(self.theta) * (1 - LAMBDA_MIN) + LAMBDA_MIN | |
| def forward(self, x): | |
| B, T, C = x.shape | |
| H, M, D = self.H, self.M, self.D | |
| a = F.softmax(self.read(x).reshape(B, T, H, M), axis=-1).transpose(0, 2, 1, 3) # (B,H,T,M) | |
| w = self.write(x).reshape(B, T, H, M) | |
| w = F.softmax(w, axis=-1) if self.write_mode == "softmax" else F.sigmoid(w) | |
| if self.gate is not None: | |
| w = w * F.sigmoid(self.gate(x)).reshape(B, T, H, 1) | |
| w = w.transpose(0, 2, 1, 3) # (B,H,T,M) | |
| v = self.value(x).reshape(B, T, H, D).transpose(0, 2, 1, 3) # (B,H,T,D) | |
| self.usage = (a.data.mean(axis=(0, 2)), w.data.mean(axis=(0, 2))) | |
| lam = self.retention() # (H, M) | |
| if self.ema: | |
| w = w * (1 - lam).reshape(1, H, 1, M) | |
| log_lam = lam.log().reshape(H, 1, M) | |
| S = self.s0.reshape(1, H, M, D) # 进入第一块时的状态（对 batch 广播） | |
| outs = [] | |
| for s in range(0, T, CHUNK): | |
| e = min(s + CHUNK, T) | |
| c = e - s | |
| ac, wc, vc = a[:, :, s:e], w[:, :, s:e], v[:, :, s:e] | |
| i = Tensor(np.arange(c, dtype=np.float64).reshape(1, c, 1)) | |
| decay = (log_lam * i).exp() # λ^i, (H,c,M) | |
| A = (ac * decay) @ (wc * (log_lam * -i).exp()).transpose(0, 1, 3, 2) # (B,H,c,c) | |
| A = A.masked_fill(np.triu(np.ones((c, c), dtype=bool), 1), 0.0) | |
| outs.append(A @ vc + (ac * decay * lam.reshape(H, 1, M)) @ S) # 块内 + 跨块 | |
| if e < T: | |
| tail = (log_lam * Tensor(np.arange(c - 1, -1, -1, dtype=np.float64).reshape(1, c, 1))).exp() | |
| S = (lam ** c).reshape(1, H, M, 1) * S + (wc * tail).transpose(0, 1, 3, 2) @ vc | |
| r = outs[0] if len(outs) == 1 else concat(outs, axis=2) | |
| return self.out(self.norm(r).transpose(0, 2, 1, 3).reshape(B, T, C)) | |
| # ---------- 递推形式（numpy，无梯度）---------- | |
| def init_state(self, B): | |
| return np.broadcast_to(self.s0.data, (B, self.H, self.M, self.D)).copy() | |
| def step(self, x, S): | |
| """x: (B, C) 一个字的输入；S: (B, H, M, D) 记忆槽。返回 (输出, 新状态)。""" | |
| B = x.shape[0] | |
| H, M, D = self.H, self.M, self.D | |
| lin = lambda layer, z: z @ layer.weight.data + layer.bias.data # noqa: E731 | |
| a = _softmax_np(lin(self.read, x).reshape(B, H, M)) | |
| w = lin(self.write, x).reshape(B, H, M) | |
| w = _softmax_np(w) if self.write_mode == "softmax" else F.sigmoid_np(w) | |
| if self.gate is not None: | |
| w = w * F.sigmoid_np(lin(self.gate, x))[:, :, None] | |
| v = lin(self.value, x).reshape(B, H, D) | |
| lam = self.retention().data | |
| if self.ema: | |
| w = w * (1 - lam)[None] | |
| self.last_addr = (a, w) # 诊断用：这一步的读地址和写地址，(B, H, M) | |
| S = lam[None, :, :, None] * S + w[..., None] * v[:, :, None, :] | |
| r = np.einsum("bhm,bhmd->bhd", a, S) | |
| r = _layernorm_np(r, self.norm).reshape(B, H * D) | |
| return lin(self.out, r), S | |
| class Block(nn.Module): | |
| def __init__(self, n_embd, n_head, n_slot, dropout=0.0, mlp_hidden=None, **mixer_opts): | |
| hidden = mlp_hidden or 4 * n_embd | |
| self.ln1 = nn.LayerNorm(n_embd) | |
| self.mix = MemoryMixer(n_embd, n_head, n_slot, **mixer_opts) | |
| self.ln2 = nn.LayerNorm(n_embd) | |
| self.fc1 = nn.Linear(n_embd, hidden) | |
| self.fc2 = nn.Linear(hidden, n_embd) | |
| self.drop = nn.Dropout(dropout) | |
| def forward(self, x): | |
| x = x + self.drop(self.mix(self.ln1(x))) | |
| return x + self.drop(self.fc2(F.gelu(self.fc1(self.ln2(x))))) | |
| def step(self, x, S): | |
| h, S = self.mix.step(_layernorm_np(x, self.ln1), S) | |
| x = x + h | |
| h = _gelu_np(_layernorm_np(x, self.ln2) @ self.fc1.weight.data + self.fc1.bias.data) | |
| return x + h @ self.fc2.weight.data + self.fc2.bias.data, S | |
| class ZeroX(nn.Module): | |
| """字符级语言模型。默认没有位置编码：位置信息来自各个槽不同的遗忘速度。 | |
| pos="sin" 只给消融实验用：额外加上论文的正弦位置编码，看看还有没有用。 | |
| mlp_hidden 默认 4 × n_embd；消融实验用它把各变体的参数量对齐。""" | |
| def __init__(self, vocab_size, n_embd=64, n_head=4, n_slot=16, n_layer=2, dropout=0.0, | |
| pos="none", mlp_hidden=None, **mixer_opts): | |
| assert pos in ("none", "sin") | |
| self.pos_kind = pos | |
| self.tok = nn.Embedding(vocab_size, n_embd) | |
| self.blocks = [Block(n_embd, n_head, n_slot, dropout, mlp_hidden, **mixer_opts) for _ in range(n_layer)] | |
| self.ln_f = nn.LayerNorm(n_embd) | |
| self.head = nn.Linear(n_embd, vocab_size) | |
| def forward(self, idx): | |
| x = self.tok(idx) | |
| if self.pos_kind == "sin": | |
| x = x + Tensor(sinusoidal_encoding(idx.shape[1], x.shape[-1])) | |
| for b in self.blocks: | |
| x = b(x) | |
| return self.head(self.ln_f(x)) | |
| def init_state(self, B): | |
| return [b.mix.init_state(B) for b in self.blocks] | |
| def step(self, idx_t, states, t=0): | |
| """idx_t: (B,) 第 t 个字。返回 (logits (B, V), 新状态)。""" | |
| x = self.tok.weight.data[idx_t] | |
| if self.pos_kind == "sin": | |
| x = x + sinusoidal_encoding(t + 1, x.shape[-1])[t] | |
| new = [] | |
| for b, S in zip(self.blocks, states): | |
| x, S = b.step(x, S) | |
| new.append(S) | |
| return _layernorm_np(x, self.ln_f) @ self.head.weight.data + self.head.bias.data, new | |
| def forward_recurrent(self, idx): | |
| """逐字递推跑完整个序列，返回 (B, T, V) 的 logits。""" | |
| idx = np.asarray(idx) | |
| states = self.init_state(idx.shape[0]) | |
| out = [] | |
| for t in range(idx.shape[1]): | |
| logits, states = self.step(idx[:, t], states, t) | |
| out.append(logits) | |
| return np.stack(out, axis=1) | |
| def state_floats(self): | |
| """推理时需要保存的状态大小（浮点数个数），和序列长度无关。""" | |
| return sum(b.mix.H * b.mix.M * b.mix.D for b in self.blocks) | |