U-Net: предсказатель шума

Урок разбирает мотор Stable Diffusion — сеть U-Net, которая на каждом шаге оценивает шум.

U-Net — нейросеть, которая по зашумлённому латентному коду, номеру шага и текстовому условию предсказывает шум, который нужно убрать.

Кто делает основную работу

В прошлых разделах мы говорили: на каждом шаге обратного процесса модель предсказывает шум. Эта модель и есть U-Net. Её вызывают по разу на каждый шаг диффузии (отсюда — чем больше steps, тем больше вызовов U-Net и тем дольше генерация).

Почему форма буквы «U»

U-Net устроена симметрично: сначала сжимающая часть постепенно уменьшает разрешение латента, улавливая общий смысл («где тут объект, какая композиция»), потом разжимающая часть восстанавливает разрешение, возвращая детали. Между ними есть skip-связи — короткие мостики, которые передают точные детали из ранних слоёв напрямую в поздние, чтобы они не потерялись при сжатии.

вход (зашумл. латент)
   \  сжатие (смысл)        разжатие (детали)  /
    \____                          ____/
         \____     skip-связи  ____/
              \__   ----->   __/
                 \__________/
                  узкое место
           выход: предсказанный шум

Как сюда попадает текст

Чтобы U-Net рисовала именно то, что просили, текстовое условие (от CLIP, см. следующий урок) подмешивается внутрь сети через механизм cross-attention: на разных уровнях U-Net «смотрит» на слова промпта и подстраивает предсказание шума под них. Грубо говоря, attention решает, к каким участкам картинки какое слово относится.

Как работает под капотом

На вход U-Net получает три вещи: текущий зашумлённый латент, номер шага (чтобы знать уровень шума) и эмбеддинги промпта. На выходе — оценка шума той же формы, что и латент. Сэмплер берёт эту оценку и обновляет латент для следующего шага. Поскольку U-Net вызывается десятки раз за генерацию, именно она потребляет основную видеопамять и время.

Частые ошибки

Думать, что U-Net выдаёт картинку. Она выдаёт оценку шума в латенте; в пиксели его потом превращает VAE.
Игнорировать номер шага. Без информации об уровне шума сеть не знала бы, насколько «сильно» расшумлять.
Считать, что текст влияет «снаружи». Текст вшит внутрь U-Net через cross-attention на каждом шаге, а не применяется в конце.

Почему U-Net потребляет больше всего ресурсов

Из всего пайплайна именно U-Net вызывают чаще всего: по разу (а с guidance — дважды) на каждый шаг. При 30 шагах и включённом guidance это около 60 проходов большой сети за одну картинку. Поэтому, когда генерация «съедает» видеопамять или идёт медленно, виновата почти всегда U-Net, а не CLIP или VAE, которые отрабатывают по одному разу.

Из этого следуют практические выводы. Хотите быстрее — уменьшайте число шагов или размер картинки (он же размер латента, с которым работает U-Net на каждом проходе). Дообучение вроде LoRA тоже добавляет работы именно U-Net, поскольку поправки вшиваются в её слои внимания. Понимая, что U-Net — главный потребитель ресурсов, легче осознанно искать баланс между скоростью и качеством.

Итог

U-Net — главная сеть, предсказывающая шум на каждом шаге диффузии.
Её «U»-форма со skip-связями ловит и общий смысл, и мелкие детали.
Текстовое условие входит в U-Net через cross-attention, направляя расшумление к смыслу промпта.