Efficient Training for Optical Computing

Abstract

Diffractive optical information processors have demonstrated significant promise in delivering high-speed, parallel, and energy efficient inference for scaling machine learning tasks. Training, however, remains a major computational bottleneck, compounded by large datasets and many simulations required for state-of-the-art classification models. The underlying linear transformations in such systems are inherently constrained to compositions of circulant and diagonal matrix factors, representing free-space propagation and phase and/or amplitude modulation of light, respectively. While theoretically established that an arbitrary linear transformation can be generated by such factors, only upper bounds on the number of factors exist, which are experimentally unfeasible. Additionally, physical parameters such as inter-layer distance, number of layers, and phase-only modulation further restrict the solution space. Without tractable analytical decompositions, prior works have implemented various constrained minimization techniques. As trainable elements occupy a small subset of the overall transformation, existing techniques incur unnecessary computational overhead, limiting scalability. In this work, we demonstrate significant reduction in training time by exploiting the structured and sparse nature of diffractive systems in training and inference. We introduce a novel backpropagation algorithm that incorporates plane wave decomposition via the Fourier transform, computing gradients across all trainable elements in a given layer simultaneously, using only change-of-basis and element wise multiplication. Given the lack of a closed-form mathematical decomposition for realizable optical architectures, this approach is not only valuable for machine learning tasks but broadly applicable for the generation of arbitrary linear transformations, wavefront shaping, and other signal processing tasks.