PixelCNN++: почему смесь логистик победила softmax на 256 классов

Оригинальный PixelCNN моделировал каждый пиксель как категориальное распределение на 256 значений и предсказывал их через softmax. Работало, но упиралось в две проблемы: параметров на выходе слишком много, а модель не знала, что значение 127 «ближе» к 128, чем к 0. PixelCNN++ (работа команды OpenAI, представлена на ICLR 2017) поменял сам способ описывать пиксель — и заодно упростил архитектуру. Ниже — что именно изменили и как это устроено на уровне кода.

Что было не так с исходным PixelCNN

Автогрессионная генерация изображений строится на простой идее: вероятность картинки раскладывается в произведение условных вероятностей пикселей, где каждый следующий пиксель зависит от всех предыдущих. Формально это цепное правило вероятности¹, а на практике — свёрточная сеть с маскированными фильтрами, которая «не подсматривает» вперёд.

Исходная модель для каждого субпикселя (отдельно R, G, B) выдавала 256-мерный softmax. Отсюда три следствия:

• Раздутый выходной слой. 256 логитов на каждый из трёх каналов — это много весов, которые тратятся впустую: соседние интенсивности почти всегда имеют близкую вероятность, но модель учит их независимо.
• Нет понятия «близости». Softmax обращается с 256 значениями как с несвязанными классами. Градиент от ошибки «предсказал 130 вместо 128» такой же, как от «предсказал 5 вместо 128».
• Разреженные градиенты. При редких значениях интенсивности часть выходов почти не обучается.

Главная замена: дискретизированная смесь логистик

Вместо категориального распределения PixelCNN++ моделирует интенсивность пикселя как непрерывную смесь логистических распределений, а затем дискретизирует её под 256 целочисленных уровней. Модель предсказывает для каждого пикселя не 256 логитов, а набор из нескольких компонент смеси, каждая со своими тремя параметрами: вес, среднее и масштаб.

Дискретизация делается через разность значений логистической CDF на границах бина: вероятность целого значения x — это интеграл плотности по интервалу шириной в один уровень вокруг x. Для крайних значений (0 и 255) интервал раскрывается до бесконечности, чтобы вероятности корректно суммировались в единицу.

Модель больше не угадывает один из 256 ящиков. Она описывает форму распределения несколькими гладкими горбами — и это распределение уже само знает, что 128 и 130 рядом.

Практический выигрыш — компактный выходной слой. Обычно берут 5–10 компонент смеси; при 10 компонентах на пиксель приходится порядка 10 весов плюс параметры средних и масштабов по каналам, что заметно меньше 256 логитов на канал.

Как учитываются цветовые каналы

Каналы R, G, B внутри одного пикселя не независимы. PixelCNN++ вводит линейную зависимость средних: предсказание зелёного корректируется на уже сгенерированный красный, а синий — на красный и зелёный. Это дешевле, чем полноценный автогрессионный проход по каналам, но сохраняет ключевую корреляцию.

Остальные модификации

Смесь логистик — заголовочная идея, но не единственная. Авторы переработали и архитектуру.

Даунсемплинг вместо ставки на дилатированные свёртки удешевляет вычисления: сеть работает на нескольких разрешениях, а пропущенная за счёт страйда информация возвращается через skip-соединения. Dropout здесь принципиален — без него глубокая PixelCNN++ переобучается на CIFAR-10.

Как это выглядит в коде

Референсная реализация авторов написана на TensorFlow, но суть функции потерь удобнее показать на упрощённом фрагменте: вероятность целого значения — разность двух сигмоид (логистическая CDF), а лосс — отрицательный логарифм смеси. Псевдокод в стиле PyTorch:

def discretized_logistic_logprob(x, means, log_scales, logit_probs):
    # x нормализован в [-1, 1]; 255 уровней -> шаг 1/127.5
    inv_scale = torch.exp(-log_scales)
    centered = x.unsqueeze(-1) - means

    plus_in  = inv_scale * (centered + 1.0 / 255.0)
    min_in   = inv_scale * (centered - 1.0 / 255.0)

    cdf_plus = torch.sigmoid(plus_in)
    cdf_min  = torch.sigmoid(min_in)

    # вероятность попадания в бин шириной 1/255
    prob = cdf_plus - cdf_min
    log_prob_mid = torch.log(torch.clamp(prob, min=1e-12))

    # смешиваем компоненты через веса logit_probs
    log_mix = log_prob_mid + torch.log_softmax(logit_probs, dim=-1)
    return torch.logsumexp(log_mix, dim=-1)

Крайние случаи (x около -1 и около +1) в полной реализации обрабатываются отдельно: там вместо разности CDF берут её значение на одном конце, а при очень узких бинах используют приближение через плотность. Без этих оговорок лосс уходит в NaN на насыщенных пикселях — типичная ловушка при попытке переписать функцию с нуля.

Что сверять при запуске

1. Нормализация входа. Пиксели приводят к диапазону около [-1, 1], иначе масштабы логистик придётся переучивать.
2. Число компонент смеси. Начните с 5–10; больше редко улучшает метрику, но замедляет обучение.
3. Обработка граничных значений 0 и 255 — обязательна, иначе сумма вероятностей не равна единице.
4. Dropout включён на обучении и выключен на сэмплировании.

Почему это важно за пределами PixelCNN

Дискретизированная смесь логистик пережила саму модель. Тот же приём для описания дискретных пиксельных значений позднее использовали в других генеративных подходах, где выход — целочисленная интенсивность, а не непрерывная величина. Идея простая и переносимая: если ваши целевые значения дискретны, но упорядочены, категориальный softmax почти всегда проигрывает дискретизированному непрерывному распределению — и по числу параметров, и по качеству градиента.

¹ Цепное правило вероятности — разложение совместной вероятности набора величин в произведение условных: p(a,b,c) = p(a)·p(b|a)·p(c|a,b). В автогрессионных моделях изображения роль величин играют пиксели в фиксированном порядке обхода.