A Differentiable Wave Optics Model for End-to-End Computational Imaging System Optimization

Abstract

End-to-end optimization, which integrates differentiable optics simulators with computational algorithms, enables the joint design of hardware and software for data-driven imaging systems. However, existing methods usually compromise physical accuracy by neglecting wave optics or off-axis effects due to the high computational cost of modeling both aberration and diffraction. This limitation raises concerns about the robustness of optimized designs. In this paper, we propose a differentiable optics simulator that accurately and efficiently models aberration and diffraction in compound optics and allows us to analyze the role and impact of diffraction in end-to-end optimization. Experimental results demonstrate that compared with ray-optics-based optimization, diffraction-aware optimization improves system robustness to diffraction blur. Through accurate wave optics modeling, we also apply the simulator to optimize the Fizeau interferometer and freeform optics elements.

Simulation Pipeline

Where we differ from prior art, and how our differentiable wave-optics pipeline works.

Prior approaches

Ray-only end-to-end: Fast and scalable, but ignores diffraction and off-axis phase—designs can be brittle under aperture/NA changes.
Paraxial / stationary PSF models: Capture limited diffraction under simplifying assumptions; accuracy degrades in compound optics and with field curvature.
ASM-based propagation: Efficient in the far/near field, but accuracy becomes sensitive when wavefront phase varies rapidly due to strong aberrations or tight sampling constraints.

Our approach: We combine aberration-aware ray tracing (for precise wavefront/phase on a reference sphere) with Rayleigh–Sommerfeld free-space propagation for physically accurate PSFs, while retaining differentiability and off-grid sampling on the wavefront map.

Rays from a point source o are traced to the reference sphere at the exit pupil (XP), yielding intersections {ρ_i} and a phase map on the wavefront. We then propagate this complex field to the sensor via Rayleigh–Sommerfeld integration to form a PSF, which feeds the image formation, loss, and gradient steps.

Results: single-figure comparison

Array view comparing reference PSFs with competitors and ours.

Grid comparison of PSFs: Reference vs Ray-only vs Hybrid vs Ours

PSFs rendered by different simulators under different conditions. Unlike existing simulators, ours avoids on-grid wavefront discretization and remains robust to defocus and large FoVs, achieving the highest accuracy and efficiency. The tuple (SSIM, ray count, time in sec.) highlights the best performance in red. As the Airy disk does not use ray-tracing, we skip its ray count and do not compare its time with others. Zoom in for details.

Simulator gallery: lens configurations, PSFs & images

Two different lenses (spherical vs. aspheric). Each row shows configuration → on-axis PSF → off-axis PSF → rendered image.

Lens configuration	On-axis PSF	Off-axis PSF	Rendered image
Spherical configuration	On-axis PSF	Off-axis PSF	Rendered image
Aspheric configuration	On-axis PSF	Off-axis PSF	Rendered image

Evolving lens architectures over iterations

Reference scene alongside lens configuration and recovered image animations.

Reference scene

Ground-truth/target image used for evaluation.

Lens layout evolution

Optimizer morphs lens groups over iterations.

Image recovery evolution

Reconstruction quality improves across iterations.

Speed Muted Sync

Acknowledgments

This project builds on the excellent work of the DeepLens and DiffOptics frameworks developed by Xinge Yang et al.. Their differentiable ray–wave model provided the foundation for extending this project toward end-to-end optimization of computational imaging systems. We gratefully acknowledge their contributions and encourage readers to also cite their work:

@article{yang2024end,
  title={End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model},
  author={Yang, Xinge and Souza, Matheus and Wang, Kunyi and Chakravarthula, Praneeth and Fu, Qiang and Heidrich, Wolfgang},
  journal={arXiv preprint arXiv:2406.00834},
  year={2024}
}

BibTeX

@article{ho2024differentiable,
  title={A Differentiable Wave Optics Model for End-to-End Computational Imaging System Optimization},
  author={Ho, Chi-Jui and Belhe, Yash and Rotenberg, Steve and Ramamoorthi, Ravi and Li, Tzu-Mao and Antipa, Nicholas},
  journal={arXiv preprint arXiv:2412.09774},
  year={2024}
}