Раскрытие SDXL Turbo: Интерпретация текст-на-изображение моделей с помощью разреженных автокодировщиков

Текст-на-изображение генерация — это быстро развивающаяся область, которая привлекает внимание благодаря моделям, таким как DALL-E, которые могут создавать высококачественные изображения на основе текстовых запросов. Сегодняшние лучшие модели в основном основаны на текст-условных диффузионных моделях, которые могут быть обучены на масштабных наборах данных, таких как LAION-5B, и генерировать фотореалистичные изображения. Однако, несмотря на их впечатляющие результаты, внутренние механизмы этих нейронных сетей остаются не до конца понятыми из-за их "черного ящика" природы.

Недавно появилось направление исследований, называемое механистической интерпретируемостью, которое стремится обратить инженерию моделей для понимания их внутреннего устройства. В этой области разреженные автокодировщики (SAE) показали свою эффективность в разложении промежуточных представлений языковых моделей (LLM) на интерпретируемые признаки, облегчая контроль и последующий анализ. Однако подобный анализ для текст-на-изображение моделей был ограничен.

В данной статье мы исследуем возможность применения SAE для извлечения интерпретируемых признаков из SDXL Turbo, недавно выпущенной открытой модели текст-на-изображение с несколькими шагами генерации. Мы обучаем SAE на обновлениях, выполняемых трансформерными блоками внутри U-net SDXL Turbo, и обнаруживаем, что их обученные признаки интерпретируемы, оказывают причинное влияние на процесс генерации и раскрывают специализацию среди блоков. Наш анализ показывает, что один блок в основном занимается композицией изображения, другой добавляет локальные детали, а третий отвечает за цвет, освещение и стиль.

Основные понятия

Разреженные автокодировщики (SAE)

Разреженные автокодировщики (SAE) представляют собой нейронные сети, которые обучаются разлагать входные данные на разреженные представления, где большинство значений признаков равно нулю. Это позволяет выявлять моносемантические (однозначные) признаки, которые легче интерпретировать и анализировать.

• Энкодер: Преобразует входные данные в разреженное представление.
• Декодер: Восстанавливает исходные данные из разреженного представления.

SAE эффективны в обработке данных, где многие признаки могут быть неактивны одновременно, что часто встречается в промежуточных представлениях нейронных сетей, особенно в LLM.

Диффузионные модели

Диффузионные модели работают, постепенно добавляя шум к исходным данным и затем обучаясь на обратном процессе — удалении шума для восстановления исходных данных. Этот процесс можно описать следующим образом:

1. Добавление шума: Начальное изображение коррумпируется гауссовым шумом.
2. Обучение на обратном процессе: Модель обучается предсказывать шум, добавленный на каждом шаге, чтобы восстановить исходное изображение.

SDXL Turbo

SDXL Turbo — это дистиллированная версия Stable Diffusion XL, которая позволяет генерировать изображения высокого качества в течение всего нескольких шагов, что значительно ускоряет процесс генерации по сравнению с оригинальной моделью.

• Архитектура: Использует U-net, где текстовое условие интегрируется через кросс-аттеншн в трансформерных блоках.
• Дистилляция: Модель обучается с использованием адверсарной игры и дистилляции шума, чтобы уменьшить количество шагов, необходимых для генерации.

Применение SAE к SDXL Turbo

Мы применили SAE к обновлениям, выполняемым трансформерными блоками в SDXL Turbo, чтобы исследовать, как они влияют на процесс генерации изображений. Вот как это было сделано:

1. Сбор данных: Мы собрали промежуточные представления (фичемапы) из SDXL Turbo на 1.5 миллиона запросов из набора данных LAION-COCO.

2. Обучение SAE: Для каждого трансформерного блока мы обучили SAE, используя методологию из (Gao et al., 2024), с функцией активации TopK и вспомогательной потерей для обработки "мертвых" признаков.

3. Анализ: Мы провели качественный и количественный анализ обученных признаков, чтобы понять их интерпретируемость и причинное влияние на генерацию изображений.

Качественный анализ

Мы использовали визуализационные техники для анализа интерпретируемости и причинного эффекта признаков:

• Пространственные активации: Визуализация карт активации признаков на изображении.
• Модуляция активации: Изменение силы активации конкретных признаков для наблюдения за изменениями в генерируемых изображениях.
• Активация на пустом контексте: Активация признаков без текстового запроса, чтобы понять их независимое влияние.

Количественный анализ

Для подтверждения качественных наблюдений мы выполнили серию экспериментов:

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

Заключение

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