Зачем сжимать диффузионные модели: проблема ресурсов и скорости
Stable Diffusion 2.1 требует минимум 6.9 ГБ VRAM только для загрузки весов в FP16. На генерацию одного изображения 512×512 уходит 8–12 секунд на RTX 3090. Развернуть такую модель в продакшене с десятками тысяч запросов в сутки - это кластер GPU и счет на тысячи долларов ежемесячно. Для мобильных устройств и веб-демок базовая версия непригодна.
Запрос аудитории очевиден: нужны облегченные версии, которые сохраняют приемлемое качество и работают быстрее. Команда Segmind предложила конкретное решение - модели SD-Small и SD-Tiny с открытым исходным кодом и весами. Сжатие выполнено методом дистилляции знаний (knowledge distillation): студенческая модель учится повторять выходы учителя на каждом уровне UNet, а не просто копирует финальное изображение. Результат - минус 35% и 55% параметров при сопоставимом визуальном качестве и двукратном ускорении инференса.
Этот подход отличается от стандартного прунинга или квантизации. Дистилляция знаний переносит поведение модели, а не просто обрезает архитектуру. Если вас интересует общая механика дистилляции языковых моделей и её бизнес-применение, разбор доступен в статье про дистилляцию LLM и снижение затрат на инференс в 25–30 раз.
Как работает дистилляция знаний в UNet: архитектура SD-Small и SD-Tiny
Архитектура Stable Diffusion построена вокруг UNet - сверточной сети с encoder-decoder структурой и skip-connections. В базовой версии UNet содержит 12 блоков: 4 down-sampling, 1 middle и 4 up-sampling с соответствующими skip-соединениями. Каждый блок состоит из residual-слоев, self-attention и cross-attention модулей. Общее число параметров UNet - около 860 млн для Stable Diffusion 2.1.
Segmind применила дистилляцию знаний на уровне промежуточных представлений (feature-based distillation). Студенческая модель обучается минимизировать разницу между своими активациями и активациями учителя на каждом оставшемся блоке. Это принципиально: модель учится воспроизводить внутренние представления, а не только финальный пиксельный результат.
Удаление блоков UNet: как выбрать, чем жертвовать
SD-Small сохраняет 8 блоков UNet вместо 12. Удалены два down-sampling блока (второй и третий) и два up-sampling блока (второй и третий). Middle-блок сохранен полностью - он отвечает за глобальную структуру изображения и семантическую связность. SD-Tiny идет дальше: оставлено 6 блоков, удалены три down-sampling и три up-sampling блока. Middle-блок снова нетронут.
Выбор основан на эмпирических тестах. Блоки, ближайшие к входу и выходу (первые и последние), критичны для детализации текстур и мелких объектов. Средние блоки обрабатывают высокоуровневую композицию. Удаление средних блоков сильнее бьет по композиционности, но меньше влияет на текстуры. Segmind выбрала компромисс: жертвуем средними блоками, сохраняем края.
Параметры моделей: SD-Small - 579 млн параметров (минус 35% от 860 млн), SD-Tiny - 387 млн параметров (минус 55%). Для сравнения: квантизация INT8 снижает размер модели вдвое, но не ускоряет вычисления на большинстве GPU без специальной поддержки. Дистилляция же дает реальное ускорение за счет сокращения операций.
Функция потерь: обучение на совпадение выходов
Loss-функция комбинирует три компонента. Первый - L2-расстояние между выходами студента и учителя на каждом сохраненном блоке UNet. Для каждого блока i вычисляется MSE между картами активаций: L_feat_i = ||f_teacher_i(x) - f_student_i(x)||². Второй компонент - стандартный diffusion loss: MSE между предсказанным шумом студента и реальным шумом. Третий - опциональный perceptual loss на основе VGG-16, добавленный для сохранения текстурного качества.
Обучение велось на датасете LAION Art Aesthetic - подмножестве LAION-5B, отфильтрованном по эстетической оценке. Это 120 млн изображений с высоким визуальным качеством, преимущественно цифровой арт, иллюстрации и фотографии с выраженной композицией. Выбор датасета объясняет стилевой уклон моделей: они лучше работают с артовыми промптами и хуже с фотореалистичными сценами, требующими точной детализации.
Тренировка SD-Small заняла 256 A100-часов, SD-Tiny - 180 A100-часов. Учитель - Stable Diffusion 2.1-base. Студенты инициализированы весами учителя с удаленными блоками, а не случайно - это ускорило сходимость в 3–4 раза.
Практический выигрыш: скорость, память и качество генерации
Главный вопрос: стоит ли сжатие потери качества? Короткий ответ - зависит от задачи. На одиночных объектах с простым фоном разница между SD-Small и базовой моделью едва заметна. На сложных сценах с несколькими взаимодействующими объектами потери становятся очевидными.
Бенчмарки инференса: цифры для SD-Small и SD-Tiny
Тесты проводились на RTX 3090 (24 ГБ VRAM) и A100 (40 ГБ). Генерация одного изображения 512×512, 50 шагов DDIM, CFG scale 7.5, batch size 1:
| Модель | Время на RTX 3090 (сек) | Время на A100 (сек) | VRAM (ГБ) |
|---|---|---|---|
| SD 2.1-base | 9.8 | 4.2 | 6.9 |
| SD-Small | 5.1 | 2.3 | 4.1 |
| SD-Tiny | 3.9 | 1.9 | 3.2 |
Ускорение SD-Small относительно базы - 92% на RTX 3090 и 83% на A100. SD-Tiny - 151% и 121% соответственно. VRAM снижена на 41% и 54%. На batch size 4 разрыв увеличивается: SD-Tiny генерирует 4 изображения за 8.1 секунды, базовая модель - за 22.3 секунды.
Качество оценивалось по FID на COCO-2017 (30K изображений). SD-Small показал FID 8.4 против 7.9 у учителя. SD-Tiny - 9.6. CLIP-score: SD-Small 0.311, SD-Tiny 0.298, учитель 0.319. Разрыв есть, но для многих продуктовых сценариев он некритичен.
Запуск моделей с diffusers: код и инструкции
Веса доступны на Hugging Face в репозиториях segmind/sd-small и segmind/sd-tiny. Лицензия - CreativeML Open RAIL-M, допускает коммерческое использование с ограничениями. Минимальный пример генерации:
from diffusers import StableDiffusionPipeline
import torch
pipe = StableDiffusionPipeline.from_pretrained(
"segmind/sd-small",
torch_dtype=torch.float16
)
pipe = pipe.to("cuda")
image = pipe(
"a cat wearing sunglasses, digital art, trending on artstation",
num_inference_steps=50,
guidance_scale=7.5
).images[0]
image.save("output.png")Для SD-Tiny замените модель на segmind/sd-tiny. Рекомендуется использовать 30–50 шагов для SD-Small и 25–40 для SD-Tiny - модель сходится быстрее из-за упрощенной архитектуры. CFG scale 7.0–8.0 дает стабильные результаты. При значениях выше 9.0 появляются артефакты: перенасыщенные цвета и потеря деталей в темных областях.
Параметр negative_prompt критичен для сжатых моделей. Без негативного промпта SD-Tiny часто генерирует размытый фон и искаженные лица. Рекомендуемый негативный промпт: «blurry, low quality, distorted, deformed, ugly, bad anatomy, watermark».
Дообучение через LoRA: адаптация под свои задачи
Полное дообучение сжатых моделей не рекомендуется: удаленные блоки не восстановить, а переобучение оставшихся ведет к катастрофическому забыванию. LoRA (Low-Rank Adaptation) добавляет обучаемые низкоранговые матрицы к attention-слоям, не меняя базовые веса. Размер адаптера - 5–15 МБ против 2+ ГБ полной модели.
Пример тренировки LoRA на пользовательском датасете из 20–50 изображений:
from diffusers import StableDiffusionPipeline
from peft import LoraConfig, get_peft_model
pipe = StableDiffusionPipeline.from_pretrained(
"segmind/sd-small",
torch_dtype=torch.float16
)
lora_config = LoraConfig(
r=8,
lora_alpha=16,
target_modules=["to_q", "to_k", "to_v", "to_out.0"],
lora_dropout=0.1
)
pipe.unet = get_peft_model(pipe.unet, lora_config)
# Далее стандартный цикл тренировки с датасетомLoRA-адаптеры для SD-Small/Tiny совместимы с большинством тренеров сообщества. Особенность: из-за удаленных блоков эффективный ранг адаптера можно увеличить до r=16 без переобучения - модель менее склонна к запоминанию конкретных примеров.
Ограничения и подводные камни: что нужно знать перед внедрением
Сжатые модели не универсальны. Четкое понимание ограничений сэкономит часы экспериментов и разочарований. Три главные проблемы: композиционность, редкие концепты и стилевой уклон.
Проблема композиционности: почему сложные промпты дают сбои
Промпт «a red cube on top of a blue sphere, studio lighting» базовая модель обрабатывает корректно: красный куб стоит на синей сфере. SD-Small в 40% случаев меняет порядок - синий куб на красной сфере, или объекты парят раздельно. SD-Tiny в 65% случаев теряет один из объектов или сливает их в бесформенную массу.
Причина - потеря емкости в удаленных блоках UNet. Cross-attention слои в этих блоках отвечают за связывание токенов промпта с пространственными позициями. Меньше блоков - меньше capacity для моделирования сложных отношений между объектами. Это фундаментальное ограничение архитектуры, а не недостаток тренировки.
Обходные пути: разбивайте сложные промпты на последовательную генерацию с img2img, используйте ControlNet с картами глубины или сегментации для фиксации композиции. Для задач, где точное расположение объектов критично (рекламные макеты, интерфейсы), используйте полную модель.
Другие ограничения
Редкие концепты страдают сильнее. Промпт «a capybara wearing a monocle» SD-Small еще генерирует, SD-Tiny часто выдает нечто среднее между капибарой и кошкой. Датасет LAION Art Aesthetic содержит мало изображений редких животных и специфических объектов.
Стилевой уклон в арт - прямое следствие датасета. Фотореалистичные промпты требуют добавления уточнений: «photorealistic, 8k, detailed skin texture, raw photo». Без них модель скатывается в цифровую живопись. Текстовые элементы в изображении (вывески, надписи) генерируются с искажениями - attention-механизм сжатых моделей хуже удерживает точную форму букв.
Пороговые рекомендации: если задача требует генерации 100+ изображений в час на одной GPU, выбирайте SD-Small. Если нужна скорость 200+ изображений и допустимы умеренные потери качества - SD-Tiny. Для ответственных применений (медицинские иллюстрации, прецизионные макеты) оставьте полную модель.
Segmind и открытый исходный код: как внедрить и участвовать
Репозиторий Segmind на GitHub содержит полный код тренировки, конфигурации и скрипты оценки. Лицензия Apache 2.0 для кода, CreativeML Open RAIL-M для весов - коммерческое использование разрешено. Сообщество активно: на момент публикации репозиторий имеет 1.2K звезд, 23 контрибьютора и 140 форков.
Веса доступны на Hugging Face с документацией и примерами. Открыты issues с баг-репортами и feature-реквестами. Команда Segmind приветствует дообученные адаптеры и бенчмарки на специфических доменах. Если вы провели тесты на своих данных - результаты помогут сообществу принимать обоснованные решения о внедрении.
Для тех, кто работает с более тяжелыми моделями на локальном железе, полезен опыт сжатия DeepSeek-V4-Flash до 54 ГБ через IQ2_XXS - смежные техники квантизации и дистилляции решают общую задачу запуска больших моделей на ограниченных ресурсах.