A generalization bound for exit wave reconstruction via deep unfolding

TL;DR

This paper introduces a neural network approach based on unfolding proximal gradient iterations for exit wave reconstruction in electron microscopy, providing theoretical generalization bounds and demonstrating improved interpretability and accuracy.

Contribution

It extends LISTA-inspired analysis to a nonlinear phase retrieval problem, deriving generalization bounds and analyzing parameter perturbation effects in deep unfolding networks.

Findings

Unfolded PGA improves reconstruction quality and interpretability.

Parameter perturbations can grow exponentially with network depth.

Theoretical bounds on generalization error are established for the nonlinear model.

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.