transformerless_lm: quantization fixes — reciprocal tiers + per-row scale

claude · claude · commit 1f08c22b3793 · 2026-05-20T23:37:10.000Z
The v1 Fibonacci-tier quantizer saturated at +1 nat loss regardless of
bit depth because:
  - tier values {1, 2, 3, 5, 8, 13, ...} have no resolution between 0
    and 1, but Gaussian weight distributions are concentrated near 0
  - global per-tensor scale wastes the tail tier values (89, 144,
    233, ...) on a single max-magnitude weight

Two fixes implemented:

1. reciprocal Fibonacci tier values (reciprocals=True)
   {0, +-1/F_max, ..., +-1/3, +-1/2, +-1, +-2, +-3, +-5, ...}
   geometric spacing crossing zero with ratio approaching phi between
   adjacent values, giving fine resolution near 0 where weights live.

2. per-row scale (scale="per_row")
   each output row of W gets its own scale factor matched to its
   own magnitude range, instead of one global scale for the entire
   tensor. Standard per-channel quantization trick, applied to the
   Fibonacci-tier setting.

train_weight_substrate.py now sweeps the cross product of:
  n_tiers in {4, 8, 16, 32} x reciprocals in {False, True} x scale in
  {"per_tensor", "per_row"}
= 16 quantizer configurations per arch, plus the 2 fp32 baselines.

The bench tells us which of the four combinations (rec * scale)
actually clears the &lt;=0.1 nat validation threshold, and at what tier
depth, separately for dense_crt and tied_substrate.

User's framing: faster inference for larger models. Each successful
fix is a step toward that. The dense_crt validating at &lt;=4 bits would
already mean 4-8x parameter-storage compression vs fp32 / fp16.
diff --git a/experiments/transformerless_lm/models_substrate.py b/experiments/transformerless_lm/models_substrate.py
@@ -52,71 +52,91 @@
 FIB_POS_UNIQUE = sorted(set(f for f in FIBONACCI if f > 0))
 
 
-def fibonacci_tier_values(n_tiers: int) -> list[float]:
-    """First n_tiers signed Fibonacci tier values, smallest magnitude first.
-
-    Tier 0 = 0 (the "fold-to-zero" attractor).
-    Tier 1 = ±1
-    Tier 2 = ±2
-    Tier 3 = ±3
-    Tier k = ±F(k) for F(k) the k-th unique positive Fibonacci number.
-
-    Returns sorted list [-F_max, ..., -1, 0, +1, ..., +F_max] suitable
-    for nearest-tier-value snapping.
+def fibonacci_tier_values(n_tiers: int, reciprocals: bool = False) -> list[float]:
+    """Signed Fibonacci tier values.
+
+    Without reciprocals (the original v1):
+        {0, ±1, ±2, ±3, ±5, ±8, ±13, ±21, ...}
+    log spacing toward infinity, no resolution between 0 and 1.
+
+    With reciprocals (v2 — fixes the "no resolution near zero" failure
+    from the v1 bench):
+        {0, ±1/F_max, ..., ±1/5, ±1/3, ±1/2, ±1, ±2, ±3, ±5, ±8, ..., ±F_max}
+    log spacing crossing zero — fine resolution near 0 where most
+    Gaussian-distributed weights actually live.
+
+    Adjacent ratios approach φ (since F(k+1)/F(k) → φ), so this is
+    the natural phi-Fibonacci tier system the substrate already uses
+    elsewhere in OMC.
     """
     fibs = FIB_POS_UNIQUE[: max(0, n_tiers - 1)]
-    return sorted([-f for f in fibs] + [0.0] + [float(f) for f in fibs])
+    pos = [float(f) for f in fibs]
+    if reciprocals:
+        pos = sorted(set(pos + [1.0 / f for f in fibs if f > 1]))
+    return sorted([-v for v in pos] + [0.0] + pos)
 
 
 def fibonacci_tier_snap(W: torch.Tensor, n_tiers: int = 8,
-                         scale: str = "per_tensor") -> tuple[torch.Tensor, int]:
+                         scale: str = "per_tensor",
+                         reciprocals: bool = False) -> tuple[torch.Tensor, int]:
     """Snap each weight in W to its nearest signed-Fibonacci tier value.
 
-    The user's Principle B: every weight folds to a Fibonacci tier;
-    high-frequency / large-magnitude weights occupy the small tiers
-    (1, 2, 3); rare ones drift to higher tiers (5, 8, 13, ...) and
-    very rare ones fold to 0 (pruned).
-
     Args:
-        W: tensor of arbitrary shape.
-        n_tiers: number of tier levels per sign (including 0).
-                  e.g. n_tiers=4 → values in {-3, -2, -1, 0, 1, 2, 3}
-                  = 7 levels = log2(7) ≈ 2.8 bits per weight.
-        scale: "per_tensor" → one scale factor over the full tensor.
+        W: tensor (1-D or 2-D).
+        n_tiers: resolution per sign (= number of distinct positive Fibonacci
+                  values, before any reciprocals).
+        scale: "per_tensor" → one global scale set by max(|W|).
+               "per_row"   → one scale per output row of a 2-D matrix
+                              (matches each row's own dynamic range; the
+                              standard per-channel quantization trick).
+        reciprocals: if True, include 1/F(k) values in the tier set —
+                      gives fine resolution near 0.
 
     Returns:
-        (W_quantized, n_unique_values_actually_used)
+        (W_quantized, n_unique_values_actually_used_avg)
     """
     tier_vals = torch.tensor(
-        fibonacci_tier_values(n_tiers), dtype=W.dtype, device=W.device
-    )                                                      # [n_levels]
-    abs_max = W.abs().max().item()
-    if abs_max == 0:
-        return W.clone(), 1
-    # Choose scale so the largest |w| maps to the largest tier value.
-    s = abs_max / max(tier_vals.abs().max().item(), 1.0)
-    # Nearest tier (in scaled units).
-    target_vals = tier_vals * s                            # [n_levels]
-    # For each element of W, find the closest target value.
-    diffs = (W.unsqueeze(-1) - target_vals).abs()           # [..., n_levels]
-    nearest = diffs.argmin(dim=-1)                          # [...]
-    W_q = target_vals[nearest]
-    n_unique = nearest.unique().numel()
-    return W_q, n_unique
+        fibonacci_tier_values(n_tiers, reciprocals=reciprocals),
+        dtype=W.dtype, device=W.device,
+    )                                                       # [n_levels]
+    max_tier = max(tier_vals.abs().max().item(), 1.0)
+
+    if scale == "per_tensor":
+        abs_max = W.abs().max().item()
+        if abs_max == 0:
+            return W.clone(), 1
+        s = abs_max / max_tier
+        target_vals = tier_vals * s
+        diffs = (W.unsqueeze(-1) - target_vals).abs()
+        nearest = diffs.argmin(dim=-1)
+        W_q = target_vals[nearest]
+        n_unique = nearest.unique().numel()
+        return W_q, n_unique
+
+    if scale == "per_row":
+        if W.dim() != 2:
+            # Fall back to per-tensor for 1-D / N-D parameters.
+            return fibonacci_tier_snap(W, n_tiers, "per_tensor", reciprocals)
+        abs_max_row = W.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-12)  # [out, 1]
+        s_row = abs_max_row / max_tier                       # [out, 1]
+        # For each row, scaled tier set is tier_vals * s_row. We need
+        # per-row argmin over [out, in, n_levels].
+        targets = tier_vals.view(1, 1, -1) * s_row.unsqueeze(-1)  # [out, 1, n_levels]
+        diffs = (W.unsqueeze(-1) - targets).abs()             # [out, in, n_levels]
+        nearest = diffs.argmin(dim=-1)                        # [out, in]
+        W_q = torch.gather(targets.expand_as(diffs), -1,
+                            nearest.unsqueeze(-1)).squeeze(-1)
+        n_unique = nearest.unique().numel()
+        return W_q, n_unique
+
+    raise ValueError(scale)
 
 
 def fibonacci_quantize_model(model: torch.nn.Module, n_tiers: int = 8,
+                              scale: str = "per_tensor",
+                              reciprocals: bool = False,
                               targets: list[str] = None) -> dict:
-    """In-place Fibonacci-tier-snap of model parameters matching `targets`.
-
-    Args:
-        model: any nn.Module.
-        n_tiers: tier resolution per sign.
-        targets: substrings; a param is quantized iff any target is in
-                  its name. Default = quantize ALL params.
-
-    Returns: dict with summary stats (params_quantized, n_unique_per_tensor).
-    """
+    """In-place Fibonacci-tier-snap of model parameters matching `targets`."""
     if targets is None:
         targets = [""]
     stats = {"params_quantized": 0, "tensors_quantized": 0,
@@ -125,7 +145,9 @@ def fibonacci_quantize_model(model: torch.nn.Module, n_tiers: int = 8,
         if not any(t in name for t in targets):
             continue
         with torch.no_grad():
-            W_q, n_unique = fibonacci_tier_snap(p.data, n_tiers=n_tiers)
+            W_q, n_unique = fibonacci_tier_snap(
+                p.data, n_tiers=n_tiers, scale=scale, reciprocals=reciprocals,
+            )
             p.data.copy_(W_q)
             stats["params_quantized"] += p.numel()
             stats["tensors_quantized"] += 1
diff --git a/experiments/transformerless_lm/train_weight_substrate.py b/experiments/transformerless_lm/train_weight_substrate.py
@@ -156,73 +156,94 @@ def main():
         print(f"\n  ✓ [arch={arch}] fp32 final_val = {final_val:.4f}")
 
         # ---- Principle B: post-hoc Fibonacci-tier quantization sweep ----
-        # Keep the trained weights in `state_dict_orig` so we can re-quantize
-        # at each n_tiers from a clean baseline.
         state_dict_orig = {k: v.clone() for k, v in model.state_dict().items()}
-        for n_tiers in tier_sweep:
+        configs = []
+        for reciprocals in [False, True]:
+            for scale in ["per_tensor", "per_row"]:
+                for n_tiers in tier_sweep:
+                    configs.append((n_tiers, reciprocals, scale))
+        for n_tiers, reciprocals, scale in configs:
             model.load_state_dict(state_dict_orig)
-            stats = fibonacci_quantize_model(model, n_tiers=n_tiers)
+            stats = fibonacci_quantize_model(
+                model, n_tiers=n_tiers,
+                reciprocals=reciprocals, scale=scale,
+            )
             eval_gen.manual_seed(args.seed + 1000)
             vq = evaluate(model, val_split, args.batch_size, args.seq_len,
                           n_batches=32, generator=eval_gen)
             n_unique_total = sum(s["n_unique_tier_values"]
                                   for s in stats["per_tensor"].values())
             n_tensors = stats["tensors_quantized"]
             avg_unique = n_unique_total / max(n_tensors, 1)
-            print(f"    n_tiers={n_tiers:>3} → val={vq:.4f}  Δ={vq - final_val:+.4f}  "
-                  f"avg_unique_tier_vals={avg_unique:.1f}", flush=True)
-            arch_record["quantized"][n_tiers] = {
+            key = f"n{n_tiers}_{'rec' if reciprocals else 'nor'}_{scale}"
+            print(f"    {key:<22} → val={vq:.4f}  Δ={vq - final_val:+.4f}  "
+                  f"avg_unique={avg_unique:.1f}", flush=True)
+            arch_record["quantized"][key] = {
+                "n_tiers": n_tiers,
+                "reciprocals": reciprocals,
+                "scale": scale,
                 "val": vq,
                 "delta": vq - final_val,
                 "params_quantized": stats["params_quantized"],
                 "avg_unique_tier_values": avg_unique,
             }
-        # Restore fp32 weights at the end so subsequent code can use the model.
         model.load_state_dict(state_dict_orig)
         results["archs"][arch] = arch_record
 
-    # ---- Summary table ----
+    # ---- Summary tables ----
     print()
-    print("=" * 96)
-    schemes = ["fp32"] + [f"q{n}" for n in tier_sweep]
-    header = f"{'arch':<18} {'attn_params':>12} {'total':>10} "
-    header += " ".join(f"{s:>10}" for s in schemes)
-    print(header)
-    print("-" * 96)
+    print("=" * 110)
+    print("FP32 BASELINES")
+    print("-" * 110)
+    print(f"{'arch':<18} {'attn_params':>12} {'total':>10} {'fp32_val':>10}")
     for arch in ["dense_crt", "tied_substrate"]:
         r = results["archs"][arch]
-        vals = [r["val_fp32"]]
-        for n_tiers in tier_sweep:
-            vals.append(r["quantized"][n_tiers]["val"])
-        cells = " ".join(f"{v:>10.4f}" for v in vals)
-        print(f"{arch:<18} {r['n_attn_params']:>12,} {r['n_params']:>10,} {cells}")
+        print(f"{arch:<18} {r['n_attn_params']:>12,} {r['n_params']:>10,} "
+              f"{r['val_fp32']:>10.4f}")
+
+    print()
+    print("=" * 110)
+    print("QUANTIZATION SWEEP — Δ vs fp32 for each arch")
+    print("(rec=with reciprocal Fibonacci tiers; nor=Fibonacci only)")
+    print("-" * 110)
+    for arch in ["dense_crt", "tied_substrate"]:
+        r = results["archs"][arch]
+        print(f"\n  {arch}  (fp32 val = {r['val_fp32']:.4f}):")
+        print(f"    {'n_tiers':>8} {'tier_set':>10} {'scale':>12} "
+              f"{'val':>10} {'Δ':>10} {'unique':>10}")
+        for key, q in r["quantized"].items():
+            print(f"    {q['n_tiers']:>8} {('rec' if q['reciprocals'] else 'nor'):>10} "
+                  f"{q['scale']:>12} {q['val']:>10.4f} {q['delta']:>+10.4f} "
+                  f"{q['avg_unique_tier_values']:>10.1f}")
 
     # ---- Interpretation ----
     print()
-    print("INTERPRETATION:")
+    print("=" * 110)
+    print("INTERPRETATION")
+    print("-" * 110)
     a_fp32 = results["archs"]["dense_crt"]["val_fp32"]
     t_fp32 = results["archs"]["tied_substrate"]["val_fp32"]
     a_attn = results["archs"]["dense_crt"]["n_attn_params"]
     t_attn = results["archs"]["tied_substrate"]["n_attn_params"]
     rel = (t_fp32 - a_fp32) / a_fp32 * 100
-    print(f"  Principle A (tied substrate vs dense_crt):")
-    print(f"    val_fp32 delta: {t_fp32 - a_fp32:+.4f} ({rel:+.1f}%)")
-    print(f"    attn param reduction: {a_attn / t_attn:.2f}x")
-    if abs(rel) < 5:
-        print(f"    → PRINCIPLE A: tied QKV trains within 5% of dense — VALIDATED")
-    else:
-        print(f"    → PRINCIPLE A: tied QKV {('lost' if rel > 0 else 'gained')} {abs(rel):.1f}% "
-              f"— {'NOT VALIDATED (too tight)' if rel > 5 else 'BEAT BASELINE'}")
-
-    print(f"\n  Principle B (Fibonacci-tier quantization on dense_crt):")
-    for n_tiers in tier_sweep:
-        delta = results["archs"]["dense_crt"]["quantized"][n_tiers]["delta"]
-        # log2(2*n_tiers - 1) bits per weight (n_tiers includes 0, then mirror)
-        # for n_tiers=8: levels = 2*8-1 = 15 → ~4 bits.
-        levels = 2 * n_tiers - 1
-        bits = (levels - 1).bit_length()
-        verdict = "VALIDATED" if delta < 0.1 else ("DEGRADED" if delta < 0.5 else "BROKEN")
-        print(f"    n_tiers={n_tiers:>3} (~{bits} bits): val delta {delta:+.4f}  → {verdict}")
+    print(f"\nPRINCIPLE A (tied substrate vs dense_crt):")
+    print(f"  val_fp32 delta: {t_fp32 - a_fp32:+.4f} ({rel:+.1f}%)")
+    print(f"  attn param reduction: {a_attn / t_attn:.2f}x")
+    verdict_a = ("VALIDATED" if abs(rel) < 5 else
+                  ("BEAT BASELINE" if rel < 0 else "NOT VALIDATED"))
+    print(f"  → PRINCIPLE A: {verdict_a}")
+
+    print(f"\nPRINCIPLE B (best quantizer per arch, vs ≤0.10 nat threshold):")
+    for arch in ["dense_crt", "tied_substrate"]:
+        r = results["archs"][arch]
+        best_key = min(r["quantized"], key=lambda k: r["quantized"][k]["delta"])
+        best = r["quantized"][best_key]
+        verdict_b = "VALIDATED" if best["delta"] < 0.10 else (
+            "USABLE" if best["delta"] < 0.30 else "BROKEN")
+        print(f"  {arch}:")
+        print(f"    best = {best_key}  (Δ={best['delta']:+.4f}, "
+              f"n_tiers={best['n_tiers']}, rec={best['reciprocals']}, "
+              f"scale={best['scale']})  → {verdict_b}")
 
     out_path = Path(__file__).parent / args.out
     with open(out_path, "w") as f:
