DiffusionGemma

DiffusionGemma는 블록 확산(block-diffusion) 인코더-디코더 언어 모델입니다. 인과(causal) 인코더가 클린 프롬프트(그리고 이전에 생성된 블록)를 KV 캐시로 읽어들이고, 양방향(bidirectional) 디코더가 해당 캐시에 크로스 어텐션을 수행하여 canvas_length 토큰 크기의 고정 크기 “캔버스”를 디노이즈(denoise)합니다. 생성 과정은 캔버스에 대한 외부 자기회귀(autoregressive) 루프와 내부 디노이즈 루프를 번갈아 수행하며, 각 단계에서 후보 토큰을 샘플링하고 균일한 손상(uniform corruption) 모드의 BlockRefinementScheduler를 통해 가장 신뢰도가 높은 토큰을 확정(commit)한 후 나머지 토큰에 다시 노이즈를 적용합니다. 모델 자체는 transformers에 DiffusionGemmaForBlockDiffusion으로 포함되어 있으며, 공개된 체크포인트는 google/diffusiongemma-26B-A4B-it입니다.

사용

import torch
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion

from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

model_id = "google/diffusiongemma-26B-A4B-it"
model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)
scheduler = BlockRefinementScheduler()

pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
pipe.model.model.decoder = torch.compile(pipe.model.model.decoder, mode="reduce-overhead", fullgraph=True)
output = pipe(
    prompt="Why is the sky blue?",
    gen_length=256,
    num_inference_steps=48,
    cache_implementation="static",
)
print(output.texts[0])

num_inference_steps는 캔버스당 디노이징 단계 수입니다(48은 공개된 체크포인트와 일치합니다). 단계 수가 적을수록 빠르지만 품질이 낮아집니다. cache_implementation="static"을 사용하면 디코더를 cudagraphs와 함께 torch.compile할 수 있습니다(정적 캐시 및 컴파일 참조). 더 간단한 동적 캐시 실행을 원하면 이 두 가지를 모두 제거하세요.

멀티턴 또는 멀티모달 입력의 경우, prompt 대신 원시 messages 대화를 전달합니다. 이는 일반적인 채팅 형식의 {"role", "content"} 딕셔너리 목록이며, 프로세서가 이를 채팅 템플릿을 통해 처리합니다:

messages = [
    {"role": "user", "content": "Why is the sky blue?"},
]
# or with an image:
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": "Describe this image."},
        ],
    },
]
output = pipe(messages=messages, gen_length=256)

단일 사용자 턴의 경우 messages를 생략하고 prompt와 함께 image를 전달할 수 있습니다. 프로세서가 이를 자동으로 모델의 이미지 입력으로 변환합니다.

스케줄러

스케줄러는 각 캔버스를 디노이즈하는 샘플러이며, 교체 가능합니다. 다른 부분을 건드리지 않고 스케줄러만 교체하여 샘플링 전략을 변경할 수 있습니다. 사용할 수 있는 스케줄러는 세 가지입니다:

• BlockRefinementScheduler (기본값): 각 단계에서 가장 확신도가 높은 토큰(즉, threshold를 초과하는 토큰과 단계별 균등 할당량)를 확정하고 나머지에 다시 노이즈를 적용합니다. editing_threshold를 사용하면 이미 확정된 토큰을 다시 수정할 수도 있습니다.
• DiscreteDDIMScheduler: 균일한 손상 과정(D3PM)의 정확한 이산 사후 분포에서 각 위치를 샘플링합니다. 매개변수가 없으며, 마지막 단계에서 예측된 토큰을 결정적으로 확정합니다.
• EntropyBoundScheduler: entropy_bound 이하의 결합 엔트로피를 유지하는 가장 낮은 엔트로피 위치를 커밋하여, 대략적으로 독립적인 토큰을 함께 수락합니다. 첫 번째 단계에서 t_max(0.8)로 시작하여 마지막 단계에서 t_min(0.4)까지 샘플링 온도를 서서히 낮추며, 이는 공개된 체크포인트의 샘플러와 일치합니다.

from diffusers import DiscreteDDIMScheduler, EntropyBoundScheduler

pipe.scheduler = DiscreteDDIMScheduler()
# or: pipe.scheduler = EntropyBoundScheduler(entropy_bound=0.1)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

Scheduler별 샘플링 파라미터(블록 정밀화 threshold/top_k, 엔트로피 상한 등)는 scheduler config에서 설정합니다:

from diffusers import BlockRefinementScheduler

pipe.scheduler = BlockRefinementScheduler.from_config(pipe.scheduler.config, threshold=0.9)

EntropyBoundScheduler는 디노이징 단계 동안 내부적으로 샘플링 온도(t_max/t_min)를 감쇠시킵니다. DiscreteDDIMScheduler와 BlockRefinementScheduler는 파이프라인에 전달되는 고정된 temperature를 사용합니다(그리디의 경우 0.0).

Predictor-corrector 샘플링

DiscreteDDIMScheduler는 Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation의 leave-one-out predictor-corrector를 지원합니다. corrector_steps개의 Gibbs sweeps를 수행하여 노이즈된 marginal의 one-coordinate conditional에서 가장 확신이 낮은 위치들을 다시 샘플링함으로써 캔버스를 정제합니다. 이 과정은 해당 marginal을 불변으로 유지하면서 추가적인 학습 비용 없이 생성 품질을 향상시킵니다. 공개된 체크포인트에서 바로 동작합니다. uniform diffusion의 경우 denoiser와 leave-one-out posterior가 closed form으로 상호 교환 가능하므로, corrector는 재학습 없이 필요한 leave-one-out 양을 복원할 수 있습니다.

corrector sweeps는 num_inference_steps 예산에 추가되는 것이 아니라 그 안에 포함됩니다. 파이프라인은 predictor 단계 수를 줄이고, 확보된 forward를 corrector에 사용하므로 모델 forward의 총 수는 num_inference_steps로 유지되며, predictor-corrector의 비용은 일반적인 ancestral sampling과 동일합니다.

from diffusers import DiscreteDDIMScheduler

pipe.scheduler = DiscreteDDIMScheduler(corrector_steps=2, corrector_k=12)
output = pipe(prompt="Why is the sky blue?", gen_length=256, num_inference_steps=48)
print(output.texts[0])

PEPT 어댑터

Denoiser는 🤗 Transformers 모델이므로, 어댑터는 diffusers의 load_lora_weights API가 아니라 해당 모델의 네이티브 PEFT 통합을 통해 로드됩니다. 이 통합은 어댑터 유형에 무관하므로, 동일한 호출로 LoRA, DoRA 또는 기타 PEFT 어댑터(예: TRL의 SFTTrainer 출력)를 로드할 수 있습니다. 모델 구성 요소에서 직접 어댑터를 관리하세요:

pipe.model.load_adapter("path/to/adapter", adapter_name="sft")  # LoRA, DoRA, ...
pipe.model.set_adapter("sft")
output = pipe(prompt="Why is the sky blue?", gen_length=256)

pipe.model.disable_adapters()  # run the base model
pipe.model.delete_adapter("sft")

Adapters는 활성화된 상태로 유지되며 병합되지 않습니다: DiffusionGemma는 인코더와 디코더의 기본 가중치를 연결(tie)하므로, 어댑터를 이 가중치에 병합하면 두 가지 분기(branch) 모두 손상됩니다.

정적 캐시 및 컴파일

pipeline은 블록마다 인코더를 한 번만 사전 입력하여 재사용 가능한 캐시(기본값은 DynamicCache)에 저장합니다. cache_implementation="static"를 전달하면 고정된 모양의 StaticCache를 대신 사용하며, 이 모양 덕분에 cudagraphs를 사용하여 디코더를 torch.compile하여 추가적인 속도 향상을 얻을 수 있습니다(pipeline은 각 스텝을 마킹하고 logits를 복제하여 cudagraph 메모리가 덮어쓰여지지 않도록 합니다). 이 구성은 사용법에 나와 있습니다. 더 간단한 동적 캐시 실행을 원하면 torch.compile 호출과 cache_implementation="static"를 모두 제거하세요.

적응형 중단

블록은 일반적으로 num_inference_steps가 모두 소진되기 전에 수렴하므로, 기본적으로 파이프라인은 각 배치 예시의 argmax 예측이 stability_threshold 단계 동안 안정적이고 토큰당 평균 엔트로피가 confidence_threshold(발행된 체크포인트에서 사용된 값인 0.005) 미만으로 떨어지면 해당 예시를 동결합니다. 모든 예시가 동결되면 디노이징 루프가 종료됩니다. 이는 동일한 품질에서 디코더 포워드를 약 절반으로 줄여 주며, 단일 스루풋 레버 중 가장 큰 것입니다. confidence_threshold=None을 전달하면 항상 전체 num_inference_steps를 실행합니다:

output = pipe(prompt="Why is the sky blue?", gen_length=256, confidence_threshold=None)  # disable adaptive stopping

콜백

Callbacks는 각 denoising 단계 이후에 실행됩니다. callback_on_step_end_tensor_inputs를 전달하여 callback_kwargs에 포함할 텐서를 선택할 수 있으며, canvas(현재 블록 토큰)와 logits가 사용 가능합니다. Callback에서 {"canvas": ...}를 반환하면 canvas가 교체됩니다.

def on_step_end(pipe, step, timestep, callback_kwargs):
    canvas = callback_kwargs["canvas"]
    # Inspect or modify `canvas` here.
    return {"canvas": canvas}


out = pipe(
    prompt="Why is the sky blue?",
    callback_on_step_end=on_step_end,
    callback_on_step_end_tensor_inputs=["canvas"],
)

DiffusionGemmaPipeline

>>> import torch
>>> from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
>>> from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline

>>> model_id = "google/diffusiongemma-26B-A4B-it"
>>> model = DiffusionGemmaForBlockDiffusion.from_pretrained(model_id, dtype=torch.bfloat16, device_map="auto")
>>> processor = AutoProcessor.from_pretrained(model_id)
>>> scheduler = BlockRefinementScheduler()

>>> pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
>>> output = pipe(prompt="Why is the sky blue?", gen_length=256)
>>> print(output.texts[0])

DiffusionGemmaPipelineOutput