From a18b1d1cd5d5a51c651e1b85ec8f73d0d24faa57 Mon Sep 17 00:00:00 2001 From: "remyx-ai[bot]" <289541483+remyx-ai[bot]@users.noreply.github.com> Date: Tue, 16 Jun 2026 13:28:30 +0000 Subject: [PATCH] Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction --- src/diffusers/pipelines/llada2/__init__.py | 2 + .../pipelines/llada2/logit_guidance.py | 158 ++++++++++++++++ .../pipelines/llada2/pipeline_llada2.py | 18 ++ tests/pipelines/llada2/test_logit_guidance.py | 171 ++++++++++++++++++ 4 files changed, 349 insertions(+) create mode 100644 src/diffusers/pipelines/llada2/logit_guidance.py create mode 100644 tests/pipelines/llada2/test_logit_guidance.py diff --git a/src/diffusers/pipelines/llada2/__init__.py b/src/diffusers/pipelines/llada2/__init__.py index 45a02e6851e2..d1156152f83a 100644 --- a/src/diffusers/pipelines/llada2/__init__.py +++ b/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 diff --git a/src/diffusers/pipelines/llada2/logit_guidance.py b/src/diffusers/pipelines/llada2/logit_guidance.py new file mode 100644 index 000000000000..a937a4804fe3 --- /dev/null +++ b/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) diff --git a/src/diffusers/pipelines/llada2/pipeline_llada2.py b/src/diffusers/pipelines/llada2/pipeline_llada2.py index c9e15e27375c..a88e3c8e872a 100644 --- a/src/diffusers/pipelines/llada2/pipeline_llada2.py +++ b/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, diff --git a/tests/pipelines/llada2/test_logit_guidance.py b/tests/pipelines/llada2/test_logit_guidance.py new file mode 100644 index 000000000000..6e28d61a9208 --- /dev/null +++ b/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()