DiffusionBlocks: Blockweises Training für generative Modelle mittels score-basierter Diffusion

DiffusionBlocks: Blockwise Training for Generative Models via Score-Based Diffusion

June 17, 2025

Autoren: Makoto Shing, Takuya Akiba

cs.AI

Zusammenfassung

Das Training großer neuronaler Netze mit end-to-end Backpropagation führt zu erheblichen Speicherengpässen, was den Zugang zu modernster KI-Forschung einschränkt. Wir schlagen DiffusionBlocks vor, ein neuartiges Trainingsframework, das Blöcke neuronaler Netze als Denoising-Operationen in einem zeitkontinuierlichen Diffusionsprozess interpretiert. Durch die Aufteilung des Netzwerks in unabhängig trainierbare Blöcke und die Optimierung der Rauschpegelzuweisungen basierend auf gleicher kumulativer Wahrscheinlichkeitsmasse erreicht unser Ansatz eine signifikante Speichereffizienz bei gleichzeitiger Beibehaltung einer wettbewerbsfähigen Leistung im Vergleich zur traditionellen Backpropagation bei generativen Aufgaben. Experimente zur Bildgenerierung und Sprachmodellierung zeigen eine speicherreduzierende Wirkung, die proportional zur Anzahl der Blöcke ist, bei gleichzeitiger Erzielung überlegener Leistung. DiffusionBlocks bietet einen vielversprechenden Weg, um den Zugang zum Training großer neuronaler Netze mit begrenzten Rechenressourcen zu demokratisieren.

English

Training large neural networks with end-to-end backpropagation creates significant memory bottlenecks, limiting accessibility to state-of-the-art AI research. We propose DiffusionBlocks, a novel training framework that interprets neural network blocks as performing denoising operations in a continuous-time diffusion process. By partitioning the network into independently trainable blocks and optimizing noise level assignments based on equal cumulative probability mass, our approach achieves significant memory efficiency while maintaining competitive performance compared to traditional backpropagation in generative tasks. Experiments on image generation and language modeling tasks demonstrate memory reduction proportional to the number of blocks while achieving superior performance. DiffusionBlocks provides a promising pathway for democratizing access to large-scale neural network training with limited computational resources.