diffusers/main/ja/api/schedulers/block_refinement.md

BlockRefinementScheduler

BlockRefinementScheduler は、離散トークン拡散におけるブロック単位の反復的リファインメントを管理します。各ステップで最も確信度の高いトークンを確定し、モデルが別のトークンを高い確信度で予測している場合、既に確定済みのトークンを編集するオプションも提供します。

このスケジューラーは、LLaDA2Pipelineで使用されます。

BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler[[diffusers.BlockRefinementScheduler]]

diffusers.BlockRefinementScheduler(block_length: int = 32, num_inference_steps: int = 32, threshold: float = 0.95, editing_threshold: float | None = None, minimal_topk: int = 1)

ソース

ブロック単位の反復的改善（コミットバイコンフィデンス）用のスケジューラー。

各ステップにおいて、スケジューラーはモデルのlogitsから候補トークンをサンプリングし、最も信頼度の高いものを確定します。1ステップあたりに確定するトークン数は、ブロック長をリファインメントステップ数で均等に分配することによって決定されます。

オプションとして編集をサポートします。すべてのマスクトークンが解決された後、モデルが正のediting_threshold（None、0.0、または負の値は編集を無効にします）を超える信頼度で別のトークンを予測した場合、トークンを置き換えることができます。

add_noise[[diffusers.BlockRefinementScheduler.add_noise]]

add_noise(original_samples: torch.LongTensor, attention_mask: torch.LongTensor, prompt_length: int, block_length: int, mask_token_id: int, generator: torch.Generator | None = None)

ソース

パラメータ:

original_samples (torch.LongTensor、shape (batch_size, seq_len)) : クリーンなトークン ID。

attention_mask (torch.LongTensor、shape (batch_size, seq_len)) : Padding mask（有効なトークンは 1、padding は 0）。

prompt_length (int) : マスクせずに保持する先頭プロンプトトークンの数。

block_length (int) : マスキング用のブロックサイズ。

mask_token_id (int) : マスクされた位置に使用するトークンID。

generator (torch.Generator, 任意) : 再現性のために使用する乱数生成器。

戻り値: tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

(noisy, noisy_rev, masked, masked_rev) — 2つの補完的なノイジーシーケンスと、それに対応するブールマスク。

半自己回帰ブロックマスクングに対して、フォワード（ノイジング）プロセスを適用します。

プロンプト以降の各ブロックにおいて、有効な（パディングでない）トークンのランダムな割合がmask_token_idに置き換えられます。2つの補完的なビュー、noisyとnoisy_revが返されます。片方のマスクされた位置は、もう片方のマスクされていない位置に対応します。

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)

ソース

パラメータ:

step_idx (int) : このブロック内の現在のリファインメントステップのインデックス。

masks_remaining (bool) : ブロック内にマスクトークンが残っているかどうか。

editing_enabled (bool) : 編集モードが有効かどうかを示します。

editing_transfer_index (torch.BoolTensor) : このステップで編集されたトークンを示します。

post_steps (int) : 現在までに実行されたマスク後の編集ステップ数。

max_post_steps (int) : マスク適用後の編集ステップの最大許容数。

finished (torch.BoolTensor) : バッチごとの終了フラグ（EOS 検出による）。

戻り値: bool

リファインメントを続行する場合はTrue、中断する場合はFalseを返します。

現在のブロックについて、内部の改善ループを継続すべきかどうかを判定します。

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)

ソース

パラメータ:

cur_x (torch.LongTensor、shape (batch_size, seq_len)) : 現在のウィンドウまでのすべてのブロックを含む、現在の完全なシーケンス。

sampled_tokens (torch.LongTensor、shape (batch_size, block_length)) : このステップでスケジューラーによってサンプリングされたトークン。

final_transfer (torch.BoolTensor、shape (batch_size, block_length)) : コミット済みと編集済みの位置のマスクを結合したもの。

finished (torch.BoolTensor、shape (batch_size,)) : 現在のバッチごとの終了フラグ。

eos_token_id (int) : EOS トークン ID。

mask_token_id (int) : マスクトークンID。

prompt_length (int) : シーケンスの先頭にあるプロンプトトークンの数。

戻り値: torch.BoolTensor

finished フラグを更新しました。

EOS トークンがコミットされた場合、バッチごとの終了フラグを更新します。

get_num_transfer_tokens[[diffusers.BlockRefinementScheduler.get_num_transfer_tokens]]

get_num_transfer_tokens(block_length: int, num_inference_steps: int)

ソース

block_lengthトークンのコミットをnum_inference_stepsステップに均等に分配します。

ステップ[[diffusers.BlockRefinementScheduler.step]]

step(model_output: torch.Tensor, timestep: int | torch.Tensor, sample: torch.LongTensor, 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 | None = None, editing_threshold: float | None = None, minimal_topk: int | None = None, prompt_mask: torch.BoolTensor | None = None, generator: torch.Generator | None = None, return_dict: bool = True)

ソース

パラメータ:

model_output (torch.Tensor、shape (batch_size, block_length, vocab_size)) : 現在のブロックに対するモデルからの生の logits。

timestep (int または torch.Tensor) : ブロックの精細化スケジュールにおける現在のステップインデックス。

sample (torch.LongTensor、shape (batch_size, block_length)) : 現在のブロックのトークン ID（マスクベースモードでは、未コミット位置のマスクトークンを含む）。

mask_token_id (int, 任意) : マスクされた位置に使用されるトークン ID。Noneの場合、スケジューラーは均一破損モードで動作します。このモードでは、コミットされた位置を内部で追跡し（timestep == 0でリセット）、コミットされていない位置を均一にランダムなトークンで再ノイズ化します。これにより、DiffusionGemma のブロックリファインメントサンプラーと一致します。

temperature (float) : サンプリング温度。

top_p (float, 任意) : ニュクレアサンプリングのカットオフ値。

top_k (int, 任意) : Top-k サンプリングの切り捨て値。

sampling_method (str) : サンプリング手法 (auto、greedy、multinomial)。

threshold (float, 任意) : トークンを確定するための信頼度閾値です。デフォルトでは設定値が使用されます。

editing_threshold (float, optional) : マスクされていないトークンを編集するための信頼度しきい値。編集を有効にするには正の値である必要があります。デフォルトでは設定値が使用されます。

minimal_topk (int, optional) : 各ステップでコミットするトークンの最小数。デフォルトでは設定値が使用されます。

prompt_mask (torch.BoolTensor, 任意) : 形状が (block_length,) のブールマスクで、True がプロンプト（編集不可）の位置を示します。

generator (torch.Generator, 任意) : サンプリング用の乱数生成器。

return_dict (bool) : BlockRefinementSchedulerOutputを返すか、タプルを返すかを指定します。

単一のリファインメントステップを実行します：logitsからサンプリングし、確信度の高いトークンを確定し、必要に応じて既存のトークンを編集します。

BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput[[diffusers.BlockRefinementSchedulerOutput]]

diffusers.BlockRefinementSchedulerOutput(prev_sample: torch.LongTensor, transfer_index: torch.BoolTensor, editing_transfer_index: torch.BoolTensor, sampled_tokens: torch.LongTensor, sampled_probs: torch.Tensor, pred_logits: torch.Tensor)

ソース

パラメータ:

prev_sample (torch.LongTensor、shape (batch_size, block_length)) : 現在のリファインメントステップ後の更新されたブロックトークン。

transfer_index (torch.BoolTensor、shape (batch_size, block_length)) : どのトークンがコミットされたかを示すブールマスク（マスク埋め込み）。

editing_transfer_index (torch.BoolTensor、shape (batch_size, block_length)) : どのトークンが編集されたか（非マスク置換）を示すブールマスク。

sampled_tokens (torch.LongTensor、shape (batch_size, block_length)) : モデルの logits からサンプリングされたトークン ID。

sampled_probs (torch.Tensor、shape (batch_size, block_length)) : サンプルされたトークンの確率。

pred_logits (torch.Tensor、形状 (batch_size, block_length, vocab_size)) : 次ステップの自己条件付けのためにそのまま渡される、denoiser の logits。

ブロックリファインメントスケジューリング用の出力クラス。