Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A generalization bound for exit wave reconstruction via deep unfolding

Moussa Atwi, Benjamin Berkels

TL;DR

The paper recasts exit wave reconstruction in transmission electron microscopy as a phase-retrieval-type inverse problem and proposes a discretized, regularized variational framework. It then develops a deep-unfolded proximal gradient network by unrolling PGA steps into layers with a learnable unitary dictionary, combining model-based structure with data-driven refinement, and analyzes stability under parameter perturbations. A key theoretical contribution is a generalization bound, of order $\mathcal{O}(\sqrt{L})$, derived via covering numbers and Rademacher complexity, that quantifies how training sample size, parameter count, and network depth affect unseen performance. Numerically, the work shows perturbations can grow exponentially with depth under the nonlinear forward model, motivating the observed trade-off between interpretability and generalization in deep-unfolded estimators for phase retrieval in TEM. Overall, the study provides a rigorous link between algorithm unrolling, nonlinear forward models, and generalization behavior in high-dimensional inverse problems.

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.

A generalization bound for exit wave reconstruction via deep unfolding

TL;DR

The paper recasts exit wave reconstruction in transmission electron microscopy as a phase-retrieval-type inverse problem and proposes a discretized, regularized variational framework. It then develops a deep-unfolded proximal gradient network by unrolling PGA steps into layers with a learnable unitary dictionary, combining model-based structure with data-driven refinement, and analyzes stability under parameter perturbations. A key theoretical contribution is a generalization bound, of order , derived via covering numbers and Rademacher complexity, that quantifies how training sample size, parameter count, and network depth affect unseen performance. Numerically, the work shows perturbations can grow exponentially with depth under the nonlinear forward model, motivating the observed trade-off between interpretability and generalization in deep-unfolded estimators for phase retrieval in TEM. Overall, the study provides a rigorous link between algorithm unrolling, nonlinear forward models, and generalization behavior in high-dimensional inverse problems.

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 . 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 . Numerical experiments support the exponential growth of parameter perturbations.

Paper Structure

This paper contains 13 sections, 14 theorems, 109 equations, 1 figure.

Table of Contents

  1. Introduction
  2. Forward Model and Variational Approach
  3. A Simplified Exit Wave Model - Phase Retrieval
  4. Full forward model:
  5. Simplified forward model:
  6. Discretization
  7. Variational Perspective
  8. Optimization via Proximal Gradient Method
  9. Deep Unfolding
  10. Unrolling the Proximal Gradient Algorithm into a Neural Network
  11. Preliminaries
  12. Exponential Growth
  13. Generalization Error Bound via Covering Numbers

Key Result

Proposition 1

Let $D \subseteq \mathcal{H}$ be nonempty, and let $T : D \to \mathcal{H}$ be an operator. Let $\alpha, \beta \in \mathbb{R}_{++}$ and $f \in \Gamma_0(\mathcal{H})$. Then,

Figures (1)

  • Figure 1: Lower bound for $K_L$ computed by numerically maximizing the ratio $\frac{\lVert f_{{\boldsymbol{\Phi}}_1}^{L}(\mathbf{G}) - f_{{\boldsymbol{\Phi}}_2}^{L}(\mathbf{G})\rVert_F}{\lVert {{\boldsymbol{\Phi}}}_1 - {{\boldsymbol{\Phi}}}_2 \rVert_{2 \to 2}}$ with respect to ${{\boldsymbol{\Phi}}}_1$ and ${{\boldsymbol{\Phi}}}_2$ for the case $N=2$, $K=1$. Here, the dictionary matrices were restricted to be rotations and the optimization was done over the rotation angles.

Theorems & Definitions (33)

  • Definition 1
  • Proposition 1: Properties of cocoercive and proximal operators
  • Proposition 2: Real variant: BaCo17, Proposition 4.12
  • Corollary 1: Real variant: BaCo17, Corollary 4.13
  • Lemma 1
  • Lemma 2: Norm Composition Be17
  • Remark 1
  • Remark 2
  • Remark 3
  • Definition 2
  • ...and 23 more