DiffusionBlocks: Treino em Blocos para Modelos Gerativos via Difusão Baseada em Pontuação

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

June 17, 2025

Autores: Makoto Shing, Takuya Akiba

cs.AI

Resumo

O treinamento de grandes redes neurais com retropropagação de ponta a ponta cria gargalos significativos de memória, limitando o acesso à pesquisa de ponta em IA. Propomos o DiffusionBlocks, uma nova estrutura de treinamento que interpreta os blocos de redes neurais como operações de remoção de ruído em um processo de difusão em tempo contínuo. Ao particionar a rede em blocos treináveis de forma independente e otimizar as atribuições de níveis de ruído com base na massa de probabilidade cumulativa igual, nossa abordagem alcança uma eficiência de memória significativa, mantendo um desempenho competitivo em comparação com a retropropagação tradicional em tarefas generativas. Experimentos em geração de imagens e modelagem de linguagem demonstram uma redução de memória proporcional ao número de blocos, ao mesmo tempo em que alcançam um desempenho superior. O DiffusionBlocks oferece um caminho promissor para democratizar o acesso ao treinamento de redes neurais em grande escala com recursos computacionais limitados.

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.