Skip to content

4 Pillars for Next-Generation AI Optical Design: From Differentiable Optics to Agentic Automation

By Xinge Yang · June 2026

For a long time I set myself no fixed research goal, choosing instead to explore as many directions as possible. Recently, AI agents have completely reshaped my workflow, which has also freed up time to step back: to summarize the work I've already done, and to think about what changes AI could bring to optical design and computational imaging as a whole. At some point the road ahead for next-generation optical design came into focus: from differentiable optics, to computational imaging, to high-performance computing, to agentic design.

Differentiable optics provides better optimization capability; computational imaging offers a new evaluation and design paradigm; high-performance computing makes every module run faster; and the agent automates the entire design flow. Fusing them into a single framework to realize fully automated next-generation optical design, and eventually optical design that self-evolves, is no longer a fantasy, but a natural outcome waiting to happen.

So I decided to write this article. It will mostly use lens design as the running example (since that's the area I've worked on most), but the ideas generalize readily to optical design at large, such as AR/VR display optics.

Differentiable Optics

By building a differentiable simulation of the optical system (ray tracing, wave propagation) and using reverse-mode gradient backpropagation, we can obtain the gradient of a design objective with respect to all system parameters. Compared with the forward finite-difference estimates used in traditional optical design software, gradient backpropagation is not only far more efficient (all gradients in a single backward pass) but also numerically more accurate, since it avoids the truncation and cancellation error of finite differences and is exact to machine precision wherever the simulation is smooth and correctly differentiated (we return to the non-smooth cases below). Reverse-mode backpropagation is not the only option here; other methods, for example the adjoint method, offer an alternative way to obtain the gradients.

Once we have the gradient information, we can apply gradient-based algorithms to optimize the optical system: first-order methods such as Adam, or Gauss–Newton-type methods such as damped least squares. On top of these, a range of optimization strategies further help, for example regularization losses to constrain manufacturability and curriculum learning to overcome the difficulty of designing from a flat-surface start. Differentiable optics has already shown remarkable optimization capability, for example automatically designing highly complex optical systems.

Implementing reverse-mode differentiable optics is actually fairly straightforward: once the physical simulation (geometric optics or wave optics) is built in an automatic-differentiation framework (such as PyTorch or JAX), gradient backpropagation comes for free. Among the open-source differentiable optics frameworks available today, our DeepLens is one of the representative works. DeepLens already provides a complete differentiable optical simulation, including geometric ray tracing, coherent ray tracing, and polarization tracing, covering both sequential and non-sequential models, as well as the transition from ray tracing to wave optics.

  • Wang, C., Chen, N., & Heidrich, W. (2022). dO: A differentiable engine for Deep Lens design of computational imaging systems. IEEE Transactions on Computational Imaging, 8, 905–916. https://doi.org/10.1109/TCI.2022.3212837
  • Yang, X., Fu, Q., & Heidrich, W. (2024). Curriculum learning for ab initio deep learned refractive optics. Nature Communications, 15, 6572. https://doi.org/10.1038/s41467-024-50835-7
  • Yang, X., Souza, M., Wang, K., Chakravarthula, P., Fu, Q., & Heidrich, W. (2024). End-to-End Hybrid Refractive-Diffractive Lens Design with Differentiable Ray-Wave Model. SIGGRAPH Asia 2024 Conference Papers. https://doi.org/10.1145/3680528.3687640 (arXiv:2406.00834)

One difficulty in implementing differentiable optics is memory blowup, but a series of engineering methods can handle it: applying gradient checkpointing across the whole pipeline, introducing lookup tables locally to reduce computation, or simply using multi-GPU parallelism. As long as you understand the autodiff framework, these engineering tricks are not hard to implement.

Another challenge is gradient continuity. Common operations (e.g., intersection, refraction, reflection) are continuous functions with physically meaningful gradients, but some operations carry no gradient information, which directly truncates gradients during backpropagation. These include binning (a ray landing in a pixel), branching (a ray randomly choosing a refraction or reflection path), and filtering (a ray clipped by an aperture).

For these operations, one common approach is reparameterization. For example, binning can be designed so that a ray's energy is distributed across the four surrounding pixels, with each share determined by the ray's distance to that pixel; for the originally non-differentiable random sampling in branching, importance sampling turns it into optimizable ray-energy information; and aperture clipping can be expressed as an integral of energy over a region.

  • Yang, X., Liu, Z., Nie, Z., & Heidrich, W. (2026). End-to-End Differentiable Design of Geometric Waveguide Displays. arXiv:2601.04370.
  • Teh, A., Gkioulekas, I., & O'Toole, M. (2024). Aperture-Aware Lens Design. ACM SIGGRAPH 2024 Conference Papers. https://doi.org/10.1145/3641519.3657398
  • Teh, A., Vicini, D., Bickel, B., Gkioulekas, I., & O'Toole, M. (2025). Automated design of compound lenses with discrete-continuous optimization. SIGGRAPH Asia 2025 Conference Papers. https://doi.org/10.1145/3757377.3763850 (arXiv:2509.23572)

Another approach is interpolation. For example, if we place the aperture at two positions d1 and d2 and compute two objective values loss1 and loss2, we can use finite differences to "estimate" the gradient of the objective with respect to aperture position. But even when d1 and d2 are very close, loss1 and loss2 can still differ greatly, because of the high non-convexity of the objective landscape: we cannot guarantee that d1 and d2 fall in the same region, and a tiny perturbation may jump the system from one region to another. Driving the separation between d1 and d2 ever closer to zero is theoretically valid, but numerical error makes the finite-difference gradient estimate increasingly biased.

People working on AI4Science like to use neural networks to fit the objective function. This is also a form of interpolation: in this example, store all the possible (d, loss) data pairs, train a network, and then backpropagate through the network to compute the gradient of the network's output (loss) with respect to its input (d). But network fitting has two problems:

(1) Network fitting often cannot be 100% accurate, and physical simulation is extremely sensitive to error. A 1% error may count as excellent on many metrics, but is unacceptable in some scenarios. For instance, if you asked an optical engineer to fabricate a lens but told them you only know the radius of curvature lies somewhere between 49 mm and 51 mm, they might find it a bit ridiculous.

(2) The choice of sampling step affects performance. As noted above, because the system is highly non-convex, even a small change in d may cause a dramatic change in loss, so what sampling step should you choose, and how much training and test data should you generate? The network will simply fit whatever data you give it, indifferently, and will not tell you whether the dataset is reasonable.

For these reasons I have always taken a very conservative attitude toward using neural networks: always keep rigorous physical simulation as a fallback to verify the results, and only use network fitting when you are sure the objective varies slowly, the problem is simple, or the accuracy requirement is limited.

  • Tseng, E., Mosleh, A., Mannan, F., St-Arnaud, K., Sharma, A., Peng, Y., Braun, A., Nowrouzezahrai, D., Lalonde, J.-F., & Heide, F. (2021). Differentiable Compound Optics and Processing Pipeline Optimization for End-to-end Camera Design. ACM Transactions on Graphics, 40(2), 18:1–18:19. https://doi.org/10.1145/3446791
  • Yang, X., Fu, Q., Elhoseiny, M., & Heidrich, W. (2023). Aberration-Aware Depth-from-Focus. ICCP 2023 / IEEE Transactions on Pattern Analysis and Machine Intelligence. https://doi.org/10.1109/TPAMI.2023.3301931
  • He, F., Zhao, D., Xu, H., Quan, T., & Zeng, S. (2025). Simulating Dual-Pixel Images From Ray Tracing For Depth Estimation. ICCV 2025. (arXiv:2503.11213)

Computational Imaging

Open your phone and take a photo, and you'll notice the picture in your album has gone through obvious post-processing (e.g., sharpening or brightening). Optical design has never been an isolated matter; we must consider image quality and the subsequent computational processing. Strictly speaking, optical imaging without subsequent computation does not exist: even a simple ISP (demosaicing, white balance, gamma correction, and so on) involves a great deal of image computation. A traditional ISP and a neural image-processing pipeline (neural ISP) are not fundamentally different; the neural ISP is simply more compute-heavy, which currently limits its use on edge devices (such as phones, smart glasses, and XR devices).

Since we're discussing next-generation optical design, let's assume every camera is already equipped with a neural ISP. We then quickly realize that the optics need not be designed so perfectly, because the network can boost image quality and compensate for some residual optical aberrations. At this point, optical design is no longer about optimizing lens parameters alone, but about jointly optimizing the whole system (optics, ISP, and image-processing network) with the final image output as the objective. This is end-to-end optimization. Hardware real estate is precious, yet users' demand for image quality keeps rising; once hardware gains hit a wall, end-to-end optimization becomes a direction optical design must explore.

The optimized system will ultimately be fabricated, so during design we must ensure that the image simulation matches the real capture as closely as possible; only then does our design have reference value. Many physical experiments have validated the accuracy of optical simulation, but image simulation falls more into the domain of computer graphics. In graphics rendering, we strive for results that are as realistic as possible, not only visually realistic, but physically correct (e.g., in energy, spectrum, and polarization); there are still many gaps in this research. Oh, and we also need the entire rendering to be differentiable, so that we can optimize via gradient backpropagation. There is excellent prior work in differentiable rendering to draw on, such as the Mitsuba 2/3 series, but meeting our requirements still demands a great deal of work.

In computational imaging, let's set aside complex 3D scenes and rendering methods for now and consider a simplified problem: given an image (the object image) and its corresponding depth map (a 2.5D scene), how do we simulate the image the camera captures? This is somewhat like a rasterization rendering process, except that each pixel of the object image must first be splatted by its own point spread function (Wikipedia says the captured image is the convolution of the object image and the PSF, but for a spatially varying PSF, splatting is the correct operation), and we must also account for the occlusion effects of a real lens.

PSF splatting addressed the optical-aberration part of image simulation; combining this with an inverse ISP (since the object image is usually an sRGB image), noise simulation, and the ISP, we obtain a physically realistic image simulation. If the simulated image is physically correct, then we can train an image-processing network directly and generalize it to real scenes without fine-tuning, closing the sim2real gap. Representative works include image denoising, optical aberration removal, and defocus deblurring.

By designing every step of the image-simulation process (PSF computation, ISP, inverse ISP, and so on) as a continuous function, the entire process becomes differentiable. After simulating the camera's captured image, we feed it to the downstream network for processing to obtain the desired output, for example object detection, or image restoration/enhancement. Building the loss function on the final output lets us support end-to-end optimization, jointly designing the network, the ISP, and the lens optics.

  • Brooks, T., Mildenhall, B., Xue, T., Chen, J., Sharlet, D., & Barron, J. T. (2019). Unprocessing Images for Learned Raw Denoising. CVPR 2019, 11036–11045.
  • Wei, K., Fu, Y., Yang, J., & Huang, H. (2020). A Physics-Based Noise Formation Model for Extreme Low-Light Raw Denoising. CVPR 2020.
  • Chen, S., Feng, H., Pan, D., Xu, Z., Li, Q., & Chen, Y. (2021). Optical Aberrations Correction in Postprocessing Using Imaging Simulation. ACM Transactions on Graphics, 40(5), 192:1–192:15. https://doi.org/10.1145/3474088 (arXiv:2305.05867)
  • Yang, X., Nguyen, C., Wang, W., Kang, K., Heidrich, W., & Li, X. (2025). Efficient depth- and spatially-varying image simulation for defocus deblur. arXiv:2507.00372.

With image simulation solved, designing the end-to-end optimization strategy is the next challenge. Neural networks typically have millions to billions of parameters, whereas optical systems usually have only dozens to hundreds, which leads to a gradient mismatch between the two parts: a bad optical system produces meaningless images that wreck network training; and if the network cannot tell the optical system "what kind of input it wants," then updating the optics is meaningless. Early works typically did one of the following:

(1) Use stronger priors. For example, pre-designing an optical parameterization suited to a particular task (such as region-partitioned DOEs for spectral encoding), introducing optical loss terms, and pretraining the optics and network before only fine-tuning them, so that end-to-end optimization starts close to a local optimum. But such methods often struggle to escape local optima, and the overall performance gain is limited.

(2) Choose a simpler model, for example a single optical element and a single field-of-view PSF, paired with a very small network, to keep the whole system as simple as possible. With fewer design degrees of freedom, end-to-end optimization becomes feasible. But overly simple models often lack practicality in the real world and struggle to demonstrate a clear advantage.

Another idea is to fix the network weights and optimize only the optics, letting the network decide what image quality is best suited to the downstream task. Because a trained network already encodes a great deal of prior knowledge about the physical world, optimizing the network tends to make it forget those priors. This is the core idea of TaskLens. With the network fixed, the gradients passed to the optics tend to be more stable and more physically interpretable, which raises the degrees of freedom of the whole end-to-end optimization: we can drop the reliance on priors and explore a much larger design space.

By analyzing optimization results across different networks and tasks, we can also distill which kinds of PSFs are better suited to which tasks. For image-restoration tasks, for example, we can analyze which kinds of optical aberrations the network can easily handle, and then selectively optimize different aberrations. In our TaskLens work, we found that a PSF with a pronounced central peak but a long tail can match the visual performance of a perfectly focused PSF while demanding far less of the optics (no need for many lens elements to eliminate residual aberrations). Although we have already distilled some empirical rules, is there a more universal criterion? That still requires further exploration.

  • Yang, X., Fu, Q., Nie, Y., & Heidrich, W. (2026). Task-Driven Lens Design. Optics Express, 34(5), 8961. https://doi.org/10.1364/OE.588912
  • Zhou, J., Chen, S., et al. (2024). Revealing the preference for correcting separated aberrations in joint optic-image design. Optics and Lasers in Engineering. (arXiv:2309.04342)
  • Côté, G., Mannan, F., Thibault, S., Lalonde, J.-F., & Heide, F. (2023). The Differentiable Lens: Compound Lens Search over Glass Surfaces and Materials for Object Detection. CVPR 2023, 20803–20812. (arXiv:2212.04441)

High-Performance Computing

By now, the amount of computation we need has grown substantially compared with traditional optical design: optical simulation, gradient backpropagation, image simulation, end-to-end optimization. These computations constantly remind us that things have slowed down again, and that the GPU is about to run out of memory. Earlier we mentioned reducing memory overhead through engineering tricks such as gradient checkpointing, but before reaching for any of those, we should ask a more fundamental question: which step is really inflating the memory overhead?

For neural networks, having millions (10^6) of parameters and on the order of 10^12 floating-point operations is not unusual, yet modern GPUs are built to run these computations at high speed. Differentiable optics also requires a large number of ray–surface intersection computations, but the total computation falls far short of the 10^12 scale. And yet, in practice, the intuitive feeling is that, even without considering the downstream image simulation and network processing, differentiable optics is still noticeably slower and consumes a lot of GPU memory. Where does the problem lie?

Common differentiable frameworks (such as PyTorch and JAX) have excellent operator support for neural networks (all sorts of GPU kernels), so that even complex computation runs blazingly fast on the GPU while consuming limited memory. For physical simulation, however, GPU operator support is very limited: often only the simplest add/subtract/multiply/divide operations are available, so a single function frequently needs to load dozens of GPU kernels to complete, and backpropagation builds an enormous computational graph, making the overall memory overhead hard to bear.

Kernel fusion is a very effective solution: fusing multiple computation steps inside a single kernel reduces data-transfer and compute overhead, noticeably accelerating computation and lowering memory use. In our AutoLens project repository, a community member contributed a Triton-based kernel-fusion implementation that brought a lens-design example from the original several hours down to about ten minutes.

Dr.Jit is excellent work, especially for differentiable ray tracing: it can fuse the entire computation into a single megakernel, achieving substantial speedups and memory savings. But for the best performance, you still need deeper kernel customization that combines algorithmic characteristics with the hardware architecture. In addition, Dr.Jit's support for physical-optics computation, its Python interface, and cross-platform support (not only NVIDIA, but also AMD and other GPU platforms) all require extra implementation.

Our latest work finds that, by exploiting the characteristics of differentiable optics and applying a series of high-performance computing methods through custom GPU kernels, the entire end-to-end pipeline can achieve over 10x speedup and 90% memory savings.* This lets differentiable optics run on a laptop, faster and more capable than existing optical software, providing a foundation for the next generation of optical design tools. At the same time, high-performance computing brings automatic lens design down to the minute level and makes near-interactive stray light analysis and near-interactive Monte Carlo tolerance analysis possible, which will, to some degree, change the entire optical design workflow.

(* These figures are based on our as-yet-unpublished work.)

Agentic Design

Optical design is not a simple lens-optimization problem, but a long chain that includes optimization, component selection, manufacturability verification, image-quality analysis, and more. From the differentiable-simulation perspective, we could integrate all the software into a single framework for end-to-end computation and optimization, which is theoretically feasible. But from another angle, simply letting an agent orchestrate each tool or piece of software and handle these problems automatically is a simpler and smarter approach.

If you think about an optical engineer's work, it too is a continual process of inferring the next step from various analysis results, in order to reach a final design that meets the required specifications. I believe automating this process with an agentic flow is a blue ocean we must explore: one that runs continuously over long stretches, drawing on design-experience documents and historical design records to arrive at the final design.

Recently I've seen some optical software add AI-plugin features that, paired with optical design documents and using either RAG or fine-tuning, provide better optical-knowledge Q&A. But Q&A offers very limited help to productivity; using an agent to automate the entire workflow is what will truly transform the industry. Early steps in this direction are beginning to appear, for example OptiAgent, a physics-driven agentic framework reported as the first attempt to apply LLMs to optical design.

  • Geng, Sun, Gao, Hu, Yi, Qian, Hu, Bai, & Wang (2026). OptiAgent: A Physics-Driven Agentic Framework for Automated Optical Design. arXiv:2602.23761.

However, pairing an existing agentic framework with traditional optical design algorithms tends to stall at some stage and fall into a dead loop. To escape such dead loops, the following methods are worth trying:

(1) Better optimization capability, escaping local optima and stalled regions directly through optimization (this is exactly the advantage of differentiable optics and our curriculum-learning method);

(2) Better heuristic operations, referencing optical design-experience documents, analyzing current lens performance, and proposing the next operation, such as adding or removing a lens element to reach a more tractable space;

(3) Smarter agent capability, launching multiple sub-agents simultaneously to explore the optimization potential of several design directions, and reinforcing the design process by analyzing historical memory.

Another known problem is that, within the agentic loop, the optical system changes uncontrollably, for example, the number of lens elements keeps growing, eventually crashing some tool and, in turn, the entire workflow. This calls for a better harness that, combined with engineering experience, keeps the agent operating within set boundaries, making the whole optimization process more controllable and reliable.

We need not only a better agentic framework, but also better tools for it to use, and the differentiable optics, computational imaging, and high-performance computing above are exactly about building better tools to enhance the agent's capability: differentiable optics provides better optimization capability, computational imaging offers a new evaluation and design paradigm, and high-performance computing makes the whole flow run faster. The agent orchestrates them together with other tools and software, and makes better decisions informed by optical design experience. Although there is no mature work yet, the overall road has become clear.

From AI-native tools to AI-native optical design software, perhaps it is time to fully reform the entire design paradigm, and even the interaction paradigm: users no longer rely on fixed UIs and menus, but directly describe their optical design goals, while the agent automatically plans and executes the path, and the system's core shifts from a control interface to Memory + Context + Runtime. Beyond this, training the agent's behavior to imitate an experienced optical engineer, ultimately reaching the goal of a self-evolving agent, opens up boundless room for imagination. Perhaps this is what truly next-generation optical design looks like.