"""前沿架构：手算规模的专家路由、线性注意力与记忆更新；不是语言模型训练。""" from argparse import ArgumentParser import hashlib import json from pathlib import Path import platform import numpy as np def softmax(values): weights = np.exp(values - values.max(axis=-1, keepdims=True)) return weights / weights.sum(axis=-1, keepdims=True) def expert(x, first, second): return np.maximum(x @ first, 0) @ second def routed(x, router, first, second, top_k=2, capacity=None): """按词元顺序分配；超过容量就丢弃该次分配，不重新归一化。""" n_experts = len(first) if not 1 <= top_k <= n_experts: raise ValueError("top_k 必须位于专家数量范围内") if capacity is not None and capacity < 0: raise ValueError("容量不能为负") ids = np.argsort(-(x @ router), axis=-1, kind="stable")[:, :top_k] scores = np.take_along_axis(x @ router, ids, axis=1) weights = softmax(scores) output = np.zeros_like(x) accepted = np.zeros(n_experts, dtype=int) requested = np.bincount(ids.ravel(), minlength=n_experts) dropped = 0 for token in range(len(x)): for slot, index in enumerate(ids[token]): if capacity is not None and accepted[index] >= capacity: dropped += 1 continue output[token] += weights[token, slot] * expert(x[token], first[index], second[index]) accepted[index] += 1 return output, ids, weights, requested, accepted, dropped def linear_recurrent(q, k, v): """q/k 已是非负特征；写入当前位置后读取，包含当前位置的因果注意力。""" state = np.zeros((k.shape[1], v.shape[1])) normalizer = np.zeros(k.shape[1]) outputs = [] for query, key, value in zip(q, k, v): state += np.outer(key, value) normalizer += key denominator = query @ normalizer if denominator <= 0: raise ValueError("当前查询没有正的匹配权重") outputs.append(query @ state / denominator) return np.array(outputs), state, normalizer def linear_explicit(q, k, v): similarity = np.tril(q @ k.T) return similarity @ v / similarity.sum(axis=1, keepdims=True) def delta_memory(k, v, decay, write): """S 的形状是 d_k × d_v；门值为手工给定，没有训练投影或门。""" state = np.zeros((k.shape[1], v.shape[1])) history = [] for key, value, alpha, beta in zip(k, v, decay, write): state = alpha * state error = value - key @ state state = state + beta * np.outer(key, error) history.append(state.copy()) return np.array(history) def scalar_state(inputs, decay): state, history = 0., [] for value in inputs: state = decay * state + value history.append(state) return history def run(): rng = np.random.default_rng(7) x = np.array([[1 + i / 10, i / 10] for i in range(6)]) router = np.array([[2., 1., -1., -2.], [0., .2, .1, 0.]]) first = rng.normal(size=(4, 2, 3)) second = rng.normal(size=(4, 3, 2)) output, ids, weights, requested, accepted, dropped = routed(x, router, first, second) # 全算后遮罩仅作独立执行路径对照，不作为“稀疏省计算”的实现。 all_outputs = np.stack([expert(x, a, b) for a, b in zip(first, second)], axis=1) mask = np.zeros((len(x), len(first))) np.put_along_axis(mask, ids, weights, axis=1) reference = np.einsum("ne,ned->nd", mask, all_outputs) limited, _, _, _, limited_load, limited_drops = routed( x, router, first, second, capacity=3) q = rng.uniform(.1, 1., size=(7, 3)) k = rng.uniform(.1, 1., size=(7, 3)) v = rng.normal(size=(7, 2)) recurrent, _, _ = linear_recurrent(q, k, v) explicit = linear_explicit(q, k, v) scores = q @ k.T scores[np.triu_indices(len(q), 1)] = -np.inf conventional = softmax(scores) @ v # 单独对照 exp(dot)，未加缩放。 keys = np.array([[1., 0.], [0., 1.], [1., 0.]]) values = np.array([[10.], [20.], [30.]]) _, state, normalizer = linear_recurrent(keys, keys, values) queries = np.eye(2) averages = (queries @ state) / (queries @ normalizer)[:, None] delta = delta_memory(keys, values, [1., 1., 1.], [1., 1., 1.]) decayed = delta_memory(keys, values, [1., 1., .5], [1., 1., 1.]) mixed_keys = np.array([[1., 0.], [2 ** -.5, 2 ** -.5]]) collision = delta_memory(mixed_keys, np.array([[10.], [30.]]), [1., 1.], [1., 1.]) return { "kind": "untrained_synthetic_mechanism_checks", "seed": 7, "moe": { "inputs": x.tolist(), "router": router.tolist(), "first": first.tolist(), "second": second.tolist(), "ids": ids.tolist(), "weights": weights.tolist(), "output": output.tolist(), "requested": requested.tolist(), "accepted": accepted.tolist(), "calls": int(accepted.sum()), "dense_calls": int(len(x) * len(first)), "total_parameters": int(first.size + second.size + router.size), "parameters_used_per_token": int(2 * first[0].size + 2 * second[0].size + router.size), "reference_max_error": float(np.max(np.abs(output - reference))), "capacity": 3, "limited_load": limited_load.tolist(), "dropped_assignments": limited_drops, "limited_output": limited.tolist(), }, "linear": { "q": q.tolist(), "k": k.tolist(), "v": v.tolist(), "recurrent": recurrent.tolist(), "explicit": explicit.tolist(), "max_error": float(np.max(np.abs(recurrent - explicit))), "softmax_max_difference": float(np.max(np.abs(recurrent - conventional))), }, "memory": { "keys": keys.tolist(), "values": values.tolist(), "average_read": averages.ravel().tolist(), "delta_read": (queries @ delta[-1]).ravel().tolist(), "decayed_read": (queries @ decayed[-1]).ravel().tolist(), "collision_read": float(collision[-1, 0, 0]), "impulse_traces": {str(a): scalar_state([1.] + [0.] * 50, a) for a in [.8, .99]}, "cache_counts": [{"tokens": n, "kv_scalars": n * 3, "normalized_linear_scalars": int(state.size + normalizer.size)} for n in [32, 128]], }, } def main(): parser = ArgumentParser(description=__doc__) parser.add_argument("--out", type=Path, help="将完整机制算例保存为新 JSON 文件") args = parser.parse_args() report = run() report["environment"] = {"python": platform.python_version(), "numpy": np.__version__, "code_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()} if args.out: args.out.parent.mkdir(parents=True, exist_ok=True) with args.out.open("x", encoding="utf-8") as file: json.dump(report, file, ensure_ascii=False, indent=2, allow_nan=False) file.write("\n") moe, linear, memory = report["moe"], report["linear"], report["memory"] print(f'专家请求数：{moe["requested"]}；实际调用 {moe["calls"]} 次，全算需 {moe["dense_calls"]} 次。') print(f'总参数 {moe["total_parameters"]}；每个词元使用 {moe["parameters_used_per_token"]} 个；稀疏与遮罩路径一致：{moe["reference_max_error"] < 1e-12}') print(f'容量为 3 时接收：{moe["limited_load"]}；丢弃 {moe["dropped_assignments"]} 次分配。') print(f'线性注意力两种计算一致：{linear["max_error"] < 1e-12}；与 softmax 对照最大差：{linear["softmax_max_difference"]:.6f}') print(f'读取 A/B：平均 {memory["average_read"]}；delta {memory["delta_read"]}；末步衰减 {memory["decayed_read"]}') print(f'非正交键更新后再读 A：{memory["collision_read"]:.6f}（原值 10）。') print(f'50 步后旧脉冲：0.8 门 {memory["impulse_traces"]["0.8"][-1]:.6f}；0.99 门 {memory["impulse_traces"]["0.99"][-1]:.6f}') print("范围：固定输入与未训练参数的机制检查，没有测量语言能力或端到端速度。") if __name__ == "__main__": main()