A Transformer-Based Approach to Facilitate Inverse Design of Achromatic Metasurfaces

Abstract

Accurate and efficient prediction of the spectral responses of metasurface unit cells is key to intelligent metasurface design. Here, we propose a Shape-integrated Dual-Spectrum-aware transformer (SiDSaT) for forward prediction of metasurface responses. Trained on a large-scale simulation dataset, SiDSaT achieves a phase mean absolute error (MAE) below 0.05 across both cylindrical and cuboidal unit cells, demonstrating strong prediction accuracy and generalization. We further embedded SiDSaT into an inverse design framework and applied it to the design of single-wavelength and broadband achromatic metalenses. Results of focusing performance and dispersion control confirm the effectiveness and versatility of SiDSaT in supporting the high-precision inverse design of metasurface optical devices. This work offers a scalable and accurate approach for intelligent metasurface design across diverse shape configurations and broadband spectral ranges.

1. Introduction

Metasurfaces are composed of sub-wavelength-scale nanostructures and enable precise control of the phase, amplitude, and polarization of light on a planar platform, thereby supporting the development of multifunctional integrated optical devices such as metalenses, holographic elements, and beam shaping systems [1,2,3,4,5,6,7,8,9,10,11,12,13,14]. Despite their significant potential in micro–nano optics, metasurfaces still face critical challenges under broadband conditions due to inherent dispersion issues. Traditional single-layer metasurfaces based on linear phase compensation can correct chromatic aberration within narrow or moderately wide spectral ranges [15,16]. However, these approaches fundamentally rely on linear approximations of monotonic dispersion, which fail to accurately capture the nonlinear dispersion behaviors of materials and nanostructures. As a result, such mismatch leads to wavefront distortions and degraded focusing performance in ultra-broadband scenarios, thus limiting their applicability in wide-spectrum optical systems [17].

To overcome this limitation, Hu et al. proposed a generalized asymptotic phase compensation strategy [18,19], which introduces tunable phase shifts along the wavelength dimension to asymptotically match the nonlinear structural dispersion. This approach expands the degree of freedom for chromatic aberration correction beyond linear approximations. However, it also significantly increases the design complexity, as it requires establishing high-dimensional mappings between desired phase profiles and structural parameters across multiple wavelengths. While traditional full-wave simulation methods (e.g., finite-difference time-domain, FDTD [10,20,21]) can provide accurate predictions, their iterative optimization process is computationally expensive and time-consuming.

Deep learning techniques have been widely integrated into the design workflow of metastructured devices in recent years to improve the simulation efficiency [22,23,24,25,26,27,28,29,30,31,32,33]. These techniques using deep neural networks (DNNs) can efficiently learn the complex mappings between structural parameters and electromagnetic (EM) responses, enabling millisecond-level forward prediction and accelerating the inverse reasoning process from target spectra to structural configurations. Previous studies have explored architectures such as fully connected networks (FCNs) [19,34,35,36], generative adversarial networks (GANs) [37], bidirectional neural networks (BNNs) [28], and neural tensor networks (NTNs) [23,30,38,39]. However, most of these methods either fail to capture long-range dependencies within spectral sequences or rely on downsampling strategies that overlook critical sharp features in the spectra, thereby compromising prediction accuracy. The transformer architecture has recently gained popularity in photonic tasks due to its powerful sequence modeling capabilities [24,40,41], which allow it to attend to global spectral features and capture long-range correlations more effectively. However, existing transformer-based approaches in this domain typically rely on deeply stacked self-attention layers to reach high accuracy, resulting in very large models and increased computational complexity. At the same time, most current deep learning frameworks do not support training on multiple heterogeneous structural types within a single model, which limits scalability and applicability in multi-shape design tasks.

In this work, we propose a transformer-enhanced deep learning framework for efficient and accurate inverse design of broadband achromatic metasurfaces. First, we designed a lightweight and extensible forward prediction network, Shape-integrated Dual-Spectrum-aware transformer (SiDSaT), for highly accurate prediction of the spectrum. SiDSaT introduces a dual-branch spectral-modeling mechanism that concurrently captures long-range dependencies and resolves sharp spectral features, thereby sustaining high prediction accuracy across wide bandwidths and overcoming the limitations of single-branch models. In addition, the shape-integration module enables a single unified model to be trained on and used for inference across multiple nanostructure geometries, eliminating the need for separate networks for each shape. Unlike prior transformer-based approaches that rely on heavy layer stacking, SiDSaT adopts a shallow, lightweight architecture that markedly reduces parameter count and computational overhead while preserving predictive accuracy. We train and evaluate SiDSaT on extensive FDTD-simulated datasets, confirming precise spectral predictions across broad wavelength ranges and multiple shape categories. We then integrate SiDSaT with a hybrid inverse design algorithm combining Particle Swarm Optimization and a Genetic Algorithm (PSO-GA), demonstrating its feasibility for both single-wavelength designs and broadband asymptotic phase-compensated metalens applications. We expect that SiDSaT can be readily extended to a wider range of practical metasurface design scenarios.

2. Theory and Methodology

To construct a unified inverse design framework for metalenses, we developed a high-precision and lightweight transformer-based forward prediction model, SiDSaT, and integrated it with an evolutionary optimization strategy, PSO-GA, as illustrated in Figure 1.

Figure 1a shows the focusing mechanism of achromatic metalenses realized by modulating the phase of incident wavefronts through a sub-wavelength nanostructure array. Figure 1b depicts the workflow of SiDSaT, which takes the structural parameters and shape type of the unit cell as inputs. After preprocessing, the Shape-integrated module extracts shape-specific features via a gated fusion mechanism. Next, the Dual-Spectrum-aware transformer performs spectral sequence modeling through two parallel branches: the Global Perception transformer (GPT) and the Local Perception transformer (LPT). The network outputs the predicted spectral response through a multi-head output module. Figure 1c illustrates the PSO-GA-based optimization framework, which embeds the pre-trained SiDSaT model to obtain the corresponding predicted phases. The predicted phases are compared with the ideal phases using a loss function that quantifies focusing performance. The result is fed back into the PSO and GA modules to jointly update the structural parameters and reference phase $r_{\lambda }$. Through this process, optimal parameters are obtained and validated by FDTD simulations, demonstrating effective focusing and dispersion control.

2.1. Forward Prediction

2.1.1. Architecture of SiDSaT

Figure 2 shows the architecture of the forward prediction network, SiDSaT. It utilizes a combination of inter-fragment cross-attention and intra-fragment self-attention mechanisms to extract a rich set of spectral features. As shown in Figure 2a, the network is composed of four main components: a preprocessing and fully connected layer (FCL) for dimensional expansion, a Shape-integrated module for handling diverse shape types, a Dual-Spectrum-aware transformer for spectral modeling, and an output head for final prediction.

The preprocessing module includes two linear layers, each tailored for a specific shape type. These modules map geometric parameters with differing semantics (e.g., radius and height for cylinders, or side length and height for squares) into a unified vector representation of size $D\times 1$, resolving inconsistencies across shape formats. After normalization and mapping, the unified features are fed into a shared three-layer residual block (FCL), expanding the embedding to a higher-dimensional representation of size $T\times 1$.

Simultaneously, as shown in Figure 2b, the Shape-integrated module encodes the shape type label using a three-layer linear embedding, producing a shape vector of the same dimension $T\times 1$. A gating mechanism then learns soft fusion weights to adaptively combine the shape prior with the geometric features. This flexible fusion strategy maintains the continuity of geometric input while integrating discrete shape information, thereby enhancing the adaptability of the model to multiple shape types. This design allows SiDSaT to be trained on and make predictions for heterogeneous structures using a single model, improving training efficiency and structural generalization compared to traditional approaches requiring separate networks per shape.

The fused features are then passed to the Dual-Spectrum-aware transformer module. The features are first normalized and reshaped into an $n\times \left({\displaystyle \frac{T}{n}}\right)$ segment format. Based on the idea that different heads focus on different spectral segments, we propose a two-stage global–local spectral attention mechanism, comprising a Global Perception transformer (GPT) and a Local Perception transformer (LPT). The GPT uses a cross-attention mechanism to capture long-range dependencies between spectral fragments, while the LPT applies self-attention within each segment to refine local detail modeling.

As illustrated in Figure 2c, the Global Perception transformer operates on the matrix formed by all spectral fragments. Through linear projections, it generates the query (Q), key (K), and value (V) matrices of identical shape $n\times \left({\displaystyle \frac{T}{n}}\right)$, which are subsequently reshaped into a multi-head attention format with dimensions $h_d\times \left({\displaystyle \frac{n}{h_d}}\right)\times \left({\displaystyle \frac{T}{n}}\right)$, where $h_d$ denotes the dimension per attention head and $m={\displaystyle \frac{n}{h_d}}$ is the total number of heads.

This setup employs a cross-fragment attention mechanism to capture dependencies among different spectral fragments, thereby enabling effective extraction of global spectral features. The attention weights $W^G(Q,K)$ and the multi-head attention output $M^G(Q,K,V)$ are computed as follows:

$$W^G\left(Q,K\right)=Softmax\left({\displaystyle \frac{Q\cdot K^{\mathrm{\top }}}{\sqrt{h_d}}}\right)$$

(1)

$$M^G\left(Q,K,V\right)=W^{G}V$$

(2)

$$Q={(Q_1,Q_2,...,Q_m)}^T,K={(K_1,K_2,...,K_m)}^T,V={(V_1,V_2,...,V_m)}^T$$

(3)

Here, the attention weights are calculated by performing scaled dot-product between the projected Q and K, followed by a softmax operation. Each query head $Q_i$ interacts exclusively with its corresponding key head $K_i$, and the outputs from all heads are concatenated to produce the final multi-head attention result $M^G(Q,K,V)$. This mechanism effectively captures the global spectral dependencies, facilitating coarse-level global spectrum modeling.

As shown in Figure 2d, the Local Perception transformer processes each spectral fragment independently. For each fragment, linear projections generate their own Q, K, and V matrices of shape $h_d\times \left({\displaystyle \frac{n}{h_d}}\right)$, where $h_d$ and $m={\displaystyle \frac{n}{h_d}}$ retain their definitions. This mechanism focuses on capturing intra-fragment (local) spectral relationships.

As described in Equations (4) and (5), the attention weights $W^L\left(Q,K\right)$ are computed similarly to the global transformer, but with two key differences: all Q, K, and V come from the same spectral fragment, and a sigmoid activation is applied instead of softmax:

$$W^L=\left(W_1^L,W_2^L,...W_n^L\right),W_i^L\left(Q,K\right)=Sigmoid\left({\displaystyle \frac{Q_i\cdot {K_i}^{\mathrm{\top }}}{\sqrt{h_d}}}\right)$$

(4)

$$M^L\left(Q,K,V\right)=W^L\odot V,M_i^L\left(Q_i,K_i,V_i\right)=W_i^L{V}_i$$

(5)

The attention output $M^L(Q,K,V)$ is obtained via element-wise (Hadamard) multiplication between the weight matrix $W^L$ and the value matrix $V$. Notably, each $W_i^L$ is computed using only the corresponding fragment’s own $Q_i$ and $K_i$, ensuring strict local context modeling. By aggregating outputs from all attention heads and fragments, the final output is constructed.

This dual-branch module separates global and local spectral modeling. The Global Perception transformer (GPT) performs cross-fragment multi-head attention over the full wavelength grid; because scaled dot-product attention links distant spectral segments, its output is most responsive where the spectrum–wavelength curve exhibits pronounced curvature (i.e., where long-range correlations across segments are strongest). By contrast, the Local Perception transformer (LPT) applies self-attention within each segment and thus acts as a local refiner, emphasizing steep slopes and narrow features that would otherwise be smeared by downsampling. This division of labor provides a physics-aligned representation: GPT targets wavelength regions pivotal for dispersion control, whereas LPT preserves sharp features relevant to fabrication-sensitive spectrum behavior.

After attention, the outputs are reshaped and projected back to the original segment format. A first residual connection is applied, followed by flattening to size $T\times 1$, then passed through a Norm and Feed Forward (NFF) layer and a second residual connection. The resulting attention-enhanced representation is sent to the output head for final spectral prediction.

Since metasurface lens design focuses primarily on phase information, the loss function for SiDSaT uses trigonometric error to mitigate the impact of phase discontinuities (see Equation (6)). The performance of the model is quantitatively evaluated using the mean absolute error (MAE) of phase on the validation set (see Equation (7)).

$$L_{\mathrm{S}\mathrm{i}\mathrm{D}\mathrm{S}\mathrm{a}\mathrm{T}}={\displaystyle \frac{1}{N}}{\sum }_{i=\mathrm{1,2},\dots ,N}[{(\mathrm{c}\mathrm{o}\mathrm{s}({Phase}_{\mathrm{p}\mathrm{r}\mathrm{e}})-\mathrm{c}\mathrm{o}\mathrm{s}({Phase}_{\mathrm{s}\mathrm{i}\mathrm{m}}))}^2+{(\mathrm{s}\mathrm{i}\mathrm{n}({Phase}_{\mathrm{p}\mathrm{r}\mathrm{e}})-\mathrm{s}\mathrm{i}\mathrm{n}({Phase}_{\mathrm{s}\mathrm{i}\mathrm{m}}))}^2]$$

(6)

$$E_{\mathrm{S}\mathrm{i}\mathrm{D}\mathrm{S}\mathrm{a}\mathrm{T}}={\displaystyle \frac{1}{N}}{\sum }_{i=\mathrm{1,2},\dots ,N}\left|{Phase}_{\mathrm{p}\mathrm{r}\mathrm{e}}-{Phase}_{\mathrm{s}\mathrm{i}\mathrm{m}}\right|$$

(7)

Here, ${Phase}_{\mathrm{p}\mathrm{r}\mathrm{e}}$ denotes the phase predicted by SiDSaT, while ${Phase}_{\mathrm{s}\mathrm{i}\mathrm{m}}$ represents the reference phase obtained from FDTD simulations.

2.1.2. Dataset Generation

To train and validate the SiDSaT model, we constructed a dataset comprising geometric parameters of metasurface unit cells and their corresponding spectral responses. The target of our metalens design is to achieve efficient focusing and broadband achromatic performance in the near-infrared regime. Titanium dioxide (TiO₂) is selected as the structural material due to its high refractive index and low absorption loss, while silicon dioxide (SiO₂) serves as the substrate owing to its high transparency and material compatibility, making it a common choice for waveguide-type structures.

The unit cell period is set to 500 nm to satisfy the zero-order diffraction condition within the 800–1600 nm wavelength range, suppress higher-order diffraction, and strike a balance between fabrication resolution and optical performance. The geometric parameters include primary scalar values such as radius (r) and width (d), along with a shape type indicator (e.g., cylindrical or cuboidal).

The spectral responses of each structure are computed using FDTD simulations across 81 discrete wavelength points in the 800–1600 nm range. To eliminate the ambiguity caused by phase wrapping, the corresponding phase responses are encoded using sine and cosine representations.

To ensure balanced sampling during training, we employ stratified splitting based on compound labels formed by grouping geometric parameters and shape types. This approach guarantees fair representation of all unit types in both the training and validation sets.

2.2. Inverse Design

The purpose of inverse design is to iteratively search for optimal nanostructure parameters that yield the desired spectral phase response across a wide wavelength range. To achieve this goal, we integrate the proposed SiDSaT forward prediction network with a hybrid optimization framework that combines Particle Swarm Optimization (PSO) and Genetic Algorithm (GA).

However, in the conducted constraint analysis, the inverse design of metasurfaces is not a purely mathematical optimization problem but is inherently constrained by the physical realizability of the phase profile and its dispersion. In broadband achromatic focusing, the target phase distribution should not only satisfy wavelength-dependent focusing conditions but also remain consistent with the nonlinear phase dispersion of nanostructured meta-units. By rigorously analyzing the compatibility between ideal phase requirements and the dispersion capacity of the meta-unit library, we delineate the feasible design space, which provides the physical foundation for the optimization framework described in the following subsection.

2.2.1. The Principle of Achromatic Metalens

When treating each metasurface unit cell as a vertical waveguide, its phase response can be expressed as follows [42]:

$${\phi }_{meta}={\displaystyle \frac{2\pi }{\lambda }}{n}_{eff}$$

(8)

In previous studies, the structural dispersion of unit cells was often assumed to be linear, which neglected the wavelength dependence of the effective refractive index [16]. As shown in Figure 3a, the phase profile of linear dispersion is typically constructed as

$${\phi }_{linear}\left(r,\lambda \right)=-{\displaystyle \frac{2\pi }{\lambda }}\left(\sqrt{r^2+f^2}-\sqrt{r_0^2+f^2}\right)$$

(9)

where $f$ is the focal length, and $r_0$ is the radial position at which the wavefront intersects the lens surface.

In achromatic design, when $r<r_0$, this approach allows approximate linear fitting of the structural dispersion, enabling compensation for chromatic aberration within a limited bandwidth. However, under one-dimensional effective medium theory, the effective refractive index generally varies with wavelength. As such, the linear compensation method becomes insufficient for broadband control of focusing dispersion. Moreover, it is nearly impossible to design a broadband focal dispersion control lens (i.e., with wavelength-dependent focal length $f_{\lambda }$) using this approach.

To better align with the nonlinear structural dispersion of metasurface units and enable arbitrary dispersion engineering across ultra-broadbands, Hu et al. introduced the asymptotic phase compensation method [18]:

$${\phi }_{asympt}\left(r,\lambda \right)=-{\displaystyle \frac{2\pi }{\lambda }}\left(\sqrt{r^2+f^2}-\sqrt{r_{\lambda }^2+f^2}\right)$$

(10)

As illustrated in Figure 3b, this method introduces a wavelength-dependent shift variable $r_{\lambda }$. Thus, the local phase relations can be freely constructed to asymptotically match nonlinear dispersion characteristics, enabling arbitrary dispersion control across broadband spectra.

Notably, Equation (10) can be decomposed as

$${\phi }_{asympt}\left(r,\lambda \right)=-{\displaystyle \frac{2\pi }{\lambda }}\left(\sqrt{r^2+f^2}-f\right)+{\phi }_{shift}\left(\lambda \right)$$

(11)

where

$${\phi }_{shift}\left(\lambda \right)={\displaystyle \frac{2\pi }{\lambda }}\left(\sqrt{r_{\lambda }^2+f^2}-f\right)$$

(12)

Here, ${\phi }_{shift}\left(\lambda \right)$ represents a phase shift term dependent only on the wavelength.

2.2.2. Constraint Analysis for Metalens Design

In practice, the phase response of each unit cell is constrained within a finite range, which directly impacts whether a target phase profile can be realized across the entire wavelength band. Thus, to ensure broadband feasibility of the desired phase distribution, the relationship between the achievable phase dispersion of the structure library and the required bandwidth must be rigorously evaluated.

For Equation (10), the required phase difference between the lens center and edge at a fixed wavelength (independent of ${\phi }_{shift}$) is

$$∆{\phi }_{asympt}\left(\lambda \right)={\phi }_{asympt}\left(0,\lambda \right)-{\phi }_{asympt}\left(R_{max},\lambda \right)={\displaystyle \frac{2\pi }{\lambda }}\left(\sqrt{{R_{max}}^2+f^2}-f\right)$$

(13)

Therefore, for two wavelengths ${\lambda }_1>{\lambda }_2$, the phase difference gap is

$$∆{\phi }_0\left({\lambda }_1,{\lambda }_2\right)=∆{\phi }_{asympt}\left({\lambda }_2\right)-∆{\phi }_{asympt}\left({\lambda }_1\right)=2\pi \left({\displaystyle \frac{1}{{\lambda }_2}}-{\displaystyle \frac{1}{{\lambda }_1}}\right)\left(\sqrt{{R_{max}}^2+f^2}-f\right)$$

(14)

To ensure the unit cell library covers the required phase dispersion range between two wavelengths, the following condition must hold (see Figure 3c):

$$\left\{\begin{array}{c}{\phi }_{asympt}\left(R_{max},{\lambda }_2\right)-{\phi }_{asympt}\left(R_{max},{\lambda }_1\right)\le ∆{\phi }_{{unit}_{min}}\left({\lambda }_1,{\lambda }_2\right)\\ ∆{\phi }_{{unit}_{max}}\left({\lambda }_1,{\lambda }_2\right)\le {\phi }_{asympt}\left(0,{\lambda }_2\right)-{\phi }_{asympt}\left(0,{\lambda }_1\right)\end{array}\right.$$

(15)

Here, $∆{\phi }_{{unit}_{min}}$ and $∆{\phi }_{{unit}_{max}}$ denote the minimum and maximum achievable phase variation in the meta-units between two wavelengths.

Summing both inequalities, we obtain

$$∆{\phi }_{uni{t}_{max}}\left({\lambda }_1,{\lambda }_2\right)-∆{\phi }_{uni{t}_{min}}\left({\lambda }_1,{\lambda }_2\right)\ge ∆{\phi }_0\left({\lambda }_1,{\lambda }_2\right)$$

(16)

Rewriting the right-hand side as a bandwidth-dependent parameter $\chi$, this yields

$$\sqrt{{R_{max}}^2+f^2}-f\le {\displaystyle \frac{∆{\phi }_{uni{t}_{max}}\left({\lambda }_1,{\lambda }_2\right)-∆{\phi }_{uni{t}_{min}}\left({\lambda }_1,{\lambda }_2\right)}{2\pi \left(\frac{1}{{\lambda }_2}-\frac{1}{{\lambda }_1}\right)}}=\chi$$

(17)

Thus, to ensure sufficient phase coverage across all wavelength pairs in the design bandwidth, the constraint becomes

$$\sqrt{{R_{max}}^2+f^2}-f\le {\chi }_{min}$$

(18)

This leads to the following constraints on focal length and lens radius:

$$\left\{\begin{array}{c}0<f,0<R_{max}\le {\chi }_{min}\\ {\displaystyle \frac{{R_{max}}^2-{{\chi }_{min}}^2}{2{\chi }_{min}}}\le f,R_{max}>{\chi }_{min}\end{array}\right.$$

(19)

And similarly for numerical aperture (NA):

$$\left\{\begin{array}{c}0<NA<1,0<R_{max}\le {\chi }_{min}\\ 0<NA<{\displaystyle \frac{2{\chi }_{min}}{{R_{max}}^2+{{\chi }_{min}}^2}},R_{max}>{\chi }_{min}\end{array}\right.$$

(20)

These derivations reveal that large-scale, high-NA, or broadband metasurfaces require a unit cell library with a sufficiently wide dispersion range $\chi$. If the library lacks adequate coverage (e.g., with a single-type structure), a trade-off between lens size, NA, and bandwidth becomes inevitable.

Since the variable $r_{\lambda }$ is wavelength-dependent and unconstrained, its range is difficult to define. However, the associated phase shift ${\phi }_{shift}\left(\lambda \right)$ remains bounded between $-\pi$ and $\pi$. Therefore, we introduce a wavelength-specific compensation term:

$${{\phi }_{shift}\left(\lambda \right)=C}_{\lambda },$$

(21)

and directly optimize it in the inverse design process.

2.2.3. PSO-GA-Based Inverse Design Strategy

The primary objective of metalens inverse design is to minimize the discrepancy between the ideal phase profile and the actual phase wavefront produced by the structural configuration at each wavelength. In this study, we adopt a phase-matching loss based on sine and cosine representations, consistent with outputs of the forward model, to avoid ambiguities caused by phase wrapping. The triangle-based error function ${\epsilon }_{tri}$ is defined as

$${\epsilon }_{tri}=\sqrt{{\sum }_{i=1}^{n_{\lambda }}{\sum }_{r=1}^{r_{max}}{\left(sin\right(\mathit{\Delta }{\phi }_{ideal})-sin(\mathit{\Delta }{\phi }_{unit}\left)\right)}^2+{\left(cos\right(\mathit{\Delta }{\phi }_{ideal})-cos(\mathit{\Delta }{\phi }_{unit}\left)\right)}^2)}$$

(22)

To solve the inverse design problem, we employ a hybrid optimization strategy that integrates Particle Swarm Optimization (PSO) and Genetic Algorithm (GA), referred to as PSO-GA, as shown in Figure 4. This approach combines the rapid local search ability of PSO with the global exploration capability of GA. The inclusion of crossover and mutation operations helps prevent PSO from falling into local minima, while the swarm’s dynamic update mechanism accelerates the convergence of GA.

The PSO component performs local search in the parameter space through iterative updates of particle velocity and position, defined as

$$\left\{\begin{array}{c}{v}_i^{t+1}=w\cdot v_i^t+c_1{r}_1(pbes{t}_i^t-x_i^t)+c_2{r}_2(gbes{t}^t-x_i^t))\\ \left(x_i^{t+1}=x_i^t+v_i^{t+1}\right)\end{array}\right.$$

(23)

where $r_1,r_2\sim U\left(\mathrm{0,1}\right)$ are random coefficients, $pbes{t}_i$ is the best-known position of particle $i$, and $gbest$ is the global best position found so far. The velocity is constrained by a maximum threshold $v_{max}=0.1\left(d_{max}-d_{min}\right)$ to prevent divergence. The inertia weight $w$ regulates memory of past velocities, while cognitive ($c_1$) and social ($c_2$) coefficients control convergence toward individual and global optima, respectively.

The GA component enhances population diversity via tournament selection, uniform crossover, and Gaussian mutation. GA is periodically invoked to intervene in the PSO process, facilitating population migration and promoting bidirectional information flow between modules.

During the optimization, the pre-trained SiDSaT forward model is utilized as a physics-aware surrogate to predict the phase response of candidate unit structures. Since the objective function is defined by the trigonometric phase error ${\mathsf{\epsilon }}_{\mathrm{t}\mathrm{r}\mathrm{i}}$, the model’s outputs (i.e., predicted sine and cosine of phase) can be directly used in the evaluation without requiring explicit phase reconstruction, thereby avoiding phase discontinuity artifacts.

3. Results and Discussion

3.1. Evaluation of SiDSaT

3.1.1. Predicted Results of Our Dataset

To evaluate the prediction accuracy of the fully trained SiDSaT model, we tested it on all samples in the dataset, as well as on additional samples randomly selected from the parameter space that were not included in the training set. The predicted and simulated spectral responses are compared, and their discrepancies are quantified using the mean absolute error (MAE).

Figure 5a,b show examples of prediction results for samples within the dataset and the predictions for unseen samples outside the dataset, respectively. Figure 5c presents the MAE of predicted phase values at each sampled wavelength from 800 nm to 1600 nm across the entire dataset. The overall MAEs for cylindrical and rectangular metasurface units are 0.0473 and 0.0446, respectively.

The proposed SiDSaT model demonstrates excellent predictive capability across varying geometric parameters (e.g., radius, width) and structural shapes (circular and rectangular). Even for unseen samples, the predicted phase spectra remain highly consistent with FDTD simulation results across the entire wavelength range, showing near-perfect overlap in the long-wavelength region. For structures with strong dispersive behavior, including those exhibiting phase discontinuities, SiDSaT still accurately captures the spectral trends, with only minor deviations near the spectral boundaries.

The MAE statistics for each wavelength point in Figure 5c reveal that prediction errors are generally low throughout the spectrum. Higher errors are mainly concentrated in the short-wavelength region (800–950 nm), which may be attributed to stronger material dispersion and rapid changes in phase response. As the wavelength increases, the MAE quickly decreases and stabilizes, indicating better generalization performance in the mid-to-long wavelength range.

In summary, SiDSaT achieves low phase prediction error across the entire wavelength range for metasurface units with various shapes and geometries. Additionally, validation on more complex and heterogeneous datasets is provided in Appendix A.1 (Figure A1) and Appendix A.2 (Figure A2). These results verify the model’s capability in learning complex high-dimensional structural relationships and accurately modeling spectral responses, providing a robust and reliable foundation for subsequent inverse design.

3.1.2. Ablation Study

To systematically evaluate the contribution of each key module in the proposed SiDSaT model to the prediction performance, we conducted ablation studies by comparing the full SiDSaT architecture against several variant models with specific components removed, as illustrated in Figure 6. Specifically, we performed ablation by individually removing the following modules: the Shape-integrated module (denoted as without Si), the Global Perception transformer (without GPT), the Local Perception transformer (without LPT), and the entire dual-transformer structure (without transformer). In addition, we included a baseline FCN model, which lacks both shape-awareness and attention mechanisms, for comparative analysis.

From Figure 6, we observe that the complete SiDSaT model (red curve) consistently achieves the lowest validation loss throughout the training process, demonstrating faster convergence and lower variance. This highlights its superior capability in learning spectral responses. In contrast, the baseline FCN (sky blue curve) exhibits significantly higher error and oscillation across the entire training trajectory, indicating that traditional fully connected architectures struggle to effectively model spectral sequences.

Notably, the inset in the upper-right corner of Figure 6 shows a magnified view of the final training stage (Epochs 3000–4000). The full model not only attains the lowest final validation loss but also maintains the smallest fluctuation range, demonstrating excellent long-term stability and robustness. Among the variant models, removing either transformer sub-module (GPT or LPT) leads to noticeable performance degradation, which confirms that global and local spectral-modeling strategies are complementary: GPT captures cross-fragment dependencies, while LPT characterizes intra-fragment variations. Together, they enhance both modeling depth and precision.

Eliminating the Shape-integrated module also causes a significant increase in prediction error, underscoring the importance of structural shape information in accurate spectral response modeling. Moreover, the model without any transformer components exhibits saturated performance in later training stages, suggesting the essential role of attention mechanisms in modeling high-dimensional structure-to-spectrum mappings.

In summary, the ablation study clearly demonstrates that each module in SiDSaT plays a critical role in performance enhancement. The full architecture surpasses all simplified variants in terms of accuracy, convergence speed, and training stability, validating the effectiveness of its overall design.

3.2. Evaluation of PSO-GA

3.2.1. Application in Single-Wavelength Metalenses

To validate the effectiveness of our SiDSaT-based inverse design framework, we performed single-wavelength metalens design experiments at key wavelengths: 850 nm, 1064 nm, and 1310 nm, corresponding to standard bands used in biological imaging, optical communication, and LiDAR systems. These wavelengths were selected due to their widespread practical applications and distinct dispersion characteristics in the near-infrared range.

The proposed metalens is constructed from a combination of multiple structural geometries. The lens diameters and focal lengths were carefully chosen to balance diffraction performance and fabrication feasibility. Specifically, lenses with 60 μm or 40 μm diameters and 50–200 μm focal lengths were used to satisfy zero-order diffraction conditions while allowing sub-wavelength resolution under fabrication constraints.

The focusing performance of the four designs is illustrated in Figure 7a–d, showing both the electric field intensity distribution and the corresponding focal plane profiles.

Table 1. Geometric specifications and focusing performance of designed single-wavelength metalenses for inverse design validation.

	Metalens 1	Metalens 2	Metalens 3	Metalens 4
Wavelength	850 nm	850 nm	1064 nm	1310 nm
Diameter	60 μm	60 μm	40 μm	40 μm
Target Focal Length	200 μm	100 μm	50 μm	50 μm
Actual Focal Length	199.2 μm	101.0 μm	48.7 μm	50.0 μm
NA	0.1483	0.2873	0.3714	0.3714
FWHM	3.058 μm	1.834 μm	1.644 μm	2.060 μm
Focus Ratio	1.04	1.21	1.12	1.14
DOF	35.760 μm	10.835 μm	7.504 μm	8.109 μm
${\eta }_{trans}$	0.6041	0.8798	0.8116	0.8428
${\eta }_{focus|trans}$	0.7644	0.7409	0.6921	0.6289
Strehl Ratio	0.9108	0.7194	0.7498	0.6314

summarizes the key metrics, including designed and actual focal length, numerical aperture (NA), full width at half maximum (FWHM), depth of focus (DOF), transmittance ${\eta }_{trans}$, focusing efficiency ${\eta }_{focus|trans}$, and focus ratio.

Notably, the focus ratio is defined as

$$\mathrm{F}\mathrm{o}\mathrm{c}\mathrm{u}\mathrm{s} \mathrm{R}\mathrm{a}\mathrm{t}\mathrm{i}\mathrm{o}={\displaystyle \frac{{\mathrm{F}\mathrm{W}\mathrm{H}\mathrm{M}}_{\mathrm{a}\mathrm{c}\mathrm{t}\mathrm{u}\mathrm{a}\mathrm{l}}}{{\mathrm{F}\mathrm{W}\mathrm{H}\mathrm{M}}_{\mathrm{d}\mathrm{i}\mathrm{f}\mathrm{f}\mathrm{r}\mathrm{a}\mathrm{c}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}-\mathrm{l}\mathrm{i}\mathrm{m}\mathrm{i}\mathrm{t}\mathrm{e}\mathrm{d}}}}={\displaystyle \frac{{\mathrm{F}\mathrm{W}\mathrm{H}\mathrm{M}}_{\mathrm{a}\mathrm{c}\mathrm{t}\mathrm{u}\mathrm{a}\mathrm{l}}}{0.514\cdot \frac{\lambda }{\mathrm{N}\mathrm{A}}}}$$

(24)

This metric quantifies the closeness of the focal spot to the diffraction limit, where a value close to 1 indicates near-diffraction-limited focusing. We also use two efficiency metrics for subsequent analyses: transmittance ${\eta }_{trans}$, the ratio of power transmitted through the metalens aperture to the incident power, and focusing efficiency ${\eta }_{focus|trans}$, the ratio of power within FWHM at $f_{actual}$ to transmitted power.

As shown in Table 1, the achieved focal lengths closely match the targets (e.g., 199.2 µm vs. 200 µm; <0.5% error across cases). The NAs span 0.15–0.37, with measured focal spot FWHM values of 1.64–3.06 µm. All four metalenses achieve focus ratios in the range of 1.04–1.21, indicating near-diffraction-limited focusing and high-fidelity phase reconstruction. The efficiency metrics ${\eta }_{trans}$ and ${\eta }_{focus|trans}$ are uniformly high across cases, while the wavefront quality is further supported by Strehl ratios of ~0.63–0.91. The depth of focus (DOF) ranges from ~7.5 to 35.8 µm, providing axial tolerance to minor defocus and fabrication variations.

Collectively, these metrics indicate that the SiDSaT-PSO-GA framework yields high-efficiency, near-diffraction-limited single-band metalenses with tightly focused beams. The close agreement between actual and target focal lengths, along with narrow focal spot sizes, further corroborates the accuracy of the SiDSaT forward model in forward prediction and the reliability of the PSO-GA optimizer in converging to physically valid geometries that reconstruct the intended phase profile.

3.2.2. Application in Broadband Metalenses

To verify the applicability of the proposed inverse design framework in broadband scenarios, we further designed an achromatic metalens capable of operating across the 1000–1400 nm spectral range. The proposed metalens is also constructed from a combination of multiple structural geometries. The design objective was to maintain an approximately constant focal length across the working bandwidth while minimizing chromatic focal shift and preserving stable focusing performance. The metalens diameter was set to 60 μm with a target focal length of 400 μm, meeting the structural requirements of typical near-infrared optical communication and imaging systems.

The designed metalens was simulated using FDTD methods. Figure 8a shows the x-z plane intensity distributions at multiple discrete wavelengths. It can be observed that the focal spot consistently remains near the target focal length across the full spectral range, exhibiting symmetric and well-confined focal profiles. This indicates that the structure achieves effective phase modulation across multiple wavelengths.

To quantitatively analyze focal length dispersion, Figure 8b plots the deviation $\Delta f$ between the actual focal length ($f_{actual}$) and the target focal length ($f_{design} = 400 \mathsf{\mu }\mathrm{m}$) for each wavelength. In the 1000–1200 nm range, the focal length remains relatively stable with minimal shift. As the wavelength increases beyond 1250 nm, a moderate focal regression is observed, yet the maximum deviation remains below 100 μm. Notably, the depth of focus (gray error bars) increases with wavelength, providing axial tolerance and mitigating the performance impact of focus shifts. This demonstrates the design’s strong capability in maintaining focal consistency, achieving quasi-achromatic behavior across the target band.

Figure 8c further analyzes the intensity at the focal plane and the variation in the full width at half maximum (FWHM) with respect to wavelength. Both at $f_{design}$ and $f_{actual}$, the focal intensity remains high, indicating good focusing efficiency and depth-of-focus tolerance. However, the FWHM curve (orange dashed line) deviates from the theoretical diffraction limit (black solid line), suggesting that although the resolution does not fully reach the diffraction limit, it remains within an acceptable range with consistent broadband-focusing trends.

System-level metrics in Figure 8d reinforce these observations: the transmittance ${\eta }_{trans}$ is consistently high with modest variation, the focusing efficiency referenced to transmitted power ${\eta }_{focus|trans}$ remains elevated, and the Strehl ratio is near-flat and high across the band. Together, the small $∆f$, near-limit FWHM, and stable efficiency/quality metrics demonstrate quasi-achromatic behavior of metalens over 1000–1400 nm. Additional quantitative comparison with previously reported achromatic metalenses and concise manufacturability assessment are provided in Appendix B and

Table A1. Comparison of Broadband Achromatic Metalenses (BAMLs).

Reference	Operating Band (nm)	Focal Length	Max Focal Shift	NA	Materials	Achromatization Strategy	Operation Mode and Polarization
Our work	1000–1400	400 μm	≈90 μm (22.5%)	0.080	TiO2 nanopillars on SiO2 substrate	Asymptotic phase compensation; Primitive meta-atom; SiDSaT-embedded PSO-GA strategy	Transmission; Polarization insensitive
	1000–1300		≈40 μm (10%)
Wang et al. [45] (2017)	1200–1680	100 μm	≈20 μm (20%)	0.268	Au multi-nanorods/SiO2 spacer/Au back reflector (on Si substrate)	Integrated Resonant Unit	Reflection; Polarization selective
Hsiao et al. [46] (2018)	420–650	150 μm	≈30 μm (20%)	0.124	Al multi-nanorods/SiO2 spacer/Al back reflector (on Si substrate)	Integrated Resonant Unit	Reflection; Polarization selective
Shrestha et al. [16] (2018)	1200–1650	200 μm	≈30 μm (15%)	0.24	a-Si nanopillars on fused silica	Complex meta-atom	Transmission; Polarization insensitive
Sun et al. [47] (2022)	450–1400	214 μm	≈25 μm (12%)	0.107	Si3N4 nanopillars on fused silica	Aperture Partition and Phase Coordination Method (APPCM, dual-lens-zone design)	Transmission; Polarization selective
Hu et al. [18] (2023)	400–1000	150 μm	≈32 μm (21%)	0.164	TiO2 nanopillars on SiO2 substrate	Asymptotic phase compensation; Complex meta-atom	Transmission; Polarization insensitive
Li & Lv [48] (2024)	1000–1250	6.5 μm	≈0 μm	0.5	Si nanopillars on SiO2 substrate	PSO-GA hybrid optimization	Transmission; Polarization selective
Song et al. [49] (2024)	430–750	50 μm	≈6 μm (12%)	0.255	LNOI thin film on LNOI substrate	Multi-Zone Taylor Expansion Method (MZTEM)	Transmission; Polarization insensitive
Hou et al. [42] (2025)	450–650	60 μm	≈3 μm (5%)	0.164	TiO2 anisotropic nanopillars on SiO2 substrate	Anisotropic nanofin phase control with PSO	Transmission; Polarization selective

.

In summary, the broadband simulation results validate the feasibility of the proposed design framework for multi-wavelength focusing tasks, highlighting its potential in broadband achromatic metalens applications.

4. Conclusions

In conclusion, we have proposed a transformer-based deep learning forward prediction network, SiDSaT, and demonstrated its feasibility for efficient and effective inverse design of achromatic metalenses. Incorporating a Shape-integrated module and dual attention mechanisms (GPT and LPT), SiDSaT can accommodate various geometric types and wavelength regimes in a single training session. We utilized SiDSaT to predict the phase for cylindrical and rectangular meta-atoms over wavelengths from 800 nm to 1600 nm, achieving mean absolute errors of 0.0473 and 0.0446, respectively. This shows that SiDSaT achieves highly accurate phase prediction across the full spectral range for various structural types. Furthermore, we conducted ablation studies to validate the proposed model’s architecture.

To verify the adaptability of SiDSaT, we integrated it into an inverse design framework (PSO-GA) and accomplished metalens designs for both single-wavelength and broadband applications. Particularly, based on the asymptotic phase compensation method, we performed a constraint analysis for the achromatic metalens design, which lowered the algorithm’s computational cost. Using the SiDSaT-embedded PSO-GA strategy, we obtained four single-wavelength metalens designs at 850 nm, 1064 nm, and 1310 nm. The resulting focal spots have FWHM values approaching the diffraction limit, confirming the framework’s high accuracy and reliability in narrowband design tasks. Moreover, we applied the framework to the design of an achromatic metalens for the 1000–1400 nm wavelength range. Simulation results show an approximately constant focal length across this operating band, demonstrating the framework’s effectiveness even when nonlinear dispersion compensation is required.

Above all, the proposed SiDSaT-embedded inverse design framework not only exhibits highly accurate modeling and efficient structure-search capabilities but also demonstrates strong adaptability and generalization in practical design tasks. In addition, by replacing traditional FDTD simulations with SiDSaT for phase prediction, the overall computational efficiency of the design pipeline improves by approximately two orders of magnitude, significantly accelerating the design process without compromising accuracy. This approach provides a powerful solution for the intelligent design of complex metasurface optical devices across multiple wavelengths and target specifications.