diffusers/pr_13872/en/api/schedulers/entropy_bound.md

EntropyBoundScheduler

The EntropyBoundScheduler commits the lowest-entropy positions whose joint entropy stays under entropy_bound, so roughly independent tokens are accepted together and the rest are renoised. It anneals its sampling temperature from t_max on the first step down to t_min on the last, matching the released checkpoint's sampler. Proposed in Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking.

This scheduler is used by DiffusionGemmaPipeline.

EntropyBoundScheduler[[diffusers.EntropyBoundScheduler]]

diffusers.EntropyBoundScheduler[[diffusers.EntropyBoundScheduler]]

Source

Entropy bound scheduler for the uniform corruption process.

At each step the scheduler samples a candidate token per position and accepts the k lowest-entropy positions such that `sum_i^k entropy_i - max(entropy_1, ..., entropy_k)

stepdiffusers.EntropyBoundScheduler.stephttps://github.com/huggingface/diffusers/blob/vr_13872/src/diffusers/schedulers/scheduling_entropy_bound.py#L118[{"name": "model_output", "val": ": torch.Tensor"}, {"name": "timestep", "val": ": int | torch.Tensor"}, {"name": "sample", "val": ": torch.LongTensor"}, {"name": "entropy_bound", "val": ": float | None = None"}, {"name": "generator", "val": ": torch.Generator | None = None"}, {"name": "return_dict", "val": ": bool = True"}]- model_output (torch.Tensor of shape (batch_size, block_length, vocab_size)) -- Raw logits from the model for the current block.

• timestep (int or torch.Tensor) -- Current step index within the denoising schedule; sets the annealed sampling temperature.
• sample (torch.LongTensor of shape (batch_size, block_length)) -- Current block token IDs.
• entropy_bound (float, optional) -- Overrides the configured entropy bound for this step.
• generator (torch.Generator, optional) -- RNG for sampling.
• return_dict (bool) -- Whether to return an EntropyBoundSchedulerOutput or a plain tuple.0

Accept the lowest-entropy positions under the entropy bound and renoise the rest.

Parameters:

entropy_bound (float, defaults to 0.1) : The maximum tolerated joint entropy of the accepted tokens. Larger values accept more tokens per step.

t_max (float, defaults to 0.8) : Sampling temperature on the first denoising step.

t_min (float, defaults to 0.4) : Sampling temperature on the last denoising step.

num_inference_steps (int, defaults to 32) : The maximum number of denoising steps.

EntropyBoundSchedulerOutput[[diffusers.EntropyBoundSchedulerOutput]]

diffusers.EntropyBoundSchedulerOutput[[diffusers.EntropyBoundSchedulerOutput]]

Source

Output class for the entropy bound scheduler.

Parameters:

prev_sample (torch.LongTensor of shape (batch_size, block_length)) : Updated block tokens after the current denoising step.

accepted_index (torch.BoolTensor of shape (batch_size, block_length)) : Boolean mask of the positions accepted (committed) in this step.

sampled_tokens (torch.LongTensor of shape (batch_size, block_length)) : Token IDs sampled from the model logits.

sampled_probs (torch.Tensor of shape (batch_size, block_length)) : Probabilities of the sampled tokens.

pred_logits (torch.Tensor of shape (batch_size, block_length, vocab_size)) : The temperature-scaled logits the candidates were drawn from, for self-conditioning the next step.