Discriminative protein sequence modelling with Latent Space Diffusion

TL;DR

The paper addresses protein sequence representation by decoupling manifold learning from distributional modelling and introducing Latent Space Diffusion (LSD), a framework that couples a sequence autoencoder with a diffusion model on its latent space. It presents two autoencoder variants, LSD-TN and LSD-NM, to produce well-distributed latent embeddings and evaluates diffusion-based representations across multiple protein-property tasks, comparing to a masked language modelling baseline and to models like ESM2 and DPLM. The key finding is that diffusion representations trained on LSD variants outperform diffusion baselines based on MLM embeddings, but still fall short of token-based reconstructive methods in overall predictive performance; the results underscore the importance of latent-space design for effective diffusion. The work lays groundwork for future improvements in latent-space construction, potential multi-modal integration (e.g., structural data), and leveraging the one-parameter diffusion representations for pre-training or fine-tuning in protein modelling tasks, with implications for more flexible generative/discriminative protein modelling pipelines.

Abstract

We explore a framework for protein sequence representation learning that decomposes the task between manifold learning and distributional modelling. Specifically we present a Latent Space Diffusion architecture which combines a protein sequence autoencoder with a denoising diffusion model operating on its latent space. We obtain a one-parameter family of learned representations from the diffusion model, along with the autoencoder's latent representation. We propose and evaluate two autoencoder architectures: a homogeneous model forcing amino acids of the same type to be identically distributed in the latent space, and an inhomogeneous model employing a noise-based variant of masking. As a baseline we take a latent space learned by masked language modelling, and evaluate discriminative capability on a range of protein property prediction tasks. Our finding is twofold: the diffusion models trained on both our proposed variants display higher discriminative power than the one trained on the masked language model baseline, none of the diffusion representations achieve the performance of the masked language model embeddings themselves.

Figures (6)

• Figure 1: The LSD model is comprised of (a) a protein sequence autoencoder which learns a latent space $z$, and (b) a diffusion model acting on this latent space. The autoencoder is trained end-to-end by balancing a reconstruction loss, between input amino acid tokens and output token logits, against a normalization loss on the distribution of latent space embeddings. We consider two variants, LSD-TN with a non-trivial normalization loss and LSD-NM with a non-trivial reconstruction loss, as described in Sec. \ref{['sec:architecture']}. The diffusion model learns to map noised latent embeddings $z_t = \cos(\pi t/2) z + \sin(\pi t/2) \varepsilon$ to their orthogonal complements $v_t = -\sin(\pi t/2) z + \cos(\pi t/2) \varepsilon$, which thereby provides an additional one-parameter family of sequence representation to that obtained at the latent space.
• Figure 2: Evaluation of the $t$-dependence of the diffusion representation for the five protein property prediction tasks: (a) LSD-TN-M, (b) LSD-NM-M. The error bars are computed from the results of 5 randomly initialized predictors.
• Figure 3: UMAP projections of the latent space learned by LSD-NM-M Diffusion model. (a) Coloured by amino acid. (b) Coloured by relative mutability jones1992rapid. (c) Coloured by hydrophobicity nature. (d) Coloured by hydrophobicity index argos1982structural. (e) Coloured by average flexibility index bhaskaran1988positional. (f) Coloured by residue accessible surface area in folded protein chothia1976nature.
• Figure 4: Detailed architecture: (a) encoder, (b) decoder (c) diffusion model.
• Figure 5: Attention Map Analysis for LSD-NM-S diffusion model. a) Average attention logits per layer, aggregated over all heads and 128 protein sequences, each consisting of 100 amino acids. b) Distribution of attention scores across different types: Context attention, Local attention, and edge-token attention.