LightRidge: An End-to-end Agile Design Framework for Diffractive Optical
  Neural Networks
To lower the barrier to diffractive optical neural networks (DONNs) design,
exploration, and deployment, we propose LightRidge, the first end-to-end
optical ML compilation framework, which consists of (1) precise and
differentiable optical physics kernels that enable complete explorations of
DONNs architectures, (2) optical physics computation kernel acceleration that
significantly reduces the runtime cost in training, emulation, and deployment
of DONNs, and (3) versatile and flexible optical system modeling and
user-friendly domain-specific-language (DSL). As a result, LightRidge framework
enables efficient end-to-end design and deployment of DONNs, and significantly
reduces the efforts for programming, hardware-software codesign, and chip
integration. Our results are experimentally conducted with physical optical
systems, where we demonstrate: (1) the optical physics kernels precisely
correlated to low-level physics and systems, (2) significant speedups in
runtime with physics-aware emulation workloads compared to the state-of-the-art
commercial system, (3) effective architectural design space exploration
verified by the hardware prototype and on-chip integration case study, and (4)
novel DONN design principles including successful demonstrations of advanced
image classification and image segmentation task using DONNs architecture and
topology.
Yingjie Li, Ruiyang Chen, Minhan Lou, Berardi Sensale-Rodriguez, Weilu Gao, Cunxi Yu
# LightRidge: An End-to-end Agile Design Framework for

###### Abstract

To lower the barrier to diffractive optical neural networks (DONNs) design, exploration, and deployment, we propose **LightRidge**, the first end-to-end optical ML compilation framework, which consists of **(1)** precise and differentiable optical physics kernels that enable complete explorations of DONNs architectures, **(2)** optical physics computation kernel acceleration that significantly reduces the runtime cost in training, emulation, and deployment of DONNs, and **(3)** versatile and flexible optical system modeling and user-friendly domain-specific-language (DSL). As a result, LightRidge framework enables efficient end-to-end design and deployment of DONNs, and significantly reduces the efforts for programming, hardware-software codesign, and chip integration. Our results are experimentally conducted with physical optical systems, where we demonstrate: **(1)** the optical physics kernels precisely correlated to low-level physics and systems, **(2)** significant speedups in runtime with physics-aware emulation workloads compared to the state-of-the-art commercial system, **(3)** effective architectural design space exploration verified by the hardware prototype and on-chip integration case study, and **(4)** novel DONN design principles including successful demonstrations of advanced image classification and image segmentation task using DONNs architecture and topology.1

Footnote 1: To appear at ASPLOS 2024.

## 1 Introduction

Deep neural networks (DNNs) has been undergoing significant growth in recent years with significant contributions in many application domains like autonomous systems, natural language processing, and health care[3, 11, 30, 51, 15, 16, 2]. However, large DNN models which produce high system throughput, usually suffers from high carbon footprint. For example, recent studies estimated 626,000 pounds of planet-warming carbon dioxide, equal to the lifetime emissions of five cars, produced in training Transformer network [50, 47]. On the other side, the embedded accelerators [49, 4, 57, 52, 46, 58], which are designed to improve resource and power efficiency, suffer from limited functionality and throughput. Thus, while there have recently seen great progress in customized accelerators that adjust the computing performance with efficiency in hardware architectures and systems, the Pareto-frontier of conventional accelerators remains the same [12, 30, 43, 19, 26, 45].

To advance the Pareto-frontier of ML systems, i.e., offering high computing performance as well as extreme power efficiency, accelerators taking advantage of optics, namely _optical neural networks_ (ONNs), have recently attracted significant interest in machine learning and hardware acceleration. The main advantages of ONNs over digital accelerators can be summarized as follows - **(1)** In optical computing systems, since the input features are encoded and carried by light, the computation and data movement will happen at the speed of light in the medium with orders of magnitude advantages in computation speed [22, 48, 23, 60, 18, 53]. **(2)** The laser implemented in the optical systems can be easily expanded with passive optical devices, such as beam splitters, to multiple channels, which means parallel computation can be easily realized with ONN systems, and the throughput of the system will be significantly increased [34, 38, 67, 33]. **(3)** The trained ONN system will be deployed with passive optical devices, which means there is no additional energy cost for all-optical inference process, thus improving the energy efficiency significantly [48, 64, 22, 17, 32, 42, 61, 66, 6]. _Diffractive Optical Neural Networks_ (DONNs) is one of the most promising research areas in ONNs, which mimic the propagation and connectivity properties of conventional neural networks, by utilizing the nature physics of light diffraction and phase modulation of coherent light [34, 32, 38, 44, 67, 7, 6]. Even though the inference of the physical DONN is all optical, the training part that leads to its design is done through digital platforms, where a precise, efficient and hardware-aware emulation engine is required.

The existing optical emulation engines, such as Mathworks BeamLab [56] and LightPipes [55], mainly focus on the emulation of the physical phenomenon while lack the key functionalities and domain-specific runtime optimizations in supporting the developments of DONNs. Specifically, it is particularly challenging for existing optical emulation frameworks to deal with DONN training and inference due to the following reasons: (1) The core emulation functions are not differentiable, which makes the backpropagation-based training hard to implement. (2) The implementation is not optimized in runtime. For example, LightPipes does not support tensor representations and operator fusion, which significantly limits the runtime performance (see Table 1). (3) There does not exist hardware/device aware emulation supports, which require significant extra efforts for correlating numerical emulations and physical deployments.
As a result, there are several critical technical barriers in design, training, exploration, and hardware deployment of DONNs, summarized as follows:

**Challenge 1:** Sufficient multi-disciplinary domain-knowledge in optical physics, fabrication, and machine learning (ML) are required for DONN system design and deployment, which puts a critical technical barrier to exploring and advancing DONN systems in real-world applications. At this point, there does not exist an end-to-end design framework that supports design and exploration for full-stack DONNs design, optimization, fabrication, and on-chip integration. Moreover, the broad architectural search space with software, optics, and fabrication hyperparameters can be an obstacle for efficient design space exploration (DSE), which also motivates the development of an end-to-end design framework.

**Challenge 2:** There have observed significant performance degradation when deploying the trained DONN model to the practical hardware, namely, there is an algorithm-hardware miscorrelation gap between the numerical modeling and the physical system. The miscorrelation gap can come from two aspects: **(1)** The imprecise numerical modeling of the DONN system, i.e., lack of the precisely implemented physics emulation intermediate representation (IR). Classic numerical models for fundamental physics kernels in DONNs such as _Finite-difference time-domain_ (FDTD) and _scalar diffraction_ modeling via _Fast Fourier Transform_ (FFT), are both verified to be sufficiently precise in the DONN system emulation [37]; **(2)** Lack of domain-specific hardware-software codesign algorithms to realize quantization-aware hardware deployment and deal with the intrinsic noise (such as fabrication variations, non-unify optical response, etc.) in optical devices. These challenges have been confirmed by Zhou et al.[67] in Figure 1, who report \(\geq\) 30% accuracy degradation while deploying the model to the physical optical system.

**Challenge 3:** Training and emulation of DONN system are challenging due to high computational cost in modeling the optical physics. For example, [34, 67, 39] reported that training 5-layer DONNs for MNIST-10 with 5 epochs takes 3-4 days (Figure 1). Besides, existing optics simulation frameworks lack runtime optimization in developing the physics kernels, nor domain-specific language (DSL) supports. Table 1 summarizes the limitations of existing frameworks for DONNs design. More importantly, the choice of numerical physics modeling has great impacts in runtime efficiency, while it is required to offer high fidelity to the hardware deployment and fabrication.

Thus, we propose _LightRidge_, an agile end-to-end framework, aiming to significantly lower the barriers to design, training, design space exploration, and hardware deployment of DONN systems. In particular, _LightRidge_ is implemented with high-performance, precise, and versatile optical physics kernels, which precisely correlate to experimental physical systems, enabling out-of-box software-to-hardware realization in an end-to-end fashion, and showing its capabilities to explore advanced DONN architectures for complex ML applications. The contributions of this paper are summarized as follows:

* We propose a novel agile physics-aware design framework _LightRidge_ for end-to-end design, exploration, and deployment for DONNs, consisting of versatile and optimized physics modeling kernels and hardware-software codesign algorithms that enable efficient and precise DONNs modeling w.r.t real-world hardware systems (Section 3).
* We propose LightRidge-DSE to accelerate the end-to-end design cycle for DONNs design, exploration, and on-chip integration, verified by our physical prototype and on-chip integration case study. Moreover, LightRidge-DSE confirms critical domain-knowledge insights [5] for designing an efficient DONN system in physics meanings (Section 4).
* We experimentally validate the effectiveness and precision of LightRidge in designing practical DONN systems and on-chip integration, via visible-range DONN prototype and end-to-end on-chip integration case study (Section 5.1-5.5).
* Furthermore, two novel advanced DONN architecture principles are developed via LightRidge to advance DONNs in complex image classification tasks, and first-ever all-optical image segmentation (Section 5.6).
* Finally, LightRidge will be released as an open-source hard

\begin{table}
\begin{tabular}{c|c|c|c|c|c} \hline \hline  & \begin{tabular}{c} Optics \\ kernels \\ \end{tabular} & DSE & \begin{tabular}{c} LoC \\ (val) \\ \end{tabular} & \begin{tabular}{c} LoC \\ (train) \\ \end{tabular} & 
\begin{tabular}{c} Runtime \\ (pre-fab) \\ \end{tabular} \\ \hline LightRidge & ✓ & ✓ & 1\(\times\) & 1\(\times\) & mins – hrs \\ \hline LightPipes[55] & ✓ & ✗ & 2\(\times\) & n/a & days \\ \hline Customized & & & & & \\ PyTorch/TF & ✗ & ✗ & 20\(\times\) & 50\(\times\) & days \\
[34, 67, 39] & & & & & & \\ \hline \hline \end{tabular}
\end{table}
Table 1: Overview comparisons of existing programming frameworks for DONNs compilation. Lines of Code (LoC) efforts are evaluated with a 5-layer DONNs [34].

Figure 1: Model performance and time-to-deployment runtime comparison between manual hardware calibration and LightRidge – (1) LightRidge reduces design cycle from days to hours with high-performance emulation kernels and DSE engine; (2) LightRidge results in significantly improved correlation in out-of-box deployment, which gets rid of expensive manual calibration processes.
ware project.2

Footnote 2: [https://lightridge.github.io/lightridge](https://lightridge.github.io/lightridge).

## 2 Diffractive Optical Neural Networks

Compared to conventional neural networks (NNs) on digital platforms, the information carrier changes from electrons to photons in DONN systems, i.e., instead of manipulating electrons between transistors to realize the computation, in DONN system, the computation is realized by manipulating the information-carried light with its physical features. Specifically, the DONN system is composed by multiple diffractive layers stacking in sequence as shown in Figure 2(a), which embed the phase modulations trained w.r.t the ML task for manipulating and encoding information on the light signal. The connection between layers is realized by the light diffraction when the light signal propagates between layers. Thus, in DONN systems, light diffraction can be considered as "neural operators" and phase change patterns can be seen as "weights", when compared with conventional NNs. However, the light-signal based DONN system requires the analog-to-digital converter to read out the prediction results. Thus, a detector is employed at the end of the system to capture the light intensity pattern for analysing and predictions. As a result, DONN systems take advantage of the light signal to encode and propagate information and its physical nature to realize the computation. Since the physical phenomenon happens by nature with light propagation, the computation happens with no extra energy cost at light speed, while the practical computation efficiency of the DONN system will be determined by the analog-to-digital conversions.

This section presents the overview of DONNs, including emulation, training, and the hardware deployment of DONN systems. To get an effective DONN model w.r.t a specific ML task, the propagation process of the light signal is emulated and the model is trained based on the emulation on digital platforms, where a precise mathematical approximation for the optical phenomenon, i.e., light diffraction and phase modulation, is required, which is illustrated in detail in Section 3.1. Each point at a given diffractive layer acts as a secondary source of the input light wave in accordance with the Huygens-Fresnel principle. The phase of the input wave is determined by the product of the input wave and the complex-valued phase modulation at that point. The diffraction space is required to generate the diffraction pattern at the receive plane. The phase modulation at each point w.r.t its location at the layer is the learnable parameter iteratively adjusted during the training process with error back-propagation method [34, 67]. The physical kernel implemented in LightRidge for DONN emulation and training is constructed with the widely used, precise and efficient mathematical approximations for scalar diffraction formulas. Finally, the trained model can be physically deployed with optical devices as shown in Figure 2(b) or on-chip integration system as shown in Figure 11, to realize the fully optical inference with few energy cost, high computation speed and high system throughput.

### DONN Emulation and Training

Enabling the precise hardware-software codesign aware emulation of the physical phenomenon happening in DONN systems, i.e., input encoding, light diffraction, phase modulation, and detector reading, is critical for the practical realization of DONN systems. There are mainly two mathematical methods for formulating light diffraction: (1) _Finite-difference time-domain_ (FDTD) method [63], which performs the full-vector differentiable numerical simulation of photonic structures by solving Maxwell's equations directly without physical approximations. It is a sophisticated and powerful method for light propagation emulation, however suffering from heavy computation efforts and heavy data dependency that prevent parallelisms in kernel developments. Specifically, FDTD requires the entire computational domain to be sufficiently fine gridded, which means the DONN system size will be expanded exponentially in the FDTD-based emulation. Since DONN systems target large-scale machine learning tasks, the FDTD-based emulation is infeasible in computation runtime and memory for DONN systems due to the system scalability. (2) _Fast Fourier Transform_ (FFT) method [21], which performs mathematical approximation based on scalar diffraction theory. It simplifies the computation with scenario-specific approximations while keeping the emulation precision. We have three widely used approximations for light diffraction in different application scenarios, i.e., _Rayleigh-Sommerfield_ approximation, _Fresnel_ approximation, and _Fraunhofer_ approximation. While both FDTD and FFT-based approximations are differentiable, FFT-based scalar diffraction modeling is more capable for large-scale DONNs emulation without size expansion requirements for fine gridded computational domain. More importantly, [37] and our physical experiments in Section 5 verify that the FFT-based approximations are sufficiently precise to close the codesign gap for the DONN system emulation. **Therefore, we

Figure 2: Overview of DONN system and hardware implementation – (a) Illustration of the DONN system, including the input plane, three diffractive layers, and a light intensity readout plane. (b) The reconfigurable optical hardware system to deploy the DONN system.
implement the FFT-based physics kernel in LightRidge as IR to provide precise and efficient DONN emulation and training (Section 3.1.1).** The phase modulation is applied to the input light wave by complex-valued matrix multiplication as illustrated in Section 3.1.2.

In our framework, the FFT based mathematical approximation for light diffraction is design to be fully differentiable from detector to laser input w.r.t the loss function acquired from the diffraction pattern captured at the detector. Specifically, during the training process, the prediction is generated according to the intensity of the diffraction pattern shown on the detector with pre-defined detector regions for different classes, where the light intensity \(I\) collected by each detector region mimics the probability of output prediction after softmax in conventional DNNs. Thus, the class whose corresponding detector region collects the highest light intensity is selected as the final prediction. With the one-hot represented ground truth class \(t\), the loss function \(L\) can be acquired with **MSELoss** as \(L=\parallel\text{Softmax}(I)-t\parallel_{2}\). Thus, the whole system is designed to be differentiable and compatible with conventional automatic differential engines.

### Hardware Deployment

The materials to deploy the trained DONN model on physical hardware need to be carefully selected, as optical devices made from different materials can have significantly different response to different laser wavelength. For example, SLMs can function as diffractive layers in the DONN system with laser wavelength in visible range; while for systems with laser wavelength in Terahertz (THz) range, a 3D printed mask with designed thickness at each pixel made with UV-curable resin is used as the diffractive layers in THz optical systems [34].

In our experimental hardware systems shown in Figure 2(b), the wavelength of the laser source is 532nm, which is in the operating range of the SLM3. Specifically, the SLM is an array of twisted nematic liquid crystal, where each pixel (liquid crystal) can be independently twisted to different angles by different applied control voltages, providing different phase modulation for the input light beam. However, such analog optical devices hardly have unified optical response to the control and can vary from each single due to fabrication errors, worsening the correlation between the numerical emulations and the hardware deployment, which highlights the importance to design precise computation kernels for emulation and hardware-software codesign algorithms for DONN systems.

Footnote 3: [https://holoeye.com/lc-2012-spatial-light-modulator/](https://holoeye.com/lc-2012-spatial-light-modulator/)

## 3 LightRidge Framework

Figure 3 shows the end-to-end design flow of DONN systems with automation provided by LightRidge. With the user-defined design specification and the targeted ML task,  the architectural and fabrication parameters such as diffraction distance, diffraction unit size, chip dimensions, etc., are selected and produced automatically by conducting fast and efficient design space exploration (DSE) with the emulation model in LightRidge, which circumvents the critical domain knowledge requirements for designing a functioning DONN model (Section 4). This exploration is enabled with our accelerated and precise computation engine, improving the runtime efficiency significantly.  When satisfying hyperparameters are acquired from the fast DSE, the emulation model will be updated with the hardware information for physical deployment, e.g., the optical response curve for SLMs w.r.t the control voltages, which is the co-design training with hardware-aware optimization.  Optical devices for practical deployment are fabricated/set w.r.t the parameters in the trained model, i.e., the phase modulation in diffractive layers. The device fabrication information is dumped and generated automatically with LightRidge.  With all components ready for deployment, a targeted all-optical DONN system can be setup for efficient and energy-saving optical inference (Section 5).  The LightRidge automation processes are all efficiently realized by the user-friendly DSL support in LightRidge.

In this section, we illustrate the physics kernel with mathematical approximation modelling for DONN systems implemented in LightRidge, and introduce the front-end DSL designed for the LightRidge compilation implementations. We also propose a novel complex-valued regularization algorithm to improve training performance.

### Physic Kernel Implementation

The DONN system functions as a neural network based on two physical phenomena, i.e., light diffraction and phase modulation. In our framework, we take FFT-based scalar diffraction theory to build our modelling kernels.

First, the continuous-wave (CW) laser source is deployed to encode the input information. The light wave is described with complex-valued numbers in physics with two properties, amplitude and phase of the wave, i.e., \(E(t)=Ae^{j\theta}\), where \(j=\sqrt{-1}\), \(A\) is the amplitude, and \(\theta\) is the phase. The input information is encoded with the intensity \(I\) of the light wave with phase initialized as 0, i.e, \(\theta=0\), \(A=I\). Then, as shown in Figure 4, the information-carried light wave is diffracted over the diffraction distance \(z\), which is emulated with mathematical diffraction approximations described in Section 3.1.1. At the diffractive layer, each diffraction unit embeds a phase modulator, where the trainable parameter, phase modulation is applied to the light signal as described in Section 3.1.2. The forward function for a multiple-layer constructed DONN system calculates diffraction and phase modulation iteratively through all stacked diffractive layers. Finally, the diffraction pattern is captured and converted to light intensity at the detector for analog-to-digital conversion for computer processing.

#### 3.1.1 Light Diffraction approximation

There are typically three approximation methods for scalar theory of diffraction, i.e, _Rayleigh-Sommerfeld approximation_, _Fresnel approximation_, and _Fraunhofer approximation_. They work under specific
application scenarios with different assumptions of the system, such as aperture size and propagation distance.

The Rayleigh-Sommerfeld is the most commonly used approximation as it requires least physical approximations on the system and is reported to give quite accurate results. The Rayleigh-Sommerfeld approximation is implemented with Equation 1 in our framework. As shown in Figure 4(a), the diffracting aperture is in the \((\xi,\eta)\) plane, and is illuminated in the positive \(z\) direction. We calculate the wavefield across the \((x,y)\) plane, which is parallel to the \((\xi,\eta)\) plane and at distance \(z\) from it. The \(z\) axis pierces both planes at their origins. Then, when \(r_{01}\gg\lambda\), the _Rayleigh-Sommerfeld approximation_ will be described as

\[U(x,y,z)=\frac{z}{j\lambda}\iint_{\Sigma}U(\xi,\eta,0)\frac{\exp(jkr_{01})}{r_ {01}^{2}}d\xi d\eta \tag{1}\]

where \(j=\sqrt{-1}\), \(U(x,y,z)\) describes the wavefield on target \((x,y)\) plane after diffraction distance \(z\) and \(U(\xi,\eta,0)\) describes the wavefield on the emission \((\xi,\eta)\) plane, \(\lambda\) is the wavelength of the input laser, \(k\) is the wave number where \(k=\frac{2\pi}{\lambda}\), \(r_{01}\) is the vector pointing from \(P_{0}\) to \(P_{1}\) and the distance \(r_{01}\) is given by

\[r_{01}=\sqrt{z^{2}+(x-\xi)^{2}+(y-\eta)^{2}} \tag{2}\]

When diffraction angle \(\theta\) shown in Figure 4 is small enough, the computation complexity can be further reduced by applying conditions to the application scenarios while maintaining the emulation accuracy. As a result, In _Fresnel approximation_, by simplifying \(r_{01}\) with binomial expansion of the square root in Equation 2 and eliminating terms but \(z\) in the \(r_{01}^{2}\) appearing in the denominator of Equation 1, it is described as

\[U(x,y,z)=\frac{e^{jkz}}{j4z}\iint_{\Sigma}U(\xi,\eta,0)\exp\{j \frac{k}{2z}[(x-\xi)^{2}+(y-\eta)^{2}]\}d\xi d\eta \tag{3}\]

In Fresnel approximation, the critical approximation happens in the approximation of the exponent, which can be seen that the spherical secondary wavelets will be replaced by wavelets with parabolic wavefronts. Thus, the condition on the distance \(z\) will be \(z^{3}\gg\frac{\pi}{4\lambda}[(x-\xi)^{2}+(y-\eta)^{2}]_{max}^{2}\), i.e., the observer (the \((x,y)\) plane) is in the near field of the aperture.

Furthermore, when \(z\gg\frac{k(\xi^{2}+\eta^{2})_{max}}{2}\) is satisfied, which means the quadratic phase factor under the integral sign in Equation 3 is approximately unity over the entire aperture, _Franuhofer approximation_ will further greatly simplify the calculations. Thus, in the far field of the aperture, the diffraction can be

Figure 4: Diffraction illustration. – (a) \((\xi,\eta)\) is the plane for diffraction aperture and illuminated by the input light beam in positive \(z\) direction, \(\Sigma\) on the plane \((\xi,\eta)\) denotes the illuminated area. \((x,y)\) plane is the target plane. \(P_{0}\) and \(P_{1}\) are illuminated points on two planes. \(\theta\) is the angle between the outward normal and the vector pointing from \(P_{1}\) to \(P_{0}\). (b) Light propagation and phase modulation through the diffractive layer w.r.t the input light wave.

Figure 3: Agile DONN design flow overview – (1) Design space exploration (DSE) w.r.t the design specification and ML task with LightRidge to automatically search best system parameters; (2) Co-design training with DSE explored parameters and physical hardware/deice parameters; (3) LightRidge backend supports for co-design fabrication; (4) Post-fabrication system integration; (5) LightRidge-DSL that simplifies (1)–(4) with user-friendly front-end APIs. Note that, (1)(2)(3)(5) are executed automatically with LightRidge and (4) is physical demonstrated in this work.
approximated as

\[U(x,y,z)=\frac{e^{jkz}e^{j\frac{kz}{z}(x^{2}+y)}}{jdz}\iint_{\Sigma}U(\xi,\eta,0) exp[-j\frac{2\pi}{\lambda z}(x\xi+y\eta)]d\xi d\eta \tag{4}\]

Thus, the diffraction process can be more generally formulated as - when an input wave resulted from \(l-1\)-th layer \((\xi,\eta)\), \(U_{l-1}(\xi,\eta,0)\) diffracts over diffraction distance \(z\) to the \(l\)-th layer (\(x,y\)), the resulted wavefield \(U_{l}^{1}(x,y,z)\) in time domain is described as

\[U_{l}^{1}(x,y,z)=\iint U_{l-1}(\xi,\eta,0)h(x-\xi,y-\eta,z)d\xi d\eta \tag{5}\]

where \(h\) is the diffraction function of free space. It can be calculated with spectral algorithm with Fast Fourier Transform (FFT) for fast and differentiable computation. By convolution theorem, the integral can be calculated with

\[\mathcal{F}_{xy}(U_{l}^{1}(x,y,z))=\mathcal{F}_{xy}(U_{l-1}(\xi,\eta,0)) \mathcal{F}_{xy}(h(x,y,z)) \tag{6}\]

\[F_{l}(\alpha,\beta,z)=F_{l-1}(\gamma,\sigma,z)H(\alpha,\beta,z) \tag{7}\]

Then, the multiplication result \(F_{l}(\alpha,\beta,z)\) will be transformed back to the time domain as \(U_{l}^{2}(x,y,z)\) by _inverse Fast Fourier Transform_ (iFFT) for phase modulation, which is the input to the phase modulation.

#### 3.1.2 Phase modulation

The phase modulation functions like _weight parameters_ in conventional neural networks and is updated iteratively during training process. Specifically, the input wave \(U_{l}^{2}(x,y)\) (for simplicity, we discard \(z\) in phase computation representation as \(z\) is not involved) can be described by its amplitude and phase, by _Euler's formula_, it can be described with a complex number in time domain, i.e.,

\[U_{l}^{2}(x,y)=A(x,y)e^{j(\theta,xy)}=A\text{cos}\theta+jA\text{sin}\theta \tag{8}\]

Where \(j=\sqrt{-1}\), \(A\) is the amplitude, \(\theta\) is the phase for the input wave; Acos\(\theta\) is the real part and Asin\(\theta\) is the imaginary part. After the phase modulation \(\phi(x,y)\), the wave function will be:

\[\begin{split} U_{l}(x,y)&=Ae^{j(\theta+\phi)}\\ &=(A\text{cos}\theta+jA\text{sin}\theta)\times(\text{cos}\phi+j \text{sin}\phi)\\ &=U_{l}^{2}(x,y)\times\phi(x,y)\end{split} \tag{9}\]

which can be realized with complex-valued matrix multiplications. \(U_{l}(x,y)\) is the input wavefunction for the forward function (Equation 5) for the \(l+1\)-th diffractive layer.

### Physics-aware Complex-valued Regularization

Considering the physics in optics, the DONN system is described and emulated with complex numbers. According to Equation 9, the training for the DONN system is more phase modulation dominated, and the intensity at the end of diffraction will decrease exponentially as the number of diffractive layers increases, which means a regularization between amplitude and phase is required to avoid gradient vanishing and explosion in the training process. With this insight, we introduce a novel regularization factor \(\gamma\) in the forward function to improve the training efficiency, which can flexibly change the gradient scales between amplitude and phase modulations. Specifically, \(\gamma\) is applied to amplitude vector \(A\) in Equation 9.

### LightRidge Framework

LightRidge framework (Table 2) consists of four major components to simplify and accelerate the process of design, exploration and deployment of the DONN system, including **a)** versatile programming modules for precise physics modeling, **b)** domain-specific neural architecture modules of DONNs, **c)** accelerated physics kernels for training and inference runtime improvements, and **d)** hardware deployment supports.

**Low-level physics modeling** - Three components are required to design a DONN model, including laser source, diffractive layers, and optical/photon detector. To model the whole physical phenomenon of DONNs, we first introduce the mathematical modeling modules for the implementation of DONN systems - **(1)** Various laser source modelings with flexible wavelength settings and beam profiles. **(2)** Precise light diffraction approximation, which falls into three categories - _Rayleigh-Sommerfeld_, which handles both far and near fields but with the highest computational complexity (Equation 1); _Fresnel_, which approximates the propagation with parabolic wavefronts, namely the near field propagation (Equation 3); _Fraunhofer_, implemented with Equation 4, approximating the propagation with planar wavefronts in the far-field [54]. **(3)** The optical/photon detector digitizes the analog light intensity to make it processable by the computer.

**Model-level APIs** - The DONN model is constructed with flexible model-level modules with LightRidge, where the architectural parameters can be used to customize the system - **(1)** the laser source module lr.laser offers precise laser customization including laser specifications such as wavelength, src_profile, etc. **(2)** The physics modeling of diffraction with trainable phase modulation is implemented in lr.layers. Two diffraction modelling with and without hardware-specific are provided with lr.layers.diffractlayer and lr.layers.diffractlayer_raw. Specifically, to deal with **challenge 2** in Section 1, lr.layers.diffractlayer employs the codesign algorithm, where the device-level information is delicately integrated in the training process with quantization methods, such as Gumbel-Softmax [25, 36, 31] and quantization-aware training (QAT) [28], applied on the trainable parameters, i.e., the phase modulation in diffractive layers, for efficient modeling-to-hardware deployment. Both modules can alternate three diffraction approximation algorithms according to user definition. Additionally, user-defined system hyperparameters such as size of diffraction unit (pixel_size), diffraction distance
(distance), the precision of the hardware implementing diffractive layers (precision) can also be customized easily with our language. **(3)** The detector is employed to capture the light intensity after propagation through the system, which is the interface component for linking training loss construction and the DONN model. In lr.layers.detector, x_loc and y_loc are lists of spatial coordinates of the detector, and the size of the detector regions is customized by det_size. **(4)** Finally, lr.models is a sequential container that stacks arbitrary numbers of customized diffractive layers in the order of light propagation in the DONN system and a detector plane. As a result, we construct a complete DONN system just like constructing a conventional neural network.

**Training support** - The DONN model is trained with conventional automatic differentiation engines in complex domain, which is supported by our differentiable physics kernels and training utility functions. Specifically, the original one-dimensional input is processed to a complex-valued input by initializing the phase information in data_to_cplex. Training parameters such as optimizer, complex-valued regularization regu_factor, loss function loss, etc, are also enabled in complex domain by lr.train.utils. The CPU and GPU accelerations are enabled by to(device). Finally, lr.train.dse enables physics-aware DSE for DONNs design and integration (Section 4).

**Hardware deployment** - To practically deploy the digitally trained model to hardware, the quantization to specific precision of the hardware (post-training quantization) is provided by lr.layers.weight_fab. More importantly, LightRidge directly supports various popular hardware deployment, e.g., SLMs for laser in visible range with lr.model.to_slm(), and 3D printed phase mask for THz systems with lr.model.to_3d_render().

## 4 Design Space Exploration

Taking advantages of LightRidge, we introduce the first explicit architectural design space exploration (DSE) engine for DONNs, namely LightRidge-DSE. As discussed earlier, the domain knowledge of optics and optical hardware are critical technical barriers to design DONNs. Therefore, there is a great need to enable an automatic DSE exploration in LightRidge, which will significantly shorten the design and hardware deployment cycle of DONNs and lower the optical domain-knowledge requirements. We propose an analytical model based DSE approach to accelerate the DSE process, where the analytical model is extracted from a ML regression model. Our main goal of the DSE engine is to provide guidance to design DONN systems under new design parameters with fabrication and chip integration requirements (e.g., fabrication technologies, chip dimension, etc.) as inputs.

**Design space of DONNs** - We consider the DONN design space from two aspects: (1) The major physical architectural design parameters of DONNs include - the diffraction unit size (the dimension of each diffractive unit), and the diffraction distance, i.e., the physical distance between the source to the first diffractive layer, layer to layer and the last layer to the detector (\(z\) in Figure 2). These two are critical architectural parameters under a fixed laser profile (wavelength). (2) The space exploration over DONNs spatial architectural parameters - system size (or system resolution) paired with

\begin{table}
\begin{tabular}{|p{42.7pt}|p{113.8pt}|p{113.8pt}|} \hline Classes & Modular Programs & Description \\ \hline \multirow{4}{*}{Low-level modeling} & Laser source \& profiles & Modeling coherent laser beams with various wavelength/profiles, e.g., Gaussian beam, Bessel beam, etc. \\ \cline{2-4}  & Diffraction approximation & High-performance tensor implementations of numerical diffraction approximations, including **Rayleigh-Sommerfeld** (Equation 1), **Fresnel** (Equation 3), and **Fraunhofer** (Equation 4). \\ \cline{2-4}  & Optical/photon detector & An photon detector to capture the light intensity and convert the analog intensity information to the digital computer-processable information. \\ \hline \multirow{4}{*}{Model-level APIs} & lr.laser & Define the laser source for the system, including laser wavelength and it profile. \\ \cline{2-4}  & \multirow{2}{*}{lr.layers APIs} & Includes modules of different types of diffractive modeling, e.g., hardware-specific layer module lr.layers.diffractlayer and general diffractive layer lr.layers.diffractlayer\_raw, that can be configured with various approximation methods, distance, diffraction unit size, etc. \\ \cline{2-4}  & lr.layers.detector & Define detector designs for various computer vision ML tasks, e.g., classification detector with coordinate information and the size of the detector regions. \\ \cline{2-4}  & lr.models & Sequential container to customize DONN system by stacking diffractive layer and detector modules. \\ \hline \multirow{4}{*}{Training} & lr.train.utils & Training utility modules including data handling (e.g., utils.data_to_cplex), complex-valued regularization, loss function, optimizer, etc. \\ \cline{2-4}  & lr.train.to(device) & Enable CPU and GPU accelerations for accelerating diffraction emulation and DONNs training. \\ \cline{2-4}  & lr.train.dse(specs) & Perform pre-fabrication design space exploration with chip integration specifications as input. \\ \cline{2-4}  & lr.layers.weight\_fab & Quantize the phase weights in the trained model w.r.t the hardware (e.g., SLMs) specifications. \\ \cline{2-4}  & lr.model.to\_system & HW deployments modules, e.g., lr.model.to\_slm(), lr.model.to\_3d\_render(), etc. \\ \hline \end{tabular}
\end{table}
Table 2: Overview of the LightRidge programming modules and partial front-end APIs. Note that we use lr to represent our integrated Python package lightridge.
* hardware/device precision, i.e., discrete phase modulation levels provided by the device, which are sensitive parameters w.r.t the performance of ML tasks. We take the physical architectural DSE as an example in this section.

**DSE features and data collection** - In our case, we show the process of conducting the DSE with the physical architectural design parameters, i.e, the diffraction unit size \(d\) and the diffraction distance \(D\), for DONN systems under different laser wavelength \(\lambda\). With fixed system size \(200\times 200\) and precision of the SLMs with 256 optical states covering [0,2\(\pi\)], we collect training data by sweeping diffraction unit size from \(10\lambda\) to \(110\lambda\) and diffraction distance \(D\) from \(0.1\)m to \(0.6\)m on a 5-layer DONN system, i.e., 121 data points, for laser wavelength \(\lambda\) of \(632\)nm and \(432\)nm.

**Analytical model based DONN DSE** - We employ a gradient boosting regression [41] model to find out a polynomial analytical model to bypass and transfer optical physics-aware DONNs DSE knowledge to new nearby \(\lambda\). Specifically, our analytical model is trained with diffraction unit size and diffraction distance exploration data points from systems with \(\lambda=632\)nm and \(\lambda=432\)nm (Figure 5 (a) and (b)) to estimate the DONNs design space in ML performance given a different laser profile with \(\lambda=532\)nm (Figure 5(c)). The regression model takes the wavelength \(\lambda\), \(d\), and \(D\) as inputs, and predicts (regression) the accuracy w.r.t MNIST dataset, trained with mean squared error (MSE) loss. The regression model is built with _n_estimators=3500, learning_rate=0.2, max_depth=3, random_state=25. The approximated prediction result from the analytical model is employed to guide DONN DSE under a new target \(\lambda\). To evaluate the analytical model based DSE strategy, we compare the predicted design space (Figure 5(c)) with the emulation verified design space (Figure 5(d)) under \(\lambda=532\)nm. The star point in Figure 5(d) shows our analytical DSE can find the best design points, which is verified by the end-to-end LightRidge development process in Section 5.

The analytical model by DSE can generalize the learnt optical formula to DONN systems with new laser wavelength while following the traditional maximum half-cone diffraction angle theory [5], i.e., the analytical model should be applied to a nearby wavelength within the applicable range by the theory of the training data. In our DSE example (Figure 5), we use the analytical model trained from \(432\) nm and \(632\) nm for predictions under \(532\) nm. However, such a analytical model trained with wavelength in visible range will not work for predictions for wavelength in other ranges, such as Infrared (IR) and Microwaves, which results in the theory violation.

**Sensitivity analysis** - We perform single parameter control variable tests for all three parameters. Our results confirm that the most sensitive parameter is the diffraction unit size, where wavelength and distance are almost equally sensitive to the accuracy performance. As shown in Table 3, by shifting the DSE explored best parameters (the star point in Figure 5(d)) +10%/+5% or -10%/-5%, we observe sharply accuracy drops on unit size (dropped to 30% accuracy on MNIST by shifting only +-5%), and are less sensitive on the other two parameters (dropped to 74% by shifting +-5%).

With the guidance of the analytical model, LightRidge-DSE finds the best architecture dimension and training parameters with several emulation iterations for selected possible parameters instead of sweeping through the grid-based search space. For example, in our case shown in Figure 5, aided by the analytical model, few emulation iterations (e.g., two emulations) instead of grid-searching over 121 data points are required for DSE, resulting in \(60\times\) speedups. On the other hand, DSE engine is able to provide general design parameters for the similar type of ML task. For example, the DSE model trained by MNIST dataset is also confirmed to be applicable to other

Figure 5: Results of architectural DSE of DONN systems w.r.t diffractive unit size, and diffraction distance under different laser wavelength (\(\lambda\)) with each grid colored according to accuracy on MNIST10. (a) and (b) are training data from emulations w.r.t design space under \(\lambda=432\)nm and \(\lambda=632\)nm for the inference model. (c) Predicted performance w.r.t design space under \(\lambda=532\)nm with the ML DSE model trained with data points from (a) and (b). (d) Grid-search validation under \(\lambda=532\)nm that verify the ML-based DSE quality. The DSE-guided setup at the star point is verified practically in Section 5.
MNIST-like datasets such as FashionMNIST [59], Kuzushiji-MNIST [8], Extension-MNIST-Letters [9].

## 5 Evaluation

In this section, we first demonstrate that LightRidge and LightRidge-DSE offer precise hardware-software correlations w.r.t real-world DONNs system realization (Figure 6) via a visible range DONN prototyping. Second, we demonstrate the effectiveness of LightRidge framework over SOTA experimental baselines [34, 67] in training performance and emulation runtime (Figure 7 - 9). Finally, we demonstrate that LightRidge and LightRidge-DSE enables comprehensive DONN system on-chip integration (Figure 11) and the capabilities to design advanced DONNs design principles, including multi-channel DONNs classifier on Place365 [65] dataset (Figure 12) and the all-optical image segmentation architecture (Figure 13). Note that experiments in Section 5.1 are physically deployed on optical hardware shown in Figure 5(a), while other results are from emulations with LightRidge.

### LightRidge and LightRidge-DSE Validations via Physical DONNs Prototyping

**Model construction via LightRidge-DSE and training -** This section demonstrates the hardware-software codesign precision and the effectiveness of LightRidge-DSE, where the parameters of the DONNs model used for physical validation experiments are automatically produced by LightRidge-DSE with the system size of \(200\times 200\). Specifically, the emulation model for DONN training is constructed with 3 sequentially stacked diffractive layers in lr.model, where each layer is defined with lr.layers.diffractlayer integrating hardware specifications: 1 the diffraction pixel size is \(36\text{um}\times 36\text{um}\); 2 the laser wavelength is \(532\text{nm}\). Consulted on DSE results shown in Figure 5(b), distance is explored to be \(\sim 0.3\text{m}\), which is further adjusted to 11 inches (\(0.28\text{m}\)) on our optical table. There are 10 pre-defined detector regions for labels placed evenly on the detector plane. The learning rate in the training process is 0.5 trained with 100 epochs for all experiments using Adam [27] with batch size 500 using MSE loss.

**Hardware prototype and validation** - Laser source CPS532 from Thorlabs, Inc. is employed to realize the physical DONN system, where SLMs (LC 2012 HOLOEYE) implement the diffractive layers. The SLMs are experimentally measured and cover a phase modulation range close to \([0,2\pi]\). The final diffraction pattern is captured on a CMOS camera (CS165MU1 Thorlabs, Inc.). To make the input easier for hardware deployment, we train and validate the model with binarized MNIST images as shown in Figure 5(a), where the trained phase modulation parameters are loaded on the SLMs.

The resulted detector patterns for the inputs are shown in Figure 5(b). The SLM used to encode input binary images is illuminated by the laser source, and the input information will be encoded on the intensity of the input laser. The intermittent practical propagation will not be available in the DONN system as the information is carried with the light beam. At the end of the system, we can capture the intensity information on the detector for the prediction. As shown in Figures 6, DONNs emulation results in LightRidge precisely match the experimental measurements, which demonstrates: (1) precise correlations between the implemented high-level modeling and low-level physics experimental system, which improves the design efficiency significantly without manual HW calibration requirements shown in Figure 1; (2) and the effectiveness of LightRidge-DSE in exploring architecture parameters, which has been further utilized for on-chip integration (Section 5.5).

### Emulation-level Evaluation

We further verify the design parameters from DSE model as discussed in Section 5.1 at emulation level. The accuracy results for image classification with MNIST [29] and FashionMNIST (FMNIST) [59] dataset are shown in Figure 7, where the baseline results are conducted on training methods in [34, 67] without the proposed physics-aware complex-valued reg

Figure 6: Evaluation results of a 3-layer DONN system in visible range explored, trained, and deployed by LightRidge – (a) The experimental system trained and deployed by LightRidge. The corresponding detector pattern from experiments and simulation results produced by LightRidge are shown in Figure 5(b); (b) Corresponding detector patterns of experimental measurements and simulation results (the simulation results generated with lr.model.prop_view()) of the 3-layer DONN system.
ularization. The inputs are encoded with the amplitude of the laser beam. To make the input fit the DONN system, we first extend the image with the original size of \(28\times 28\) in MNIST10 and FMNIST datasets to \(200\times 200\) in SLM resolution, and transfer the original one-dimensional image to complex-valued image in the emulation. **With the regularization factor \(\gamma\) implemented, our training algorithm has a significant advantage in training less complex DONN models.** For example, when the DONN model is implemented with only one diffractive layer (D=1), the accuracy is 31% (34%) improved for MNIST (FMNIST) dataset, compared with the baseline. Additionally, our algorithm can achieve a similar accuracy performance (0.98 for MNIST, 0.89 for FMNIST) for DONN systems regardless of its complexity, i.e., the number of diffractive layers implemented in the system, by adjusting \(\gamma\) for the model training. However, according to the discussion in [34], the performance of DONNs with fewer number of layers are fundamentally limited by the optical physics, which is opposite to our accuracy results.

To understand the increase of accuracy, we analyze the robustness of the DONNs trained with complex-valued regularization. Specifically, we explore the confidence of the predictions acquired by the system, by adding random uniform noise at the detector phase with upper bound 1%, 3%, and 5% intensity noise. As a result, for both datasets, **as the depth of DONNs increases, the prediction confidence increases, while the prediction accuracy with no noise applied are all relatively the same.** For example, there is no accuracy degradation on five-layer DONNs for MNIST, and less than 1% degradation on FMNIST with up to 5% applied noise. However, for single-layer DONNs, the accuracy drops 63% for MNIST and 54% for FMNIST with 1% noise applied, and drops to 0 when applied noise increases to 3% and 5%.

### LightRidge Runtime Evaluation

Runtime efficiency of emulating DONNs is crucial in simulation, training, and exploration. Thus, optimizing runtime performance is another key contribution in LightRidge framework. We first analyze the DONN workloads, where we identify that the majority (\(\geq\)90%) of the runtime complexity comes from the numerical modeling of light diffraction. Thus, the major optimization efforts should lie over the diffraction kernels. Second, to effectively utilize the modern computing platforms, we aim to maximize the parallelism from the fundamental physics modeling, which is the main reason of implementing scalar diffraction modeling instead of FDTD in the computation kernel as mentioned earlier in Sections 1 and 2. The diffraction approximation functions with scalar diffraction modeling (Equations 1 - 4) can be breakdown into three major tensor-level operators: complex-domain 2-D FFT (FFT2), inverse 2-D FFT (iFFT2), and complex matrix multiplications (Complex MM). Based on the analysis and kernel breakdowns, we take advantages of modern CPU and GPU platforms by incorporating efficient complex-tensor datatypes and operators. For CPU, the diffraction kernel is optimized via Intel Math Kernel Library (MKL-DNN) complex kernels with AVX-512 support; for GPU, cuFFT, cuFFTW, and cuTENSOR libraries with efficient complex-domain FFTs and MM are deployed.

To demonstrate the runtime improvements, we compare the runtime of our proposed framework with up-to-date commercial version LightPipes(2021), running various emulation loads, i.e., {1,3,5,7,10}-layer DONNs while system resolution sweeps from 100\(\times\)100 to 500\(\times\)500. All LightPipe-CPU and LightRidge-CPU results are conducted on Intel Xeon Gold 6230 20x CPU. To make fair GPU comparisons, we re-implement the kernels in LightPipes with cupy [40], and runtime results are collected on Nvidia 3090 Ti GPU platform.

Figure 8 shows LightRidge consistently outperforms LightPipes on both CPU and GPU back

Figure 8: LightRidge runtime speedups over LightPipes with various DONNs system sizes – (a) CPU speedups. (b) GPU speedups.

Figure 7: Confidence evaluation of DONNs trained with complex-domain regularization under various system complexity. Baseline results are conducted on methods in [34, 67] without noise assumption.

Figure 9: Runtime speedups breakdown with 5-layer 500\(\times\)500 DONNs. FFT2, iFFT2, and Complex MM are the main operators for DONN numerical modelling.
ure (a)a shows at most 6.4\(\times\) speedup of LightRidge-CPU over LightPipe-CPU at \(depth=5,systemsize=500^{2}\). Figure (b)b shows at most 12\(\times\) speedup of LightRidge-GPU over LightPipe-GPU at \(depth=1,systemsize=500^{2}\). To understand the runtime speedups offered by LightRidge, we provide normalized speedups breakdown analysis w.r.t LightPipe CPU/GPU, shown in Figure 9 with 5-layer DONNs workload. We observe that the 6.4\(\times\) CPU runtime speedups are contributed from the FFT2 (11\(\times\)), iFFT2 (10\(\times\)), and Complex MM kernels (4\(\times\)); similarly for GPU, 8.6\(\times\) overall speedups are primarily contributed from the three kernels with 7\(\times\), 7\(\times\), and 12\(\times\) speedups, respectively.

Furthermore, we evaluate capability of LightRidge of training large DONNs systems (Figure 10). The runtime is acquired on a single Nvidia 3090 Ti GPU. We can see that LightRidge handles 30-layer DONNs training in \(\sim\) 280 seconds per epoch, with input image resolution at 500\({}^{2}\). Besides, we observe runtime increases almost linear w.r.t the DONNs depth, while there is a runtime jump when the system size increases beyond 300\({}^{2}\), mainly due to the limited resource on a single GPU. This posts strong motivations for further CUDA optimization and multiple-GPU training supports in future work.

### Performance Comparison between DONNs and conventional NNs

Compared with conventional NN models on digital platforms, the current optical-devices-deployed DONNs systems at this early stage suffer from accuracy performance degradation while feature with significantly improved energy efficiency. As shown in Table 4, we evaluate two conventional NNs including a MLP, which consist of two linear layers with hidden size of 128, and the input image is flattened as one-dimensional tensor, i.e., MLP (40000 \(\rightarrow\) 128 \(\rightarrow\) 10); and a CNN, which consists of two Conv2D, where the kernel size of both layers is set as (5, 5) and 32 filters for the first layer and 64 filters for the second layer with stride and padding being 2, two MaxPooling2D, where kernel size is set as (3, 3) with stride 2, followed by two linear layers. Additionally, we deploy the conventional NNs on different digital platforms including Nvidia GPU 2080 Ti, Nvidia GPU 3090 Ti, Intel Xeon 6230 20x CPU, and Google EdgeTPU [62].

As a result, the conventional NNs can produce the accuracy performance of 0.99/0.99 for MNIST, and 0.91/0.91 for FashionMNIST with the MLP and the CNN model, respectively, while DONNs systems reach the accuracy performance of 0.98/0.89 for MNIST/FashionMNIST, which shows 1% accuracy performance degradation. For practical realization with DONNs systems, we take the prototype in Figure (a)a as an example, the power of a CW 532nm laser source is \(\sim\) 5mW. The diffractive layers are passive optical devices and requires no extra energy for computation. Then the power consumption at the CMOS detector is \(\sim\) 1 W (max) @ 1000 fps with the system size of 200\(\times\) 200. Thus, the power efficiency for the DONN system can be estimated as 995fps/Watt. The corresponding energy efficiency results for conventional NNs on various digital platforms are shown in Table 4, which shows the DONN system is roughly 2 orders more efficient than desktop CPU and GPU, 1 orders than digital edge devices with batch size as 1. The energy efficiency provided by DONN systems can be more significant when dealing with more complex ML tasks (e.g., applications in Section 5.6) as the computation part (with passive optical devices) consumes zero power. Note that DONNs energy efficiency can be further optimized with integrated fabrication and high-end detector.

Therefore, the DONN system shows its great potential in completing ML task much more energy-efficiently than conventional NNs. However, the degradation of accuracy performance and the challenges in deploying the practical inference systems call for more future works in broad disciplinaries, such as complex-domain training algorithms, domain-specific co-designing, and optics, which also highlights the potential of our framework.

### On-chip DONNs Integration via LightRidge

The bulky 3D free-space DONN systems can be integrated as a 3D monolithic on-chip DONNs via 3D additive fabrication [13, 14, 20, 35], e.g., galvo-dithered two-photon nanolithography [20], electron beam lithography overlay process [35], etc. Such monolithic on-chip DONNs can be integrated in a hybrid computing system, with DONNs performing as an optical co-processor hosted by central processor via system interconnects (e.g., PCIe 4.0). The host processor controls the laser encoding for loading images and the results collection with the co-processor interconnects, illustrated in Figure 11. Each diffractive layer is a thin film, where the trained phase information is encoded with the thickness of the material used for layer fabrications. Between diffractive layers, the optical clear adhesive is employed to provide free-space light propa

\begin{table}
\begin{tabular}{|c|c|c|c|} \hline
**Platform** & **fps/Watt** & **Accuracy** \\ \hline
**GPU 2080 Ti** & 3.3 (301\(\times\)) & 3.8 (261\(\times\)) & \\ \hline
**GPU 3090 Ti** & 2.4 (414\(\times\)) & 1.7(585\(\times\)) & \\ \hline
**CPU Xeon** & 1.5(663\(\times\)) & 2.0 (497\(\times\)) & \\ \hline
**XPU (EdgeTPU)** & 23(43\(\times\)) & 26 (38\(\times\)) & \\ \hline
**Our DONNs** & \multirow{2}{*}{995} & \multirow{2}{*}{0.98} & \multirow{2}{*}{0.89} \\
**prototype** & & & \\ \hline \end{tabular}
\end{table}
Table 4: Energy efficiency (fps/Watt) and accuracy comparisons between DONNs systems and conventional NNs.

Figure 10: Large-scale DONNs training runtime.
gation, whose thickness is the diffraction distance. Diffractive layers and optical clear adhesive are stacked sequentially to construct an on-chip DONN system. The final prediction is captured on the detector, with Analog-Digital-Converter (ADC), I/O interface, and memory buffers integrated on the peripheral circuits. An example of aforementioned real-world DONN on-chip integration is realized by [35]. However, due to the three challenges we discussed earlier, the design cycle could take months to year efforts. LightRidge framework can significantly simplifies the end-to-end on-chip design process, demonstrated by an case study as follows.

**Case study** - We target a 5-layer DONN system integration under wavelength 532nm for a CMOS detector chip (CS165MU1 from Thorlabs, Inc.), shown in Figure 11, where the CMOS chip defines the pixel size of 3.45um. The key for on-chip integration is to search for valid fabrication parameters with high prediction performance w.r.t ML tasks. Therefore, following the four steps of LightRidge design flow (Figure 3), we first deploy LightRidge-DSE to explore the 3D fabrication dimension, including distance, resolution, and diffraction unit size. According to the emulation results in Section 4 and Figure 5(c), when we fix the wavelength as 532nm and the diffraction unit size (pixel size of the CMOS chip) as 3.45um, considering image classification as ML tasks (e.g., MNIST), LightRidge-DSE returns the diffraction distance of 532um, and the resolution 200\(\times\)200, to fit the CMOS chip. Thus, the DONNs fabrication dimension is finalized as \(690\text{um}\times 690\text{um}\times 2660\text{um}\), where 2660um is the height, and flat chip dimension is \(690\times 690\text{um}^{2}\), which align with the chip fabrication procedure in [35]. Next, after training completed, each layers will be fabricated w.r.t the phase parameters optimized by the codesign stage via nano-printing on the targeted CMOS detector chip. The integrated DONNs can be then used as a co-processor via ADCs and I/O integrated with the CMOS detector chip, where the pre-fabrication design process takes less than a day via LightRidge.

### Advanced DONN Architectures

With the design capabilities of LightRidge and LightRidge-DSE verified by physical optical systems, we further explore the potentials of DONN systems with more advanced architectures dealing with complex computer vision tasks. Specifically, we propose and evaluate **(1)** a multi-channel DONN architecture implemented with diffractive layers to deal with RGB image classifications, and **(2)** the first-ever optical image segmentation demonstration using DONNs with _optical skip connection_, to deal with image segmentation and potentially other image-to-image synthesis tasks.

#### 5.6.1 All-optical RGB image classification

To deal with more complex datasets in image classification, e.g., Place365 [65], a high-resolution RGB image dataset, we propose a multi-channel RGB-DONNs architecture. As shown in Figure 12, three optical channels are employed in the DONN system to deal with 'R', 'G', 'B' channels separately in the original image. The input laser beam is split with the beam splitter into three beams and reflected with mirrors into three channels to encode the corresponding input information. Note that the image information is encoded with light intensity at the encoding layer for each channel, in which case each channel takes a gray-scaled image as input and propagates through five diffractive layers. Each channels is constructed with the same system parameters in Section 5.1 expect for changing to 5 diffractive layers. The output laser beams from all channels are projected to a single detector, where the light intensity is merged for final prediction. Similar to the detector design for classification shown in Figure 2, a single detector collects the intensity of the output within each pre-defined detector region and produces the predicted class by argmax. All three channels are trained w.r.t the same shared loss function. The training setups are the same as discussed in Section 5.1.

The emulation accuracy results for image classification with Place365 are shown in Table 5, including top-1, top-3, and

Figure 11: Monolithic on-chip DONNs design and overall hybrid architecture system integration.

\begin{table}
\begin{tabular}{|c|c|c|c|} \hline
**Places365[65]** & _Top-1_ & _Top-3_ & _Top-5_ \\ \hline Our (Fig. 12) & **0.52** & **0.73** & **0.84** \\ \hline Baseline [67] & 0.23 & 0.48 & 0.67 \\ \hline \end{tabular}
\end{table}
Table 5: Classification accuracy on Places365 (standard, 256-by-256) with _type of environment_ as classes.

Figure 12: Multi-channel RGB-DONNs architecture for image classification using Places365 dataset [65]. The ”R/G/B” channels are all encoded as three gray-scaled images using 532 nm laser source.
top-5 accuracy. The baseline is the emulation accuracy from the DONN model trained with the algorithm in [67]. We can see the model trained with our compiler has better accuracy performance than the baseline in all accuracy matrix (29% improvement for top-1 accuracy, 25% improvement for top-3 accuracy, 17% improvement for top-5 accuracy), and ours outperforms the baseline most at the top-1 accuracy.

#### 5.6.2 All-Optical image segmentation

Image segmentation is an interesting and challenging task in modern computer vision tasks, which has a great impact on autonomous physical systems such as autonomous driving, robotics, etc. Unlike image classification tasks, image segmentation is a process of generating representations of an image into specific image-to-image objectives. While in DONN classification systems, we observe that the system (output detector in particular) is not fully utilized, as only a given number of small detector regions are used for classification. As the DONN system propagates the input image w.r.t trained phase modulations in the full spatial dimension of the system, it is expected to be able to deal with image-to-image based tasks. Thus, we design and demonstrate the first-ever all-optical image segmentation.

Figure 12(a) includes the proposed a 5-layer DONN system, where we introduce two innovations in DONN architecture: 1) optical skip connection, which is inspired from the residual block design in conventional ResNet [24] architecture. It aims to smooth the gradient descent for better training performance and also is involved in inference for better detailed segmentation. Since the light signal is aggressively diffracted during the propagation, the optical skip connection can help to restore some features from less-diffracted inputs, making the model prediction be aware of the original information, which is confirmed to introduce better image segmentation performance with our results; and 2) layer normalization [1] before the detector plane, which is only employed in the training process for better training performance of the DONN by smoothing the training gradients. The dataset we demonstrate here is selected from CityScapes dataset [10], where the images are converted to gray-scaled images and resized to \(350\times 350\). We use binary labels in this case study to generate segmentation masks for _buildings_ and others. The baseline is the results from the DONN model construction without optical skip connection and the training method without layer normalization proposed in [34, 67]. The system parameters and training setups are the same as discussed in Section 5.1 expect for the system size changing to \(350\times 350\) and the model structure changing to Figure 12(a). The results shown in Figure 12(b) demonstrate that the advanced model trained with LightRidge outperforms the baseline in both edge detection and segmentation with significant clarity improvements on small object segmentation. These advanced DONNs architecture design and validation demonstrate the generalizability and power of LightRidge in exploring new architectural designs and applications.

## 6 Conclusion and Future Work

This work presents an agile end-to-end design framework LightRidge that enables seamless design-to-deployment of DONNs. LightRidge accelerates and simplifies the design, exploration, and on-chip integration by offering highly versatile and runtime efficient programming modules, and DSE (LightRidge-DSE) engine to construct and train the DONN systems in a wide range of optical settings. The high-performance physics emulation kernels are optimized for runtime efficiency, and verified together with hardware-software codesign algorithm on our visible-range prototype. Additionally, two advanced DONN architectural designs constructed with LightRidge show the capabilities and generalizability of LightRidge for various ML tasks and system design explorations. We believe our framework LightRidge will enable collaborative and open-source research in optical accelerators for not only ML tasks, but also other optics-related research areas such as optical structure emulation, chip fabrication (lithography), meta-material exploration, etc.

In the future, we will further optimize the runtime efficiency of LightRidge, including realizing high-performance CUDA kernel optimization and multiple-GPU computation. We also expect more functionality to be integrated in the framework and our hardware prototype. For example, the non-linearity in DONN systems, which can be realized by nonlinear optical materials (crystals, polymers, graphane, etc.), is an important implementation for more complex DONNs systems.

Figure 13: Image segmentation demonstrations using CityScapes [10] datasets with (a) a novel advanced DONNs architecture with optical skip connection and layer normalization for improving training efficiency, and (b) evaluations and comparisons to SOTA baselines [34, 67].