diff --git a/README.md b/README.md index e8610783099c8f22968722274152b0c6c09745e9..0978a606bbfac682ba9a6817c7fc0dea5be7e8f9 100644 --- a/README.md +++ b/README.md @@ -40,19 +40,43 @@ Answers are unique but only ~10 characters with shallow, local dependencies. One | english | 95.3 | 95.2 | -0.1 | | finnish | 50.9 | 55.8 | +4.9 | | finnish | 48.6 | 56.6 | +8.0 | +| finnish | 50.9 | 55.8 | +4.9 | +| finnish | 52.9 | 57.9 | +4.9 | +| finnish | 48.6 | 56.6 | +8.0 | +| finnish | 54.3 | 61.7 | +7.4 | | georgian | 94.3 | 94.1 | -0.2 | | georgian | 95.4 | 95.3 | -0.1 | +| georgian | 94.3 | 94.1 | -0.2 | +| georgian | 93.8 | 94.0 | +0.1 | +| georgian | 95.4 | 95.3 | -0.1 | +| georgian | 93.9 | 93.9 | +0.0 | | german | 78.9 | 79.3 | +0.4 | | german | 75.6 | 79.5 | +3.9 | +| german | 78.9 | 79.3 | +0.4 | +| german | 77.9 | 78.0 | +0.1 | +| german | 75.6 | 79.5 | +3.9 | +| german | 76.4 | 78.9 | +2.5 | | german | 77.2 | 78.3 | +1.0 | | navajo | 29.5 | 33.8 | +4.3 | | navajo | 31.4 | 40.2 | +8.9 | +| navajo | 29.5 | 33.8 | +4.3 | +| navajo | 24.3 | 29.0 | +4.8 | +| navajo | 31.4 | 40.2 | +8.9 | +| navajo | 24.6 | 31.7 | +7.1 | | russian | 72.7 | 72.0 | -0.7 | | russian | 69.1 | 71.7 | +2.6 | +| russian | 72.7 | 72.0 | -0.7 | +| russian | 71.9 | 72.0 | +0.1 | +| russian | 69.1 | 71.7 | +2.6 | +| russian | 74.1 | 75.9 | +1.8 | | 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 | 80.2 | 81.2 | +1.0 | +| turkish | 84.5 | 85.1 | +0.7 | +| turkish | 77.5 | 82.2 | +4.7 | +| turkish | 79.9 | 82.2 | +2.3 | | turkish | 76.4 | 81.8 | +5.4 | ## Table 3 — Non-unique answers: the limit no training removes @@ -78,4 +102,24 @@ Every apparent failure at 10+ digits was a learning-rate artifact, not a limit o | 10 | 0.0003 | 500 | 47.3 | 1/3 | | 10 | 0.0003 | 4000 | 100.0 | 3/3 | +## Table 6 — Minimum depth for one-pass addition + +The account in THEORY.md predicts minimum depth growing as log2(n), since the carry is a prefix scan. A ripple-carry account predicts linear growth; an impossibility account predicts no depth suffices. + +| digits | 2L | 3L | 4L | 6L | 8L | log2(n) | min depth >90% | +|---|---|---|---|---|---|---|---| +| **4** | 100 | 100 | 100 | 100 | 100 | 2.0 | 2 | +| **8** | 99 | 100 | – | – | – | 3.0 | 2 | + +## Table 7 — Which component carries the gain (inflection) + +| language | mdlm | matched | distill | full | (gap 1→32) | +|---|---|---|---|---|---| +| finnish | 51.4 | – | – | 51.9 | +7.7 / – / – / +4.9 | +| georgian | 94.6 | – | – | 94.1 | -0.0 / – / – / -0.0 | +| german | 76.0 | – | – | 78.4 | +3.2 / – / – / +0.2 | +| navajo | 28.0 | – | – | 26.9 | +8.0 / – / – / +4.5 | +| russian | 71.6 | – | – | 72.3 | +2.2 / – / – / -0.3 | +| turkish | 78.7 | – | – | 82.3 | +3.5 / – / – / +0.8 | + diff --git a/figures/fig3_depth_scaling.png b/figures/fig3_depth_scaling.png new file mode 100644 index 0000000000000000000000000000000000000000..71369b22d386385dbd1b2c2760ce44638397fb05 Binary files /dev/null and b/figures/fig3_depth_scaling.png differ diff --git a/figures/summary.md b/figures/summary.md index ba55ab1cdb6e0742d6d37d4fcbabfb0f08f43235..f0626a6a39b8596f42ebeb1971b63abf2107f77b 100644 --- a/figures/summary.md +++ b/figures/summary.md @@ -23,19 +23,43 @@ Answers are unique but only ~10 characters with shallow, local dependencies. One | english | 95.3 | 95.2 | -0.1 | | finnish | 50.9 | 55.8 | +4.9 | | finnish | 48.6 | 56.6 | +8.0 | +| finnish | 50.9 | 55.8 | +4.9 | +| finnish | 52.9 | 57.9 | +4.9 | +| finnish | 48.6 | 56.6 | +8.0 | +| finnish | 54.3 | 61.7 | +7.4 | | georgian | 94.3 | 94.1 | -0.2 | | georgian | 95.4 | 95.3 | -0.1 | +| georgian | 94.3 | 94.1 | -0.2 | +| georgian | 93.8 | 94.0 | +0.1 | +| georgian | 95.4 | 95.3 | -0.1 | +| georgian | 93.9 | 93.9 | +0.0 | | german | 78.9 | 79.3 | +0.4 | | german | 75.6 | 79.5 | +3.9 | +| german | 78.9 | 79.3 | +0.4 | +| german | 77.9 | 78.0 | +0.1 | +| german | 75.6 | 79.5 | +3.9 | +| german | 76.4 | 78.9 | +2.5 | | german | 77.2 | 78.3 | +1.0 | | navajo | 29.5 | 33.8 | +4.3 | | navajo | 31.4 | 40.2 | +8.9 | +| navajo | 29.5 | 33.8 | +4.3 | +| navajo | 24.3 | 29.0 | +4.8 | +| navajo | 31.4 | 40.2 | +8.9 | +| navajo | 24.6 | 31.7 | +7.1 | | russian | 72.7 | 72.0 | -0.7 | | russian | 69.1 | 71.7 | +2.6 | +| russian | 72.7 | 72.0 | -0.7 | +| russian | 71.9 | 72.0 | +0.1 | +| russian | 69.1 | 71.7 | +2.6 | +| russian | 74.1 | 75.9 | +1.8 | | 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 | 80.2 | 81.2 | +1.0 | +| turkish | 84.5 | 85.1 | +0.7 | +| turkish | 77.5 | 82.2 | +4.7 | +| turkish | 79.9 | 82.2 | +2.3 | | turkish | 76.4 | 81.8 | +5.4 | ## Table 3 — Non-unique answers: the limit no training removes @@ -61,3 +85,23 @@ Every apparent failure at 10+ digits was a learning-rate artifact, not a limit o | 10 | 0.0003 | 500 | 47.3 | 1/3 | | 10 | 0.0003 | 4000 | 100.0 | 3/3 | +## Table 6 — Minimum depth for one-pass addition + +The account in THEORY.md predicts minimum depth growing as log2(n), since the carry is a prefix scan. A ripple-carry account predicts linear growth; an impossibility account predicts no depth suffices. + +| digits | 2L | 3L | 4L | 6L | 8L | log2(n) | min depth >90% | +|---|---|---|---|---|---|---|---| +| **4** | 100 | 100 | 100 | 100 | 100 | 2.0 | 2 | +| **8** | 99 | 100 | – | – | – | 3.0 | 2 | + +## Table 7 — Which component carries the gain (inflection) + +| language | mdlm | matched | distill | full | (gap 1→32) | +|---|---|---|---|---|---| +| finnish | 51.4 | – | – | 51.9 | +7.7 / – / – / +4.9 | +| georgian | 94.6 | – | – | 94.1 | -0.0 / – / – / -0.0 | +| german | 76.0 | – | – | 78.4 | +3.2 / – / – / +0.2 | +| navajo | 28.0 | – | – | 26.9 | +8.0 / – / – / +4.5 | +| russian | 71.6 | – | – | 72.3 | +2.2 / – / – / -0.3 | +| turkish | 78.7 | – | – | 82.3 | +3.5 / – / – / +0.8 | + diff --git a/results/g14-d4-L2-s0.json b/results/g14-d4-L2-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..180b7201394c39240dd957926a03bd507d7ca685 --- /dev/null +++ b/results/g14-d4-L2-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 2, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L2-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": 6.882920638720194, + "compute": { + "params": 3557376, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L2-s1.json b/results/g14-d4-L2-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..c377689e70df6f78578be9e3a45d3cfe310e42e7 --- /dev/null +++ b/results/g14-d4-L2-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 2, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L2-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": 6.816111731529236, + "compute": { + "params": 3557376, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L3-s0.json b/results/g14-d4-L3-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..30b6e65ea18ed86a9c1e45fd48afce45cce24038 --- /dev/null +++ b/results/g14-d4-L3-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 3, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L3-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": 9.519309445222218, + "compute": { + "params": 5330304, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L3-s1.json b/results/g14-d4-L3-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..e68f81254e5fb534a239ecc383b0fa416519972a --- /dev/null +++ b/results/g14-d4-L3-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 3, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L3-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": 9.63311020930608, + "compute": { + "params": 5330304, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L4-s0.json b/results/g14-d4-L4-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..864bb160442917cb968f1ba51c931f85e19dd6a8 --- /dev/null +++ b/results/g14-d4-L4-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 4, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L4-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": 12.398685284455617, + "compute": { + "params": 7103232, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L4-s1.json b/results/g14-d4-L4-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..683494dd108a6a6526918ded24f10750db894c1b --- /dev/null +++ b/results/g14-d4-L4-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 4, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L4-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": 12.252778033415476, + "compute": { + "params": 7103232, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L6-s0.json b/results/g14-d4-L6-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..ddb0159804f980b7b685c66bcbf3784717fd14ee --- /dev/null +++ b/results/g14-d4-L6-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 6, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L6-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": 16.965274214744568, + "compute": { + "params": 10649088, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L6-s1.json b/results/g14-d4-L6-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..bb3ae8469669c5cc43b585713a7c968c4d140fa4 --- /dev/null +++ b/results/g14-d4-L6-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 6, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d4-L6-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": 17.125181821982064, + "compute": { + "params": 10649088, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L8-s0.json b/results/g14-d4-L8-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..e7b9bd65fdca4d6fb68f646862211ea83f00f2a0 --- /dev/null +++ b/results/g14-d4-L8-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "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/g14-d4-L8-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": 21.700168172518413, + "compute": { + "params": 14194944, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d4-L8-s1.json b/results/g14-d4-L8-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..aec660ef3c12c9dd721e8d087a5468815879a1ba --- /dev/null +++ b/results/g14-d4-L8-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 4, + "seq_len": 16, + "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/g14-d4-L8-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": 21.729793294270834, + "compute": { + "params": 14194944, + "train_tokens": 122880000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L2-s0.json b/results/g14-d8-L2-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..0f20cc0798d02b8546c04e6ce855812b47a1cb69 --- /dev/null +++ b/results/g14-d8-L2-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 2, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L2-s0" + }, + "accuracy_by_passes": { + "1": 0.99, + "2": 0.9975, + "3": 1.0, + "4": 1.0, + "6": 1.0, + "8": 1.0, + "12": 1.0, + "16": 1.0 + }, + "minutes": 7.338654283682505, + "compute": { + "params": 3557376, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L2-s1.json b/results/g14-d8-L2-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..a60df0a3f64ee9c6e941e54309d431adf37aa86c --- /dev/null +++ b/results/g14-d8-L2-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 2, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L2-s1" + }, + "accuracy_by_passes": { + "1": 0.9975, + "2": 1.0, + "3": 1.0, + "4": 1.0, + "6": 1.0, + "8": 1.0, + "12": 1.0, + "16": 1.0 + }, + "minutes": 7.30421595176061, + "compute": { + "params": 3557376, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L3-s0.json b/results/g14-d8-L3-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..4380775c34335599b68067923f72be66d5f5c21d --- /dev/null +++ b/results/g14-d8-L3-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 3, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L3-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": 9.942399954795837, + "compute": { + "params": 5330304, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L3-s1.json b/results/g14-d8-L3-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..7019a4734671e6962c5072dab2f44a415b9e418c --- /dev/null +++ b/results/g14-d8-L3-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 3, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L3-s1" + }, + "accuracy_by_passes": { + "1": 0.9975, + "2": 1.0, + "3": 0.9975, + "4": 0.9975, + "6": 0.9975, + "8": 0.9975, + "12": 0.9975, + "16": 0.9975 + }, + "minutes": 9.788207534948985, + "compute": { + "params": 5330304, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L4-s0.json b/results/g14-d8-L4-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..7258ef60887f1ffb1b82d2f2166d40c85a481dc8 --- /dev/null +++ b/results/g14-d8-L4-s0.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 4, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 0, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L4-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": 12.345258478323618, + "compute": { + "params": 7103232, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g14-d8-L4-s1.json b/results/g14-d8-L4-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..3ac933858af2b7f6c8bf1981d3bdffbd35519a89 --- /dev/null +++ b/results/g14-d8-L4-s1.json @@ -0,0 +1,39 @@ +{ + "args": { + "mode": "full", + "K_max": 16, + "K": 4, + "lam": 1.0, + "digits": 8, + "seq_len": 28, + "n_train": 800000, + "n_eval": 400, + "eval_bs": 100, + "d": 384, + "layers": 4, + "heads": 6, + "pe": "rope", + "bs": 128, + "steps": 60000, + "lr": 0.0001, + "warmup": 4000, + "seed": 1, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g14-d8-L4-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": 12.423506192366283, + "compute": { + "params": 7103232, + "train_tokens": 215040000, + "fwd_per_step": 2 + } +} \ No newline at end of file diff --git a/results/g15-finnish-full-s0.json b/results/g15-finnish-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..5b12332aefec3a2299019df76ec0c2668bd01cb9 --- /dev/null +++ b/results/g15-finnish-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "finnish", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-finnish-full-s0" + }, + "accuracy_by_passes": { + "1": 0.5093256353021978, + "2": 0.511461195054945, + "4": 0.5276978880494505, + "8": 0.5495793269230769, + "16": 0.5583683894230769, + "32": 0.5583683894230769, + "ar": 0.05099587912087912 + }, + "minutes": 4.275957576433817, + "compute": { + "params": 4769792, + "answer_len_mean": 15.12968635559082 + } +} \ No newline at end of file diff --git a/results/g15-finnish-full-s1.json b/results/g15-finnish-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..8208f04d24376c7ef6363b71d44aa12193d07ecc --- /dev/null +++ b/results/g15-finnish-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "finnish", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-finnish-full-s1" + }, + "accuracy_by_passes": { + "1": 0.5292217548076923, + "2": 0.5335250686813187, + "4": 0.549547132554945, + "8": 0.5708490728021978, + "16": 0.5776850103021978, + "32": 0.5786615728021978, + "ar": 0.05548162774725275 + }, + "minutes": 4.345337744553884, + "compute": { + "params": 4769792, + "answer_len_mean": 15.12968635559082 + } +} \ No newline at end of file diff --git a/results/g15-finnish-mdlm-s0.json b/results/g15-finnish-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..4fcc05f75be997255b7c2638be856ff466101855 --- /dev/null +++ b/results/g15-finnish-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "finnish", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-finnish-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.4858237465659341, + "2": 0.5041959993131868, + "4": 0.5329133756868132, + "8": 0.545640882554945, + "16": 0.5636160714285714, + "32": 0.5655691964285714, + "ar": 0.05392556662087912 + }, + "minutes": 2.4385203957557677, + "compute": { + "params": 4769792, + "answer_len_mean": 15.12968635559082 + } +} \ No newline at end of file diff --git a/results/g15-finnish-mdlm-s1.json b/results/g15-finnish-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..b7c83c111c664968a940eb8a89bba6e91a9b6cd9 --- /dev/null +++ b/results/g15-finnish-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "finnish", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-finnish-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.5428936298076923, + "2": 0.5407902644230769, + "4": 0.5704842032967032, + "8": 0.5990191449175823, + "16": 0.6165972699175823, + "32": 0.6165972699175823, + "ar": 0.06173806662087912 + }, + "minutes": 2.4639089981714886, + "compute": { + "params": 4769792, + "answer_len_mean": 15.12968635559082 + } +} \ No newline at end of file diff --git a/results/g15-georgian-full-s0.json b/results/g15-georgian-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..c03785840c10fd382719138dbd843d8b9312b8d8 --- /dev/null +++ b/results/g15-georgian-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "georgian", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-georgian-full-s0" + }, + "accuracy_by_passes": { + "1": 0.9432091346153846, + "2": 0.9434344951923077, + "4": 0.9402794471153846, + "8": 0.9424579326923077, + "16": 0.9412560096153846, + "32": 0.9412560096153846, + "ar": 0.1258263221153846 + }, + "minutes": 4.286119385560354, + "compute": { + "params": 4766208, + "answer_len_mean": 9.932000160217285 + } +} \ No newline at end of file diff --git a/results/g15-georgian-full-s1.json b/results/g15-georgian-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..ec669f2f57637f9c2bc54c1c7c3e8766417028f8 --- /dev/null +++ b/results/g15-georgian-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "georgian", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-georgian-full-s1" + }, + "accuracy_by_passes": { + "1": 0.9383263221153846, + "2": 0.9385516826923077, + "4": 0.9385516826923077, + "8": 0.9395282451923077, + "16": 0.9395282451923077, + "32": 0.9395282451923077, + "ar": 0.12387319711538461 + }, + "minutes": 4.304093313217163, + "compute": { + "params": 4766208, + "answer_len_mean": 9.932000160217285 + } +} \ No newline at end of file diff --git a/results/g15-georgian-mdlm-s0.json b/results/g15-georgian-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..2f6d75988bae508625684b43b0db7eb838757d2b --- /dev/null +++ b/results/g15-georgian-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "georgian", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-georgian-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.9539513221153846, + "2": 0.9522235576923077, + "4": 0.9502704326923077, + "8": 0.9532001201923077, + "16": 0.9532001201923077, + "32": 0.9532001201923077, + "ar": 0.1258263221153846 + }, + "minutes": 2.4665067235628766, + "compute": { + "params": 4766208, + "answer_len_mean": 9.932000160217285 + } +} \ No newline at end of file diff --git a/results/g15-georgian-mdlm-s1.json b/results/g15-georgian-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..4b5bc9e4bb6b0540a6a087d3080cbfd0914434c9 --- /dev/null +++ b/results/g15-georgian-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "georgian", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-georgian-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.9385516826923077, + "2": 0.9375751201923077, + "4": 0.9375751201923077, + "8": 0.9395282451923077, + "16": 0.9385516826923077, + "32": 0.9385516826923077, + "ar": 0.12387319711538461 + }, + "minutes": 2.4444581945737203, + "compute": { + "params": 4766208, + "answer_len_mean": 9.932000160217285 + } +} \ No newline at end of file diff --git a/results/g15-german-full-s0.json b/results/g15-german-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..b0401b9443513a1ad0cfaa814337f8f42ed334e7 --- /dev/null +++ b/results/g15-german-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "german", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-german-full-s0" + }, + "accuracy_by_passes": { + "1": 0.7893359375, + "2": 0.792265625, + "4": 0.7903125, + "8": 0.79296875, + "16": 0.79296875, + "32": 0.79296875, + "ar": 0.0776171875 + }, + "minutes": 4.36189177831014, + "compute": { + "params": 4776448, + "answer_len_mean": 10.743975639343262 + } +} \ No newline at end of file diff --git a/results/g15-german-full-s1.json b/results/g15-german-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..8e041308c20e997195fdd9f9cb253914a149f156 --- /dev/null +++ b/results/g15-german-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "german", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-german-full-s1" + }, + "accuracy_by_passes": { + "1": 0.7788671875, + "2": 0.7705078125, + "4": 0.7734375, + "8": 0.7795703125, + "16": 0.7795703125, + "32": 0.7795703125, + "ar": 0.0724609375 + }, + "minutes": 4.429418742656708, + "compute": { + "params": 4776448, + "answer_len_mean": 10.743975639343262 + } +} \ No newline at end of file diff --git a/results/g15-german-mdlm-s0.json b/results/g15-german-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..fa59399d6f22da6758313c08ee360b10a3a775a5 --- /dev/null +++ b/results/g15-german-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "german", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-german-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.7555859375, + "2": 0.7710546875, + "4": 0.7808203125, + "8": 0.7911328125, + "16": 0.7950390625, + "32": 0.7950390625, + "ar": 0.076640625 + }, + "minutes": 2.526710093021393, + "compute": { + "params": 4776448, + "answer_len_mean": 10.743975639343262 + } +} \ No newline at end of file diff --git a/results/g15-german-mdlm-s1.json b/results/g15-german-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..3f051b46566087e1cd8625108b74291e6823294a --- /dev/null +++ b/results/g15-german-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "german", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-german-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.7641015625, + "2": 0.76421875, + "4": 0.7815234375, + "8": 0.786953125, + "16": 0.7886328125, + "32": 0.7886328125, + "ar": 0.0769140625 + }, + "minutes": 2.5262390971183777, + "compute": { + "params": 4776448, + "answer_len_mean": 10.743975639343262 + } +} \ No newline at end of file diff --git a/results/g15-navajo-full-s0.json b/results/g15-navajo-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..f1d9322891ec2163c232a6d2e7161aa5f5b15170 --- /dev/null +++ b/results/g15-navajo-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "navajo", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-navajo-full-s0" + }, + "accuracy_by_passes": { + "1": 0.29529747596153844, + "2": 0.30814302884615385, + "4": 0.3108473557692308, + "8": 0.33743990384615385, + "16": 0.33841646634615385, + "32": 0.33841646634615385, + "ar": 0.06332632211538461 + }, + "minutes": 4.317290306091309, + "compute": { + "params": 4765696, + "answer_len_mean": 10.005999565124512 + } +} \ No newline at end of file diff --git a/results/g15-navajo-full-s1.json b/results/g15-navajo-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..a423536517b1bd982dba344c0357270d303d3161 --- /dev/null +++ b/results/g15-navajo-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "navajo", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-navajo-full-s1" + }, + "accuracy_by_passes": { + "1": 0.24263822115384615, + "2": 0.25045072115384615, + "4": 0.265625, + "8": 0.2872596153846154, + "16": 0.2901893028846154, + "32": 0.2901893028846154, + "ar": 0.0625751201923077 + }, + "minutes": 4.282497262954712, + "compute": { + "params": 4765696, + "answer_len_mean": 10.005999565124512 + } +} \ No newline at end of file diff --git a/results/g15-navajo-mdlm-s0.json b/results/g15-navajo-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..1dc89f80306c0714a2e32356f121568253662b9c --- /dev/null +++ b/results/g15-navajo-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "navajo", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-navajo-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.3137019230769231, + "2": 0.3273737980769231, + "4": 0.37169471153846156, + "8": 0.4004657451923077, + "16": 0.4024188701923077, + "32": 0.4024188701923077, + "ar": 0.08473557692307693 + }, + "minutes": 2.4744708220163982, + "compute": { + "params": 4765696, + "answer_len_mean": 10.005999565124512 + } +} \ No newline at end of file diff --git a/results/g15-navajo-mdlm-s1.json b/results/g15-navajo-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..4c556b0fe16e8642f12e5d94aa047b461b6d9235 --- /dev/null +++ b/results/g15-navajo-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "navajo", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-navajo-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.24579326923076922, + "2": 0.2675030048076923, + "4": 0.29725060096153844, + "8": 0.3133263221153846, + "16": 0.3172325721153846, + "32": 0.3172325721153846, + "ar": 0.06475360576923077 + }, + "minutes": 2.43903698126475, + "compute": { + "params": 4765696, + "answer_len_mean": 10.005999565124512 + } +} \ No newline at end of file diff --git a/results/g15-russian-full-s0.json b/results/g15-russian-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..f2b44872e73a834cebb4b4931932b34ac881c1b2 --- /dev/null +++ b/results/g15-russian-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "russian", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-russian-full-s0" + }, + "accuracy_by_passes": { + "1": 0.7271740845959596, + "2": 0.7232678345959596, + "4": 0.7213147095959596, + "8": 0.7218868371212122, + "16": 0.7199337121212122, + "32": 0.7199337121212122, + "ar": 0.03911182133838384 + }, + "minutes": 4.331714260578155, + "compute": { + "params": 4770304, + "answer_len_mean": 11.276381492614746 + } +} \ No newline at end of file diff --git a/results/g15-russian-full-s1.json b/results/g15-russian-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..6cf9635d9f793d8de1f703c09bceebd3867a6889 --- /dev/null +++ b/results/g15-russian-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "russian", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-russian-full-s1" + }, + "accuracy_by_passes": { + "1": 0.7187894570707071, + "2": 0.7197660195707071, + "4": 0.7158597695707071, + "8": 0.7190755208333333, + "16": 0.7200520833333333, + "32": 0.7200520833333333, + "ar": 0.03813525883838384 + }, + "minutes": 4.2800087610880535, + "compute": { + "params": 4770304, + "answer_len_mean": 11.276381492614746 + } +} \ No newline at end of file diff --git a/results/g15-russian-mdlm-s0.json b/results/g15-russian-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..76263ef95eb583969e1b6555c1464874e0817c68 --- /dev/null +++ b/results/g15-russian-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "russian", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-russian-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.6910412720959596, + "2": 0.7068339646464646, + "4": 0.7132161458333333, + "8": 0.7106908933080809, + "16": 0.7165502683080809, + "32": 0.7165502683080809, + "ar": 0.035896070075757576 + }, + "minutes": 2.446616017818451, + "compute": { + "params": 4770304, + "answer_len_mean": 11.276381492614746 + } +} \ No newline at end of file diff --git a/results/g15-russian-mdlm-s1.json b/results/g15-russian-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..edcd67339174048ab9740b3e52618833d28615d2 --- /dev/null +++ b/results/g15-russian-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "russian", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-russian-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.7412997159090909, + "2": 0.7453243371212122, + "4": 0.7542810921717171, + "8": 0.7633562184343434, + "16": 0.7591639046717171, + "32": 0.7591639046717171, + "ar": 0.03618213383838384 + }, + "minutes": 2.493031601111094, + "compute": { + "params": 4770304, + "answer_len_mean": 11.276381492614746 + } +} \ No newline at end of file diff --git a/results/g15-turkish-full-s0.json b/results/g15-turkish-full-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..93b9f6363ebed52e7641cb66cc437e1e4838f763 --- /dev/null +++ b/results/g15-turkish-full-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "turkish", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-turkish-full-s0" + }, + "accuracy_by_passes": { + "1": 0.8020938340053764, + "2": 0.7859753024193549, + "4": 0.7945228494623656, + "8": 0.8103998655913979, + "16": 0.8119854670698925, + "32": 0.8119854670698925, + "ar": 0.10095346102150538 + }, + "minutes": 4.366493237018585, + "compute": { + "params": 4772864, + "answer_len_mean": 13.684530258178711 + } +} \ No newline at end of file diff --git a/results/g15-turkish-full-s1.json b/results/g15-turkish-full-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..c4c18f29bd656309b999c3bf873a625b6959d3f7 --- /dev/null +++ b/results/g15-turkish-full-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "turkish", + "size": "medium", + "mode": "full", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-turkish-full-s1" + }, + "accuracy_by_passes": { + "1": 0.844705561155914, + "2": 0.8401902721774194, + "4": 0.841775873655914, + "8": 0.8499558971774194, + "16": 0.8512999831989247, + "32": 0.8512999831989247, + "ar": 0.1005859375 + }, + "minutes": 4.334676361083984, + "compute": { + "params": 4772864, + "answer_len_mean": 13.684530258178711 + } +} \ No newline at end of file diff --git a/results/g15-turkish-mdlm-s0.json b/results/g15-turkish-mdlm-s0.json new file mode 100644 index 0000000000000000000000000000000000000000..7ba7ee4a3195b7c34a82924c905203bf3ef9fa42 --- /dev/null +++ b/results/g15-turkish-mdlm-s0.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "turkish", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 0, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-turkish-mdlm-s0" + }, + "accuracy_by_passes": { + "1": 0.7749915994623656, + "2": 0.7828040994623656, + "4": 0.8090557795698925, + "8": 0.8204070060483871, + "16": 0.8217510920698925, + "32": 0.8217510920698925, + "ar": 0.09607064852150538 + }, + "minutes": 2.545037277539571, + "compute": { + "params": 4772864, + "answer_len_mean": 13.684530258178711 + } +} \ No newline at end of file diff --git a/results/g15-turkish-mdlm-s1.json b/results/g15-turkish-mdlm-s1.json new file mode 100644 index 0000000000000000000000000000000000000000..e79ada6c88871ddba071a3bde19598fd05e6617f --- /dev/null +++ b/results/g15-turkish-mdlm-s1.json @@ -0,0 +1,36 @@ +{ + "args": { + "lang": "turkish", + "size": "medium", + "mode": "mdlm", + "seq_len": 48, + "d": 256, + "layers": 6, + "heads": 8, + "pe": "rope", + "bs": 64, + "steps": 15000, + "lr": 0.0001, + "warmup": 1000, + "K": 4, + "lam": 1.0, + "seed": 1, + "data_root": "/scratch/zt1/project/msml612/user/govind02/pcmd/data/sig", + "eval_bs": 128, + "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g15-turkish-mdlm-s1" + }, + "accuracy_by_passes": { + "1": 0.7992901545698925, + "2": 0.8051495295698925, + "4": 0.8152826780913979, + "8": 0.8168682795698925, + "16": 0.8221186155913979, + "32": 0.8221186155913979, + "ar": 0.10193002352150538 + }, + "minutes": 2.530856684843699, + "compute": { + "params": 4772864, + "answer_len_mean": 13.684530258178711 + } +} \ No newline at end of file diff --git a/slurm/finalize.sh b/slurm/finalize.sh index 207d7170d35f77a37e30be1b6bb868323e2f81ed..3f850dc44f60bfecc0d845835d2cab9a0c5a8288 100644 --- a/slurm/finalize.sh +++ b/slurm/finalize.sh @@ -72,16 +72,30 @@ table = summary.read_text() if summary.exists() else "_(pending)_" (stage / "README.md").write_text(f"""--- license: apache-2.0 -tags: [masked-diffusion, discrete-diffusion, length-generalization, positional-encoding, from-scratch] +tags: [masked-diffusion, discrete-diffusion, parallel-decoding, text8, from-scratch] --- -# Length Generalization in Masked Diffusion Language Models +# Two Failure Modes of Parallel Decoding in Masked Diffusion -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. +Parallel decoding in masked diffusion loses accuracy for two different reasons +that need opposite fixes: -**No pretrained weights or tokenizers are used anywhere.** +1. **The marginals are untrained at high mask ratio.** A K-pass decode only ever + evaluates the model at mask ratios {{1, (K-1)/K, ..., 1/K}}, while uniform-t + training spends most of its capacity on nearly-complete states that + high-parallelism decoding never visits. **Fixable only at training time.** +2. **The committed positions carry real mutual information.** Parallel decoding + samples the product of marginals; the truth is the joint. The discarded + dependence is the total correlation of the committed set, upper-bounded by + its summed conditional entropy. **Fixable only at inference, by committing + less.** + +Prior work applies inference-time fixes to both, which is why it flatlines on +arithmetic - where mode 2 is identically zero and all the loss is mode 1. + +Everything is trained from scratch: no pretrained weights, no pretrained +tokenizer, and the autoregressive evaluator used to score generated text is +itself trained from scratch on the same corpus. Code: https://github.com/GIND123/MSML612 diff --git a/slurm/g14.sh b/slurm/g14.sh new file mode 100644 index 0000000000000000000000000000000000000000..b175f19964b92c14f2794ad480a837cd2f61bb52 --- /dev/null +++ b/slurm/g14.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Grid 14 - THE THEORY TEST, now at a learning rate that is not broken. +# +# THEORY.md predicts the minimum depth for one-pass decoding grows as log2(n), +# because addition's carry is a prefix scan computable by carry lookahead in +# O(log n) depth, and constant-depth transformers sit in TC0 which contains +# addition. +# +# log-depth (ours) : min depth ~ 2, 3, 4, 4 for 4, 8, 12, 16 digits +# ripple-carry : min depth ~ 4, 8, 12, 16 +# impossibility : no depth suffices +# +# The first attempt at this sweep was cancelled on discovering it would have run +# at lr 3e-4 with 500 warmup, where deep models fail for optimisation reasons and +# would have faked a clean "depth does not help" curve. +#SBATCH -J g14 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=48g +#SBATCH -t 08:00:00 +#SBATCH -o logs/g14-%A_%a.out +#SBATCH --array=0-47%12 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +DIG=(4 8 12 16); LAYERS=(2 3 4 6 8 12); SEEDS=(0 1) +I=$SLURM_ARRAY_TASK_ID +D=${DIG[$((I / 12))]}; L=${LAYERS[$(((I % 12) / 2))]}; S=${SEEDS[$((I % 2))]} +SL=$(( 3 * D + 4 )) +NAME="g14-d$D-L$L-s$S" +echo "=== $NAME ===" +$SCR/hf_env/bin/python train_add.py --mode full --digits $D --seq_len $SL --seed $S \ + --n_train 800000 --n_eval 400 --eval_bs 100 \ + --d 384 --layers $L --heads 6 --bs 128 --steps 60000 --lr 1e-4 --warmup 4000 \ + --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME diff --git a/slurm/g15.sh b/slurm/g15.sh new file mode 100644 index 0000000000000000000000000000000000000000..52d6046f12c8aa88604511d2b977f9b395eb3898 --- /dev/null +++ b/slurm/g15.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Grid 15 - component ablation on inflection. Which part carries the gain, and +# does the answer hold across typologically diverse languages rather than just +# on addition? +#SBATCH -J g15 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=32g +#SBATCH -t 06:00:00 +#SBATCH -o logs/g15-%A_%a.out +#SBATCH --array=0-47%12 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +LANGS=(german turkish finnish russian navajo georgian) +MODES=(mdlm matched distill full); SEEDS=(0 1) +I=$SLURM_ARRAY_TASK_ID +L=${LANGS[$((I / 8))]}; M=${MODES[$(((I % 8) / 2))]}; S=${SEEDS[$((I % 2))]} +NAME="g15-$L-$M-s$S" +echo "=== $NAME ===" +$SCR/hf_env/bin/python train_sig.py --lang $L --mode $M --size medium --seed $S \ + --steps 15000 --d 256 --layers 6 --lr 1e-4 --warmup 1000 \ + --data_root $SCR/pcmd/data/sig --out $SCR/pcmd/runs/$NAME diff --git a/slurm/g16.sh b/slurm/g16.sh new file mode 100644 index 0000000000000000000000000000000000000000..44020854568a03e770d9b44bf6f6372932b25db4 --- /dev/null +++ b/slurm/g16.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# Grid 16 - THE DEFINITIVE FRONTIER, one consistent configuration throughout. +# +# Grid 11 (10 layers, 60k steps) understated the method: it gave 75% at sixteen +# digits, while grid 13 (12 layers, 150k steps) reached 99.9% at TWENTY digits +# with the baseline at exactly zero. Mixing configurations across digit counts +# makes the curve meaningless, so this reruns every length at one setting. +#SBATCH -J g16 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 16:00:00 +#SBATCH -o logs/g16-%A_%a.out +#SBATCH --array=0-29%12 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +DIG=(12 16 20 24 28); MODES=(mdlm full); SEEDS=(0 1 2) +I=$SLURM_ARRAY_TASK_ID +D=${DIG[$((I / 6))]}; M=${MODES[$(((I % 6) / 3))]}; S=${SEEDS[$((I % 3))]} +SL=$(( 3 * D + 4 )) +NAME="g16-d$D-$M-s$S" +echo "=== $NAME ===" +$SCR/hf_env/bin/python train_add.py --mode $M --digits $D --seq_len $SL --seed $S \ + --n_train 1500000 --n_eval 400 --eval_bs 100 \ + --d 512 --layers 12 --heads 8 --bs 128 --steps 150000 --lr 1e-4 --warmup 6000 \ + --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME diff --git a/slurm/g16b.sh b/slurm/g16b.sh new file mode 100644 index 0000000000000000000000000000000000000000..567096deb456e60750fdf42036723c7c5687e5ba --- /dev/null +++ b/slurm/g16b.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# Grid 16 - THE DEFINITIVE FRONTIER, one consistent configuration throughout. +# +# Grid 11 (10 layers, 60k steps) understated the method: it gave 75% at sixteen +# digits, while grid 13 (12 layers, 150k steps) reached 99.9% at TWENTY digits +# with the baseline at exactly zero. Mixing configurations across digit counts +# makes the curve meaningless, so this reruns every length at one setting. +#SBATCH -J g16 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 16:00:00 +#SBATCH -o logs/g16-%A_%a.out +#SBATCH --array=0-11%6 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +DIG=(12 16); MODES=(mdlm full); SEEDS=(0 1 2) +I=$SLURM_ARRAY_TASK_ID +D=${DIG[$((I / 6))]}; M=${MODES[$(((I % 6) / 3))]}; S=${SEEDS[$((I % 3))]} +SL=$(( 3 * D + 4 )) +NAME="g16-d$D-$M-s$S" +echo "=== $NAME ===" +$SCR/hf_env/bin/python train_add.py --mode $M --digits $D --seq_len $SL --seed $S \ + --n_train 1500000 --n_eval 400 --eval_bs 100 \ + --d 512 --layers 12 --heads 8 --bs 128 --steps 150000 --lr 1e-4 --warmup 6000 \ + --K 4 --lam 1.0 --out $SCR/pcmd/runs/$NAME diff --git a/slurm/g17.sh b/slurm/g17.sh new file mode 100644 index 0000000000000000000000000000000000000000..2c6f45399771969f7f3df93b675eb2ac8bd13457 --- /dev/null +++ b/slurm/g17.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Grid 17 - COMPUTE-MATCHED CONTROL. +# +# `full` evaluates the model twice per training step (matched-schedule term plus +# high-ratio term), so at equal step counts it spends roughly twice the training +# FLOPs of the baseline. Any reviewer will ask whether the gap is the method or +# simply the extra compute. This gives the BASELINE the larger budget: 300k +# steps of plain MDLM against 150k steps of the method, which is if anything +# generous to the baseline since the method's second pass is on a shared batch. +#SBATCH -J g17 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 24:00:00 +#SBATCH -o logs/g17-%A_%a.out +#SBATCH --array=0-8%6 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +DIG=(12 16 20); SEEDS=(0 1 2) +I=$SLURM_ARRAY_TASK_ID +D=${DIG[$((I / 3))]}; S=${SEEDS[$((I % 3))]} +SL=$(( 3 * D + 4 )) +NAME="g17-d$D-mdlm2x-s$S" +echo "=== $NAME : baseline with DOUBLE the training steps ===" +$SCR/hf_env/bin/python -u train_add.py --mode mdlm --digits $D --seq_len $SL --seed $S \ + --n_train 1500000 --n_eval 400 --eval_bs 100 \ + --d 512 --layers 12 --heads 8 --bs 128 --steps 300000 --lr 1e-4 --warmup 6000 \ + --out $SCR/pcmd/runs/$NAME diff --git a/slurm/g17b.sh b/slurm/g17b.sh new file mode 100644 index 0000000000000000000000000000000000000000..b572dcf3740228cef70afa995c4041a7cdbc78a5 --- /dev/null +++ b/slurm/g17b.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Grid 17 - COMPUTE-MATCHED CONTROL. +# +# `full` evaluates the model twice per training step (matched-schedule term plus +# high-ratio term), so at equal step counts it spends roughly twice the training +# FLOPs of the baseline. Any reviewer will ask whether the gap is the method or +# simply the extra compute. This gives the BASELINE the larger budget: 300k +# steps of plain MDLM against 150k steps of the method, which is if anything +# generous to the baseline since the method's second pass is on a shared batch. +#SBATCH -J g17 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 24:00:00 +#SBATCH -o logs/g17-%A_%a.out +#SBATCH --array=0-3%4 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +DIG=(16 20); SEEDS=(0 1) +I=$SLURM_ARRAY_TASK_ID +D=${DIG[$((I / 2))]}; S=${SEEDS[$((I % 2))]} +SL=$(( 3 * D + 4 )) +NAME="g17-d$D-mdlm2x-s$S" +echo "=== $NAME : baseline with DOUBLE the training steps ===" +$SCR/hf_env/bin/python -u train_add.py --mode mdlm --digits $D --seq_len $SL --seed $S \ + --n_train 1500000 --n_eval 400 --eval_bs 100 \ + --d 512 --layers 12 --heads 8 --bs 128 --steps 300000 --lr 1e-4 --warmup 6000 \ + --out $SCR/pcmd/runs/$NAME diff --git a/slurm/h2h.sh b/slurm/h2h.sh new file mode 100644 index 0000000000000000000000000000000000000000..a0be54b2668b6e8b9852b8d223f996657fb13ee2 --- /dev/null +++ b/slurm/h2h.sh @@ -0,0 +1,17 @@ +#!/bin/bash +#SBATCH -J h2h +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 4 +#SBATCH --mem=32g +#SBATCH -t 3:00:00 +#SBATCH -o logs/h2h-%j.out +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +for D in 12 16; do + $SCR/hf_env/bin/python -u beat_baselines.py --runs $SCR/pcmd/runs --digits $D \ + --n_eval 400 --eval_bs 100 --out $SCR/pcmd/figures/beat_d$D.json +done diff --git a/slurm/seqchk.sh b/slurm/seqchk.sh new file mode 100644 index 0000000000000000000000000000000000000000..126f41f44a5271d32739ef99fd2bf87b4e09f644 --- /dev/null +++ b/slurm/seqchk.sh @@ -0,0 +1,15 @@ +#!/bin/bash +#SBATCH -J seqchk +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 4 +#SBATCH --mem=32g +#SBATCH -t 2:00:00 +#SBATCH -o logs/seqchk-%j.out +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +$SCR/hf_env/bin/python -u seq_check.py --runs $SCR/pcmd/runs --pattern 'g13-*' \ + --out $SCR/pcmd/figures/seq_check_g13.json diff --git a/slurm/t8.sh b/slurm/t8.sh new file mode 100644 index 0000000000000000000000000000000000000000..6b9b8dfd55f084d0a47ff536bf4c052029b2f95e --- /dev/null +++ b/slurm/t8.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# The capstone experiment: does the training-side fix compose with the +# inference-side fix on data that contains BOTH failure modes at once? +#SBATCH -J t8 +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 16:00:00 +#SBATCH -o logs/t8-%A_%a.out +#SBATCH --array=0-5 +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +MODES=(mdlm matched full); SEEDS=(0 1) +I=$SLURM_ARRAY_TASK_ID +M=${MODES[$((I / 2))]}; S=${SEEDS[$((I % 2))]} +NAME="t8-$M-s$S" +echo "=== $NAME ===" +$SCR/hf_env/bin/python -u train_text8.py --mode $M --seed $S \ + --root $SCR/pcmd/data/text8 --seq_len 256 \ + --d 512 --layers 12 --heads 8 --bs 128 --steps 60000 --lr 1e-4 --warmup 4000 \ + --K 8 --lam 1.0 --n_gen 256 --gen_bs 64 \ + --evallm $SCR/pcmd/runs/evallm --out $SCR/pcmd/runs/$NAME diff --git a/slurm/t8eval.sh b/slurm/t8eval.sh new file mode 100644 index 0000000000000000000000000000000000000000..41077bd4c318d094a5358d779a5e888750a4e851 --- /dev/null +++ b/slurm/t8eval.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# The shared judge: one autoregressive character LM, trained from scratch on +# text8, used to score every diffusion run's samples. Trained once so that all +# systems are measured against exactly the same yardstick. +#SBATCH -J t8eval +#SBATCH -p gpu-h100 +#SBATCH --gpus=h100:1 +#SBATCH -c 8 +#SBATCH --mem=64g +#SBATCH -t 8:00:00 +#SBATCH -o logs/t8eval-%j.out +source /etc/profile +SCR=/scratch/zt1/project/msml612/user/$USER +module load python/3.10.10 +export PYTHONPATH= +cd $SCR/pcmd/src +$SCR/hf_env/bin/python -u train_eval_lm.py --root $SCR/pcmd/data/text8 \ + --seq_len 256 --d 512 --layers 8 --heads 8 --bs 128 \ + --steps 50000 --lr 3e-4 --warmup 2000 --out $SCR/pcmd/runs/evallm diff --git a/src/audit.py b/src/audit.py index 8f2f8d986d3d83f58c4e3658024a5cae5afc5415..7a13c9b23fd933d14f71bc0ab0f2ffd6a7babb89 100644 --- a/src/audit.py +++ b/src/audit.py @@ -102,6 +102,76 @@ for f in glob.glob(f"{runs}/*/result.json"): oor.append((os.path.basename(os.path.dirname(f)), k, v)) check("all accuracies within [0,1]", not oor, str(oor[:3])) +print("\n=== 7. decoders: adaptive commitment rules behave as specified ===") +# These three rules produce the headline text8 numbers, so their contract is +# checked directly rather than trusted: every masked position must end up +# filled, decoding must always make progress (no stall), and the entropy budget +# must actually behave monotonically in the budget - otherwise the +# quality/compute curve it traces would be meaningless. +try: + from text8 import Text8Tokenizer + t8 = Text8Tokenizer() + torch.manual_seed(0) + L, B = 32, 8 + dm = Transformer(len(t8), 64, 2, 2, "rope", causal=False, max_len=L + 8) + dm.eval() + blank = torch.full((B, L), t8.mask, dtype=torch.long) + + for name, fn in [ + ("fixed-K K=1", lambda b: dfn.fixed_k_decode(dm, b, t8, "cpu", 1)), + ("fixed-K K=8", lambda b: dfn.fixed_k_decode(dm, b, t8, "cpu", 8)), + ("confidence tau=0.99", + lambda b: dfn.confidence_threshold_decode(dm, b, t8, "cpu", 0.99)), + ("entropy-budget B=0.1", + lambda b: dfn.entropy_budget_decode(dm, b, t8, "cpu", 0.1)), + ("entropy-budget B=inf", + lambda b: dfn.entropy_budget_decode(dm, b, t8, "cpu", 1e9))]: + x, nfe = fn(blank) + check(f"{name}: no [MASK] survives", not bool((x == t8.mask).any()), + f"{int((x == t8.mask).sum())} left, nfe={nfe}") + check(f"{name}: emits only real symbols", bool((x < len(t8) - 1).all())) + + _, n1 = dfn.fixed_k_decode(dm, blank, t8, "cpu", 1) + check("fixed-K K=1 uses exactly one forward pass", n1 == 1, f"nfe={n1}") + _, nb = dfn.entropy_budget_decode(dm, blank, t8, "cpu", 1e9) + check("unbounded entropy budget commits in one pass", nb == 1, f"nfe={nb}") + + # a tighter budget must never be cheaper than a looser one + nfes = [dfn.entropy_budget_decode(dm, blank, t8, "cpu", b)[1] + for b in (0.01, 0.5, 2.0, 8.0, 1e9)] + check("entropy budget is monotone: tighter budget costs >= passes", + all(nfes[i] >= nfes[i + 1] for i in range(len(nfes) - 1)), str(nfes)) + + # partial contexts must be left alone: only `fillable` positions may change + part = blank.clone() + part[:, :L // 2] = 3 + fill = torch.zeros_like(part, dtype=torch.bool) + fill[:, L // 2:] = True + y, _ = dfn.entropy_budget_decode(dm, part, t8, "cpu", 0.5, fillable=fill) + check("decoding never overwrites the given context", + bool((y[:, :L // 2] == 3).all())) +except Exception as e: + check("decoder contract", False, f"{type(e).__name__}: {e}") + +print("\n=== 8. text8 splits are disjoint and standard ===") +try: + from text8 import load_text8, N_TRAIN, N_VALID, N_TEST + root = os.environ.get("T8ROOT", "data/text8") + tr8, va8, te8, tk8 = load_text8(root, 256) + check("standard 90M/5M/5M split", (N_TRAIN, N_VALID, N_TEST) == + (90_000_000, 5_000_000, 5_000_000)) + check("vocabulary is 27 symbols + [MASK]", len(tk8) == 28, f"{len(tk8)}") + check("no window straddles a split boundary", + tr8.size + va8.size + te8.size <= 100_000_000, + f"{tr8.size + va8.size + te8.size} chars covered") + seen = {bytes(r) for r in tr8[::997]} + dup = sum(bytes(r) in seen for r in te8[::997]) + check("no test window appears in the training sample", dup == 0, f"{dup} shared") +except FileNotFoundError: + print(" [skip] text8 not downloaded here") +except Exception as e: + check("text8 loader", False, f"{type(e).__name__}: {e}") + print(f"\n=== AUDIT: {len(FAIL)} failure(s) ===") for f in FAIL: print(" !", f) diff --git a/src/beat_baselines.py b/src/beat_baselines.py index 78ec3306d8d2a5e524b1874b0e8a46f11fb2a327..10cd3546723231d070b3335a0bbd8c7aea14c1c9 100644 --- a/src/beat_baselines.py +++ b/src/beat_baselines.py @@ -59,7 +59,7 @@ def run(model, steps, strategy, tau=None): if not masked.any(): break logits = model(x) - probs = logits.softmax(-1) + probs = dfn.real_probs(logits, tok) conf, pred = probs.max(-1) if strategy == "random": score = torch.rand_like(conf) @@ -85,8 +85,44 @@ def run(model, steps, strategy, tau=None): rem = (x == tok.mask) & ab if rem.any() and pred is not None: x[rem] = pred[rem] - accs.append(exact_match(x.cpu(), xb.cpu(), ab.cpu(), tok)) - return float(np.mean(accs)) + accs.extend(per_example(x.cpu(), xb.cpu(), ab.cpu())) + return np.array(accs, dtype=float) + + +def per_example(pred, gold, amask): + """Correct/incorrect per evaluation item, not a batch mean. + + A paired test needs to know which items each system got right, on the SAME + items. Batch means throw that away. + """ + out = [] + for p, g, a in zip(pred.tolist(), gold.tolist(), amask.tolist()): + pa = [p[i] for i in range(len(a)) if a[i]] + ga = [g[i] for i in range(len(a)) if a[i]] + out.append(float(pa == ga)) + return out + + +def paired_test(ours, theirs, n_perm=20000, seed=0): + """Paired permutation test on the same evaluation items. + + The two systems are scored on identical inputs, so the items are paired and + a paired test is the right one: under the null that the systems are + interchangeable, the sign of each item's difference is exchangeable. + Reports the observed difference, a bootstrap 95% interval on it, and the + two-sided p-value. + """ + d = ours - theirs + obs = float(d.mean()) + rng = np.random.default_rng(seed) + flips = rng.choice([-1.0, 1.0], size=(n_perm, len(d))) + null = (flips * d).mean(axis=1) + p = float((np.abs(null) >= abs(obs) - 1e-12).mean()) + idx = rng.integers(0, len(d), size=(10000, len(d))) + bs = d[idx].mean(axis=1) + return {"diff": obs, "ci95": [float(np.percentile(bs, 2.5)), + float(np.percentile(bs, 97.5))], + "p_value": p, "n": int(len(d))} base_dirs = sorted(glob.glob(f"{a.runs}/g11-d{a.digits}-mdlm-s*")) @@ -99,6 +135,7 @@ results = {"digits": a.digits, "inference_time_on_baseline": {}, "ours": {}} print(f"\n=== {a.digits}-digit addition: can any inference-time method rescue " f"the baseline? ===") print(f"{'strategy':<22}{'1 pass':>9}{'2':>8}{'4':>8}{'8':>8}{'16':>8}") +ITEMS = {} for strat, tau in [("confidence", None), ("entropy", None), ("entropy-gated", 0.5), ("margin", None), ("random", None)]: row = [] @@ -111,16 +148,31 @@ for strat, tau in [("confidence", None), ("entropy", None), ("entropy-gated", 0. results["inference_time_on_baseline"][strat] = row print(f"{strat:<22}" + "".join(f"{v*100:8.1f}" for v in row)) -row = [] +row, ours_items = [], [] for steps in (1, 2, 4, 8, 16): - vals = [] - for d in ours_dirs: - m, _ = load(d) - vals.append(run(m, steps, "confidence")) - row.append(float(np.mean(vals))) + vals = [run(load(d)[0], steps, "confidence") for d in ours_dirs] + per = np.mean(np.stack(vals), axis=0) + row.append(float(per.mean())); ours_items.append(per) results["ours"]["confidence"] = row print(f"{'OURS (training-time)':<22}" + "".join(f"{v*100:8.1f}" for v in row)) +print(f"\n=== is the difference significant? (paired permutation test) ===") +print("Our method at ONE pass against each inference-time strategy at its BEST") +print("budget, scored on identical evaluation items.\n") +print(f"{'vs strategy':<22}{'their best':>12}{'ours @1':>10}{'diff':>9}" + f"{'95% CI':>18}{'p':>10}") +results["significance"] = {} +for strat, rows in ITEMS.items(): + best_i = int(np.argmax([r.mean() for r in rows])) + st = paired_test(ours_items[0], rows[best_i]) + results["significance"][strat] = dict(st, their_best_budget=[1, 2, 4, 8, 16][best_i]) + print(f"{strat:<22}{rows[best_i].mean()*100:11.1f}%{ours_items[0].mean()*100:9.1f}%" + f"{st['diff']*100:+8.1f}" + f" [{st['ci95'][0]*100:+.1f},{st['ci95'][1]*100:+.1f}]" + f"{st['p_value']:10.2g}") +print(f"\n(n = {len(ours_items[0])} paired evaluation items; " + f"p is two-sided, 20000 permutations)") + best_inf = max(v[0] for v in results["inference_time_on_baseline"].values()) best_inf_any = max(max(v) for v in results["inference_time_on_baseline"].values()) print(f"\nbest inference-time method at 1 pass : {best_inf*100:.1f}%") diff --git a/src/collect.py b/src/collect.py index fbea377d2aa758d5b757b9a3bd96746506751dc0..a22cfb83e7e1f661c611f254907bd733abd45330 100644 --- a/src/collect.py +++ b/src/collect.py @@ -6,7 +6,7 @@ The organising variable across all three task families is DEPENDENCY DEPTH: deep + unique (addition) -> one pass fails, and the method fixes it any + non-unique (probes) -> information-theoretic limit, unfixable """ -import json, glob, os +import json, glob, math, os import numpy as np import matplotlib matplotlib.use("Agg") @@ -162,6 +162,178 @@ if g10: f"{sum(x>0.9 for x in v)}/{len(v)} |") md.append("") + +# ---- Table 5: the definitive frontier, one configuration throughout --------- +# Grid 11 mixed configurations across digit counts, which makes the curve a +# statement about the configurations as much as about the method. Grid 16 reruns +# every length at a single setting so the frontier means one thing. +g16 = runs("g16-") +if g16: + md.append("## Table 5 — Addition frontier at ONE fixed configuration\n") + md.append("12 layers, 150k steps, lr 1e-4, 3 seeds, identical at every " + "length, so the curve reflects the method and not the config.\n") + md.append("| digits | MDLM baseline | matched (ours) | seeds >90% |") + md.append("|---|---|---|---|") + for d in sorted({r["args"]["digits"] for r in g16}): + cells, rel = [], [] + for m in ("mdlm", "full"): + rs = runs("g16-", digits=d, mode=m) + if not rs: + cells.append("–"); rel.append("–"); continue + v = [r["accuracy_by_passes"]["1"] for r in rs] + lo, hi = boot(v) + cells.append(f"{np.mean(v)*100:.1f} [{lo:.0f},{hi:.0f}]") + rel.append(f"{sum(x>0.9 for x in v)}/{len(v)}") + md.append(f"| **{d}** | {cells[0]} | {cells[1]} | {rel[0]} → {rel[1]} |") + md.append("") + +# ---- Table 6: the depth sweep, the theory's falsifiable prediction ---------- +g14 = runs("g14-") +if g14: + md.append("## Table 6 — Minimum depth for one-pass addition\n") + md.append("The account in THEORY.md predicts minimum depth growing as " + "log2(n), since the carry is a prefix scan. A ripple-carry account " + "predicts linear growth; an impossibility account predicts no " + "depth suffices.\n") + lays = sorted({r["args"]["layers"] for r in g14}) + md.append("| digits | " + " | ".join(f"{l}L" for l in lays) + + " | log2(n) | min depth >90% |") + md.append("|---" * (len(lays) + 3) + "|") + for d in sorted({r["args"]["digits"] for r in g14}): + cells, mind = [], None + for l in lays: + rs = runs("g14-", digits=d, layers=l) + if not rs: + cells.append("–"); continue + mu = np.mean([r["accuracy_by_passes"]["1"] for r in rs]) * 100 + cells.append(f"{mu:.0f}") + if mu > 90 and mind is None: + mind = l + md.append(f"| **{d}** | " + " | ".join(cells) + + f" | {math.log2(d):.1f} | {mind or '–'} |") + md.append("") + + fig, ax = plt.subplots(figsize=(6.4, 4.2)) + ds = sorted({r["args"]["digits"] for r in g14}) + obs = [] + for d in ds: + mind = None + for l in lays: + rs = runs("g14-", digits=d, layers=l) + if rs and np.mean([r["accuracy_by_passes"]["1"] for r in rs]) > 0.9: + mind = l; break + obs.append(mind) + ok = [(d, m) for d, m in zip(ds, obs) if m] + if ok: + ax.plot([d for d, _ in ok], [m for _, m in ok], "o-", color="#1a5276", + label="measured minimum depth") + ax.plot(ds, [math.log2(d) for d in ds], "--", color="#117a65", + label="log2(n) — prefix-scan prediction") + ax.plot(ds, [d / 4 for d in ds], ":", color="#c0392b", + label="n/4 — ripple-carry prediction") + ax.set_xlabel("operand digits"); ax.set_ylabel("layers needed for >90% at one pass") + ax.set_title("Does required depth grow like log n or like n?") + ax.legend(fontsize=8); ax.grid(alpha=.3) + fig.tight_layout(); fig.savefig(f"{OUT}/fig3_depth_scaling.png", dpi=160) + +# ---- Table 7: component ablation on inflection ----------------------------- +g15 = runs("g15-") +if g15: + md.append("## Table 7 — Which component carries the gain (inflection)\n") + modes = ["mdlm", "matched", "distill", "full"] + langs = sorted({r["args"]["lang"] for r in g15}) + md.append("| language | " + " | ".join(modes) + " | (gap 1→32) |") + md.append("|---" * (len(modes) + 2) + "|") + for L in langs: + cells, gaps = [], [] + for m in modes: + rs = runs("g15-", lang=L, mode=m) + if not rs: + cells.append("–"); gaps.append("–"); continue + one = np.mean([r["accuracy_by_passes"]["1"] for r in rs]) * 100 + many = np.mean([r["accuracy_by_passes"]["32"] for r in rs]) * 100 + cells.append(f"{one:.1f}"); gaps.append(f"{many-one:+.1f}") + md.append(f"| {L} | " + " | ".join(cells) + " | " + " / ".join(gaps) + " |") + md.append("") + +# ---- Table 8: text8, where the two fixes have to compose ------------------- +t8 = [r for r in R if r["name"].startswith("t8-")] +if t8: + md.append("## Table 8 — text8: quality against compute\n") + md.append("Generative perplexity under a from-scratch autoregressive " + "character evaluator, against the number of function evaluations " + "actually used. Entropy and distinct-4-gram rate are reported " + "beside it because generative perplexity alone is gamed by " + "degenerate repetition; the real-text reference row is the " + "target, not zero.\n") + ref = t8[0].get("reference", {}) + md.append(f"Real text8 reference — ppl {ref.get('gen_ppl', float('nan')):.2f}, " + f"entropy {ref.get('entropy', float('nan')):.3f}, " + f"distinct-4 {ref.get('distinct4', float('nan')):.3f}\n") + md.append("| training | decoding | NFE | gen ppl | entropy | distinct-4 | valid BPC |") + md.append("|---|---|---|---|---|---|---|") + by = {} + for r in t8: + by.setdefault(r["args"]["mode"], []).append(r) + for mode in ("mdlm", "matched", "full"): + for r in by.get(mode, [])[:1]: + for key, row in sorted(r.get("decode", {}).items()): + if row.get("gen_ppl") is None: + continue + md.append(f"| {mode} | {key} | {row['nfe']:.1f} | " + f"{row['gen_ppl']:.2f} | {row['entropy']:.3f} | " + f"{row['distinct4']:.3f} | {r.get('valid_bpc', float('nan')):.4f} |") + md.append("") + + # quality/compute frontier: the plot the claim lives or dies on + fig, ax = plt.subplots(figsize=(6.8, 4.4)) + style = {"mdlm": ("#c0392b", "uniform-t training"), + "matched": ("#e67e22", "matched training"), + "full": ("#1a5276", "matched + high-ratio")} + for mode, rs in by.items(): + c, lbl = style.get(mode, ("#555", mode)) + for fam, mk, ls in (("fixedK_", "o", "-"), ("entbudget_", "s", "--")): + pts = [] + for r in rs: + for k, row in r.get("decode", {}).items(): + if k.startswith(fam) and row.get("gen_ppl"): + pts.append((row["nfe"], row["gen_ppl"])) + if pts: + pts.sort() + ax.plot([p[0] for p in pts], [p[1] for p in pts], marker=mk, + ms=4, ls=ls, color=c, alpha=.85, + label=f"{lbl}, {'fixed-K' if fam[0]=='f' else 'entropy budget'}") + if ref.get("gen_ppl"): + ax.axhline(ref["gen_ppl"], color="k", lw=1, ls=":", label="real text8") + ax.set_xscale("log", base=2); ax.set_yscale("log") + ax.set_xlabel("forward passes (NFE)"); ax.set_ylabel("generative perplexity") + ax.set_title("Quality against compute: do the two fixes compose?") + ax.legend(fontsize=7); ax.grid(alpha=.3) + fig.tight_layout(); fig.savefig(f"{OUT}/fig4_text8_quality_vs_compute.png", dpi=160) + +# ---- Table 9: failed to parallelise, or failed to learn? ------------------- +sc = os.path.join(OUT, "seq_check_g13.json") +if os.path.exists(sc): + rows = json.load(open(sc)) + md.append("## Table 9 — Did the baseline fail to parallelise, or fail to learn?\n") + md.append("A 20-digit answer occupies 21 slots, so the standard 16-pass grid " + "is not a sequential decode. Re-decoded at one token per pass, the " + "two explanations separate.\n") + md.append("| digits | mode | @1 pass | fully sequential | reading |") + md.append("|---|---|---|---|---|") + agg = {} + for r in rows: + agg.setdefault((r["digits"], r["mode"]), []).append(r) + for k in sorted(agg): + v = agg[k] + one = np.mean([x["acc_1pass"] for x in v]) * 100 + seq = np.mean([x["acc_sequential"] for x in v]) * 100 + read = ("never learned the task" if seq < 5 else + "learned it, cannot parallelise" if one < seq - 5 else + "learned it and parallelises") + md.append(f"| {k[0]} | {k[1]} | {one:.1f} | {seq:.1f} | {read} |") + md.append("") + open(f"{OUT}/summary.md", "w").write("\n".join(md) + "\n") print("\n".join(md)) print(f"\nwrote figures + summary to {OUT}") diff --git a/src/diffusion.py b/src/diffusion.py index d3b4b2e2226fd0d49193d383ac9785703ce03e84..b3777098be7bdea3a200a115b07b2e3db336c996 100644 --- a/src/diffusion.py +++ b/src/diffusion.py @@ -421,7 +421,7 @@ def cond_decode(model, x0, amask, tok, steps, device): if not masked.any(): break logits = model(x) - conf, pred = logits.softmax(-1).max(-1) + conf, pred = real_probs(logits, tok).max(-1) conf = conf.masked_fill(~masked, -1e9) left = int(n_ans * (s - 1) / steps) for b in range(B): @@ -469,3 +469,225 @@ def pmd_progressive(model, x0, amask, tok, progress, K_max=16, lam_max=1.0, if lam_now > 0: loss = loss + lam_now * sequential_distill_loss(model, x0, amask, tok) return loss, K_now, lam_now + + +# ------------------------------------------------- the unified formulation --- +# Two failure modes make parallel decoding lose accuracy, and they need opposite +# fixes. Keeping them apart is the whole argument of this project. +# +# (i) THE MARGINALS ARE UNTRAINED AT HIGH MASK RATIO. +# A K-pass decode only ever evaluates the model at mask ratios +# {1, (K-1)/K, ..., 1/K}. Uniform-t training spends most of its capacity +# on nearly-complete states that high-parallelism decoding never visits. +# No inference-time reordering can repair a marginal that was never fit; +# this is fixable ONLY at training time, and `cond_mdlm_loss(t_dist= +# "matched", K=...)` is the fix. At K=1 the matched schedule degenerates +# to t=1, which is exactly `sequential_distill_loss` - the addition +# experiments are the K=1 special case of one principle. +# +# (ii) THE COMMITTED POSITIONS CARRY REAL MUTUAL INFORMATION. +# Parallel decoding samples the PRODUCT of marginals; the truth is the +# joint. The discarded dependence is the total correlation of the +# committed set S, +# TC(S) = sum_i H(x_i | c) - H(S | c) <= sum_i H(x_i | c), +# so the summed conditional entropy of what we commit is an upper bound +# on the error we incur by committing it. If every committed conditional +# is a point mass the bound is zero and independent sampling reproduces +# the joint exactly. No training objective can remove this term without +# abandoning the data distribution, so it is fixable ONLY at inference, +# by committing less. `entropy_budget_decode` is that fix. +# +# Prior work applies inference-time fixes to both, which is why it flatlines on +# addition: there, every conditional is a point mass, (ii) is identically zero, +# and all of the loss is (i). The claim here is that the two compose - and text, +# which contains deterministic and high-entropy positions in the same sequence, +# is where that composition is tested. + +@torch.no_grad() +def _entropy(logits): + p = logits.softmax(-1) + return -(p * p.clamp_min(1e-9).log()).sum(-1) + + +@torch.no_grad() +def _commit(x, take, samp): + return torch.where(take, samp, x) + + +def real_probs(logits, tok): + """Distribution over REAL tokens only. + + The absorbing state is an input symbol, never an output: in the SUBS + parameterisation the denoiser places zero probability on [MASK]. Sampling + from the raw softmax instead lets a position be "committed" as [MASK], so it + stays masked, decoding silently stalls, and the adaptive rules stop being + monotone in their budget. The audit caught exactly this. + """ + logits = logits.clone() + logits[..., tok.mask] = -float("inf") + if getattr(tok, "pad", tok.mask) != tok.mask: + logits[..., tok.pad] = -float("inf") + return logits.softmax(-1) + + +@torch.no_grad() +def entropy_budget_decode(model, x, tok, device, budget=0.5, max_steps=256, + temperature=1.0, fillable=None): + """FIX (ii): commit a set whose summed conditional entropy stays under `budget`. + + At each pass, rank the still-masked positions by conditional entropy and + commit the longest low-entropy prefix whose cumulative entropy is within + `budget` nats. By the bound above, `budget` directly caps the dependence + thrown away at that pass, so it is an interpretable dial rather than a tuned + heuristic: budget -> 0 recovers near-sequential decoding, budget -> infinity + recovers one-shot parallel decoding, and the curve between them is the + quality/compute trade-off. + + Returns (x, nfe); nfe is data-dependent, which is the point of an adaptive + rule - cheap sequences finish in fewer passes. + """ + x = x.clone().to(device) + can = (x == tok.mask) if fillable is None else fillable.to(device) + B, L = x.shape + ar = torch.arange(L, device=device) + nfe = 0 + for _ in range(max_steps): + masked = (x == tok.mask) & can + has = masked.any(1) + if not has.any(): + break + logits = model(x) + nfe += 1 + if temperature != 1.0: + logits = logits / temperature + probs = real_probs(logits, tok) + H = -(probs * probs.clamp_min(1e-9).log()).sum(-1) + samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L) + + # sort masked positions by entropy; unmasked go last via +inf + Hm = H.masked_fill(~masked, float("inf")) + order = Hm.argsort(dim=1) + cum = Hm.gather(1, order).cumsum(1) + k = (cum <= budget).sum(1) + k = torch.where(has, k.clamp_min(1), torch.zeros_like(k)) # always progress + sel = ar[None, :] < k[:, None] + take = torch.zeros_like(masked).scatter_(1, order, sel) & masked + x = _commit(x, take, samp) + return x, nfe + + +@torch.no_grad() +def confidence_threshold_decode(model, x, tok, device, tau=0.9, max_steps=256, + temperature=1.0, fillable=None): + """The published inference-time competitor (confidence-threshold family). + + Commit every masked position whose top-1 probability exceeds `tau`. This is + the rule the 2026 parallel-decoding accelerators use; it is included so the + comparison is against the actual algorithm rather than a top-k heuristic. It + is a sensible rule and on text it works - the argument here is not that it + is wrong, but that it addresses failure mode (ii) only and therefore cannot + recover accuracy lost to (i). + """ + x = x.clone().to(device) + can = (x == tok.mask) if fillable is None else fillable.to(device) + B, L = x.shape + nfe = 0 + for _ in range(max_steps): + masked = (x == tok.mask) & can + has = masked.any(1) + if not has.any(): + break + logits = model(x) + nfe += 1 + if temperature != 1.0: + logits = logits / temperature + probs = real_probs(logits, tok) + conf = probs.max(-1).values + samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L) + + take = (conf >= tau) & masked + # a row that clears nothing still has to advance, or decoding stalls + stuck = has & ~take.any(1) + if stuck.any(): + best = conf.masked_fill(~masked, -float("inf")).argmax(1) + take[stuck, best[stuck]] = True + x = _commit(x, take, samp) + return x, nfe + + +@torch.no_grad() +def fixed_k_decode(model, x, tok, device, steps, temperature=1.0, fillable=None): + """Plain K-pass decoding: reveal an equal share of positions per pass, + highest-confidence first. The non-adaptive reference point, and the setting + in which the addition results are reported.""" + x = x.clone().to(device) + can = (x == tok.mask) if fillable is None else fillable.to(device) + B, L = x.shape + total = int(can.sum(1).max()) + ar = torch.arange(L, device=device) + nfe = 0 + for s in range(steps, 0, -1): + masked = (x == tok.mask) & can + has = masked.any(1) + if not has.any(): + break + logits = model(x) + nfe += 1 + if temperature != 1.0: + logits = logits / temperature + probs = real_probs(logits, tok) + conf = probs.max(-1).values.masked_fill(~masked, -float("inf")) + samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L) + + left = int(total * (s - 1) / steps) + k = (masked.sum(1) - left).clamp_min(0) + k = torch.where(has, k.clamp_min(1), torch.zeros_like(k)) + order = conf.argsort(dim=1, descending=True) + sel = ar[None, :] < k[:, None] + take = torch.zeros_like(masked).scatter_(1, order, sel) & masked + x = _commit(x, take, samp) + return x, nfe + + +def uncond_mdlm_loss(model, x0, tok, t_dist="uniform", K=None): + """Unconditional masked-diffusion loss with the schedule as a free choice. + + This is `cond_mdlm_loss` with every position maskable, written separately + because text has no prompt/answer split. `t_dist="matched"` is FIX (i): + draw t from the ratios a K-pass decode actually visits instead of from + U(0,1). The 1/t weighting of the ELBO is kept so that `t_dist="uniform"` + remains the standard MDLM objective and the baseline is exactly the + published one. + """ + B, L = x0.shape + dev = x0.device + if t_dist == "matched" and K: + t = torch.randint(1, K + 1, (B,), device=dev).float() / K + else: + t = torch.rand(B, device=dev).clamp_min(1.0 / L) + + m = torch.rand(B, L, device=dev) < t[:, None] + m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True + xt = torch.where(m, torch.full_like(x0, tok.mask), x0) + logits = model(xt) + ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1), + reduction="none").view(B, L) + return ((1.0 / t) * (ce * m).sum(1) / L).mean() + + +def text_loss(model, x0, tok, mode="mdlm", K=8, lam=1.0): + """Training objectives for text8. + + mdlm standard uniform-t MDLM (the published baseline) + matched FIX (i): schedule matched to a K-pass decode + full matched schedule plus an explicit high-ratio term + + `full` adds one extra evaluation at t drawn from the top of the schedule, + which is where a low-NFE decode spends its first and most damaging commits. + """ + if mode == "mdlm": + return uncond_mdlm_loss(model, x0, tok, "uniform") + if mode == "matched": + return uncond_mdlm_loss(model, x0, tok, "matched", K) + base = uncond_mdlm_loss(model, x0, tok, "matched", K) + return base + lam * uncond_mdlm_loss(model, x0, tok, "matched", 2) diff --git a/src/seq_check.py b/src/seq_check.py new file mode 100644 index 0000000000000000000000000000000000000000..f509a498beed3b403d6a16e7b810e60a00e0a80b --- /dev/null +++ b/src/seq_check.py @@ -0,0 +1,93 @@ +"""Did the baseline fail to PARALLELISE, or fail to LEARN? + +At 20 digits the baseline scores 0.0% at one pass and 0.0% at sixteen. It is +tempting to read that as a very large parallel-decoding gap, but the standard +pass grid stops at sixteen while a 20-digit answer occupies 21 slots - so that +number was never a fully sequential decode, and the two explanations are not +distinguished by it: + + (a) the baseline learned the task and merely cannot commit it in parallel + -> it should succeed once decoding is fully sequential; + (b) the baseline never learned the task at all + -> it fails at every budget, and the claim is about sample efficiency + rather than about parallelism. + +These are different papers. This script decodes every run at a budget equal to +the number of answer positions, which is genuinely one token per pass, and +reports accuracy at that budget next to the one-pass number. +""" +import argparse, glob, json, os, sys +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(__file__)) +from data import build_addition +from model import Transformer +import diffusion as dfn +from train_add import exact_match + +p = argparse.ArgumentParser() +p.add_argument("--runs", default="runs") +p.add_argument("--pattern", default="g13-*") +p.add_argument("--n_eval", type=int, default=200) +p.add_argument("--eval_bs", type=int, default=50) +p.add_argument("--out", default="figures/seq_check.json") +a = p.parse_args() + +dev = "cuda" if torch.cuda.is_available() else "cpu" +rows = [] + +for run_dir in sorted(glob.glob(os.path.join(a.runs, a.pattern))): + rf = os.path.join(run_dir, "result.json") + mf = os.path.join(run_dir, "model.pt") + if not (os.path.exists(rf) and os.path.exists(mf)): + continue + r = json.load(open(rf)) + cfg = r["args"] + d, SL = cfg["digits"], cfg["seq_len"] + + X, A, tok = build_addition(a.n_eval, d, SL, seed=91_000, exact=True) + X, A = torch.from_numpy(X), torch.from_numpy(A) + n_slots = int(A[0].sum()) + + m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"], + causal=False, max_len=SL + 8).to(dev) + m.load_state_dict(torch.load(mf, map_location=dev)) + m.eval() + + accs = {} + for steps in (1, n_slots): + vals = [] + for i in range(0, len(X), a.eval_bs): + xb, ab = X[i:i + a.eval_bs], A[i:i + a.eval_bs] + out = dfn.cond_decode(m, xb, ab, tok, steps, dev) + vals.append(exact_match(out.cpu(), xb, ab, tok)) + accs[steps] = float(np.mean(vals)) + + rows.append({"run": os.path.basename(run_dir), "digits": d, + "mode": cfg["mode"], "seed": cfg["seed"], "slots": n_slots, + "acc_1pass": accs[1], "acc_sequential": accs[n_slots]}) + print(f"{rows[-1]['run']:<22} {d:>3}d {cfg['mode']:<6} " + f"1-pass {accs[1]*100:6.2f} fully-sequential ({n_slots} passes) " + f"{accs[n_slots]*100:6.2f}", flush=True) + +print("\n=== does the baseline learn the task at ANY budget? ===") +agg = {} +for r in rows: + agg.setdefault((r["digits"], r["mode"]), []).append(r) +print(f"{'digits':>7} {'mode':<7}{'1 pass':>9}{'sequential':>12} verdict") +for k in sorted(agg): + v = agg[k] + one = np.mean([x["acc_1pass"] for x in v]) * 100 + seq = np.mean([x["acc_sequential"] for x in v]) * 100 + if seq < 5: + verdict = "never learned the task (sample efficiency, not parallelism)" + elif one < seq - 5: + verdict = "learned it, cannot parallelise (a true parallelism gap)" + else: + verdict = "learned it and parallelises" + print(f"{k[0]:>7} {k[1]:<7}{one:8.1f} {seq:11.1f} {verdict}") + +os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True) +json.dump(rows, open(a.out, "w"), indent=2) +print(f"\nwrote {a.out}") diff --git a/src/text8.py b/src/text8.py new file mode 100644 index 0000000000000000000000000000000000000000..897cf5b8fa58a3dccd5a12bdfaad3caeae853fbe --- /dev/null +++ b/src/text8.py @@ -0,0 +1,91 @@ +"""text8: the standard character-level language-modelling benchmark. + +100M characters of cleaned English Wikipedia, vocabulary of 27 symbols (a-z and +space). The conventional split is the first 90M characters for training, the +next 5M for validation and the last 5M for test; it is used unchanged here so +bits-per-character is comparable with published numbers. + +This is the task that separates the two failure modes of parallel decoding. +Addition has unique answers, so every conditional is a point mass and all of the +damage comes from marginals that were never trained at high mask ratios. Natural +text has both kinds of position at once: closing a word, a suffix or a function +word is nearly deterministic given the context, while the choice of the next +content word carries real entropy. A method that claims to work "across all +things" has to handle a sequence in which both appear together. +""" +import os +import numpy as np + +URL = "http://mattmahoney.net/dc/text8.zip" +VOCAB = " abcdefghijklmnopqrstuvwxyz" # 27 symbols, exactly as distributed + +# the conventional text8 split, in characters +N_TRAIN, N_VALID, N_TEST = 90_000_000, 5_000_000, 5_000_000 + + +class Text8Tokenizer: + """Character vocabulary plus the absorbing [MASK] state. + + [MASK] is appended after the 27 real characters, so ids 0..26 are the data + symbols and id 27 is the absorbing state. No padding symbol is needed: every + training sequence is a fixed-length contiguous window of the corpus. + """ + + def __init__(self): + self.itos = list(VOCAB) + ["[MASK]"] + self.stoi = {c: i for i, c in enumerate(self.itos)} + self.mask = self.stoi["[MASK]"] + self.pad = self.mask # unused; kept for interface compatibility + + def __len__(self): + return len(self.itos) + + def encode(self, s): + return np.array([self.stoi[c] for c in s if c in self.stoi], dtype=np.int64) + + def decode(self, ids): + return "".join(self.itos[int(i)] for i in ids) + + +def _raw(root): + path = os.path.join(root, "text8") + if not os.path.exists(path): + raise FileNotFoundError( + f"{path} not found. Fetch it once with:\n" + f" mkdir -p {root} && cd {root} && curl -sLO {URL} && unzip -q text8.zip") + with open(path, "r") as f: + return f.read() + + +def load_text8(root="data/text8", seq_len=256): + """Return (train, valid, test) integer arrays of shape (n_windows, seq_len). + + Windows are non-overlapping and contiguous, which keeps the evaluation + splits strictly disjoint from training: no window can straddle a split + boundary, so no test character is ever seen during training. + """ + tok = Text8Tokenizer() + text = _raw(root) + if len(text) < N_TRAIN + N_VALID + N_TEST: + raise ValueError(f"text8 is {len(text)} chars, expected 100,000,000") + + ids = tok.encode(text) + tr = ids[:N_TRAIN] + va = ids[N_TRAIN:N_TRAIN + N_VALID] + te = ids[N_TRAIN + N_VALID:N_TRAIN + N_VALID + N_TEST] + + def windows(a): + n = len(a) // seq_len + return a[:n * seq_len].reshape(n, seq_len) + + return windows(tr), windows(va), windows(te), tok + + +def batches(arr, bs, seed=0, shuffle=True): + """Infinite batch iterator over pre-windowed data.""" + rng = np.random.default_rng(seed) + n = len(arr) + while True: + idx = rng.permutation(n) if shuffle else np.arange(n) + for i in range(0, n - bs + 1, bs): + yield arr[idx[i:i + bs]] diff --git a/src/train_eval_lm.py b/src/train_eval_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..9c86477f2db34c5b8ae7a977be2546ff78110b9f --- /dev/null +++ b/src/train_eval_lm.py @@ -0,0 +1,94 @@ +"""Autoregressive evaluator for generative perplexity on text8. + +Generative perplexity - the likelihood an independent language model assigns to +our samples - is the standard quality metric for diffusion language models, +because the diffusion model's own ELBO says nothing about the quality of what it +actually generates at a low number of function evaluations. + +The evaluator is trained here from scratch on the same corpus rather than using +a pretrained GPT-2, for two reasons. The project's constraint is that everything +is trained from scratch, and a character-level evaluator trained on text8 itself +is the correct judge for text8 samples - a word-piece model trained on web text +would score character-level artifacts through a tokenizer that never sees them. + +The evaluator is trained ONCE and shared by every diffusion run, so all systems +are judged by exactly the same yardstick. +""" +import argparse, json, math, os, time +import numpy as np +import torch +import torch.nn.functional as F + +from model import Transformer +from text8 import load_text8, batches + +p = argparse.ArgumentParser() +p.add_argument("--root", default="data/text8") +p.add_argument("--seq_len", type=int, default=256) +p.add_argument("--d", type=int, default=512) +p.add_argument("--layers", type=int, default=8) +p.add_argument("--heads", type=int, default=8) +p.add_argument("--bs", type=int, default=128) +p.add_argument("--steps", type=int, default=60000) +p.add_argument("--lr", type=float, default=3e-4) +p.add_argument("--warmup", type=int, default=2000) +p.add_argument("--seed", type=int, default=0) +p.add_argument("--out", default="runs/evallm") +a = p.parse_args() + +torch.manual_seed(a.seed) +np.random.seed(a.seed) +dev = "cuda" if torch.cuda.is_available() else "cpu" +os.makedirs(a.out, exist_ok=True) + +tr, va, te, tok = load_text8(a.root, a.seq_len) +print(f"text8: train {tr.shape}, valid {va.shape}, test {te.shape}, vocab {len(tok)}", + flush=True) + +model = Transformer(len(tok), a.d, a.layers, a.heads, "rope", + causal=True, max_len=a.seq_len + 8).to(dev) +print(f"evaluator params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M", flush=True) +opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01, betas=(0.9, 0.95)) +sched = torch.optim.lr_scheduler.LambdaLR( + opt, lambda s: min((s + 1) / a.warmup, max(0.0, (a.steps - s) / max(1, a.steps - a.warmup)))) + +it = batches(tr, a.bs, seed=a.seed) +t0 = time.time() +for step in range(a.steps): + x = torch.from_numpy(next(it)).to(dev) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): + logits = model(x) + # next-character prediction: position i predicts position i+1 + loss = F.cross_entropy(logits[:, :-1].reshape(-1, len(tok)), + x[:, 1:].reshape(-1)) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step(); sched.step() + if step % 2000 == 0: + print(f"step {step:6d} loss {loss.item():.4f} " + f"bpc {loss.item()/math.log(2):.3f} {time.time()-t0:.0f}s", flush=True) + + +@torch.no_grad() +def test_bpc(arr, n=200): + model.eval() + tot, cnt = 0.0, 0 + for i in range(0, min(n * a.bs, len(arr)), a.bs): + x = torch.from_numpy(arr[i:i + a.bs]).to(dev) + if len(x) < 2: + continue + logits = model(x) + ce = F.cross_entropy(logits[:, :-1].reshape(-1, len(tok)), + x[:, 1:].reshape(-1), reduction="sum") + tot += float(ce); cnt += x[:, 1:].numel() + model.train() + return tot / cnt / math.log(2) + + +bpc = test_bpc(te) +print(f"\nevaluator test BPC: {bpc:.4f}", flush=True) +torch.save(model.state_dict(), os.path.join(a.out, "model.pt")) +json.dump({"args": vars(a), "test_bpc": bpc, "vocab": len(tok)}, + open(os.path.join(a.out, "result.json"), "w"), indent=2) +print(f"saved to {a.out}", flush=True) diff --git a/src/train_sig.py b/src/train_sig.py index 64c73d9d6b837567be25f8afccf163df4393ce1c..fd922822db8ee027cbf8fea3d292321fc67a8c26 100644 --- a/src/train_sig.py +++ b/src/train_sig.py @@ -23,7 +23,7 @@ def get_args(): p = argparse.ArgumentParser() p.add_argument("--lang", default="english") p.add_argument("--size", default="medium", choices=["low", "medium", "high"]) - p.add_argument("--mode", choices=["ar", "mdlm", "full", "prog"], default="mdlm", + p.add_argument("--mode", choices=["ar", "mdlm", "matched", "distill", "full", "prog"], default="mdlm", help="ar = autoregressive baseline; full = ours") p.add_argument("--seq_len", type=int, default=48) p.add_argument("--d", type=int, default=256) diff --git a/src/train_text8.py b/src/train_text8.py new file mode 100644 index 0000000000000000000000000000000000000000..ea5f61bee68fd9d9b0a89367656ffb9a027f36f7 --- /dev/null +++ b/src/train_text8.py @@ -0,0 +1,199 @@ +"""Masked diffusion on text8, and the experiment that tests the unified claim. + +Addition isolates failure mode (i): every conditional is a point mass, so all of +the loss from parallel decoding comes from marginals that were never trained at +high mask ratios, and no inference-time rule recovers any of it. The free-choice +probes isolate failure mode (ii): the conditionals are maximum-entropy, nothing +is trainable, and committing less is the only remedy. + +Natural text contains both in the same sequence. Finishing a word, a suffix or a +function word is nearly deterministic; choosing the next content word is not. So +text is where the two fixes have to compose, and this script measures whether +they do, as a 2x2: + + fixed-K commit entropy-budgeted commit + uniform-t training (a) (b) + matched training (c) (d) + +The prediction is (d) < (b), (c) < (a) on generative perplexity at equal compute, +and that the two gains are largely independent - one fixes what the other cannot +touch. + +Metrics, all standard for this benchmark: + * validation bits-per-character (the MDLM bound) - likelihood, comparable to + published text8 numbers, and the check that the schedule change does not + buy sample quality by wrecking the model. + * generative perplexity under a held-out autoregressive evaluator, against + the number of function evaluations actually used. + * entropy and distinct-n-gram rate of the samples, reported alongside, because + generative perplexity alone is trivially gamed by degenerate repetition. +""" +import argparse, json, math, os, time +import numpy as np +import torch +import torch.nn.functional as F + +from model import Transformer +from text8 import load_text8, batches +import diffusion as dfn + + +def get_args(): + p = argparse.ArgumentParser() + p.add_argument("--mode", choices=["mdlm", "matched", "full"], default="mdlm") + p.add_argument("--K", type=int, default=8, help="parallelism the schedule is matched to") + p.add_argument("--lam", type=float, default=1.0) + p.add_argument("--root", default="data/text8") + p.add_argument("--seq_len", type=int, default=256) + p.add_argument("--d", type=int, default=512) + p.add_argument("--layers", type=int, default=12) + p.add_argument("--heads", type=int, default=8) + p.add_argument("--bs", type=int, default=64) + p.add_argument("--steps", type=int, default=80000) + p.add_argument("--lr", type=float, default=1e-4) + p.add_argument("--warmup", type=int, default=4000) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--n_gen", type=int, default=256, help="samples per decoding setting") + p.add_argument("--gen_bs", type=int, default=64) + p.add_argument("--evallm", default="runs/evallm") + p.add_argument("--out", default="runs/t8dev") + return p.parse_args() + + +a = get_args() +torch.manual_seed(a.seed) +np.random.seed(a.seed) +dev = "cuda" if torch.cuda.is_available() else "cpu" +os.makedirs(a.out, exist_ok=True) + +tr, va, te, tok = load_text8(a.root, a.seq_len) +print(f"text8: train {tr.shape} valid {va.shape} test {te.shape} vocab {len(tok)}", + flush=True) + +model = Transformer(len(tok), a.d, a.layers, a.heads, "rope", + causal=False, max_len=a.seq_len + 8).to(dev) +print(f"denoiser params: {sum(q.numel() for q in model.parameters())/1e6:.1f}M", flush=True) +opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=0.01, betas=(0.9, 0.95)) +sched = torch.optim.lr_scheduler.LambdaLR( + opt, lambda s: min((s + 1) / a.warmup, max(0.0, (a.steps - s) / max(1, a.steps - a.warmup)))) + +it = batches(tr, a.bs, seed=a.seed) +t0 = time.time() +for step in range(a.steps): + x = torch.from_numpy(next(it)).to(dev) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=(dev == "cuda")): + loss = dfn.text_loss(model, x, tok, a.mode, a.K, a.lam) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step(); sched.step() + if step % 2000 == 0: + print(f"step {step:6d} loss {loss.item():.4f} {time.time()-t0:.0f}s", flush=True) + +model.eval() +res = {"args": vars(a)} + +# ---- likelihood: the standard text8 number --------------------------------- +vb = [torch.from_numpy(va[i:i + a.bs]) for i in range(0, min(64 * a.bs, len(va)), a.bs)] +res["valid_bpc"] = dfn.nll_bpc(model, vb, tok, dev, mc=8) +print(f"\nvalidation BPC (MDLM bound): {res['valid_bpc']:.4f}", flush=True) + +# ---- the evaluator --------------------------------------------------------- +ev = None +ev_path = os.path.join(a.evallm, "model.pt") +if os.path.exists(ev_path): + ec = json.load(open(os.path.join(a.evallm, "result.json")))["args"] + ev = Transformer(len(tok), ec["d"], ec["layers"], ec["heads"], "rope", + causal=True, max_len=ec["seq_len"] + 8).to(dev) + ev.load_state_dict(torch.load(ev_path, map_location=dev)) + ev.eval() + print(f"loaded evaluator from {a.evallm}", flush=True) +else: + print(f"WARNING: no evaluator at {ev_path}; generative perplexity skipped", flush=True) + + +@torch.no_grad() +def gen_ppl(x): + """Per-character perplexity the held-out autoregressive model assigns.""" + tot, cnt = 0.0, 0 + for i in range(0, len(x), a.gen_bs): + b = x[i:i + a.gen_bs].to(dev) + ce = F.cross_entropy(ev(b)[:, :-1].reshape(-1, len(tok)), + b[:, 1:].reshape(-1), reduction="sum") + tot += float(ce); cnt += b[:, 1:].numel() + return math.exp(tot / cnt) + + +def sample_stats(x): + """Entropy and distinct-n-gram rate: the degeneracy guard. + + A model that emits the same low-perplexity string every time wins on + generative perplexity while being useless, so the unigram entropy and the + fraction of distinct 4-grams are reported next to it. Reference values from + real text8 windows are computed below for comparison. + """ + f = np.bincount(x.reshape(-1).numpy(), minlength=len(tok)).astype(float) + f = f[f > 0] / f.sum() + ent = float(-(f * np.log(f)).sum()) + grams = set() + tot = 0 + for row in x.numpy(): + for i in range(len(row) - 3): + grams.add(tuple(row[i:i + 4])); tot += 1 + return ent, len(grams) / max(tot, 1) + + +ref = torch.from_numpy(te[:a.n_gen]) +res["reference"] = {"gen_ppl": gen_ppl(ref) if ev else None, + "entropy": sample_stats(ref)[0], + "distinct4": sample_stats(ref)[1]} +print(f"real text8 reference: ppl {res['reference']['gen_ppl']}, " + f"entropy {res['reference']['entropy']:.3f}, " + f"distinct-4 {res['reference']['distinct4']:.3f}", flush=True) + + +def blank(): + return torch.full((a.gen_bs, a.seq_len), tok.mask, dtype=torch.long) + + +def evaluate(tag, fn): + """Run one decoding rule over n_gen samples and record quality vs compute.""" + outs, nfes = [], [] + for _ in range(max(1, a.n_gen // a.gen_bs)): + x, nfe = fn(blank()) + outs.append(x.cpu()); nfes.append(nfe) + x = torch.cat(outs) + ent, d4 = sample_stats(x) + row = {"nfe": float(np.mean(nfes)), + "gen_ppl": gen_ppl(x) if ev else None, + "entropy": ent, "distinct4": d4, + "sample": tok.decode(x[0][:120].tolist())} + print(f" {tag:<26} nfe {row['nfe']:6.1f} ppl {row['gen_ppl']:9.2f} " + f"H {ent:.3f} d4 {d4:.3f}", flush=True) + return row + + +print("\n=== decoding: quality against compute ===", flush=True) +res["decode"] = {} + +print("fixed-K (non-adaptive reference):", flush=True) +for K in (1, 2, 4, 8, 16, 32, 64, 128, 256): + res["decode"][f"fixedK_{K}"] = evaluate( + f"fixed-K K={K}", + lambda b, K=K: dfn.fixed_k_decode(model, b, tok, dev, K)) + +print("confidence-threshold (published inference-time method):", flush=True) +for tau in (0.5, 0.7, 0.9, 0.95, 0.99): + res["decode"][f"conf_{tau}"] = evaluate( + f"confidence tau={tau}", + lambda b, tau=tau: dfn.confidence_threshold_decode(model, b, tok, dev, tau)) + +print("entropy-budget (ours, fix (ii)):", flush=True) +for B in (0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0): + res["decode"][f"entbudget_{B}"] = evaluate( + f"entropy-budget B={B}", + lambda b, B=B: dfn.entropy_budget_decode(model, b, tok, dev, B)) + +torch.save(model.state_dict(), os.path.join(a.out, "model.pt")) +json.dump(res, open(os.path.join(a.out, "result.json"), "w"), indent=2) +print(f"\nsaved to {a.out}", flush=True)