Differentiable Wave Optics

• Differentiable wave optics are computational methods that model light propagation accurately via simulations of diffraction, interference, and aberrations.
• They integrate automatic differentiation with precise forward models—such as FFT-based convolutions and phase modulations—using GPU-accelerated frameworks like PyTorch and TensorFlow.
• Applications span computational photography, microscopy, metasurface design, AR displays, and astronomical instruments, driving system co-design and inverse optimization.

Differentiable wave optics denotes a class of computational methods and software frameworks that enable end-to-end gradient-based optimization of optical systems governed by wave propagation physics. These methods integrate accurate physically based modeling of interference, diffraction, and aberration with efficient and exact automatic differentiation (autodiff), supporting joint optimization of physical device parameters and downstream algorithmic pipelines. Differentiable wave-optics simulators have been developed for compound refractive systems, hybrid refractive-diffractive assemblies, metasurfaces, multilayer thin-film structures, and large-scale astronomical interferometers, powered by implementations in GPU-accelerated autodiff frameworks such as PyTorch, TensorFlow, or JAX. They facilitate not only system-level co-design but also high-dimensional calibration, Fisher information–maximizing design, and topology optimization of nontrivial optical devices (Ho et al., 2024, Page et al., 2020, Yang et al., 2024, 2207.14780, Yang et al., 7 Jan 2026, Desdoigts et al., 2024, Desdoigts et al., 2024, Zhu et al., 2023).

1. Theoretical Foundations and Forward Models

Differentiable wave-optics simulators employ first-principles modeling of the propagation and modulation of optical fields. A typical physical forward model involves:

• Field Construction: The initial field in the input plane is described as $E_0(x) = A(x) \exp[i\phi(x)]$, with $A(x)$ the pupil transmission and $\phi(x)$ the phase profile, parameterized by low–order bases (e.g., Zernike) or pixelwise surface profiles (Desdoigts et al., 2024).
• Propagation Operators: Free-space propagation is handled by Rayleigh–Sommerfeld, Fresnel, or angular spectrum integrals. The field $U(x, y, z)$ at propagation distance $z$ is computed via FFT-based convolutions or transfer functions $H(f_x, f_y; z)$ in the spatial frequency domain (2207.14780, Desdoigts et al., 2024).
• Mask and Lens Modeling: Modulations by thin lenses, metasurfaces, diffractive optical elements (DOE), or multilayer films are implemented as complex transfer functions or locally periodic cell-wise phase/amplitude models (Yang et al., 2024, 2207.14780, Zhu et al., 2023).
• Imaging Formation: The system’s point spread function (PSF) is typically $PSF(x) = |E(x)|^2.$ Shift-variant systems employ local PSF sampling and interpolation; the full sensor image forms via superposition or convolution with the scene (Ho et al., 2024).
• Hybrid Ray–Wave Coupling: For compound or hybrid systems, geometric ray tracing computes the field at critical surfaces, which is then used as the input to downstream wave propagation steps. The ray–wave handoff involves coherent phased deposition of rays onto a grid, followed by scalar diffraction (Yang et al., 2024, Ho et al., 2024).

2. Differentiability and Automatic Differentiation

All constituent modules—geometric operations, phase accumulations, ray–surface intersections, FFTs, convolution, amplitude/phase mask application—are implemented as autodiff-compatible computational graphs in PyTorch, TensorFlow, or JAX. Key features include:

• Exact Gradients: Every step in the pipeline, including OPL (optical path length) accumulation, phase modulation by DOE or metasurfaces, Monte Carlo integration, and convolutions, is fully differentiable without black-box approximations (Ho et al., 2024, Desdoigts et al., 2024).
• Backpropagation: Gradients naturally flow from scalar or vector losses (e.g., image RMSE, perceptual scores, Fisher information) back to device parameters, with efficient batched evaluation and memory management.
• Hybrid Differentiability: Systems coupling both ray and wave stages (e.g., refractive-diffractive hybrids, waveguide displays) maintain differentiability across both domains by ensuring ray-to-wave deposition, polarization transforms, matrix-chain products, and non-sequential sampling steps are autodiff-friendly (Yang et al., 2024, Yang et al., 7 Jan 2026).

3. System-Level Co-Optimization Pipelines

Differentiable wave-optics enables end-to-end learning designs where both the optical configuration and downstream algorithms are co-optimized via stochastic gradient descent. Architectures include:

• Compound Optics Co-Design: Physical optics simulators are paired with vision or reconstruction networks (U-Nets, ResNets, NAFNet). Losses are formed on the final task output (e.g., classification accuracy, scene fidelity), and gradients are propagated through the entire stack (Ho et al., 2024, Yang et al., 2024).
• Calibration and Self-Consistent Inference: Models are calibrated using known (ground-truth) or in-situ data by fitting high-dimensional parameterizations (phase, flat-field, scattering properties) directly with respect to data-driven losses (Page et al., 2020, Desdoigts et al., 2024).
• Fisher Information Optimization: Higher-order autodiff (e.g., Hessian, Jacobian) is leveraged for experimental design, maximizing figures of merit such as minimum variance or determinant of Fisher matrices, crucial for precision astronomical hardware and exoplanet missions (Desdoigts et al., 2024).

4. Empirical Performance, Validation, and Practical Implementation

Empirical studies confirm substantial functional and robustness gains from joint wave-optics-aware optimization:

• Performance Gains: Systems trained with physically accurate wave models outperform ray-only–trained designs by 16–72% on reconstruction metrics and by up to 48.5→23% accuracy on classification when evaluated with wave-exact PSFs. Catastrophic failure is observed for ray-only designs under true wave blur (Ho et al., 2024).
• Robustness: Wave-trained networks avoid overfitting to fragile, high-frequency structures erased by diffraction, instead learning blur-invariant features (Ho et al., 2024).
• Validation: Differentiable simulators match full scalar-diffraction theory at the pixel level, outperforming commercial optomechanical software in handling edge diffraction, DOE discontinuities, and large-angle aberrations (Yang et al., 2024).
• Numerical Efficiency: Acceleration techniques—PSF subsampling, batched FFT-based convolutions, grouped kernel launches—enable tractable runtimes for compound and high-NA systems (Ho et al., 2024, 2207.14780).
• Scalability: Memory and compute strategies, including multi-GPU model/data parallelism, reparameterization tricks, and layer-wise gradient checkpointing, allow optimization of $\sim$10⁶–10⁹ parameters for high-dimensional systems (e.g., telescope phases, metasurface tiles, thin-film stacks) (Desdoigts et al., 2024, Yang et al., 7 Jan 2026).

5. Domains of Application

Differentiable wave optics underpins advances in numerous subfields of computational imaging and instrument design:

• Computational Photography: End-to-end optimization of compact cameras, hybrid refractive-diffractive smartphone optics, and extended depth-of-field imaging (Yang et al., 2024, Ho et al., 2024).
• Microscopy and Bioimaging: Calibration and inverse design of fluorescence microscopes, programmable phase masks, and engineered PSFs for microscopy (Page et al., 2020).
• Metasurface and Flat Optics: Design and co-optimization of nanostructured metasurfaces, with cellwise differentiable RCWA or neural surrogates, for spatial, spectral, and polarization filtering (2207.14780, Zhu et al., 2023).
• Augmented Reality and Display: Topology and parameter optimization of geometric waveguide displays, multi-layer coatings, and non-sequential light transport subject to multi-objective throughput and uniformity metrics (Yang et al., 7 Jan 2026).
• Astronomical Instrumentation: Large-scale phase-retrieval/calibration, diffractive mask design targeting Cramér–Rao–bound precision (e.g., Toliman exoplanet mission), and covariance-constrained system design (Desdoigts et al., 2024, Desdoigts et al., 2024).

6. Methodological Comparisons and Limitations

Differentiable wave-optics models stand apart from classical ray optics and non-differentiable simulation pipelines:

• Accuracy: Only rigorous wave-optics modeling correctly handles diffraction, edge effects, discontinuous phase, and high-NA aberrations. Pure ray-tracing and paraxial models fail in these regimes (Ho et al., 2024, Yang et al., 2024).
• Hybrid Modeling: Coupling of ray and wave stages extends tractability to complex hybrid systems but requires careful management of field handoff and differentiable deposition (Yang et al., 2024, Zhu et al., 2023).
• Computational Cost: Wave-based simulations are 2–4× slower per PSF sample than ray only, and memory demand is substantial (single 10×10 RGB PSF grid with autodiff can reach 35 GB) (Ho et al., 2024, Yang et al., 2024).
• Numerical Stability: Double-precision arithmetic and tailored precision casting are often essential to maintain stability for phase-sensitive backpropagation (Yang et al., 2024).
• Domain-Specific Limitations: Certain approaches (e.g., D-Flat for metasurfaces, ray–wave frameworks for hybrid diffractive-refractive assemblies) may be limited to specific topology or DOE placement.

7. Future Prospects and Extensions

The landscape of differentiable wave optics continues to expand:

• Full-Maxwell Solvers: Ongoing work incorporates direct Maxwell equations solvers and RCWA-based cell libraries for even higher-fidelity modeling of nanophotonic effects (2207.14780).
• Topology Optimization: Layer-pruning and continuous-to-discrete parameterization enable fully automated discovery of optimal multilayer or hybrid structures, including manufacturability constraints (Yang et al., 7 Jan 2026).
• Hierarchical and Multi-Physics Coupling: Emergent frameworks integrate thermal, polarization, and nonlinear effects, as well as multi-modal coupling between optics and electronics.
• Machine Learning Synergy: Neural networks are leveraged as both surrogate optical solvers and downstream restoration modules, enabling superior performance in high-noise, high-aberration regimes (2207.14780, Ho et al., 2024).
• Open-Source and Reproducibility: Packages such as dLux (JAX), D-Flat (TensorFlow), and WaveBlocks (PyTorch) are public, facilitating rapid experimentation and methodological cross-pollination (Desdoigts et al., 2024, 2207.14780, Page et al., 2020).

Differentiable wave-optics frameworks thus now underpin state-of-the-art research in system co-design, inverse problems, and automated optical engineering, establishing a rigorous computational foundation for next-generation optical sciences.