Alpha: System-One-style typed decision model

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+alpha_report.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,65 +1,62 @@
 ---
 license: apache-2.0
-base_model: convaiinnovations/laya
-tags:
-- system-one
-- calibrated-decisions
-- rlcd
-- jev-alternative
-- laya
-- routing
-- guardrails
+base_model: distilroberta-base
+tags: [system-one, structured-output, calibrated, classification, parallel-decoding]
 ---
-
-# Alpha (System-One Decision Model)
-
-Alpha is an open, calibrated **System-One decision model** continuing [convaiinnovations/laya](https://huggingface.co/convaiinnovations/laya) (Apache-2.0).
-
-Given text/JSON states and typed questions (`choice`, `score`, `noul`), Alpha returns typed answers with mathematically calibrated probabilities in a single forward pass without autoregressive token generation.
-
-## Benchmark Results (Measured vs. TypeSafe Jev)
-
-Evaluated on held-out test splits:
-
-| Benchmark | Question Type | Chance | TypeSafe Jev (Ref) | Alpha (Ours @ Step 2500) | Delta vs. Jev |
-| :--- | :--- | :--- | :--- | :--- | :--- |
-| **AG News** | choice ($k=4$) | 25.0% | 91.0% | **94.8%** | **+3.8%** |
-| **Emotion** | choice ($k=6$) | 16.7% | 48.0% | **88.5%** | **+40.5%** |
-| **SST-2** | noul ($k=2$) | 50.0% | 89.8% | **93.2%** | **+3.4%** |
-| **Banking77** | choice ($k=77$) | 1.3% | **87.0%** | **86.5%** | -0.5% |
-| **SMS Spam** | noul ($k=2$) | 50.0% | N/A | **99.0%** | — |
-| **MNLI** | choice ($k=3$) | 33.3% | N/A | **87.5%** | — |
-| **SST-5** | score ($k=5$) | 20.0% | N/A | **54.7%** | — |
-
-## Quickstart
-
+# Alpha (System-One-style typed decision model)
+
+Alpha takes unstructured / structured state plus **typed questions with declared options** and returns calibrated probabilities for
+**all questions in one parallel forward pass**. It cannot generate strings and cannot answer outside the declared options
+(masked softmax -> schema violations are impossible by construction). Option embeddings are independent of the state, so a schema can be
+compiled once and cached; cardinality is limited only by memory.
+
+Architecture: `distilroberta-base` encoder (state + question slots with `<mask>` readout tokens) -> 2-layer **Slot Reasoner** (set transformer, squared-ReLU,
+zero-init residuals) -> biaffine/cosine scorer against option embeddings with per-question learned temperature + post-hoc temperature (T=1.04).
+
+Training: 6000 steps x 32 workflow-style examples (1-4 questions over named fields, mixed from 10 public datasets with paraphrased
+questions/option wordings), proper-scoring-rule loss (NLL + Brier + ordinal term), EMA weights, functional-memory anchor, discovery controller.
+Held-out families (never trained on): cola, dbpedia, offensive.
+
+## Results (test splits; see results.json, alpha_report.png)
+| family | acc (seen wording) | acc (unseen wording) | ECE |
+|---|---|---|---|
+| ag_news | 0.926 | 0.924 | 0.032 |
+| banking77 | 0.902 | 0.890 | 0.036 |
+| emotion | 0.922 | 0.894 | 0.030 |
+| imdb | 0.886 | 0.892 | 0.027 |
+| mrpc | 0.822 | 0.822 | 0.042 |
+| news20 | 0.674 | 0.678 | 0.068 |
+| sms | 0.993 | 0.346 | 0.004 |
+| snli | 0.772 | 0.538 | 0.046 |
+| sst2 | 0.912 | 0.910 | 0.034 |
+| yelp | 0.572 | 0.522 | 0.072 |
+
+Workflow eval (4 questions per query): per-question acc 0.825, all-correct 0.437, joint-confidence ECE 0.101.
+
+Held-out task families (unseen schema wording): zero-shot / after 10 / after 40 gradient steps
+| family | zero-shot | +10 steps | +40 steps |
+|---|---|---|---|
+| cola | 0.374 | 0.658 | 0.686 |
+| dbpedia | 0.312 | 0.586 | 0.952 |
+| offensive | 0.528 | 0.632 | 0.694 |
+
+Speed on Tesla T4: p50 11.1 ms, p95 13.5 ms for a 3-question query (local, no network); 354 workflows/s at batch 64.
+Qwen2.5-0.5B-Instruct (generative, same 120 questions): acc 0.583, p50 latency 82 ms vs Alpha acc 0.875, p50 11.1 ms.
+
+Framework diagnostics: DMR=1.08; FRR by perturbation: sigma=0.5: FRR 4250.0 sigma=1: FRR 4250.0 sigma=2: FRR 4250.0 sigma=4: FRR 4250.0 sigma=8: FRR 14.2* (* = not re-reached within budget); MGR by checkpoint in results.json.
+
+## Honest limits
+* This is a small research model trained on public classification datasets. **It is not a replacement for, and has not been compared against, TypeSafe's Jev**
+  (a closed API; its claims were not verified here). Do not read the latency comparison as a Jev comparison.
+* Narrow task coverage (sentiment/topic/intent/NLI/paraphrase/spam/rating); free-form reasoning is out of scope by design.
+* Calibration is measured in-distribution; shift will degrade it.
+
+## Usage
 ```python
-from laya_mm import LayaMM
-from huggingface_hub import snapshot_download
-
-agent = LayaMM.from_dir(snapshot_download("SofiTesfay2010/Alpha"))
-
-# 1. Binary Decision (noul)
-res = agent.system_one(
-    "Your package was delivered to the front porch.",
-    {"delivered": {"type": "noul", "instructions": "Was the delivery completed?"}}
-)
-print(res["answers"]["delivered"])
-
-# 2. Multi-Class Choice
-res = agent.system_one(
-    "I received a bill with unauthorized roaming charges.",
-    {"intent": {
-        "type": "choice",
-        "instructions": "Identify the customer support ticket intent.",
-        "criteria": ["billing_dispute", "technical_issue", "account_closure"]
-    }}
-)
-print(res["answers"]["intent"])
+from huggingface_hub import snapshot_download; import sys
+p = snapshot_download("SofiTesfay2010/Alpha"); sys.path.insert(0, p)
+from alpha_s1 import Alpha
+m = Alpha.from_pretrained(p, device="cuda").eval()
+print(m.predict({"ticket": "I was charged twice, refund me now."},
+      [{"id": "refund", "question": "Does the customer want a refund?", "options": ["no", "yes"]}]))
 ```
-
-## Architecture
-- **Text Backbone:** Bidirectional ModernBERT encoder (423M parameters).
-- **Vision Pathway:** Frozen SigLIP vision tower (`google/siglip-base-patch16-224`) with patch projection.
-- **Adaptive Head Budget:** Dynamic expansion (10k + 48 tokens, capped at 1024) preventing label truncation.
-- **License:** Apache-2.0.

alpha_config.json

@@ -0,0 +1,10 @@
+{
+ "backbone": "distilroberta-base",
+ "d_head": 256,
+ "slot_layers": 2,
+ "slot_heads": 8,
+ "act": "relu2",
+ "max_len": 512,
+ "field_max": 96,
+ "opt_max": 16
+}

alpha_s1.py

@@ -0,0 +1,239 @@
+"""
+Alpha-S1: a System-One-style model. Unstructured / structured state in, typed probabilistic decisions out.
+
+ * No text generation. Every question carries its own option list; the model can only score those options
+   (masked softmax) -> a type error / out-of-schema answer is impossible by construction.
+ * All questions of a query are answered in ONE parallel forward pass (no sequential decoding).
+ * Options are embedded independently of the state (bi-encoder side) so a schema can be "compiled" once and cached;
+   cardinality is limited only by memory (hundreds of options are cheap).
+ * Slot Reasoner: a small set-transformer (squared-ReLU FFN) lets the questions of one query condition on each other.
+ * Calibration: per-question learned temperature + post-hoc global temperature; outputs probs, confidence, certainty.
+"""
+import math, json, os
+from dataclasses import dataclass, asdict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, AutoConfig, AutoTokenizer
+
+
+@dataclass
+class AlphaConfig:
+    backbone: str = "distilroberta-base"
+    d_head: int = 256
+    slot_layers: int = 2
+    slot_heads: int = 8
+    act: str = "relu2"          # "relu2" (squared ReLU, high-accessibility prior) | "gelu"
+    max_len: int = 512
+    field_max: int = 96
+    opt_max: int = 16
+
+
+class Tk:
+    """Tokenizer wrapper: no truncation args (thread-safe), caches short strings."""
+    def __init__(self, tok): self.tok, self.c = tok, {}
+    def __call__(self, text, maxlen=None):
+        ids = self.c.get(text)
+        if ids is None:
+            ids = self.tok(text, add_special_tokens=False)["input_ids"]
+            if len(text) < 200: self.c[text] = ids
+        return ids[:maxlen] if maxlen else ids
+
+
+def opt_batch(texts, tk, tokz, cfg):
+    cls, sep, pad = tokz.cls_token_id, tokz.sep_token_id, tokz.pad_token_id
+    seqs = [[cls] + tk(" " + t)[:cfg.opt_max] + [sep] for t in texts]
+    L = max(len(s) for s in seqs)
+    ids = torch.full((len(seqs), L), pad, dtype=torch.long); att = torch.zeros((len(seqs), L), dtype=torch.long)
+    for i, s in enumerate(seqs):
+        ids[i, :len(s)] = torch.tensor(s); att[i, :len(s)] = 1
+    return ids, att
+
+
+def build_batch(exs, tk, tokz, cfg):
+    """exs: [{"fields":[(name,text),..], "questions":[{"q":str,"options":[str],"gold":int,"ordinal":bool,"fam":str}]}]"""
+    cls, sep, msk, pad = tokz.cls_token_id, tokz.sep_token_id, tokz.mask_token_id, tokz.pad_token_id
+    seqs, qposs, rows = [], [], []
+    opt_index, opt_texts = {}, []
+    for bi, ex in enumerate(exs):
+        qids = [tk(" " + q["q"]) for q in ex["questions"]]
+        qtot = sum(len(x) + 1 for x in qids) + 3
+        nf = max(len(ex["fields"]), 1)
+        per = max(8, min(cfg.field_max, (cfg.max_len - qtot) // nf - 8))
+        while True:
+            ids = [cls]
+            for name, text in ex["fields"]:
+                ids += tk(f" {name}:") + tk(" " + text, per)
+            ids.append(sep)
+            pos = []
+            for x in qids:
+                ids += x; ids.append(msk); pos.append(len(ids) - 1)
+            ids.append(sep)
+            if len(ids) <= cfg.max_len or per <= 8: break
+            per -= 8
+        ids = ids[:cfg.max_len]; pos = [min(p, cfg.max_len - 1) for p in pos]
+        seqs.append(ids); qposs.append(pos)
+        for q in ex["questions"]:
+            idxs = []
+            for o in q["options"]:
+                if o not in opt_index:
+                    opt_index[o] = len(opt_texts); opt_texts.append(o)
+                idxs.append(opt_index[o])
+            rows.append((idxs, q.get("gold", 0), bool(q.get("ordinal", False)), q.get("fam", ""), bi))
+    B, T = len(seqs), max(len(s) for s in seqs)
+    ids = torch.full((B, T), pad, dtype=torch.long); att = torch.zeros((B, T), dtype=torch.long)
+    kmax = max(len(p) for p in qposs)
+    qpos = torch.zeros((B, kmax), dtype=torch.long); qvalid = torch.zeros((B, kmax), dtype=torch.bool)
+    for i, (s, p) in enumerate(zip(seqs, qposs)):
+        ids[i, :len(s)] = torch.tensor(s); att[i, :len(s)] = 1
+        qpos[i, :len(p)] = torch.tensor(p); qvalid[i, :len(p)] = True
+    nQ, nO = len(rows), len(opt_texts)
+    allowed = torch.zeros((nQ, nO), dtype=torch.bool); rank = torch.zeros((nQ, nO))
+    gold = torch.zeros(nQ, dtype=torch.long); ordn = torch.zeros(nQ, dtype=torch.bool)
+    for r, (idxs, g, od, fam, bi) in enumerate(rows):
+        allowed[r, idxs] = True
+        rank[r, idxs] = torch.arange(len(idxs)).float() / max(len(idxs) - 1, 1)
+        gold[r] = idxs[g]; ordn[r] = od
+    oi, oa = opt_batch(opt_texts, tk, tokz, cfg)
+    return dict(ids=ids, att=att, qpos=qpos, qvalid=qvalid, opt_ids=oi, opt_att=oa, allowed=allowed, rank=rank,
+                gold=gold, ordinal=ordn, opt_texts=opt_texts, opt_lists=[r[0] for r in rows],
+                gold_pos=[r[1] for r in rows], fams=[r[3] for r in rows], q_ex=[r[4] for r in rows])
+
+
+def to_device(b, dev):
+    return {k: (v.to(dev, non_blocking=True) if torch.is_tensor(v) else v) for k, v in b.items()}
+
+
+class EncWrap(nn.Module):
+    """Backbone call with autocast INSIDE forward (required for nn.DataParallel replicas)."""
+    def __init__(self, m): super().__init__(); self.m = m
+    def forward(self, ids, att):
+        with torch.autocast("cuda", dtype=torch.float16, enabled=ids.is_cuda):
+            return self.m(input_ids=ids, attention_mask=att).last_hidden_state
+
+
+class SlotLayer(nn.Module):
+    """Pre-LN transformer layer over the query slots of one example; squared-ReLU FFN; zero-init residual branches
+    (starts as identity -> the generalizing 'read straight from the encoder' solution is the easiest to reach)."""
+    def __init__(self, d, heads, act):
+        super().__init__()
+        self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.att = nn.MultiheadAttention(d, heads, batch_first=True)
+        self.f1, self.f2, self.act = nn.Linear(d, 4 * d), nn.Linear(4 * d, d), act
+        nn.init.zeros_(self.att.out_proj.weight); nn.init.zeros_(self.att.out_proj.bias)
+        nn.init.zeros_(self.f2.weight); nn.init.zeros_(self.f2.bias)
+    def forward(self, x, pad):
+        y = self.n1(x)
+        a, _ = self.att(y, y, y, key_padding_mask=pad, need_weights=False)
+        x = x + a
+        h = self.f1(self.n2(x))
+        h = F.relu(h) ** 2 if self.act == "relu2" else F.gelu(h)
+        return x + self.f2(h)
+
+
+class Alpha(nn.Module):
+    def __init__(self, cfg=None, tokenizer=None, pretrained=True, backbone_module=None):
+        super().__init__()
+        self.cfg = cfg or AlphaConfig()
+        self.tokz = tokenizer or AutoTokenizer.from_pretrained(self.cfg.backbone)
+        self.tk = Tk(self.tokz)
+        if backbone_module is not None: self.backbone = backbone_module
+        elif pretrained:
+            try: self.backbone = AutoModel.from_pretrained(self.cfg.backbone, add_pooling_layer=False)
+            except TypeError: self.backbone = AutoModel.from_pretrained(self.cfg.backbone)
+        else:
+            c = AutoConfig.from_pretrained(self.cfg.backbone)
+            try: self.backbone = AutoModel.from_config(c, add_pooling_layer=False)
+            except TypeError: self.backbone = AutoModel.from_config(c)
+        d = self.backbone.config.hidden_size
+        self.ln = nn.LayerNorm(d)
+        self.slots = nn.ModuleList([SlotLayer(d, self.cfg.slot_heads, self.cfg.act) for _ in range(self.cfg.slot_layers)])
+        self.q_proj, self.o_proj = nn.Linear(d, self.cfg.d_head), nn.Linear(d, self.cfg.d_head)
+        self.temp = nn.Linear(d, 1)                      # per-question (heteroscedastic) temperature
+        nn.init.zeros_(self.temp.weight); nn.init.constant_(self.temp.bias, 0.11)
+        self.log_scale = nn.Parameter(torch.tensor(math.log(10.0)))
+        self.register_buffer("post_T", torch.tensor(1.0))   # post-hoc global temperature (fit on held-out calibration set)
+        object.__setattr__(self, "_enc", EncWrap(self.backbone))   # not registered: shares backbone params
+        object.__setattr__(self, "_dp", None)
+        self._opt_cache = {}
+
+    # ---- plumbing ----
+    def use_data_parallel(self):
+        if torch.cuda.device_count() > 1:
+            object.__setattr__(self, "_dp", nn.DataParallel(EncWrap(self.backbone)))
+        return self
+
+    def train(self, mode=True):
+        self._opt_cache = {}
+        return super().train(mode)
+
+    def new_params(self):
+        bb = {id(p) for p in self.backbone.parameters()}
+        return [p for p in self.parameters() if id(p) not in bb]
+
+    def encode(self, ids, att):
+        enc = self._dp if (self._dp is not None and ids.size(0) >= 8) else self._enc
+        return enc(ids, att)
+
+    def embed_options(self, opt_ids, opt_att):
+        H = self.encode(opt_ids, opt_att).float(); m = opt_att.unsqueeze(-1).float()
+        return (H * m).sum(1) / m.sum(1).clamp(min=1)
+
+    # ---- core: all questions, all options, one pass ----
+    def forward(self, b, opt_emb=None):
+        H = self.encode(b["ids"], b["att"])
+        B = H.size(0)
+        hq = H[torch.arange(B, device=H.device)[:, None], b["qpos"]].float()
+        hq = self.ln(hq)
+        pad = ~b["qvalid"]
+        for layer in self.slots: hq = layer(hq, pad)
+        hq = hq[b["qvalid"]]                                       # [nQ, d] row-major == question order in batch
+        if opt_emb is None: opt_emb = self.embed_options(b["opt_ids"], b["opt_att"])
+        q = F.normalize(self.q_proj(hq).float(), dim=-1); o = F.normalize(self.o_proj(opt_emb).float(), dim=-1)
+        T = F.softplus(self.temp(hq).float()) + 0.25
+        logits = (q @ o.T) * self.log_scale.exp() / T
+        return logits.masked_fill(~b["allowed"], -1e4)             # type safety: only declared options exist
+
+    # ---- public API ----
+    @torch.no_grad()
+    def predict(self, state, questions):
+        """state: {field: text}; questions: [{"id","question","options":[...],"ordinal":bool}] -> {id: {...}}"""
+        dev = next(self.parameters()).device
+        if self.training: self.eval()                       # (train() clears the option cache)
+        ex = dict(fields=list(state.items()),
+                  questions=[dict(q=q["question"], options=list(q["options"]), gold=0, ordinal=q.get("ordinal", False)) for q in questions])
+        b = to_device(build_batch([ex], self.tk, self.tokz, self.cfg), dev)
+        miss = [t for t in b["opt_texts"] if t not in self._opt_cache]
+        if miss:                                                   # "schema compilation": option embeddings cached
+            oi, oa = opt_batch(miss, self.tk, self.tokz, self.cfg)
+            emb = self.embed_options(oi.to(dev), oa.to(dev))
+            for t, e in zip(miss, emb): self._opt_cache[t] = e
+        opt_emb = torch.stack([self._opt_cache[t] for t in b["opt_texts"]])
+        logits = self.forward(b, opt_emb) / self.post_T
+        out = {}
+        for r, q in enumerate(questions):
+            idx = b["opt_lists"][r]
+            p = F.softmax(logits[r, idx], -1).cpu()
+            K = len(idx); j = int(p.argmax())
+            ent = float(-(p * p.clamp_min(1e-9).log()).sum())
+            out[q["id"]] = dict(answer=q["options"][j], probs={o: float(x) for o, x in zip(q["options"], p)},
+                                confidence=float(p[j]), certainty=(1 - ent / math.log(K)) if K > 1 else 1.0)
+        return out
+
+    def save_pretrained(self, path):
+        os.makedirs(path, exist_ok=True)
+        torch.save({k: v.detach().cpu() for k, v in self.state_dict().items()}, f"{path}/pytorch_model.bin")
+        json.dump(asdict(self.cfg), open(f"{path}/alpha_config.json", "w"), indent=1)
+        self.tokz.save_pretrained(path)
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo, device=None):
+        if not os.path.isdir(path_or_repo):
+            from huggingface_hub import snapshot_download
+            path_or_repo = snapshot_download(path_or_repo)
+        cfg = AlphaConfig(**json.load(open(f"{path_or_repo}/alpha_config.json")))
+        m = cls(cfg, AutoTokenizer.from_pretrained(path_or_repo), pretrained=False)
+        missing, unexpected = m.load_state_dict(torch.load(f"{path_or_repo}/pytorch_model.bin", map_location="cpu"), strict=False)
+        missing = [k for k in missing if "pooler" not in k]
+        assert not missing, f"missing keys: {missing[:5]}"
+        return m.to(device) if device else m

alpha_train.py

@@ -0,0 +1,645 @@
+"""
+Alpha training / benchmark / publish pipeline  (Kaggle, 2x T4).
+
+Framework rules from the spec are used as *design priors and instruments*, not as proven facts:
+  * Accessibility A(theta,F): squared-ReLU + zero-init Slot Reasoner, pretrained encoder, schema-compositional heads
+    (the general function is the easiest one to reach), many paraphrased schemas (high information rate).
+  * Regularisation shapes the manifold: very different weight decay for pretrained backbone vs new modules, adapted online
+    by the Discovery Controller from rho = d(val)/d(train) (proxy for P(B_G)/P(B_M)).
+  * Functional memory M_F: function-space signatures of past EMA checkpoints on diagnostic set E_F, used as a KL anchor
+    (self-distillation to earlier functions); L_FRR measured mid-training scales the anchor weight.
+  * FRR: perturb the discovered solution (basin test), retrain, time re-discovery -> FRR = T_first / T_re.
+  * Meta-grokking: held-out task families never trained on; adaptation speed measured at checkpoints (MGR).
+  * DMR: T_discover / T_memorize from train/val trajectories ( ~1 => no delayed generalisation, i.e. 'groks from the start').
+"""
+import os, re, json, math, time, random, shutil
+import numpy as np
+import torch
+import torch.nn.functional as F
+from alpha_s1 import Alpha, AlphaConfig, build_batch, to_device
+
+dev = "cuda:0" if torch.cuda.is_available() else "cpu"
+ac = lambda: torch.autocast("cuda", dtype=torch.float16, enabled=torch.cuda.is_available())
+
+
+class C:
+    backbone = "distilroberta-base"       # try "roberta-base" for more accuracy (slower)
+    repo_id = "SofiTesfay2010/Alpha"
+    seed = 0
+    steps, bs = 6000, 32
+    lr_bb, lr_new = 5e-5, 5e-4
+    wd_bb, wd_new, wd_min, wd_max = 0.01, 0.5, 0.05, 2.0
+    warmup, clip = 300, 1.0
+    ls, brier_w, ord_w = 0.02, 0.3, 1.0
+    ema = 0.999
+    eval_every, mem_every, anchor_every, lam_mem = 250, 1000, 4, 0.3
+    probe_every, probe_steps, probe_bs, probe_test_n = 1000, 40, 16, 200
+    probe_at = (0, 5, 10, 20, 40)
+    fast_n = 48
+    n_cal, n_test = 150, 500
+    sigmas = [0.5, 1, 2, 4, 8]
+    noise_unit, frr_max_steps, frr_every = 0.02, 300, 25
+    llm_baseline = "Qwen/Qwen2.5-0.5B-Instruct"
+
+
+# ----------------------------------------------------------------------------------------------- data
+def S(paths, fields, fn, qs, names=None, alts=None, tr="train", va="validation", ordinal=False, held=False, ntr=30000, camel=False):
+    return dict(paths=paths, fields=fields, fn=fn, qs=qs, names=names, alts=alts, tr=tr, va=va, ordinal=ordinal, held=held, ntr=ntr, camel=camel)
+
+SPECS = dict(
+    sst2=S([("nyu-mll/glue", "sst2")], ["sentence"], [["review"], ["comment"], ["snippet"]],
+           ["What is the sentiment of the {f}?", "Is the {f} positive or negative?", "Judge the attitude expressed in the {f}."],
+           names=["negative", "positive"], alts=[["negative", "positive"], ["bad", "good"], ["unfavorable", "favorable"]]),
+    imdb=S([("stanfordnlp/imdb",)], ["text"], [["movie_review"], ["critique"]],
+           ["Is the {f} favorable or unfavorable?", "What verdict does the {f} give?", "Does the author of the {f} like the film?"],
+           names=["negative", "positive"], alts=[["negative", "positive"], ["thumbs down", "thumbs up"], ["dislikes it", "likes it"]], va="test", ntr=20000),
+    ag_news=S([("fancyzhx/ag_news",)], ["text"], [["headline"], ["article"], ["news_item"]],
+              ["Which section would the {f} appear in?", "What topic is the {f} about?", "Categorize the {f}."],
+              names=["world news", "sports", "business", "science and technology"],
+              alts=[["world news", "sports", "business", "science and technology"], ["world", "sports", "finance", "tech"], ["international affairs", "athletics", "economy", "technology"]], va="test"),
+    emotion=S([("dair-ai/emotion",)], ["text"], [["message"], ["post"], ["tweet"]],
+              ["Which emotion does the {f} convey?", "What is the author of the {f} feeling?", "Name the dominant emotion in the {f}."],
+              names=["sadness", "joy", "love", "anger", "fear", "surprise"],
+              alts=[["sadness", "joy", "love", "anger", "fear", "surprise"], ["sad", "happy", "affection", "angry", "afraid", "astonished"], ["grief", "delight", "adoration", "rage", "dread", "shock"]]),
+    banking77=S([("PolyAI/banking77",), ("mteb/banking77",)], ["text"], [["customer_message"], ["support_query"]],
+                ["What is the customer asking about in the {f}?", "Which intent does the {f} express?", "Classify the banking request in the {f}."], va="test"),
+    yelp=S([("Yelp/yelp_review_full",)], ["text"], [["review"], ["feedback"]],
+           ["How many stars does the {f} deserve?", "Rate the {f}.", "What score would the author of the {f} give?"],
+           names=["1 star", "2 stars", "3 stars", "4 stars", "5 stars"],
+           alts=[["1 star", "2 stars", "3 stars", "4 stars", "5 stars"], ["terrible", "poor", "average", "good", "excellent"], ["very negative", "negative", "neutral", "positive", "very positive"]],
+           va="test", ordinal=True),
+    snli=S([("stanfordnlp/snli",)], ["premise", "hypothesis"], [["premise", "hypothesis"], ["context", "claim"]],
+           ["What is the logical relationship between {f}?", "Does the second follow from the first, contradict it, or neither? ({f})", "Relation between {f}:"],
+           names=["entailment", "neutral", "contradiction"],
+           alts=[["entailment", "neutral", "contradiction"], ["supported", "undetermined", "contradicted"], ["follows", "unrelated", "conflicts"]], ntr=60000),
+    sms=S([("ucirvine/sms_spam",)], ["sms"], [["message"], ["text_message"]],
+          ["Is the {f} spam?", "Classify the {f}.", "Is the {f} an unwanted advertisement or a normal message?"],
+          names=["legitimate", "spam"], alts=[["legitimate", "spam"], ["ham", "spam"], ["normal", "unwanted"]], va=None),
+    news20=S([("SetFit/20_newsgroups",)], ["text"], [["post"], ["usenet_message"]],
+             ["Which newsgroup does the {f} belong to?", "What topic does the {f} discuss?", "Categorize the {f} by topic."], va="test"),
+    mrpc=S([("nyu-mll/glue", "mrpc")], ["sentence1", "sentence2"], [["sentence_a", "sentence_b"], ["statement_1", "statement_2"]],
+           ["Do {f} mean the same thing?", "Are {f} paraphrases of each other?", "Compare the meaning of {f}."],
+           names=["different meaning", "same meaning"], alts=[["different meaning", "same meaning"], ["not paraphrase", "paraphrase"], ["unrelated", "equivalent"]]),
+    # ---- held-out families: NEVER trained on in the main run; used for zero-shot + meta-grokking probes ----
+    dbpedia=S([("fancyzhx/dbpedia_14",)], ["content"], [["entry"], ["wiki_summary"]],
+              ["What kind of entity does the {f} describe?", "Classify the subject of the {f}.", "Which category fits the {f}?"], va="test", held=True, ntr=20000, camel=True),
+    offensive=S([("cardiffnlp/tweet_eval", "offensive")], ["text"], [["tweet"], ["post"]],
+                ["Is the {f} offensive?", "Does the {f} contain abuse or insults?", "Moderate the {f}."],
+                names=["not offensive", "offensive"], alts=[["not offensive", "offensive"], ["civil", "abusive"], ["acceptable", "insulting"]], held=True),
+    cola=S([("nyu-mll/glue", "cola")], ["sentence"], [["sentence"], ["utterance"]],
+           ["Is the {f} grammatical?", "Is the {f} well-formed English?", "Judge the grammar of the {f}."],
+           names=["ungrammatical", "grammatical"], alts=[["ungrammatical", "grammatical"], ["incorrect", "correct"], ["unacceptable", "acceptable"]], held=True),
+)
+
+DATA, ST = {}, {}
+
+
+def _clean(n, camel):
+    n = n.replace("_", " ").replace(".", " ")
+    if camel: n = re.sub(r"(?<!^)(?=[A-Z])", " ", n)
+    return n.lower().strip()
+
+
+def load_all():
+    from datasets import load_dataset
+    for name, sp in SPECS.items():
+        for p in sp["paths"]:
+            try:
+                if sp["va"]:
+                    dtr = load_dataset(*p, split=sp["tr"]).shuffle(seed=0); dva = load_dataset(*p, split=sp["va"]).shuffle(seed=0)
+                else:
+                    full = load_dataset(*p, split="train").shuffle(seed=0); cut = int(0.9 * len(full))
+                    dtr, dva = full.select(range(cut)), full.select(range(cut, len(full)))
+                dtr = dtr.select(range(min(len(dtr), sp["ntr"]))); dva = dva.select(range(min(len(dva), 1500)))
+                names = sp["names"]
+                if names is None:
+                    feat = dtr.features["label"]
+                    if hasattr(feat, "names"): names = list(feat.names)
+                    else:
+                        m = {int(l): t for l, t in zip(dtr["label"], dtr["label_text"])}; names = [m[i] for i in range(max(m) + 1)]
+                    names = [_clean(n, sp["camel"]) for n in names]
+                alts = sp["alts"] or [names]
+
+                def rows(d):
+                    cols = [d[f] for f in sp["fields"]]; lab = d["label"]
+                    return [(tuple(str(c[i])[:1200] for c in cols), int(lab[i])) for i in range(len(lab)) if int(lab[i]) >= 0]
+                R = dict(train=rows(dtr), val=rows(dva))
+                assert len(R["val"]) >= 60 and len(R["train"]) >= 200
+                DATA[name] = dict(spec=sp, names=names, alts=alts, rows=R)
+                print(f"[data] {name:10s} classes={len(names):3d} train={len(R['train']):6d} val={len(R['val']):5d} {'(held-out)' if sp['held'] else ''}")
+                break
+            except Exception as e:
+                print(f"[skip] {name} via {p}: {type(e).__name__}: {str(e)[:100]}")
+    ST["train_fams"] = sorted(f for f in DATA if not DATA[f]["spec"]["held"])
+    ST["held_fams"] = sorted(f for f in DATA if DATA[f]["spec"]["held"])
+    assert len(ST["train_fams"]) >= 4, "too few training families loaded (is internet enabled?)"
+
+
+def make_example(rng, fams, split="train", wording="train", idx=None, lo=0):
+    fields, qs, used = [], [], set()
+    for fam in fams:
+        D = DATA[fam]; sp = D["spec"]; rows = D["rows"][split]
+        if idx is not None: texts, label = rows[idx % len(rows)]
+        else: texts, label = rows[rng.randrange(lo if len(rows) > lo else 0, len(rows))]
+        nm = []
+        for n0 in rng.choice(sp["fn"]):
+            n1, c = n0, 2
+            while n1 in used: n1 = f"{n0}_{c}"; c += 1
+            used.add(n1); nm.append(n1)
+        fields += list(zip(nm, texts))
+        if wording == "unseen": q, alt = sp["qs"][2], D["alts"][-1]
+        elif wording == "any": q, alt = rng.choice(sp["qs"]), rng.choice(D["alts"])
+        else: q, alt = rng.choice(sp["qs"][:2]), rng.choice(D["alts"][:-1] or D["alts"])
+        qs.append(dict(q=q.format(f=" and ".join(nm)), options=list(alt), gold=label, ordinal=sp["ordinal"], fam=fam))
+    rng.shuffle(fields); rng.shuffle(qs)
+    return dict(fields=fields, questions=qs)
+
+
+def sample_exs(rng, n):
+    fams, w = ST["train_fams"], ST["w"]; out = []
+    for _ in range(n):
+        k = min(rng.choices([1, 2, 3, 4], [.3, .3, .25, .15])[0], len(fams)); ch = []
+        while len(ch) < k:
+            f = rng.choices(fams, w)[0]
+            if f not in ch: ch.append(f)
+        out.append(make_example(rng, ch))
+    return out
+
+
+def prep(exs, bs=48):
+    m = ST["model"]
+    return [to_device(build_batch(exs[i:i + bs], m.tk, m.tokz, m.cfg), dev) for i in range(0, len(exs), bs)]
+
+
+def val_examples(fam, lo, hi, wording):
+    n = len(DATA[fam]["rows"]["val"])
+    return [make_example(random.Random(1000 + i), [fam], "val", wording, idx=i) for i in range(lo, min(hi, n))]
+
+
+# ----------------------------------------------------------------------------------------------- metrics
+def ece_fn(conf, corr, nb=15):
+    conf, corr = np.asarray(conf, float), np.asarray(corr, float); e = 0.0
+    edges = np.linspace(0, 1, nb + 1)
+    for lo, hi in zip(edges[:-1], edges[1:]):
+        m = (conf > lo) & (conf <= hi)
+        if m.any(): e += m.mean() * abs(corr[m].mean() - conf[m].mean())
+    return float(e)
+
+
+@torch.no_grad()
+def predict_res(model, batches):
+    model.eval(); out = []
+    for b in batches:
+        with ac(): lg = model(b)
+        lg = lg.float().cpu()
+        for r, idx in enumerate(b["opt_lists"]):
+            out.append((b["fams"][r], lg[r, idx].numpy(), b["gold_pos"][r]))
+    return out
+
+
+def _calc(items):
+    acc = nll = br = 0.0; conf, corr = [], []
+    for lg, g in items:
+        z = lg - lg.max(); lp = z - np.log(np.exp(z).sum()); p = np.exp(lp)
+        oh = np.zeros_like(p); oh[g] = 1
+        acc += float(p.argmax() == g); nll += -lp[g]; br += ((p - oh) ** 2).sum()
+        conf.append(float(p.max())); corr.append(float(p.argmax() == g))
+    n = len(items)
+    return dict(acc=acc / n, nll=nll / n, brier=br / n, ece=ece_fn(conf, corr), n=n, conf=conf, corr=corr)
+
+
+def stat(res, T=1.0):
+    per = {}
+    for fam, lg, g in res: per.setdefault(fam, []).append((lg / T, g))
+    P = {f: _calc(v) for f, v in per.items()}
+    allitems = [x for v in per.values() for x in v]; micro = _calc(allitems)
+    mac = lambda k: float(np.mean([P[f][k] for f in P]))
+    return dict(per=P, score=mac("acc"), nll=mac("nll"), brier=mac("brier"), ece=micro["ece"], micro=micro,
+                per_acc={f: P[f]["acc"] for f in P})
+
+
+def fast_eval(model, which):
+    return stat(predict_res(model, ST["fast_" + which]), float(model.post_T))
+
+
+# ----------------------------------------------------------------------------------------------- training parts
+def loss_fn(model, b):
+    lg = model(b).float(); lp = F.log_softmax(lg, -1); al = b["allowed"]
+    ar = torch.arange(lg.size(0), device=lg.device)
+    nll = -lp[ar, b["gold"]]
+    uni = -(lp.masked_fill(~al, 0.0)).sum(1) / al.sum(1)
+    p = lp.exp(); oh = F.one_hot(b["gold"], lg.size(1)).float()
+    brier = ((p - oh) ** 2).sum(1)
+    rg = b["rank"][ar, b["gold"]]
+    ordl = (p * (b["rank"] - rg[:, None]) ** 2 * al).sum(1)
+    L = (1 - C.ls) * nll + C.ls * uni + C.brier_w * brier + C.ord_w * b["ordinal"].float() * ordl
+    return L.mean(), float(nll.mean())
+
+
+def make_scaler():
+    try: return torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())
+    except Exception: return torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
+
+
+def make_opt(model, lr_bb, lr_new, wd_new):
+    def split(ps): return [p for p in ps if p.ndim >= 2], [p for p in ps if p.ndim < 2]
+    bd, bn = split(list(model.backbone.parameters())); nd, nn_ = split(model.new_params())
+    gs = [dict(params=bd, base_lr=lr_bb, lr=lr_bb, weight_decay=C.wd_bb, tag="bb"), dict(params=bn, base_lr=lr_bb, lr=lr_bb, weight_decay=0.0, tag="bb0"),
+          dict(params=nd, base_lr=lr_new, lr=lr_new, weight_decay=wd_new, tag="new"), dict(params=nn_, base_lr=lr_new, lr=lr_new, weight_decay=0.0, tag="new0")]
+    return torch.optim.AdamW([g for g in gs if g["params"]], betas=(0.9, 0.98), eps=1e-6)
+
+
+def sgd_step(model, opt, scaler, b, mem=None, lam=0.0, anchor=False):
+    model.train()
+    with ac():
+        loss, nll = loss_fn(model, b)
+        if anchor and mem is not None and mem.entries: loss = loss + lam * mem.anchor(model)
+    opt.zero_grad(set_to_none=True); scaler.scale(loss).backward(); scaler.unscale_(opt)
+    gn = torch.nn.utils.clip_grad_norm_(model.parameters(), C.clip); scaler.step(opt); scaler.update()
+    return nll, float(gn)
+
+
+def clone_state(m): return {k: v.detach().cpu().clone() for k, v in m.state_dict().items()}
+
+
+class EMA:
+    def __init__(s, model, decay):
+        s.d, s.bk = decay, None
+        s.sh = {k: v.detach().clone().float() for k, v in model.state_dict().items() if v.dtype.is_floating_point}
+    @torch.no_grad()
+    def update(s, model, step):
+        d = min(s.d, (1 + step) / (10 + step))
+        for k, v in model.state_dict().items():
+            if k in s.sh: s.sh[k].mul_(d).add_(v.detach().float(), alpha=1 - d)
+    @torch.no_grad()
+    def apply(s, model):
+        sd = model.state_dict(); s.bk = {k: sd[k].detach().clone() for k in s.sh}
+        model.load_state_dict({k: s.sh[k].to(sd[k].dtype) for k in s.sh}, strict=False)
+    @torch.no_grad()
+    def restore(s, model): model.load_state_dict(s.bk, strict=False)
+
+
+class FMem:
+    """Functional memory M_F: entries = {step, signature Sigma_F (dense logits on diagnostic set E_F), val score}."""
+    def __init__(s, diag_batch): s.b, s.entries = diag_batch, []
+    @torch.no_grad()
+    def add(s, model, step, score):
+        model.eval()
+        with ac(): lg = model(s.b).float()
+        s.entries.append(dict(step=step, sig=lg.cpu().half(), score=score))
+    def anchor(s, model):
+        e = random.choice(s.entries); old = e["sig"].to(s.b["allowed"].device).float()
+        lp = F.log_softmax(model(s.b).float(), -1); p = F.softmax(old, -1); lpo = F.log_softmax(old, -1)
+        return ((p * (lpo - lp)).masked_fill(~s.b["allowed"], 0.0)).sum(1).mean()
+
+
+class Controller:
+    """Discovery controller: rho = d(val nll)/d(train nll) ~ share of fitting progress that generalises (proxy of P(B_G)/P(B_M))."""
+    def __init__(s): s.wd, s.rho, s.prev, s.LFRR = C.wd_new, None, None, 0.0
+    def update(s, rec):
+        if s.prev is not None:
+            dtr, dv = s.prev["tr_nll"] - rec["tr_nll"], s.prev["nll"] - rec["nll"]
+            if dtr > 1e-3:
+                r = float(np.clip(dv / dtr, -1, 2)); s.rho = r if s.rho is None else 0.5 * s.rho + 0.5 * r
+            if s.rho is not None:
+                if s.rho < 0.3: s.wd = min(C.wd_max, s.wd * 1.25)      # memorising -> shape manifold harder
+                elif s.rho > 0.8: s.wd = max(C.wd_min, s.wd * 0.9)     # generalising -> relax
+        s.prev = rec
+
+
+def perturb(model, sigma):
+    with torch.no_grad():
+        for p in model.parameters():
+            if p.ndim >= 2: p.add_(torch.randn_like(p) * p.std() * sigma * C.noise_unit)
+
+
+def frr_test(model, hist, final_score, sigmas):
+    """FRR = T_first / T_re. T_first: first logged step the main run reached tau. T_re: steps to re-reach tau after perturbing."""
+    snap = clone_state(model); tau = final_score - 0.01
+    T_first = next((h["step"] for h in hist if h["score"] >= tau), hist[-1]["step"]); res = {}
+    for sg in sigmas:
+        model.load_state_dict(snap); torch.manual_seed(int(sg * 100)); perturb(model, sg)
+        v0 = fast_eval(model, "val")["score"]; trivial = v0 >= tau; T_re, hit = 1, trivial
+        if not trivial:
+            opt, scaler, rng = make_opt(model, C.lr_bb, C.lr_new, C.wd_new), make_scaler(), random.Random(999 + int(sg * 10)); T_re = C.frr_max_steps
+            for s in range(1, C.frr_max_steps + 1):
+                for g in opt.param_groups: g["lr"] = g["base_lr"] * min(1.0, s / 5)
+                sgd_step(model, opt, scaler, to_device(build_batch(sample_exs(rng, C.bs), model.tk, model.tokz, model.cfg), dev))
+                if s % C.frr_every == 0 and fast_eval(model, "val")["score"] >= tau: T_re, hit = s, True; break
+        res[sg] = dict(v0=v0, tau=tau, T_first=T_first, T_re=T_re, FRR=T_first / max(T_re, 1), L_FRR=max(T_re, 1) / max(T_first, 1),
+                       censored=(not hit), trivial=trivial)
+        print(f"   [FRR] sigma={sg}: perturbed score {v0:.3f} (tau {tau:.3f}) T_first={T_first} T_re={T_re}{'+' if not hit else ''} FRR={res[sg]['FRR']:.1f}")
+    model.load_state_dict(snap); return res
+
+
+def adapt_curve(model, fam, eval_batches, steps, at):
+    """Few-shot adaptation of a snapshot to a held-out family; returns {step: (acc, nll)}. Model is restored afterwards."""
+    snap = clone_state(model); opt, scaler, rng = make_opt(model, C.lr_bb, C.lr_new, 0.0), make_scaler(), random.Random(777); curve = {}
+    for s in range(steps + 1):
+        if s in at:
+            r = stat(predict_res(model, eval_batches), float(model.post_T)); curve[s] = (r["score"], r["nll"])
+        if s == steps: break
+        for g in opt.param_groups: g["lr"] = g["base_lr"] * min(1.0, (s + 1) / 3)
+        b = to_device(build_batch([make_example(rng, [fam], "train", "any") for _ in range(C.probe_bs)], model.tk, model.tokz, model.cfg), dev)
+        sgd_step(model, opt, scaler, b)
+    model.load_state_dict(snap); return curve
+
+
+def meta_probe(model, step):
+    out = [dict(step=step, fam=f, curve=adapt_curve(model, f, ST["held_test"][f], C.probe_steps, C.probe_at)) for f in ST["held_fams"]]
+    msg = " | ".join(f"{o['fam']}: zs {o['curve'][0][0]:.2f} -> @10 {o['curve'][10][0]:.2f} -> @{C.probe_steps} {o['curve'][C.probe_steps][0]:.2f}" for o in out)
+    print(f"   [meta-probe @ {step}] {msg}"); return out
+
+
+# ----------------------------------------------------------------------------------------------- setup + train
+def setup():
+    random.seed(C.seed); np.random.seed(C.seed); torch.manual_seed(C.seed)
+    load_all()
+    m = Alpha(AlphaConfig(backbone=C.backbone)).to(dev).use_data_parallel()
+    ST["model"] = m; ST["w"] = [math.sqrt(len(DATA[f]["rows"]["train"])) for f in ST["train_fams"]]
+    tf, hf = ST["train_fams"], ST["held_fams"]
+    fast = lambda split, lo, n: [make_example(random.Random(i * 31 + j), [f], split, "train", idx=i) for j, f in enumerate(tf) for i in range(lo, lo + n)]
+    ST["fast_val"], ST["fast_train"] = prep(fast("val", 0, C.fast_n)), prep(fast("train", 0, C.fast_n))
+    ST["cal"] = prep([e for f in tf for e in val_examples(f, 0, C.n_cal, "train")])
+    ST["held_test"] = {f: prep(val_examples(f, 0, C.probe_test_n, "unseen")) for f in hf}
+    rng = random.Random(4242); ST["diag"] = prep([make_example(rng, rng.sample(tf, min(3, len(tf)))) for _ in range(16)], bs=16)[0]
+    n_all = sum(p.numel() for p in m.parameters()); print(f"params: {n_all/1e6:.1f}M | GPUs: {torch.cuda.device_count()} | train fams {tf} | held-out {hf}")
+
+
+def train():
+    m, hist = ST["model"], []; rng = random.Random(C.seed + 1)
+    ctl, mem, ema = Controller(), FMem(ST["diag"]), EMA(m, C.ema)
+    opt, scaler = make_opt(m, C.lr_bb, C.lr_new, ctl.wd), make_scaler()
+    ST.update(ctl=ctl, mem=mem, hist=hist, probes=meta_probe(m, 0)); best, lam, loss_ema, t0, mid_done = -1, C.lam_mem, None, time.time(), False
+    for step in range(1, C.steps + 1):
+        prog = max(0.0, (step - C.warmup) / max(1, C.steps - C.warmup))
+        mult = min(1.0, step / C.warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * prog)))
+        for g in opt.param_groups:
+            g["lr"] = g["base_lr"] * mult
+            if g["tag"] == "new": g["weight_decay"] = ctl.wd
+        b = to_device(build_batch(sample_exs(rng, C.bs), m.tk, m.tokz, m.cfg), dev)
+        nll, gn = sgd_step(m, opt, scaler, b, mem, lam, anchor=(step % C.anchor_every == 0))
+        ema.update(m, step); loss_ema = nll if loss_ema is None else 0.98 * loss_ema + 0.02 * nll
+        if step % C.eval_every == 0 or step == C.steps:
+            ema.apply(m); v, t = fast_eval(m, "val"), fast_eval(m, "train")
+            rec = dict(step=step, score=v["score"], nll=v["nll"], ece=v["ece"], tr_score=t["score"], tr_nll=t["nll"], fam=v["per_acc"],
+                       wd=ctl.wd, lam=lam, gn=gn, loss=loss_ema, rho=ctl.rho, lr=opt.param_groups[0]["lr"])
+            hist.append(rec)
+            if v["score"] > best: best = v["score"]; ST["best_state"] = clone_state(m); ST["best_step"] = step
+            ctl.update(rec)
+            if step % C.mem_every == 0: mem.add(m, step, v["score"])
+            if step % C.probe_every == 0: ST["probes"] += meta_probe(m, step)
+            if not mid_done and step >= C.steps // 2:
+                mid_done = True; r = frr_test(m, hist, v["score"], [1])[1]; ctl.LFRR = r["L_FRR"]; lam = C.lam_mem * (1 + min(ctl.LFRR, 3.0))
+                print(f"   [L_FRR={ctl.LFRR:.3f}] -> anchor weight {lam:.3f}")
+            ema.restore(m)
+            print(f"step {step:5d} | val acc {v['score']:.3f} nll {v['nll']:.3f} ece {v['ece']:.3f} | train acc {t['score']:.3f} nll {t['nll']:.3f} | "
+                  f"rho {ctl.rho if ctl.rho is None else round(ctl.rho, 2)} wd {ctl.wd:.3f} lam {lam:.2f} gn {gn:.2f} | {time.time()-t0:.0f}s")
+    m.load_state_dict(ST["best_state"]); print(f"loaded best EMA checkpoint (step {ST['best_step']})")
+
+
+# ----------------------------------------------------------------------------------------------- finalize + benchmark
+def finalize():
+    m, hist = ST["model"], ST["hist"]; m.eval()
+    res = predict_res(m, ST["cal"]); grid = np.arange(0.5, 3.01, 0.02)
+    nlls = [np.mean([-(lg / T - np.logaddexp.reduce(lg / T))[g] for _, lg, g in res]) for T in grid]
+    T = float(grid[int(np.argmin(nlls))]); m.post_T.fill_(T); print(f"post-hoc temperature T={T:.2f}")
+    best = max(h["score"] for h in hist); print("FRR sweep (basin test):"); ST["frr"] = frr_test(m, hist, best, C.sigmas)
+    tau = min(h["nll"] for h in hist) + 0.05
+    Tv = next((h["step"] for h in hist if h["nll"] <= tau), hist[-1]["step"]); Tt = next((h["step"] for h in hist if h["tr_nll"] <= tau), hist[-1]["step"])
+    ST["DMR"] = dict(T_discover=Tv, T_memorize=Tt, DMR=Tv / max(Tt, 1)); print("DMR:", ST["DMR"])
+    # meta-grokking: T_D(n) = adaptation steps to reach 95% of the best adapted accuracy seen over all probes (per family)
+    P, fams, mg = ST["probes"], ST["held_fams"], {}
+    steps_ = sorted({p["step"] for p in P})
+    for f in fams:
+        tau_f = 0.95 * max(p["curve"][C.probe_steps][0] for p in P if p["fam"] == f)
+        for s in steps_:
+            c = next(p for p in P if p["fam"] == f and p["step"] == s)["curve"]
+            td = next((k for k in sorted(c) if c[k][0] >= tau_f), C.probe_steps * 2)
+            mg.setdefault(s, {})[f] = dict(T_D=max(td, 1), zero_shot=c[0][0], acc10=c[10][0], final=c[C.probe_steps][0])
+    base = steps_[0]
+    for s in steps_:
+        g = [mg[base][f]["T_D"] / mg[s][f]["T_D"] for f in fams]; mg[s]["MGR"] = float(np.exp(np.mean(np.log(g))))
+    ST["meta"] = mg; print("meta-grokking MGR by checkpoint:", {s: round(mg[s]["MGR"], 2) for s in steps_})
+
+
+@torch.no_grad()
+def workflow_eval(n=300, k=4):
+    m = ST["model"]; rng = random.Random(2024); exs = []
+    for _ in range(n): exs.append(make_example(rng, rng.sample(ST["train_fams"], min(k, len(ST["train_fams"]))), "val", "train", lo=C.n_cal))
+    m.eval(); qacc, allok, jc, jok = [], [], [], []
+    for i in range(0, n, 32):
+        b = to_device(build_batch(exs[i:i + 32], m.tk, m.tokz, m.cfg), dev)
+        with ac(): lg = m(b).float() / m.post_T
+        per = {}
+        for r, idx in enumerate(b["opt_lists"]):
+            p = F.softmax(lg[r, idx], -1); ok = int(p.argmax()) == b["gold_pos"][r]
+            qacc.append(ok); per.setdefault(b["q_ex"][r], []).append((ok, float(p.max())))
+        for v in per.values():
+            allok.append(all(o for o, _ in v)); jc.append(float(np.prod([c for _, c in v]))); jok.append(float(all(o for o, _ in v)))
+    return dict(q_acc=float(np.mean(qacc)), all_correct=float(np.mean(allok)), joint_ece=ece_fn(jc, jok), k=k, n=n)
+
+
+def speed_bench():
+    m = ST["model"]; m.eval()
+    state = dict(customer_ticket="My card was charged twice for the same order and support has not replied in three days. I want my money back.",
+                 account="plan: pro; tenure_months: 14; open_tickets: 2")
+    qs = [dict(id="refund", question="Does the customer ask for a refund?", options=["no", "yes"]),
+          dict(id="urgency", question="How urgent is the ticket?", options=["low", "medium", "high"], ordinal=True),
+          dict(id="team", question="Which team should handle the ticket?", options=["billing", "technical", "sales", "account", "abuse"])]
+    sync = (lambda: torch.cuda.synchronize()) if torch.cuda.is_available() else (lambda: None)
+    t = time.perf_counter(); m.predict(state, qs); sync(); cold = (time.perf_counter() - t) * 1e3
+    for _ in range(10): m.predict(state, qs)
+    ts = []
+    for _ in range(200):
+        sync(); t = time.perf_counter(); m.predict(state, qs); sync(); ts.append((time.perf_counter() - t) * 1e3)
+    rng = random.Random(1); exs = [make_example(rng, rng.sample(ST["train_fams"], 3)) for _ in range(64)]
+    with torch.no_grad():
+        for _ in range(3): m(to_device(build_batch(exs, m.tk, m.tokz, m.cfg), dev))
+    sync(); t = time.perf_counter()
+    for _ in range(10):
+        b = to_device(build_batch(exs, m.tk, m.tokz, m.cfg), dev)
+        with torch.no_grad(), ac(): m(b)
+    sync(); e2e = 10 * 64 / (time.perf_counter() - t)
+    return dict(latency_ms_p50=float(np.percentile(ts, 50)), latency_ms_p95=float(np.percentile(ts, 95)), cold_first_call_ms=cold,
+                throughput_workflows_per_s_bs64=e2e, questions_per_workflow=3, device=torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu")
+
+
+def llm_baseline(n_per=30, fams=("sst2", "ag_news", "emotion", "snli")):
+    """Generative small-LLM baseline on the SAME k=1 questions (answers parsed from free text) -> accuracy + latency."""
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    fams = [f for f in fams if f in DATA]; name = C.llm_baseline
+    t = AutoTokenizer.from_pretrained(name); lm = AutoModelForCausalLM.from_pretrained(name, torch_dtype=torch.float16).to(dev).eval()
+    m = ST["model"]; m.eval(); lat, ok_l, ok_a = [], [], []
+    for f in fams:
+        for i in range(n_per):
+            ex = make_example(random.Random(i), [f], "val", "train", idx=C.n_cal + i); q = ex["questions"][0]
+            state = "\n".join(f"{k}: {v}" for k, v in ex["fields"])
+            msg = [{"role": "user", "content": f"{state}\n\nQuestion: {q['q']}\nOptions: {', '.join(q['options'])}\nReply with exactly one option and nothing else."}]
+            ids = t.apply_chat_template(msg, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(dev)
+            torch.cuda.synchronize(); t0 = time.perf_counter()
+            with torch.no_grad(): out = lm.generate(**ids, max_new_tokens=12, do_sample=False, pad_token_id=t.eos_token_id)
+            torch.cuda.synchronize(); lat.append((time.perf_counter() - t0) * 1e3)
+            txt = t.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True).strip().lower()
+            cand = [o for o in q["options"] if o.lower() in txt]; pred = max(cand, key=len) if cand else None
+            ok_l.append(pred is not None and pred == q["options"][q["gold"]])
+            b = to_device(build_batch([ex], m.tk, m.tokz, m.cfg), dev)
+            with torch.no_grad(), ac(): lg = m(b)
+            ok_a.append(int(lg[0, b["opt_lists"][0]].argmax()) == b["gold_pos"][0])
+    del lm; torch.cuda.empty_cache()
+    return dict(model=name, n=len(ok_l), llm_acc=float(np.mean(ok_l)), alpha_acc=float(np.mean(ok_a)), llm_latency_ms_p50=float(np.percentile(lat, 50)),
+                llm_latency_ms_p95=float(np.percentile(lat, 95)), families=fams)
+
+
+def benchmark(run_llm=True):
+    m = ST["model"]; m.eval(); T = float(m.post_T); R = dict(post_T=T, quality={}, calib={})
+    for wording in ("train", "unseen"):
+        allr = []
+        for f in ST["train_fams"]:
+            res = predict_res(m, prep(val_examples(f, C.n_cal, C.n_cal + C.n_test, wording))); allr += res
+            s = stat(res, T)["per"][f]; R["quality"].setdefault(wording, {})[f] = {k: s[k] for k in ("acc", "nll", "ece", "brier", "n")}
+        s0, s1 = stat(allr, 1.0), stat(allr, T)
+        R["quality"][wording]["_macro"] = dict(acc=s1["score"], nll=s1["nll"], ece=s1["ece"], brier=s1["brier"])
+        R["calib"][wording] = dict(ece_before_T=s0["ece"], ece_after_T=s1["ece"], conf=s1["micro"]["conf"], corr=s1["micro"]["corr"])
+    R["held_out"] = {}
+    for f in ST["held_fams"]:
+        eb = prep(val_examples(f, C.n_cal, C.n_cal + C.n_test, "unseen"))
+        c = adapt_curve(m, f, eb, C.probe_steps, C.probe_at); R["held_out"][f] = dict(zero_shot_acc=c[0][0], adapted_acc={k: v[0] for k, v in c.items()}, n=len(val_examples(f, C.n_cal, C.n_cal + C.n_test, "unseen")))
+    R["workflow"] = workflow_eval()
+    R["speed"] = speed_bench()
+    R["type_safety"] = dict(schema_violations=0, note="outputs are a masked softmax over the declared options; an out-of-schema answer cannot be produced")
+    if run_llm:
+        try: R["llm_baseline"] = llm_baseline()
+        except Exception as e: R["llm_baseline"] = dict(error=f"{type(e).__name__}: {str(e)[:200]}")
+    R["frr"] = ST["frr"]; R["dmr"] = ST["DMR"]; R["meta"] = {str(k): v for k, v in ST["meta"].items()}
+    ST["R"] = R
+    print(json.dumps({k: v for k, v in R.items() if k not in ("calib", "meta")}, indent=1, default=float)[:6000]); return R
+
+
+def plots(path="alpha_report.png"):
+    import matplotlib; import matplotlib.pyplot as plt
+    H, R = ST["hist"], ST["R"]; st = [h["step"] for h in H]
+    fig, ax = plt.subplots(2, 4, figsize=(22, 9)); ax = ax.ravel()
+    ax[0].plot(st, [h["score"] for h in H], label="val acc"); ax[0].plot(st, [h["tr_score"] for h in H], label="train acc"); ax[0].set_title("Learning curve (EMA weights)"); ax[0].legend()
+    ax[1].plot(st, [h["nll"] for h in H], label="val nll"); ax[1].plot(st, [h["tr_nll"] for h in H], label="train nll"); ax[1].axvline(ST["DMR"]["T_discover"], color="g", ls=":"); ax[1].axvline(ST["DMR"]["T_memorize"], color="r", ls=":")
+    ax[1].set_title(f"DMR = {ST['DMR']['DMR']:.2f}  (T_discover/T_memorize)"); ax[1].legend()
+    ax[2].plot(st, [h["wd"] for h in H], label="wd (new modules)"); ax[2].plot(st, [h["rho"] if h["rho"] is not None else 0 for h in H], label="rho"); ax[2].plot(st, [h["gn"] for h in H], label="grad norm"); ax[2].legend(); ax[2].set_title("Discovery controller")
+    sg = list(R["frr"]); ax[3].bar([str(s) for s in sg], [min(R["frr"][s]["FRR"], 100) for s in sg]); ax[3].axhline(1, color="k", ls="--"); ax[3].set_yscale("log"); ax[3].set_title("FRR vs perturbation sigma (basin test)")
+    mg = ST["meta"]; ss = sorted(k for k in mg); ax[4].plot(ss, [mg[s]["MGR"] for s in ss], "o-"); ax[4].axhline(1, color="k", ls="--"); ax[4].set_title("Meta-grokking rate MGR (held-out families)")
+    for f in ST["held_fams"]: ax[5].plot(ss, [mg[s][f]["zero_shot"] for s in ss], "o-", label=f"{f} zero-shot"); ax[5].plot(ss, [mg[s][f]["acc10"] for s in ss], "s--", label=f"{f} +10 steps")
+    ax[5].legend(fontsize=7); ax[5].set_title("Unseen families: zero-shot / after 10 steps")
+    conf, corr = np.array(R["calib"]["train"]["conf"]), np.array(R["calib"]["train"]["corr"]); e = np.linspace(0, 1, 11)
+    xs, ys = [], []
+    for lo, hi in zip(e[:-1], e[1:]):
+        mk = (conf > lo) & (conf <= hi)
+        if mk.sum() > 5: xs.append(conf[mk].mean()); ys.append(corr[mk].mean())
+    ax[6].plot([0, 1], [0, 1], "k--"); ax[6].plot(xs, ys, "o-"); ax[6].set_title(f"Reliability (ECE {R['calib']['train']['ece_after_T']:.3f})")
+    fams = ST["train_fams"]; x = np.arange(len(fams))
+    ax[7].bar(x - .2, [R["quality"]["train"][f]["acc"] for f in fams], .4, label="seen wording"); ax[7].bar(x + .2, [R["quality"]["unseen"][f]["acc"] for f in fams], .4, label="unseen wording")
+    ax[7].set_xticks(x); ax[7].set_xticklabels(fams, rotation=60); ax[7].legend(); ax[7].set_title("Test accuracy per family")
+    plt.tight_layout(); plt.savefig(path, dpi=110); plt.show()
+
+
+# ----------------------------------------------------------------------------------------------- publish
+def get_token():
+    t = os.environ.get("HF_TOKEN")
+    if not t:
+        try:
+            from kaggle_secrets import UserSecretsClient
+            t = UserSecretsClient().get_secret("HF_TOKEN")
+        except Exception: pass
+    return t
+
+
+def model_card(R):
+    q, sp, wf, h = R["quality"], R["speed"], R["workflow"], R["held_out"]
+    rows = "\n".join(f"| {f} | {q['train'][f]['acc']:.3f} | {q['unseen'][f]['acc']:.3f} | {q['train'][f]['ece']:.3f} |" for f in ST["train_fams"])
+    hrows = "\n".join(f"| {f} | {v['zero_shot_acc']:.3f} | {v['adapted_acc'][10]:.3f} | {v['adapted_acc'][C.probe_steps]:.3f} |" for f, v in h.items())
+    llm = R.get("llm_baseline", {}); llm_s = (f"Qwen2.5-0.5B-Instruct (generative, same {llm['n']} questions): acc {llm['llm_acc']:.3f}, p50 latency {llm['llm_latency_ms_p50']:.0f} ms "
+                                              f"vs Alpha acc {llm['alpha_acc']:.3f}, p50 {sp['latency_ms_p50']:.1f} ms.") if "llm_acc" in llm else "LLM baseline not run."
+    fr = " ".join(f"sigma={s}: FRR {v['FRR']:.1f}{'*' if v['censored'] else ''}" for s, v in R["frr"].items())
+    return f"""---
+license: apache-2.0
+base_model: {C.backbone}
+tags: [system-one, structured-output, calibrated, classification, parallel-decoding]
+---
+# Alpha (System-One-style typed decision model)
+
+Alpha takes unstructured / structured state plus **typed questions with declared options** and returns calibrated probabilities for
+**all questions in one parallel forward pass**. It cannot generate strings and cannot answer outside the declared options
+(masked softmax -> schema violations are impossible by construction). Option embeddings are independent of the state, so a schema can be
+compiled once and cached; cardinality is limited only by memory.
+
+Architecture: `{C.backbone}` encoder (state + question slots with `<mask>` readout tokens) -> 2-layer **Slot Reasoner** (set transformer, squared-ReLU,
+zero-init residuals) -> biaffine/cosine scorer against option embeddings with per-question learned temperature + post-hoc temperature (T={R['post_T']:.2f}).
+
+Training: {C.steps} steps x {C.bs} workflow-style examples (1-4 questions over named fields, mixed from {len(ST['train_fams'])} public datasets with paraphrased
+questions/option wordings), proper-scoring-rule loss (NLL + Brier + ordinal term), EMA weights, functional-memory anchor, discovery controller.
+Held-out families (never trained on): {', '.join(ST['held_fams'])}.
+
+## Results (test splits; see results.json, alpha_report.png)
+| family | acc (seen wording) | acc (unseen wording) | ECE |
+|---|---|---|---|
+{rows}
+
+Workflow eval ({wf['k']} questions per query): per-question acc {wf['q_acc']:.3f}, all-correct {wf['all_correct']:.3f}, joint-confidence ECE {wf['joint_ece']:.3f}.
+
+Held-out task families (unseen schema wording): zero-shot / after 10 / after {C.probe_steps} gradient steps
+| family | zero-shot | +10 steps | +{C.probe_steps} steps |
+|---|---|---|---|
+{hrows}
+
+Speed on {sp['device']}: p50 {sp['latency_ms_p50']:.1f} ms, p95 {sp['latency_ms_p95']:.1f} ms for a 3-question query (local, no network); {sp['throughput_workflows_per_s_bs64']:.0f} workflows/s at batch 64.
+{llm_s}
+
+Framework diagnostics: DMR={R['dmr']['DMR']:.2f}; FRR by perturbation: {fr} (* = not re-reached within budget); MGR by checkpoint in results.json.
+
+## Honest limits
+* This is a small research model trained on public classification datasets. **It is not a replacement for, and has not been compared against, TypeSafe's Jev**
+  (a closed API; its claims were not verified here). Do not read the latency comparison as a Jev comparison.
+* Narrow task coverage (sentiment/topic/intent/NLI/paraphrase/spam/rating); free-form reasoning is out of scope by design.
+* Calibration is measured in-distribution; shift will degrade it.
+
+## Usage
+```python
+from huggingface_hub import snapshot_download; import sys
+p = snapshot_download("{C.repo_id}"); sys.path.insert(0, p)
+from alpha_s1 import Alpha
+m = Alpha.from_pretrained(p, device="cuda").eval()
+print(m.predict({{"ticket": "I was charged twice, refund me now."}},
+      [{{"id": "refund", "question": "Does the customer want a refund?", "options": ["no", "yes"]}}]))
+```
+"""
+
+
+def publish(out_dir="alpha_out"):
+    from huggingface_hub import HfApi
+    tok = get_token(); assert tok, "Set a Kaggle secret named HF_TOKEN (Add-ons > Secrets) or env var HF_TOKEN."
+    R = ST["R"]; m = ST["model"]; m.save_pretrained(out_dir)
+    shutil.copy("alpha_s1.py", out_dir); shutil.copy("alpha_train.py", out_dir)
+    if os.path.exists("alpha_report.png"): shutil.copy("alpha_report.png", out_dir)
+    slim = {k: v for k, v in R.items() if k != "calib"}; json.dump(slim, open(f"{out_dir}/results.json", "w"), indent=1, default=float)
+    json.dump(dict(hist=ST["hist"], probes=ST["probes"]), open(f"{out_dir}/training_log.json", "w"), default=float)
+    open(f"{out_dir}/README.md", "w").write(model_card(R))
+    api = HfApi(token=tok); api.create_repo(C.repo_id, exist_ok=True, private=False)
+    api.upload_folder(folder_path=out_dir, repo_id=C.repo_id, commit_message="Alpha: System-One-style typed decision model")
+    print(f"published -> https://huggingface.co/{C.repo_id}")
+
+
+def demo():
+    m = ST["model"]; m.eval()
+    state = dict(review="The battery died after two days and support never answered. Total waste of money.", account="plan: pro; tenure_months: 14")
+    qs = [dict(id="sentiment", question="Is the review positive or negative?", options=["negative", "positive"]),
+          dict(id="stars", question="How many stars does the review deserve?", options=["1 star", "2 stars", "3 stars", "4 stars", "5 stars"], ordinal=True),
+          dict(id="topic", question="What is the review about?", options=["hardware", "billing", "support", "shipping"])]
+    print(json.dumps(m.predict(state, qs), indent=1))
+
+
+def chat():
+    m = ST["model"]; print("Alpha structured console. Empty text to quit.")
+    while True:
+        text = input("state text> ").strip()
+        if not text: break
+        q = input("question> ").strip(); opts = [o.strip() for o in input("options (comma separated)> ").split(",") if o.strip()]
+        r = m.predict(dict(input=text), [dict(id="q", question=q, options=opts)])["q"]
+        print(f"  answer={r['answer']}  confidence={r['confidence']:.3f}  certainty={r['certainty']:.3f}\n  probs={ {k: round(v, 3) for k, v in r['probs'].items()} }")

pytorch_model.bin

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

results.json

@@ -0,0 +1,435 @@
+{
+ "post_T": 1.0399999618530273,
+ "quality": {
+  "train": {
+   "ag_news": {
+    "acc": 0.926,
+    "nll": 0.23406508564949036,
+    "ece": 0.0323681704998016,
+    "brier": 0.11545059829950333,
+    "n": 500
+   },
+   "banking77": {
+    "acc": 0.902,
+    "nll": 0.3612114489078522,
+    "ece": 0.035957046896219286,
+    "brier": 0.15479955077171326,
+    "n": 500
+   },
+   "emotion": {
+    "acc": 0.922,
+    "nll": 0.15576711297035217,
+    "ece": 0.030244563996791818,
+    "brier": 0.08809012919664383,
+    "n": 500
+   },
+   "imdb": {
+    "acc": 0.886,
+    "nll": 0.2880212068557739,
+    "ece": 0.027435769557952863,
+    "brier": 0.1726732850074768,
+    "n": 500
+   },
+   "mrpc": {
+    "acc": 0.8217054263565892,
+    "nll": 0.3895823657512665,
+    "ece": 0.042235060486682595,
+    "brier": 0.25083833932876587,
+    "n": 258
+   },
+   "news20": {
+    "acc": 0.674,
+    "nll": 1.0884853601455688,
+    "ece": 0.06790356469154359,
+    "brier": 0.4469330608844757,
+    "n": 500
+   },
+   "sms": {
+    "acc": 0.9926470588235294,
+    "nll": 0.038313742727041245,
+    "ece": 0.0037884642096127656,
+    "brier": 0.015446269884705544,
+    "n": 408
+   },
+   "snli": {
+    "acc": 0.772,
+    "nll": 0.5716506838798523,
+    "ece": 0.045889752805233,
+    "brier": 0.3297711908817291,
+    "n": 500
+   },
+   "sst2": {
+    "acc": 0.912,
+    "nll": 0.24116064608097076,
+    "ece": 0.034141219377517684,
+    "brier": 0.14030593633651733,
+    "n": 500
+   },
+   "yelp": {
+    "acc": 0.572,
+    "nll": 0.9883294105529785,
+    "ece": 0.0719212611913681,
+    "brier": 0.5307314991950989,
+    "n": 500
+   },
+   "_macro": {
+    "acc": 0.8380352485180118,
+    "nll": 0.4356587529182434,
+    "ece": 0.01943733276175781,
+    "brier": 0.2245039939880371
+   }
+  },
+  "unseen": {
+   "ag_news": {
+    "acc": 0.924,
+    "nll": 0.25852325558662415,
+    "ece": 0.041514035522937774,
+    "brier": 0.12416581809520721,
+    "n": 500
+   },
+   "banking77": {
+    "acc": 0.89,
+    "nll": 0.4005417227745056,
+    "ece": 0.06322333976626397,
+    "brier": 0.16635079681873322,
+    "n": 500
+   },
+   "emotion": {
+    "acc": 0.894,
+    "nll": 0.38785520195961,
+    "ece": 0.12037511637806893,
+    "brier": 0.1763131022453308,
+    "n": 500
+   },
+   "imdb": {
+    "acc": 0.892,
+    "nll": 0.28929969668388367,
+    "ece": 0.023690259099006624,
+    "brier": 0.17336110770702362,
+    "n": 500
+   },
+   "mrpc": {
+    "acc": 0.8217054263565892,
+    "nll": 0.4847447872161865,
+    "ece": 0.16866437339967533,
+    "brier": 0.3047117590904236,
+    "n": 258
+   },
+   "news20": {
+    "acc": 0.678,
+    "nll": 1.0626460313796997,
+    "ece": 0.05303302364051344,
+    "brier": 0.44085004925727844,
+    "n": 500
+   },
+   "sms": {
+    "acc": 0.34558823529411764,
+    "nll": 1.1533381938934326,
+    "ece": 0.3848064363587137,
+    "brier": 0.8449604511260986,
+    "n": 408
+   },
+   "snli": {
+    "acc": 0.538,
+    "nll": 0.9769825339317322,
+    "ece": 0.11520019578933716,
+    "brier": 0.5887723565101624,
+    "n": 500
+   },
+   "sst2": {
+    "acc": 0.91,
+    "nll": 0.22101609408855438,
+    "ece": 0.02428796315193177,
+    "brier": 0.1290440559387207,
+    "n": 500
+   },
+   "yelp": {
+    "acc": 0.522,
+    "nll": 1.1097854375839233,
+    "ece": 0.09259637647867204,
+    "brier": 0.6057519912719727,
+    "n": 500
+   },
+   "_macro": {
+    "acc": 0.7415293661650708,
+    "nll": 0.6344732642173767,
+    "ece": 0.03709241313372004,
+    "brier": 0.35542815923690796
+   }
+  }
+ },
+ "held_out": {
+  "cola": {
+   "zero_shot_acc": 0.374,
+   "adapted_acc": {
+    "0": 0.374,
+    "5": 0.652,
+    "10": 0.658,
+    "20": 0.686,
+    "40": 0.686
+   },
+   "n": 500
+  },
+  "dbpedia": {
+   "zero_shot_acc": 0.312,
+   "adapted_acc": {
+    "0": 0.312,
+    "5": 0.488,
+    "10": 0.586,
+    "20": 0.876,
+    "40": 0.952
+   },
+   "n": 500
+  },
+  "offensive": {
+   "zero_shot_acc": 0.528,
+   "adapted_acc": {
+    "0": 0.528,
+    "5": 0.674,
+    "10": 0.632,
+    "20": 0.656,
+    "40": 0.694
+   },
+   "n": 500
+  }
+ },
+ "workflow": {
+  "q_acc": 0.825,
+  "all_correct": 0.43666666666666665,
+  "joint_ece": 0.10074469861705296,
+  "k": 4,
+  "n": 300
+ },
+ "speed": {
+  "latency_ms_p50": 11.066616000334761,
+  "latency_ms_p95": 13.510893649527132,
+  "cold_first_call_ms": 173.55092199977662,
+  "throughput_workflows_per_s_bs64": 354.2243768052664,
+  "questions_per_workflow": 3,
+  "device": "Tesla T4"
+ },
+ "type_safety": {
+  "schema_violations": 0,
+  "note": "outputs are a masked softmax over the declared options; an out-of-schema answer cannot be produced"
+ },
+ "llm_baseline": {
+  "model": "Qwen/Qwen2.5-0.5B-Instruct",
+  "n": 120,
+  "llm_acc": 0.5833333333333334,
+  "alpha_acc": 0.875,
+  "llm_latency_ms_p50": 82.17461200001708,
+  "llm_latency_ms_p95": 150.7793488504376,
+  "families": [
+   "sst2",
+   "ag_news",
+   "emotion",
+   "snli"
+  ]
+ },
+ "frr": {
+  "0.5": {
+   "v0": 0.8416666666666666,
+   "tau": 0.83375,
+   "T_first": 4250,
+   "T_re": 1,
+   "FRR": 4250.0,
+   "L_FRR": 0.00023529411764705883,
+   "censored": false,
+   "trivial": true
+  },
+  "1": {
+   "v0": 0.8479166666666667,
+   "tau": 0.83375,
+   "T_first": 4250,
+   "T_re": 1,
+   "FRR": 4250.0,
+   "L_FRR": 0.00023529411764705883,
+   "censored": false,
+   "trivial": true
+  },
+  "2": {
+   "v0": 0.8375,
+   "tau": 0.83375,
+   "T_first": 4250,
+   "T_re": 1,
+   "FRR": 4250.0,
+   "L_FRR": 0.00023529411764705883,
+   "censored": false,
+   "trivial": true
+  },
+  "4": {
+   "v0": 0.8416666666666668,
+   "tau": 0.83375,
+   "T_first": 4250,
+   "T_re": 1,
+   "FRR": 4250.0,
+   "L_FRR": 0.00023529411764705883,
+   "censored": false,
+   "trivial": true
+  },
+  "8": {
+   "v0": 0.7979166666666667,
+   "tau": 0.83375,
+   "T_first": 4250,
+   "T_re": 300,
+   "FRR": 14.166666666666666,
+   "L_FRR": 0.07058823529411765,
+   "censored": true,
+   "trivial": false
+  }
+ },
+ "dmr": {
+  "T_discover": 3250,
+  "T_memorize": 3000,
+  "DMR": 1.0833333333333333
+ },
+ "meta": {
+  "0": {
+   "cola": {
+    "T_D": 5,
+    "zero_shot": 0.46,
+    "acc10": 0.665,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 80,
+    "zero_shot": 0.115,
+    "acc10": 0.045,
+    "final": 0.045
+   },
+   "offensive": {
+    "T_D": 80,
+    "zero_shot": 0.645,
+    "acc10": 0.64,
+    "final": 0.64
+   },
+   "MGR": 1.0
+  },
+  "1000": {
+   "cola": {
+    "T_D": 5,
+    "zero_shot": 0.55,
+    "acc10": 0.635,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.34,
+    "acc10": 0.525,
+    "final": 0.93
+   },
+   "offensive": {
+    "T_D": 40,
+    "zero_shot": 0.64,
+    "acc10": 0.64,
+    "final": 0.745
+   },
+   "MGR": 1.5874010519681994
+  },
+  "2000": {
+   "cola": {
+    "T_D": 20,
+    "zero_shot": 0.4,
+    "acc10": 0.61,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.305,
+    "acc10": 0.7,
+    "final": 0.955
+   },
+   "offensive": {
+    "T_D": 40,
+    "zero_shot": 0.625,
+    "acc10": 0.69,
+    "final": 0.755
+   },
+   "MGR": 1.0
+  },
+  "3000": {
+   "cola": {
+    "T_D": 10,
+    "zero_shot": 0.4,
+    "acc10": 0.665,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.245,
+    "acc10": 0.64,
+    "final": 0.975
+   },
+   "offensive": {
+    "T_D": 80,
+    "zero_shot": 0.59,
+    "acc10": 0.685,
+    "final": 0.65
+   },
+   "MGR": 1.0
+  },
+  "4000": {
+   "cola": {
+    "T_D": 5,
+    "zero_shot": 0.41,
+    "acc10": 0.665,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.305,
+    "acc10": 0.71,
+    "final": 0.975
+   },
+   "offensive": {
+    "T_D": 40,
+    "zero_shot": 0.57,
+    "acc10": 0.625,
+    "final": 0.74
+   },
+   "MGR": 1.5874010519681994
+  },
+  "5000": {
+   "cola": {
+    "T_D": 5,
+    "zero_shot": 0.4,
+    "acc10": 0.665,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.355,
+    "acc10": 0.68,
+    "final": 0.945
+   },
+   "offensive": {
+    "T_D": 80,
+    "zero_shot": 0.575,
+    "acc10": 0.635,
+    "final": 0.715
+   },
+   "MGR": 1.2599210498948732
+  },
+  "6000": {
+   "cola": {
+    "T_D": 10,
+    "zero_shot": 0.395,
+    "acc10": 0.665,
+    "final": 0.665
+   },
+   "dbpedia": {
+    "T_D": 40,
+    "zero_shot": 0.355,
+    "acc10": 0.695,
+    "final": 0.94
+   },
+   "offensive": {
+    "T_D": 40,
+    "zero_shot": 0.555,
+    "acc10": 0.66,
+    "final": 0.77
+   },
+   "MGR": 1.2599210498948732
+  }
+ }
+}

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "RobertaTokenizer",
+  "trim_offsets": true,
+  "unk_token": "<unk>"
+}

training_log.json

@@ -0,0 +1 @@
+{"hist": [{"step": 250, "score": 0.42499999999999993, "nll": 1.3864185810089111, "ece": 0.07763973646021138, "tr_score": 0.4333333333333333, "tr_nll": 1.3839653730392456, "fam": {"ag_news": 0.8333333333333334, "banking77": 0.0, "emotion": 0.3958333333333333, "imdb": 0.5, "mrpc": 0.6666666666666666, "news20": 0.125, "sms": 0.8958333333333334, "snli": 0.22916666666666666, "sst2": 0.4583333333333333, "yelp": 0.14583333333333334}, "wd": 0.5, "lam": 0.3, "gn": 4.74530553817749, "loss": 1.4470987440866285, "rho": null, "lr": 4.166666666666667e-05}, {"step": 500, "score": 0.4875, "nll": 1.2226603031158447, "ece": 0.059230080638856944, "tr_score": 0.4645833333333333, "tr_nll": 1.2430381774902344, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.0625, "emotion": 0.3958333333333333, "imdb": 0.5, "mrpc": 0.7291666666666666, "news20": 0.2916666666666667, "sms": 1.0, "snli": 0.375, "sst2": 0.4583333333333333, "yelp": 0.10416666666666667}, "wd": 0.5, "lam": 0.3, "gn": 4.94678258895874, "loss": 1.3388411353066019, "rho": null, "lr": 4.986344023965386e-05}, {"step": 750, "score": 0.5041666666666667, "nll": 1.1497793197631836, "ece": 0.05984721723943949, "tr_score": 0.4937500000000001, "tr_nll": 1.1202466487884521, "fam": {"ag_news": 0.8958333333333334, "banking77": 0.25, "emotion": 0.3958333333333333, "imdb": 0.5208333333333334, "mrpc": 0.6458333333333334, "news20": 0.3541666666666667, "sms": 1.0, "snli": 0.375, "sst2": 0.4583333333333333, "yelp": 0.14583333333333334}, "wd": 0.45, "lam": 0.3, "gn": 3.086398124694824, "loss": 1.2649021407320382, "rho": 1.1620062206951038, "lr": 4.931150598363494e-05}, {"step": 1000, "score": 0.6145833333333333, "nll": 0.9992086291313171, "ece": 0.08922299118712546, "tr_score": 0.5999999999999999, "tr_nll": 0.9506429433822632, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.2916666666666667, "emotion": 0.5625, "imdb": 0.7708333333333334, "mrpc": 0.6666666666666666, "news20": 0.4791666666666667, "sms": 1.0, "snli": 0.25, "sst2": 0.7708333333333334, "yelp": 0.3958333333333333}, "wd": 0.405, "lam": 0.3, "gn": 7.461432456970215, "loss": 1.1085828546282157, "rho": 0.8777702575734148, "lr": 4.834611651233304e-05}, {"step": 1250, "score": 0.7041666666666667, "nll": 0.7773005366325378, "ece": 0.060513773995141164, "tr_score": 0.6979166666666666, "tr_nll": 0.7319567203521729, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.3958333333333333, "emotion": 0.8333333333333334, "imdb": 0.8333333333333334, "mrpc": 0.7083333333333334, "news20": 0.5416666666666666, "sms": 0.9791666666666666, "snli": 0.4166666666666667, "sst2": 0.9375, "yelp": 0.4791666666666667}, "wd": 0.36450000000000005, "lam": 0.3, "gn": 4.511319160461426, "loss": 0.87960810418064, "rho": 0.8827748724431335, "lr": 4.6985571585149876e-05}, {"step": 1500, "score": 0.7229166666666667, "nll": 0.6783604025840759, "ece": 0.036758976274480416, "tr_score": 0.7374999999999999, "tr_nll": 0.627739429473877, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.5208333333333334, "emotion": 0.9166666666666666, "imdb": 0.8541666666666666, "mrpc": 0.6875, "news20": 0.5625, "sms": 1.0, "snli": 0.22916666666666666, "sst2": 0.9583333333333334, "yelp": 0.5416666666666666}, "wd": 0.32805000000000006, "lam": 0.3, "gn": 7.843607425689697, "loss": 0.8132543158575171, "rho": 0.9487538569865493, "lr": 4.5255661461418854e-05}, {"step": 1750, "score": 0.7458333333333333, "nll": 0.6375653743743896, "ece": 0.03906969136248029, "tr_score": 0.75625, "tr_nll": 0.5771324038505554, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.625, "emotion": 0.8958333333333334, "imdb": 0.8541666666666666, "mrpc": 0.7291666666666666, "news20": 0.5833333333333334, "sms": 1.0, "snli": 0.3125, "sst2": 0.9791666666666666, "yelp": 0.5625}, "wd": 0.2952450000000001, "lam": 0.3, "gn": 7.016384601593018, "loss": 0.7663959374513087, "rho": 0.9490588800899795, "lr": 4.3189178024614896e-05}, {"step": 2000, "score": 0.7729166666666666, "nll": 0.5967954397201538, "ece": 0.025282762618735424, "tr_score": 0.7583333333333333, "tr_nll": 0.5407644510269165, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.7291666666666666, "emotion": 0.9166666666666666, "imdb": 0.875, "mrpc": 0.7291666666666666, "news20": 0.5833333333333334, "sms": 1.0, "snli": 0.3958333333333333, "sst2": 0.9583333333333334, "yelp": 0.625}, "wd": 0.2657205000000001, "lam": 0.3, "gn": 9.855315208435059, "loss": 0.7069250917348392, "rho": 0.87758640404572, "lr": 4.0825293184962056e-05}, {"step": 2250, "score": 0.7875, "nll": 0.5712440013885498, "ece": 0.04493665626893442, "tr_score": 0.75625, "tr_nll": 0.5333563685417175, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.8125, "emotion": 0.8958333333333334, "imdb": 0.8541666666666666, "mrpc": 0.75, "news20": 0.5416666666666666, "sms": 1.0, "snli": 0.4791666666666667, "sst2": 0.9583333333333334, "yelp": 0.6666666666666666}, "wd": 0.23914845000000007, "lam": 0.3, "gn": 8.997612953186035, "loss": 0.6946176960095765, "rho": 0.9993132683013181, "lr": 3.8208816343334156e-05}, {"step": 2500, "score": 0.8104166666666666, "nll": 0.5312714576721191, "ece": 0.05300602965677776, "tr_score": 0.78125, "tr_nll": 0.5068899393081665, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.7916666666666666, "emotion": 0.875, "imdb": 0.8541666666666666, "mrpc": 0.7708333333333334, "news20": 0.5833333333333334, "sms": 1.0, "snli": 0.5833333333333334, "sst2": 1.0, "yelp": 0.6875}, "wd": 0.21523360500000008, "lam": 0.3, "gn": 8.83574390411377, "loss": 0.6346616129986626, "rho": 1.499656634150659, "lr": 3.5389344991836974e-05}, {"step": 2750, "score": 0.8083333333333332, "nll": 0.49808335304260254, "ece": 0.04100577929057176, "tr_score": 0.7875, "tr_nll": 0.47457775473594666, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.7708333333333334, "emotion": 0.875, "imdb": 0.8541666666666666, "mrpc": 0.75, "news20": 0.6041666666666666, "sms": 1.0, "snli": 0.6875, "sst2": 0.9583333333333334, "yelp": 0.625}, "wd": 0.19371024450000007, "lam": 0.3, "gn": 6.362133979797363, "loss": 0.58079190026933, "rho": 1.5049839023584053, "lr": 3.242032455214346e-05}, {"step": 3000, "score": 0.8291666666666666, "nll": 0.47499704360961914, "ece": 0.04109802747455738, "tr_score": 0.8020833333333333, "tr_nll": 0.4411352574825287, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.8333333333333334, "emotion": 0.8958333333333334, "imdb": 0.875, "mrpc": 0.7916666666666666, "news20": 0.625, "sms": 1.0, "snli": 0.75, "sst2": 0.9375, "yelp": 0.625}, "wd": 0.17433922005000008, "lam": 0.3, "gn": 9.050947189331055, "loss": 0.5507215846613893, "rho": 1.2660459721916015, "lr": 2.9358035273127483e-05}, {"step": 3250, "score": 0.8229166666666667, "nll": 0.45618176460266113, "ece": 0.03542887705067793, "tr_score": 0.8125, "tr_nll": 0.4283953309059143, "fam": {"ag_news": 0.9375, "banking77": 0.875, "emotion": 0.9166666666666666, "imdb": 0.8541666666666666, "mrpc": 0.7708333333333334, "news20": 0.6041666666666666, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9166666666666666, "yelp": 0.625}, "wd": 0.15690529804500009, "lam": 0.3001, "gn": 7.896519184112549, "loss": 0.5361467864703793, "rho": 0.9781872430895555, "lr": 2.6260525391993023e-05}, {"step": 3500, "score": 0.8229166666666666, "nll": 0.44824162125587463, "ece": 0.04027308725441499, "tr_score": 0.8333333333333333, "tr_nll": 0.41801562905311584, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.8541666666666666, "emotion": 0.9166666666666666, "imdb": 0.8541666666666666, "mrpc": 0.7916666666666666, "news20": 0.5833333333333334, "sms": 1.0, "snli": 0.7708333333333334, "sst2": 0.9375, "yelp": 0.6041666666666666}, "wd": 0.14121476824050008, "lam": 0.3001, "gn": 7.859862804412842, "loss": 0.5212966933069149, "rho": 1.2275311193516978, "lr": 2.3186510781715892e-05}, {"step": 3750, "score": 0.825, "nll": 0.44514527916908264, "ece": 0.0378269159855942, "tr_score": 0.8416666666666668, "tr_nll": 0.41067466139793396, "fam": {"ag_news": 0.9166666666666666, "banking77": 0.8333333333333334, "emotion": 0.9375, "imdb": 0.8541666666666666, "mrpc": 0.7708333333333334, "news20": 0.5833333333333334, "sms": 1.0, "snli": 0.7708333333333334, "sst2": 0.9375, "yelp": 0.6458333333333334}, "wd": 0.12709329141645007, "lam": 0.3001, "gn": 4.868711471557617, "loss": 0.5031073578348193, "rho": 0.9962497321208292, "lr": 2.0194261942894628e-05}, {"step": 4000, "score": 0.83125, "nll": 0.43428102135658264, "ece": 0.04779160485292472, "tr_score": 0.8479166666666667, "tr_nll": 0.4039950370788574, "fam": {"ag_news": 0.9375, "banking77": 0.8333333333333334, "emotion": 0.9375, "imdb": 0.8541666666666666, "mrpc": 0.7708333333333334, "news20": 0.6458333333333334, "sms": 1.0, "snli": 0.75, "sst2": 0.9375, "yelp": 0.6458333333333334}, "wd": 0.12709329141645007, "lam": 0.3001, "gn": 6.261903762817383, "loss": 0.5046875660977973, "rho": 0.709019548630384, "lr": 1.7340499438004994e-05}, {"step": 4250, "score": 0.8375, "nll": 0.4272522032260895, "ece": 0.037357999570667745, "tr_score": 0.8416666666666666, "tr_nll": 0.397662878036499, "fam": {"ag_news": 0.9375, "banking77": 0.8541666666666666, "emotion": 0.9375, "imdb": 0.8541666666666666, "mrpc": 0.7916666666666666, "news20": 0.6458333333333334, "sms": 1.0, "snli": 0.75, "sst2": 0.9583333333333334, "yelp": 0.6458333333333334}, "wd": 0.11438396227480506, "lam": 0.3001, "gn": 10.974567413330078, "loss": 0.47385876453106374, "rho": 1.167748460619184, "lr": 1.4679318706019013e-05}, {"step": 4500, "score": 0.8375, "nll": 0.42399367690086365, "ece": 0.038577316980808976, "tr_score": 0.84375, "tr_nll": 0.39106932282447815, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.8541666666666666, "emotion": 0.9166666666666666, "imdb": 0.8541666666666666, "mrpc": 0.7916666666666666, "news20": 0.6666666666666666, "sms": 1.0, "snli": 0.75, "sst2": 0.9583333333333334, "yelp": 0.625}, "wd": 0.10294556604732455, "lam": 0.3001, "gn": 6.629716396331787, "loss": 0.4687045077363484, "rho": 1.1388838315746999, "lr": 1.2261164638420832e-05}, {"step": 4750, "score": 0.8395833333333332, "nll": 0.42487725615501404, "ece": 0.034516340897728995, "tr_score": 0.84375, "tr_nll": 0.38701558113098145, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.8541666666666666, "mrpc": 0.7916666666666666, "news20": 0.6458333333333334, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6458333333333334}, "wd": 0.0926510094425921, "lam": 0.3001, "gn": 7.820315361022949, "loss": 0.45272200221767106, "rho": 0.8165412590434078, "lr": 1.013187535438278e-05}, {"step": 5000, "score": 0.8395833333333332, "nll": 0.4237828850746155, "ece": 0.0383824546821416, "tr_score": 0.8458333333333334, "tr_nll": 0.3849530816078186, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.875, "mrpc": 0.7916666666666666, "news20": 0.625, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6458333333333334}, "wd": 0.11581376180324011, "lam": 0.3001, "gn": 6.347634315490723, "loss": 0.4564408298368683, "rho": 0.2992874578055718, "lr": 8.331813301137644e-06}, {"step": 5250, "score": 0.8395833333333332, "nll": 0.4211656451225281, "ece": 0.04211962077145778, "tr_score": 0.8541666666666667, "tr_nll": 0.378847599029541, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.8541666666666666, "mrpc": 0.7916666666666666, "news20": 0.6458333333333334, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6458333333333334}, "wd": 0.11581376180324011, "lam": 0.3001, "gn": 6.547802448272705, "loss": 0.4516936281062966, "rho": 0.4149458703356097, "lr": 6.8951001502612065e-06}, {"step": 5500, "score": 0.8416666666666666, "nll": 0.419617235660553, "ece": 0.04883364103734491, "tr_score": 0.85, "tr_nll": 0.37390393018722534, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.875, "mrpc": 0.7916666666666666, "news20": 0.6458333333333334, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6458333333333334}, "wd": 0.11581376180324011, "lam": 0.3001, "gn": 7.733129501342773, "loss": 0.43449673415669426, "rho": 0.42180815916788295, "lr": 5.8489699930418664e-06}, {"step": 5750, "score": 0.84375, "nll": 0.41734427213668823, "ece": 0.042368609768648896, "tr_score": 0.8479166666666668, "tr_nll": 0.37107983231544495, "fam": {"ag_news": 0.9375, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.875, "mrpc": 0.7916666666666666, "news20": 0.6666666666666666, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6666666666666666}, "wd": 0.11581376180324011, "lam": 0.3001, "gn": 9.12564754486084, "loss": 0.430490410395239, "rho": 0.3675093773257097, "lr": 5.213253095656177e-06}, {"step": 6000, "score": 0.84375, "nll": 0.4162411093711853, "ece": 0.04647985563303032, "tr_score": 0.8458333333333334, "tr_nll": 0.36772745847702026, "fam": {"ag_news": 0.9583333333333334, "banking77": 0.875, "emotion": 0.9375, "imdb": 0.875, "mrpc": 0.7708333333333334, "news20": 0.6666666666666666, "sms": 1.0, "snli": 0.7291666666666666, "sst2": 0.9583333333333334, "yelp": 0.6666666666666666}, "wd": 0.11581376180324011, "lam": 0.3001, "gn": 8.71489429473877, "loss": 0.4269631357885679, "rho": 0.5861776263692953, "lr": 5e-06}], "probes": [{"step": 0, "fam": "cola", "curve": {"0": [0.46, 0.6933133602142334], "5": [0.665, 0.6618085503578186], "10": [0.665, 0.6377608776092529], "20": [0.665, 0.6427217721939087], "40": [0.665, 0.6488895416259766]}}, {"step": 0, "fam": "dbpedia", "curve": {"0": [0.115, 2.635910749435425], "5": [0.045, 2.6632566452026367], "10": [0.045, 2.696634531021118], "20": [0.045, 2.638427734375], "40": [0.045, 2.638758897781372]}}, {"step": 0, "fam": "offensive", "curve": {"0": [0.645, 0.6905834674835205], "5": [0.64, 0.6852070093154907], "10": [0.64, 0.6692162156105042], "20": [0.64, 0.6649163961410522], "40": [0.64, 0.7305782437324524]}}, {"step": 1000, "fam": "cola", "curve": {"0": [0.55, 0.6800724864006042], "5": [0.665, 0.6440149545669556], "10": [0.635, 0.6801218390464783], "20": [0.665, 0.6502220630645752], "40": [0.665, 0.6403223276138306]}}, {"step": 1000, "fam": "dbpedia", "curve": {"0": [0.34, 2.2151966094970703], "5": [0.45, 1.7422538995742798], "10": [0.525, 1.1715236902236938], "20": [0.795, 0.5880295634269714], "40": [0.93, 0.17617054283618927]}}, {"step": 1000, "fam": "offensive", "curve": {"0": [0.64, 0.6211206912994385], "5": [0.66, 0.6001800298690796], "10": [0.64, 0.6511775851249695], "20": [0.665, 0.7224280834197998], "40": [0.745, 1.5302973985671997]}}, {"step": 2000, "fam": "cola", "curve": {"0": [0.4, 1.2561464309692383], "5": [0.355, 1.106208324432373], "10": [0.61, 0.6751940846443176], "20": [0.665, 0.6598563194274902], "40": [0.665, 0.6479304432868958]}}, {"step": 2000, "fam": "dbpedia", "curve": {"0": [0.305, 2.196328639984131], "5": [0.575, 1.3923637866973877], "10": [0.7, 0.8865454792976379], "20": [0.86, 0.5139257311820984], "40": [0.955, 0.13148292899131775]}}, {"step": 2000, "fam": "offensive", "curve": {"0": [0.625, 0.6901967525482178], "5": [0.655, 0.6434993743896484], "10": [0.69, 0.5594972372055054], "20": [0.71, 0.6698970198631287], "40": [0.755, 0.5170150995254517]}}, {"step": 3000, "fam": "cola", "curve": {"0": [0.4, 1.342832326889038], "5": [0.62, 0.6754884123802185], "10": [0.665, 0.6443837881088257], "20": [0.665, 0.6436806321144104], "40": [0.665, 0.6559992432594299]}}, {"step": 3000, "fam": "dbpedia", "curve": {"0": [0.245, 2.18424916267395], "5": [0.455, 1.5412259101867676], "10": [0.64, 1.0229398012161255], "20": [0.885, 0.5090384483337402], "40": [0.975, 0.0869322121143341]}}, {"step": 3000, "fam": "offensive", "curve": {"0": [0.59, 0.8314429521560669], "5": [0.66, 0.8946431875228882], "10": [0.685, 0.6292023062705994], "20": [0.68, 0.6085357069969177], "40": [0.65, 0.5966840982437134]}}, {"step": 4000, "fam": "cola", "curve": {"0": [0.41, 1.2136904001235962], "5": [0.635, 0.6475875973701477], "10": [0.665, 0.6374750733375549], "20": [0.665, 0.6445882320404053], "40": [0.665, 0.6596666574478149]}}, {"step": 4000, "fam": "dbpedia", "curve": {"0": [0.305, 2.2257626056671143], "5": [0.575, 1.4236860275268555], "10": [0.71, 0.984937310218811], "20": [0.845, 0.6737918853759766], "40": [0.975, 0.105534628033638]}}, {"step": 4000, "fam": "offensive", "curve": {"0": [0.57, 0.8187274932861328], "5": [0.595, 0.7490072846412659], "10": [0.625, 0.72652268409729], "20": [0.645, 0.619418740272522], "40": [0.74, 0.8135643005371094]}}, {"step": 5000, "fam": "cola", "curve": {"0": [0.4, 1.2741897106170654], "5": [0.64, 0.6492821574211121], "10": [0.665, 0.6398915648460388], "20": [0.665, 0.6444342136383057], "40": [0.665, 0.6411840319633484]}}, {"step": 5000, "fam": "dbpedia", "curve": {"0": [0.355, 2.148179531097412], "5": [0.585, 1.378219485282898], "10": [0.68, 1.0227428674697876], "20": [0.87, 0.4382198750972748], "40": [0.945, 0.20549844205379486]}}, {"step": 5000, "fam": "offensive", "curve": {"0": [0.575, 0.9428966045379639], "5": [0.675, 0.7043898701667786], "10": [0.635, 0.6242003440856934], "20": [0.695, 0.5720301270484924], "40": [0.715, 0.5336439609527588]}}, {"step": 6000, "fam": "cola", "curve": {"0": [0.395, 1.3663846254348755], "5": [0.625, 0.6492061018943787], "10": [0.665, 0.6497192978858948], "20": [0.665, 0.6383249759674072], "40": [0.665, 0.6384083032608032]}}, {"step": 6000, "fam": "dbpedia", "curve": {"0": [0.355, 2.106869697570801], "5": [0.435, 1.3792909383773804], "10": [0.695, 1.0408819913864136], "20": [0.895, 0.3912413418292999], "40": [0.94, 0.190094992518425]}}, {"step": 6000, "fam": "offensive", "curve": {"0": [0.555, 0.9884564280509949], "5": [0.65, 0.7950121164321899], "10": [0.66, 0.5845534801483154], "20": [0.695, 0.9688383340835571], "40": [0.77, 0.5907872319221497]}}]}