Benchmark-selected step 1000: Open SLM average 33.88%

README.md

@@ -21,8 +21,8 @@ Benchmark-selected checkpoint from a 7.85M-parameter
 Channel-Mixing Attention generalist language model using the unmodified native
 GPT-S 4,096-token tokenizer. It has no arithmetic-specific token splitting,
 place embeddings, role embeddings, or inference-time equation detection. It
-was selected at step 0 for an Open SLM Leaderboard-style average of
-30.66%. WikiText normalized BPB is reported separately
+was selected at step 1,000 for an Open SLM Leaderboard-style average of
+33.88%. WikiText normalized BPB is reported separately
 and is not used for selection.
 
 ## Loading
@@ -69,7 +69,7 @@ ordinary causal-language-model text, with no task-specific model features or
 tokenizer transformations. All four training sources are streamed natural-text
 corpora.
 
-## Zero-shot evaluation at step 0
+## Zero-shot evaluation at step 1,000
 
 The four lm-eval tasks use normalized accuracy when supplied by lm-eval
 0.4.12, with float32 weights and softmax. ArithMark
@@ -77,17 +77,17 @@ uses float32 weights and its official raw continuation log-likelihood-sum rule.
 
 | Benchmark | Accuracy |
 |---|---:|
-| HellaSwag | 26.31% |
-| ARC-Easy | 25.55% |
-| ARC-Challenge | 23.46% |
-| PIQA | 49.56% |
-| ArithMark-2 | 22.28% |
-| ARC mean | 24.51% |
-| Open SLM Leaderboard-style average | 30.66% |
+| HellaSwag | 27.81% |
+| ARC-Easy | 32.11% |
+| ARC-Challenge | 21.42% |
+| PIQA | 55.06% |
+| ArithMark-2 | 25.88% |
+| ARC mean | 26.76% |
+| Open SLM Leaderboard-style average | 33.88% |
 
 The average is `(HellaSwag + mean(ARC-Easy, ARC-Challenge) + PIQA +
 ArithMark-2) / 4`, matching the Open SLM Leaderboard formula.
 
-WikiText-103 validation at this step: loss 8.3686, perplexity
-4309.76, normalized BPB 3.7838 over 358,911
+WikiText-103 validation at this step: loss 3.5111, perplexity
+33.48, normalized BPB 1.5875 over 358,911
 scored tokens and 1,145,226 normalized UTF-8 bytes.

benchmark_results.json

@@ -1,45 +1,45 @@
 {
-  "step": 0,
+  "step": 1000,
   "parameters": 7849161,
   "lm_eval_version": "0.4.12",
   "num_fewshot": 0,
   "leaderboard_accuracy": {
-    "arc_easy": 0.25547138047138046,
-    "arc_challenge": 0.23464163822525597,
-    "hellaswag": 0.26309500099581756,
-    "piqa": 0.4956474428726877
+    "arc_easy": 0.32112794612794615,
+    "arc_challenge": 0.21416382252559726,
+    "hellaswag": 0.27813184624576776,
+    "piqa": 0.5505984766050055
   },
   "lm_eval_metric_policy": "acc_norm when available, otherwise acc",
   "selection_metric": "open_slm_leaderboard_average",
   "arithmark_2.0": {
-    "acc": 0.2228,
-    "correct": 557,
+    "acc": 0.2588,
+    "correct": 647,
     "total": 2500,
     "by_operator_count": {
       "1": {
-        "acc": 0.2208,
-        "correct": 276,
+        "acc": 0.248,
+        "correct": 310,
         "total": 1250
       },
       "2": {
-        "acc": 0.21866666666666668,
-        "correct": 164,
+        "acc": 0.2866666666666667,
+        "correct": 215,
         "total": 750
       },
       "3": {
-        "acc": 0.234,
-        "correct": 117,
+        "acc": 0.244,
+        "correct": 122,
         "total": 500
       }
     }
   },
-  "arc_average": 0.2450565093483182,
-  "open_slm_leaderboard_average": 0.30664973830420583,
+  "arc_average": 0.2676458843267717,
+  "open_slm_leaderboard_average": 0.33879405179438626,
   "average_formula": "(hellaswag + mean(arc_easy, arc_challenge) + piqa + arithmark_2.0) / 4",
   "validation": {
-    "loss": 8.368637380723772,
-    "perplexity": 4309.759498478308,
-    "normalized_bpb": 3.7837711245079073,
+    "loss": 3.511079413898325,
+    "perplexity": 33.4843918019419,
+    "normalized_bpb": 1.58748913326839,
     "tokens": 358911,
     "normalized_utf8_bytes": 1145226
   }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:767c860fea9367ec00d2e2b3d271f29fa6d926be279518e1fe15fc59a0b37b0f
+oid sha256:d3afa91d8c126df0efdf49b944fee4f838aadf2f6d5681ae9b1569fc5eee2942
 size 36330428

modeling_cma.py

@@ -171,15 +171,10 @@ class CMA(nn.Module):
         k = F.normalize(k.float(), dim=-1).to(v.dtype)
         scale = self.logit_scale.clamp(max=math.log(100.0)).exp()
         scale = scale.view(1, self.h, 1, 1)
-        routed = F.scaled_dot_product_attention(
-            q * scale.to(q.dtype),
-            k,
-            v,
-            attn_mask=self.bias.unsqueeze(0).to(q.dtype),
-            dropout_p=0.0,
-            is_causal=False,
-            scale=1.0,
-        )
+        attn_logits = (q @ k.transpose(-2, -1)) * scale.to(q.dtype)
+        attn_logits = attn_logits + self.bias.unsqueeze(0).to(q.dtype)
+        attn = F.softmax(attn_logits, dim=-1, dtype=torch.float32).to(v.dtype)
+        routed = attn @ v
         route_signal = (
             torch.einsum(
                 "bhnd,hd->bhn", q.float(), self.route_gate_weight.float()
@@ -191,11 +186,6 @@ class CMA(nn.Module):
         contribution = route_coeff.unsqueeze(-1) * (routed - v)
         y = v + contribution
         if self.record_diagnostics:
-            attn_logits = (
-                q.float() @ k.float().transpose(-2, -1)
-            ) * scale.float()
-            attn_logits = attn_logits + self.bias.float().unsqueeze(0)
-            attn = F.softmax(attn_logits, dim=-1, dtype=torch.float32)
             probs = attn.float().clamp_min(1e-9)
             entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
             diagonal_mass = probs.diagonal(dim1=-2, dim2=-1).mean()
@@ -241,10 +231,9 @@ class Block(nn.Module):
         )
 
     def forward(self, x, cos, sin):
-        # Expose the full attention -> residual -> norm -> CMA handoff to
-        # Inductor as one graph so pointwise boundaries can be fused.
-        x = torch.add(x, self.attn(self.n1(x), cos, sin))
-        return torch.add(x, self.mix(self.n2(x)))
+        x = x + self.attn(self.n1(x), cos, sin)
+        x = x + self.mix(self.n2(x))
+        return x
 
 
 class CMAModel(nn.Module):

training_config.json

@@ -51,8 +51,5 @@
   "data_seed": 1337,
   "shuffle_buffer": 50000,
   "tokenize_batch_size": 64,
-  "compile": true,
-  "compile_mode": "max-autotune",
-  "compile_fullgraph": true,
-  "compile_dynamic": false
+  "compile": true
 }