End-to-end differentiability and tensor processing unit computing to accelerate materials’ inverse design

Numerical simulations, such as density functional theory and molecular dynamics, are valuable for predicting material properties. However, applying these simulations to inverse design – identifying optimal materials with desired properties – is hampered by high computational costs and the lack of differentiability in traditional simulations. The computational burden prevents thorough exploration of the design space, and the absence of differentiability hinders integration with gradient-based optimization methods. To overcome these limitations, researchers often replace simulations with differentiable surrogate machine learning models, but this approach can be inaccurate and limit the discovery of novel material designs. This paper introduces a novel approach that avoids surrogate models by using end-to-end differentiable simulations within a deep generative pipeline. This is exemplified through the inverse design of a porous matrix with a targeted sorption isotherm, implemented using a differentiable lattice-based density functional theory (LDFT) code in TensorFlow. The resulting pipeline leverages the power of TPUs to accelerate the computationally intensive simulations.

Existing inverse materials design methods often employ surrogate machine learning models (e.g., autoencoders, GANs, generative inverse design networks) as differentiable predictors. These models approximate the relationship between material structure and properties, enabling gradient-based optimization. However, these methods require training both the generator and predictor, which can be complex and lead to limitations in discovering new, non-intuitive designs. The accuracy and generalizability of the surrogate model significantly influence the generator's capability. Recent advancements in automatic differentiation have led to differentiable programming platforms (TensorFlow, JAX, TaiChi) that are being used in differentiable simulation applications; however, their application in material science simulations remains relatively unexplored. This paper seeks to address the shortcomings of surrogate models by directly utilizing differentiable simulations within an inverse design framework.

The researchers developed an end-to-end differentiable simulator based on LDFT. They modeled a porous matrix as a square N-by-N lattice (or cubic N³ lattice for 3D simulations), where each pixel (voxel) can be solid (ηi = 0) or a pore (ηi = 1). The water density (pi) in each pore is calculated iteratively using LDFT, which involves a convolution operation. This LDFT equation was decomposed into differentiable computational layers in TensorFlow: an input layer, a CONV layer (representing the convolution), and an output layer. Multiple CONV layers (M iterations) are used to achieve convergence in water density at each relative humidity (RH). The entire process is differentiable, allowing for gradient backpropagation. The accuracy of the differentiable simulator was compared against a reference (undifferentiable) simulator, showing a negligible loss (<0.36%) with sufficient iterations (M ≥ 100). For inverse design, a generator model, structured as dual, parallel deconvolution blocks (one for low-RH and one for high-RH), was integrated with the simulator. The generator takes a target sorption isotherm as input, produces a porous matrix structure, and the simulator calculates the actual sorption curve. Gradient backpropagation optimizes the generator to minimize the difference between the input and output sorption curves. The simulator's convolutional layers have fixed weights, avoiding the complexities of simultaneous optimization of the generator and predictor. The training utilized a large dataset of 6,400,000 synthetic sorption isotherms (generated by a self-defined function mimicking real-world characteristics), with a validation set of 8769 curves generated by the simulator from randomly generated porous matrix grids. TPUs were used to accelerate the training process, significantly outperforming a GPU, especially at large batch sizes. The approach was then extended to 3D porous matrices and to the inverse design of materials with hysteresis behavior by modifying the convolution layers and loss functions accordingly.

The differentiable simulator accurately replicated the results of the non-differentiable LDFT simulation, with minimal loss. The generator-simulator pipeline successfully generated porous matrix structures that closely matched the target sorption isotherms (average prediction loss of 3% for 2D and 2.9% for 3D). TPU computing significantly accelerated the training process compared to GPU computing, especially for large batch sizes and grid sizes. The approach was successfully extended to 3D porous matrix inverse design, albeit with increased computational demands. The generation of porous solids with target hysteresis was achieved, demonstrating the method's adaptability to complex properties. The difference between the target and the output sorption curve is minimal, demonstrating the success of the proposed methodology. In the case of 3D models, the average percentage loss is around 2.9%, showing a slight increase in the error but still demonstrating good performance.

This work successfully demonstrates a robust inverse design pipeline based on end-to-end differentiable simulations. The key advantage lies in directly training the generator using the simulator as a predictor, bypassing the limitations of surrogate machine learning models. This eliminates the need for a predefined dataset and enables accurate extrapolation beyond the training space. The use of TPUs significantly accelerates the training, highlighting the potential of TPUs in scientific computing beyond deep learning applications. While the current study utilizes a simplified sorption model, it has significant implications for accelerating materials discovery and design. It shows that differentiable simulations are a promising tool for machine learning pipelines, and it suggests that TPUs can efficiently handle complex scientific simulations, expanding their applicability beyond deep learning. The successful design of porous structures with tailored sorption isotherms offers potential for applications like CO2 capture, gas separation, and drug delivery.

This research successfully developed a robust pipeline for inverse materials design using differentiable simulations. The method’s strength lies in directly training the generator using a differentiable simulator, eliminating the reliance on surrogate models and improving accuracy and extrapolation capability. TPU computing substantially accelerates the training. While the research used a simplified sorption model, the methodology holds promise for various materials design applications, highlighting the value of differentiable simulations and the potential of TPUs in scientific computation.

The study utilizes a simplified model of porous matrices (2D and then 3D lattices) and a simplified sorption model. Extending the approach to more complex materials and properties will require further research and potentially more computationally intensive simulations. The complexity of the 3D model limits the maximum size of the grid that can be efficiently processed, even with TPUs. The accuracy of the hysteresis prediction is slightly lower than for the simple sorption isotherms, which might be due to the complexities of hysteresis behavior in disordered porous materials.