Unified, Verifiable Neural Simulators for Electromagnetic Wave Inverse Problems

TL;DR

UCMax presents a unified, verifiably accurate neural surrogate for electromagnetic scattering that scales to thousands of degrees of freedom and arbitrary illumination, using attentional multi-conditioning and non-recurrent intermediate-state supervision. It offers $O(1)$-time predictions for intermediate timesteps and a computable inference-time error bound, enabling robust performance guarantees. Trained on randomized data, UCMax generalizes to optical tomography, beam shaping in random media, and freeform photonic inverse design, delivering up to $96\%$ speedups with accuracy comparable to full-wave FDTD. The approach extends to time-domain PDEs beyond electromagnetics and provides a practical blueprint for verifiable neural surrogates in complex wave problems.

Abstract

Simulators based on neural networks offer a path to orders-of-magnitude faster electromagnetic wave simulations. Existing models, however, only address narrowly tailored classes of problems and only scale to systems of a few dozen degrees of freedom (DoFs). Here, we demonstrate a single, unified model capable of addressing scattering simulations with thousands of DoFs, of any wavelength, any illumination wavefront, and freeform materials, within broad configurable bounds. Based on an attentional multi-conditioning strategy, our method also allows non-recurrent supervision on and prediction of intermediate physical states, which provides improved generalization with no additional data-generation cost. Using this O(1)-time intermediate prediction capability, we propose and prove a rigorous, efficiently computable upper bound on prediction error, allowing accuracy guarantees at inference time for all predictions. After training solely on randomized systems, we demonstrate the unified model across a suite of challenging multi-disciplinary inverse problems, finding strong efficacy and speed improvements up to 96% for problems in optical tomography, beam shaping through volumetric random media, and freeform photonic inverse design, with no problem-specific training. Our findings demonstrate a path to universal, verifiably accurate neural surrogates for existing scattering simulators, and our conditioning and training methods are directly applicable to any PDE admitting a time-domain iterative solver.

Figures (5)

• Figure 1: Diagram of the UCMax model and conditioning approach. (a) Angle, wavelength, and timestep information are encoded, then fed into the model along with refractive index information. The orange boxes indicate a sinusoidal encoding layer followed by an MLP. (b) Model architecture. Each labeled square denotes a convolutional residual block followed by an attentional computation, with channel dimensions equal to a multiple of hyperparameter $X$, here chosen as 64. Light blue boxes incorporate a downsampling operation, green incorporate upsampling, and pink incorporate neither. (c) Predictions for all inputs are computed in a single batch, with no recurrence.
• Figure 2: (left) Diagram of the error bound calculation method. (right) The error upper bound (in blue) and true error (in orange) at $T=300$ for different values of $K$. Note that the bound is asymptotic to the true error as $K$ approaches $T$.
• Figure 3: (top) Diagram of tomographic reconstruction. An initialization is iteratively updated based on tomographic readings from a ground truth simulation. (a-c) Target refractive indices for microspheres, multi-material nanopillars, and a randomized multi-material mosaic. (d-f) FDTD reconstructed refractive indices. (g-i) UCMax reconstructed refractive indices.
• Figure 4: (top) Diagram of the inverse design process. (a-c) Refractive indices for the FDTD-designed spectral splitter, EM fields of the light with $\lambda=450$nm coupling to +80 degrees, and EM fields of light with $\lambda = 500$nm coupling to -80 degrees, respectively. (d-f) The same images for a UCMax-designed splitter. (i-k) Results for an FDTD-designed multi-wavelength coupler, with $\lambda=450$nm and $\lambda=500$nm both coupled to +80 degree outputs (j-m) The same results for a UCMax-designed coupler. Graphs (g), (h), (o), and (p) compare the electromagnetic field amplitudes along a vertical strip one voxel to the right of the designed devices, showing target field (green), FDTD-designed field (orange), and UCMax-designed field (blue).
• Figure 5: (top) Diagram of the wavefront shaping process. Note that the amplitudes $A$ and the phases $\phi$ are the only parameters being learned. (a-c), (d-f), (g-h), and (j-k) show the average power of the electromagnetic fields for an FDTD-designed source, UCMax-designed source, and a graph that compares optical power at a vertical strip positioned at the chosen focal location, respectively. Runtimes are noted for each design.