Sakana AI's DiffusionBlocks trains one block at a time, claiming 1/B memory with end-to-end parity

ICLR 2026 work recasts block-wise updates as reverse diffusion, reporting comparable results in vision, image generation, and language while storing activations for a single block.

By · Published

Why it matters

Training is increasingly limited by memory, not just compute. A block-wise method that avoids keeping the full network in memory could let teams train bigger models on fewer or smaller GPUs.

Sakana AI's DiffusionBlocks trains one block at a time, claiming 1/B memory with end-to-end parity — ICLR 2026 work recasts block-wise updates as reverse diffusion, reporting comparable results in vision, image generation, and language whil

Sakana AI researchers Makoto Shing and Takuya Akiba, with Masanori Koyama of the University of Tokyo, introduced DiffusionBlocks, a block-wise neural network training framework that reinterprets each block's update as the reverse step of a diffusion process. The team announced the work in an X post and a project page, with the paper on arXiv and OpenReview. The work was presented at ICLR 2026.

https://x.com/SakanaAILabs/status/2059648778051924281

DiffusionBlocks partitions a residual network into B blocks and trains each block independently, storing activations for only one block at a time. The authors say this reduces training memory by a factor of B while matching the performance of end-to-end backpropagation across architectures.

In experiments spanning image classification, image generation, and text generation, the paper reports results comparable to end-to-end training while using a fraction of the memory. The method also extends to recurrent-depth models, replacing K-iteration backpropagation through time with a single forward pass during training while retaining the original K iterations at inference.

If these results hold more broadly, DiffusionBlocks could ease a central training bottleneck by decoupling memory growth from model depth, enabling larger models or more modest hardware to be used during training.

What it could change for builders and capital allocators:

  • Training budgets: If parity holds, the reported 1/B activation memory profile can translate into fewer high-memory GPUs for a given model or deeper models on the same hardware. That lowers the barrier to custom pretraining and long-context fine-tuning for smaller teams.
  • Hardware strategy: Memory relief may expand the viable set of accelerators and instances. Teams could pilot deeper models on lower-memory cards during training, while keeping inference unchanged.
  • Throughput trade-offs: The headline benefit is memory. Founders should measure wall-clock time, throughput, and convergence under their workloads, since the paper centers on memory rather than speed.
  • Scaling risk: Diligence should focus on external replication, stability at larger scales, optimizer compatibility, and how block-wise updates interact with data/model parallelism in real training pipelines.
  • Iterative architectures: The recurrent-depth extension suggests potential savings where K is large, with a single forward pass during training and K iterations retained at inference. That could simplify training loops for iterative models if results generalize.
  • Tooling and ecosystem: Watch for framework integrations and vendor support. The teams that productize scheduler tooling, checkpoints, and distributed training support for block-wise updates could see early demand if the method proves robust.

Near-term actions: run side-by-side baselines on your models and data, track validation parity and activation memory, and document any throughput or stability deltas. For investors, diligence teams integrating this into training stacks or cloud offerings, with an eye to reproducibility and customer workloads rather than benchmarks alone.

Reader comments

Conversation for this story loads after sign-in.