Policy-Mask: обученная политика выбирает признаки замороженной диффузионной модели под заданный бюджет каналов
Диффузионные модели можно использовать не только для генерации картинок, но и как источник признаков (промежуточных представлений) для задач распознавания: сегментации, классификации, поиска соответствий между изображениями. Авторы препринта отмечают, что выбрать такие признаки из замороженной модели непросто: полезный сигнал распределён по траектории удаления шума, по глубине сети и по каналам. На практике нужно выбрать сразу три вещи: шаг по времени (timestep), блок сети и подмножество каналов.
Существующие методы, по словам авторов, либо задают эти выборы вручную, либо оптимизируют лишь несколько гиперпараметров извлечения признаков, чаще всего только шаг по времени, а затем применяют одно и то же правило для всех входов и всех бюджетов.
Авторы предлагают Policy-Mask: политику, зависящую от бюджета и обучаемую напрямую на целевой задаче. Она совместно выбирает все три компонента (шаг, блок и набор каналов) при заранее заданном бюджете по числу каналов. Отдельная техническая деталь: без нормализации вклад оценки каналов растёт вместе с числом выбранных каналов и начинает подавлять обновления для шага и блока. Авторы нормируют его на размер подмножества, благодаря чему три «головы» политики остаются в сопоставимых масштабах, а совместное обучение стабилизируется.
При одинаковых бюджетах Policy-Mask, как заявляют авторы, превосходит базовые методы с фиксированным извлечением и методы с фильтрацией признаков по трём типам задач. Приведены итоговые значения: 38,9 mIoU на ADE20K (сегментация), 75,4% top-1 на CIFAR-100 при 128 измерениях (классификация) и 76,3 PCK@0,1 на SPair-71k (семантические соответствия).
Покомпонентные абляции (отключение частей метода по одной) показывают, что выбор блока заметно влияет на итоговую точность. Наибольший зафиксированный результат по бюджету каналов приходится на 256 сохранённых каналов; авторы видят в этом практичную рабочую точку для компактных диффузионных представлений.
Наконец, метод адаптирован к диффузионным трансформерам, включая PixArt-Σ и SD3.5 MMDiT: вместо каналов U-Net там используются скрытые размерности токенов изображения.
Ключевые факты
- Policy-Mask, политика, обученная на целевой задаче и совместно выбирающая шаг денойзинга, блок и подмножество каналов замороженной диффузионной модели при заданном бюджете каналов.
- Нормировка вклада каналов на размер подмножества удерживает три «головы» политики в сопоставимых масштабах и стабилизирует совместное обучение.
- Заявленные результаты: 38,9 mIoU на ADE20K, 75,4% top-1 на CIFAR-100 при 128 измерениях, 76,3 PCK@0,1 на SPair-71k.
- Абляции показывают, что выбор блока существенно влияет на точность; наивысший зафиксированный результат по бюджету каналов, при 256 каналах.
- Метод адаптирован к диффузионным трансформерам (PixArt-Σ, SD3.5 MMDiT) заменой каналов U-Net на скрытые размерности токенов изображения.
Почему это важно
Признаки замороженных диффузионных моделей используют в задачах распознавания, но способ их извлечения обычно подбирают вручную или настраивают лишь по шагу времени. Работа предлагает вместо этого обучаемый выбор сразу по трём осям (шаг, блок, каналы) с явным ограничением на размер представления. Это шаг к более компактным признакам, но пока это препринт с заявленными авторами результатами.
Кому это важно
Исследователям компьютерного зрения, которые берут готовые диффузионные модели как «экстрактор признаков» для сегментации, классификации и поиска соответствий, а также тем, кого интересуют компактные представления с ограничением по размерности.
Как это применить
Из описания видно, что идея переносима: бюджет каналов задаётся заранее, а политика подбирает шаг, блок и каналы под конкретную задачу. Для диффузионных трансформеров предложена замена каналов U-Net на скрытые размерности токенов изображения. Авторы называют 256 каналов практичной рабочей точкой. Сведений о выпуске кода, затратах вычислений и времени инференса в описании нет.
Можно ли доверять
Все утверждения исходят от самих авторов препринта и приведены в кратком описании. В нём не названы авторы и организации, не указаны значения базовых методов, поэтому величина улучшения не видна. Не названа и диффузионная модель, на которой получены результаты на ADE20K, CIFAR-100 и SPair-71k. Результатов для PixArt-Σ и SD3.5 MMDiT не приводится. Независимой проверки в источнике нет.
Риски и подводные камни
Без базовых цифр и названия основной модели сравнить метод с другими нельзя. Бюджет каналов для ADE20K и SPair-71k не указан. Максимальный результат по бюджету каналов при 256 каналах отражает лишь зафиксированные в экспериментах значения. Перенос на диффузионные трансформеры заявлен, но подтверждающих чисел нет.