File size: 10,953 Bytes
5c6c847 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """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)
|