Inverse design of two-dimensional materials with invertible neural networks

In materials discovery, the goal is to select and test candidates that best satisfy a targeted functional property while maintaining broad diversity in the search space. High-throughput experiments or computational modeling (e.g., DFT) are commonly used to rapidly screen candidates, but the vastness of chemical and materials spaces makes exhaustive evaluation costly. Machine learning offers surrogate models for property prediction and generative models for proposing candidates with specified target properties. Generative modeling enables inverse design by sampling a continuous latent space that maps to unique materials. Inverse problems entail learning an inverse mapping x = f−1(y) from properties to design parameters; while forward mappings can be computed (e.g., via DFT), reverse mappings must be inferred. Prior approaches such as VAEs learn approximate inverse mappings but can face accuracy and stability challenges. Here, invertible neural networks (INNs) and conditional INNs (cINNs) are leveraged to learn forward and exact inverse mappings within a single model, potentially improving stability and avoiding issues such as GAN mode collapse. The proposed framework, MatDesINNe, integrates INNs with down-selection and gradient-based localization to generate high-fidelity candidates for targeted properties. The study demonstrates MatDesINNe on band gap engineering of monolayer MoS2 using a seven-dimensional design space (six lattice degrees for tensile/compressive and shear strain, and an out-of-plane electric field) to target specific electronic band gaps. The approach aims to deliver near-chemical accuracy, diverse candidate generation, and insights into metal–insulator transitions relevant to neuromorphic applications.

The paper situates inverse materials design within a growing body of data-driven discovery strategies that rely on ML for accelerated screening and generation of candidates. Generative models (VAEs, GANs) have shown promise for inverse molecular and materials design but often suffer from mode collapse (GANs) and limited invertible expressivity (VAEs). Mixture density networks (MDNs) provide a probabilistic inverse model but can be limited by the Gaussian mixture assumption in complex, non-Gaussian distributions. Invertible neural networks (INNs) and conditional INNs (cINNs) have emerged as robust alternatives for inverse problems, offering intrinsic invertibility, stable training, and the ability to sample full posteriors p(x|y). The study builds upon these advances to address band gap engineering in 2D materials, referencing strain and electric-field modulation literature that demonstrates experimental feasibility of tuning band gaps in dichalcogenides like MoS2. It also notes ongoing challenges such as developing invertible, rotationally and translationally invariant representations for crystal structures, which limit fully structural generative approaches and motivate a focus on parametric design spaces (strain, field) in this work.

Design space and data generation: The design parameters comprise six lattice degrees of freedom (a, b, c, α, β, γ) representing tensile/compressive and shear strain, and an out-of-plane electric field E. Sampling ranges: ±20% around equilibrium lattice parameters (a = b = 3.16 Å, c = 12.30 Å, α = β = 90°, γ = 120°) and E from −1 to +1 V/Å. Approximately 10,799–11,000 DFT calculations were performed to map the design space to the target property (band gap). DFT details: VASP with PBE (GGA) exchange-correlation, PAW method, 500 eV cutoff, spin polarization, Monkhorst–Pack 9×9×1 k-point grid. The computed PBE band gap at equilibrium is 1.97 eV, close to the experimental 1.90 eV, suitable for demonstration. MatDesINNe framework: Four steps: (1) training, (2) inference, (3) down-selection, (4) localization. Training: INN trained bidirectionally with forward L2 loss and backward MMD loss; the best forward model is saved as a surrogate. cINN trained by maximum likelihood with y as a conditioning input to affine coupling layers in both directions. Inference: sample latent z ~ N(0, I) and invert to obtain posterior samples x ~ p(x|y) for a given target y. Down-selection: filter out-of-distribution candidates (outside training parameter ranges) and those predicted by the surrogate to be far from the target. Localization: starting from posterior samples, perform gradient-based optimization (Adam) using the saved forward surrogate and automatic differentiation to minimize loss between predicted and target property, thereby localizing to exact inverse solutions. Invertible neural networks: INNs augment the forward mapping with a latent z to encode information lost in the forward process. The inverse is x = h(y, z; Θ), with z sampled from a Gaussian prior. Training losses: forward L2 loss Ly = ||y − yi||^2 and backward MMD loss on z; or a maximum likelihood objective incorporating log-determinant of Jacobian. cINN directly maps x to z conditioned on y and is trained via maximum likelihood L(z) = 0.5||z||^2 − log|det(Jz)|. Affine coupling layers enable efficient Jacobian determinant computation. Localization from posterior samples: (i) Prior exploration: draw posterior samples x from p(x|y) via invertible mapping from z. (ii) Gradient estimation: compute gradients of loss with respect to x via automatic differentiation using the surrogate f(x). (iii) Gradient descent: update x_{t+1} = x_t − γ g (Adam optimizer) to reach localized solutions near the true inverse. Implementation details: PyTorch + FrEIA for invertible blocks. Network: 6 invertible blocks; each block uses 2 fully connected layers with 256 neurons and ReLU activations. Optimizer: Adam, learning rate 1e−3, weight decay 1e−5. Noise stabilization during training with Gaussian perturbations σ1 = 5e−3 and σ2 = 2e−3. Dataset: ~11,000 DFT samples, 80% train / 20% test, batch size 500, 1000 epochs. Hardware: Ubuntu 18.04, 16 Xeon CPUs, NVIDIA Quadro RTX 6000; training time ~10–15 minutes per model on a single GPU. Baselines: CVAE with 3-layer encoder/decoder (128 hidden nodes each), latent dimension 3, Adam 1e−3, batch 500, 1000 epochs; trained with ELBO-style loss. MDN with 3 Gaussian components; categorical and mixture networks are 3-layer NNs with hidden dimension equal to input dimension; trained by likelihood maximization with Adam 1e−3, batch 500, 1000 epochs. Validation and benchmarking: Surrogate-validated performance assessed by generating 10,000 samples per model for targets Eg = 0, 0.5, 1.0 eV and computing MAE vs the surrogate. Computational time measured on CPU for generating 10,000 samples. Best model (MatDesINNe-cINN) further validated with DFT on 200 samples per target to obtain absolute performance relative to DFT.

- Data characteristics: Within the sampled strain and field ranges, most configurations reduce the band gap from the equilibrium 1.97 eV; the majority become metallic (Eg = 0), creating a highly imbalanced dataset. - Surrogate-validated benchmarking (10,000 samples per model): • MDN MAE (eV): Eg=0: 0.184; 0.5: 0.421; 1.0: 0.840. CVAE: 0.064; 0.461; 0.973. INN: 0.038; 0.527; 0.835. cINN: 0.063; 0.219; 0.193. MatDesINNe-INN: 0.038; 0.512; 0.321. MatDesINNe-cINN: 0.020; 0.013; 0.015. • Generation time for 10,000 samples (s): MDN ~5.3; CVAE ~5.5; INN ~5.9; cINN ~5.8–6.0; MatDesINNe variants ~28–34 (due to localization step). All are orders of magnitude faster than DFT. - Absolute DFT-validated performance (200 samples per target): MAE ≈ 0.008 eV (Eg=0), 0.111 eV (Eg=0.5), 0.067 eV (Eg=1). Outliers are rare: ~3% have non-zero gaps for Eg=0 target; <2% have zero gaps for Eg=0.5 and Eg=1 targets. - Speedup vs direct DFT screening: Generating 10,000 candidates with target Eg=1±0.1 eV by DFT would require ~1.03×10^7 s, over five orders of magnitude slower than MatDesINNe (~31 s CPU for generation/localization, excluding training/data generation). - Diversity and coverage: Generated samples reproduce the distribution of the training data across in-plane and out-of-plane strains, avoiding collapse to narrow regions; UMAP embeddings show broad, well-matched coverage. - Uniqueness: Percentage of generated samples unique relative to training data (tolerance average > 0.01 or max > 0.01 in parameter difference): Eg=0: 99.96% / 99.98%; Eg=0.5: 98.61% / 98.34%; Eg=1.0: 90.82% / 85.00%. - Physical insights: Band gap is more sensitive to shear strain than to tensile/compressive strain; few Eg=1 samples exist when average shear strain magnitude exceeds ~5%. Higher absolute strain correlates with increased probability of Eg=0, especially for shear. Eg=0.5 and 1 share much of the strain space, but Eg=1 diminishes at higher tensile strain compared to Eg=0.5. - Metal–insulator transition (MIT) mapping: UMAP of the 7D space reveals local regions where small perturbations in strain/electric field induce transitions between Eg=0 and non-zero gaps. Example transitions: Eg 0↔0.5 via ~7% y-axis tensile strain and 0.82 V/Å field change; Eg 0↔1 via ~6% y-axis tensile and ~7% z-axis compressive strain. These regions suggest pathways for fast, energy-efficient switching relevant to neuromorphic devices.

The study demonstrates that invertible neural networks, when integrated into the MatDesINNe framework with down-selection and gradient-based localization, can effectively solve inverse design problems in materials science. By learning a bidirectional mapping between design parameters and properties, the method generates high-fidelity candidates that meet target band gaps with near-chemical accuracy, addressing the one-to-many nature of inverse problems by sampling full posteriors p(x|y). The approach outperforms MDN, CVAE, and base INN/cINN baselines, particularly under severe class imbalance (many metallic states). The computational efficiency—minutes for thousands of candidates versus weeks of DFT—makes inverse design tractable in high-dimensional spaces. The generated candidates maintain diversity and uniqueness while providing interpretable insights: band gap sensitivity is dominated by shear strain, and low-strain regions favor non-zero gaps, offering design principles and guidance for efficient sampling. Mapping MIT boundaries in reduced-dimensional embeddings highlights regions where small perturbations trigger transitions, informing strategies for neuromorphic switching or, conversely, for avoiding unintended MIT. The framework is general and can be applied to other materials and properties, enabling closed-loop integration with autonomous experimentation by proposing high-quality candidates on the fly and learning from outcomes.

MatDesINNe, an inverse design framework based on invertible neural networks with down-selection and localization, enables rapid generation of diverse, high-fidelity materials candidates that achieve targeted electronic band gaps in monolayer MoS2. It delivers near-chemical accuracy (≤0.1 eV MAE in DFT validation), significant speedups over direct DFT screening, and robust performance in imbalanced, complex design spaces. Beyond candidate generation, the approach provides physical insights, such as the dominant role of shear strain in band gap tuning and identification of regions facilitating metal–insulator transitions. The framework is materials-agnostic and extendable to other properties and design spaces, offering a pathway to integrate ML-driven inverse design with autonomous experimentation. Future directions include adopting more expressive invertible architectures (e.g., neural spline flows), improving surrogate fidelity, and developing invertible, symmetry-invariant crystal representations to incorporate atomic and compositional degrees of freedom into the generative process.

- Data imbalance: The majority of sampled configurations are metallic (Eg=0), posing challenges for generative modeling; careful training and localization are required to avoid bias. - Surrogate dependence: DFT-validated performance is bounded by the accuracy of the surrogate used for localization; discrepancies arise when the surrogate deviates from true DFT values. - Model expressivity: Affine coupling layers in the INN/cINN may have limited expressivity for complex distributions; more expressive normalizing flows (e.g., neural spline flows) could further improve performance. - DFT functional limitations: PBE-GGA underestimates band gaps in general; while adequate here (MoS2 equilibrium gap close to experiment), broader application may require higher-level methods for training labels. - Scope of design space: The study focuses on parametric strain and field; atomic/compositional variations are not included due to the absence of suitable invertible, symmetry-invariant crystal representations. - Practical constraints: Strain ranges (±20%) may exceed experimental limits; such feasibility constraints are deferred to the down-selection stage and not hard-coded into training.