Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction

[remyx-ai]

Drafted by an autonomous discovery loop — Remyx ranks recent arXiv papers against this team's research interest and shipping history; Claude Code selects the candidate most directly implementable against this repo from the lookback window and drafts it.

Recommended paper: Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction
Confidence: 🟢 high (Remyx relevance 0.93)
Research interest: diffusers
Implementation by: Claude Code as autonomous agent

---

Why this paper for this team

While the iteration thread focuses on AnyFlow for video, diffusers' 'Modular Pipelines' (Stage 2) and general model integration (Stage 1) imply a broad interest in controllable generation across modalities. This paper, GILC, offers a plug-and-play, training-free framework for controllable generation with discrete diffusion models, directly applicable to models like LLaDA2. It provides an efficient way to estimate guidance signals for both differentiable and non-differentiable reward functions. This is a natural next step that aligns with diffusers' 'Training and Fine-tuning' research interest, enabling users to exert fine-grained control over discrete outputs without expensive retraining or complex setups, enhancing the utility of discrete models within the library.

Why this candidate (selected from the lookback pool)

The repo ships a LLaDA2 discrete-diffusion pipeline whose refinement loop produces per-step token logits and commits confident tokens — exactly the substrate GILC operates on. GILC is a training-free, plug-and-play logit-correction step (logits in → reward-guided logits out, same shape) that slots between logit production (L466) and scheduler.step (L468) with no contract change to the model or scheduler; the existing call site stays and the new guidance module is wired in. No existing guidance lives in the discrete path (the guiders/ module only serves continuous-diffusion CFG variants), so this is a net-new module anchored on a verified existing forward path rather than a speculative add-alongside.

License & code availability

🟡 No code repository surfaced — couldn't fetch a LICENSE to evaluate. Worth confirming the paper has an open release before investing in adoption.

• Code / model: no repository or model URL surfaced in the paper, recommendation envelope, or arxiv abstract page.
• License: (none detected) (class: no-code-link, compat: 0.30)

Suggested experiment

Integrate GILC as a plug-and-play guidance mechanism for the LLaDA2 Discrete Diffusion Pipeline in diffusers. Implement its Jacobian-free logit correction for a specific controlled generation task (e.g., generating sequences with desired length or specific syntactic properties). Compare its performance against unguided generation or simpler inference-time guidance methods, verifying its training-free nature.

---

What this PR delivers

Call site: LLaDA2Pipeline.__call__ in src/diffusers/pipelines/llada2/pipeline_llada2.py — the new logit_guidanceis applied toblock_logitsinside the refinement loop, beforeself.scheduler.step.

Delivers (from the paper):

• A training-free, plug-and-play reward-guidance step (GILC) for the LLaDA2 discrete-diffusion pipeline: corrects clean-prediction logits in place (logits in -> reward-guided logits out, same shape) between logit production and scheduler.step, with no contract change to model or scheduler.
• RewardLogitGuidance implementing the Jacobian-free correction in two modes: gradient (one ascent step on expected reward, l_v + g*p_v*(r_v - E_p[r])) and tilt (exponential reward-tilt, l_v + g*(r_v - E_p[r])), accepting differentiable OR non-differentiable per-candidate rewards.
• Wired into the existing LLaDA2Pipeline.__call__ via a new logit_guidance kwarg (accepts a RewardLogitGuidance or any [B,L,V]->[B,L,V] callable; None disables) — so unguided generation is unchanged and guidance is opt-in at the pipeline's public entry point.
• from_token_rewards convenience constructor for sparse token_id -> reward control (e.g. boosting EOS for shorter outputs or penalizing disallowed tokens), the concrete length/syntax-control use case the spec suggested.
• active_only masking so the correction touches only still-masked positions, leaving already-committed context untouched.

Intentionally out of scope (not needed for this contribution):

• GILC's multi-step variational proxy that re-runs the denoiser to score the fully-denoised sequence (would require extra denoiser forward passes / a variational estimator); the single-step in-place correction is delivered instead.
• The paper's domain reward models for DNA / protein / molecule generation (require external trained scorers and datasets) — the module is reward-agnostic and lets the caller supply any reward tensor/callable.
• Any quantitative benchmark vs. fine-tuning or unguided baselines (the spec's 'suggested experiment' — needs eval harness/datasets), left to the user.
• Schedulers/models other than LLaDA2's block-refinement loop — the wiring is anchored on that one verified forward path.

This delivers GILC's core in-pipeline value: a training-free, plug-and-play logit-correction step that lets a user steer LLaDA2 discrete-diffusion generation toward a per-token reward (e.g. length/syntax control) with no retraining, via a Jacobian-free correction that supports non-differentiable rewards. It is reachable from the product: the pre-existing LLaDA2Pipeline.__call__ now invokes it through a new optional logit_guidance kwarg, and behavior is identical to before when guidance is None. Intentionally scoped out as unnecessary for that result are the paper's multi-step variational proxy (extra denoiser passes), its domain-specific reward models (DNA/protein/molecule scorers), and a quantitative eval harness — all of which are deliberate boundaries, with the reusable correction core kept reward-agnostic so those can be layered on later.

Test results

ℹ️ Tests could not run in CI — the runner lacks this repo's dependencies (a collection/import error, not a code failure). Run the suite locally to validate.

ages have valid Python names.
Traceback:
/opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
examples/controlnet/test_controlnet.py:23: in <module>
    from test_examples_utils import ExamplesTestsAccelerate, run_command  # noqa: E402
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
examples/test_examples_utils.py:23: in <module>
    from accelerate.utils import write_basic_config
E   ModuleNotFoundError: No module named 'accelerate'
=========================== short test summary info ============================
ERROR examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py
ERROR examples/consistency_distillation/test_lcm_lora.py
ERROR examples/controlnet/test_controlnet.py
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 3 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
3 errors in 0.91s

---

Want eval-on-every-PR? Outrider Validate (coming soon, paid tier) runs your benchmark suite against this diff and posts the results as a PR comment. Design partner pilot is open — join the waitlist.

Opened by the Remyx Recommendation orchestrator.