diffusers/pr_14741/en/api/schedulers/block_refinement.md

BlockRefinementScheduler

The BlockRefinementScheduler manages block-wise iterative refinement for discrete token diffusion. At each step it commits the most confident tokens and optionally edits already-committed tokens when the model predicts a different token with high confidence.

This scheduler is used by LLaDA2Pipeline.

For the uniform corruption process, which has no mask token, use UniformRefinementScheduler instead.

This scheduler follows the shared discrete diffusion scheduler contract: a decreasing float corruption level in (0, 1], step(model_output, timestep, sample), sampling knobs on the config, and a DiscreteSchedulerOutput return.

BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, mask_token_id: int | None = None, temperature: float = 0.0, top_p: float | None = None, top_k: int | None = None, sampling_method: str = 'auto', threshold: float = 0.95, editing_threshold: float | None = None)

Source

Parameters:

block_length (int, defaults to 32) : The block size this scheduler is configured for. Pipelines read it as their default block size; the commit quota itself is taken from the width of sample.

num_inference_steps (int, defaults to 32) : The number of refinement steps the commit quota is spread across.

mask_token_id (int, optional) : Token ID marking an undecided position. Required by step(); it lives in the config because it is a property of the tokenizer, matching AmusedScheduler.

temperature (float, defaults to 0.0) : Sampling temperature applied to the logits when drawing candidates. 0.0 takes the argmax. The confidence driving the quota is measured on the unscaled distribution, so threshold and editing_threshold do not move with this value.

top_p (float, optional) : Nucleus sampling cutoff.

top_k (int, optional) : Top-k sampling cutoff.

sampling_method (str, defaults to "auto") : One of "auto", "greedy", "multinomial". "auto" draws multinomially when temperature != 0.

threshold (float, defaults to 0.95) : Confidence above which a masked position commits even if the quota is already met.

editing_threshold (float, optional) : Confidence above which an already-resolved position is overwritten with a different predicted token. Must be positive to enable editing; None, 0.0, or negative disables it.

Scheduler for block-wise iterative refinement (commit-by-confidence).

At each step, the scheduler samples candidate tokens from model logits and commits those with the highest confidence. The number of tokens to commit per step is determined by evenly distributing the block length across the number of refinement steps.

Optionally supports editing: after all mask tokens are resolved, tokens can be replaced if the model predicts a different token with confidence above a positive editing_threshold (None, 0.0, or negative disables editing).

This scheduler models the absorbing (masked) corruption process. For the uniform process, where every position always holds a real token and there is no mask token, use UniformRefinementScheduler.

add_noise[[diffusers.BlockRefinementScheduler.add_noise]]

add_noise(original_samples: torch.LongTensor, timesteps: float | torch.Tensor, generator: torch.Generator | None = None, **kwargs)

Source

Parameters:

original_samples (torch.LongTensor of shape (batch_size, seq_len)) : Clean token IDs.

timesteps (float or torch.Tensor) : Masking probability. A scalar applies to the whole batch; a tensor of shape (batch_size,) or (batch_size, 1) gives a per-example rate.

generator (torch.Generator, optional) : RNG for reproducibility.

Returns: tuple[torch.LongTensor, torch.BoolTensor]

the noisy tokens and the boolean mask of noised positions.

Apply the forward (noising) process: replace each position with mask_token_id with probability timesteps.

timesteps is the corruption level in [0, 1], so it is the expected masking fraction — 1 masks everything, 0 masks nothing. The caller chooses it, matching every other add_noise in the library.

check_block_should_continue[[diffusers.BlockRefinementScheduler.check_block_should_continue]]

check_block_should_continue(step_idx: int, masks_remaining: bool, editing_enabled: bool, editing_transfer_index: torch.BoolTensor, post_steps: int, max_post_steps: int, finished: torch.BoolTensor)

Source

Parameters:

step_idx (int) : Current refinement step index within this block.

masks_remaining (bool) : Whether any mask tokens remain in the block.

editing_enabled (bool) : Whether editing mode is active.

editing_transfer_index (torch.BoolTensor) : Which tokens were edited in this step.

post_steps (int) : Number of post-mask editing steps taken so far.

max_post_steps (int) : Maximum allowed post-mask editing steps.

finished (torch.BoolTensor) : Per-batch finished flags (from EOS detection).

Returns: bool

True if refinement should continue, False to break.

Determine whether the inner refinement loop should continue for the current block.

check_eos_finished[[diffusers.BlockRefinementScheduler.check_eos_finished]]

check_eos_finished(cur_x: torch.LongTensor, sampled_tokens: torch.LongTensor, final_transfer: torch.BoolTensor, finished: torch.BoolTensor, eos_token_id: int, mask_token_id: int, prompt_length: int)

Source

Parameters:

cur_x (torch.LongTensor of shape (batch_size, seq_len)) : Current full sequence including all blocks up to the current window.

sampled_tokens (torch.LongTensor of shape (batch_size, block_length)) : Tokens sampled by the scheduler in this step.

final_transfer (torch.BoolTensor of shape (batch_size, block_length)) : Combined mask of committed and edited positions.

finished (torch.BoolTensor of shape (batch_size,)) : Current per-batch finished flags.

eos_token_id (int) : EOS token ID.

mask_token_id (int) : Mask token ID.

prompt_length (int) : Number of prompt tokens at the start of the sequence.

Returns: torch.BoolTensor

Updated finished flags.

Update per-batch finished flags when EOS tokens are committed.

get_num_transfer_tokens[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

Source

Evenly distribute block_length token commits across num_inference_steps steps.

Deprecated: the per-step quota is now computed inline in step() from step_index, so there is no schedule tensor to build.

index_for_timestep[[diffusers.BlockRefinementScheduler.index_for_timestep]]

index_for_timestep(timestep: Union[float, torch.FloatTensor], schedule_timesteps: Optional[torch.FloatTensor] = None)

Source

Parameters:

timestep (float or torch.FloatTensor) : The timestep to find the index for.

schedule_timesteps (torch.FloatTensor, optional) : The schedule timesteps to validate against. If None, the scheduler's timesteps are used.

Returns: int

The index of the timestep.

Get the index for the given timestep.

set_begin_index[[diffusers.BlockRefinementScheduler.set_begin_index]]

set_begin_index(begin_index: int = 0)

Source

Parameters:

begin_index (int, defaults to 0) : The begin index for the scheduler.

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

set_timesteps[[diffusers.BlockRefinementScheduler.set_timesteps]]

set_timesteps(num_inference_steps: int, device: str | torch.device | None = None, **kwargs)

Source

Parameters:

num_inference_steps (int) : The number of refinement steps.

device (str or torch.device, optional) : The device the timesteps should be moved to.

Set the discrete timestep grid, as the decreasing corruption level t in [0, 1].

The grid matches set_timesteps() — 1.0 down to 1 / num_inference_steps — so the discrete schedulers are interchangeable in a pipeline loop. timesteps is the public loop variable; the commit quota is derived from the integer step_index so it stays exact for any num_inference_steps.

step[[diffusers.BlockRefinementScheduler.step]]

step(model_output: torch.Tensor, timestep: float | torch.Tensor, sample: torch.LongTensor, generator: torch.Generator | None = None, return_dict: bool = True, **kwargs)

Source

Parameters:

model_output (torch.Tensor of shape (batch_size, block_length, vocab_size)) : Raw logits from the model for the current block.

timestep (float or torch.Tensor) : The current corruption level, one entry of ~BlockRefinementScheduler.timesteps.

sample (torch.LongTensor of shape (batch_size, block_length)) : Current block token IDs, with mask_token_id at the positions still undecided.

generator (torch.Generator, optional) : RNG for sampling.

return_dict (bool) : Whether to return a DiscreteSchedulerOutput or a plain tuple.

Perform a single refinement step: sample from logits, commit confident masked positions, and optionally edit already-resolved ones.

step_edit[[diffusers.BlockRefinementScheduler.step_edit]]

step_edit(model_output: torch.Tensor, sample: torch.LongTensor, generator: torch.Generator | None = None, return_dict: bool = True)

Source

Parameters:

model_output (torch.Tensor of shape (batch_size, block_length, vocab_size)) : Raw logits from the model for the current block.

sample (torch.LongTensor of shape (batch_size, block_length)) : Current block token IDs, with every position resolved.

generator (torch.Generator, optional) : RNG for sampling.

return_dict (bool) : Whether to return a DiscreteSchedulerOutput or a plain tuple.

Overwrite already-resolved positions whose prediction is both different and confident.

This is the post-mask refinement phase: once no mask_token_id remains there is nothing left to unmask, so the step is a confidence-thresholded overwrite rather than a diffusion step and takes no timestep. It also does not advance step_index, so a pipeline can run as many sweeps as it likes after exhausting the schedule.

DiscreteSchedulerOutput[[diffusers.DiscreteSchedulerOutput]]

diffusers.DiscreteSchedulerOutput[[diffusers.DiscreteSchedulerOutput]]

diffusers.DiscreteSchedulerOutput(prev_sample: LongTensor, pred_original_sample: LongTensor, sampled_probs: Tensor, pred_logits: Tensor, committed_mask: BoolTensor, edited_mask: typing.Optional[torch.BoolTensor] = None)

Source

Parameters:

prev_sample (torch.LongTensor of shape (batch_size, sequence_length)) : Computed sample (x_{t-1}) of the previous timestep, as token IDs. prev_sample should be used as the next model input in the denoising loop.

pred_original_sample (torch.LongTensor of shape (batch_size, sequence_length)) : The predicted clean token at each position, i.e. the discrete analog of x0. Sampled from the model distribution rather than taken as its argmax, unless the sampling configuration is greedy.

sampled_probs (torch.Tensor of shape (batch_size, sequence_length)) : Probability of each token in pred_original_sample under the unmodified denoiser distribution, so that confidence thresholds mean the same thing across schedulers and across sampling temperatures.

pred_logits (torch.Tensor of shape (batch_size, sequence_length, vocab_size)) : The distribution the tokens were actually drawn from, after any shaping the scheduler applies (temperature, top-k, top-p, or a schedule-dependent annealing). Returned because the scheduler owns that shaping, so it is the only holder of this tensor; pipelines that self-condition the denoiser on its own prediction need exactly this and not the raw logits they passed in.

committed_mask (torch.BoolTensor of shape (batch_size, sequence_length)) : Positions that adopted their predicted token in this step. For schedulers that resample every position each step, this is every position.

edited_mask (torch.BoolTensor of shape (batch_size, sequence_length), optional) : Subset of positions that overwrote an already-committed token. None for schedulers with no editing concept.

Base class for the output of a discrete diffusion scheduler's step function.

Discrete diffusion operates on token IDs rather than continuous latents, so a step yields the token predictions and the per-position bookkeeping of which positions were decided, alongside the prev_sample hand-off common to every scheduler.