Active learning of deep surrogates for PDEs: Application to metasurface design

Surrogate models for partial-differential equations are widely used in the design of meta-materials to rapidly evaluate the behavior of composable components. However, the training cost of accurate surrogates by machine learning can rapidly increase with the number of variables. For photonic-device models, we find that this training becomes especially challenging as design regions grow larger than the optical wavelength. We present an active learning algorithm that reduces the number of training points by more than an order of magnitude for a neural-network surrogate model of optical-surface components compared to random samples. Results show that the surrogate evaluation is over two orders of magnitude faster than a direct solve, and we demonstrate how this can be exploited to accelerate large-scale engineering optimization.


Incident planewave
optical structures include ultra-compact sensors, imaging, and spectroscopy devices used in cell phone cameras and in medical applications 2 . As the metamaterials become larger in scale and as the manufacturing capabilities improve, there is a pressing need for scalable computational design tools.
In this work, surrogate models were used to rapidly evaluate the effect of each metamaterial components during device design 3 , and machine learning is an attractive technique for such models [4][5][6][7] . However, in order to exploit improvements in nano-manufacturing capabilities, components have an increasing number of design parameters and training the surrogate models (using brute-force numerical simulations) becomes increasingly expensive. The question then becomes: How can we obtain an accurate model from minimal training data? We present a new active-learning (AL) approach-in which training points are selected based on an error measure ( Fig. 3)-that can reduce the number of training points by more than an order of magnitude for a neural-network (NN) surrogate model of partial-differential equations (PDEs). Further, we show how such a surrogate can be exploited to speed up large-scale engineering optimization by > 100×.
In particular, we apply our approach to the design of optical metasurfaces: large (10 2 -10 6 wavelengths λ) aperiodic nanopattered ( λ) structures that perform functions such as compact lens- Metasurface design can be performed by breaking the surface into unit cells with a few parameters each ( Fig. 1) via domain-decomposition approximations 3, 9 , learning a "surrogate" model that predicts the transmitted optical field through each unit as a function of an individual cell's parameters, and optimizing the total field (e.g. the focal intensity) as a function of the parameters of every unit cell 3 (Sec. 2). This makes metasurfaces an attractive application for machine learning (Sec. 4) because the surrogate unit-cell model is re-used millions of times during the design process, amortizing the cost of training the model based on expensive "exact" Maxwell solves sampling many unit-cell parameters. For modeling the effect of a 1-4 unit-cell parameters, Chebyshev polynomial interpolation can be very effective 3 , but encounters an exponential "curse of dimensionality" with more parameters 10,11 . In this paper, we find that a NN can be trained with orders of magnitude fewer Maxwell solves for the same accuracy with ∼ 10 parameters, even for the most challenging case of multi-layer unit cells many wavelengths (> 10λ) thick (Sec. 5).
In contrast, we show that subwavelength-diameter design regions (considered by several other authors 4-7, 12, 13 ) require orders of magnitude fewer training points for the same number of parameters (Sec. 3), corresponding to the physical intuition that wave propagation through subwavelength regions is effectively determined by a few "homogenized" parameters 14  Ref. 3 introduced an optimization approach to metasurface design using Chebyshev-polynomial surrogate model, which was subsequently extended to topology optimization (∼ 10 3 parameters per cell) with "online" Maxwell solvers 24 . Metasurface modeling can also be composed with signal/image-processing stages for optimized "end-to-end design" 25,26 . Previous work demonstrated NN surrogate models in optics for a few parameters [27][28][29] , or with more parameters in deeply subwavelength design regions 4,12 . As we will show in Sec. 3, deeply subwavelength regions pose a vastly easier problem for NN training than parameters spread over larger diameters. Another approach involves generative design, again typically for subwavelength 6,7 or wavelength-scale unit cells 30 , in some cases in conjunction with larger-scale models 5,12,13 . A generative model is essentially the inverse of a surrogate function: instead of going from geometric parameters to performance, it takes the desired performance as an input and produces the geometric structure, but the mathematical challenge appears to be closely related to that of surrogates. where N is the number of training samples) 37,38 . To our knowledge, the work presented here is the first to achieve training time efficiency (we show an order of magnitude reduction sample complexity), design time efficiency (the actively learned surrogate model is at least two orders of magnitude faster than solving Maxwell's equations), and realistic large-scale designs (due to our optimization framework 3 ), all in one package.

Metasurfaces and surrogate models
In this section, we present the neural-network surrogate model used in this paper, for which we adopt the metasurface design formulation from Ref. 3. The first step of this approach is to divide the metasurface into unit cells with a few geometric parameters p each. For example, Fig. 1(left) shows several possible unit cells: (a) a rectangular pillar ("fin") etched into a 3d dielectric slab 39 (two parameters); (b) an H-shaped hole (four parameters) in a dielectric slab 4 ; or a (c) multilayered 2d unit cell with ten holes of varyings widths considered in this paper. As depicted in Fig. 1(right), a metasurface consists of an array of these unit cells. The second step is to solve for the transmitted field (from an incident planewave) independently for each unit cell using approximate boundary conditions 3,24,39,40 , in our case a locally periodic approximation (LPA) based on the observation that optimal structures often have parameters that mostly vary slowly from one unit cell to the next 3 . (Other approximate boundary conditions are also possible 9 .) For a subwavelength period, the LPA transmitted far field is entirely described by a single number-the complex transmission coefficient t(p). One can then compute the field anywhere above the metasurface by convolving these approximate transmitted fields with a known Green's function, a near-to-farfield transformation 41 . Finally, any desired function of the transmitted field, such as the focal-point intensity, can be optimized as a function of the geometric parameters of each unit cell 3 .
In this way, optimizing an optical metasurface is built on top of evaluating the function t(p) (transmission through a single unit cell as a function of its geometric parameters) thousands or even millions of times-once for every unit cell, for every step of the optimization process. Although it is possible to solve Maxwell's equations "online" during the optimization process, allowing one to use thousands of parameters p per unit cell requires substantial parallel computing clusters 24 .
Alternatively, one can solve Maxwell's equations "offline" (before metasurface optimization) in order to fit t(p) to a surrogate modelt which can subsequently be evaluated rapidly during metasurface optimization (perhaps for many different devices). For similar reasons, surrogate (or "reduced-order") models are attractive for any design problem involving a composite of many components that can be modeled separately 6,7,42 .
The key challenge of the surrogate approach is to increase the number of design parameters, especially in non-subwavelength regions as discussed in Sec. 3.
In this paper, the surrogate model for each of the real and imaginary parts of the complex transmission is an ensemble of J = 5 independent neural networks (NNs) with the same training data but different random "batches" 43 on each training step. Each of NN i is trained to output a prediction µ i (p) and an error estimate σ i (p) for every set of parameters p. To obtain these µ i and σ i from training data y(p) (from brute-force "offline" Maxwell solves) we minimize 31 : over the parameters Θ i of NN i. Equation (2) is motivated by problems in which y was sampled from a Gaussian distribution for each p, in which case µ i and σ 2 i could be interpreted as mean and hetero-skedastic variance, respectively 31 . Although our exact function t(p) is smooth and noisefree, we find that Eq. (2) still works well to estimate the fitting error, as demonstrated in Sec. 4.
Each NN is composed of an input layer with 13 nodes (10 nodes for the geometry parameterization and 3 nodes for the one-hot encoding 43 of three frequencies of interest), three fully-connected hidden layers with 256 rectified linear units (ReLU 43 ), and one last layer containing one unit with a scaled hyperbolic-tangent activation function 43 (for µ i ) and one unit with a softplus activation function 43 (for σ i ). Given this ensemble of J NNs, the final prediction µ * (for the real or imaginary part of t(p)) and its associated error estimate σ * are amalgamated as 31 : Physically, for extremely sub-wavelength volumes the waves only "see" an averaged effective medium 14 , so there are effectively only a few independent design parameters regardless of the number of geometric degrees of freedom. Quantitatively, we find that the Hessian of the trained surrogate model (second-derivative matrix) in the smallest unit-cell case is dominated by only two singular values-consistent with a function that effectively has only two free parameters-with the other singular values being more than 100× smaller in magnitude; for the other two cases, many more training points would be required to accurately resolve the smallest Hessian singular values. A unit cell with large design-volume diameter ( λ) is much harder to train, because the dimensionality of the design parameters is effectively much larger.
4 Active-learning algorithm Here, we present an algorithm to choose training points that is significantly better at reducing the error than choosing points at random. As described below, we select the training points where the estimated model error is largest, given the estimated error σ * from Sec. 2.
The algorithm used to train each of the real and imaginary parts is outlined in Fig. 3 and Algorithm 1. Initially we choose n init uniformly distributed random points p 1 , p 2 , ..., p n init to train a first iterationt 0 (p) over 50 epochs 43 . Then, given the model at iteration i, we evaluatet i (p) (which is orders of magnitude faster than the Maxwell solver) at M K points sampled uniformly at random and choose the K points that correspond to the largest σ 2 * . We perform the expensive Maxwell solves only for these K points, and add the newly labeled data to the training set. We traint i+1 (p) with the newly augmented training set. We repeat this process T times.
Essentially, the method works because the error estimate σ * is updated every time the model is retrained with an augmented dataset. In this way, model tells us where it does poorly by setting a large σ * for parameters p where the estimation would be bad in order to minimize Eq. (2). Order-of-magnitude reduction in training data We compared the fractional errors of a NN surrogate model trained using uniform random samples with an identical NN trained using an activelearning approach, in both cases modeling the complex transmission of a multi-layer unit cell with ten independent parameters (Fig. 4(right)). With the notation of Sec. 4, the baseline corresponds to T = 0, and n init equal to the total number of training points. This corresponds to no active learning at all, because the n init points are chosen at random. In the case of active learning, n init = 2000, M = 4, and we computed for K = 500, 1000, 2000, 4000, 8000, 16000, 32000, 64000, and 128000. Although three orders of magnitude on the log-log plot is too small to determine if the apparent linearity indicates a power law, Fig. 4(left) shows that the lower the desired fractional error, the greater the reduction in training cost compared to the baseline algorithm; the slope of the active-learning fractional error (−0.2) is about 30% steeper that that of baseline (−0.15). The active-learning algorithm achieves a reasonable fractional error of 0.07 in twelve times less points than the baseline, which corresponds to more than one order of magnitude saving in training data (much less expensive Maxwell solves). This advantage would presumably increase for a lower error tolerance, though computational costs prohibited us from collecting orders of magnitude more training data to explore this in detail. For comparison and completeness, Fig. 4(left) shows fractional errors using Chebyshev interpolation (for the blue frequency only). Chebyshev interpolation has a much worse fractional error for a similar number of training points. Chebyshev interpolation suffers from the "curse of dimensionality"-the number of training points is exponential with the number of variables. The two fractional errors shown are for three and four interpolation points in each of the dimensions, respectively. In contrast, NNs are known to mitigate the "curse of dimensionality" 44 .
Application to metalens design We used both surrogates models to design a multiplexer-an optical device that focuses different wavelength at different points in space. The actively learned surrogate model results in a design that much more closely matches a numerical validation than the baseline surrogate (Fig. 5). As explained in Sec. 2, we replace a Maxwell's equations solver with a surrogate model to rapidly compute the optical transmission through each unit cell; a similar surrogate approached could be used for optimizing many other complex physical systems. In the case of our two-dimensional unit cell, the surrogate model is two orders of magnitude faster than solving Maxwell's equations with a finite difference frequency domain (FDFD) solver 45  solvers are much more costly while a surrogate model remains the same.
The surrogate model is evaluated millions of times during a meta-structure optimization.
We used the actively learned surrogate model and the baseline surrogate model (random training samples), in both cases with 514000 training points, and we optimized a ten-layer metastructure with 100 unit cells of period 400 nm for a multiplexer application-where three wavelengths (blue: 405 nm, green: 540 nm, and red: 810 nm) are focused on three different focal spots (−10 µm, 60 µm), (0, 60 µm), and (+10 µm, 60 µm), respectively. The diameter is 40 µm and the focal length is 60 µm, which corresponds to a numerical aperture of 0.3. Our optimization scheme tends to yield results robust to manufacturing errors 3 for two reasons: first, we optimize for the worst case of the three focal spot intensities, using an epigraph formulation 3 ; second, we compute the average intensity from an ensemble of surrogate models that can be thought of as a Gaussian distributiont(p) = µ * (p) + σ * (p) with ∼ N (0, 1), and µ * and σ * are defined in Eq. (3) and Eq. (4), respectively, where G is a Green's function that generates the far-field from the sources of the metastructure 3 .
The resulting optimized structure for the active-learning surrogate is shown in Fig. 5(bottom).
In order to compare the surrogate models, we validate the designs by computing the optimal unit cell fields directly using a Maxwell solver instead of using the surrogate model. This is computationally easy because it only needs to be done once for each of the 100 unit cells instead of millions of times during the optimization. The focal lines-the field intensity along a line parallel to the two-dimensional metastructure and passing through the focal spots-resulting from the validation are exact solutions to Maxwell's equations assuming the locally periodic approximation (Sec. 2). Fig. 5(top) shows the resulting focal lines for the active-learning and baseline surrogate models. A multiplexer application requires similar peak intensity for each of the focal spots, which is achieved using worst case optimization 3 . Fig. 5(top) shows that the actively learned surrogate has ≈ 3× smaller error in the focal intensity compared to the baseline surrogate model. This result shows that not only is the active-learning surrogate more accurate than the baseline surrogate for 514000 training points, but also the results are more robust using the active-learning surrogatethe optimization does not drive the parameters towards regions of high inaccuracy of the surrogate model. Note that we limited the design to a small overall diameter (100 unit cells) mainly to ease visualization (Fig. 5(bottom)), and we find that this design can already yield good focusing performance despite the small diameter. In earlier work, we have already demonstrated that our optimization framework is scalable to designs that are orders of magnitudes larger 46 .
Previous work, such as Ref. 47-in a different approach to active-learning that does not quantify uncertainty-suggested iteratively adding the optimum design points to the training set (re-optimizing before each new set of training points is added). However, we did not find this approach to be beneficial in our case. In particular, we tried adding the data generated from LPA validations of the optimal design parameters, in addition to the points selected by our active learning algorithm, at each training iteration, but we found that this actually destabilized the learning and resulted in designs qualitatively worse than the baseline. By exploiting validation points, it seems that the active learning of the surrogate tends to explore less of the landscape of the complex transmission function, and hence leads to poorer designs. Such exploitation-exploration trade-offs are known in the active-learning literature 32 .

Concluding remarks
In this paper, we present an active-learning algorithm for composite materials which reduces the training time of the surrogate model for a physical response, by at least one order of magnitude.
The simulation time is reduced by at least two orders of magnitude using the surrogate model compared to solving the partial differential equations numerically. While the domain-decomposition method used here is the locally periodic approximation and the partial differential equations are the Maxwell equations, the proposed approach is directly applicable to other domain-decomposition methods (e.g. overlapping domain approximation 9 ) and other partial differential equations or ordinary differential equations 48 .
We used an ensemble of NNs for interpolation in a regime that is seldom considered in the machine-learning literature-when the data is obtained from a smooth function rather than noisy measurements. In this regime, it would be instructive to have a deeper understanding of the relationship between NNs and traditional approximation theory (e.g. with polynomials and rational functions 10,11 ). For example, the likelihood maximization of our method forces σ * to go to zero whent(p) = t(p). Although this allows us to simultaneously obtain a prediction µ * and an error estimate σ * , there is a drawback. In the interpolation regime (when the surrogate is fully determined), σ * would become identically zero even if the surrogate does not match the exact model away from the training points. In contrast, interpolation methods such as Chebyshev polynomials yield a meaningful measure of the interpolation error even for exact interpolation of the training data 10,11 .
In the future, we plan to separate the estimation model and the model for the error measure using a meta-learner architecture 36 , with expectation that the meta-learner will produce a more accurate error measure and further improve training time. We believe that the method presented in this paper will greatly extend the reach of surrogate-model based optimization of composite materials and other applications requiring moderate-accuracy high-dimensional interpolation.

Methods
Training-data computation The complex transmission coefficients were computed in parallel using an open-source finite difference frequency-domain solver for Helmholtz equation 49  in the background. In the normal unit cell, the period of the cell is 400 nm, the height of the ten holes is fixed to 304 nm and their widths varies between 60 nm and 340 nm, each hole is separated by 140 nm of substrate. In the small unit cell, the period of the cell is 400 nm, the height of the ten holes is 61 nm, and their widths varies between 60 nm and 340 nm, there is no separation between the holes. The smallest unit cell is the same as the small unit cell shrunk ten times (period of 40 nm, ten holes of heigth 6.1 nm and width varying between 6 nm and 34 nm). (μ * (p) +σ * (p) )dr Σ G(µ * (p) + σ * (p) )dr , = Ḡμ * Gµ * + 2 Ḡσ * Gσ * + 2 Re Ḡμ * Gσ * ,

Metalens design problem
where the(·) notation denotes the complex conjugate, the notations Σ (·)dr and G(r, r ) are simplified to p, and G, and the notation p( r ) is dropped for clarity. From the linearity of expectation, where we used that E( ) = 0 and E( 2 ) = 1.
Active-learning architecture and training The ensemble of NN was implemented using Py-Torch 51 on a 3.5 GHz 6-Core Intel Xeon E5 processor. We trained an ensemble of 5 NN for each surrogate models. Each NN is composed of an input layer with 13 nodes (10 nodes for the geometry parameterization and 3 nodes for the one-hot encoding 43  and v true is where | · | is the L2-norm for complex vectors.

Data Availability
The data that support the findings of this study are available from the corresponding author upon reasonable request.