Middle output regularized end-to-end optimization for computational imaging

Abstract

Optical coding is an essential technique in computational imaging (CI) that allows high-dimensional signal sensing through post-processed coded projections to decode the underlying signal. Currently, optical coding elements (OCEs) are optimized in an end-to-end (E2E) manner where a set of layers (encoder) of a deep neural network models the OCE while the rest of the network (decoder) performs a given computational task. However, while the training performance of the whole network is acceptable, the encoder layers can be flawed, leading to deficient OCE designs. This flawed performance of the encoder is originated from factors such as the loss function of the network not considering the intermedium layers separately, as the output at those layers is unknown. Second, the encoder suffers from a vanishing gradient since the encoder takes place in the first layers. Third, the proper estimation of the gradient in these layers is constrained to satisfy physical limitations. In this work, we propose a middle output regularized E2E optimization, where a set of regularization functions is used to overcome the flawed optimization of the encoder. The significant advantage of our regularization is that it does not require additional knowledge of the encoder and can be applied to most optical sensing instruments in CI. Instead, the regularization exploits some prior knowledge about the computational task, the statistical properties of the output of the encoder (measurements), and the sensing model. Specifically, we propose three types of regularizers: the first one is based on statistical divergences of the measurements, the second depends only on the variance of the measurements, and the last one is a structural regularizer promoting low rankness and sparsity of the set of measurements. We validated the proposed training procedure in two representative CI systems, a single-pixel camera and a coded aperture snapshot spectral imager, showing significant improvement with respect to non-regularized designs.