High-performance, differentiable light diffraction for optimization and ML-driven optical systems

Abstract

Accurate free-space diffraction modeling is essential for modern microscopy, lithography, and holography. However, as optical systems grow in complexity, a gap remains between traditional simulation tools and modern machine learning pipelines requiring millions of differentiable forward propagations. We introduce a unified, fully differentiable framework for light diffraction that enables high-throughput propagation with fine-grained parameter control. By leveraging the Chirp Z-Transform and Bluestein algorithm, our approach allows for sub-pixel accurate propagation to rotated regions of interest and differentiable rotated cross-section evaluation. All parameters are differentiable and batch-evaluable, enabling real-time x-z simulation and robust phase retrieval. Optimized via Tensor Core-accelerated complex convolutions and kernel fusion, the framework achieves a 3× speedup and 2.5× memory reduction. This work provides a high-performance foundation for large-scale optical engineering and next-generation, optimization-driven optical systems.