GeoSPRINT распределяет шаги диффузии по геометрии траектории

Арпита Джоши из The Scripps Research Institute представила GeoSPRINT для ускорения диффузионных моделей в препринте, который не проходил рецензирования. По замерам автора, метод сохраняет качество DDIM при сокращении числа вычислений нейросети на 25%. Работа важна командам, которые уже используют диффузионную модель и хотят снизить стоимость генерации без дообучения.

Что сделали

Генерация диффузионной моделью состоит из последовательных шагов: на каждом модель оценивает, как изменить текущее скрытое представление изображения. Один такой вызов называется вычислением нейронной функции, или NFE. Чем меньше NFE требуется для результата заданного качества, тем ниже потенциальные задержка и вычислительная стоимость.

Обычное расписание DDIM выбирает шаги равномерно и не учитывает форму конкретной траектории удаления шума. GeoSPRINT сначала прогоняет набор эталонных траекторий по полному расписанию, а затем ищет участки, где последовательные скрытые состояния лежат почти на одной прямой. Такие состояния считаются геометрически избыточными: соседние сохранённые точки уже описывают направление движения.

Расстояние точки от прямой или более общего аффинного подпространства вычисляется через QR-разложение матрицы направлений. На испытанных моделях автору оказалось достаточно проверки на прямолинейность. Результаты нескольких эталонных прогонов объединяются в одно универсальное расписание, поэтому анализ не нужно повторять для каждого нового изображения.

Расписание сочетает два сигнала. Первый распределяет шаги по логарифмическому отношению сигнала к шуму и учитывает ошибку численного решения. Второй отражает кривизну траектории: больше шагов выделяется там, где её направление меняется быстрее. Чистый выбор точек только по кривизне оставлял слишком большие промежутки, поэтому GeoSPRINT смешивает оба распределения.

Дополнительно работа вводит показатель проекции траектории αtraj. Он измеряет долю вариации, которая не восстанавливается из соседних сохранённых точек. Нулевое значение соответствует идеально прямой траектории; малое значение означает, что удалённые шаги почти не добавляли нового направления. Показатель можно использовать отдельно от ускорения — например, для сравнения прямолинейности траекторий разных вариантов модели или планировщика.

Что показали

Качество автор оценивает через FID — расстояние между распределениями признаков сгенерированных и эталонных изображений; меньшее значение считается лучшим. На CIFAR-10 GeoSPRINT снизил FID относительно равномерного DDIM на 0,7–1,1 пункта в испытанном диапазоне бюджетов. Это означает, что преимущество наблюдалось не только при одном специально выбранном расписании.

На LSUN Church FID уменьшился с 1,48 до 1,26. Для Stable Diffusion v1.5 максимальная разница с DDIM составила 1,93 пункта. Автор также получила преимущество над DPM-Solver++ на части бюджетов CIFAR-10, хотя GeoSPRINT использовал более простой решатель первого порядка. В рамках этого эксперимента выбор моментов вычисления оказался важнее повышения порядка решателя.

Отдельные испытания показали, что распределение только по отношению сигнала к шуму уступало смешанному варианту. Следовательно, результат нельзя свести к замене равномерной временной шкалы на уже известную шкалу решателя: вклад вносит именно профиль кривизны, полученный из эталонных траекторий.

Ограничения

Метод проверяли только на изображениях: CIFAR-10 размером 32×32, LSUN Church размером 256×256 и скрытых представлениях Stable Diffusion v1.5, соответствующих изображениям 512×512. Работа не показывает, сохраняется ли эффект для видео, аудио и других типов диффузионных моделей.

Универсальное расписание требует заранее построить полные эталонные траектории. Для единичных генераций эта подготовка может стоить дороже сэкономленных вычислений; подход рассчитан на сценарий, где расходы распределяются по серии запросов. В работе сравнивается число NFE, но не приведено измерение итоговой задержки и стоимости на GPU, поэтому уменьшение числа шагов нельзя напрямую считать таким же снижением производственных расходов.

Параметр смешивания настраивали на CIFAR-10, а для Stable Diffusion использовали фиксированный уровень управления текстовым условием. Не показано, останется ли одно расписание оптимальным при других настройках модели и распределении запросов. Сравнение с DPM-Solver++ проведено на CIFAR-10, но не расширено на модели высокого разрешения; эмпирического сравнения с методом SDM, анализирующим локальную жёсткость дифференциального уравнения, также нет.

Для Stable Diffusion FID считался относительно изображений, созданных полным расписанием DDIM, а не относительно набора реальных изображений. Этот тест измеряет близость к результатам исходного решателя, но не доказывает улучшение абсолютного визуального качества.

Что это значит

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

Планы стоит менять, если генерация выполняется сериями на одном контрольном снимке модели и офлайн-подготовка расписания амортизируется по многим запросам. В таком случае анализ расписания разумно поставить перед более дорогим дообучением. αtraj также даёт диагностический сигнал: он помогает увидеть, какие модели и настройки создают более прямые, потенциально лучше сокращаемые траектории.

Работа не даёт оснований менять архитектуру или обещать пропорциональное снижение счёта за GPU. Перед внедрением нужно извлечь расписание на собственных запросах и уровне управления условием, затем сравнить одинаковые бюджеты NFE по задержке, стоимости и продуктовым показателям качества. Пока результаты показывают перспективность оптимизации расписания, а не готовую универсальную замену существующих решателей.

Источники

• GeoSPRINT: Geometric Redundancy-Aware Step Pruning for Inference in Diffusion Trajectories