Variational Autoencoders Improve Energy Minimization Stability

By Dr Silpaja Chandrasekar, PhDReviewed by Susha Cheriyedath, M.Sc.Aug 1 2024

In a paper published in the journal Machine Learning Science and Technology, researchers address instabilities in energy minimization using deep neural networks (DNN) by employing variational autoencoders (VAEs). By creating a compressed and regular representation of ground-state density profiles, VAEs reduce numerical instabilities and variational biases. Tests on one-dimensional and three-dimensional models show accurate energy estimations and density profiles with minimal errors. Additionally, transfer learning with pre-trained VAEs proves effective for different potentials.

Related Work

Past work in density functional theory (DFT) has struggled with the instability of machine learning (ML) models for energy minimization, especially with DNN and gradient descent methods. Challenges in DFT include severe instabilities in machine learning models during energy minimization, particularly when using NNs and gradient descent methods. Small inaccuracies in functionals can amplify, resulting in noisy density profiles and large errors.

Optimized DL-DFT Approach

In this study, a deep learning-based DFT (DL-DFT) method is developed and tested using VAEs for single-particle Hamiltonians, focusing on both 1D and challenging 3D models. DFT aims to map the ground-state density profile to the ground-state energy, with the VAE learning a compressed encoding of realistic density profiles. This approach enables accurate energy minimization by avoiding excessive constraints and numerical instabilities arising from traditional DNN and gradient descent methods.

VAEs are employed to create a regular and compressed latent space for density profiles. The encoder network processes density profiles to a latent representation while the decoder network reconstructs the profiles from this latent space. Training involves minimizing a combined reconstruction and regularization loss, with hyperparameters such as the latent space dimension and the regularization factor (β) playing critical roles in model stability and convergence. The latent space ensures that only realistic profiles are explored during gradient descent, reducing the risk of instabilities and variational property violations.

The DL-functional is trained to map density profiles to their corresponding energy values using supervised learning. Optimization of this function involves gradient descent within the latent space of the VAE, where the decoder manages constraints such as normalization and positivity. The gradient descent method is adjusted based on the regularization loss and the latent space dimension, enabling stable and accurate minimization. This approach effectively reduces numerical artifacts and improves performance, as demonstrated by accurate energy and density profile predictions in various test models.

TVAE Accuracy Validation

The initial analysis assesses the VAE's accuracy in reconstructing density profiles, focusing on two 1D potentials: Gaussian and speckle. The goal is to determine whether the VAE can replicate input density profiles accurately without losing significant information. The hyperparameter β, which balances the regularization and reconstruction loss, is crucial in this process.

A latent space dimension of 4 is used for the Gaussian potential, while a larger dimension is chosen for the speckle potential due to its increased variability. The team measured the reconstruction accuracy using the integrated absolute density difference, with results indicating that reconstruction errors decrease as β increases. However, excessively high β values lead to distortions in the reconstructed profiles.

The VAE, trained to encode density profiles, is then utilized to guide the energy minimization of the DL function through gradient descent. This process ensures the minimization converges to the ground state and avoids spurious constraints and numerical instabilities. Analysis of 1D Gaussian and speckle potentials reveals that an overestimated latent space dimension or inappropriate β values can result in positive or negative energy discrepancies. Proper tuning of these parameters generally results in energy discrepancies below the threshold of chemical accuracy, with visual examples demonstrating the impact of β on the accuracy of reconstructed density profiles.

The study also extends to a challenging 3D Gaussian-scatterer potential. Here, the analysis shows that energy discrepancies after gradient descent remain below chemical accuracy, and the accuracy of the reconstructed density profiles is confirmed through slices at different z-coordinates. This extension highlights the VAE's ability to handle complex, higher-dimensional systems effectively, reproducing density profiles accurately in three dimensions.

Lastly, the transferability of trained VAEs is evaluated by applying them to modified 3D testbeds with different numbers of scatterers. The results demonstrate that VAEs trained on one dataset can solve novel problems with other features while maintaining chemical accuracy. They suggest that VAEs trained on specific datasets can generalize to new issues, enhancing their applicability in diverse DL-DFT scenarios.

Conclusion

To sum up, this investigation demonstrated that β-VAEs effectively created compressed and regular density profile representations for diverse physical systems. The combination of VAE and convolutional models enabled efficient and stable energy minimization of DL energy functionals, leveraging automatic differentiation and performance. The study analyzed various testbed models, including 3D Hamiltonians, and confirmed that suitable parameters could be identified to avoid variational property violations. Additionally, transfer learning was feasible, allowing pre-trained VAEs to address novel potentials.