Residual D<sup>2</sup>NN: training diffractive deep neural networks via learnable light shortcuts
Hongkun Dou, Yue Deng, Tao Yan, Huaqiang Wu, Xing Lin, Qionghai Dai
Abstract
The diffractive deep neural network ( <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> </mml:mrow> </mml:math> ) has demonstrated its importance in performing various all-optical machine learning tasks, e.g., classification, segmentation, etc. However, deeper <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">s</mml:mi> </mml:mrow> </mml:math> that provide higher inference complexity are more difficult to train due to the problem of gradient vanishing. We introduce the residual <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">s</mml:mi> </mml:mrow> </mml:math> (Res- <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> </mml:mrow> </mml:math> ), which enables us to train substantially deeper diffractive networks by constructing diffractive residual learning blocks to learn the residual mapping functions. Unlike the existing plain <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">s</mml:mi> </mml:mrow> </mml:math> , Res- <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">s</mml:mi> </mml:mrow> </mml:math> contribute to the design of a learnable light shortcut to directly connect the input and output between optical layers. Such a shortcut offers a direct path for gradient backpropagation in training, which is an effective way to alleviate the gradient vanishing issue on very deep diffractive neural networks. Experimental results on image classification and pixel super-resolution demonstrate the superiority of Res- <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">s</mml:mi> </mml:mrow> </mml:math> over the existing plain <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML" display="inline"> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:msup> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">D</mml:mi> </mml:mrow> <mml:mn>2</mml:mn> </mml:msup> </mml:mrow> <mml:mrow class="MJX-TeXAtom-ORD"> <mml:mi mathvariant="normal">N</mml:mi> <mml:mi mathvariant="normal">N</mml:mi> </mml:mrow> </mml:math> architectures.