Sakana AI представила DiffusionBlocks: нейросети теперь можно обучать по блокам и забыть про дикие требования к памяти
Команда Sakana AI выкатила на ICLR 2026 работу, которая решает одну из самых больных проблем в обучении больших моделей: память. Чем глубже сеть, тем больше GPU нужно держать под полный форвард и беквард, и это растёт линейно по глубине. Отсюда вечные OOM, градиент чекпоинтинг, шардинг и прочие костыли.
Идея DiffusionBlocks простая и при этом очень красивая. Сеть режут на блоки и обучают каждый блок отдельно. В памяти лежит только один блок, а не весь граф вычислений. Но фокус не в нарезке, а в том, как именно авторы задают каждому блоку цель.
Каждому блоку явно присваивают роль: слегка довести представление до цели, ближе, чем это сделал предыдущий. Это поведение один в один повторяет то, что делает диффузионная модель на очередном шаге. Благодаря этой аналогии каждый блок оптимизирует свой локальный объектив и обучается независимо от остальных.
Авторы проверили метод на пяти разных архитектурах: ViT, DiT, masked diffusion, autoregressive transformers и recurrent‑depth transformers. Во всех случаях качество остаётся на уровне обычного end‑to‑end обучения, при этом памяти нужна в несколько раз меньше. На картинке к анонсу Sakana рисует 8x сокращение памяти.
Отдельный бонус ждёт тех, кто работает с looped или recurrent‑depth трансформерами. Обычно их учат через дорогую backpropagation through time, то есть разворачивают все итерации и ведут градиент по всей этой цепочке. В DiffusionBlocks всё это сворачивается в один форвардный проход во время обучения. Это прямой удар по вычислительной стоимости BPTT и по объёму активаций, которые нужно хранить.
Что это даёт на практике. Можно обучать более глубокие модели на том же железе, легче распараллеливать обучение, блоки реально подходят под независимый тренинг на разных узлах и можно экспериментировать с отдельными блоками, не ломая всю сеть. Для исследователей это похоже на быстрый обход того ограничения, которое долго считали фундаментальным.
Источники и ссылки и разбор