v4: same recipe trained for 48,000 steps on 480k requests (Decision Index 0.3 public 30.90); v3 kept at tag v3

README.md

@@ -8,14 +8,15 @@ datasets:
 tags:
 - typed-decisions
 - decision-index
-- cross-encoder
 - decision-model
+- cross-encoder
 ---
 
 # tasksource-decider-nano
 
-This repository was `tasksource/tasksource-jev-nano-v1`. Its `main` (tag `v3`) is the joint cross-encoder trained for
-24,000 steps. The same model trained for 12,000 steps (public index 26.77) stays at tag
+This repository was `tasksource/tasksource-jev-nano-v1`. Its `main` (tag `v4`) is the joint cross-encoder trained for
+48,000 steps. The same recipe at 24,000 steps (public index 29.05) stays at tag
+[`v3`](https://huggingface.co/tasksource/tasksource-decider-nano/tree/v3), at 12,000 steps (26.77) at tag
 [`v2`](https://huggingface.co/tasksource/tasksource-decider-nano/tree/v2), and the previous multi-vector model at tag
 [`v1`](https://huggingface.co/tasksource/tasksource-decider-nano/tree/v1).
 
@@ -25,9 +26,9 @@ generate text.
 
 | | |
 |---|---|
-| Decision Index 0.3, public index | **29.05** |
+| Decision Index 0.3, public index | **30.90** |
 | Requests | 140,171 `ok`, 7 `unsupported` (pairs above 8,192 tokens) |
-| Latency, 1× NVIDIA A10 23 GB, one request at a time | median 73 ms · mean 96 ms · p80 95 ms |
+| Latency, 1× NVIDIA A10 23 GB, one request at a time | median 52 ms · mean 65 ms · p80 63 ms |
 | Complete results | [runs/tasksource-decider-nano](https://huggingface.co/datasets/tasksource/decision-index-results/tree/main/runs/tasksource-decider-nano) |
 
 ## Architecture
@@ -67,8 +68,8 @@ python -m decision_index pipeline --engine decision_index_engine:DeciderEngine \
 
 ## Training
 
-24,000 steps, batches of 8 requests, learning rate 2e-5, inputs up to 2,048 tokens, seed 7. The loss is soft
-cross-entropy over each question's options. The training mix has 240,000 requests:
+48,000 steps, batches of 8 requests, learning rate 2e-5, inputs up to 2,048 tokens, seed 7. The loss is soft
+cross-entropy over each question's options. The training mix has 480,000 requests:
 
 | Share | Source |
 |---:|---|
@@ -87,10 +88,10 @@ request.
 
 | Area | Skill | Raw |
 |---|---:|---:|
-| Knowledge & Reasoning | 6.6 | 29.8 |
-| Language Understanding | 33.2 | 51.5 |
-| Retrieval & Classification | 47.1 | 61.3 |
-| Tools & Automation | 43.3 | 48.2 |
-| Arts & Human Taste | 14.1 | 43.2 |
+| Knowledge & Reasoning | 7.8 | 30.8 |
+| Language Understanding | 37.7 | 54.7 |
+| Retrieval & Classification | 49.4 | 63.3 |
+| Tools & Automation | 43.6 | 48.3 |
+| Arts & Human Taste | 12.8 | 42.7 |
 
-Raw index 46.41; kit `9eb2dbe`, edition 0.3, 140,171 requests scored.
+Raw index 47.86; kit `9eb2dbe`, edition 0.3, 140,171 requests scored.

decider.py

@@ -106,6 +106,7 @@ class Decider:
         self.net = net.to(self.device).eval()
         self.layout = self.config.get("layout", "pair")
         self.pool = self.config.get("pool", "opt")
+        self.attn = self.config.get("attn", "cached")  # shared layout: "cached" | "full" | "auto"
         self.window = getattr(cfg, "sliding_window", None)
 
     @classmethod
@@ -149,19 +150,35 @@ class Decider:
             start += len(opts)
         return out
 
+    def _full(self, n_questions: int, state_len: int, auto_extra: int = 4096) -> bool:
+        """cached: the state reads only itself and is shared by all questions of a request; full: one row per
+        question, the state also reads its question; auto: full when that costs <= auto_extra extra tokens."""
+        if self.attn == "auto":
+            return (n_questions - 1) * state_len <= auto_extra
+        return self.attn == "full"
+
     def _score_shared(self, items, row_extra: int = 2048) -> List[np.ndarray]:
         torch, tok = self.torch, self.tok
         rows, oid, prefix = [], 0, {}
-        for s, q, opts in items:  # one row group per distinct state; rows of max(2P, P + row_extra) tokens
+        runs = []  # questions per request = consecutive items with the same state
+        for s, _, _ in items:
+            if runs and runs[-1][0] == s:
+                runs[-1][1] += 1
+            else:
+                runs.append([s, 1])
+        nq = [n for _, n in runs for _ in range(n)]
+        first = [k == 0 for _, n in runs for k in range(n)]
+        for (s, q, opts), n_q, new_req in zip(items, nq, first):  # rows of max(2P, P + row_extra) tokens
             p = prefix.setdefault(s, tok(s)["input_ids"])
+            full = self._full(n_q, len(p))
             qi = tok(f"Question: {q}\nOption:", add_special_tokens=False)["input_ids"]
             ob = [tok(" " + o, add_special_tokens=False)["input_ids"] + [tok.sep_token_id] for o in opts]
             longest = len(p) + len(qi) + max(map(len, ob))
             if longest > self.max_length:
                 raise TooLong(f"a (state, question, option) pair needs {longest} tokens > {self.max_length}")
             cap = max(2 * len(p), len(p) + row_extra)
-            if not rows or rows[-1][0] is not p or rows[-1][2] + len(qi) + len(ob[0]) > cap:
-                rows.append([p, [], len(p)])
+            if not rows or rows[-1][0] is not p or new_req or full or rows[-1][2] + len(qi) + len(ob[0]) > cap:
+                rows.append([p, [], len(p), full])
             k = 0
             while k < len(ob):
                 row = rows[-1]
@@ -170,7 +187,7 @@ class Decider:
                     take.append((oid + k + len(take), ob[k + len(take)]))
                     size += len(take[-1][1])
                 if row[1] and row[2] + size > cap:
-                    rows.append([p, [], len(p)])
+                    rows.append([p, [], len(p), full])
                     continue
                 row[1].append((qi, take))
                 row[2] += size
@@ -193,7 +210,8 @@ class Decider:
                 qix = torch.full((B, L), -1, dtype=torch.long)
                 oix = torch.full((B, L), -1, dtype=torch.long)
                 spans = []
-                for r, (p, blocks, _) in enumerate(batch):
+                rowfull = torch.tensor([r[3] for r in batch], dtype=torch.bool)
+                for r, (p, blocks, _, _) in enumerate(batch):
                     ids[r, :len(p)], pos[r, :len(p)], kind[r, :len(p)] = torch.tensor(p), torch.arange(len(p)), 0
                     c = len(p)
                     for b, (qi, take) in enumerate(blocks):
@@ -208,11 +226,12 @@ class Decider:
                             kind[r, c:c + m], qix[r, c:c + m], oix[r, c:c + m] = 2, b, o
                             spans.append((o, r, c, m, len(p), (r, len(p) + sum(len(bb[0]) + sum(len(y) for _, y in bb[1]) for bb in blocks[:b])), n))
                             c += m
-                ids, pos, kind, qix, oix = (t.to(self.device) for t in (ids, pos, kind, qix, oix))
+                ids, pos, kind, qix, oix, rowfull = (t.to(self.device) for t in (ids, pos, kind, qix, oix, rowfull))
                 ki, kj = kind[:, :, None], kind[:, None, :]
                 allow = (kj == 0) & (ki >= 0)
                 allow |= (kj == 1) & (ki >= 1) & (qix[:, :, None] == qix[:, None, :])
                 allow |= (kj == 2) & (ki == 2) & (oix[:, :, None] == oix[:, None, :])
+                allow |= rowfull[:, None, None] & (ki == 0) & (kj == 1)  # full rows: the state reads its question
                 allow |= torch.eye(L, dtype=torch.bool, device=self.device)[None]
                 masks = {"full_attention": allow[:, None]}
                 if self.window is not None:

decider_config.json

@@ -4,19 +4,19 @@
   "base": "cross-encoder/ettin-reranker-150m-v1",
   "train_max_length": 2048,
   "inference_max_length": 8192,
-  "step": 24000,
+  "step": 48000,
   "train": {
     "arch": "joint",
     "opt_pool": "marker",
     "shuffle_options": false,
-    "manifest": "runs/manifests/mix_all.jsonl",
+    "manifest": "runs/manifests/mix_all_480k.jsonl",
     "extra_manifest": null,
     "extra_frac": 0.0,
     "lr": 2e-05,
     "batch_requests": 8,
     "max_k": 32,
     "seed": 7,
-    "total_steps": 24000,
+    "total_steps": 48000,
     "max_length": 2048,
     "pack": 0,
     "init": null,
@@ -28,5 +28,7 @@
     "tok_budget": 12288
   },
   "layout": "shared",
-  "pool": "opt"
+  "pool": "opt",
+  "attn": "cached",
+  "weights_dtype": "bf16"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ca302526bc031a820677508630f72806c7f757fe25191d5eb1e154f658191717
-size 596074556
+oid sha256:e26e29f12339163b473f985270ebb4b9b3b90ee3b7ffd488715d45550842bef4
+size 298044578