"""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)