Учёные нашли скрытые «регистры смысла» в диффузионных трансформерах и ускорили генерацию на 20%
Диффузионные трансформеры (DiT) для генерации изображений по тексту одновременно обрабатывают текстовые и изобразительные токены, но что именно происходит внутри них во время денойзинга (шаг за шагом устранения шума при генерации картинки) было плохо понятно. Команда исследователей во главе с Маохуа Ли предложила каузальный (причинный) фреймворк интерпретируемости для крупных DiT: он сочетает разложение внимания (attention) с целевыми вмешательствами, точечным отключением или изменением отдельных диапазонов токенов, голов внимания и слоёв, чтобы проверить, что реально влияет на результат, а не просто коррелирует с ним.
С помощью этого инструмента авторы разделили токены на два типа: те, что несут содержание самого промпта, и служебные структурные токены шаблона, которые пайплайн генерации добавляет вокруг текста запроса. Оказалось, что структурные токены сами по себе почти не несут специфики промпта на выходе текстового энкодера. Но внутри DiT они неожиданно становятся главными точками притяжения внимания «от изображения к тексту» (attention sinks) и именно они каузально удерживают идентичность объекта на протяжении всей генерации, то есть работают как неявные «регистры смысла», которые сама модель незаметно завела для хранения контекста.
При этом свою роль эти токены получают не напрямую от промпта. Смысл сначала «впечатывается» в латентные представления изображения, а уже оттуда считывается обратно в служебные токены шаблона, а не передаётся из токенов промпта в токены шаблона напрямую. Иными словами, токены, которые на входе кодируют смысл, не те же самые токены, что удерживают этот смысл в процессе генерации.
На основе находки авторы предложили не требующее дообучения правило прунинга (обрезки) модели: головы внимания, которые сильнее всего внимательны к токенам промпта, оказались избыточными для итогового результата. Их отключение убирает 20% вычислений внимания (attention FLOPs) и снижает качество всего на 1,4 пункта по бенчмарку GenEval (тесту на соответствие сгенерированного изображения текстовому описанию). Дополнительно работа показывает, как вычисления в DiT организованы по головам и глубине сети: смысловая маршрутизация отделена от визуального синтеза, а генерация проходит путь от формирования идентичности объекта к её распространению и уточнению деталей.
Ключевые факты
- Предложен каузальный фреймворк интерпретируемости для крупных диффузионных трансформеров (DiT): разложение внимания плюс целевые вмешательства по диапазонам токенов, головам и слоям
- Служебные структурные токены шаблона почти не несут смысла промпта на выходе энкодера, но внутри DiT становятся главными точками притяжения внимания и каузально удерживают идентичность объекта, работают как неявные «регистры смысла»
- Смысл передаётся косвенно: сначала из промпта в латентные представления изображения, а уже потом считывается обратно в токены шаблона, а не напрямую из токенов промпта
- На основе находки создано не требующее дообучения правило прунинга: отключение голов внимания, сильнее всего смотрящих на промпт, убирает 20% attention FLOPs при падении качества всего на 1,4 пункта по GenEval
- Показано, как вычисления в DiT распределены по головам и глубине: смысловая маршрутизация отделена от визуального синтеза, генерация идёт от формирования идентичности объекта к её распространению и уточнению
Почему это важно
Это одна из первых работ, которая объясняет не «что делает» диффузионный трансформер, а «как именно» он это делает изнутри, с помощью причинных экспериментов, а не догадок по картинкам внимания. Находка ломает интуитивное предположение: токены, которые на входе несут смысл (слова промпта), не обязаны быть теми же токенами, что удерживают этот смысл во время генерации, эту роль неожиданно взяли на себя служебные токены шаблона.
Кому это важно
Разработчикам и инженерам моделей генерации изображений по тексту на архитектуре DiT (диффузионный трансформер), тем, кто строит и оптимизирует такие пайплайны для продакшена. Полезно и исследователям интерпретируемости генеративных моделей, которые ищут методы разбирать внутреннюю механику трансформеров, а не только языковых моделей.
Как это применить
Главный практический результат, правило прунинга (обрезки) без дообучения: можно отключить головы внимания, которые сильнее всего сфокусированы на токенах промпта, поскольку они оказались избыточны для сохранения смысла картинки. Это убирает 20% вычислений внимания (attention FLOPs) при инференсе, снижая качество генерации лишь на 1,4 пункта по бенчмарку GenEval, то есть прирост скорости почти без заметной потери результата.
Можно ли доверять
Работа опирается не на визуализацию внимания (что часто вводит в заблуждение), а на каузальные вмешательства, точечное отключение конкретных токенов, голов и слоёв с измерением результата, что методологически надёжнее чистой корреляции. Приведены конкретные количественные результаты (20% FLOPs, 1,4 пункта GenEval), которые можно проверить. При этом это препринт без указания рецензируемой публикации, поэтому выводы стоит воспринимать как предварительные до независимого воспроизведения.
Риски и подводные камни
Метод прунинга протестирован в рамках конкретной методологии и конкретных моделей DiT; неясно, насколько правило переносится на другие архитектуры, чекпоинты и типы контента без потери качества сверх заявленных 1,4 пункта. Есть и общий риск интерпретируемости: находки о «регистрах смысла» получены на основе вмешательств в конкретных экспериментальных условиях и могут не полностью описывать поведение модели на реальных разнообразных промптах.
«Головы внимания, которые сильнее всего сфокусированы на токенах промпта, оказались избыточными: их отсечение убирает 20% вычислений внимания (attention FLOPs) при потере всего 1,4 пункта по бенчмарку GenEval»
— из статьи Маохуа Ли и соавторов