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

src/diffusers/pipelines/llada2/__init__.py

@@ -22,6 +22,7 @@
 
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
+    _import_structure["logit_guidance"] = ["RewardLogitGuidance"]
     _import_structure["pipeline_llada2"] = ["LLaDA2Pipeline", "LLaDA2PipelineOutput"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
@@ -32,6 +33,7 @@
     except OptionalDependencyNotAvailable:
         from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
+        from .logit_guidance import RewardLogitGuidance
         from .pipeline_llada2 import LLaDA2Pipeline, LLaDA2PipelineOutput
 else:
     import sys

src/diffusers/pipelines/llada2/logit_guidance.py

@@ -0,0 +1,158 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Training-free, plug-and-play reward guidance for discrete-diffusion logits.
+
+This implements a focused slice of *Gradient-Informed Logit Correction* (GILC),
+"Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit
+Correction" (https://arxiv.org/abs/2606.06303). GILC steers a discrete-diffusion
+sampler toward a reward without any retraining by correcting the clean-prediction
+logits in place: logits in -> reward-guided logits out, same shape.
+
+The correction is *Jacobian-free*: the reward function is evaluated on the
+candidate-token axis directly (so it may be non-differentiable), and the guidance
+signal is the analytic gradient of the expected reward with respect to the logits
+under the model's own clean-prediction distribution. No gradient is taken through
+the denoising network, which is what makes the step stable in the high-dimensional
+discrete space.
+
+For clean-prediction logits ``l`` with ``p = softmax(l)`` and a per-candidate
+reward ``r`` (a vector over the vocabulary, broadcast over batch/position):
+
+- ``mode="gradient"`` (default) takes one gradient-ascent step on the expected
+  reward ``E_p[r] = sum_v p_v r_v``. Its gradient w.r.t. logit ``l_v`` is
+  ``p_v (r_v - E_p[r])``, so the corrected logit is
+  ``l_v + guidance_scale * p_v * (r_v - E_p[r])``.
+- ``mode="tilt"`` applies the exponential reward-tilt ``p'(v) ∝ p(v) exp(g r_v)``,
+  i.e. ``l_v + guidance_scale * (r_v - E_p[r])`` (centering only re-scales, since
+  softmax is shift-invariant).
+
+Both forms are training-free and accept differentiable or non-differentiable
+rewards. Out of scope here (and not needed for the in-pipeline value) is GILC's
+multi-step variational proxy that re-runs the denoiser to score the *fully*
+denoised sequence, plus the paper's domain reward models (DNA / protein /
+molecule). The in-place logit correction is the reusable core.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Mapping
+
+import torch
+
+
+class RewardLogitGuidance:
+    """Reward-guided logit correction for the LLaDA2 block-refinement loop.
+
+    An instance is callable as ``guidance(logits) -> corrected_logits`` and is
+    meant to be applied to the clean-prediction logits of a refinement step, just
+    before they are handed to the scheduler. The output keeps the input shape and
+    dtype, so it is a drop-in correction that changes no other contract.
+
+    Args:
+        reward (`torch.Tensor` or `Callable`):
+            Per-candidate reward. Either a tensor broadcastable to the logits
+            (e.g. shape `[vocab]`, `[1, 1, vocab]`, or the full `[batch, seq,
+            vocab]`), or a callable mapping the logits to such a tensor. Larger
+            reward favors the corresponding token. May be non-differentiable.
+        guidance_scale (`float`, defaults to `1.0`):
+            Strength of the correction. `0.0` is a no-op (unguided generation).
+        mode (`str`, defaults to `"gradient"`):
+            `"gradient"` for the GILC gradient-of-expected-reward step, or
+            `"tilt"` for the exponential reward-tilt.
+        active_only (`bool`, defaults to `True`):
+            When the still-masked token positions are known (passed as `tokens` /
+            `mask_token_id`), restrict the correction to those positions so
+            already-committed context is left untouched.
+    """
+
+    def __init__(
+        self,
+        reward: "torch.Tensor | Callable[[torch.Tensor], torch.Tensor]",
+        guidance_scale: float = 1.0,
+        mode: str = "gradient",
+        active_only: bool = True,
+    ):
+        if mode not in {"gradient", "tilt"}:
+            raise ValueError(f"`mode` must be 'gradient' or 'tilt', got {mode!r}.")
+        if not callable(reward) and not torch.is_tensor(reward):
+            raise TypeError("`reward` must be a torch.Tensor or a callable returning one.")
+        self.reward = reward
+        self.guidance_scale = float(guidance_scale)
+        self.mode = mode
+        self.active_only = active_only
+
+    @classmethod
+    def from_token_rewards(
+        cls,
+        token_rewards: "Mapping[int, float] | torch.Tensor",
+        vocab_size: int,
+        default: float = 0.0,
+        **kwargs,
+    ) -> "RewardLogitGuidance":
+        """Build guidance from a sparse map of `token_id -> reward`.
+
+        Convenient for length / syntax control, e.g. boosting an end-of-sequence
+        token to favor shorter outputs, or penalizing a set of disallowed tokens.
+        """
+        if torch.is_tensor(token_rewards):
+            reward = token_rewards.to(dtype=torch.float32).reshape(-1)
+            if reward.numel() != vocab_size:
+                raise ValueError(
+                    f"`token_rewards` has {reward.numel()} entries but `vocab_size` is {vocab_size}."
+                )
+        else:
+            reward = torch.full((vocab_size,), float(default), dtype=torch.float32)
+            for token_id, value in token_rewards.items():
+                if not 0 <= int(token_id) < vocab_size:
+                    raise ValueError(f"token id {token_id} out of range for vocab_size={vocab_size}.")
+                reward[int(token_id)] = float(value)
+        return cls(reward, **kwargs)
+
+    def _reward_tensor(self, logits: torch.Tensor) -> torch.Tensor:
+        reward = self.reward(logits) if callable(self.reward) else self.reward
+        if not torch.is_tensor(reward):
+            raise TypeError("Reward callable must return a torch.Tensor.")
+        reward = reward.to(device=logits.device, dtype=torch.float32)
+        # Broadcast-check against the logits shape; let torch raise on a real mismatch.
+        return reward.expand_as(logits) if reward.shape != logits.shape else reward
+
+    def __call__(
+        self,
+        logits: torch.Tensor,
+        *,
+        tokens: torch.Tensor | None = None,
+        mask_token_id: int | None = None,
+    ) -> torch.Tensor:
+        """Return reward-corrected logits with the same shape and dtype as `logits`."""
+        if self.guidance_scale == 0.0:
+            return logits
+        if logits.ndim != 3:
+            raise ValueError(f"`logits` must be `[batch, seq, vocab]`, got shape {tuple(logits.shape)}.")
+
+        reward = self._reward_tensor(logits)
+        probs = torch.softmax(logits.float(), dim=-1)
+        expected = (probs * reward).sum(dim=-1, keepdim=True)
+        centered = reward - expected
+
+        if self.mode == "gradient":
+            correction = self.guidance_scale * probs * centered
+        else:  # "tilt"
+            correction = self.guidance_scale * centered
+
+        if self.active_only and tokens is not None and mask_token_id is not None:
+            active = (tokens == mask_token_id).unsqueeze(-1)
+            correction = torch.where(active, correction, torch.zeros_like(correction))
+
+        return logits + correction.to(logits.dtype)

src/diffusers/pipelines/llada2/pipeline_llada2.py

@@ -24,6 +24,7 @@
 from ...schedulers import BlockRefinementScheduler
 from ...utils import BaseOutput, logging, replace_example_docstring
 from ..pipeline_utils import DiffusionPipeline
+from .logit_guidance import RewardLogitGuidance
 
 
 logger = logging.get_logger(__name__)
@@ -264,6 +265,7 @@ def __call__(
         eos_token_id: int | None = None,
         mask_token_id: int | None = None,
         generator: torch.Generator | None = None,
+        logit_guidance: RewardLogitGuidance | Callable[[torch.Tensor], torch.Tensor] | None = None,
         output_type: str = "text",
         return_dict: bool = True,
         callback_on_step_end: Callable[[int, int, dict], None]
@@ -327,6 +329,12 @@ def __call__(
                 Mask token ID to use for the template.
             generator (`torch.Generator`, *optional*):
                 RNG for sampling.
+            logit_guidance ([`RewardLogitGuidance`] or `Callable`, *optional*):
+                Training-free, plug-and-play reward guidance applied to each step's clean-prediction logits
+                before the scheduler commits tokens (a focused implementation of GILC, arXiv:2606.06303). A
+                [`RewardLogitGuidance`] steers generation toward a per-token reward via a Jacobian-free logit
+                correction; any callable mapping block logits `[batch, block_length, vocab]` to corrected logits
+                of the same shape is also accepted. `None` disables guidance.
             output_type (`str`, defaults to `"text"`):
                 Output format. `"text"` decodes sequences into strings (requires a tokenizer). `"seq"` returns raw
                 token ID sequences only.
@@ -465,6 +473,16 @@ def __call__(
                 logits = self.model(block_x, attention_mask=block_attn_mask, position_ids=block_position_ids).logits
                 block_logits = logits[:, -block_length:, :]
 
+                # Training-free plug-and-play reward guidance (GILC): correct the clean-prediction
+                # logits in place before the scheduler commits tokens. Same shape in, same shape out.
+                if logit_guidance is not None:
+                    if isinstance(logit_guidance, RewardLogitGuidance):
+                        block_logits = logit_guidance(
+                            block_logits, tokens=block_tokens, mask_token_id=mask_token_id
+                        )
+                    else:
+                        block_logits = logit_guidance(block_logits)
+
                 scheduler_output = self.scheduler.step(
                     model_output=block_logits,
                     timestep=step_idx,

tests/pipelines/llada2/test_logit_guidance.py

@@ -0,0 +1,171 @@
+# Copyright 2025 HuggingFace Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for training-free reward guidance (GILC) wired into the LLaDA2 pipeline.
+
+The integration tests drive the real `LLaDA2Pipeline.__call__` (a non-new module)
+so the wiring edit at the refinement-loop call site is actually exercised.
+"""
+
+import inspect
+import unittest
+
+import torch
+
+from diffusers import BlockRefinementScheduler, LLaDA2Pipeline
+from diffusers.pipelines.llada2 import RewardLogitGuidance
+
+
+class _DummyModelOutput:
+    def __init__(self, logits):
+        self.logits = logits
+
+
+class _DummyCausalLM(torch.nn.Module):
+    """Position-dependent logits so unguided top-k commits are deterministic."""
+
+    def __init__(self, vocab_size: int):
+        super().__init__()
+        self.vocab_size = int(vocab_size)
+        self.register_buffer("_device_anchor", torch.empty(0))
+
+    @property
+    def dtype(self):
+        return torch.float32
+
+    @property
+    def device(self):
+        return self._device_anchor.device
+
+    def forward(self, input_ids, attention_mask=None, position_ids=None, **kwargs):
+        batch_size, seq_len = input_ids.shape
+        logits = torch.zeros((batch_size, seq_len, self.vocab_size), device=input_ids.device, dtype=torch.float32)
+        positions = torch.arange(seq_len, device=input_ids.device, dtype=torch.float32).view(1, seq_len, 1)
+        token_ids = (torch.arange(seq_len, device=input_ids.device) % (self.vocab_size - 2)).view(1, seq_len, 1)
+        logits.scatter_(2, token_ids.expand(batch_size, -1, -1), 1.0 + positions.expand(batch_size, -1, -1) * 0.1)
+        return _DummyModelOutput(logits=logits)
+
+
+def _make_pipeline():
+    return LLaDA2Pipeline(model=_DummyCausalLM(vocab_size=32), scheduler=BlockRefinementScheduler())
+
+
+_GEN_KWARGS = {
+    "use_chat_template": False,
+    "gen_length": 24,
+    "block_length": 8,
+    "num_inference_steps": 8,
+    "temperature": 0.0,
+    "threshold": 2.0,  # force top-k commits
+    "minimal_topk": 1,
+    "editing_threshold": 0.0,  # disable editing for a deterministic comparison
+    "eos_early_stop": False,
+    "mask_token_id": 31,
+    "eos_token_id": None,
+    "output_type": "seq",
+}
+
+
+class RewardLogitGuidanceUnitTest(unittest.TestCase):
+    def test_shape_and_dtype_preserved(self):
+        logits = torch.randn(2, 5, 7)
+        guidance = RewardLogitGuidance(torch.zeros(7), guidance_scale=1.0)
+        out = guidance(logits)
+        self.assertEqual(out.shape, logits.shape)
+        self.assertEqual(out.dtype, logits.dtype)
+
+    def test_zero_scale_is_noop(self):
+        logits = torch.randn(1, 3, 6)
+        guidance = RewardLogitGuidance(torch.arange(6, dtype=torch.float32), guidance_scale=0.0)
+        self.assertTrue(torch.equal(guidance(logits), logits))
+
+    def test_gradient_correction_raises_favored_logit(self):
+        # Uniform logits -> the correction direction is exactly the centered reward.
+        logits = torch.zeros(1, 1, 4)
+        reward = torch.tensor([0.0, 0.0, 10.0, 0.0])
+        out = RewardLogitGuidance(reward, guidance_scale=1.0, mode="gradient")(logits)
+        delta = (out - logits).reshape(-1)
+        self.assertEqual(int(delta.argmax()), 2)
+        self.assertGreater(delta[2].item(), 0.0)
+
+    def test_tilt_matches_logprob_shift(self):
+        # In tilt mode the distribution equals softmax(logits + scale * reward).
+        logits = torch.randn(1, 2, 5)
+        reward = torch.randn(5)
+        out = RewardLogitGuidance(reward, guidance_scale=0.8, mode="tilt")(logits)
+        got = torch.softmax(out, dim=-1)
+        expected = torch.softmax(logits + 0.8 * reward, dim=-1)
+        self.assertTrue(torch.allclose(got, expected, atol=1e-5))
+
+    def test_active_only_skips_committed_positions(self):
+        logits = torch.zeros(1, 2, 4)
+        tokens = torch.tensor([[31, 5]])  # position 0 masked, position 1 committed
+        guidance = RewardLogitGuidance(torch.tensor([0.0, 9.0, 0.0, 0.0]), guidance_scale=1.0)
+        out = guidance(logits, tokens=tokens, mask_token_id=31)
+        self.assertTrue((out[0, 0] != logits[0, 0]).any())  # masked position corrected
+        self.assertTrue(torch.equal(out[0, 1], logits[0, 1]))  # committed position untouched
+
+    def test_from_token_rewards(self):
+        guidance = RewardLogitGuidance.from_token_rewards({3: 5.0, 1: -2.0}, vocab_size=6)
+        reward = guidance.reward
+        self.assertEqual(reward.shape, (6,))
+        self.assertEqual(reward[3].item(), 5.0)
+        self.assertEqual(reward[1].item(), -2.0)
+        self.assertEqual(reward[0].item(), 0.0)
+
+
+class RewardLogitGuidancePipelineTest(unittest.TestCase):
+    def test_call_signature_exposes_logit_guidance(self):
+        # The wiring edit must surface the new parameter on the public pipeline.
+        params = inspect.signature(LLaDA2Pipeline.__call__).parameters
+        self.assertIn("logit_guidance", params)
+
+    def test_guidance_steers_committed_tokens(self):
+        target_token = 10
+        input_ids = torch.tensor([[5, 6, 7, 8], [1, 2, 3, 4]], dtype=torch.long)
+
+        unguided = _make_pipeline().to("cpu")(input_ids=input_ids, **_GEN_KWARGS).sequences
+
+        # `tilt` mode is the direct log-prob shift, which reliably overrides the model's prior.
+        guidance = RewardLogitGuidance.from_token_rewards(
+            {target_token: 100.0}, vocab_size=32, guidance_scale=1.0, mode="tilt"
+        )
+        guided = _make_pipeline().to("cpu")(input_ids=input_ids, logit_guidance=guidance, **_GEN_KWARGS).sequences
+
+        guided_hits = int((guided == target_token).sum())
+        unguided_hits = int((unguided == target_token).sum())
+
+        # Guidance should dominate the committed tokens and clearly beat the unguided baseline.
+        self.assertGreater(guided_hits, unguided_hits)
+        self.assertGreaterEqual(guided_hits, int(0.9 * guided.numel()))
+
+    def test_callable_guidance_is_invoked(self):
+        # A plain callable (logits -> logits) is also accepted at the call site.
+        calls = {"n": 0}
+
+        def reward_fn(block_logits):
+            calls["n"] += 1
+            bias = torch.zeros(block_logits.shape[-1])
+            bias[3] = 50.0
+            return block_logits + bias
+
+        input_ids = torch.tensor([[5, 6, 7, 8]], dtype=torch.long)
+        out = _make_pipeline().to("cpu")(input_ids=input_ids, logit_guidance=reward_fn, **_GEN_KWARGS).sequences
+
+        self.assertGreater(calls["n"], 0)
+        self.assertEqual(out.shape, (1, 24))
+
+
+if __name__ == "__main__":
+    unittest.main()