transformerless_lm: lazy training applied to FibGen seed components

Two substrate-aligned implementations of "lazy loading at training":

  (1) LAZY_DROPOUT (lazy_tier_dropout=True):
      Bernoulli mask on each seed component with keep_prob = 1/sqrt(tier).
      Tier-1 components (smallest Fibonacci index, k_i=k_j=0) are
      always active; tier-k components active 1/sqrt(k) of steps.
      Eval rescales by keep_prob so train E[output] matches eval output.

  (2) TIER_LR_SCALE (apply_tier_lr_scale post-backward):
      Keep all components in the forward, but scale each component's
      gradient by 1/sqrt(tier) before optimizer.step(). Low-tier
      components learn fast; high-tier learn slowly. Deterministic,
      no train/eval mismatch.

train_lazy_subsim.py runs a 3-arm bench at d=128:
  subsim_baseline       - vanilla Subsim (the +5.7% gap to dense)
  subsim_lazy_dropout   - variant (1)
  subsim_tier_lr        - variant (2)

Both lazy variants are designed to make the trained model PRUNABLE
post-hoc: prune high-tier components first, since they carry less
learned signal. This delivers the "35B in 8GB" framing through
deployment-time component pruning.

100-step smoke: lazy_dropout converges but ~2x slower per val unit
(harder optimization with stochastic masking). tier_lr_scale untested
in the smoke; expected to be smoother because it's deterministic.

Bench will reveal which variant best balances training stability with
post-training pruneability.

Author: claude

experiments/transformerless_lm/models_fibgen.py

@@ -87,18 +87,43 @@ class FibGenLinear(nn.Module):
 
     def __init__(self, in_features: int, out_features: int, K: int = 16,
                  mode: str = "separable",
-                 bias: bool = True, init_scale: float = 0.1):
+                 bias: bool = True, init_scale: float = 0.1,
+                 lazy_tier_dropout: bool = False):
         super().__init__()
         self.in_features = in_features
         self.out_features = out_features
         self.K = min(K, len(FIBONACCI))
         if mode not in ("separable", "cross"):
             raise ValueError(f"unknown mode: {mode}")
         self.mode = mode
+        self.lazy_tier_dropout = lazy_tier_dropout
         n_components = self.K if mode == "separable" else self.K * self.K
         self.seed = nn.Parameter(
             torch.randn(n_components, 4) * (init_scale / max(1, math.sqrt(n_components)))
         )
+
+        # Fibonacci tier per seed component, used for lazy-tier dropout.
+        # Lower tier = more important = active more often.
+        if mode == "separable":
+            # Component k → tier (k+1). F(tier) = Fibonacci number.
+            tiers_int = [i + 1 for i in range(self.K)]
+        else:
+            # Cross-mode pair (k_i, k_j) → tier max(k_i, k_j) + 1.
+            # Pair (0, 0) is tier 1 (most important, always active).
+            # Pair (31, 31) is tier 32 (rarely active under 1/F(32) probability).
+            tiers_int = [max(k_i, k_j) + 1
+                         for k_i in range(self.K) for k_j in range(self.K)]
+        # Two substrate-aligned schemes available on this buffer:
+        # (1) lazy_tier_dropout=True   -> mask seed via Bernoulli(tier_keep_probs)
+        # (2) gradient-scale via tier_lr_scale (applied by training loop)
+        keep_probs = torch.tensor(
+            [1.0 / math.sqrt(t) for t in tiers_int], dtype=torch.float,
+        )
+        self.register_buffer("tier_keep_probs", keep_probs)
+        # tier-weighted learning rate: low-tier components get full LR, high-tier
+        # get reduced LR proportional to 1/sqrt(tier). Apply by multiplying
+        # seed.grad by this buffer BEFORE optimizer.step().
+        self.register_buffer("tier_lr_scale", keep_probs.unsqueeze(-1))
         if bias:
             self.bias = nn.Parameter(torch.zeros(out_features))
         else:
@@ -151,6 +176,29 @@ def generate_W(self) -> torch.Tensor:
             return cached
         return self._compute_W()
 
+    def _maybe_lazy_seed(self) -> torch.Tensor:
+        """Returns the seed (optionally masked by Fibonacci-tier dropout).
+
+        Substrate-native lazy LOADING applied to the seed itself:
+          - Tier 1 components are always active (full participation)
+          - Tier-k components active with probability 1/sqrt(k)
+          - Only active components contribute to this step's forward;
+            only they receive gradient on backward.
+
+        Magnitude matching: at training the mask is Bernoulli; at eval
+        we scale the seed by the per-component keep_prob so the
+        EXPECTED forward output during training matches the deterministic
+        forward at eval. This avoids the magnitude crash that pure-mask
+        without scaling caused.
+        """
+        if not self.lazy_tier_dropout:
+            return self.seed
+        if self.training:
+            mask = torch.bernoulli(self.tier_keep_probs)        # [n_components]
+            return self.seed * mask.unsqueeze(-1)
+        # eval: deterministic, scaled by keep_prob to match training E[seed]
+        return self.seed * self.tier_keep_probs.unsqueeze(-1)
+
     def _forward_compressed(self, x: torch.Tensor) -> torch.Tensor:
         """Substrate-native forward: compute y = W·x WITHOUT materializing W.
 
@@ -166,8 +214,9 @@ def _forward_compressed(self, x: torch.Tensor) -> torch.Tensor:
         K-dim projected x, then projected back.
         """
         # x: [B, T, in_features]
+        seed = self._maybe_lazy_seed()
         if self.mode == "separable":
-            a, b, c, d = self.seed[:, 0], self.seed[:, 1], self.seed[:, 2], self.seed[:, 3]
+            a, b, c, d = seed[:, 0], seed[:, 1], seed[:, 2], seed[:, 3]
             # Project x into Fibonacci-basis along input axis: [B, T, K]
             x_cos = x @ self.cos_j                        # [B, T, K]
             x_sin = x @ self.sin_j                        # [B, T, K]
@@ -185,8 +234,8 @@ def _forward_compressed(self, x: torch.Tensor) -> torch.Tensor:
             return y
         # cross mode: seed [K, K, 4] mixing matrix
         K = self.K
-        seed = self.seed.view(K, K, 4)
-        a, b, c, d = seed[..., 0], seed[..., 1], seed[..., 2], seed[..., 3]
+        seed_cross = seed.view(K, K, 4)
+        a, b, c, d = seed_cross[..., 0], seed_cross[..., 1], seed_cross[..., 2], seed_cross[..., 3]
         x_cos = x @ self.cos_j                            # [B, T, K]
         x_sin = x @ self.sin_j
         # K×K mixing in seed space:

experiments/transformerless_lm/models_subsim.py

@@ -49,13 +49,16 @@ class SubstrateSimilarityAttention(nn.Module):
     """
 
     def __init__(self, d_model: int, K: int = 32, seq_len: int = 128,
-                 fibgen_K: int = 32, mode: str = "cross"):
+                 fibgen_K: int = 32, mode: str = "cross",
+                 lazy_tier_dropout: bool = False):
         super().__init__()
         self.d_model = d_model
         self.K = K
-        self.W_sig = FibGenLinear(d_model, K, K=fibgen_K, mode=mode, bias=False)
-        self.W_v = FibGenLinear(d_model, d_model, K=fibgen_K, mode=mode, bias=False)
-        self.W_out = FibGenLinear(d_model, d_model, K=fibgen_K, mode=mode, bias=False)
+        kw = dict(K=fibgen_K, mode=mode, bias=False,
+                   lazy_tier_dropout=lazy_tier_dropout)
+        self.W_sig = FibGenLinear(d_model, K, **kw)
+        self.W_v = FibGenLinear(d_model, d_model, **kw)
+        self.W_out = FibGenLinear(d_model, d_model, **kw)
         # Standard causal mask; substrate-distance attention is dense in
         # principle. Could also use Fibonacci-offset mask for sparsity.
         mask = torch.tril(torch.ones(seq_len, seq_len))
@@ -81,13 +84,16 @@ class SubsimBlock(nn.Module):
     """Substrate-similarity attention + FibGen FFN."""
 
     def __init__(self, d_model: int, seq_len: int, K: int = 32,
-                 fibgen_K: int = 32, mode: str = "cross"):
+                 fibgen_K: int = 32, mode: str = "cross",
+                 lazy_tier_dropout: bool = False):
         super().__init__()
         self.attn = SubstrateSimilarityAttention(d_model, K=K, seq_len=seq_len,
-                                                   fibgen_K=fibgen_K, mode=mode)
+                                                   fibgen_K=fibgen_K, mode=mode,
+                                                   lazy_tier_dropout=lazy_tier_dropout)
         # FFN with FibGen weights (separate K for FFN if desired)
-        self.w1 = FibGenLinear(d_model, 4 * d_model, K=fibgen_K, mode=mode)
-        self.w2 = FibGenLinear(4 * d_model, d_model, K=fibgen_K, mode=mode)
+        kw = dict(K=fibgen_K, mode=mode, lazy_tier_dropout=lazy_tier_dropout)
+        self.w1 = FibGenLinear(d_model, 4 * d_model, **kw)
+        self.w2 = FibGenLinear(4 * d_model, d_model, **kw)
         self.ln1 = nn.LayerNorm(d_model)
         self.ln2 = nn.LayerNorm(d_model)
 
@@ -108,15 +114,16 @@ class SubsimLM(nn.Module):
 
     def __init__(self, vocab_size: int, d_model: int, n_blocks: int,
                  seq_len: int, K: int = 32, fibgen_K: int = 32,
-                 mode: str = "cross"):
+                 mode: str = "cross", lazy_tier_dropout: bool = False):
         super().__init__()
         self.seq_len = seq_len
         self.K = K
         self.embed = nn.Embedding(vocab_size, d_model)
         pe = self._crt_pe(seq_len, d_model)
         self.register_buffer("pe", pe)
         self.blocks = nn.ModuleList([
-            SubsimBlock(d_model, seq_len, K=K, fibgen_K=fibgen_K, mode=mode)
+            SubsimBlock(d_model, seq_len, K=K, fibgen_K=fibgen_K, mode=mode,
+                          lazy_tier_dropout=lazy_tier_dropout)
             for _ in range(n_blocks)
         ])
         self.ln_f = nn.LayerNorm(d_model)

experiments/transformerless_lm/train_lazy_subsim.py

@@ -0,0 +1,176 @@
+"""Lazy training applied to FibGen seed components.
+
+Two substrate-aligned variants tested:
+
+  (1) LAZY_DROPOUT: Bernoulli mask on each FibGen seed component.
+      keep_prob = 1/sqrt(tier) so low-tier (small Fibonacci index)
+      components active near-always, high-tier components active
+      stochastically. Eval rescales by keep_prob to match expected
+      training magnitudes. This is "lazy loading at the seed level":
+      each step uses only a substrate-defined subset of components.
+
+  (2) TIER_LR_SCALE: keep all components active in the forward, but
+      scale each component's GRADIENT by 1/sqrt(tier) before
+      optimizer.step(). Low-tier components learn fast (full LR),
+      high-tier learn slowly. Over training, low-tier components
+      accumulate more signal. Deterministic, no train/eval mismatch.
+
+Both share the substrate intent ("fold to respected tier") but
+differ in implementation. We also include the pure-baseline Subsim
+for direct comparison.
+
+The deployment payoff (orthogonal to which training scheme wins):
+post-training, prune high-tier components and measure perplexity
+loss. The lazy-trained model should prune more gracefully because
+high-tier components were either inactive (variant 1) or had small
+learned magnitudes (variant 2).
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).parent))
+from corpus import make_dataset
+from models import make_model
+from models_subsim import SubsimLM
+from models_fibgen import FibGenLinear
+from train_distractor_mix import build_distractor_stream
+from lazy_data import fib_positions_in_window, get_fib_strided_batch
+
+
+def evaluate(model, val_split, batch_size, window, fib_positions, generator,
+              n_batches=16):
+    model.eval()
+    losses = []
+    with torch.no_grad():
+        for _ in range(n_batches):
+            x, y = get_fib_strided_batch(val_split, batch_size, window,
+                                           fib_positions, generator)
+            logits = model(x)
+            losses.append(F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)), y.reshape(-1)).item())
+    model.train()
+    return sum(losses) / len(losses)
+
+
+def apply_tier_lr_scale(model: torch.nn.Module):
+    """For each FibGenLinear, multiply seed.grad by tier_lr_scale.
+    Tier-1 components get full grad; tier-k get grad * 1/sqrt(k)."""
+    for m in model.modules():
+        if isinstance(m, FibGenLinear) and m.seed.grad is not None:
+            m.seed.grad.mul_(m.tier_lr_scale)
+
+
+def train_one(name, model, train_split, val_split, args, fib_positions,
+               apply_lr_scale: bool = False):
+    torch.manual_seed(args.seed)
+    gen = torch.Generator(); gen.manual_seed(args.seed + 1)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"\n[train {name}] params={n_params:,}  "
+          f"apply_lr_scale={apply_lr_scale}", flush=True)
+    t0 = time.time()
+    best_val = float("inf")
+    best_step = -1
+    eval_every = 200
+    val_hist = []
+    for step in range(args.steps):
+        x, y = get_fib_strided_batch(train_split, args.batch_size, args.seq_len,
+                                       fib_positions, gen)
+        logits = model(x)
+        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+        optimizer.zero_grad(); loss.backward()
+        if apply_lr_scale:
+            apply_tier_lr_scale(model)
+        optimizer.step()
+        if step % eval_every == 0 or step == args.steps - 1:
+            vl = evaluate(model, val_split, args.batch_size, args.seq_len,
+                          fib_positions, gen)
+            val_hist.append((step, vl, time.time() - t0))
+            marker = ""
+            if vl < best_val:
+                best_val = vl
+                best_step = step
+                marker = " ← BEST"
+            print(f"  step {step:5d}  val={vl:.4f}  ({time.time()-t0:.1f}s){marker}",
+                  flush=True)
+    return {"name": name, "n_params": n_params, "best_val": best_val,
+             "best_step": best_step, "wall_time": time.time() - t0,
+             "val_history": val_hist}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--steps", type=int, default=2500)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--seq-len", type=int, default=128)
+    parser.add_argument("--d-model", type=int, default=128)
+    parser.add_argument("--n-blocks", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=3e-4)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--distractor-frac", type=float, default=0.20)
+    parser.add_argument("--out", type=str, default="results_lazy_subsim.json")
+    args = parser.parse_args()
+
+    chars, stoi, itos, encoded = make_dataset(seq_len=args.seq_len,
+                                                 source="tinyshakespeare")
+    vocab_size = len(chars)
+    train_split, val_split = build_distractor_stream(
+        encoded, args.distractor_frac, args.seq_len, args.seed,
+    )
+    fib_positions = fib_positions_in_window(args.seq_len)
+
+    results = {}
+
+    # 1. Baseline Subsim (no lazy)
+    m = SubsimLM(vocab_size=vocab_size, d_model=args.d_model,
+                  n_blocks=args.n_blocks, seq_len=args.seq_len,
+                  K=32, fibgen_K=32, mode="cross",
+                  lazy_tier_dropout=False)
+    results["subsim_baseline"] = train_one(
+        "subsim_baseline", m, train_split, val_split, args, fib_positions,
+    )
+
+    # 2. Subsim + lazy seed dropout
+    m = SubsimLM(vocab_size=vocab_size, d_model=args.d_model,
+                  n_blocks=args.n_blocks, seq_len=args.seq_len,
+                  K=32, fibgen_K=32, mode="cross",
+                  lazy_tier_dropout=True)
+    results["subsim_lazy_dropout"] = train_one(
+        "subsim_lazy_dropout", m, train_split, val_split, args, fib_positions,
+    )
+
+    # 3. Subsim + tier-weighted gradient scaling
+    m = SubsimLM(vocab_size=vocab_size, d_model=args.d_model,
+                  n_blocks=args.n_blocks, seq_len=args.seq_len,
+                  K=32, fibgen_K=32, mode="cross",
+                  lazy_tier_dropout=False)
+    results["subsim_tier_lr"] = train_one(
+        "subsim_tier_lr", m, train_split, val_split, args, fib_positions,
+        apply_lr_scale=True,
+    )
+
+    # Summary
+    print()
+    print("=" * 84)
+    print(f"{'config':<24} {'params':>10} {'best_val':>10} {'best_step':>10} "
+          f"{'wall':>10}")
+    print("-" * 84)
+    for name, r in results.items():
+        print(f"{name:<24} {r['n_params']:>10,} {r['best_val']:>10.4f} "
+              f"{r['best_step']:>10} {r['wall_time']:>9.1f}s")
+
+    out_path = Path(__file__).parent / args.out
+    with open(out_path, "w") as f:
+        json.dump(results, f, indent=2, default=str)
+    print(f"\nWrote {out_path}")
+
+
+if __name__ == "__main__":
+    main()