A generalization bound for exit wave reconstruction via deep unfolding

Abstract: Transmission Electron Microscopy enables high-resolution imaging of materials, but the resulting images are difficult to interpret directly. One way to address this is exit wave reconstruction, i.e., the recovery of the complex-valued electron wave at the specimen's exit plane from intensity-only measurements. This is an inverse problem with a nonlinear forward model. We consider a simplified forward model, making the problem equivalent to phase retrieval, and propose a discretized regularized variational formulation. To solve the resulting non-convex problem, we employ the proximal gradient algorithm (PGA) and unfold its iterations into a neural network, where each layer corresponds to one PGA step with learnable parameters. This unrolling approach, inspired by LISTA, enables improved reconstruction quality, interpretability, and implicit dictionary learning from data. We analyze the effect of parameter perturbations and show that they can accumulate exponentially with the number of layers $L$. Building on proof techniques of Behboodi et al., originally developed for LISTA, i.e., for a linear forward model, we extend the analysis to our nonlinear setting and establish generalization error bounds of order $\mathcal{O}(\sqrt{L})$. Numerical experiments support the exponential growth of parameter perturbations.