DiffusionBlocks: 스코어 기반 확산을 통한 생성 모델의 블록 단위 학습

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

June 17, 2025

저자: Makoto Shing, Takuya Akiba

cs.AI

초록

종단 간 역전파를 사용하여 대규모 신경망을 학습시키는 것은 상당한 메모리 병목 현상을 초래하여 최신 AI 연구에 대한 접근성을 제한합니다. 우리는 DiffusionBlocks라는 새로운 학습 프레임워크를 제안합니다. 이 프레임워크는 신경망 블록을 연속 시간 확산 과정에서의 잡음 제거 작업으로 해석합니다. 네트워크를 독립적으로 학습 가능한 블록으로 분할하고 동일한 누적 확률 질량을 기반으로 잡음 수준 할당을 최적화함으로써, 우리의 접근 방식은 생성 작업에서 전통적인 역전파와 비교하여 경쟁력 있는 성능을 유지하면서도 상당한 메모리 효율성을 달성합니다. 이미지 생성 및 언어 모델링 작업에 대한 실험은 블록 수에 비례하여 메모리 사용량을 줄이면서도 우수한 성능을 달성함을 보여줍니다. DiffusionBlocks는 제한된 계산 자원으로 대규모 신경망 학습에 대한 접근을 민주화할 수 있는 유망한 경로를 제공합니다.

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.