Upload folder using huggingface_hub

README.md

@@ -1,3 +1,98 @@
 ---
-license: mit
+license: apache-2.0
+datasets:
+  - HuggingFaceFW/fineweb
+language:
+  - en
+library_name: pytorch
+tags:
+  - looped-transformer
+  - weight-tying
+  - recurrent-depth
+  - small-language-model
 ---
+
+# SimLoop 1+loop x1+1 (10.0M)
+
+A **looped** (weight-tied recurrent) Qwen3-style transformer trained from
+scratch on FineWeb under a hard budget of **<= 10M parameters** and
+**<= 100M training tokens**.
+
+One middle block is applied **K = 1** times; the layers around it are
+ordinary unlooped layers:
+
+```
+embed -> [1 layer] -> ( 1 looped layer ) x K -> [1 layer] -> RMSNorm -> tied head
+```
+
+| | |
+|---|---|
+| parameters | **9,962,496** (incl. embeddings, input/output tied) |
+| d_model / d_mlp | 384 / 720 |
+| heads (GQA) | 6 query / 2 key-value, head_dim 64 |
+| context | 512 tokens |
+| vocabulary | 16,384 byte-level BPE trained on FineWeb |
+| training tokens | 60,014,592 |
+| primitives | RMSNorm, RoPE, SwiGLU, GQA, QK-norm |
+
+## Results
+
+| k | CE (nats) | perplexity | bits-per-byte |
+|---|---|---|---|
+| 0 | 5.4088 | 223.35 | 1.8913 |
+| 1 | 4.3966 | 81.17 | 1.5374 **<- best** |
+| 2 | 4.5995 | 99.43 | 1.6083 |
+| 3 | 4.9413 | 139.95 | 1.7278 |
+| 4 | 5.2659 | 193.62 | 1.8413 |
+| 5 | 5.5476 | 256.63 | 1.9399 |
+| 6 | 5.7900 | 327.01 | 2.0246 |
+| 7 | 6.0000 | 403.41 | 2.0980 |
+| 8 | 6.1835 | 484.69 | 2.1622 |
+| 9 | 6.3452 | 569.77 | 2.2188 |
+| 10 | 6.4887 | 657.65 | 2.2689 |
+| 11 | 6.6166 | 747.41 | 2.3137 |
+| 12 | 6.7313 | 838.20 | 2.3537 |
+
+Reference points on the same validation split: a context-free **unigram** model scores CE 7.5476 (ppl 1896.10); **uniform** over the vocabulary scores CE 9.7041.
+
+
+`k` is the number of applications of the looped block at inference. The model
+is weight-tied, so **any k can be run**; the table is a single checkpoint
+evaluated at every depth.
+
+**Measured behaviour:** TRAINED UNLOOPED (K=1): the k>1 rows show how a model trained at one application behaves when looped anyway, not whether looping pays.
+
+- first application of the looped block buys **+1.0122** nats (k=0 -> k=1)
+- all further applications buy **+0.0000** nats (k=1 -> k=1)
+
+
+## Usage
+
+```python
+import torch
+from simloop.stack import StackConfig, StackedLoop
+
+ck = torch.load("model.pt", map_location="cpu", weights_only=False)
+model = StackedLoop(StackConfig(**ck["model_cfg"]))
+model.load_state_dict(ck["model"]); model.eval()
+
+ids = torch.tensor([[1, 2, 3]])            # from tokenizer.json
+logits = model(ids, K=1)           # try other K: the block is tied
+```
+
+See `load_example.py`. The tokenizer is a `tokenizers` BPE:
+`Tokenizer.from_file("tokenizer.json")`.
+
+## Honest limitations
+
+- Trained on 60,014,592 tokens at ~10M parameters. It is a research artifact
+  for studying looped depth, **not** a useful general-purpose language model.
+- Perplexity is tokenizer-dependent; **bits-per-byte** is the comparable
+  number and is reported above.
+- The looped block **saturates**: past the depth listed as best above, extra
+  applications make cross-entropy worse, not better. This is measured, not
+  assumed, and is the central finding of the project.
+- English-only, no instruction tuning, no safety filtering beyond FineWeb's.
+
+Full experimental record, including every failed experiment:
+https://github.com/brkdrd/SimLoop

config.json

@@ -0,0 +1,62 @@
+{
+  "arch": "stack",
+  "model": {
+    "vocab_size": 16384,
+    "d_model": 384,
+    "n_pre": 1,
+    "n_loop": 1,
+    "n_post": 1,
+    "n_heads": 6,
+    "n_kv_heads": 2,
+    "head_dim": 64,
+    "d_mlp": 720,
+    "max_seq_len": 512,
+    "rope_theta": 10000.0,
+    "dropout": 0.0,
+    "norm_eps": 1e-06,
+    "init_std": 0.02,
+    "scale_resid_init": true
+  },
+  "train": {
+    "K": 1,
+    "k_schedule": "const",
+    "k_end": 1,
+    "k_decay_start_frac": 0.5,
+    "k_decay_end_frac": 0.9,
+    "k_min": 1,
+    "drop_eps": 0.01,
+    "eval_k_max": 8,
+    "batch_size": 16,
+    "grad_accum": 2,
+    "seq_len": 512,
+    "max_tokens": 60000000,
+    "lr": 0.001,
+    "min_lr_ratio": 0.1,
+    "warmup_steps": 100,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "ce_weight": 1.0,
+    "teacher_run": "",
+    "teacher_k": 0,
+    "init_from_run": "",
+    "distill_weight": 0.0,
+    "distill_space": "hidden",
+    "distill_depths": "first",
+    "distill_start_frac": 0.0,
+    "distill_tau": 1.0,
+    "anytime_ce_weight": 0.0,
+    "grad_checkpoint": false,
+    "budget_locked": true,
+    "log_every": 20,
+    "eval_every": 150,
+    "val_batches": 32,
+    "ckpt_every": 1000,
+    "best_metric": "val_ce",
+    "seed": 1337
+  },
+  "K": 1,
+  "params": 9962496,
+  "training_tokens": 60014592
+}

eval_results.json

@@ -0,0 +1,181 @@
+{
+  "ckpt": "/kaggle/working/out/stack_k1_60m/best.pt",
+  "arch": "stack",
+  "K_train": 1,
+  "k_max": 12,
+  "val_tokens": 1539584,
+  "params": 9962496,
+  "uniform_ce": 9.704060527839234,
+  "unigram_ce": 7.547556945473469,
+  "best_k": 1,
+  "gain_first_loop": 1.012203666907749,
+  "gain_extra_loops": 0.0,
+  "mean_rel_step": 0.28043451470633346,
+  "verdict": "TRAINED UNLOOPED (K=1): the k>1 rows show how a model trained at one application behaves when looped anyway, not whether looping pays.",
+  "rows": [
+    {
+      "k": 0,
+      "ce": 5.408761565218267,
+      "ppl": 223.35480596749395,
+      "bpb": 1.8913016483504717,
+      "delta_ce": null,
+      "flops_per_token": 18259968
+    },
+    {
+      "k": 1,
+      "ce": 4.396557898310518,
+      "ppl": 81.17098845673196,
+      "bpb": 1.5373606508401174,
+      "delta_ce": 1.012203666907749,
+      "flops_per_token": 21098496,
+      "std_ratio": 0.9512100219726562,
+      "erank": 172.63084411621094,
+      "step_cos": 0.7777038216590881,
+      "rel_step": 0.8602083921432495,
+      "state_norm": 15.174788475036621
+    },
+    {
+      "k": 2,
+      "ce": 4.5994899281640365,
+      "ppl": 99.4335844337849,
+      "bpb": 1.6083206437954694,
+      "delta_ce": -0.20293202985351844,
+      "flops_per_token": 23937024,
+      "std_ratio": 0.9564117789268494,
+      "erank": 152.0776138305664,
+      "step_cos": 0.9423013627529144,
+      "rel_step": 0.5801940560340881,
+      "state_norm": 21.41031837463379
+    },
+    {
+      "k": 3,
+      "ce": 4.941250652976713,
+      "ppl": 139.94516299387544,
+      "bpb": 1.727825379655199,
+      "delta_ce": -0.3417607248126764,
+      "flops_per_token": 26775552,
+      "std_ratio": 0.9550213515758514,
+      "erank": 126.67342376708984,
+      "step_cos": 0.9775733947753906,
+      "rel_step": 0.3683442771434784,
+      "state_norm": 27.632020950317383
+    },
+    {
+      "k": 4,
+      "ce": 5.265907488037353,
+      "ppl": 193.6219386774754,
+      "bpb": 1.8413493351660124,
+      "delta_ce": -0.32465683506064025,
+      "flops_per_token": 29614080,
+      "std_ratio": 0.9511829912662506,
+      "erank": 109.17454528808594,
+      "step_cos": 0.9875146746635437,
+      "rel_step": 0.27501752972602844,
+      "state_norm": 33.76785087585449
+    },
+    {
+      "k": 5,
+      "ce": 5.547618621125902,
+      "ppl": 256.625704638506,
+      "bpb": 1.9398563083325306,
+      "delta_ce": -0.28171113308854867,
+      "flops_per_token": 32452608,
+      "std_ratio": 0.945938229560852,
+      "erank": 88.04708099365234,
+      "step_cos": 0.9916830360889435,
+      "rel_step": 0.2244347706437111,
+      "state_norm": 39.811710357666016
+    },
+    {
+      "k": 6,
+      "ce": 5.789976763812497,
+      "ppl": 327.00542592830055,
+      "bpb": 2.024602575888849,
+      "delta_ce": -0.24235814268659528,
+      "flops_per_token": 35291136,
+      "std_ratio": 0.9405999779701233,
+      "erank": 86.94815063476562,
+      "step_cos": 0.9938283264636993,
+      "rel_step": 0.19342376291751862,
+      "state_norm": 45.64400100708008
+    },
+    {
+      "k": 7,
+      "ce": 5.99996355549616,
+      "ppl": 403.4140909984358,
+      "bpb": 2.0980294334891267,
+      "delta_ce": -0.2099867916836633,
+      "flops_per_token": 38129664,
+      "std_ratio": 0.9329713881015778,
+      "erank": 64.25496482849121,
+      "step_cos": 0.9951082766056061,
+      "rel_step": 0.1726725548505783,
+      "state_norm": 53.134456634521484
+    },
+    {
+      "k": 8,
+      "ce": 6.18351374157727,
+      "ppl": 484.6920503676553,
+      "bpb": 2.1622121054934955,
+      "delta_ce": -0.18355018608110996,
+      "flops_per_token": 40968192,
+      "std_ratio": 0.9266401827335358,
+      "erank": 56.75149154663086,
+      "step_cos": 0.9959644377231598,
+      "rel_step": 0.157696433365345,
+      "state_norm": 59.62197303771973
+    },
+    {
+      "k": 9,
+      "ce": 6.345236095432273,
+      "ppl": 569.7718943785347,
+      "bpb": 2.218762158723424,
+      "delta_ce": -0.16172235385500233,
+      "flops_per_token": 43806720,
+      "std_ratio": 0.9184004366397858,
+      "erank": 48.32929611206055,
+      "step_cos": 0.9965897798538208,
+      "rel_step": 0.14616148173809052,
+      "state_norm": 67.28807830810547
+    },
+    {
+      "k": 10,
+      "ce": 6.488679977838169,
+      "ppl": 657.6546712541515,
+      "bpb": 2.268920711280938,
+      "delta_ce": -0.14344388240589634,
+      "flops_per_token": 46645248,
+      "std_ratio": 0.9087829887866974,
+      "erank": 38.35821723937988,
+      "step_cos": 0.9970792233943939,
+      "rel_step": 0.13675487786531448,
+      "state_norm": 75.25280380249023
+    },
+    {
+      "k": 11,
+      "ce": 6.616616099706154,
+      "ppl": 747.4116465669175,
+      "bpb": 2.3136566078914447,
+      "delta_ce": -0.1279361218679851,
+      "flops_per_token": 49483776,
+      "std_ratio": 0.8985040187835693,
+      "erank": 32.420021057128906,
+      "step_cos": 0.9974798560142517,
+      "rel_step": 0.12871047109365463,
+      "state_norm": 83.94036483764648
+    },
+    {
+      "k": 12,
+      "ce": 6.731254219889919,
+      "ppl": 838.1978914258499,
+      "bpb": 2.3537425431010153,
+      "delta_ce": -0.11463812018376451,
+      "flops_per_token": 52322304,
+      "std_ratio": 0.8866880834102631,
+      "erank": 26.89052963256836,
+      "step_cos": 0.9978146255016327,
+      "rel_step": 0.12159556895494461,
+      "state_norm": 93.48854446411133
+    }
+  ]
+}

load_example.py

@@ -0,0 +1,15 @@
+"""Minimal loader. Requires: torch, tokenizers."""
+import torch
+from tokenizers import Tokenizer
+from simloop.stack import StackConfig, StackedLoop
+
+ck = torch.load("model.pt", map_location="cpu", weights_only=False)
+model = StackedLoop(StackConfig(**ck["model_cfg"]))
+model.load_state_dict(ck["model"])
+model.eval()
+
+tok = Tokenizer.from_file("tokenizer.json")
+ids = torch.tensor([tok.encode("The capital of France is").ids])
+with torch.no_grad():
+    logits = model(ids, K=1)
+print(tok.decode([int(logits[0, -1].argmax())]))

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4ee0d100eb055908914969dc5f291a500eb073e8c30a41ee632f49316ac4c2a
+size 39862696

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e1ecc44d47c74e4a057f099f3095b18e2236665637563710345728b30e98086
+size 39853464

run_config.json

@@ -0,0 +1,55 @@
+{
+  "arch": "stack",
+  "model": {
+    "vocab_size": 16384,
+    "d_model": 384,
+    "n_pre": 1,
+    "n_loop": 1,
+    "n_post": 1,
+    "n_heads": 6,
+    "n_kv_heads": 2,
+    "head_dim": 64,
+    "d_mlp": 720,
+    "max_seq_len": 512,
+    "rope_theta": 10000.0,
+    "dropout": 0.0,
+    "norm_eps": 1e-06,
+    "init_std": 0.02,
+    "scale_resid_init": true
+  },
+  "train": {
+    "K": 1,
+    "k_schedule": "const",
+    "k_end": 1,
+    "k_decay_start_frac": 0.5,
+    "k_decay_end_frac": 0.9,
+    "k_min": 1,
+    "drop_eps": 0.01,
+    "eval_k_max": 8,
+    "batch_size": 16,
+    "grad_accum": 2,
+    "seq_len": 512,
+    "max_tokens": 60000000,
+    "lr": 0.001,
+    "min_lr_ratio": 0.1,
+    "warmup_steps": 100,
+    "weight_decay": 0.1,
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "grad_clip": 1.0,
+    "distill_weight": 0.0,
+    "distill_space": "hidden",
+    "distill_depths": "first",
+    "distill_start_frac": 0.0,
+    "distill_tau": 1.0,
+    "anytime_ce_weight": 0.0,
+    "grad_checkpoint": false,
+    "log_every": 20,
+    "eval_every": 150,
+    "val_batches": 32,
+    "ckpt_every": 1000,
+    "best_metric": "val_ce",
+    "seed": 1337,
+    "budget_locked": true
+  }
+}

simloop/model.py

@@ -0,0 +1,205 @@
+"""SimLoop model: Qwen3-style looped transformer.
+
+A single block of cfg.n_layers transformer layers (pre-RMSNorm, RoPE, SwiGLU,
+GQA, QK-norm) is applied K times to the embedded input. Input/output
+embeddings are tied. Dropout inside the block doubles as the stochastic
+augmentation for the siamese consistency objective: two forward passes of the
+same input in train mode give two "views" (SimCSE-style).
+
+Loop states are exposed via forward_loops(..., collect=True) so a future
+early-exit head can read intermediate states.
+"""
+
+from dataclasses import dataclass, replace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+
+@dataclass
+class SimLoopConfig:
+    vocab_size: int = 16384
+    d_model: int = 384
+    n_layers: int = 2       # layers inside the looped block
+    n_heads: int = 6
+    n_kv_heads: int = 2     # GQA
+    head_dim: int = 64
+    d_mlp: int = 1024
+    max_seq_len: int = 512
+    rope_theta: float = 10000.0
+    dropout: float = 0.1    # the augmentation for the siamese objective
+    # Where the augmentation noise enters:
+    #   "input" - dropout applied ONCE to the embedding, per branch; the
+    #             loop itself is deterministic. Branch divergence is then a
+    #             fixed input perturbation, so consistency across depth
+    #             measures whether the loop CONTRACTS it.
+    #   "loop"  - standard transformer dropout inside every layer, resampled
+    #             at every loop application. Divergence accumulates with
+    #             depth and contraction cannot be measured.
+    aug_mode: str = "loop"
+    norm_eps: float = 1e-6
+    predictor_hidden: int = 0  # >0: SimSiam-style MLP predictor (train-only)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x):
+        dtype = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x * self.weight.float()).to(dtype)
+
+
+def _rope_tables(head_dim, max_len, theta):
+    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+    t = torch.arange(max_len, dtype=torch.float32)
+    freqs = torch.outer(t, inv_freq)  # [T, head_dim/2]
+    return freqs.cos(), freqs.sin()
+
+
+def _apply_rope(x, cos, sin):
+    # x: [B, H, T, hd]; cos/sin: [max_T, hd/2]
+    T = x.shape[-2]
+    cos = cos[:T].view(1, 1, T, -1)
+    sin = sin[:T].view(1, 1, T, -1)
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
+    return out.flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.n_heads, self.n_kv, self.hd = cfg.n_heads, cfg.n_kv_heads, cfg.head_dim
+        self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.head_dim, bias=False)
+        self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=False)
+        self.o_proj = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.d_model, bias=False)
+        self.q_norm = RMSNorm(cfg.head_dim, cfg.norm_eps)  # Qwen3 QK-norm
+        self.k_norm = RMSNorm(cfg.head_dim, cfg.norm_eps)
+        self.dropout = cfg.dropout
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.q_norm(self.q_proj(x).view(B, T, self.n_heads, self.hd))
+        k = self.k_norm(self.k_proj(x).view(B, T, self.n_kv, self.hd))
+        v = self.v_proj(x).view(B, T, self.n_kv, self.hd)
+        q, k, v = (t.transpose(1, 2) for t in (q, k, v))  # [B, H, T, hd]
+        q, k = _apply_rope(q, cos, sin), _apply_rope(k, cos, sin)
+        rep = self.n_heads // self.n_kv
+        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
+        o = F.scaled_dot_product_attention(
+            q, k, v, is_causal=True,
+            dropout_p=self.dropout if self.training else 0.0)
+        return self.o_proj(o.transpose(1, 2).reshape(B, T, -1))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.gate = nn.Linear(cfg.d_model, cfg.d_mlp, bias=False)
+        self.up = nn.Linear(cfg.d_model, cfg.d_mlp, bias=False)
+        self.down = nn.Linear(cfg.d_mlp, cfg.d_model, bias=False)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+class Layer(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.ln_attn = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.attn = Attention(cfg)
+        self.ln_mlp = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.mlp = SwiGLU(cfg)
+        self.drop = nn.Dropout(cfg.dropout)
+
+    def forward(self, x, cos, sin):
+        x = x + self.drop(self.attn(self.ln_attn(x), cos, sin))
+        x = x + self.drop(self.mlp(self.ln_mlp(x)))
+        return x
+
+
+class PredictorMLP(nn.Module):
+    """SimSiam-style bottleneck predictor h, applied to post-norm states.
+    Training-time only — discarded at inference, exactly like SimSiam's h.
+    Restores the predictor asymmetry the extra-loop predictor lacks (g=f)."""
+
+    def __init__(self, d_model, hidden):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, hidden)
+        self.ln = nn.LayerNorm(hidden)
+        self.fc2 = nn.Linear(hidden, d_model)
+
+    def forward(self, x):
+        return self.fc2(F.gelu(self.ln(self.fc1(x))))
+
+
+class SimLoop(nn.Module):
+    def __init__(self, cfg: SimLoopConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        # input mode: noise lives at the embedding, the loop is deterministic
+        loop_cfg = replace(cfg, dropout=0.0) if cfg.aug_mode == "input" else cfg
+        self.in_drop = nn.Dropout(cfg.dropout if cfg.aug_mode == "input" else 0.0)
+        self.layers = nn.ModuleList(Layer(loop_cfg) for _ in range(cfg.n_layers))
+        self.norm_out = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.predictor = (PredictorMLP(cfg.d_model, cfg.predictor_hidden)
+                          if cfg.predictor_hidden > 0 else None)
+        cos, sin = _rope_tables(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        self.apply(self._init)
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, std=0.02)
+
+    def embed_aug(self, x):
+        """Embedded input with the branch augmentation applied. Call once PER
+        BRANCH so each draws its own mask. A no-op unless aug_mode="input",
+        so aug_mode="loop" behaves exactly as before."""
+        return self.in_drop(self.embed(x))
+
+    def step(self, z):
+        """One loop iteration: the whole block applied once."""
+        for layer in self.layers:
+            z = layer(z, self.rope_cos, self.rope_sin)
+        return z
+
+    def forward_loops(self, z0, n_steps, grad_checkpoint=False, collect=False):
+        """Apply the block n_steps times.
+
+        Returns (z_final, states); states = [z0, z1, ..., z_n] if collect,
+        else None. Checkpointing preserves RNG state, so recomputed dropout
+        masks match the original forward.
+        """
+        states = [z0] if collect else None
+        z = z0
+        for _ in range(n_steps):
+            if grad_checkpoint and self.training and torch.is_grad_enabled():
+                z = checkpoint(self.step, z, use_reentrant=False)
+            else:
+                z = self.step(z)
+            if collect:
+                states.append(z)
+        return z, states
+
+    def post(self, z):
+        """Final RMSNorm — the space where consistency is matched (per token)
+        and where the LM head reads."""
+        return self.norm_out(z)
+
+    def logits(self, z):
+        return F.linear(self.post(z), self.embed.weight)  # tied head
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

simloop/stack.py

@@ -0,0 +1,158 @@
+"""SimLoop-Stack (E9+): an ordinary Qwen3-style transformer stack in which
+ONE middle block is applied K times (weight-tied); every other layer is a
+plain unlooped layer.
+
+    embed -> [pre layers] -> ( loop block ) x K -> [post layers] -> norm -> tied head
+
+This is deliberately NOT the phase-1 SimLoop architecture (`simloop/model.py`,
+where every layer lives inside the loop and there is no readout stack). It is
+kept as a separate class so every phase-1 result stays bit-reproducible.
+
+The point of the arrangement is that K is the ONLY thing that differs between
+the baseline and the looped run: same parameters, same initial values (init
+draws do not depend on K), same data order. So any difference in the learning
+curve is attributable to the extra loop applications and nothing else.
+
+Two properties the experiment depends on:
+
+* k = 0 is a real forward pass. `enter()` -> `readout()` skips the looped
+  block entirely, so "the loops stopped mattering" can be told apart from
+  "the loop learned to be the identity" (the trivial minimiser of any
+  deep-to-shallow distillation term).
+* Weight tying means one checkpoint can be evaluated at any k, including k
+  larger than the K it was trained at.
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from simloop.model import RMSNorm, Layer, _rope_tables
+
+
+@dataclass
+class StackConfig:
+    vocab_size: int = 16384
+    d_model: int = 384
+    n_pre: int = 1          # unlooped layers before the loop
+    n_loop: int = 1         # layers inside the looped (weight-tied) block
+    n_post: int = 1         # unlooped layers after the loop
+    n_heads: int = 6
+    n_kv_heads: int = 2     # GQA
+    head_dim: int = 64
+    d_mlp: int = 720        # 1.875x d_model: the largest that fits 3 blocks
+                            # + a tied 16k embedding under the 10M cap
+    max_seq_len: int = 512
+    rope_theta: float = 10000.0
+    dropout: float = 0.0    # 0: <1 epoch of FineWeb, and it would confound K
+    norm_eps: float = 1e-6
+    init_std: float = 0.02
+    # Scale residual-branch output init by 1/sqrt(2*depth). depth counts
+    # PARAMETER blocks, not loop applications, on purpose: making it depend
+    # on K would give the baseline and the looped run different initial
+    # weights and destroy the controlled comparison.
+    scale_resid_init: bool = True
+
+
+class StackedLoop(nn.Module):
+    def __init__(self, cfg: StackConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.pre_layers = nn.ModuleList(Layer(cfg) for _ in range(cfg.n_pre))
+        self.loop_layers = nn.ModuleList(Layer(cfg) for _ in range(cfg.n_loop))
+        self.post_layers = nn.ModuleList(Layer(cfg) for _ in range(cfg.n_post))
+        self.norm_out = RMSNorm(cfg.d_model, cfg.norm_eps)
+        cos, sin = _rope_tables(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        self._init_weights()
+
+    # ---------------------------------------------------------------- init
+    def _init_weights(self):
+        for m in self.modules():
+            if isinstance(m, (nn.Linear, nn.Embedding)):
+                nn.init.normal_(m.weight, std=self.cfg.init_std)
+        if not self.cfg.scale_resid_init:
+            return
+        depth = self.cfg.n_pre + self.cfg.n_loop + self.cfg.n_post
+        std = self.cfg.init_std / math.sqrt(2.0 * depth)
+        for layer in self.all_layers():
+            nn.init.normal_(layer.attn.o_proj.weight, std=std)
+            nn.init.normal_(layer.mlp.down.weight, std=std)
+
+    def all_layers(self):
+        return list(self.pre_layers) + list(self.loop_layers) + list(self.post_layers)
+
+    # ------------------------------------------------------------- forward
+    def enter(self, x):
+        """Token ids -> the state that enters the looped block (k = 0)."""
+        h = self.embed(x)
+        for layer in self.pre_layers:
+            h = layer(h, self.rope_cos, self.rope_sin)
+        return h
+
+    def loop_once(self, h):
+        """One application of the looped block."""
+        for layer in self.loop_layers:
+            h = layer(h, self.rope_cos, self.rope_sin)
+        return h
+
+    def readout(self, h):
+        """Loop-exit state -> final normalised hidden state."""
+        for layer in self.post_layers:
+            h = layer(h, self.rope_cos, self.rope_sin)
+        return self.norm_out(h)
+
+    def logits_at(self, h):
+        """Logits from a loop-exit state at ANY k (tied head)."""
+        return F.linear(self.readout(h), self.embed.weight)
+
+    def run(self, x, K, collect=True, grad_checkpoint=False):
+        """Returns (h_K, states) with states = [h_0, h_1, ..., h_K], h_k the
+        loop-exit state after k applications. Collecting is free: autograd
+        already holds these tensors."""
+        h = self.enter(x)
+        states = [h] if collect else None
+        for _ in range(K):
+            if grad_checkpoint and self.training and torch.is_grad_enabled():
+                h = checkpoint(self.loop_once, h, use_reentrant=False)
+            else:
+                h = self.loop_once(h)
+            if collect:
+                states.append(h)
+        return h, states
+
+    def forward(self, x, K):
+        h, _ = self.run(x, K, collect=False)
+        return self.logits_at(h)
+
+    # --------------------------------------------------------------- meta
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def flops_per_token(self, K, seq_len=None, backward=True):
+        """Analytic FLOPs/token. Matmuls counted as 2*params (multiply+add);
+        attention scores counted causally (half the full T x T grid).
+        backward=True applies the usual 3x (fwd + bwd ~ 2x fwd)."""
+        c = self.cfg
+        seq_len = seq_len or c.max_seq_len
+        qo = 2 * c.d_model * c.n_heads * c.head_dim          # q_proj + o_proj
+        kv = 2 * c.d_model * c.n_kv_heads * c.head_dim       # k_proj + v_proj
+        mlp = 3 * c.d_model * c.d_mlp                        # gate + up + down
+        per_layer = 2 * (qo + kv + mlp) + 2 * seq_len * c.n_heads * c.head_dim
+        n_applied = c.n_pre + K * c.n_loop + c.n_post
+        head = 2 * c.d_model * c.vocab_size
+        f = n_applied * per_layer + head
+        return f * (3 if backward else 1)
+
+    def describe(self, K):
+        c = self.cfg
+        return (f"{c.n_pre}pre + {c.n_loop}loop x K={K} + {c.n_post}post "
+                f"= {c.n_pre + K * c.n_loop + c.n_post} layer applications, "
+                f"{c.n_pre + c.n_loop + c.n_post} parameter blocks, "
+                f"d_model={c.d_model} d_mlp={c.d_mlp} vocab={c.vocab_size}")