DiffusionBlocks – Block-Wise NN Training via Diffusion Interpretation

Researchers have developed DiffusionBlocks, a framework that partitions transformer neural networks into independently trainable blocks to reduce memory requirements proportionally while maintaining competitive performance. The method, detailed in a paper accepted at the 2026 International Conference on Learning Representations, enables block-wise training by interpreting the process through a diffusion model lens. The official implementation supports image classification with Vision Transformers and is available for public use.

We propose , a principled framework that partitions transformers into independently trainable blocks, reducing memory requirements proportionally while maintaining competitive performance across diverse architectures and tasks.DiffusionBlocks This is an official implementation of DiffusionBlocks on image classification using Vision Transformers ViT . Please install uv https://docs.astral.sh/uv/getting-started/installation/ . Then, run: Install dependencies uv sync make sure to login huggingface and wandb uv run huggingface-cli login uv run wandb login We conducted our experiments in the following environment: Python Version 3.12 and CUDA Version 12.2 H100. The model checkpoints are saved in logs folder. Baseline ViT : uv run main.py train cifar100 --model type vit DiffusionBlocks: uv run main.py train cifar100 --model type dblock NOTE: the total epochs in DiffusionBlocks is multiplied by the number of blocks to align the total number of iterations with the baseline as one step in DiffusionBlocks corresponds to training for one block. Details In the base setting, we don't reply on techniques such as heavy data augmentation. In case you want to see the performance with heavy data augmentation and learning rate scheduler, run as follows: Baseline ViT : BATCH SIZE=128 EPOCHS=1000 POSTFIX="-rand-augment" WARMUP STEPS=3900 MODEL TYPE="dblock" srun uv run main.py train cifar100 \ --model type $MODEL TYPE \ --batch size $BATCH SIZE --num epochs $EPOCHS --postfix=$POSTFIX \ --scheduler type cosine with min lr --num warmup steps $WARMUP STEPS --lr 5e-4 \ --scheduler specific kwargs '{"min lr": 5e-5}' \ --add rand aug DiffusionBlocks: BATCH SIZE=128 EPOCHS=1000 POSTFIX="-rand-augment" WARMUP STEPS=$ 3900 3 3 indicates the number of blocks MODEL TYPE="dblock" srun uv run main.py train cifar100 \ --model type $MODEL TYPE \ --batch size $BATCH SIZE --num epochs $EPOCHS --postfix=$POSTFIX \ --scheduler type cosine with min lr --num warmup steps $WARMUP STEPS --lr 5e-4 \ --scheduler specific kwargs '{"min lr": 5e-5}' \ --add rand aug Baseline ViT : CKPT PATH="logs/path-to-last.ckpt" uv run main.py test cifar100 --model type vit --ckpt path $CKPT DiffusionBlocks: CKPT PATH="logs/path-to-last.ckpt" uv run main.py test cifar100 --model type dblock --ckpt path $CKPT The implementation of Vision Transformer in vit.py /SakanaAI/DiffusionBlocks/blob/main/vit.py is based on HuggingFace Transformers https://github.com/huggingface/transformers . And, the implementation of EDM is based on Stability-AI/generative-models https://github.com/Stability-AI/generative-models . We are grateful for their work. To cite our work, please use the following BibTeX: @inproceedings{shing2026diffusionblocks, title = {DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation}, author. = {Makoto Shing and Masanori Koyama and Takuya Akiba}, booktitle = {The Fourteenth International Conference on Learning Representations}, year = {2026}, url = {https://openreview.net/forum?id=pwVSmK71cS} }