Диффузионный трансформер на ESP32-S3: генерация лиц по тексту за 5,4 секунды
Разработчик портировал латентную диффузионную модель на микроконтроллер ESP32-S3, доказав, что генерация изображений по текстовому описанию возможна без облачных сервисов и мощных GPU. В статье он подробно описал, как уместить диффузионный трансформер в 512 КБ оперативной памяти и добиться приемлемого качества на разрешении 64×64.
Эксперимент, опубликованный в потоке «Разработка», демонстрирует полноценный запуск диффузионной модели на микроконтроллере ESP32-S3, стоимость отладочной платы с которым составляет несколько сотен рублей. Автор задался целью не просто отправить промпт в облако, а честно прогнать все вычисления диффузии непосредственно на чипе. В результате на устройстве с двумя ядрами по 240 МГц и 512 КБ оперативной памяти удалось сгенерировать лицо размером 64×64 пикселя по текстовому запросу вроде «улыбающаяся девушка со светлыми волосами» всего за 5,4 секунды. Это достижение контрастирует с требованиями современных моделей: Stable Diffusion нуждается в видеокарте с 4–8 ГБ памяти, а Flux — в 12–24 ГБ, тогда как ESP32-S3 не имеет ни тензорных ядер, ни сопоставимых ресурсов.
Ключевая проблема заключалась в объёме памяти: веса только текстового энкодера CLIP из Stable Diffusion занимают 250 МБ, что в 500 раз превышает всю оперативную память микроконтроллера. Поэтому автор отказался от готовых моделей и спроектировал архитектуру с нуля, ориентируясь на жёсткие ограничения железа. Были приняты радикальные меры: домен сужен до генерации лиц (датасет CelebA из 200 тысяч фотографий с 40 атрибутами), разрешение снижено до 64×64, а диффузия проводится не в пиксельном пространстве, а в сжатом латенте 16×16×4, что в 12 раз уменьшает объём данных. Вместо U-Net используется диффузионный трансформер (DiT) на 64 токенах, что позволяет свести все вычисления к одному ядру int8 × int8 → int32. Текстовый энкодер также собственный — крошечный трансформер на 300 тысяч параметров со словарём из 45 «слов». Применено int8-квантование весов и активаций с аккумулятором int32, что даёт не только экономию памяти, но и ускорение за счёт векторных инструкций PIE, перемножающих 16 пар int8 за инструкцию. Для ускорения генерации используется rectified flow вместо классического DDPM, что позволяет получать приемлемые результаты уже за 8 шагов вместо 50–1000. Дополнительно применён трюк с classifier-free guidance: две ветви склеиваются в один батч, чтобы веса читались из флеш-памяти один раз вместо двух.
Технические детали реализации раскрывают, как именно удалось достичь таких показателей. Автоэнкодер сжимает изображение 64×64×3 в латент 16×16×4 и обратно; энкодер используется только при обучении и содержит 1,9 млн параметров, тогда как декодер, который будет работать на ESP, имеет всего 179 тысяч параметров и состоит исключительно из свёрток 3×3, ReLU и nearest-апсемплинга без нормализаций. Для обучения декодера применяется комбинация L1-лосса и перцептивного LPIPS с небольшим KL-штрафом для гладкости латентного пространства. Денойзер — это DiT, который режет латент на патчи 2×2, получая 64 токена по 16 чисел, и обрабатывает их стандартным трансформером. Чтобы сократить число параметров, автор использовал adaLN-single: вместо отдельной проекции для каждого блока применяется одна общая проекция плюс крошечная обучаемая таблица сдвига на блок, что экономит 1,1 млн параметров. LayerNorm лишён параметров, так как масштаб и сдвиг приходят из adaLN. Текстовое условие передаётся через пулинг, а не cross-attention. Все эти оптимизации позволили уложиться в бюджет памяти и вычислительной мощности микроконтроллера.
Контекст эксперимента важен для понимания тренда: генеративные модели становятся всё более требовательными к ресурсам, и запуск даже упрощённой диффузии на микроконтроллере выглядит вызовом современным подходам. Ранее подобные задачи решались только на серверах или мощных потребительских GPU, а попытки портировать нейросети на микроконтроллеры ограничивались простыми классификаторами или генерацией небольших изображений с помощью GAN. Например, существуют проекты по запуску MobileNet на ESP32, но диффузионные модели с трансформерами на таких устройствах не встречались. Данный проект демонстрирует, что при сильном сужении домена и агрессивной оптимизации можно добиться приемлемого качества генерации на оборудовании стоимостью несколько сотен рублей. Это открывает путь к созданию автономных устройств, способных генерировать изображения без подключения к интернету, что может быть востребовано в embedded-системах, робототехнике и IoT.
Для российского рынка микроэлектроники и разработки встраиваемых систем этот эксперимент имеет особое значение. ESP32-S3 — популярный микроконтроллер, доступный по цене и широко используемый в локальных проектах, а способность запускать на нём генеративные модели может стимулировать появление отечественных решений в области edge-вычислений. Учитывая тренд на импортозамещение и ограниченный доступ к высокопроизводительным GPU, подобные оптимизации становятся не просто академическим интересом, а практической необходимостью. Реакция отрасли пока сдержанная: в комментариях к статье обсуждаются детали квантования и возможность масштабирования на другие задачи, но массового внедрения ждать рано. Тем не менее, проект показывает, что даже на слабом железе можно решать задачи, которые ещё недавно казались уделом дата-центров.
Сравнение с альтернативами подчёркивает уникальность подхода. Если рассматривать облачные сервисы генерации изображений, они обеспечивают высокое качество и разнообразие, но требуют постоянного подключения и платы за API. Локальные решения на базе Raspberry Pi или Jetson Nano справляются с диффузией, но потребляют на порядок больше энергии и стоят дороже. ESP32-S3 же демонстрирует рекордную эффективность по соотношению цена-производительность, хотя и ограничен доменом лиц и разрешением 64×64. По сравнению с GAN-ами, которые тоже могут генерировать лица на микроконтроллерах, диффузионные модели обеспечивают более стабильное качество и лучше следуют текстовому описанию. Однако достичь такого результата удалось ценой сильного упрощения архитектуры и обучения на узком датасете, что ограничивает применение.
Перспективы развития этого направления выглядят многообещающе, но остаётся ряд открытых вопросов. Сможет ли автор увеличить разрешение до 128×128 без потери скорости? Можно ли расширить домен за пределы лиц, например, генерировать простые объекты или пейзажи? Насколько сложно портировать модель на другие микроконтроллеры, например, ESP32-C3 или STM32? Также неясно, как поведёт себя модель при квантовании в int4 или при использовании разреженных вычислений. Пока же эксперимент доказывает, что генеративный ИИ может работать на устройствах с экстремально ограниченными ресурсами, и это вдохновляет на дальнейшие исследования в области эффективных архитектур и методов сжатия. В будущем подобные технологии могут найти применение в автономных камерах, игрушках и системах дополненной реальности, где важна приватность и отсутствие задержек связи.