Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据
- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线 - 修复 prepare_memory_router_dataset 候选池重建缺陷（mega 家族 3568x 加速，输出逐字节相同） - 修复 v5 被破坏的拒答与多跳标签（train 未知样本 319 -> 16319，multi_hop 平均正例 1.00 -> 2.00） - 同存储预算下 V2-128 v6 逐轴 22/22 通过：Top-1 41.12% -> 94.62%，未知拒答 0.00% -> 100.00% - 记录三条被实测推翻的显然优化（logits_to_keep=1 反而慢 55%、XL 容量未带来收益） - 记忆手术跨架构可移植性 14/14，读写关闭时与原生模型逐位相同
This commit is contained in:
WpyQwq
@@ -0,0 +1,175 @@
|
||||
"""Train the attribute-coverage head: "which attribute is this question asking about?"
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
The runtime answers ~75% of unanswerable paraphrased questions with an invented value.
|
||||
A score threshold cannot fix that: measured on the same frozen keys, the best fitted head
|
||||
over nine score-geometry features reaches AUC 0.61 (and ~0.5 for three of four scorers),
|
||||
because 24 same-shape candidates look equally plausible for any question.
|
||||
|
||||
What works instead is a different question: *which attribute is being asked about, and does
|
||||
the bank actually hold it?* This script trains that head and saves it as a deployable
|
||||
artifact.
|
||||
|
||||
Labels are recovered exactly, not guessed: every query in this dataset is one of the
|
||||
authored paraphrases, so the asked-about attribute is looked up from the paraphrase table
|
||||
rather than taken from the episode metadata (which records the answerable *target* and is
|
||||
therefore wrong for abstention episodes).
|
||||
|
||||
Usage::
|
||||
|
||||
python -m V2_dpskw.train_attribute_head --data-dir data/zero_overlap ^
|
||||
--feature-cache H:\\Memory\\nm_cache\\nm_zero_overlap\\feature_cache ^
|
||||
--output-dir checkpoints/memory_attribute_head
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .make_zero_overlap_paraphrase_data import ATTRIBUTE_PARAPHRASES
|
||||
|
||||
|
||||
def text_key(text: str) -> str:
|
||||
return hashlib.sha1(str(text).encode("utf-8", "replace")).hexdigest()
|
||||
|
||||
|
||||
def paraphrase_table() -> dict[str, str]:
|
||||
table = {}
|
||||
for attribute, paraphrases in ATTRIBUTE_PARAPHRASES:
|
||||
for paraphrase in paraphrases:
|
||||
table[paraphrase] = attribute
|
||||
return table
|
||||
|
||||
|
||||
def load(path: Path) -> list[dict]:
|
||||
rows = []
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
for line in handle:
|
||||
line = line.strip()
|
||||
if line:
|
||||
rows.append(json.loads(line))
|
||||
return rows
|
||||
|
||||
|
||||
def softmax_fit(X: np.ndarray, y: np.ndarray, classes: int, *, steps: int = 6000,
|
||||
lr: float = 0.5, l2: float = 1e-3):
|
||||
W = np.zeros((X.shape[1], classes))
|
||||
b = np.zeros(classes)
|
||||
n = len(y)
|
||||
onehot = np.zeros((n, classes))
|
||||
onehot[np.arange(n), y] = 1.0
|
||||
for _ in range(steps):
|
||||
logits = X @ W + b
|
||||
logits -= logits.max(axis=1, keepdims=True)
|
||||
probs = np.exp(logits)
|
||||
probs /= probs.sum(axis=1, keepdims=True)
|
||||
grad = probs - onehot
|
||||
W -= lr * (X.T @ grad / n + l2 * W)
|
||||
b -= lr * grad.mean(axis=0)
|
||||
return W, b
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--data-dir", default="data/zero_overlap")
|
||||
parser.add_argument("--feature-cache", default=r"H:\Memory\nm_cache\nm_zero_overlap\feature_cache")
|
||||
parser.add_argument("--output-dir", default="checkpoints/memory_attribute_head")
|
||||
parser.add_argument("--report", default="attribute_head_training.json")
|
||||
args = parser.parse_args()
|
||||
|
||||
cache = Path(args.feature_cache)
|
||||
bank = np.load(cache / "features.f16.npy", mmap_mode="r")
|
||||
lookup = json.loads((cache / "index.json").read_text(encoding="utf-8"))
|
||||
table = paraphrase_table()
|
||||
attributes = sorted({name for name, _ in ATTRIBUTE_PARAPHRASES})
|
||||
index_of = {name: position for position, name in enumerate(attributes)}
|
||||
|
||||
def build(rows):
|
||||
X, y, rows_used = [], [], []
|
||||
for row in rows:
|
||||
attribute = table.get(row["query"])
|
||||
if attribute is None:
|
||||
raise SystemExit(f"query is not one of the authored paraphrases: {row['query']!r}")
|
||||
key = text_key(row["query"])
|
||||
if key not in lookup:
|
||||
raise SystemExit(f"query missing from the feature bank: {row['query']!r}")
|
||||
X.append(np.asarray(bank[lookup[key]], dtype=np.float32))
|
||||
y.append(index_of[attribute])
|
||||
rows_used.append(row)
|
||||
return np.array(X), np.array(y), rows_used
|
||||
|
||||
train_X, train_y, _ = build(load(Path(args.data_dir) / "train.jsonl"))
|
||||
eval_X, eval_y, eval_rows = build(load(Path(args.data_dir) / "eval.jsonl"))
|
||||
|
||||
train_Xn = train_X / (np.linalg.norm(train_X, axis=1, keepdims=True) + 1e-6)
|
||||
eval_Xn = eval_X / (np.linalg.norm(eval_X, axis=1, keepdims=True) + 1e-6)
|
||||
W, b = softmax_fit(train_Xn, train_y, len(attributes))
|
||||
|
||||
def probabilities(X):
|
||||
logits = X @ W + b
|
||||
logits -= logits.max(axis=1, keepdims=True)
|
||||
p = np.exp(logits)
|
||||
return p / p.sum(axis=1, keepdims=True)
|
||||
|
||||
eval_probs = probabilities(eval_Xn)
|
||||
eval_pred = eval_probs.argmax(axis=1)
|
||||
answerable = np.array([bool(row.get("positive_indices")) for row in eval_rows])
|
||||
|
||||
# Coverage = the asked-about attribute is actually present in this episode's bank.
|
||||
predicted_names = [attributes[position] for position in eval_pred]
|
||||
present_sets = [{c.get("attribute") for c in row["candidates"] if c.get("attribute")}
|
||||
for row in eval_rows]
|
||||
covered = np.array([name in present for name, present in zip(predicted_names, present_sets)])
|
||||
|
||||
report = {
|
||||
"built_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
|
||||
"attributes": attributes,
|
||||
"classes": len(attributes),
|
||||
"chance_pct": round(100.0 / len(attributes), 2),
|
||||
"train_queries": int(train_y.size),
|
||||
"eval_queries": int(eval_y.size),
|
||||
"attribute_accuracy_train_pct": round(100 * float(np.mean(
|
||||
probabilities(train_Xn).argmax(axis=1) == train_y)), 2),
|
||||
"attribute_accuracy_eval_pct": round(100 * float(np.mean(eval_pred == eval_y)), 2),
|
||||
"attribute_accuracy_eval_answerable_pct": round(100 * float(np.mean(
|
||||
(eval_pred == eval_y)[answerable])), 2),
|
||||
"coverage_rule": {
|
||||
"answerable_covered_pct": round(100 * float(np.mean(covered[answerable])), 2),
|
||||
"abstention_covered_pct": round(100 * float(np.mean(covered[~answerable])), 2),
|
||||
"unknown_refusal_pct": round(100 * float(np.mean(~covered[~answerable])), 2),
|
||||
"known_false_refusal_pct": round(100 * float(np.mean(~covered[answerable])), 2),
|
||||
"counts": {
|
||||
"tp": int(np.sum(~covered & ~answerable)),
|
||||
"fn": int(np.sum(covered & ~answerable)),
|
||||
"fp": int(np.sum(~covered & answerable)),
|
||||
"tn": int(np.sum(covered & answerable)),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
out_dir = Path(args.output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
torch.save({
|
||||
"weight": torch.tensor(W, dtype=torch.float32),
|
||||
"bias": torch.tensor(b, dtype=torch.float32),
|
||||
"attributes": attributes,
|
||||
"format_version": 1,
|
||||
"input": "l2_normalised frozen query key (hidden_size)",
|
||||
}, out_dir / "attribute_head.pt")
|
||||
(out_dir / "attribute_head_meta.json").write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
Path(args.report).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(json.dumps(report, ensure_ascii=False, indent=2), flush=True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user