VAE Latent Landscape

The Variational Autoencoder (VAE) is the generative backbone of ALSEBO. It learns a smooth, continuous latent space from a protein MSA that can be sampled to produce novel, biologically plausible sequences [Ziegler2023].

Conceptual background

A standard autoencoder maps an input x (a one-hot encoded sequence) to a compressed representation z (latent code) and back. A VAE adds a probabilistic twist: the encoder outputs a distribution over z (a Gaussian parameterised by μ and σ) rather than a single point. During training the model is forced to keep that distribution close to a standard normal, which makes the latent space smooth and interpolatable.

Sequence (one-hot)
      │
      ▼
┌─────────────┐     ┌─────────────┐
│   Encoder   │────►│  μ,  σ      │  (latent parameters)
└─────────────┘     └──────┬──────┘
                           │  reparameterisation trick:  z = μ + ε·σ
                           ▼
                    ┌─────────────┐
                    │   Decoder   │────► Reconstructed sequence
                    └─────────────┘

The reparameterisation trick (z = μ + ε·σ, ε ~ 𝒩(0,1)) makes the sampling step differentiable so the whole network can be trained end-to-end with gradient descent.

Architecture

The ALSEBO VAE (alsebo.VAE.model) uses a fully-connected architecture:

Component	Details
Input	One-hot encoded sequence flattened to shape `(num_aa_types × sequence_length,)`; 23 amino acid categories (20 standard + selenocysteine, gap, pyrrolysine).
Encoder	Two dense hidden layers (2 × hidden_units → hidden_units) with BatchNormalization and configurable activation (default `relu`), followed by two linear heads for μ and σ.
Latent space	2-D by default (`dim_latent_vars=2`), making it directly visualisable as a 2-D map of sequence diversity.
Decoder	Mirror of the encoder; output reshaped to `(num_aa_types, sequence_length)` and passed through a Softmax to produce a per-position amino acid probability distribution.
Loss	Reconstruction loss (binary cross-entropy) + KL divergence between the encoder distribution and 𝒩(0, I).

Training

Train the VAE on an MSA FASTA file using the run_vae.py script:

python src/alsebo/VAE/run_vae.py \
    path/to/msa.fasta \
    path/to/save/model \
    path/to/logs/

Key hyperparameters in run_vae.py:

Parameter	Default	Effect
`dim_latent_vars`	2	Dimensionality of the latent space. 2-D is recommended for visualisation and latent featurization.
`num_hidden_units`	`seq_len × 3`	Capacity of the encoder/decoder. Increase for longer or more diverse MSAs.
`EPOCHS`	1000	Training iterations; `EarlyStopping(patience=10)` terminates early if loss plateaus.
`BATCH_SIZE`	16	Sequences per gradient update.

Generating the sequence space

After training, sample uniformly from the 2-D latent space and decode each point to obtain a pool of candidate sequences. The resulting FASTA file should have the latent coordinates (z0, z1) written into each sequence header so that the "latent" featurization method can read them directly:

>1.5991983967935868 -2.1042084168336674
MKVLILGAGFIGSELTARLHESTGDNVKVFCLVRDN...

This is the generated_seqs.fasta consumed by alsebo.seq_space.generate_seq_space().

Latent landscape as a functional map

The key insight from Ziegler et al. [Ziegler2023] is that the VAE latent space is not merely a compression — it organises sequences by phylogenetic relationships and functional properties. Regions of high Hamiltonian energy (as computed by DCA) correspond to fitness barriers, while low-energy basins correspond to functional clusters. ALSEBO exploits this structure by initialising the BO search from a diverse, landscape-covering set of sequences rather than a random one.

References

[Ziegler2023] (1,2)

Ziegler, C., Martin, J., Sinner, C., & Morcos, F. (2023). Latent generative landscapes as maps of functional diversity in protein sequence space. Nature Communications, 14, 2222. https://doi.org/10.1038/s41467-023-37958-z