Automated push from Zaratan

results/g1-copy-mdlm-s0.json

@@ -13,6 +13,7 @@
     "objective": "mdlm",
     "k_frac": 0.25,
     "correct_weight": 1.0,
+    "ramp_frac": 0.3,
     "task": "copy",
     "coupling": 4,
     "seq_len": 128,
@@ -31,7 +32,7 @@
     "eval_bs": 100,
     "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-mdlm-s0"
   },
-  "minutes": 3.757214069366455,
+  "minutes": 3.7576344490051268,
   "compute": {
     "params": 10660608,
     "train_tokens": 163840000,

results/g1-copy-mdlm-s1.json

@@ -13,6 +13,7 @@
     "objective": "mdlm",
     "k_frac": 0.25,
     "correct_weight": 1.0,
+    "ramp_frac": 0.3,
     "task": "copy",
     "coupling": 4,
     "seq_len": 128,
@@ -31,7 +32,7 @@
     "eval_bs": 100,
     "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-mdlm-s1"
   },
-  "minutes": 3.7269402066866557,
+  "minutes": 3.7286835312843323,
   "compute": {
     "params": 10660608,
     "train_tokens": 163840000,

results/g1-copy-selfcorrupt-s0.json

@@ -1,18 +1,19 @@
 {
   "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
+    "1": 0.346,
+    "2": 0.346,
+    "4": 0.602,
+    "8": 0.6639999999999999,
+    "16": 0.6460000000000001,
+    "32": 0.64,
+    "64": 0.6200000000000001,
+    "128": 0.6140000000000001
   },
   "args": {
     "objective": "selfcorrupt",
     "k_frac": 0.25,
     "correct_weight": 1.0,
+    "ramp_frac": 0.3,
     "task": "copy",
     "coupling": 4,
     "seq_len": 128,
@@ -31,7 +32,7 @@
     "eval_bs": 100,
     "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-selfcorrupt-s0"
   },
-  "minutes": 4.29505068063736,
+  "minutes": 4.341549547513326,
   "compute": {
     "params": 10660608,
     "train_tokens": 163840000,

results/g1-copy-selfcorrupt-s1.json

@@ -1,18 +1,19 @@
 {
   "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
+    "1": 0.024,
+    "2": 0.024,
+    "4": 0.21800000000000003,
+    "8": 0.28400000000000003,
+    "16": 0.306,
+    "32": 0.33199999999999996,
+    "64": 0.354,
+    "128": 0.352
   },
   "args": {
     "objective": "selfcorrupt",
     "k_frac": 0.25,
     "correct_weight": 1.0,
+    "ramp_frac": 0.3,
     "task": "copy",
     "coupling": 4,
     "seq_len": 128,
@@ -31,7 +32,7 @@
     "eval_bs": 100,
     "out": "/scratch/zt1/project/msml612/user/govind02/pcmd/runs/g1-copy-selfcorrupt-s1"
   },
-  "minutes": 4.269979015986125,
+  "minutes": 4.30184254248937,
   "compute": {
     "params": 10660608,
     "train_tokens": 163840000,

results/g2-agree-mdlm-s0.json

@@ -0,0 +1,46 @@
+{
+  "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": "mdlm",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-agree-mdlm-s0"
+  },
+  "minutes": 3.6784939726193744,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g2-agree-mdlm-s1.json

@@ -0,0 +1,46 @@
+{
+  "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": "mdlm",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-agree-mdlm-s1"
+  },
+  "minutes": 3.7346819678942365,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g2-copy-coord-s0.json

@@ -0,0 +1,46 @@
+{
+  "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": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-copy-coord-s0"
+  },
+  "minutes": 5.439730898539225,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g2-copy-coord-s1.json

@@ -0,0 +1,46 @@
+{
+  "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": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-copy-coord-s1"
+  },
+  "minutes": 5.407112606366476,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g2-copy-mdlm-s0.json

@@ -0,0 +1,46 @@
+{
+  "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": "mdlm",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-copy-mdlm-s0"
+  },
+  "minutes": 3.6734960357348125,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

results/g2-copy-mdlm-s1.json

@@ -0,0 +1,46 @@
+{
+  "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": "mdlm",
+    "alpha": 0.5,
+    "beta": 1.0,
+    "group": 4,
+    "coord_every": 4,
+    "k_frac": 0.25,
+    "correct_weight": 1.0,
+    "ramp_frac": 0.3,
+    "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/g2-copy-mdlm-s1"
+  },
+  "minutes": 3.669589972496033,
+  "compute": {
+    "params": 10660608,
+    "train_tokens": 163840000,
+    "train_flops_approx": 10479804088320000,
+    "fwd_per_step": 1
+  }
+}

slurm/grid1b.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# Grid 1 — learnability gate and the first quality-vs-parallelism curves.
+#
+# Doubles as a gate: if a model cannot solve a probe even at FULLY SEQUENTIAL
+# decoding (128 passes for a 128-token canvas), the task is not learnable at
+# this scale and no parallel-decoding claim built on it would mean anything.
+# Launching a large grid before establishing this is the mistake that cost a
+# day on the previous project.
+#SBATCH -J g1
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=64g
+#SBATCH -t 04:00:00
+#SBATCH -o logs/g1-%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=
+PY=$SCR/hf_env/bin/python
+cd $SCR/pcmd/src
+
+I=$SLURM_ARRAY_TASK_ID
+TASKS=(copy); OBJ=(mdlm selfcorrupt); SEEDS=(0 1)
+T=${TASKS[$((I / 4))]}
+O=${OBJ[$(((I % 4) / 2))]}
+S=${SEEDS[$((I % 2))]}
+
+NAME="g1-$T-$O-s$S"
+echo "=== $NAME ==="
+$PY train.py --task $T --objective $O --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 \
+  --out $SCR/pcmd/runs/$NAME

slurm/grid2.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Grid 2 — the method vs its ablations vs the baseline, on all three probes.
+#
+# Objectives
+#   mdlm        ELBO only (standard masked diffusion)
+#   selfcorrupt ELBO + self-conditioning        (component 1)
+#   coord       ELBO + coordination             (component 2, the core)
+#   cmd         ELBO + both                     (full method)
+#
+# The probes make the claim falsifiable: `copy` has independent answer
+# positions, so the method must NOT hurt there; `agree` and `parity` have
+# jointly-dependent positions, which is where it must win.
+#SBATCH -J g2
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=64g
+#SBATCH -t 08:00:00
+#SBATCH -o logs/g2-%A_%a.out
+#SBATCH --array=0-23%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); OBJ=(mdlm selfcorrupt coord cmd); SEEDS=(0 1)
+T=${TASKS[$((I / 8))]}
+O=${OBJ[$(((I % 8) / 2))]}
+S=${SEEDS[$((I % 2))]}
+
+NAME="g2-$T-$O-s$S"
+echo "=== $NAME ==="
+$PY train.py --task $T --objective $O --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 \
+  --alpha 0.5 --beta 1.0 --group 4 --coord_every 4 --ramp_frac 0.3 \
+  --out $SCR/pcmd/runs/$NAME

slurm/sm.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+#SBATCH -J smoke
+#SBATCH -p gpu-h100
+#SBATCH --gpus=h100:1
+#SBATCH -c 8
+#SBATCH --mem=32g
+#SBATCH -t 00:25:00
+#SBATCH -o logs/smoke-%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 O in mdlm selfcorrupt coord cmd; do
+  echo "--- $O ---"
+  $SCR/hf_env/bin/python train.py --objective $O --task agree --coupling 4 \
+    --seq_len 32 --n_train 2000 --n_eval 40 --d 64 --layers 2 --heads 2 --bs 16 \
+    --steps 60 --group 4 --coord_every 4 --eval_bs 20 --out /tmp/sm$O 2>&1 | tail -2
+done

src/diffusion.py

@@ -37,7 +37,7 @@ def mdlm_loss(model, x0, tok, t_min=None, fixed_t=None):
 
 
 def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
-                        correct_weight=1.0):
+                        correct_weight=1.0, ramp=1.0):
     """OUR METHOD. Train on states that parallel decoding actually produces.
 
     1. mask as usual and take a NO-GRAD pass to get the current marginals
@@ -51,6 +51,12 @@ def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
 
     The model is therefore optimised for the distribution it meets at inference
     rather than for ground-truth context it will never see.
+
+    `ramp` in [0,1] scales the corruption rate. It must start near zero: a
+    randomly-initialised model samples noise, so corrupting from step 0 trains
+    on pure garbage context and the model never learns the task at all (it
+    scored 0% even on copy, which the baseline solves trivially). This is the
+    same reason scheduled sampling anneals its mixing probability.
     """
     B, L = x0.shape
     dev = x0.device
@@ -65,7 +71,7 @@ def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
         samp = torch.distributions.Categorical(probs=probs).sample()
 
     # commit a random subset of masked positions from the model's own samples
-    commit = m & (torch.rand(B, L, device=dev) < k_frac)
+    commit = m & (torch.rand(B, L, device=dev) < k_frac * ramp)
     x_self = torch.where(commit, samp, xt)
     still = m & ~commit
 
@@ -73,11 +79,11 @@ def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
     ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                          reduction="none").view(B, L)
 
-    denom = still.sum(1).clamp_min(1)
-    main = ((1.0 / t) * (ce * still).sum(1) / denom).mean()
-    # repair term: committed positions are supervised against ground truth too
-    rep = (ce * commit).sum(1) / commit.sum(1).clamp_min(1)
-    return main + correct_weight * rep.mean()
+    # Normalise by L exactly as mdlm_loss does, so the two objectives are on the
+    # same scale and any difference in results is the method, not the loss size.
+    main = ((1.0 / t) * (ce * still).sum(1) / L).mean()
+    rep = ((1.0 / t) * (ce * commit).sum(1) / L).mean()
+    return main + correct_weight * rep
 
 
 # --------------------------------------------------------------- decoding ---
@@ -157,3 +163,79 @@ def nll_bpc(model, batches, tok, device, mc=8):
             tot += float(((1.0 / t) * (ce * m).sum(1) / L).sum())
             n += B
     return tot / max(n, 1) / torch.log(torch.tensor(2.0)).item()
+
+
+def coordination_loss(model, x0, tok, group=4, t_min=None):
+    """CORE OF THE METHOD: make parallel sampling reproduce sequential sampling.
+
+    The measured failure: on jointly-constrained-but-individually-free positions
+    (our `agree` probe), each position's likelihood-optimal marginal is uniform
+    over the valid values, so sampling them independently is guaranteed to be
+    inconsistent - 0% accuracy at every parallelism, 100% only when decoded one
+    at a time. No decoding schedule can fix that, because the marginals
+    themselves carry no coordination.
+
+    So we change the marginals. For a random subset S:
+
+      teacher  : decode S SEQUENTIALLY with no gradient, each position
+                 conditioned on the ones already chosen. Because each step sees
+                 the previous choices, the result is a coherent joint sample.
+      student  : the marginals S gets from ONE parallel forward pass.
+      loss     : cross-entropy of the teacher's joint sample under the student's
+                 parallel marginals.
+
+    The model is its own teacher, so this needs no external model and works in
+    pretraining from scratch. It deliberately trades likelihood (marginals move
+    off the true conditional, toward a coordinated mode) for parallel
+    decodability - a trade the standard ELBO can never make.
+
+    COST: the teacher costs `group` sequential forward passes, so `group` is the
+    knob that keeps training affordable. An early version coordinated over 32
+    positions and made every step 32x more expensive, which simply never
+    finished. Combined with applying the term every `coord_every` steps, the
+    amortised overhead is roughly 1 + group/coord_every.
+    """
+    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)
+
+    # choose, per example, a subset S of masked positions to coordinate over
+    score = torch.rand(B, L, device=dev).masked_fill(~m, -1.0)
+    k = min(group, L)
+    S = score.topk(k, dim=1).indices                      # (B, k)
+
+    # ---- teacher: reveal S one position at a time, conditioning as we go ----
+    with torch.no_grad():
+        x_seq = xt.clone()
+        for j in range(k):
+            idx = S[:, j]
+            logits = model(x_seq)
+            pick = logits[torch.arange(B, device=dev), idx].argmax(-1)
+            x_seq[torch.arange(B, device=dev), idx] = pick
+        target = x_seq.gather(1, S)                       # the joint sample
+
+    # ---- student: one parallel pass over the same masked state --------------
+    logits = model(xt)
+    sel = logits.gather(1, S[..., None].expand(-1, -1, logits.size(-1)))
+    return F.cross_entropy(sel.reshape(-1, sel.size(-1)), target.reshape(-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):
+    """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)
+    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

src/train.py

@@ -16,12 +16,23 @@ import diffusion as dfn
 
 def get_args():
     p = argparse.ArgumentParser()
-    p.add_argument("--objective", choices=["mdlm", "selfcorrupt"], default="mdlm",
-                   help="mdlm = standard marginal training; selfcorrupt = ours")
+    p.add_argument("--objective", choices=["mdlm", "selfcorrupt", "coord", "cmd"],
+                   default="mdlm",
+                   help="mdlm=ELBO baseline; selfcorrupt=+self-conditioning; "
+                        "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("--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")
     p.add_argument("--k_frac", type=float, default=0.25,
                    help="fraction of masked positions committed from the model's "
                         "own parallel samples during self-corrupted training")
     p.add_argument("--correct_weight", type=float, default=1.0)
+    p.add_argument("--ramp_frac", type=float, default=0.3,
+                   help="fraction of training over which self-corruption ramps "
+                        "from 0 to --k_frac; corrupting from step 0 trains on a "
+                        "random model's noise and prevents learning entirely")
     p.add_argument("--task", choices=["text8", "copy", "agree", "parity"], default="agree")
     p.add_argument("--coupling", type=int, default=4,
                    help="probe tasks: how many positions are bound together")
@@ -142,10 +153,18 @@ def main():
         x0 = train[idx].to(dev)
         opt.zero_grad(set_to_none=True)
         with torch.autocast(**amp):
-            loss = (dfn.self_corrupted_loss(model, x0, tok, a.k_frac,
-                                            correct_weight=a.correct_weight)
-                    if a.objective == "selfcorrupt"
-                    else dfn.mdlm_loss(model, x0, tok))
+            ramp = min(1.0, step / max(1, a.ramp_frac * a.steps))
+            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)
+            elif a.objective == "coord":
+                loss = dfn.cmd_loss(model, x0, tok, 0.0, a.beta, a.k_frac,
+                                    ramp, a.group, use_coord=uc)
+            elif a.objective == "selfcorrupt":
+                loss = dfn.cmd_loss(model, x0, tok, a.alpha, 0.0, a.k_frac, ramp)
+            else:
+                loss = dfn.mdlm_loss(model, x0, tok)
         loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         opt.step(); sched.step()