Physics-Informed Neural Networks for Optimal Beam Shaping in Flat Optics

Technical Summary: Physics-Informed Neural Networks for Optimal Beam Shaping in Flat Optics

Problem Statement
The paper addresses the challenge of designing phase-only flat optical elements (such as SLMs or metasurfaces) that reshape an incident optical beam into a prescribed target intensity distribution. Conventional approaches, such as the Gerchberg–Saxton (GS) algorithm, rely on iterative Fourier optics and often yield irregular phase profiles with speckle-like fluctuations in the reconstructed intensity. Alternatively, formulating beam shaping as an energy-conserving ray-mapping problem leads to the Monge–Ampère equation. However, this equation is fully nonlinear and potentially degenerate, making it difficult for standard finite-difference or finite-element solvers to converge, particularly for non-cylindrically symmetric targets.

Methodology
The authors propose a Physics-Informed Neural Network (PINN) framework to solve the Monge–Ampère beam-shaping equations directly. The core methodology involves the following steps:

1. Theoretical Formulation:

   • The problem is modeled using a generalized Snell's law where the phase gradient determines the outgoing ray direction.
   • Conservation of optical power enforces a relationship between the source and target intensities via the Jacobian of the ray map, resulting in a second-order nonlinear partial differential equation (PDE): the Monge–Ampère equation.
   • The framework handles two distinct geometries:
     • Finite-distance: Rays are mapped to a specific target plane at distance $z$. The equation involves complex coefficients dependent on the phase gradient and propagation distance.
     • Far-field: The target is defined by an angular distribution (direction cosines). Here, the target coordinates are directly the phase gradients, simplifying the algebraic structure of the equation.

2. Normalization and Network Architecture:

   • To address conditioning issues arising from disparate physical scales, the authors normalize coordinates and phase variables onto a dimensionless square domain $[-1, 1]^2$.
   • A fully connected feedforward neural network with four hidden layers (150 neurons each) and smooth $\tanh$ activation functions represents the normalized phase potential $\bar{\Phi}(\bar{x}, \bar{y})$.
   • Smooth activations are essential because the PDE residual requires second-order derivatives.

3. Training Process:

   • The network is trained to minimize the pointwise residual of the Monge–Ampère equation. The loss function is the mean squared error of the power conservation violation across a set of collocation points.
   • Crucially, this is not a data-driven surrogate; the training signal is the optical PDE itself. No precomputed phase labels or ray correspondences are used.
   • Automatic differentiation (via PyTorch) computes the first and second derivatives required for the Hessian determinant.
   • Because the mapped target coordinates depend on the trainable phase, the target intensity is sampled using differentiable bilinear resampling, allowing gradients to propagate through the optical mapping.
   • Optimization employs Adam followed by full-batch L-BFGS to refine the solution near sharp features.

4. Validation:

   • The learned phase profiles are validated independently using scalar diffraction simulations (angular-spectrum propagation), distinct from the ray-based training objective.

Key Results
The authors demonstrate the method's efficacy through two numerical experiments:

• Finite-Distance Shaping: A Gaussian beam is reshaped into a smoothed star-shaped flat-top pattern at a distance of 10 cm. The PINN-generated phase produces a target intensity that closely matches the prescribed profile, with discrepancies concentrated only near the boundaries due to diffraction limits.
• Far-Field Shaping: A Gaussian beam is shaped into a "DT" logo pattern in the far field. The method successfully reproduces the rectangular support and internal peaks with a smooth phase potential.
• Comparison with Gerchberg–Saxton (GS): When applied to the same far-field logo target, the PINN outperforms the GS algorithm.
  • Visual Quality: GS produces significant speckle and high-frequency oscillations, whereas the PINN yields a cleaner profile with suppressed high-frequency structure.
  • Quantitative Metrics: The PINN achieves a mean-squared error (MSE) of $\approx 3.59 \times 10^6$ compared to $2.08 \times 10^8$ for GS. The relative root-mean-square error is 2.97% for PINN versus 22.6% for GS.
  • Efficiency: The energy efficiency (power within the target support) is 99.98% for PINN compared to 95.90% for GS.

Significance and Claims
The paper claims to present the first application of PINNs to beam-shaping problems in flat optics. Its primary contributions are:

1. Solving Nonlinear PDEs: It provides a robust method for solving the fully nonlinear Monge–Ampère equation for beam shaping, a task where traditional numerical solvers often struggle with convergence.
2. Generalization: The approach extends previous demonstrations of cylindrically symmetric metasurface beam shaping to arbitrary, non-cylindrically symmetric targets.
3. Superior Fidelity: The results indicate that PINN-derived phase profiles offer higher accuracy, lower speckle, and better energy efficiency compared to conventional phase-retrieval methods like Gerchberg–Saxton.
4. Wave-Optics Validity: The agreement between the ray-based training and independent scalar diffraction simulations confirms that the learned ray maps remain effective even when wave effects are considered.

The authors conclude that this PINN approach provides superior accuracy and fidelity for flat-optics beam-shaping tasks, offering a viable alternative to iterative Fourier methods for complex target distributions.