Combining variational autoencoders and physical bias for improved microscopy data analysis^∗

A VAE [30] is a deep generative probabilistic model that belongs to the family of probabilistic graphical models and variational Bayesian methods. The purpose of a VAE model is to infer (latent space learning) the factors of variability of the large-scale complex data. The two primary components of any VAE models are the encoder (stochastic inference model) and decoder (generative model). The encoder consists of the convoluted neural network layers, ${f_e},\,$ with $e$ being the parameters (weights and biases) and an activation function at each layer, followed by the flattening of data. Thus, given an N-dimensional input in the real space, ${\boldsymbol{{X}}}$, the encoder projects the input into a reduced n-dimensional (where $n < N$) latent variables, ${\boldsymbol{{Z}}}$, following a distribution $p({\boldsymbol{{Z}}}|{f_e}\left( {\boldsymbol{{X}}} \right))$. The decoder structure is the reverse of the encoder, which is another convoluted neural network layers, ${f_d}$ with $d$ being the parameters (weights and biases). From the latent maps, the data de-flattened and fit into the neural network with an activation function, and finally we score the observed data against the Bernoulli (Ber) likelihood parametrized by the decoder output into a reconstructed high-dimensional real space (similar to the input space). Here we model the decoder as outputting parameters (probabilities) of a Bernoulli distribution and the 'observation' (which is equivalent to the reconstructed output) is considered as a sample from this Bernoulli distribution. This is more in line with the underlying probabilistic model of the VAE and makes the probabilistic nature of the model more explicit. For more details, we refer to this introductory tutorial on VAE implementation as a deep probabilistic model [42]. Thus, given any sampling in the low-dimensional latent space, the decoder reconstructs the n-dimensional latent variables into the respective estimated N-dimensional input, $\overline {\boldsymbol{{X}}}$, following a likelihood distribution $p(\overline {\boldsymbol{{X}}} |{f_d}\left( {\boldsymbol{{Z}}} \right))$. This encoding-decoding process of the VAE models need to be optimized such that the models can best learn the training data with minimizing the loss of information. The loss function of a VAE, ${\mathcal{L}_{{\text{vae}}}}$, which we minimize, can be defined as the sum of reconstruction error and KL divergence. The reconstruction error $\varphi$ can be chosen as mean square error, the cross-entropy error. Here, we have considered the reconstruction loss as negative log likelihood function. The KL divergence [36, 37] ${D_{{\text{KL}}}}(p(z\left| {{f_e}\left( x \right))} \right||p\left( z \right))$ is the distance loss between the prior $p\left( z \right)$ distribution (usually chosen as standard Gaussian) and the posterior $p(z|{f_e}\left( x \right))$ distributions of the latent representation from data. VAE models have been implemented to materials systems on various tasks like classification, feature or pattern recognition, prediction though unsupervised, semi-, or supervised learning [41, 43–46].

$$\begin{equation}{\mathcal{L}_{{\text{vae}}}} = {\text{ }}\varphi + {\text{ }}{D_{{\text{KL}}}}\left( {p({\boldsymbol{{Z}}}\left| {{f_e}\left( {\boldsymbol{{X}}} \right))} \right|\left| {p\left( {\boldsymbol{{Z}}} \right)} \right.} \right)\end{equation} \tag{ 1 }$$

where $\varphi$ is the reconstruction loss, ${D_{{\text{KL}}}}(p(.\left| {x)} \right||p\left( . \right))$ is the KL divergence.

This general or vanilla VAE, as stated in equation (1), is not robust to the variability (e.g. rotation, translation, scale) of the data during the training process. However, the extension of this vanilla VAE has been attempted to develop VAE models which enforce rotational, translation or scale in-variances individually or together, while learning the individual continuous latent representations or jointly learning both continuous and discrete latent representation from the data, as implemented in pyroVED package in Python [47]. Here, we considered a specific example of shift-invariant variational autoencoder model (sh-VAE), as demonstrated in figure 1. As stated, sh-VAE is one of the extensions to vanilla VAE where the shift or translational variability of the data is enforced during model training. The encoding-decoding architecture of sh-VAE is very similar to vanilla VAE as described above, except the latent maps are separated into the conventional latent maps and special maps associated with the positional variances of the object of interest in the data. Here, the latent maps consist of latent vectors, sampled from normal distribution, and the translational transformation matrix, where the model learns the transformation matrix in an unsupervised way such that the learned latent maps are invariant to the translational transformation matrix. To describe the training process in sh-VAE, we use the first $d$ columns of the latent matrix or the special latent variables to shift the co-ordinate grid of the input data by $k\Delta t$ where $0 \lt k \unicode{x2A7D} 1$ represents the belief of the degree of disorder in the system, and $\Delta t$ is the matrix of first $d$ columns. This enhancement is based on the Galilean principle that mechanical equations of motion are invariant under the $t{^{^{\prime}}} = t + \Delta t$ transformation, which is not accounted for by the classical VAE. Then the shifted co-ordinate grid is augmented with the remaining columns of the latent variables or the conventional latent variables before passing to the decoder. The loss function of a beta sh-VAE, ${\mathcal{L}_{{\text{sh - vae}}}}$, can be mathematically written as equation (2)

$$\begin{equation}{\mathcal{L}_{{\text{sh}} - {\text{vae}}}} = {\text{ }}\varphi + {\text{ }}\beta \left( i \right){D_{{\text{KL}}}}\left( {p({\boldsymbol{{Z}}}\left| {{f_e}\left( {\boldsymbol{{X}}} \right))} \right|\left| {p\left( {\boldsymbol{{Z}}} \right)} \right.} \right)\end{equation} \tag{ 2 }$$

Zoom In Zoom Out Reset image size

where $\beta \left( i \right)$ is the continuous scale factor of KL divergence at $i{\text{th}}$ training cycle of sh-VAE. The scale factor encourages a better disentanglement of the data, thus provides better learning [48]. For more details on the development of the sh-VAE, readers are referred to [41]. Previously, sh-VAE has been implemented for unsupervised discovery of features, patterns, and order parameters from complex experimental microscopic (STM, STEM) images, where key features of several training images have positional variance[41]. However, the sh-VAE model (or any other VAE models) does not consider any physical characteristics of the data, which can reduce the accuracy of learning of such complex images, compromising physics relevant knowledge extraction.

Here, our enhance the performance of the VAE models by introducing inferential bias that allow to encode known physical laws and constraints. Thus, we propose a physics driven VAE (phy-VAE) model to facilitate the training process for the analysis of the atomic and mesoscopic images by combining the existing capability of VAEs to discover the factors of variability in the data under the stated invariances and the new capability of learning any physics-based global criteria to refine the discovery of valuable information from such complex data. It is to be noted that in the paper, though we have focused on the enhancement of earlier described sh-VAE to phy-VAE where we aim to inject known physical behavior of a system during the training process or the latent space learning, this enhancement is well suited for any chosen VAE models as per the data since none of the VAE models possesses any prior physical knowledge of the system or guarantee to converge learning in accordance to the known physics. In the next part, we will thus focus on the extension of the shift-VAE to the phy-VAE where we will focus on the enhancement (proposed contribution) part of the modelling. The other setups like encoder, latent map generation and decoder are same as described earlier in this section.

Figure 2 shows the general approach to integrate the physics driven criteria to a VAE model and develop phy-VAE. As an extension to sh-VAE, the loss function can be rewritten as below equation (3).

$$\begin{equation}{\mathcal{L}_{{\text{phy}} - {\text{vae}}}} = \left( {\varphi + {\text{ }}\beta \left( i \right){D_{{\text{KL}}}}\left( {p({\boldsymbol{{Z}}}\left| {{f_e}\left( {\boldsymbol{{X}}} \right))} \right|\left| {p\left( {\boldsymbol{{Z}}} \right)} \right.} \right)} \right)\left( {w + {{ \Psi }}} \right)\end{equation} \tag{ 3 }$$

Zoom In Zoom Out Reset image size

where ${{\Psi }}$ is a physic driven loss, $w$ is the slack (a small value) added to new physics driven loss with $0 \unicode{x2A7D} w \unicode{x2A7D} 0.5$, to avoid numerical error. In this paper, we derived two loss functions which serves the goal of better learning the phase boundaries with maximizing the smoothness of the learned spatial maps of the encoded latent variables. While the reconstruction and KL divergence loss learn the factors of variability of the data, the new loss component imposes the physical constraints of the data.

Here, the first smoothing loss for ${{\Psi }}$ is computed through Scharr edge detection filter. As such, it imposes the condition that the total length of the edges within latent image should be minimal, corresponding to the minimization of the interface length due surface tension. The second smoothing loss is computed through peak intensity of the Fourier transform of the spatial latent maps. According to the Parseval theorem, the total intensity of the Fourier image is equal to the roughness spectrum of the real-space image. The intensity of the Fourier image is partitioned between the peaks corresponding to the crystal lattice, the background corresponding to the disordered features (would be a ring if there was characteristic length scale), and the zero-wavelength peak corresponding to the large-scale features in the image. Additionally, the Moire effects due to the beating between the sampling grid and the crystallographic lattice will result in the clear Fourier peaks. By reducing the intensity of the Fourier peaks in latent images, we minimize these effects, leading to minimization of the smoothing loss. For the purpose of simplicity, we call these two losses as SL1 and SL2, which we augment to equation (3) as ${{\Psi }}$ to minimize. Interestingly, if we refer to equation (3), we have multiplied ${{\Psi }}$ with the conventional loss instead to addition. The decision is from our extensive analysis to find the best way to augment the physics loss. Though minimizing the loss functions when augmented in either adding or multiplying produces similar results, however the maximization of the loss function when ${{\Psi }}$ is added does not work as expected. As it turned out, with maximization, did not reduce the smoothness of the latent space. As we know, for the optimization problem, the regulation of the results should follow as we minimize or maximize, which worked for multiplying ${{\Psi }}$. We will further clarify during the analysis in the result section. Algorithms 1 and 2 provides the detail workflow to compute SL1 and SL2.

Algorithm 1. SL1: edge magnitude of latent images (considering latent dim = 2).

1.   Generate the latent maps ${\boldsymbol{{Z}}} = \left( {{{\boldsymbol{{z}}}_{\boldsymbol{{1}}}},\,{{\boldsymbol{{z}}}_{\boldsymbol{{2}}}}} \right)$, given the current trained model, ${\boldsymbol{\Lambda }}$ and data ${\boldsymbol{d}}$ 2.   Normalize ${\boldsymbol{{Z}}} = \left( {{{\boldsymbol{{z}}}_{\boldsymbol{{1}}}},\,{{\boldsymbol{{z}}}_{\boldsymbol{{2}}}}} \right)$. 3.   Use a denoising filter to the latent maps. Here we use bilateral filter with proper tuning as per the data. The reason we choose this filter as it an edge-preserving, denoising filter. We considered the function 'denoise_bilateral' under scikit-image Python library for bilateral filter. Let us annotate the denoised latent maps as the output from bilateral filter as ${\boldsymbol{{Z}}} = \left( {{{\boldsymbol{{z}}}_{\boldsymbol{{1}}}},\,{{\boldsymbol{{z}}}_{\boldsymbol{{2}}}}} \right).$ 4.   Given the denoised latent maps ${\boldsymbol{{Z}}} = \left( {{{\boldsymbol{{z}}}_{\boldsymbol{{1}}}},\,{{\boldsymbol{{z}}}_{\boldsymbol{{2}}}}} \right),{\boldsymbol{{\,}}}$find the edge magnitude. Here we used Scharr transform, implementing the scikit-image Python library function 'filters.scharr'. The function outputs the edge maps (array) of $\left( {{{\boldsymbol{{z}}}_{\boldsymbol{{1}}}},\,{{\boldsymbol{{z}}}_{\boldsymbol{{2}}}}} \right)$ as ${{\boldsymbol{{e}}}_{{z_1}}},{\boldsymbol{{\,}}}{{\boldsymbol{{e}}}_{{z_2}}}{\boldsymbol{{\,}}}$respectively. 5.   Calculate loss: Compute the mean of the sum of the edge maps of denoised latent images as per equation (4), $$\begin{equation}{{\Psi }} = {{{\Psi }}_{{\text{SL}}1}} = \frac{{{{\mathop \sum \nolimits }}{e_{{z_1}}} + {{\mathop \sum \nolimits }}{e_{{z_2}}}}}{2}.\end{equation} \tag{ 4 }$$

Here, like the conventional VAE loss computation, these new losses are also calculated at each epoch of training in batches, considering the same batch data $\left( {\boldsymbol{{d}}} \subset {\boldsymbol{{D}}} \right)$ in calculating VAE loss. ${\boldsymbol{{D}}}$ is the full training data. The batch file, instead of full training data, is considered for computation due to the reasonable memory usage and computational cost of model training for practical implementations. However, due to using the batch file (low resolution latent maps), we approximated SL1 or SL2 instead of exact calculation with full training data (high resolution latent maps). In additional to the loss functions, we have also optionally integrated a visualization function (refer to figure 2) to track the boundary segmentation of the latent images at the current state of the phy-VAE training. We will discuss more on the boundary segmentation techniques implemented in the result section.

In algorithms 1 and 2, the process of generating the latent maps (step 1) by VAE is different, depending on the specific VAE model (e.g. vanilla VAE, rVAE, sh-VAE) we are implementing for a given problem. It is to be noted the phy-VAE is universal to integrate with any VAE models, however, for the demonstration we implemented here with sh-VAE. In this case, for the sh-VAE to encode any positional information of input image patches, we store their original coordinates in the image space and use them to construct a 'map' where we assign a value associated with a latent variable z to each coordinate. Here the first two column of the z matrix is used to shift the co-ordinate grid by $k\Delta t,\,\Delta t = z\left[ {:,\,:2} \right],\,0 \lt k \unicode{x2A7D} 1$, and then the transformed grid is concatenated with the remaining columns of the latent variables before passing to the decoder. The details of the workflow of the latent maps generation in sh-VAE can be found in [41].

For NiO-LSMO, we considered denoising factor = 0.1 for bilateral filtering during SL1 computation. For model training with either SLI or SL2, we considered 50 epochs for model training with $\beta \left( i \right)$ trajectory linearly increasing from 0.05 to 1 as the model completes 50 epochs. We considered the hyperparameter settings as done in [41] with a good initial guess through manual tweaking, for training the standard sh-VAE and phy-VAE models, to justify the comparison. It is to be noted, another approach to aim at improved (physically meaningful) data analysis would be to focus on optimizing the hyper-parameters, particularly the $\beta \left( i \right)$ trajectory. However, tuning this high-dimensional iterative dependent hyper-parameter would be a very time-consuming task (given the sh-VAE model is computationally very expensive) which need to be done prior to the training. However, one approach we took earlier in tuning such trajectory hyperparameter is through latent Bayesian optimization approach [35] but would still take a significant amount of time to conduct such optimization task. Here, bypassing the need and the mentioned challenges of prior tuning of particularly high-dimensional hyperparameters (e.g. $\beta \left( i \right)$), the proposed approach is more direct utilization of the known physical behavior of the system where the VAE model can train the latent space simultaneously in accordance with the applied known physical constraints. In this case, we see minimizing the objectives (SL1 or SL2) on latent map, ${z_1}$, oversimplifies the learning such that the domain boundary edges disappear. Therefore, we only optimize the latent map, ${z_2}$, following algorithms 1 and 2 as the phy-VAE model learns ${z_1}$ efficiently with the existing VAE loss function.

After several investigations, we considered the following approach to the visualization (ref figure 2) of the image segmentation of the latent maps of the full training data as the training progresses and the final learned latent maps. As we compute the losses with batch data, thus, it is essential to track the performance with whole data. First, we applied a threshold to the latent maps to smooth the full domain if there are negligible difference between the maximum and the minimum latent values. We set this threshold as 1 × 10⁻³. Once the above threshold is passed, then after normalizing, we denoise the latent maps with bilateral filter, setting the denoising factor = 0.1. Then we applied Otsu threshold and Chan-Vese methods, considering the inbuild functions in Python scikit-image library [50, 51], to enhance the boundaries between foreground and background object in the images, and image segmentation respectively. The reason for the selection of these methods, and this sequence of approach to ensure proper defining edges of the latent maps.

Figure 4 compares the results of the learned latent maps phy-VAE, at different slack values associated to SL1 loss. The same training with phy-VAE models with considering SL2 for physics driven loss is provided in figure S2 in supplementary material. We can see that sh-VAE yields highly satisfactory results (figure S1(a)) in segmenting the phases between NiO and LSMO and therefore the purpose of the phy-VAE is to preserve the existing training of shape reconstruction, without degrading due to enforcing physical constraints. Comparing figure 4, phy-VAE attain the similar phase boundaries of both the latent maps, and while we optimize only the latent map, ${z_2}$, we see it keeps the actual domain boundary shape and does not affect the boundary shape of ${z_1}$. Thus, we do not see any sign of distortion or discontinuity over the learned latent manifolds and latent distributions. We find similar performance for the phy-VAE models considering SL2 objective function (figure S2). Here, in the NiO-LSMO system, the known factors of variation are the phase composition, meaning the single-phase regions on rock salt NiO and perovskite LSMO. Based on general physical considerations, we expect the boundary between the phases to be smooth (minimization of the interfacial free energy). The unknown factors of variation here can include strains, or variation in the imaging conditions due to mis-tilt. In both cases, we expect these to be smooth on the level of the unit cells. Hence, the interpretation of the latent space and quality of representation disentanglement is established based on whether these constraints are satisfied. We also explore the latent distributions—the bimodal distribution will suggest that the system has only two components, whereas more complex distributions suggest either problems in VAE, or presence of additional factors of variability.

Zoom In Zoom Out Reset image size

Moving forward to more complex BFO system (BiFeO₃), during computation of SL1, we considered denoising factor (for bilateral filtering) as the standard deviation of the low-resolution latent images. For model training with either SLI or SL2, we considered 50 epochs for model training with $\beta \left( i \right)$ trajectory linearly increasing from 0.1 to 1 as the model completes 50 epochs. For the visualization part, we followed the similar approach to the previous case study, except considering denoising factor = 0.3, and considering multi-Otsu threshold instead of binary Otsu and Chan-Vese for the multi-label domain segmentation. For the BiFeO₃, the situation is more complex. Here, we still have two phases—the ferroelectric BiFeO₃ and SrTiO₃ substrate that can be visualized by human eye. However, within BiFeO₃ we additionally have ferroelectric domains. These are not detectable by human eye and can be visualized via physics-based analysis (atom finding). The goal for ph-VAE is to discover these based on images only and explore if there are additional (unknown) factors of variability. Again, based on general physics knowledge we postulate that domain walls should be sharp, whereas other factors of variability should change slowly on the length scale of a unit cell. From figure 5, we can now see the limitation of sh-VAE performance as we observe several scattered edges and grains within the phase which deviates the physical behavior that typically order parameters in physical systems are postulated to change smoothly on the length scale of the unit cell. It is to be noted that we followed the same visualization process for image segmentation of the sh-VAE learned latent maps as for the phy-VAE, and thus explain our earlier statement of requirement of enforcing the physics driven knowledge in phy-VAE in spite of the image analysis procedure. Figure 6 shows the latent maps of phy-VAE at different training epochs (following figures 6(a)–(e)), considering for SL1, $w = 0.1$. Looking at the edge detection image of ${z_2}$ in figure 6(b) (iteration 5), we see the first indication of the SrTiO₃ substrate phase (blue region) which gets more refined as the model trained with more iterations (referring 6(c)–(f)). Subsequently, the large right region of the edge detection image of ${z_2}$ for those figures is the ferroelectric phase. Similarly, we can also see the first indication of the domain boundary between the SrTiO₃ substrate phase and ferroelectric phase in the edge detection image of ${z_1}$ in figure 6(d), which gets refined with progression of the training (referring 6(e) and (f)). Within the ferroelectric phase (referring 6(c)), we see the formation of the curved boundary to segment the ferroic variants, which gets refined as the training progress (referring figures 6(d)–(f)) and at the same time minimizing all the small edges and grains within the variants of the ferroelectric segments, following the physical behavior as the model trained after iteration 50 (referring figure 6(f)). Figure 7 compares the results of the learned latent maps phy-VAE, at different slack values associated to SL1 loss. The same training with phy-VAE models with considering SL2 for physics driven loss is provided in figure S3 in supplementary material. We clearly see the phy-VAE with both physics driven losses outperformed sh-VAE (figure 5) in restoring the physical behavior. However, interestingly, we see for figure 7(a), the learned latent map 2 does not detect the ferroic variants while the respective latent map 1 provides well defined boundaries. This can happen due to the fact that latent map 1 learn all the factors of variability of training data in reconstructing the phase, as we can see from the encoded latent distributions of the same. Likewise, for all the cases, we see the learned latent manifolds and distributions does not have any distortion or discontinuities, unlike the case for training with sh-VAE. To further clarify the reasoning of multiplying ${{\Psi }}$ in equation (3), we provided figure S4 in supplementary material to showcase an example of the regulation of the physics driven loss functions towards the trained latent maps as we maximize those, where we see the increments of boundary edges, distortion, and discontinuities over the latent space.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

From the application of phy-VAE on previous two material systems, we see the model on one hand perverse the goodness of the phase extraction over the latent maps for both physics driven losses, and on the other hand able to enforce the physical constraints to learn the data in terms of relevant physical behavior of the system. Therefore, we attempt to implement the phy-VAE model into the graphene system where the domain phase is more complex (refer figure 3(c)) for interpretation. Here, during computation of SL1, we considered the same setup as for BFO system. For model training with either SLI or SL2, we considered 200 epochs for model training with $\beta \left( i \right)$ trajectory linearly increasing from 0.01 to 0.5 as the model completes 150 epochs and then kept constant at 0.5 for the remaining 50 training epochs. The image segmentation process also remains the same as for BFO system. Figure 8 shows the results of the learned latent maps with phy-VAE, considering SL1 and SL2 losses. Finally, for graphene we see complex patterns associated with long-range quantum interference effects due to defects (see figure 8(a)). In this case, we explore the phy-VAE behavior on an unknown system. We segment the images for edge or phase boundary detection with considering 2, 3 and 4 phases. In the figure, the coloring of the segmented phases is as per the provided colormap where 'L' denoted low values and 'H' denoted high values. We see the edges and corners of the latent maps has the domain with color representing smaller values, which gradually increases as the distance of latent variables from the edges of latent maps increases, with possibly new domain regions. Similarly, we find no distorted latent manifolds or discontinuous latent distributions. Finally, in this paper, though we have focused on the meaningful understanding of the domain phases and its boundaries through image segmentation in the latent maps, there are possible more key information that are left to be unearthed from the data through other visualization approach for potential scientific discoveries, which will be considered for future scope of study.

Zoom In Zoom Out Reset image size