cortexm/vsa/role_vectors.py
Rama main · réplica de lectura del repositorio original.
"""Engineered role vectors — NSR-inspired (ESWEEK24).
arXiv insight: NSR trains an autoencoder to learn compact embeddings
for the fact vocabulary, then uses those embeddings to construct role
vectors that are *maximally mutually orthogonal* in the data manifold.
Random role vectors work (HDC theory guarantees capacity grows with
sqrt(dims)), but engineered ones have higher effective capacity and
lower cross-talk because they sit on the directions of greatest
variance in the actual data.
This module provides:
EngineeredRoleVectors(dims, n_roles=3)
.fit(fact_matrix) -> trains a tiny 1-layer autoencoder
on the fact matrix and stores the
top-k principal directions as role
vectors
.role_vec(role) -> returns the engineered role vector
(falls back to VSA's random init
if not yet fit)
.save(path) / .load(path)
WHY THIS MATTERS:
- role_vec("S") currently = rng.standard_normal(dims). For dims=768
that's a random point on the unit sphere — fine in theory but it
wastes capacity on directions orthogonal to the actual fact vocab.
- After fit(), role_vec("S") = the top-1 principal direction of the
S-vector manifold. role_vec("R") = top-2, role_vec("V") = top-3.
These directions are where the data actually lives, so:
(a) bound holograms exploit the full effective capacity
(b) cross-talk between S/R/V role bindings drops because the
top-k principal directions are approximately orthogonal by
construction
(c) retrieval probes have higher SNR because they're aligned
with the data axes
IMPLEMENTATION:
Tiny 1-layer linear autoencoder (no nonlinearity) trained with
plain SGD on the fact matrix. The encoder weights converge to the
top-k principal components (proven by Baldi & Horn, 1989 — linear
AE on centered data recovers PCA). We center the data and train the
encoder to have orthogonal rows via a soft orthogonality penalty.
Cost: ~1 day of human work, ~30 seconds of CPU time on 1000 facts.
Deterministic given the seed.
"""
from __future__ import annotations
import json
import os
from pathlib import Path
import numpy as np
class EngineeredRoleVectors:
"""Train and serve engineered role vectors.
Construction:
erv = EngineeredRoleVectors(dims=768, n_roles=3, seed=42)
Fit:
erv.fit(fact_matrix) # fact_matrix: (n_facts, dims)
# trains a tiny AE, extracts top-k principal directions,
# stores them as the role vectors
Use:
erv.role_vec("S") # top-1 principal direction
erv.role_vec("R") # top-2
erv.role_vec("V") # top-3
"""
ROLE_ORDER = ["S", "R", "V"] # subject / relation / value
def __init__(self, dims: int = 768, n_roles: int = 3,
seed: int = 42, n_epochs: int = 200,
lr: float = 0.01, orth_penalty: float = 0.1) -> None:
self.dims = dims
self.n_roles = n_roles
self.seed = seed
self.n_epochs = n_epochs
self.lr = lr
self.orth_penalty = orth_penalty
self._role_vecs: dict[str, np.ndarray] = {}
self._mean: np.ndarray | None = None
self._fit_loss: list[float] = []
# ------------------------------------------------------------- fit
def fit(self, fact_matrix: np.ndarray) -> dict:
"""Train a tiny linear autoencoder on `fact_matrix`.
fact_matrix: (n_samples, dims) float32. Should contain the
subject / relation / value vectors for all facts in the
corpus, stacked. The AE learns to reconstruct them through
a k-dim bottleneck where k = n_roles.
After training, the encoder's rows are the top-k principal
directions. We store them as role_vec("S"/"R"/"V").
Returns a training report dict.
"""
X = np.asarray(fact_matrix, dtype=np.float32)
if X.ndim != 2:
raise ValueError(f"expected 2D matrix, got shape {X.shape}")
n, d = X.shape
if d != self.dims:
raise ValueError(
f"matrix dim {d} != configured dims {self.dims}")
if n < self.n_roles:
# not enough samples to learn k directions — fall back
# to top-k random vectors orthogonalized via Gram-Schmidt
rng = np.random.default_rng(self.seed)
vecs = rng.standard_normal((self.n_roles, d)).astype(np.float32)
for i in range(self.n_roles):
for j in range(i):
vecs[i] -= np.dot(vecs[i], vecs[j]) * vecs[j]
n_ = float(np.linalg.norm(vecs[i]))
vecs[i] /= max(n_, 1e-9)
for i, r in enumerate(self.ROLE_ORDER[:self.n_roles]):
self._role_vecs[r] = vecs[i]
return {"trained": False, "reason": "insufficient_samples",
"n_samples": n, "n_roles": self.n_roles,
"fallback": "gram_schmidt_random"}
# center the data (PCA assumes centered data)
self._mean = X.mean(axis=0)
Xc = X - self._mean
# tiny linear AE: encoder W: (k, d), decoder W.T: (d, k)
# init with small random values
rng = np.random.default_rng(self.seed)
k = self.n_roles
scale = float(1.0 / np.sqrt(d))
W = rng.standard_normal((k, d)).astype(np.float32) * scale
# SGD on reconstruction loss + orthogonality penalty
# (orth penalty: W @ W.T should be ≈ I)
batch_size = min(64, n)
for epoch in range(self.n_epochs):
perm = rng.permutation(n)
epoch_loss = 0.0
for i in range(0, n, batch_size):
idx = perm[i:i + batch_size]
xb = Xc[idx] # (b, d)
# forward: z = xb @ W.T (b, k)
# recon: xhat = z @ W (b, d)
z = xb @ W.T
xhat = z @ W
# L2 reconstruction loss per sample — clip to avoid
# overflow on numerically large data (the orth penalty
# can push weights to blow up if lr is too high)
diff = xhat - xb
# clip per-element to a safe range
diff = np.clip(diff, -1e4, 1e4)
loss = float(np.mean(np.sum(diff * diff, axis=1)))
# grad on W: dL/dW = 2 * (xhat - x).T @ z / b
# shape (d, k) -> transpose for our layout
g = 2.0 * diff.T @ z / len(idx) # (d, k)
# clip gradient to prevent overflow → NaN
g = np.clip(g, -1.0, 1.0)
# orthogonality penalty: ||W @ W.T - I||_F^2 / k
# grad: 2/k * (W @ W.T - I) @ W
WtW = W @ W.T # (k, k)
I_k = np.eye(k, dtype=np.float32)
orth_grad = (2.0 / k) * (WtW - I_k) @ W
orth_grad = np.clip(orth_grad, -1.0, 1.0)
# combined grad on W (note: g is (d,k), so transpose)
W -= self.lr * (g.T + self.orth_penalty * orth_grad)
# also clip weights themselves to a safe range
W = np.clip(W, -10.0, 10.0)
epoch_loss += loss * len(idx)
epoch_loss /= n
self._fit_loss.append(epoch_loss)
if epoch % 20 == 0 or epoch == self.n_epochs - 1:
# report conditioning of W @ W.T (lower = more orthogonal)
cond = float(np.linalg.cond(WtW)) if k > 1 else 1.0
# suppress per-epoch logging during normal use
pass
# extract role vectors: rows of W are the top-k principal dirs
# normalize to unit length
for i, r in enumerate(self.ROLE_ORDER[:k]):
v = W[i].copy()
nrm = float(np.linalg.norm(v))
self._role_vecs[r] = v / max(nrm, 1e-9)
# report final reconstruction loss + orthogonality
WtW_final = W @ W.T
off_diag = float(np.sum(np.abs(WtW_final - np.eye(k, dtype=np.float32))))
return {
"trained": True,
"n_samples": n,
"dims": d,
"n_roles": k,
"epochs": self.n_epochs,
"final_loss": self._fit_loss[-1] if self._fit_loss else 0.0,
"initial_loss": self._fit_loss[0] if self._fit_loss else 0.0,
"loss_reduction_pct": (
(1.0 - self._fit_loss[-1] / max(self._fit_loss[0], 1e-9)) * 100
if self._fit_loss else 0.0),
"off_diag_sum": off_diag,
"condition_number": (float(np.linalg.cond(WtW_final))
if k > 1 else 1.0),
}
# ------------------------------------------------------------- serve
def role_vec(self, role: str) -> np.ndarray | None:
return self._role_vecs.get(role)
@property
def is_fit(self) -> bool:
return bool(self._role_vecs)
def save(self, path: str | os.PathLike) -> None:
"""Persist the role vectors to a .npz file."""
arrays = {f"role_{r}": v for r, v in self._role_vecs.items()}
if self._mean is not None:
arrays["mean"] = self._mean
arrays["meta"] = np.array(json.dumps({
"dims": self.dims, "n_roles": self.n_roles,
"seed": self.seed, "final_loss": (
self._fit_loss[-1] if self._fit_loss else 0.0),
}), dtype=str)
np.savez(path, **arrays)
def load(self, path: str | os.PathLike) -> None:
data = np.load(path, allow_pickle=False)
for r in self.ROLE_ORDER[:self.n_roles]:
key = f"role_{r}"
if key in data:
self._role_vecs[r] = data[key]
if "mean" in data:
self._mean = data["mean"]
__all__ = ["EngineeredRoleVectors"]