Robust global optimization of atomic structures via a learning loss-informed on-the-fly firefly algorithm

DFT

All DFT calculations were performed using the Vienna Ab initio Simulation Package (VASP)^[28,29] with the projector augmented wave (PAW) method^[30] and the GGA proposed by Perdew, Burke, and Ernzerhof (PBE)^[31] as the exchange-correlation (xc) functional. A planewave kinetic energy cutoff of 500 eV was used. The convergence criteria for the electronic step were set to 10^-5 eV, while structural relaxations were considered converged once the norm of forces on each atom was reduced to less than 10^-2 eV/Å.

For the O/Cu(111) surface oxide calculations d, the DFT data was adopted from the database developed in our previous work^[16]. All Cu(111) slab models in the database consist of four atomic layers, where the bottom two layers were fixed to represent a bulk-like region, and a minimum vacuum spacing of 14 Å was used along the surface-normal direction. A Γ-centered k-point grid with a spacing of 0.15 Å^-1 was used consistently throughout all the structures within the database. For the gold cluster calculations, a cubic unit cell with dimensions of 20 Å along each axis was used, and all atoms are set apart from the unit cell boundary more than 5 Å to avoid the unwanted interactions between repeated unit cells. The 5d¹⁰ 6s¹ orbitals of gold atoms are considered as valence states, and all calculations were performed at the Γ-point. Dispersion corrections (e.g., vdW or MBD) were not applied, consistent with the aim of this work to assess robustness against target-potential limitations. Additional computational details (i.e., input settings) are provided in the Supplementary Materials.

Random structure generation and structure conditioning

For each structure, atomic coordinates (x, y, z) were randomly sampled from a uniform probability distribution within predefined regions or unit cells according to the required number of atoms. Because both random initialization and subsequent firefly algorithm (FA) update steps can generate candidate structures with unphysical close contacts, we apply a conditioning step to remove severe overlaps before energy/force evaluation. Such unrealistically short interatomic distances can destabilize calculations and increase computational costs, and this issue becomes more pronounced when using machine learning-based energy/force predictions due to the lack of training data in such regions.

To mitigate this, we employed a conditioning procedure applied immediately after random structure generation and after each FA update. In this procedure, atoms closer than a predefined threshold distance were repelled from each other before energy/force evaluation and subsequent local relaxation. We define this threshold using the sum of covalent radii^[32] as a physically motivated length scale, and we set it to 80% of the summed radii as a practical compromise. This choice removes clearly unphysical close contacts while avoiding an overly conservative constraint that would unnecessarily inflate initial separations and bias the search. We emphasize that this conditioning is not enforced during the subsequent local relaxation under the target potential and therefore does not impose a hard lower bound on final bond lengths. We utilized Hooke’s law to define the repulsive force and corresponding energy whenever two atoms approach within the threshold^[33] as follows,

Here, interatomic distances, r_ij, were computed using the neighbor analysis function implemented in the atomic simulation environment (ASE^[34]), and structural relaxation was conducted to minimize the defined repulsive force. For this study, conditioning thresholds were set as follows: 2.18 Å for gold cluster and 1.58, 2.11, and 1.06 Å for Cu–O, Cu–Cu, and O–O, respectively, for O/Cu(111) surface oxides. As the defined conditioning forces and energies were not simultaneously combined with actual potential energies or atomic forces, a consistent spring constant of k was applied as 1 eV/ Ų throughout this study.

Visualization of PES

To analyze the results of global optimization, the PES was post-processed and visualized. All structures generated during the global optimization were collected and used to train a variational autoencoder (VAE)-based architecture^[35,36], reducing dimensionality to a 2D latent space for visualization. In defining global descriptors from symmetry-invariant local descriptors assigned to each atom, an attention pooling layer^[37] was implemented. This layer computes attention scores from each atom’s local symmetry invariant descriptor and applies a softmax function to aggregate these into global descriptors. The training was conducted in two steps: initially, the attention pooling layer and VAE architecture were trained simultaneously to minimize the reconstruction error; subsequently, only the VAE parameters were optimized based on the evidence lower bound (ELBO)^[35]. This approach ensured that attention scores highlighted structurally significant distinctions, and resulted in a continuous representation within the VAE latent space as well.

Implementation of the FA for atomic structure optimization

The FA operates according to three key principles: (i) all fireflies are treated as unisex, meaning any individual can be attracted to any other regardless of “gender”; (ii) mutual attractiveness is proportional to perceived brightness, which decays exponentially with distance; and (iii) each firefly moves toward the one it perceives as most attractive, or performs a random walk if no firefly is identified as attractive. When applied to atomic structure optimization, each firefly represents an atomic configuration, and the evolution of configurations toward the next generation is governed by^[38]:

where x_i^t denotes the Cartesian coordinates of agent i at generation t (a 3N representation for N atoms), x_j^t is the configuration of a more attractive agent j, and ε is a random displacement. The parameter α sets the step size of the random displacement term for agent i, and in atomic structure optimization it can be chosen analogously to the perturbation magnitude used in basin hopping. In this work, we use α =0.5 Å. The parameter β₀ is the attractiveness at zero distance and sets the maximum scale of the attractiveness term. It is commonly treated as a fixed constant close to unity^[38], and we use β₀ =0.9. The parameter γ controls the distance-dependent decay of attractiveness, β(r_ij) = β₀exp(-γr_ij²), and thus sets the effective interaction range between agents. A larger γ strengthens the distance decay and limits interactions to nearby agents, whereas a smaller γ maintains longer-ranged interactions. Accordingly, we keep α and β₀ fixed across systems and adjust γ when the distance scale or energy-landscape characteristics change. Following prior FA studies that report practical population size in the range of 15-50^[38], we use 30 agents to balance exploration diversity against the cost of local relaxations and energy/force evaluations.

In practice, the displacement term (x_j^t - x_i^t) in Equation (2) cannot be taken as a direct subtraction of the raw 3N coordinate vectors, because two configurations that represent the same motif may differ by atom-index permutations and symmetry-related transformations. To obtain a meaningful displacement, we define (x_j^t - x_i^t) only after establishing a consistent atom-to-atom correspondence between the two structures using our symmetry-invariant local descriptor-based atom index matching.

To effectively implement the FA with a proper association of atom indices, our method begins by selecting the most similar atom pair (a₁, b₁) from two distinct configurations, Structures A and B, respectively. This selection is based on a symmetry-invariant descriptor, which is obtained by extracting the “0e” channel from the symmetry-equivariant node features after the message-passing blocks. The optimal pair is then identified by minimizing the difference between their descriptors:

where ϕ_a^A, ϕ_b^B are symmetry-invariant descriptors of a₁ and b₁, respectively.

Subsequently, we calculate neighbor lists $$ \mathcal{N}_{a_1}^{A} $$ and $$ \mathcal{N}_{b_1}^{B} $$ of a₁ and b₁ within a predefined radius cutoff. In our implementation, we set a default cutoff of r_cut = 3 Å. If fewer than three neighbor candidates are available, r_cut is increased in 0.5 Å increments until at least three neighbors are obtained.

From these neighbor lists, we select two pairs (a₂, b₂) and (a₃, b₃) by identifying the combination that minimizes the sum of their descriptor distances:

To ensure physical consistency, we further applied geometric criteria by comparing lengths of triangles formed by atom triplets (a₁, a₂, a₃) and (b₁, b₂, b₃), with tolerance of 0.5 Å. If no triplet satisfying the triangle tolerance is found after testing candidate combinations, only translation is applied.

We then compute the rotation matrix, R, required to align vectors ($$ \overrightarrow{b_{1}b_{2}} $$, $$ \overrightarrow{b_{1}b_{3}} $$) from Structure B with vectors ($$ \overrightarrow{a_{1}a_{2}} $$, $$ \overrightarrow{a_{1}a_{3}} $$) from Structure A, including improper rotations, encompassing inversion and reflection symmetries. Depending on structural periodicity, we adapt the selection and application of R as follows:

• For structures with periodic boundary conditions (PBCs), we first identify the closest available rotation matrix from the list of possible rotations calculated through spglib^[39].

• If spglib provides only the identity matrix, we first convert the rotation matrix to fractional coordinates with the cell matrix, L, as given by:

then round this fractional rotation matrix, R_frac, to the nearest integer values, and check its orthogonality for rigorous lattice transformation.

• For non-PBC systems, the computed rotation matrix is used directly.

Once the appropriate rotation is determined, we apply it to initially align atom b₁ in Structure B with a₁ in Structure A. This provides a starting point for a full permutation alignment using the Hungarian algorithm^[40]. We then perform a translation of the entire structure to minimize atom-wise root-mean-square deviation (RMSD), with an additional rotation applied for non-PBC systems. Finally, we calculate the atom-wise displacement vectors required to move each atom in Structure A to its corresponding position in the aligned Structure B.

This symmetry-invariant descriptor-informed structural alignment provides a well-defined displacement for atomic structures within our FA-based NARA framework.

Integration of learning loss into graph neural network interatomic potential

As depicted in Figure 1, the main module comprises multiple (usually 3 to 5) message-passing blocks, followed by an energy prediction block, which yields per-atom energies from symmetry-invariant “0e” scalar channels of node features. For periodic systems, we construct the pairwise distance matrix by explicitly including across-boundary atom pairs to ensure consistency under PBCs, as detailed in Supplementary Figures 1 and 2. Atomic forces are computed by differentiating the predicted energies with respect to atomic positions. So, the main module’s loss, $$ \mathcal{L} $$_target, is defined as follows^[41],

Figure 1. Architecture of GNN used in this work consists of a main module responsible for predicting energy and forces, and a loss prediction module that estimates the predicted loss, $$ \hat{\mathcal{L}} $$. In each message passing block, the invariant component, the “0e” channel, is extracted and subsequently concatenated after passing through a FC layer. The loss prediction module employs standard nonlinear functions, ReLU and Softplus, which introduce nonlinearity and enforce stable positive outputs. GNN: Graph neural network; FC: fully connected; ReLU: rectified linear unit; GT: ground truth.

where $$ \hat{E} $$ is the predicted energy, E is the reference DFT energy, N is the number of atoms, and λ_E, λ_F represent weighting factors for energy and forces, respectively.

To integrate learning loss into this main prediction module, we attached a loss-prediction submodule after minor modification. Specifically, instead of employing global average pooling layers, we extracted symmetry-invariant (“0e”) node features from each message-passing block and concatenated these features to predict the loss corresponding to the deviation of atomic forces between prediction and reference. The submodule was trained using the margin ranking loss approach described in the original paper^[42], pairing two atomic environments within each batch to ensure consistency between the trends of predicted loss, $$ \hat{\mathcal{L}} $$ and the ranking of target loss, $$ \mathcal{L} $$, defined by Equation (8)

where y equals +1 if $$ \mathcal{L}_{i} $$ > $$ \mathcal{L}_{j} $$, and -1 otherwise, with ξ set to 1 as the predefined margin. Therefore, the total loss for the integrated graph neural network (GNN) model, comprising both the main module and loss prediction submodule, is defined as^[42],

The uncertainty-quantification performance of the learning loss approach, together with an assessment of whether introducing the learning loss submodule affects energy and force prediction, is reported in Supplementary Figures 3 and 4. The associated computational overhead is summarized in Supplementary Table 1.

NARA framework for robust global optimization of atomic structures

Evolution of atomic structures in the FA

The second limitation concerns the structural evolution step in the FA, which requires a consistent atom-to-atom mapping to define meaningful displacement vectors. In prior works^[43,44], non-periodic structures were aligned by translating to a common center-of-mass, applying rigid-body rotation using the Kabsch algorithm^[49], and performing permutation matching via the Hungarian algorithm^[40]. However, under PBCs, these methods often neglect symmetry operations such as rotation, reflection, or inversion. As illustrated in Figure 2, when Structure B is symmetry-transformed from Structure A, the resulting displacement vectors become inaccurate. Although identical structures can be filtered using symmetry-invariant fingerprints, such global approaches fail for cases where structures are mostly similar but contain localized differences, leading to a physically inaccurate representation of atomic displacements and reducing the efficiency of structural evolution.

Figure 2. Figure 2. (A) Example of structural matching between Structure A and Structure B during the FA. Structure B is modified using symmetry operations to minimize the RMSD relative to Structure A; (B) Flowchart of structure evolution in the FA. FA: Firefly algorithm; RMSD: root-mean-square deviation.

To resolve this, we incorporated per-atom symmetry-invariant descriptors extracted from the GNN, enabling atom matching based on local environments rather than atomic coordinates. The descriptor-informed approach ensures that interpolations between structures remain physically meaningful and symmetry-consistent even under PBC. Figure 2 outlines the procedure, while detailed algorithms and mathematical formulations are provided in the Methods. All implementations are available in the NARA code (see Code availability).

On-the-fly FA searches

As previously discussed, we implemented an active learning framework as depicted in Figure 3, to enhance the robustness of our machine learning-accelerated global optimization. This framework automatically performs calculations using the target potential (e.g., DFT) based on the uncertainty quantifications provided by the learning loss submodule attached to the GNN main module, subsequently retraining the GNN to ensure reliability. Here, a critical aspect of this framework is deciding which structures to recalculate.

Figure 3. The active learning procedure implemented in NARA. The green-colored boxes in the middle column represent energy evaluations performed using the GNN, while the orange-colored boxes in the right column correspond to energy evaluations based on the target potential (e.g., DFT). NARA: Nature-inspired algorithm for robust atomic structure search; GNN: graph neural network; DFT: density-functional theory; SCF: self-consistent field calculations; LCB: lower-confidence-bound.

In molecular dynamics simulations accelerated by machine learning interatomic potentials (MLIPs), a common active learning strategy is recalculating structures whose uncertainty exceeds a predefined threshold^[45-47,50]. Even though it is effective for improving calculation stability, this approach does not align well with the goal of global optimization, which focuses primarily on identifying lower-energy structures. Indeed, results from previous work^[51] indicate a significant drop in success rates when retraining is driven predominantly by uncertainty alone. Therefore, we adopted a lower-confidence-bound (LCB) approach^[51,52] to balance thermodynamic stability and structural uncertainty. Specifically, structures that update the lowest LCB values are selected for recalculation. The LCB for a given structure X is defined as follows,

(10) $$ E_{\rm LCB}(X) = \hat{E}_{\rm GNN}-\kappa \mathcal{S}\left(\hat{\mathcal{L}}\right) \quad . $$

Here, $$ \hat{E} $$_GNN is the energy predicted by the GNN model, and κ is a weighting factor balancing uncertainty and stability. $$ \hat{\mathcal{L}} $$ is the per-atom uncertainty proxy (in arbitrary unit) learned from the trend of per-atom force-error loss, and S is a scaling function that maps the summed per-atom uncertainties to an energy-error scale by a linear least-squares fit to the observed absolute energy error, as follows,

(11) $$ |{\hat{E}-E}|=\mathcal{S}\left(\hat{\mathcal{L}}\right)=\sum_{k}\sum_{i \in j_{k}}\left(a_{k}\hat{\mathcal{L}}_{i}+b_{k}\right)\quad , $$

where the index k denotes the element type, and j_k represents the set of atoms corresponding to element type k. The coefficients a_k and b_k are the parameters obtained from linear fitting for each element.

With the defined LCB, our active learning procedure follows these steps: In the absence of an initial dataset, we start by generating n_init random structures to build an initial GNN model (GNN₀). We then generate other random structures, relax them with loose convergence criteria (using a convergence criteria ten times looser than the standard force convergence criteria), and employ k-means clustering to select representative structures. These structures are fully optimized using the target potential to form the initial training dataset for GNN (GNN₁). If an initial dataset already exists, it directly constructs the first GNN model (GNN₁). Subsequently, the trained GNN guides structural evolution within the FA through iterative generations. Each generational evolution comprises five stages, as detailed in the orange box (lower left panel of Figure 3): (a) Assess brightness (-E) for all agents; (b) Calculate a distance matrix among agents based on the differences between global descriptors, obtained by summing the per-atom features; (c) Evaluate each agent’s most attractive partner by implementing brightness attenuation as a function of distance, ensuring minimal interaction between distant agents; (d) Agents move towards their respective most attractive neighboring agents, generating new structures according to the FA’s evolution rule; (e) After relaxing newly generated structures, their LCBs are evaluated. If new minima are identified, those agents are frozen and do not participate in communication with other agents. Once the number of frozen agents reaches a predefined threshold, recalculations using the target potential are performed. The resulting structures are then added to the training and validation datasets, followed by retraining of the GNN. This iterative process continues until the global best solution remains unchanged for a predefined patience limit, P_max.

Through this comprehensive procedure, our framework enables robust, machine learning-based global optimization even without an initial dataset. We validated its effectiveness in the gold cluster systems which is discussed later in this work.

Application: the “8” surface oxide on Cu(111)

We evaluated the efficiency of our proposed FA-based NARA framework in identifying the global minimum structure of the previously found “8” surface oxide on Cu(111) reported in our earlier work^[16]. Basin hopping (BH) was chosen as the comparison method because it is a well-established global optimization approach for atomistic structure searches, including recent applications to complex surface systems^[11]. It also serves as a conceptually clear reference for the non-communicating limit of the FA when the attractiveness term becomes negligible at large γ. We therefore use BH as a representative reference method for evaluating the efficiency of our FA-based framework in this case study. The same GNN model trained earlier for uncertainty quantification was employed here. Since our previous study confirmed the capability to accurately find the global minimum using this dataset, active learning was turned off, and only the GNN model was used for energy/force evaluations without the attached submodule.

For the comparison, BH was configured with a maximum atom-wise displacement of 0.5 Å, employing a Metropolis acceptance criterion at a temperature of T = 300 K. The FA was set up with 30 firefly agents, using parameters α =0.5, β₀ = 0.9, γ = 0.02. The trajectory of FA agents on the PES is visualized in Figure 4A. The PES was constructed by training a VAE-based neural network with local attention scores, using all structures encountered during the FA search, mapping them into a 2D latent space. Energy levels are indicated by colors, and the known global minimum, “8”@2023t is marked with a white star as a reference. Initially, at the first generation (Gen. 1), the agents, marked with black stars, are located in higher-energy regions since they are just local optimized structures after random initialization. From the second generation (Gen. 2) onward, the evolutionary trajectories which are depicted by black arrows, illustrate agents moving toward lower energy regions. Note that the evolutionary trajectories include random perturbations (αε), meaning agents do not move straight toward the adjacent minimum.

Figure 4. (A) Scatter plot on the PES generated by an attention-based VAE, colored according to relative energy compared to the global minimum. Black stars indicate agent positions in FA searches, with arrows tracing their evolutionary paths; (B) Performance comparison of 24 independent single BH runs (colored in blue) and 24 independent single FA searches (with 30 agents, colored in red) based on the number of local optimizations vs. energy difference from the global minimum; (C) Comparison between 24 independent FA and parallel BH runs. Both consist of 30 agents in each generation, and start from the same initial structures generated from FA runs. In (B) and (C), lines with markers indicate the mean and shaded regions indicate the standard deviation over 24 independent runs. PES: Potential energy surface; VAE: variational autoencoder; FA: firefly algorithm; BH:

To statistically evaluate the FA efficiency, we performed several independent FA and BH runs, comparing the mean and standard deviation of energies as a function of the number of local optimizations. Each BH step corresponds to a single local optimization, and the mean and standard deviation of 24 independent BH runs are plotted in blue markers and shaded regions in Figure 4B, respectively. In contrast, since each generation of FA search contains 30 local optimizations and the agent with the lowest energy is selected, the number of local optimizations is represented by multiplication of the number of generations by 30. The mean and standard deviation of 24 independent FA is plotted in Figure 4B as red markers and shaded regions, respectively. Since the results of FA are far below those of BH, this indicates that a single FA run will likely achieve more efficient global optimization than a single BH run, given the same number of local optimization calculations.

However, this direct comparison between a single BH run and a single FA run may be influenced by two underlying factors. One is indeed the efficiency of the FA algorithm, and the other is the advantage of using multiple agents simultaneously. Since the best agent was selected from the population in each generation, this likely reduced the influence of randomness. To remove all other factors than the algorithm’s efficiency, we utilized a parallelized BH method, simultaneously running 30 independent BH calculations per generation, similar to FA.

Through this setup, the modified BH run becomes technically equivalent to the FA when the inter-agent communication term (β) is removed and the Metropolis acceptance criterion is adopted. By initializing each BH run with the same random structures generated from the FA runs, we effectively eliminated all factors other than inter-agent communication (or mating). A total of 24 independent parallel-BH runs were conducted, resulting in 720 independent BH runs overall. As shown in Figure 4C, although the parallel-BH results (shown in green) begin from the same starting point at generation 0, they exhibit broader standard deviations (wider shaded regions) and slower convergence (higher-positioned markers) compared to the FA. This strongly indicates that inter-agent communication within the FA significantly enhances robustness in global optimization, consistent with previous findings^[53], and highlights the advantage of such communication (or mating) in evolutionary algorithms over single-agent stochastic global optimization methods such as BH.

Application: gold cluster

The previously discussed “8” surface oxide on Cu(111) system exhibited a relatively simple PES, lacking clearly distinguishable local minima apart from the global minimum. Therefore, to highlight the advantages of the FA algorithm, we selected the gold cluster (n_Au = 11) system as the target system, which has been systematically analyzed by Goldsmith et al.^[26]. This system serves as an excellent example demonstrating how the choice of target potential can significantly affect the predicted global minimum structure. While the global minimum predicted using the PBE exchange-correlation (xc) functional is a planar edge-capped elongated hexagon structure, the vdW or MBD corrected xc functional displayed a structurally distinct global minimum, with a non-planar trigonal prism structure, which is also supported by experimental results. In this case, even with the use of state-of-the-art active learning combined with highly efficient and robust global optimization algorithms, relying solely on the PBE xc functional as the target potential leads to the identification of only the planar structure. In fact, the more efficient the global optimization becomes, the more likely it is to overlook the non-planar configuration under a fixed target potential. To address this limitation, one could either adopt a more accurate target potential, such as a vdW-corrected functional [e.g., hybrid Heyd-Scuseria-Ernzerhof (HSE) functional including MBD effects^[26]], or perform multiple global optimization runs to statistically sample metastable configurations. However, since the FA algorithm is inherently designed to identify multimodal solutions beyond the global minimum, it is expected to capture both planar and non-planar configurations even with the PBE xc functional in a single run, as schematically illustrated in Figure 5A.

To validate this expectation, we applied NARA to the gold cluster (n_Au = 11) system using an active learning-based global optimization scheme. For constructing the initial dataset, we set n_init = 20 for single-point calculations and n_{force_init} = 2 for force evaluations to configure our first GNN model (GNN₁). During the FA searches, NARA called the 13 additional short local optimization calculations based on target potential (i.e., DFT with PBE xc), based on the LCB criterion. The GNN architecture employed irreducible representations of “32 × 0e + 32 × 0o + 32 × 1e + 32 × 1o”, 3 message passing blocks, and a radius cutoff 5 Å. The FA parameters were set to α = 0.5, β₀ = 0.9, γ = 0.05. Compared to “8” surface oxide on Cu(111) case, a larger γ was used to strengthen the distance decay in the attractiveness term, making the interactions more short-ranged. The optimization terminated at Generation 241 when the global minimum remained unchanged for the predefined maximum patience, P_max = 75. The agents from the final generation were overlaid on the PES created from all previously encountered gold cluster structures, as illustrated in Figure 5B. Interestingly, several distinct structures were found, including both the planar structure (global minimum when using PBE xc) and the non-planar structure (global minimum when including vdW corrections), demonstrating the robustness and insensitivity of the FA-based NARA algorithm to target potentials.

Figure 5. (A) Schematic illustration showing robustness to error in the target potential. In E_PBE, Structure B represents the global minimum, but with vdW corrections, Structure A becomes more stable. Conventional GO algorithms always tend toward Structure B, the minimum of the chosen target potential, whereas the FA explores both minima, capturing both A and B (shaded region); (B) Agents from the final generation from a single NARA active learning run are shown on the PES. Top-left structure: non-planar global minimum (capped trigonal prism). Bottom-left structure: planar global minimum (edge-capped elongated hexagon). vdW: Van der Waals; GO: global optimization; FA: firefly algorithm; NARA: nature-inspired algorithm for robust atomic structure search; PES: potential energy surface; PBE: Perdew, Burke, and Ernzerhof.