Automated push from Zaratan

results/g3-gam0.0-elbo0.05-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo0.05-s0"
+  },
+  "minutes": 5.432604718208313,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.0-elbo0.05-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo0.05-s1"
+  },
+  "minutes": 5.529767870903015,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.0-elbo0.2-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo0.2-s0"
+  },
+  "minutes": 5.417752428849538,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.0-elbo0.2-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo0.2-s1"
+  },
+  "minutes": 5.443442535400391,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.0-elbo1.0-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo1.0-s0"
+  },
+  "minutes": 5.452999484539032,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.0-elbo1.0-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.0-elbo1.0-s1"
+  },
+  "minutes": 5.407823407649994,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo0.05-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo0.05-s0"
+  },
+  "minutes": 8.475962579250336,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo0.05-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo0.05-s1"
+  },
+  "minutes": 8.52948188384374,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo0.2-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo0.2-s0"
+  },
+  "minutes": 8.450251412391662,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo0.2-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo0.2-s1"
+  },
+  "minutes": 8.430337603886922,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo1.0-s0.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo1.0-s0"
+  },
+  "minutes": 8.355311985810598,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.1-elbo1.0-s1.json

@@ -0,0 +1,48 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.1-elbo1.0-s1"
+  },
+  "minutes": 8.506091515223185,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo0.05-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo0.05-s0"
+  },
+  "minutes": 10.564957304795582,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo0.05-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo0.05-s1"
+  },
+  "minutes": 10.484416369597117,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo0.2-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo0.2-s0"
+  },
+  "minutes": 10.366264355182647,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo0.2-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo0.2-s1"
+  },
+  "minutes": 10.578846287727355,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo1.0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo1.0-s0"
+  },
+  "minutes": 10.441694112618764,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam0.5-elbo1.0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam0.5-elbo1.0-s1"
+  },
+  "minutes": 10.469638713200887,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo0.05-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo0.05-s0"
+  },
+  "minutes": 10.54262443780899,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo0.05-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo0.05-s1"
+  },
+  "minutes": 10.598351208368937,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo0.2-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo0.2-s0"
+  },
+  "minutes": 10.403599965572358,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo0.2-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo0.2-s1"
+  },
+  "minutes": 10.530384941895802,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo1.0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo1.0-s0"
+  },
+  "minutes": 10.438667662938435,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam1.0-elbo1.0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 1.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam1.0-elbo1.0-s1"
+  },
+  "minutes": 10.58419309457143,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo0.05-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo0.05-s0"
+  },
+  "minutes": 10.587692022323608,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo0.05-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.05,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo0.05-s1"
+  },
+  "minutes": 10.614048687616984,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo0.2-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo0.2-s0"
+  },
+  "minutes": 10.523026792208354,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo0.2-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 0.2,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo0.2-s1"
+  },
+  "minutes": 10.575358454386393,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo1.0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo1.0-s0"
+  },
+  "minutes": 10.310235365231831,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g3-gam2.0-elbo1.0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g3-gam2.0-elbo1.0-s1"
+  },
+  "minutes": 10.305338247617085,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-agree-gam0.5-dw0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 0.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-agree-gam0.5-dw0-s0"
+  },
+  "minutes": 6.659133970737457,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-agree-gam0.5-dw0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 0.0,
+    "2": 0.0,
+    "4": 0.0,
+    "8": 0.0,
+    "16": 0.0,
+    "32": 0.0,
+    "64": 0.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "agree",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-agree-gam0.5-dw0-s1"
+  },
+  "minutes": 6.71096578836441,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam0.5-dw0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam0.5-dw0-s0"
+  },
+  "minutes": 6.716651705900828,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam0.5-dw0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam0.5-dw0-s1"
+  },
+  "minutes": 6.62840025027593,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam0.5-dw1-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam0.5-dw1-s0"
+  },
+  "minutes": 8.691880404949188,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam0.5-dw1-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 0.5,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam0.5-dw1-s1"
+  },
+  "minutes": 8.653803006807964,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam2.0-dw0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 2.0,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam2.0-dw0-s0"
+  },
+  "minutes": 6.72961677312851,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam2.0-dw0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 2.0,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam2.0-dw0-s1"
+  },
+  "minutes": 6.645028352737427,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam2.0-dw1-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam2.0-dw1-s0"
+  },
+  "minutes": 8.62372316122055,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam2.0-dw1-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 2.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam2.0-dw1-s1"
+  },
+  "minutes": 8.700080172220867,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam8.0-dw0-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 8.0,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam8.0-dw0-s0"
+  },
+  "minutes": 6.665486351648966,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam8.0-dw0-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 8.0,
+    "dep_weighted": 0,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam8.0-dw0-s1"
+  },
+  "minutes": 6.602849837144216,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam8.0-dw1-s0.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 8.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 0,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam8.0-dw1-s0"
+  },
+  "minutes": 8.781805845101674,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g4-copy-gam8.0-dw1-s1.json

@@ -0,0 +1,49 @@
+{
+  "accuracy_by_passes": {
+    "1": 1.0,
+    "2": 1.0,
+    "4": 1.0,
+    "8": 1.0,
+    "16": 1.0,
+    "32": 1.0,
+    "64": 1.0,
+    "128": 1.0
+  },
+  "args": {
+    "objective": "coord",
+    "alpha": 0.5,
+    "beta": 0.0,
+    "gamma": 8.0,
+    "dep_weighted": 1,
+    "elbo_w": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.2,
+    "task": "copy",
+    "coupling": 4,
+    "seq_len": 128,
+    "n_train": 200000,
+    "n_eval": 500,
+    "d": 384,
+    "layers": 6,
+    "heads": 6,
+    "pe": "rope",
+    "bs": 64,
+    "steps": 20000,
+    "lr": 0.0003,
+    "warmup": 500,
+    "seed": 1,
+    "data_root": "data",
+    "eval_bs": 100,
+    "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g4-copy-gam8.0-dw1-s1"
+  },
+  "minutes": 8.761348100503286,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

slurm/g4.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# Grid 4 - the method: dependence-weighted entropy regularization.
+#
+# Prediction, and the thing that makes it falsifiable:
+#   agree/parity  (positions dependent)   -> accuracy at high parallelism RISES
+#   copy          (positions independent) -> UNCHANGED at 100%, because measured
+#                                            dependence is ~0 so no penalty applies
+# The un-weighted ablation (dep_weighted=0) should hurt `copy`, which is how we
+# show the weighting is doing the work rather than the entropy term alone.
+#SBATCH -J g4
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=64g
+#SBATCH -t 06:00:00
+#SBATCH -o logs/g4-%A_%a.out
+#SBATCH --array=0-35%12
+source /etc/profile
+SCR=/scratch/zt1/project/msml612/user/$USER
+module load python/3.10.10
+export PYTHONPATH=
+PY=$SCR/hf_env/bin/python
+cd $SCR/pcmd/src
+
+I=$SLURM_ARRAY_TASK_ID
+TASKS=(copy agree parity); GAMMA=(0.5 2.0 8.0); DW=(1 0); SEEDS=(0 1)
+T=${TASKS[$((I / 12))]}
+G=${GAMMA[$(((I % 12) / 4))]}
+D=${DW[$(((I % 4) / 2))]}
+S=${SEEDS[$((I % 2))]}
+
+NAME="g4-$T-gam$G-dw$D-s$S"
+echo "=== $NAME ==="
+$PY train.py --task $T --objective coord --coupling 4 --seed $S \
+  --seq_len 128 --n_train 200000 --n_eval 500 --eval_bs 100 \
+  --d 384 --layers 6 --heads 6 --bs 64 --steps 20000 --lr 3e-4 \
+  --beta 0.0 --gamma $G --dep_weighted $D --elbo_w 1.0 --ramp_frac 0.2 \
+  --out $SCR/pcmd/runs/$NAME

slurm/grid3.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Grid 3 — trace the likelihood <-> parallel-decodability trade-off.
+#
+# Grid 2 showed coordination alone is cancelled by the ELBO: the sequential
+# teacher picks an arbitrary value for a free group, so averaged over batches
+# its target is still uniform, while the ELBO actively pulls back to uniform.
+# The weight IS the experiment. gamma penalises marginal entropy (explicit
+# symmetry breaking); elbo_w down-weights the term fighting it.
+#SBATCH -J g3
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=64g
+#SBATCH -t 06:00:00
+#SBATCH -o logs/g3-%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=
+PY=$SCR/hf_env/bin/python
+cd $SCR/pcmd/src
+
+I=$SLURM_ARRAY_TASK_ID
+GAMMA=(0.0 0.1 0.5 1.0 2.0); ELBOW=(1.0 0.2 0.05); SEEDS=(0 1)
+G=${GAMMA[$((I / 6))]}
+E=${ELBOW[$(((I % 6) / 2))]}
+S=${SEEDS[$((I % 2))]}
+
+NAME="g3-gam$G-elbo$E-s$S"
+echo "=== $NAME ==="
+$PY train.py --task agree --objective coord --coupling 4 --seed $S \
+  --seq_len 128 --n_train 200000 --n_eval 500 --eval_bs 100 \
+  --d 384 --layers 6 --heads 6 --bs 64 --steps 20000 --lr 3e-4 \
+  --beta 1.0 --gamma $G --elbo_w $E --group 4 --coord_every 4 --ramp_frac 0.2 \
+  --out $SCR/pcmd/runs/$NAME

src/diffusion.py

@@ -224,18 +224,96 @@ def coordination_loss(model, x0, tok, group=4, t_min=None):
     return F.cross_entropy(sel.reshape(-1, sel.size(-1)), target.reshape(-1))
 
 
+def entropy_penalty(model, x0, tok, t_min=None):
+    """Break the symmetry that makes free groups undecodable in parallel.
+
+    Where a group must agree but the value is free, the likelihood-optimal
+    marginal is uniform over the valid values, and independent sampling from
+    uniform marginals can never coordinate. Penalising marginal entropy pushes
+    each position toward a single mode; because the model is shared across the
+    group, the same mode is chosen at each member and parallel sampling becomes
+    self-consistent. This is the term that is *supposed* to cost likelihood -
+    it is the explicit price of parallel decodability.
+    """
+    B, L = x0.shape
+    dev = x0.device
+    t_min = t_min if t_min is not None else 1.0 / L
+    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
+    m = torch.rand(B, L, device=dev) < t[:, None]
+    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
+    p = model(xt).softmax(-1)
+    ent = -(p * p.clamp_min(1e-9).log()).sum(-1)
+    return (ent * m).sum() / m.sum().clamp_min(1)
+
+
 def cmd_loss(model, x0, tok, alpha=0.5, beta=1.0, k_frac=0.25, ramp=1.0,
-             group=4, use_coord=True):
+             group=4, use_coord=True, gamma=0.0, elbo_w=1.0,
+             dep_weighted=True):
     """Coordinated Masked Diffusion: ELBO + self-conditioning + coordination.
 
     The ELBO term is always present. The two additions are auxiliaries, not
     replacements: an earlier version used self-conditioning INSTEAD of the ELBO
     and scored 35% where the plain baseline scored 100%.
     """
-    total = mdlm_loss(model, x0, tok)
+    total = elbo_w * mdlm_loss(model, x0, tok)
+    if gamma > 0 and ramp > 0:
+        total = total + gamma * ramp * (
+            dependence_weighted_entropy(model, x0, tok) if dep_weighted
+            else entropy_penalty(model, x0, tok))
     if alpha > 0 and ramp > 0:
         total = total + alpha * self_corrupted_loss(
             model, x0, tok, k_frac=k_frac, correct_weight=1.0, ramp=ramp)
     if beta > 0 and ramp > 0 and use_coord:
         total = total + beta * ramp * coordination_loss(model, x0, tok, group=group)
     return total
+
+
+def dependence_weighted_entropy(model, x0, tok, t_min=None):
+    """THE METHOD. Sharpen the marginals exactly where positions are dependent.
+
+    Why entropy is the right quantity. Two independent draws from a marginal p
+    agree with probability sum_v p(v)^2 = exp(-H2(p)), the Renyi-2 entropy. So
+    minimising H2 directly maximises the probability that positions sampled
+    independently - which is precisely what parallel decoding does - land on the
+    same choice. For a group that must agree but whose value is free, this is
+    the only way to make parallel sampling self-consistent: the likelihood
+    optimum (uniform over valid values) is guaranteed to disagree.
+
+    Why it must be WEIGHTED. Sharpening everywhere would destroy likelihood
+    wherever high entropy is correct, which is most of natural text, and would
+    also break `copy`, where positions are already independent and the baseline
+    is perfect. So the penalty is scaled by a measured dependence:
+
+        reveal one masked position, and see how far every other position's
+        belief moves. A position whose distribution shifts a lot depends on the
+        revealed one; a position that does not move is conditionally
+        independent and is left alone.
+
+    d_i = KL(p_i-after-reveal || p_i-before), normalised per example. The
+    penalty is d_i * H2(p_i), so the likelihood cost is paid only where it buys
+    coordination.
+    """
+    B, L = x0.shape
+    dev = x0.device
+    t_min = t_min if t_min is not None else 1.0 / L
+    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
+    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)
+
+    with torch.no_grad():
+        p0 = model(xt).softmax(-1)
+        samp = torch.distributions.Categorical(probs=p0).sample()
+        # reveal one masked position per example and measure the ripple
+        pick = torch.rand(B, L, device=dev).masked_fill(~m, -1.0).argmax(1)
+        x1 = xt.clone()
+        ar = torch.arange(B, device=dev)
+        x1[ar, pick] = samp[ar, pick]
+        p1 = model(x1).softmax(-1)
+        d = (p1 * (p1.clamp_min(1e-9).log() - p0.clamp_min(1e-9).log())).sum(-1)
+        d = (d * m).clamp_min(0)
+        d = d / d.amax(1, keepdim=True).clamp_min(1e-6)      # per-example scale
+
+    p = model(xt).softmax(-1)
+    h2 = -(p.pow(2).sum(-1).clamp_min(1e-9)).log()           # Renyi-2 entropy
+    return (d * h2 * m).sum() / m.sum().clamp_min(1)

src/train.py

@@ -22,6 +22,14 @@ def get_args():
                         "coord=+coordination; cmd=full method (all three)")
     p.add_argument("--alpha", type=float, default=0.5, help="self-conditioning weight")
     p.add_argument("--beta", type=float, default=1.0, help="coordination weight")
+    p.add_argument("--gamma", type=float, default=0.0,
+                   help="marginal-entropy penalty: the explicit price paid in "
+                        "likelihood to buy parallel decodability")
+    p.add_argument("--dep_weighted", type=int, default=1,
+                   help="1 = sharpen only where positions are measurably dependent; "
+                        "0 = sharpen everywhere (ablation)")
+    p.add_argument("--elbo_w", type=float, default=1.0,
+                   help="down-weight the ELBO so coordination is not cancelled by it")
     p.add_argument("--group", type=int, default=4)
     p.add_argument("--coord_every", type=int, default=4,
                    help="apply the coordination term every N steps; the teacher\n                         costs `group` sequential passes, so this amortises it")
@@ -157,10 +165,14 @@ def main():
             uc = (step % a.coord_every == 0)
             if a.objective == "cmd":
                 loss = dfn.cmd_loss(model, x0, tok, a.alpha, a.beta, a.k_frac,
-                                    ramp, a.group, use_coord=uc)
+                                    ramp, a.group, use_coord=uc,
+                                    gamma=a.gamma, elbo_w=a.elbo_w,
+                                    dep_weighted=bool(a.dep_weighted))
             elif a.objective == "coord":
                 loss = dfn.cmd_loss(model, x0, tok, 0.0, a.beta, a.k_frac,
-                                    ramp, a.group, use_coord=uc)
+                                    ramp, a.group, use_coord=uc,
+                                    gamma=a.gamma, elbo_w=a.elbo_w,
+                                    dep_weighted=bool(a.dep_weighted))
             elif a.objective == "selfcorrupt":
                 loss = dfn.cmd_loss(model, x0, tok, a.alpha, 0.0, a.k_frac, ramp)
             else: