End-to-end differentiable design of geometric waveguide displays

TL;DR

This work presents the first end-to-end differentiable optimization framework for geometric waveguide that couples non-sequential Monte Carlo polarization ray tracing with a differentiable transfer-matrix thin-film solver and expands the scope of differentiable optics for next-generation optical design.

Abstract

Geometric waveguides are a promising architecture for optical see-through augmented reality displays, but their performance is severely bottlenecked by the difficulty of jointly optimizing non-sequential light transport and polarization-dependent multilayer thin-film coatings. Here we present the first end-to-end differentiable optimization framework for geometric waveguide that couples non-sequential Monte Carlo polarization ray tracing with a differentiable transfer-matrix thin-film solver. A differentiable Monte Carlo ray tracer avoids the exponential growth of deterministic ray splitting while enabling gradients backpropagation from eyebox metrics to design parameters. With memory-saving strategies, we optimize more than one thousand layer-thickness parameters and billions of non-sequential ray-surface intersections on a single multi-GPU workstation. Automated layer pruning is achieved by starting from over-parameterized stacks and driving redundant layers to zero thickness under discrete manufacturability constraints, effectively performing topology optimization to discover optimal coating structures. On a representative design, starting from random initialization within thickness bounds, our method increases light efficiency from 4.1\% to 33.5\% and improves eyebox and FoV uniformity by $\sim$17$\times$ and $\sim$11$\times$, respectively. Furthermore, we jointly optimize the waveguide and an image preprocessing network to improve perceived image quality. Our framework not only enables system-level, high-dimensional coating optimization inside the waveguide, but also expands the scope of differentiable optics for next-generation optical design.

Figures (4)

• Figure 1: Illustration of geometric waveguide architecture and our proposed differentiable optimization.a The GWG employs partially reflective mirror arrays to redirect light from the display engine to the user's eye pupil. The exit pupil is replicated and expanded relative to the input pupil (display panel). b We establish an end-to-end differentiable simulation spanning PRMA geometry and multilayer coatings on each mirror, enabling system-level optimization. c A representative GWG achieves a 100$\times$ pupil expansion, increasing from a 1.6 mm $\times$ 1.6 mm display panel to 16 mm $\times$ 16 mm at the eyebox with a FoV of 38$^\circ$$\times$ 38$^\circ$. d At each partially reflective mirror, light is either reflected or transmitted. We use differentiable non-sequential Monte Carlo ray tracing to simulate these paths within the GWG. At the eyebox region, output mirror arrays couple light out across a large region. e Each mirror is coated with multilayer thin films that modulate polarization. f A differentiable transfer-matrix solver heavens1960optical computes effective Fresnel coefficients for multilayer coatings. In the forward pass, rays carry gradient information through the sampled paths. In backpropagation, gradients of the eyebox image with respect to film thicknesses are computed automatically, enabling gradient-based optimization.
• Figure 2: Evaluation of end-to-end differentiable optimization for GWG coatings.a Light efficiency and uniformity for the initial design, a genetic-algorithm baseline and our differentiable optimization, measured by eyebox throughput $\bar{I}$ (fraction of input power reaching the eyebox), $\text{CV}_{\text{FoV}}$ and $\text{CV}_{\text{eyebox}}$ (all metrics averaged over five Monte Carlo runs with different random seeds). b Loss curves for differentiable optimization and the genetic-algorithm baseline. Gradient-based optimization converges faster and reaches a lower loss. c Thickness distribution of the optimized multilayer stack. Several layers are driven towards zero thickness, indicating that they can be removed in the final design. The discrete optimization strategy supports an over-parameterized starting stack followed by pruning during optimization. d Simulated pupil response ($256 \times 256$) at the centre of the eyebox, displayed image and RGB channels ($512 \times 512$) for the initial design, the genetic algorithm and differentiable optimization. Differentiable optimization increases brightness and reduces FoV and eyebox non-uniformities.
• Figure 3: Memory-saving strategies for large-scale differentiable optimization of the GWG system.a The pupil image at the eyebox is partitioned into FoV patches and distributed across multiple GPUs. The patches are assembled into a full-FoV image to compute the loss; in backpropagation, gradients are computed on the full image and then scattered back to the corresponding GPUs. b Differentiable Monte Carlo ray tracing uses a two-pass intersection strategy. In the first pass, we compute ray-surface intersections without tracking gradients and record the intersected surfaces. In the second pass, we recompute only those intersections in differentiable mode. c We use gradient checkpointing to decouple backpropagation through ray tracing and through the thin-film solver. Gradients are first backpropagated to the effective Fresnel coefficients and stored; we then backpropagate through the multilayer solver to update layer thicknesses using stored intermediates. d Peak GPU memory usage with these strategies, enabling large-scale differentiable optimization on a single workstation.
• Figure 4: End-to-end co-design of image preprocessing and GWG coatings.a A neural network processes the displayed image before emitted into the GWG. The network and coating parameters are optimized jointly to improve perceived image quality at the eyebox. b PSNR and SSIM measured on validation set. Both coating-only optimization and end-to-end co-design improve image quality, while end-to-end co-design further compensates residual artefacts and improves image quality. c Example outputs. Unoptimized coatings produce dark images with non-uniformities, coating optimization improves brightness and reduces global non-uniformities, and end-to-end co-design further compensates residual artefacts and improves image quality.