A recent breakthrough in AI research from a Japanese lab called Sakana AI proposes a novel method to drastically reduce the memory requirements for training large language models (LLMs) by two to three times, potentially up to six times. This method, called diffusion blocks, reinterprets the training process of transformer layers as a diffusion process, where each block of layers is treated as a denoising step. By splitting a transformer into smaller blocks and training each block independently with its own noise level, the approach breaks the traditional end-to-end training constraint, significantly lowering memory usage during training without shrinking the model itself.
Traditional training of LLMs demands enormous GPU memory because the entire model, including weights, gradients, activations, and optimizer states, must be stored and updated simultaneously. For example, training a 10 billion parameter model in full precision can require around 160 GB of memory, which is far beyond the capacity of most consumer GPUs. While inference can use quantization to reduce memory, training is more sensitive and typically requires full precision, making memory savings difficult. Diffusion blocks address this by allowing training on smaller chunks sequentially or in parallel, reducing memory load and communication overhead in distributed training.
The researchers tested diffusion blocks on various models, including vision transformers for image classification, diffusion transformers for image generation, mask diffusion language models, and autoregressive language models. Results showed promising performance with minimal accuracy loss or even slight improvements in some cases, such as image generation tasks where diffusion blocks achieved better FID scores than standard training. However, these experiments were conducted on relatively small models (up to a few hundred million parameters), so the scalability to billion-parameter models remains to be proven.
An important insight from the research is that there is a trade-off between the number of blocks and their effectiveness. While more blocks mean greater memory savings, each block must have enough layers to perform effective denoising. Splitting a model into too many small blocks weakens performance. The researchers suggest that for very deep models, such as those with 128 layers, dividing into blocks with a minimum number of layers could yield substantial memory savings, potentially up to 20 times less memory usage. This also reduces communication bottlenecks in distributed training, making large-scale training more accessible and efficient.
Beyond standard transformers, diffusion blocks also enhance training for loop transformers, which reuse the same block multiple times to simulate depth. This method allows training loop transformers with a single forward pass per step, further reducing memory and computation costs. Although diffusion blocks do not reduce memory requirements for inference due to the need for the entire model to be present for causal generation, the technique holds great promise for accelerating AI research and democratizing access to training large models on more affordable hardware. The research is still in early stages, and updates on scaling this method to larger models are eagerly awaited.
Useful Links
