Diffusion Schrödinger Bridges with enhanced posterior sampling for metasurface inverse design

Metasurface inverse design is challenged by the intricate relationship between structural parameters and electromagnetic responses, as well as the high dimensionality of the optimization space. Local models, while commonly employed, quickly become infeasible for complex and locally coupled structures. Conventional iterative optimization techniques, on the other hand, are computationally intensive, time-consuming, and susceptible to convergence in local minima. This study explores a versatile generative methodology based on enhanced posterior sampling within the Schrödinger Bridge framework. By decomposing posterior sampling into amplitude and directional contributions, we effectively integrated different kind of posterior sampling. This approach is further supported by refined training strategies to enhance performance and reduce the complexity of hyperparameter optimization. The proposed framework demonstrates exceptional accuracy and robustness, representing a significant advancement in metasurface design. Notably, it enables high-precision inverse design for large-scale configurations of up to $350 \times 350$ pillar arrays, despite being trained on significantly smaller arrays of $23 \times 23$ pillars.

💡 Research Summary

This paper tackles the challenging inverse design problem of metasurfaces, where the goal is to find a structural configuration that yields a prescribed electromagnetic far‑field response. Traditional local models and Bayesian optimization quickly become infeasible as the number of design parameters grows into the hundreds or thousands, and conventional iterative methods are computationally expensive and prone to local minima. To overcome these limitations, the authors adopt the Diffusion Schrödinger Bridge (DSB) framework—a recent generalization of diffusion models that solves an entropy‑regularized stochastic optimal transport problem, thereby connecting an initial distribution to a target distribution through a time‑continuous diffusion process.

A key innovation is the introduction of enhanced posterior sampling strategies that decompose the guidance term into amplitude and directional components. Three novel techniques are presented: (1) Robust Posterior Sampling, which combines a central loss (distance between the desired far‑field and a surrogate predictor) with a robustness loss that penalizes sensitivity to small perturbations in the design parameters; (2) Spherical Gaussian constraint, which normalizes the magnitude of the guidance vector while preserving its direction, mitigating the Jensen’s gap that grows with dimensionality; and (3) Ring Gaussian constraint, which similarly bounds the amplitude but allows a ring‑shaped distribution of guidance magnitudes to encourage diversity. The authors also propose a weighted averaging scheme for Monte‑Carlo samples that emphasizes samples close to the current mean, further stabilizing the optimization.

Training incorporates a consistency loss and conditions the score network on the target far‑field amplitude. Unlike conditional diffusion models that inject the condition at every denoising step, DSBs embed the condition only once at the initial state, yet the authors demonstrate that adding the condition also to the score prediction improves performance. Experiments on a beam‑shaping metasurface composed of dielectric pillars (spacing λ/2, variable radii) show that the proposed DSB with enhanced posterior sampling achieves very high R² values (up to 0.966 for a 23 × 23 pillar array) when evaluated with full‑wave FDTD simulations. Importantly, the model trained on 23 × 23 arrays generalizes to much larger designs—98 × 98 and even 350 × 350 pillar configurations—while maintaining R² above 0.94, indicating strong scalability.

A comprehensive ablation study compares raw posterior sampling, Monte‑Carlo sampling, robust sampling, and the two Gaussian‑constrained variants, both within the DSB framework and against standard diffusion models. Across all metrics, DSBs consistently outperform diffusion models, and the constrained sampling methods provide the best trade‑off between accuracy and computational cost. While robust and Monte‑Carlo sampling increase memory and runtime linearly with the number of samples, the authors find that using 16–32 samples per step offers a practical balance.

Overall, the work demonstrates that diffusion‑based Schrödinger bridges, when equipped with carefully designed posterior sampling and amplitude constraints, can efficiently solve high‑dimensional inverse design problems in photonics. The ability to train on modest datasets yet reliably generate large‑scale metasurface layouts opens the door to rapid prototyping and potentially real‑time design loops for advanced optical components. Future directions include extending the framework to multi‑physics metasurfaces, incorporating fabrication tolerances directly into the robustness term, and deploying the method in an end‑to‑end design‑fabrication pipeline.