diffusers/main/ko/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)

소스

블록 단위 반복적 정제(신뢰도 기반 커밋)를 위한 스케줄러

각 단계에서 스케줄러는 모델 로짓에서 후보 토큰을 샘플링하고 가장 높은 신뢰도를 가진 토큰을 확정합니다. 단계별로 확정할 토큰의 수는 블록 길이를 리파인먼트 단계 수로 균등하게 나누어 결정됩니다.

선택적으로 편집을 지원합니다: 모든 mask 토큰이 해결된 후, 모델이 양의 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)) : Clean token ID.

attention_mask (torch.LongTensor, shape (batch_size, seq_len)) : 패딩 마스크 (유효한 토큰은 1, 패딩은 0).

prompt_length (int) : 마스킹하지 않고 유지할 앞부분 프롬프트 토큰의 개수.

block_length (int) : 마스킹을 위한 블록 크기.

mask_token_id (int) : 마스킹된 위치에 사용할 토큰 ID.

generator (torch.Generator, 선택 사항) : 재현성을 위한 RNG.

반환값: tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor, torch.BoolTensor]

(noisy, noisy_rev, masked, masked_rev) — 두 개의 상보적 노이즈 시퀀스와 이에 해당하는 불리언 마스크입니다.

반전파(noising) 과정을 반자동 회귀 블록 마스킹에 적용합니다.

프롬프트 이후의 각 블록에서, 유효한(패딩이 아닌) 토큰의 무작위 비율을 mask_token_id로 대체합니다. 두 개의 상보적인 뷰가 반환됩니다: 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 단계에 걸쳐 균등하게 분배합니다.

step[[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 기반 모드에서는 미커밋된 위치의 mask 토큰을 포함합니다).

mask_token_id (int, 선택 사항) : 마스킹된 위치에 사용되는 토큰 ID입니다. None인 경우, 스케줄러는 균일한 손상 모드에서 실행됩니다. 이 모드에서는 내부적으로 커밋된 위치를 추적하고(timestep == 0에서 초기화됨) 커밋되지 않은 위치를 균일하게 무작위로 선택된 토큰으로 재노이즈하여, DiffusionGemma의 블록 리파인먼트 샘플러와 일치합니다.

temperature (float) : 샘플링 온도.

top_p (float, 선택 사항) : Nucleus sampling cutoff.

top_k (int, 선택 사항) : Top-k 샘플링 절단 값.

sampling_method (str) : 샘플링 방법 (auto, greedy, multinomial).

threshold (float, 선택 사항) : 토큰 커밋을 위한 신뢰도 임계값. 기본값은 config 값입니다.

editing_threshold (float, 선택 사항) : 마스크가 아닌 토큰을 편집할 때 사용할 신뢰도 임계값입니다. 편집을 활성화하려면 양수여야 합니다. 기본값은 config 값입니다.

minimal_topk (int, 선택 사항) : 각 단계에서 커밋할 최소 토큰 수. 기본값은 config 값입니다.

prompt_mask (torch.BoolTensor, 선택 사항) : (block_length,) 모양의 불리언 마스크로, True는 프롬프트(편집 불가) 위치를 나타냅니다.

generator (torch.Generator, 선택 사항) : 샘플링을 위한 RNG.

return_dict (bool) : BlockRefinementSchedulerOutput 또는 튜플을 반환할지 여부.

…

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)) : 현재 refinement 단계 이후 업데이트된 block 토큰.

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, shape (batch_size, block_length, vocab_size)) : 다음 단계의 self-conditioning을 위해 전달되는 denoiser logits입니다.

블록 리파인먼트 스케줄링을 위한 출력 클래스.