Generative Embeddings of Brain Collective Dynamics Using Variational Autoencoders

We consider the problem of encoding pairwise correlations between coupled dynamical systems in a low-dimensional latent space based on few distinct observations. We use variational autoencoders (VAEs) to embed temporal correlations between coupled nonlinear oscillators that model brain states in the wake-sleep cycle into a two-dimensional manifold. Training a VAE with samples generated using two different parameter combinations results in an embedding that encodes the repertoire of collective dynamics, as well as the topology of the underlying connectivity network. We first follow this approach to infer the trajectory of brain states measured from wakefulness to deep sleep from the two end points of this trajectory; then, we show that the same architecture was capable of representing the pairwise correlations of generic Landau-Stuart oscillators coupled by complex network topology.

Figure 1

VAE architecture. The inputs are correlation matrices $C_{ij}$ obtained from the model [Eq. (1)] fitted to wakefulness and $N3$ sleep. The input layer has 8100 units, followed by an intermediate layer with 1028 neurons and a two-dimensional latent space. The next two layers reverse the encoding process, yielding a matrix $C_{ij}^{*}$ for each $z_1$, $z_2$ pair in the latent space. The bottom panel presents input matrices $C_{ij}$ (above diagonal) and their reconstructed versions $C_{ij}$ (below diagonal) for the model fitted to wakefulness and $N3$ sleep.

Figure 2

The latent space obtained from wakefulness and $N3$ sleep contains the orderly progression of intermediate brain states. (a) Latent space representation obtained by encoding the test set (wakefulness and $N3$ sleep) and the encoding obtained for the two intermediate states that were not used to train the VAE ($N1$ and $N2$ sleep). (b) Correlation matrices obtained by decoding an exhaustive exploration of the latent space variables $z_1$ and $z_2$. (c) Latent space divided into regions with maximal similarity to wakefulness, $N1$, $N2$, and $N3$ correlation matrices.

Figure 3

Latent space variables can be extrapolated to reproduce increased structure-function coupling as a signature of unconsciousness. (a) An exhaustive exploration of the SSIM between the decoded correlation matrices and the empirical structural connectivity matrix. High $z_1$ and low $z_2$ maximize this similarity. The red rectangle indicates the range of $z_1$ and $z_2$ reproduced in Fig. 2. (b) The empirical structural connectivity (left) and the best connectivity matrix reconstructed from the latent space (right) with the lower triangular part representing the matrices thresholded at 0.2.

Figure 4

Relationship between latent space variables ($z_1$, $z_2$) and the parameters of the homogeneous model ($a$, $G$).