Optical training of large-scale Transformers and deep neural networks with direct feedback alignment

Ziao Wang,Kilian Müller,Matthew Filipovich,Julien Launay,Ruben Ohana,Gustave Pariente,Safa Mokaadi,Charles Brossollet,Fabien Moreau,Alessandro Cappelli,Iacopo Poli,Igor Carron,Laurent Daudet,Florent Krzakala,Sylvain Gigan

2024-09-01

Abstract:Modern machine learning relies nearly exclusively on dedicated electronic hardware accelerators. Photonic approaches, with low consumption and high operation speed, are increasingly considered for inference but, to date, remain mostly limited to relatively basic tasks. Simultaneously, the problem of training deep and complex neural networks, overwhelmingly performed through backpropagation, remains a significant limitation to the size and, consequently, the performance of current architectures and a major compute and energy bottleneck. Here, we experimentally implement a versatile and scalable training algorithm, called direct feedback alignment, on a hybrid electronic-photonic platform. An optical processing unit performs large-scale random matrix multiplications, which is the central operation of this algorithm, at speeds up to 1500 TeraOps. We perform optical training of one of the most recent deep learning architectures, including Transformers, with more than 1B parameters, and obtain good performances on both language and vision tasks. We study the compute scaling of our hybrid optical approach, and demonstrate a potential advantage for ultra-deep and wide neural networks, thus opening a promising route to sustain the exponential growth of modern artificial intelligence beyond traditional von Neumann approaches.

Emerging Technologies,Disordered Systems and Neural Networks,Machine Learning,Applied Physics,Optics

What problem does this paper attempt to address?

The paper attempts to address the following key issues:

1. **Application of Photonic Computing in Large-Scale Neural Network Training**: Traditional electronic hardware (such as CPUs and GPUs) faces bottlenecks when handling large-scale random matrix multiplication, especially in training large neural networks. This paper proposes a method based on optical processing units (OPUs) that can efficiently perform large-scale random matrix multiplication, thereby accelerating the training process.

2. **Application of Direct Feedback Alignment (DFA) Algorithm**: Compared to traditional backpropagation (BP), the DFA algorithm achieves parallel updates by directly projecting errors to each layer, reducing training time. This paper demonstrates how to apply the DFA algorithm to a photonic computing platform and validates its effectiveness on large-scale models.

3. **Training of Large-Scale Models**: This paper is the first to demonstrate the capability of using a photonic computing platform combined with the DFA algorithm to train Transformer models and fully connected neural networks with over 1 billion parameters, proving that this method achieves good performance on language and vision tasks.

4. **Scalability Study of Ultra-Large-Scale Models**: The study investigates the time complexity of training neural networks with different dimensions and layers using the DFA algorithm combined with photonic hardware. It finds that for extremely large-scale models, the optical training method has significant speed advantages.

In summary, this paper aims to present a new method that combines novel hardware (optical processors) and advanced algorithms (DFA) to overcome the limitations of existing electronic hardware in training large-scale neural networks, particularly in terms of speed and energy consumption.