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

Abstract

Photonic Neural Networks (PNNs) can accelerate machine learning workloads by implementing Matrix-Vector Multiplications (MVMs) in integrated photonic circuits, but existing simulation and training frameworks scale poorly to large Photonic Unitary Matrix (PUM) meshes because they explicitly construct and manipulate dense transfer matrices. This work introduces the Slicing method, which models PUM meshes as sequences of local 2×2 operations organized into computational windows and computes forward and backward propagation using only localized matrix–vector updates with linear complexity in the number of active cells. The method is implemented in LuxIA, an open-source PyTorch-based framework for end-to-end Photonic Neural Network (PNN) simulation and training that supports multiple mesh architectures and datasets. A formal analysis shows that Slicing reduces per-pass work by one degree (from quartic to cubic) in the worst-case. Experiments on Clements, Fldzhyan, and Bell-optimized meshes trained on Iris, Digits, MNIST, and Olivetti Faces show that LuxIA matches the training dynamics and task accuracy of existing tools while substantially improving training efficiency: on large meshes and batches, LuxIA achieves up to 4.7× lower training time and more than an order-of-magnitude reduction in Graphics Processing Unit (GPU) memory compared with conventional transfer-matrix frameworks, and it remains within the memory budget where competing tools fail.