DiffusionGemma

DiffusionGemma — это блочно-диффузионная модель языка с архитектурой «кодировщик-декодировщик». Причинный кодировщик считывает чистый промпт (а также ранее сгенерированные блоки) в KV-кэш, а двунаправленный декодер выполняет денуазинг фиксированного «холста» из canvas_length токенов, обращаясь к этому кэшу через кросс-внимание. Генерация чередует внешний авто-регрессивный цикл по холстам с внутренним циклом денуазинга, на каждом шаге которого отбираются кандидаты токенов, наиболее уверенные из них фиксируются через BlockRefinementScheduler в режиме равномерного повреждения, а остальные подвергаются повторному шумовому воздействию. Сама модель находится в 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" позволяет декодеру компилироваться с помощью torch.compile с использованием cudagraphs (см. Статический кэш и компиляция); для более простого запуска с динамическим кэшем уберите оба параметра.

Для многораундовых или мультимодальных входов передайте исходный разговор messages вместо prompt. Это список словарей {"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 и передать image вместе с prompt; процессор автоматически преобразует их во входные данные изображения для модели.

Планировщики

Планировщик — это сэмплер, который выполняет денуазинг каждого холста, и он заменяем: вы можете заменить его, чтобы изменить стратегию сэмплирования, не изменяя ничего другого. Доступны три планировщика:

• 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])

Параметры выборки, специфичные для планировщика (порог threshold/top_k для уточнения блока, ограничение энтропии и т. д.), задаются в конфигурации планировщика:

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 для жадного выбора).

Сэмплирование с предсказателем-корректором

DiscreteDDIMScheduler поддерживает предиктор-корректор leave-one-out из работы Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation. Он уточняет холст с помощью corrector_steps шагов Гиббса, которые пересэмплируют позиции с наименьшей уверенностью из однокоординатного условного распределения шумного маргинального распределения, что оставляет это маргинальное распределение инвариантным и улучшает генерацию без дополнительных затрат на обучение. Он работает непосредственно с опубликованной контрольной точкой: для uniform diffusion денойзер и leave-one-out апостериорное распределение взаимозаменяемы в замкнутой форме, поэтому корректор восстанавливает необходимые ему leave-one-out величины без какого-либо переобучения.

Шаги корректора включаются в бюджет num_inference_steps, а не добавляются сверх него: пайплайн выполняет меньше шагов предиктора и тратит высвобожденные проходы на корректоры, поэтому общее число проходов модели остаётся равным num_inference_steps, а стоимость предиктора-корректора совпадает со стоимостью обычного анцестрального сэмплирования.

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])

PEFT-адаптеры

Денойзер представляет собой модель 🤗 Transformers, поэтому адаптеры загружаются через его нативную интеграцию PEFT, а не через API load_lora_weights из diffusers. Поскольку эта интеграция не зависит от типа адаптера, те же вызовы загружают LoRA, DoRA или любой другой адаптер PEFT (например, результат работы SFTTrainer из TRL). Управляйте адаптерами непосредственно на компоненте модели:

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")

Адаптеры остаются активными и не объединяются: DiffusionGemma связывает базовые веса энкодера и декодера, поэтому слияние адаптера с ними повредит обе ветви.

Статический кэш и компиляция

Пайплайн выполняет предварительное заполнение энкодера один раз на блок в переиспользуемый кэш (по умолчанию это DynamicCache). Передача cache_implementation="static" использует вместо него StaticCache с фиксированной формой, чьи размеры позволяют применить torch.compile к декодеру с cudagraphs для дополнительного ускорения (пайплайн помечает каждый шаг и клонирует логиты, чтобы память cudagraph не была перезаписана); именно такая конфигурация показана в разделе Usage. Чтобы выполнить более простой прогон с динамическим кэшем, удалите и вызов 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) выполняются после каждого шага удаления шума. Передайте callback_on_step_end_tensor_inputs, чтобы выбрать, какие тензоры будут включены в callback_kwargs; доступны canvas (текущие токены блока) и logits. Верните из обработчика {"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