PaperlayerView exact revision on arXiv

Source-linked AI summary

DELA: Learning Controllable Edits for Diffusion Language Models

Xinsong Feng, Chenan Wang, Leshu Tang, Johnny Liu, Dong Li, Da Li, Haipeng Chen

arXiv:2610.32096v1cs.AI

TL;DR

Token-level denoising schedules leave draft-level refinement implicit in diffusion language models. DELA adds a controller for explicit edits while keeping the denoising backbone frozen, and improves mathematical reasoning and code generation by up to 5.5 percentage points over direct diffusion decoding.

  • Problem

    Diffusion language models largely control generation through token-level denoising schedules, leaving explicit editing of clean drafts undercontrolled.

    Source: §1 · pp. 1–2

  • Method

    DELA trains an external controller to select draft edits while a frozen DLM denoises masked content, using staged training and rollout-informed candidate selection.

    Source: §1 · p. 2, §3.2 · p. 6

  • Results

    DELA improves mathematical reasoning and code generation by up to 5.5 percentage points over direct diffusion decoding while keeping the backbone fixed.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DELA offers a plug-in approach to explicit draft-level refinement without modifying the pretrained denoising backbone.

    Source: §1 · p. 2

    The inference margin gate limits deviations from the default editor action but does not itself favor STOP over an edit.

    Source: §3.2 · p. 7

Abstract

from arXiv · show

Diffusion language models (DLMs) support bidirectional denoising and iterative refinement, but generation remains largely controlled by token-level masking schedules. This leaves draft-level refinement implicit and limits explicit control over when to insert missing content, delete redundant spans, revisit uncertain regions, or stop editing. We introduce Diffusion Editing Language Adapter (DELA), a plug-in framework that enables controllable draft refinement with a frozen DLM. DELA operates on clean, variable-length drafts using four explicit editing actions: Insert, Delete, Remask, and Stop. The controller determines when and where to edit the draft, while the frozen DLM fills masked regions through denoising, separating high-level edit control from token-level generation. A naive single-phase editor would be computationally expensive because arbitrary interior edits can disrupt prefix stability and limit reliable KV-cache reuse. To address this, DELA uses a two-phase decoding procedure: a fast, append-only drafting phase that supports KV-cache reuse, followed by a global refinement phase that enables targeted non-monotonic edits to the full draft. We train the controller using a staged curriculum that begins with learning local edit priors from synthetic recovery tasks and mined actions that improve draft quality. Subsequent stages refine edit selection through rollout-based outcome preferences and alignment with downstream task feedback. Experiments on mathematical reasoning and code generation show that DELA achieves gains of up to 5.5 percentage points over direct diffusion decoding while keeping the backbone frozen. DELA thus provides a modular approach to explicit draft-level control in diffusion language models.

1 INTRODUCTION

DELA adds explicit draft-level editing to frozen diffusion language models, addressing structural errors that token-level denoising schedules do not control. Its controller separates edit decisions from token generation, using a staged curriculum and two-phase decoding to support refinement.

1 INTRODUCTION

  • Token-level denoising schedules leave draft-level refinement implicit, despite structural errors such as omitted reasoning steps, redundant content, and inconsistencies.Source: §1 · p. 1
  • Prior variable-length approaches incorporate structural edits into generative dynamics, often requiring changes to model architecture, training objectives, or backbone parameters.Source: §1 · p. 2
  • A lightweight controller applies INSERT, DELETE, REMASK, and STOP to clean drafts while a frozen DLM denoises masked content.Source: §1 · p. 2
  • The controller curriculum progresses from synthetic edit recovery and mined improving actions to rollout preferences and downstream task feedback.Source: §1 · p. 2
  • DELA uses append-only, cache-friendly drafting before global refinement, avoiding non-monotonic edits throughout generation that disrupt prefix stability and KV-cache reuse.Source: §1 · p. 2
  • DELA improves mathematical reasoning and code generation by up to 5.5 percentage points over direct diffusion decoding while keeping the backbone fixed.Source: §1 · p. 2

2 RELATED WORK

Related work establishes masked diffusion as DELA’s token-level generation foundation and explores scaling, efficient decoding, and variable-length editing. DELA instead places structural control in an external editor over a frozen denoiser.

Diffusion language models

  • Masked DLMs corrupt tokens into [MASK] and predict clean tokens at masked positions; DELA uses a pretrained masked model as a frozen denoiser.Source: §2 · p. 2

Scaling and efficient generation with DLMs

  • DLM efficiency work includes blockwise generation and KV-cache reuse, approximate caching, adaptive token ordering, early stopping, and inference-time remasking.Source: §2 · p. 3

Variable-length and edit-based generation

  • Variable-length and edit-based models incorporate operations such as insertion, deletion, or expansion into diffusion or flow dynamics, whereas DELA separates edit control from denoising.Source: §2 · p. 3

3 DELA: LEARNING CONTROLLABLE EDITS

DELA makes clean-draft refinement an explicit control problem: an editor selects structural actions, and a frozen DLM realizes masked content. Its two-phase procedure combines cache-friendly drafting with global editing, while a staged curriculum trains the controller to propose and select useful edits.

3.1 EDIT-CONDITIONED REFINEMENT

  • The editor selects among INSERT, DELETE, REMASK, and STOP; a frozen DLM denoises masked drafts, while STOP returns the current sequence without another denoising transition.Source: §3 · p. 3, §3.1 · p. 4
  • Fast drafting appends masked blocks and locally remasks uncertain tokens, preserving finalized prefix KV states for reuse as later blocks are added.Source: §3.1 · p. 4
  • Global refinement enables targeted insertion, deletion, remasking, or stopping after the initial draft, reserving non-monotonic edits for that phase.Source: §3.1 · p. 4

3.2 LEARNING THE EDIT CONTROLLER

  • Direct reward-based controller learning is difficult because actions combine operation choice, localization, remasking, and stopping, with effects visible only after denoising and possibly multiple steps.Source: §3.2 · p. 4
  • Stage 1 uses synthetic corruptions to teach inverse INSERT or DELETE edits, REMASK localization, and STOP on uncorrupted sequences, then adds mined actions that improve drafts under the frozen denoiser.Source: §3.2 · pp. 5–6
  • Stage 2 evaluates executable candidates with short-horizon rollouts and trains both the selector and editor from outcome preferences; Stage 3 replaces general-text scores with downstream task rewards.Source: §3.2 · pp. 4, 6, 7
  • At inference, the selector scores candidate actions without target text or rollout evaluation, and a margin gate determines whether to override the default action.Source: §3.2 · p. 7

4 EXPERIMENTS

Experiments evaluate DELA as a controller over frozen diffusion backbones on mathematical reasoning and code generation. DELA improves over direct decoding, while comparisons and analyses examine inference cost, controller selection, editing behavior, and efficiency.

4.1 EXPERIMENTAL SETUP

  • Evaluation uses frozen LLaDA-1.5 and additionally tests Dream-7B; tasks include GSM8K, MATH, HumanEval, and MBPP, with DIRECT, editor, selector, oracle, and REMASK-only variants.The initial draft uses the same frozen backbone and decoding budget across methods; the ORACLE uses the task evaluator only as an upper bound on candidate selection.Source: §4 · p. 7, §4.1 · pp. 7–8

4.2 MAIN RESULTS

  • DELA consistently improves over direct decoding; the further gain from EDITOR+SELECTOR shows that choosing which proposed edits to execute matters.The remaining gap to ORACLE indicates that useful candidates are sometimes available but not selected by the learned selector.Source: §4.2 · p. 8

4.3 COMPARISON WITH INFERENCE-TIME REFINEMENT

  • On HumanEval, DELA-REMASK reaches 40.9% pass@1 with five additional NFEs over DIRECT, while full DELA reaches 42.7% with eight additional NFEs.Tolerator-style refinement and ReMDM both reach 38.4% pass@1 with additional denoising; the comparison does not isolate learned editing from extra compute.Source: §4.3 · p. 8
  • On frozen Dream-7B, DELA improves HumanEval from 41.5% to 43.9% and MBPP from 57.8% to 60.6%.The gains provide initial evidence of applicability across diffusion backbones beyond LLaDA.Source: §4.3 · pp. 8–9

4.4 CONTROLLER AND EDITING ANALYSIS

  • Stage 1 edits every HumanEval example, breaking 43 previously correct drafts while fixing none; Stage 2 recovers DIRECT performance, and Stage 3 aligns selection with task success.The results indicate that useful edit proposals alone are insufficient; the controller must also learn which edits to execute.Source: §4.4 · p. 9
  • On HumanEval, DELA performs 24 REMASK, 2 INSERT, and 1 DELETE actions, editing 15 of 164 outputs (9.1%); edited drafts change a median of 12 tokens.DELA-REMASK reaches 40.9%, compared with 42.7% for the full action space; the analysis characterizes the edits as sparse, targeted repair.Source: §4.4 · p. 9
  • In a batch-size-one profile, two-phase drafting and refinement takes 0.461 s per sample, compared with 0.823 s when global edit control runs after every generated block.Delaying non-monotonic refinement avoids repeatedly invoking the global controller during generation.Source: §4.4 · p. 9

5 CONCLUSION

DELA brings explicit draft-level editing to frozen diffusion language models through an external editor and outcome-guided training. Experiments report gains over direct diffusion decoding without backbone retraining, while evaluation remains limited to two backbones and a focused task set.

5 CONCLUSION

  • DELA learns an external editor to refine clean drafts through structural actions, with a staged curriculum that incorporates outcome-guided signals for downstream task performance.The pretrained diffusion backbone remains fixed.Source: §5 · p. 9
  • The study covers two diffusion backbones and a focused task set; broader evaluation across scales and generation settings remains future work.Source: §5 · p. 9

A.1 WHY SINGLETON-REMASK CONSISTENCY WORKS

The singleton-remask objective is exactly conditional pseudo-log-likelihood, and under stated assumptions its optimum recovers the data joint distribution. The result supports a theoretical interpretation of the practical multi-token remask score, not a claim that it is exact joint likelihood or improves downstream outcomes.

A.1 WHY SINGLETON-REMASK CONSISTENCY WORKS

  • Singleton-remask training exactly maximizes the data distribution’s conditional pseudo-log-likelihood in the fixed-length, one-position-remasked setting.Qgen is a component of the task-agnostic score Fgen; practical scoring uses more general remask patterns.Source: §A.1 · pp. 12, 14
  • The population objective is bounded above by a parameter-independent constant, with equality exactly when the model matches the data full-conditionals.The match holds for every position and q⋆-almost every context and partially observed sequence.Source: §A.1 · p. 12
  • With realizable model full-conditionals compatible with a unique strictly positive joint distribution, an objective maximizer recovers the data joint distribution.Under those assumptions, matching all full-conditionals determines the same joint distribution as joint maximum likelihood.Source: §A.1 · p. 13
  • The result does not establish that practical multi-token or blockwise remasking is exact normalized joint likelihood or that its selected edits improve downstream outcomes.Those variants are conditional-consistency surrogates; the theorem concerns fixed-length singleton remasking.Source: §A.1 · p. 14

B.1 EXPERIMENTAL CONFIGURATION

Experiments keep the LLaDA-1.5 denoising backbone frozen and train only DELA’s edit-control modules. DIRECT shares DELA’s initial-draft pipeline, while DELA adds global-refinement computation.

B.1 EXPERIMENTAL CONFIGURATION

  • DELA keeps LLaDA-1.5 frozen and trains only its proposal editor and selector; the initial draft uses gen_length = 256, steps = 256, and block_length = 16.DIRECT uses the same initial-draft pipeline, including local per-block remasking, and returns the draft before global refinement.Source: §B.1 · p. 14
  • DIRECT uses 256 NFEs for shared drafting, while DELA uses eight additional NFEs for global refinement.NFE counts frozen-backbone calls across drafting and any subsequent refinement.Source: §B.1 · p. 14
  • At final task evaluation, the selector uses learned logits only; Stage 3 disables test-time rollout reranking.Rollout rewards, verifier feedback, target answers, and oracle candidate scores are reserved for training or diagnostic upper bounds.Source: §B.1 · p. 14

B.2 COMPUTE RESOURCES

The append-only drafting phase and global refinement are paired with a staged training curriculum. Experiments used a frozen LLaDA-1.5 backbone on an eight-GPU workstation.

B.2 COMPUTE RESOURCES

  • Main experiments used a workstation with 8 NVIDIA RTX PRO 6000 GPUs, while keeping LLaDA-1.5 frozen and updating only DELA edit-control modules.Training and evaluation used the same machine unless otherwise stated.Source: §B.2 · pp. 14–15

B.3 TRAINING ALGORITHM

  • The training curriculum progresses from synthetic edit imitation and mined improving-action distillation to rollout-guided general improvement and task-specific alignment.The final stage replaces the general score with task reward and trains using task-level preferences.Source: §B.3 · p. 15

B.4 DECODING ALGORITHM

  • Decoding first drafts by append-only block insertion with local remasking, then globally refines using the full edit-action space.During refinement, the selector overrides the editor’s default action only when its score margin exceeds the action-specific gate; STOP ends refinement.Source: §B.4 · pp. 15–16

C.1 STAGE 1: SYNTHETIC RECOVERY

Stage 1 teaches the editor executable recovery actions on synthetic corruptions, then mined-positive distillation shifts proposals toward edits that improve drafts. Stage 2 further improves selection among executable edits, though useful candidates are not always identified.

C.1 STAGE 1: SYNTHETIC RECOVERY

  • Synthetic recovery trains the editor to map structural corruptions to recovery actions, including operation type, edit position, and span length.Deleting a span induces INSERT; adding a distractor span induces DELETE. These are action-grammar diagnostics, not evidence of downstream task improvement.Source: §C.1 · p. 16, Table 6 · p. 17

C.2 STAGE 1: PROPOSAL QUALITY

  • Mined-positive distillation raises improving-proposal coverage from 38.3% to 84.4% and top-action draft improvement from 50.0% to 87.5%.Mined-positive measures whether candidate sets contain a rollout-mined improving action; Top-1 improves measures whether the editor’s highest-probability action improves the completed draft.Source: §C.2 · p. 17

C.3 STAGE 2: SELECTION DIAGNOSTICS

  • Stage 2 improves average reward and recovery rate over directly executing Stage 1 actions, but the gap to the oracle shows useful candidates are not always selected.The selector is evaluated on held-out editing states; downstream results indicate Stage 2 primarily prevents harmful over-editing before task-specific alignment.Source: §C.3 · p. 17

C.4 ADDITIONAL INFERENCE-TIME COMPARISONS

Additional HumanEval comparisons find that alternative block sizes and token ordering do not explain DELA’s gains, while refinement sensitivity is limited in the tested setting. Executed edits are sparse and predominantly local.

C.4 ADDITIONAL INFERENCE-TIME COMPARISONS

  • Block size 32 modestly improves quality, larger blocks eventually degrade it, and random ordering performs poorly; Tolerator-style refinement remains below DELA-REMASK and full DELA.Tolerator-style refinement improves over DIRECT but requires additional denoising evaluations. All methods use the same frozen LLaDA-1.5 backbone.Source: §C.4 · p. 17, Table 9 · p. 18

C.5 INFERENCE-TIME SENSITIVITY

  • Performance is stable across γop ∈[0, 0.5]; two refinement rounds solve 70 problems versus 68 with one, while a third adds no gain.Disabling selector-guided overrides (γop = ∞) reduces performance to the EDITOR-ONLY result of 64 correct solutions.Source: §C.5 · p. 18

C.6 EDIT BEHAVIOR AND DIFFICULTY

  • Only 9.1% of HumanEval drafts receive an edit, and the median edited example changes 12 tokens; REMASK dominates, while INSERT and DELETE are occasional length corrections.The edits are characterized as sparse and local, rather than general-purpose rewriting.Source: §C.6 · p. 18

C.7 TWO-PHASE DECODING AND CACHE PROFILING

Prefix reuse reduces computation during fast drafting, while global editing throughout drafting incurs repeated controller overhead. The profiling comparison uses frozen prefixes, so it measures cache reuse cost rather than the effect of that approximation against full bidirectional recomputation.

C.7 TWO-PHASE DECODING AND CACHE PROFILING

  • The cache comparison holds finalized prefix states fixed, so it measures reuse cost rather than the effect of frozen-prefix approximation against full bidirectional recomputation.Cached and uncached execution are compared under the same frozen-prefix condition.Source: §C.7 · p. 18
  • Prefix reuse cuts processed query positions by 7.4× and attention pairs by 8.5×; wall-clock gains rise from 1.14× at 256 tokens to 1.62× at 1,024 tokens.The profiling covers 256-token fast drafting with batch size one; at 512 tokens, the wall-clock benefit is 1.25×.Source: §C.7 · pp. 18–19, Table 12 · p. 20
  • The two-phase design avoids repeated global-controller overhead while enabling prefix reuse; the reported 1.8× overhead is conservative because the controller selects STOP at block boundaries.The forced-edit configuration illustrates the larger cost of non-monotonic refinement throughout draft construction.Source: §C.7 · p. 19, Table 13 · p. 20

Related papers

Loading 2610.32096v1…