Layer-Parallel Training for Transformers

2026-01-13

AI Summary

This work addresses the limited parallelizability of conventional Transformers as depth increases, which hinders training efficiency for large-scale models. The authors introduce, for the first time, a multilayer parallel-in-time algorithm into Transformer training by modeling the network as a neural ordinary differential equation (neural ODE), enabling cross-layer parallelism in both forward and backward passes. To ensure stable convergence while maximizing computational efficiency, they further propose an error-monitoring mechanism with adaptive switching between serial and parallel execution modes. Experiments on BERT, GPT-2, Vision Transformers (ViT), and machine translation architectures demonstrate that the method significantly enhances parallel scalability and training speed for deep models without compromising pretraining or fine-tuning accuracy.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

Abstract

We present a new training methodology for transformers using a multilevel, layer-parallel approach. Through a neural ODE formulation of transformers, our application of a multilevel parallel-in-time algorithm for the forward and backpropagation phases of training achieves parallel acceleration over the layer dimension. This dramatically enhances parallel scalability as the network depth increases, which is particularly useful for increasingly large foundational models. However, achieving this introduces errors that cause systematic bias in the gradients, which in turn reduces convergence when closer to the minima. We develop an algorithm to detect this critical transition and either switch to serial training or systematically increase the accuracy of layer-parallel training. Results, including BERT, GPT2, ViT, and machine translation architectures, demonstrate parallel-acceleration as well as accuracy commensurate with serial pre-training while fine-tuning is unaffected.

Problem

• layer-parallel training
• gradient bias
• convergence degradation
• Transformer
• parallel-in-time

Innovation

• layer-parallel training
• neural ODE
• parallel-in-time
• Transformer
• multilevel algorithm