1{
2  "op": "replace",
3  "filename": "LLaDA/custom_demask_eval.py",
4  "start_line": 59,
5  "end_line": 151,
6  "content": "class DemaskDecoder:\n    \"\"\"Adaptive Composite Demasking with Consistency Tracking (ACD-CT).\n\n    Novel algorithmic contributions over KLASS:\n    1. Margin-Confidence Fusion: positions ranked by conf * (p1 - p2),\n       combining prediction strength with decisiveness.\n    2. Argmax Consistency: lightweight stability via prediction constancy\n       across steps, replacing full-distribution KL divergence.\n    3. Top-k Annealing Sampler: for fully-parallel text decoding, samples\n       from top-k tokens with decaying temperature for diversity.\n    \"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.history_length = history_length\n        self.sample_k = 10\n        self.sample_temp_start = 0.9\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n        is_parallel = (block_length == gen_length)\n        prev_pred = torch.full_like(x, -1)\n        cons = torch.zeros_like(x, dtype=torch.int)\n        used = 0\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n            num_xfer = get_num_transfer_tokens(\n                (x[:, bs:be] == mid), steps_per_block)\n            cons[:, bs:be] = 0\n            prev_pred[:, bs:be] = -1\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                blk = torch.zeros_like(mask_idx)\n                blk[:, bs:be] = True\n                mask_idx = mask_idx & blk\n                if not mask_idx.any():\n                    break\n                logits = model(x).logits\n                p = F.softmax(logits.to(torch.float64), dim=-1)\n                x0g = torch.argmax(p, dim=-1)\n                cf = torch.gather(p, -1, x0g.unsqueeze(-1)).squeeze(-1)\n                t2v, _ = torch.topk(p, 2, dim=-1)\n                mg = t2v[..., 0] - t2v[..., 1]\n                score = cf * mg\n                if is_parallel and self.sample_k > 1:\n                    prog = step / max(steps_per_block - 1, 1)\n                    ts = max(0.1, self.sample_temp_start * (1 - prog))\n                    sv, si = torch.topk(logits.float(), self.sample_k, -1)\n                    sp = F.softmax(sv / ts, dim=-1)\n                    sh = sp.shape[:-1]\n                    idx = torch.multinomial(sp.reshape(-1, self.sample_k), 1)\n                    x0 = torch.gather(\n                        si, -1, idx.reshape(*sh, 1)).squeeze(-1)\n                else:\n                    x0 = x0g\n                same = (x0g == prev_pred)\n                cons = torch.where(same, cons + 1, torch.ones_like(cons))\n                prev_pred = x0g.clone()\n                stable = cons >= self.history_length\n                ready = stable & (cf > self.conf_threshold) & mask_idx\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                for j in range(ready.shape[0]):\n                    rdy = torch.where(ready[j])[0]\n                    if len(rdy) > 0:\n                        xfer[j, rdy] = True\n                    else:\n                        s = score[j].clone()\n                        s[~mask_idx[j]] = -float(\"inf\")\n                        nm = mask_idx[j].sum().item()\n                        k = min(int(num_xfer[j, step].item()), nm)\n                        if k > 0:\n                            _, tk = torch.topk(s, k)\n                            xfer[j, tk] = True\n                x = torch.where(xfer, x0, x)\n                used += 1\n        return x, used\n"
7}

1{
2  "op": "replace",
3  "filename": "LLaDA/custom_demask_eval.py",
4  "start_line": 59,
5  "end_line": 148,
6  "content": "class DemaskDecoder:\n    \"\"\"Adaptive Composite Demasking with KL Stability (ACD-KL).\n\n    Novel algorithmic contributions over KLASS:\n    1. Margin-Confidence Fusion Score: fallback position ranking uses\n       score = conf * margin (product of top-1 prob and p1-p2 gap),\n       rewarding both prediction strength and decisiveness.  This avoids\n       committing to positions where the model is confident but\n       indecisive (small gap between top-1 and top-2 candidates).\n    2. Top-k Tempered Sampling: for fully-parallel text decoding,\n       samples from top-k (k=10) with temperature > 1 for controlled\n       diversity, improving MAUVE/entropy while bounding perplexity.\n    \"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n        self.sample_k = 10\n        self.sample_temp = 1.5  # >1 for flatter distribution (diversity)\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n        is_parallel = (block_length == gen_length)\n        V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                       else model.config.vocab_size\n        kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                              dtype=torch.float64, device=x.device)\n        p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                             device=x.device)\n        used = 0\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n            num_xfer = get_num_transfer_tokens(\n                (x[:, bs:be] == mid), steps_per_block)\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                block_m = torch.zeros_like(mask_idx)\n                block_m[:, bs:be] = True\n                mask_idx = mask_idx & block_m\n                if not mask_idx.any():\n                    break\n                logits = model(x).logits\n                p_curr = F.softmax(logits.to(torch.float64), dim=-1)\n                x0g = torch.argmax(p_curr, dim=-1)\n                conf = torch.gather(\n                    p_curr, -1, x0g.unsqueeze(-1)).squeeze(-1)\n                # Novel: margin-confidence fusion score\n                t2v, _ = torch.topk(p_curr, 2, dim=-1)\n                mg = t2v[..., 0] - t2v[..., 1]\n                score = conf * mg\n                # KL stability tracking (KLASS)\n                eps = 1e-12\n                kl = (p_curr * (torch.log(p_curr + eps)\n                               - torch.log(p_prev + eps))).sum(-1)\n                kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                kl_hist[..., -1] = kl\n                p_prev = p_curr.clone()\n                # Novel: top-k tempered sampling for text diversity\n                if is_parallel:\n                    sv, si = torch.topk(\n                        logits.float(), self.sample_k, dim=-1)\n                    sp = F.softmax(sv / self.sample_temp, dim=-1)\n                    sh = sp.shape[:-1]\n                    idx = torch.multinomial(\n                        sp.reshape(-1, self.sample_k), 1)\n                    x0 = torch.gather(\n                        si, -1, idx.reshape(*sh, 1)).squeeze(-1)\n                else:\n                    x0 = x0g\n                # KLASS stability criterion\n                if step >= self.history_length - 1:\n                    stable = torch.all(\n                        kl_hist < self.kl_threshold, dim=-1)\n                else:\n                    stable = torch.zeros_like(conf, dtype=torch.bool)\n                ready = stable & (conf > self.conf_threshold) & mask_idx\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                for j in range(ready.shape[0]):\n                    rdy = torch.where(ready[j])[0]\n                    if len(rdy) > 0:\n                        xfer[j, rdy] = True\n                    else:\n                        s = score[j].clone()\n                        s[~mask_idx[j]] = -float(\"inf\")\n                        nm = mask_idx[j].sum().item()\n                        k = min(int(num_xfer[j, step].item()), nm)\n                        if k > 0:\n                            _, tk = torch.topk(s, k)\n                            xfer[j, tk] = True\n                x = torch.where(xfer, x0, x)\n                used += 1\n        return x, used\n"
7}

1{
2  "op": "replace",
3  "filename": "LLaDA/custom_demask_eval.py",
4  "start_line": 59,
5  "end_line": 165,
6  "content": "class DemaskDecoder:\n    \"\"\"Adaptive Composite Demasking with Mode-Adaptive Stability (ACD-MAS).\n\n    Novel algorithmic contributions:\n    1. Margin-Confidence Fusion Score: fallback position ranking uses\n       score = conf * (p1-p2), rewarding both strength and decisiveness.\n       Avoids committing to positions where top-1 and top-2 are close.\n    2. Mode-Adaptive Stability: semi-AR accuracy tasks use KL-divergence\n       stability (conservative, proven by KLASS). Parallel text tasks\n       use lightweight argmax consistency (fewer steps, similar quality).\n    \"\"\"\n\n    def __init__(self, mask_id: int, temperature: float = 0.0,\n                 conf_threshold: float = 0.9, kl_threshold: float = 0.01,\n                 history_length: int = 2):\n        self.mask_id = mask_id\n        self.temperature = temperature\n        self.conf_threshold = conf_threshold\n        self.kl_threshold = kl_threshold\n        self.history_length = history_length\n\n    @torch.no_grad()\n    def decode(self, model, input_ids, gen_length: int, steps: int,\n               block_length: int):\n        mid = self.mask_id\n        x = torch.full((1, input_ids.shape[1] + gen_length), mid,\n                       dtype=torch.long, device=model.device)\n        x[:, :input_ids.shape[1]] = input_ids.clone()\n        assert gen_length % block_length == 0\n        num_blocks = gen_length // block_length\n        assert steps % num_blocks == 0\n        steps_per_block = steps // num_blocks\n        is_par = (block_length == gen_length)\n        if is_par:\n            prev_pred = torch.full_like(x, -1)\n            cons = torch.zeros_like(x, dtype=torch.int)\n        else:\n            V = model.lm_head.out_features if hasattr(model, \"lm_head\") \\\n                                           else model.config.vocab_size\n            kl_hist = torch.zeros((1, x.shape[1], self.history_length),\n                                  dtype=torch.float64, device=x.device)\n            p_prev = torch.zeros((1, x.shape[1], V), dtype=torch.float64,\n                                 device=x.device)\n        used = 0\n        for b in range(num_blocks):\n            bs = input_ids.shape[1] + b * block_length\n            be = bs + block_length\n            num_xfer = get_num_transfer_tokens(\n                (x[:, bs:be] == mid), steps_per_block)\n            if is_par:\n                cons[:, bs:be] = 0\n                prev_pred[:, bs:be] = -1\n            for step in range(steps_per_block):\n                mask_idx = (x == mid)\n                blk = torch.zeros_like(mask_idx)\n                blk[:, bs:be] = True\n                mask_idx = mask_idx & blk\n                if not mask_idx.any():\n                    break\n                logits = model(x).logits\n                p = F.softmax(logits.to(torch.float64), dim=-1)\n                x0 = torch.argmax(p, dim=-1)\n                cf = torch.gather(p, -1, x0.unsqueeze(-1)).squeeze(-1)\n                t2v, _ = torch.topk(p, 2, dim=-1)\n                mg = t2v[..., 0] - t2v[..., 1]\n                score = cf * mg\n                if is_par:\n                    same = (x0 == prev_pred)\n                    cons = torch.where(\n                        same, cons + 1, torch.ones_like(cons))\n                    prev_pred = x0.clone()\n                    stable = cons >= self.history_length\n                else:\n                    eps = 1e-12\n                    kl = (p * (torch.log(p + eps)\n                              - torch.log(p_prev + eps))).sum(-1)\n                    kl_hist = torch.roll(kl_hist, -1, dims=-1)\n                    kl_hist[..., -1] = kl\n                    p_prev = p.clone()\n                    if step >= self.history_length - 1:\n                        stable = torch.all(\n                            kl_hist < self.kl_threshold, dim=-1)\n                    else:\n                        stable = torch.zeros_like(\n                            cf, dtype=torch.bool)\n                ready = stable & (cf > self.conf_threshold) & mask_idx\n                xfer = torch.zeros_like(x0, dtype=torch.bool)\n                for j in range(ready.shape[0]):\n                    rdy = torch.where(ready[j])[0]\n                    if len(rdy) > 0:\n                        xfer[j, rdy] = True\n                    else:\n                        s = score[j].clone()\n                        s[~mask_idx[j]] = -float(\"inf\")\n                        nm = mask_idx[j].sum().item()\n                        k = min(int(num_xfer[j, step].item()), nm)\n                        if k > 0:\n                            _, tk = torch.topk(s, k)\n                            xfer[j, tk] = True\n                x = torch.where(xfer, x0, x)\n                used += 1\n        return x, used\n"
7}

1{
2  "n": 3
3}