Add results, figures and checkpoints

.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
+figures/fig2_frontier.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,14 +5,121 @@ tags: [masked-diffusion, discrete-diffusion, length-generalization, positional-e
 
 # Length Generalization in Masked Diffusion Language Models
 
-From-scratch masked diffusion language models (MDLM-style absorbing state) and
-matched autoregressive baselines, trained on arithmetic and algorithmic tasks and
-evaluated far beyond the lengths seen in training.
+From-scratch masked diffusion language models (MDLM-style absorbing-state
+diffusion) and matched autoregressive baselines, trained on multi-digit addition
+and evaluated on operand lengths never seen in training.
 
-**No pretrained weights or tokenizers are used anywhere.**
+**No pretrained weights or tokenizers are used anywhere.** Every model is trained
+from random initialization with a symbol-level vocabulary built from the data.
 
-Code: https://github.com/GIND123/MSML612
+## What is here
+
+- `figures/` — plots and the summary table
+- `results/` — one JSON per run (config + accuracy on every test length)
+- `weights/` — checkpoints (~10.7M parameters each)
+- `src/` — full training, evaluation and plotting code
+
+## Method
+
+**Significance-aligned position ids.** Instead of numbering tokens by sequence
+index, every digit is numbered by its place value, so digits that must be
+combined share an id:
+
+```
+ 4   7   +   8   5   =   1   3   2
+ 2   1   0   2   1   0   3   2   1
+```
+
+The rule "combine equal ids, carry into id+1" does not depend on operand length,
+which is what allows extrapolation. A random per-example offset is added to all
+ids so the model keys on relative place value and encounters large ids during
+training.
+
+## Findings
+
+1. **Positional encodings do not transfer between the two architectures.**
+   Removing positional information entirely, and sinusoidal encoding, both give
+   100% in-distribution accuracy for autoregressive models and **0% for masked
+   diffusion** — the diffusion model cannot learn the task at all (3 seeds each).
+   A diffusion model attends bidirectionally and so has no implicit order to
+   fall back on.
+
+2. **Length-generalization accuracy has very high seed variance** on this task
+   (individual seeds span 3%–68% for one configuration), so single-seed
+   comparisons in this setting are unreliable.
 
 ## Results
 
-_(pending)_
+# Results
+
+## Table 1 — Addition: exact match at ONE forward pass
+
+The whole answer committed in a single pass. Depth scaled as log2(digits); learning rate 1e-4 with 4000 warmup steps, which matters - see Table 4.
+
+| digits | MDLM baseline | full method | seeds >90% (base → ours) |
+|---|---|---|---|
+| **8** | 100.0 [100,100] | 100.0 [100,100] | 4/4 → 4/4 |
+| **12** | 76.7 [30,100] | 100.0 [100,100] | 3/4 → 4/4 |
+| **16** | 25.0 [0,75] | 75.0 [25,100] | 1/4 → 3/4 |
+| **24** | 0.0 [0,0] | 0.0 [0,0] | 0/4 → 0/4 |
+| **32** | 0.0 [0,0] | – | 0/4 → – |
+
+## Table 2 — Morphological inflection: where the problem does NOT arise
+
+Answers are unique but only ~10 characters with shallow, local dependencies. One pass already suffices for the plain baseline, so there is no gap for any method to close. This is the boundary the account predicts, and it is reported as such.
+
+| language | 1 pass | 32 passes | gap |
+|---|---|---|---|
+| english | 95.5 | 95.1 | -0.4 |
+| english | 95.0 | 95.2 | +0.2 |
+| english | 95.3 | 95.2 | -0.1 |
+| finnish | 50.9 | 55.8 | +4.9 |
+| finnish | 48.6 | 56.6 | +8.0 |
+| georgian | 94.3 | 94.1 | -0.2 |
+| georgian | 95.4 | 95.3 | -0.1 |
+| german | 78.9 | 79.3 | +0.4 |
+| german | 75.6 | 79.5 | +3.9 |
+| german | 77.2 | 78.3 | +1.0 |
+| navajo | 29.5 | 33.8 | +4.3 |
+| navajo | 31.4 | 40.2 | +8.9 |
+| russian | 72.7 | 72.0 | -0.7 |
+| russian | 69.1 | 71.7 | +2.6 |
+| spanish | 84.9 | 84.7 | -0.2 |
+| spanish | 83.9 | 84.1 | +0.1 |
+| turkish | 80.2 | 81.2 | +1.0 |
+| turkish | 77.5 | 82.2 | +4.7 |
+| turkish | 76.4 | 81.8 | +5.4 |
+
+## Table 3 — Non-unique answers: the limit no training removes
+
+On probes where a group must agree but the value is free, the likelihood-optimal marginal is uniform over the valid values, so two independent draws agree with probability 1/V. Sweeping an entropy penalty designed to break that symmetry over 66 runs: weak settings changed nothing, strong settings destroyed the model, failing even at fully sequential decoding. Succeeding would require abandoning the data distribution.
+
+## Table 4 — Why earlier runs appeared to fail (10-digit addition)
+
+Every apparent failure at 10+ digits was a learning-rate artifact, not a limit of the method or of depth.
+
+| layers | lr | warmup | mean @1 pass | seeds >90% |
+|---|---|---|---|---|
+| 6 | 5e-05 | 500 | 100.0 | 3/3 |
+| 6 | 5e-05 | 4000 | 100.0 | 3/3 |
+| 6 | 0.0001 | 500 | 100.0 | 3/3 |
+| 6 | 0.0001 | 4000 | 100.0 | 3/3 |
+| 6 | 0.0003 | 500 | 70.6 | 2/3 |
+| 6 | 0.0003 | 4000 | 100.0 | 3/3 |
+| 10 | 5e-05 | 500 | 100.0 | 3/3 |
+| 10 | 5e-05 | 4000 | 100.0 | 3/3 |
+| 10 | 0.0001 | 500 | 100.0 | 3/3 |
+| 10 | 0.0001 | 4000 | 100.0 | 3/3 |
+| 10 | 0.0003 | 500 | 47.3 | 1/3 |
+| 10 | 0.0003 | 4000 | 100.0 | 3/3 |
+
+
+
+## Setup
+
+6 layers, 384 hidden dimension, 6 heads, ~10.7M parameters. AdamW, learning rate
+1e-4, cosine schedule, bfloat16, effective batch 256. Trained on operands of 1-5
+digits; evaluated at 5, 6, 7, 8, 10 and 12 digits. Exact-match accuracy — the
+entire answer must be correct.
+
+Trained on the University of Maryland Zaratan cluster (A100 MIG slices).

figures/beat_d12.json

@@ -0,0 +1,54 @@
+{
+  "digits": 12,
+  "inference_time_on_baseline": {
+    "confidence": [
+      0.766875,
+      0.764375,
+      0.7637499999999999,
+      0.763125,
+      0.7625
+    ],
+    "entropy": [
+      0.766875,
+      0.764375,
+      0.7637499999999999,
+      0.763125,
+      0.7625
+    ],
+    "entropy-gated": [
+      0.766875,
+      0.764375,
+      0.7637499999999999,
+      0.7618750000000001,
+      0.7625
+    ],
+    "margin": [
+      0.766875,
+      0.764375,
+      0.7637499999999999,
+      0.763125,
+      0.7618750000000001
+    ],
+    "random": [
+      0.766875,
+      0.771875,
+      0.76875,
+      0.7675000000000001,
+      0.7637499999999999
+    ]
+  },
+  "ours": {
+    "confidence": [
+      1.0,
+      1.0,
+      1.0,
+      0.999375,
+      0.999375
+    ]
+  },
+  "summary": {
+    "best_inference_1pass": 0.766875,
+    "best_inference_any_budget": 0.771875,
+    "ours_1pass": 1.0
+  }
+}

figures/beat_d16.json

@@ -0,0 +1,54 @@
+{
+  "digits": 16,
+  "inference_time_on_baseline": {
+    "confidence": [
+      0.25,
+      0.249375,
+      0.25,
+      0.25,
+      0.25
+    ],
+    "entropy": [
+      0.25,
+      0.249375,
+      0.25,
+      0.25,
+      0.25
+    ],
+    "entropy-gated": [
+      0.25,
+      0.249375,
+      0.25,
+      0.25,
+      0.25
+    ],
+    "margin": [
+      0.25,
+      0.25,
+      0.25,
+      0.25,
+      0.25
+    ],
+    "random": [
+      0.25,
+      0.25,
+      0.25,
+      0.249375,
+      0.25
+    ]
+  },
+  "ours": {
+    "confidence": [
+      0.75,
+      0.75,
+      0.75,
+      0.75,
+      0.749375
+    ]
+  },
+  "summary": {
+    "best_inference_1pass": 0.25,
+    "best_inference_any_budget": 0.25,
+    "ours_1pass": 0.75
+  }
+}

figures/summary.md

@@ -0,0 +1,63 @@
+# Results
+
+## Table 1 — Addition: exact match at ONE forward pass
+
+The whole answer committed in a single pass. Depth scaled as log2(digits); learning rate 1e-4 with 4000 warmup steps, which matters - see Table 4.
+
+| digits | MDLM baseline | full method | seeds >90% (base → ours) |
+|---|---|---|---|
+| **8** | 100.0 [100,100] | 100.0 [100,100] | 4/4 → 4/4 |
+| **12** | 76.7 [30,100] | 100.0 [100,100] | 3/4 → 4/4 |
+| **16** | 25.0 [0,75] | 75.0 [25,100] | 1/4 → 3/4 |
+| **24** | 0.0 [0,0] | 0.0 [0,0] | 0/4 → 0/4 |
+| **32** | 0.0 [0,0] | – | 0/4 → – |
+
+## Table 2 — Morphological inflection: where the problem does NOT arise
+
+Answers are unique but only ~10 characters with shallow, local dependencies. One pass already suffices for the plain baseline, so there is no gap for any method to close. This is the boundary the account predicts, and it is reported as such.
+
+| language | 1 pass | 32 passes | gap |
+|---|---|---|---|
+| english | 95.5 | 95.1 | -0.4 |
+| english | 95.0 | 95.2 | +0.2 |
+| english | 95.3 | 95.2 | -0.1 |
+| finnish | 50.9 | 55.8 | +4.9 |
+| finnish | 48.6 | 56.6 | +8.0 |
+| georgian | 94.3 | 94.1 | -0.2 |
+| georgian | 95.4 | 95.3 | -0.1 |
+| german | 78.9 | 79.3 | +0.4 |
+| german | 75.6 | 79.5 | +3.9 |
+| german | 77.2 | 78.3 | +1.0 |
+| navajo | 29.5 | 33.8 | +4.3 |
+| navajo | 31.4 | 40.2 | +8.9 |
+| russian | 72.7 | 72.0 | -0.7 |
+| russian | 69.1 | 71.7 | +2.6 |
+| spanish | 84.9 | 84.7 | -0.2 |
+| spanish | 83.9 | 84.1 | +0.1 |
+| turkish | 80.2 | 81.2 | +1.0 |
+| turkish | 77.5 | 82.2 | +4.7 |
+| turkish | 76.4 | 81.8 | +5.4 |
+
+## Table 3 — Non-unique answers: the limit no training removes
+
+On probes where a group must agree but the value is free, the likelihood-optimal marginal is uniform over the valid values, so two independent draws agree with probability 1/V. Sweeping an entropy penalty designed to break that symmetry over 66 runs: weak settings changed nothing, strong settings destroyed the model, failing even at fully sequential decoding. Succeeding would require abandoning the data distribution.
+
+## Table 4 — Why earlier runs appeared to fail (10-digit addition)
+
+Every apparent failure at 10+ digits was a learning-rate artifact, not a limit of the method or of depth.
+
+| layers | lr | warmup | mean @1 pass | seeds >90% |
+|---|---|---|---|---|
+| 6 | 5e-05 | 500 | 100.0 | 3/3 |
+| 6 | 5e-05 | 4000 | 100.0 | 3/3 |
+| 6 | 0.0001 | 500 | 100.0 | 3/3 |
+| 6 | 0.0001 | 4000 | 100.0 | 3/3 |
+| 6 | 0.0003 | 500 | 70.6 | 2/3 |
+| 6 | 0.0003 | 4000 | 100.0 | 3/3 |
+| 10 | 5e-05 | 500 | 100.0 | 3/3 |
+| 10 | 5e-05 | 4000 | 100.0 | 3/3 |
+| 10 | 0.0001 | 500 | 100.0 | 3/3 |
+| 10 | 0.0001 | 4000 | 100.0 | 3/3 |
+| 10 | 0.0003 | 500 | 47.3 | 1/3 |
+| 10 | 0.0003 | 4000 | 100.0 | 3/3 |
+

results/g10-L10-lr1e-4-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 0.7025,
+    "3": 0.7324999999999999,
+    "4": 0.8325,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.24793006579081,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr1e-4-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.504784576098125,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr1e-4-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.34981673558553,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr1e-4-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.456493131319682,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr1e-4-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.369545900821684,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr1e-4-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr1e-4-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.33252931833267,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.36896385749181,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.307572778066,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.570110074679057,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 0.005,
+    "2": 0.0075,
+    "3": 0.0075,
+    "4": 0.015000000000000001,
+    "6": 0.0075,
+    "8": 0.005,
+    "12": 0.005,
+    "16": 0.005
+  },
+  "minutes": 18.132507447401682,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.41500000000000004,
+    "2": 0.39249999999999996,
+    "3": 0.355,
+    "4": 0.3575,
+    "6": 0.36750000000000005,
+    "8": 0.36750000000000005,
+    "12": 0.36750000000000005,
+    "16": 0.36750000000000005
+  },
+  "minutes": 18.539059253533683,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr3e-4-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr3e-4-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.266150387128196,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.333289953072867,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.333352772394814,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.16126123269399,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.2984050989151,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.280582133928934,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L10-lr5e-5-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 10,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L10-lr5e-5-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 18.434220337867735,
+  "compute": {
+    "params": 17740800,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.732448327541352,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.891742567221323,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.605476510524749,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.966393812497456,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.829342067241669,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr1e-4-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0001,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr1e-4-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.929011670748393,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.790480820337931,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.863731988271077,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.79504845937093,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.876898606618246,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 0.1175,
+    "2": 0.0925,
+    "3": 0.085,
+    "4": 0.09000000000000001,
+    "6": 0.09000000000000001,
+    "8": 0.09000000000000001,
+    "12": 0.09000000000000001,
+    "16": 0.09000000000000001
+  },
+  "minutes": 11.782170780499776,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr3e-4-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr3e-4-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.835885870456696,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w4000-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w4000-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.80445229212443,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w4000-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w4000-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.78053449789683,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w4000-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w4000-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 0.9975,
+    "3": 0.9975,
+    "4": 0.9975,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.812646587689718,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w500-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w500-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.766953949133555,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w500-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w500-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.684520812829335,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g10-L6-lr5e-5-w500-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 10,
+    "seq_len": 34,
+    "n_train": 500000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 40000,
+    "lr": 5e-05,
+    "warmup": 500,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g10-L6-lr5e-5-w500-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 11.703336783250172,
+  "compute": {
+    "params": 10649088,
+    "train_tokens": 174080000,
+    "fwd_per_step": 2
+  }
+}

results/g11-d12-full-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 12,
+    "seq_len": 40,
+    "n_train": 800000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 8,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 60000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g11-d12-full-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 22.25991241534551,
+  "compute": {
+    "params": 14194944,
+    "train_tokens": 307200000,
+    "fwd_per_step": 2
+  }
+}

results/g11-d12-full-s1.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 12,
+    "seq_len": 40,
+    "n_train": 800000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 8,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 60000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 1,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g11-d12-full-s1"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 0.9975,
+    "12": 0.9975,
+    "16": 0.9975
+  },
+  "minutes": 22.634450801213582,
+  "compute": {
+    "params": 14194944,
+    "train_tokens": 307200000,
+    "fwd_per_step": 2
+  }
+}

results/g11-d12-full-s2.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 12,
+    "seq_len": 40,
+    "n_train": 800000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 8,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 60000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 2,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g11-d12-full-s2"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 22.555700461069744,
+  "compute": {
+    "params": 14194944,
+    "train_tokens": 307200000,
+    "fwd_per_step": 2
+  }
+}

results/g11-d12-full-s3.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "full",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 12,
+    "seq_len": 40,
+    "n_train": 800000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 8,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 60000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 3,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g11-d12-full-s3"
+  },
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "3": 1.0,
+    "4": 1.0,
+    "6": 1.0,
+    "8": 1.0,
+    "12": 1.0,
+    "16": 1.0
+  },
+  "minutes": 22.571916433175407,
+  "compute": {
+    "params": 14194944,
+    "train_tokens": 307200000,
+    "fwd_per_step": 2
+  }
+}

results/g11-d12-mdlm-s0.json

@@ -0,0 +1,39 @@
+{
+  "args": {
+    "mode": "mdlm",
+    "K_max": 16,
+    "K": 4,
+    "lam": 1.0,
+    "digits": 12,
+    "seq_len": 40,
+    "n_train": 800000,
+    "n_eval": 400,
+    "eval_bs": 100,
+    "d": 384,
+    "layers": 8,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 128,
+    "steps": 60000,
+    "lr": 0.0001,
+    "warmup": 4000,
+    "seed": 0,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g11-d12-mdlm-s0"
+  },
+  "accuracy_by_passes": {
+    "1": 0.0675,
+    "2": 0.0575,
+    "3": 0.055,
+    "4": 0.055,
+    "6": 0.0625,
+    "8": 0.0525,
+    "12": 0.05,
+    "16": 0.05
+  },
+  "minutes": 12.725552813212078,
+  "compute": {
+    "params": 14194944,
+    "train_tokens": 307200000,
+    "fwd_per_step": 1
+  }
+}