LuxIA: A Lightweight Unitary matriX-based Framework Built on an Iterative Algorithm for Photonic Neural Network Training

TL;DR

The paper tackles the scalability bottleneck in training photonic neural networks caused by transfer-matrix computations. It introduces the Slicing Method, a windowed, sequential processing approach that enables efficient gradient-based training by localizing computations within computational windows and using a uniform cell abstraction. Integrated into LuxIA, this method achieves significant memory and time savings, outperforming existing tools across multiple meshes and standard datasets (e.g., MNIST, Digits, Olivetti) while preserving training dynamics and accuracy. The work demonstrates practical improvements in speed and scalability, facilitating broader exploration and deployment of photonic neural hardware for AI tasks.

Abstract

PNNs present promising opportunities for accelerating machine learning by leveraging the unique benefits of photonic circuits. However, current state of the art PNN simulation tools face significant scalability challenges when training large-scale PNNs, due to the computational demands of transfer matrix calculations, resulting in high memory and time consumption. To overcome these limitations, we introduce the Slicing method, an efficient transfer matrix computation approach compatible with back-propagation. We integrate this method into LuxIA, a unified simulation and training framework. The Slicing method substantially reduces memory usage and execution time, enabling scalable simulation and training of large PNNs. Experimental evaluations across various photonic architectures and standard datasets, including MNIST, Digits, and Olivetti Faces, show that LuxIA consistently surpasses existing tools in speed and scalability. Our results advance the state of the art in PNN simulation, making it feasible to explore and optimize larger, more complex architectures. By addressing key computational bottlenecks, LuxIA facilitates broader adoption and accelerates innovation in AI hardware through photonic technologies. This work paves the way for more efficient and scalable photonic neural network research and development.

Figures (13)

• Figure 1: (a) Hardware implementation on a photonic accelerator and (b) a fully connected layer of a . In (a), the input vector $\mathbf{x}$ is converted to the optical domain for computation, and the result is converted back to the electronic domain to apply the activation function, producing the output $\mathbf{y}$.
• Figure 2: A is illustrated in the gray rectangle, composed of two beam splitters and two phase shifters ($\theta$ and $\phi$). This full serves as the single building block in various mesh architectures, such as the Clements mesh Clements2016. In contrast, the dotted orange rectangle highlights the simpler single building block used in the Fldzhyan mesh Fldzhyan2020, composed of a beam splitter followed by a phase shifter.
• Figure 3: Fldzhyan mesh structure showing (a) single block, (b) mesh layers, and (c) architecture.
• Figure 4: Decomposition of the Fldzhyan mesh structure into sequential processing windows. Physical blocks ($B_{i,k}$) are mapped to computational cells ($\text{CELL}_{w,c}$) for processing.
• Figure 5: used for the benchmark experiments. The circuit consists of two sets of Clements meshes followed by a photodetector layer. The mesh size varies depending on the experiment: for Subsection \ref{['subsec:results_frameworks']}, the size is fixed to $N=64$; for Subsection \ref{['subsec:hardware_comparison']}, the size depends on the scenario. For the Batch Size-Dependent scenario, the size is fixed to $N=800$, while for the N-Dependent experiment, $N$ ranges from 100 to 900 in steps of 200.