Diffusion on language model encodings for protein sequence generation

## Abstract

Abstract Protein design requires a deep understanding of the inherent complexities of the protein universe.
While many efforts lean towards conditional generation or focus on specific families of proteins, the foundational task of unconditional generation remains underexplored and undervalued.
Here, we explore this pivotal domain, introducing DiMA, a model that leverages continuous diffusion on embeddings derived from the protein language model, ESM-2, to generate amino acid sequences.
DiMA surpasses leading solutions, including autoregressive transformer-based and discrete diffusion models, and we quantitatively illustrate the impact of the design choices that lead to its superior performance.
We extensively evaluate the quality, diversity, distribution similarity, and biological relevance of the generated sequences using multiple metrics across various modalities.
Our approach consistently produces novel, diverse protein sequences that accurately reflect the inherent structural and functional diversity of the protein space.
This work advances the field of protein design and sets the stage for conditional models by providing a robust framework for scalable and high-quality protein sequence generation.

## 1 Introduction

Generative modeling of proteins is gaining traction as a key area in academic research, potentially reshaping bioinformatics, synthetic biology, and protein-based therapeutics (Wu et al., ; Ovchinnikov & Huang, ).
A key part of this research area is the focus on the unconditional generation of protein sequences or 3D models.
Despite the increasing emphasis on conditional generation and family-specific fine-tuning (Madani et al., ; Sevgen et al., ), the foundational step of unconditional generation remains a challenging yet vital aspect.
The reason is simple: proficiency in unconditional generation provides a solid groundwork for more specialized and nuanced conditional generation, followed by subsequent fine-tuning.

Proteins can be represented in two distinct ways: via their linear amino acid sequence and their three-dimensional (3D) structure.
This duality offers both a sequential and a spatial perspective, and the interplay between these representations is fundamental in protein science. The amino acid sequence dictates the 3D arrangement, defining the unique folding pattern that shapes the protein’s overall form and function.

Recent protein generation methods, particularly diffusion models, show promising results (Watson et al., ; Lin & AlQuraishi, ; Lee et al., ; Wu et al., ; Alamdari et al., ). However, a critical limitation remains: many methods struggle to effectively generate proteins over the broad spectrum of the pan-protein space. Instead, these approaches often resort to workarounds such as devising pipelines to selectively pick satisfactory results or narrowing their scope to specific protein families, leaving a significant portion of the protein universe unexplored and constraining practical applications.

Figure: Figure 1: Overview of the proposed diffusion model on pLM embeddings for protein sequence generation. During training, we use ESM-2 to encode amino acids sequences to continuous representations. Then we train the diffusion model to reconstruct corrupted embeddings. During inference, we start from random Gaussian embedding and perform iterative refinement to generate a protein embedding. Then the decoder maps embedding to a sequence of amino acids.
Refer to caption: extracted/5452544/figures/pic1/pic1_v3_2_cropped.png

Current state-of-the-art models for generating 3D protein structures rely on scarce 3D protein data. This introduces biases and limitations due to the finite and non-representative nature of these data. The vast size difference (three orders of magnitude) between non-redundant datasets of protein 3D structures and amino acid sequences favors sequence data for building high-generalizing generative models. Additionally, many proteins, including important functional ones, like enzymes and transcription factors, carry segments lacking fixed structure, called intrinsically disordered regions (Dyson & Wright, ; Shukla et al., ). And since the structure of these proteins is not determined, they do not appear in structural databases. Consequently, deep learning models trained on 3D structures do not consider a vast class of intrinsically disordered proteins. Sequence representation of proteins naturally overcomes this issue, and the extensive pool of available protein sequence data broadens the potential coverage of the protein landscape.

In this paper, we propose DiMA, a diffusion model for generating protein sequences using protein language model (pLM) embeddings.
As shown in Figure 1, we employ ESM-2 pLM (Lin et al., ) as the encoder to obtain continuous representations of protein sequences, on which we train a denoising diffusion model. During inference, iterative refinement is performed, and the resulting embedding is decoded to amino acid sequence.
We extensively evaluate the quality and diversity of generation, the ability of the model to capture the distribution of training data, and the functional features of natural proteins. We use multiple evaluation metrics across three modalities: protein sequences, 3D-structures, and pLM features.

The main contributions of our work can be summarized as follows:

- •
We introduce DiMA, a diffusion-based generative model for protein sequence design. DiMA uses a continuous diffusion approach through the embeddings of a transformer-based protein language model.
- •
We demonstrate that DiMA outperforms other approaches for generating amino acid sequences in terms of the quality and diversity of the generated sequences. We also show that our model accurately reflects various characteristics of training data. This establishes that continuous diffusion using pLM latent representations is a viable approach for text generation in the domain of protein sequences.
- •
Through a thorough ablation study, we reveal the impact of our architectural design choices and implemented techniques for training and sampling.

The code is available at [https://anonymous.4open.science/r/DiMA-2FBF](https://anonymous.4open.science/r/DiMA-2FBF).

## 2 Related Work

Diffusion generative models, introduced by Sohl-Dickstein et al. (), have gained attention for their remarkable results in image (Ho et al., ; Song et al., , ), and speech generation (Chen et al., ; Popov et al., ).
Due to their impressive generative quality, some studies have extended the application of diffusion models to the text domain. Hoogeboom et al. () and Austin et al. () proposed multinomial diffusion for discrete data corruption.
Subsequently, other works (Li et al., ; Lin et al., ; Gulrajani & Hashimoto, ; Han et al., ; Strudel et al., ; Gao et al., ) adapted Gaussian diffusion to sequence learning by embedding discrete data into continuous space.
Yuan et al. () extended the text diffusion model to the sequence-to-sequence setting.
Ye et al. () conducted a study on the discrepancy of the text embedding space, demonstrating that the diffusion task at small noise scales is trivial.
Zhang et al. () implemented latent text diffusion inside a VAE with an autoregressive decoder.
Lovelace et al. () utilized diffusion models to generate a fixed-length latent representation, mapped into a high-dimensional space with the reconstruction network before being fed into an autoregressive decoder to generate text.

In the domain of protein science, deep learning has emerged as a transformative tool.
s, pre-trained on extensive protein sequence datasets, provide representations widely employed in various tasks (Elnaggar et al., ; Lin et al., ).
Generative models for protein sequences, exemplified by recent advancements, enhance predictions of proteins with improved properties and functions (Wu et al., ; Ovchinnikov & Huang, ).
Simultaneously, progress in the sequence-to-structure domain, as seen in models like AlphaFold (Jumper et al., ) and ESMFold (Lin et al., ), enables the prediction of 3D protein conformation from amino acid sequences.
Models such as ProteinMPNN (Dauparas et al., ) or ESM-IF1 (Hsu et al., ) predict an amino acid sequence given a specific 3D structure, effectively reverse engineering the process.
These achievements establish a link between linear amino acid arrangements and spatial structures, advancing our understanding of protein folding and dynamics. Bidirectional conversions, from sequence to structure and vice versa, offer a unique validation opportunity. Aligning outputs from generative, sequence-to-structure, or inverse folding models ensures coherent assessments of robustness, accuracy, and the quality of generated proteins.
This approach has become the standard in assessing the quality of generative models for protein sequences (Alamdari et al., ; Madani et al., ) and 3D-structures (Lin & AlQuraishi, ; Madani et al., ; Watson et al., ).

## 3 Continuous diffusion on LM representations of protein sequences

### 3.1 Overview

Figure 1 outlines our framework of latent diffusion on the pLM representations.
The proposed method comprises three parts.
The first part is a pre-trained single-sequence encoder ($\mathcal{E}$) that learns a meaningful latent space corresponding to the original protein space.
The second part is a diffusion model ($\mathcal{F}$) that generates vectors of protein latent space from a Gaussian noise.
The third part is a decoder ($\mathcal{D}$) that maps generated latent into the sequence of amino acids.

### 3.2 Training phase

We use a pre-trained transformer-based pLM (ESM-2 (Lin et al., )) as an encoder.
It is a lightweight encoder with approximately 8M parameters and an embedding size of 320.
ESM-2 was trained on a masked language modeling objective on 65 million amino acid sequences and showed high results in a range of structure prediction tasks.

Encoder maps the sequence of amino acids $y=[y_{1},...,y_{s}]$ of length $s$ to the latent vectors $x=[x_{1},...,x_{s}]\in R^{s\times d}$, $x=\mathcal{E}(y)$, $d=320$.
Then, we employ dimension normalization to encourage each component of a single vector in the sequence $x$ to have zero mean and unit variance $z_{0}=Normalize(x)$.
This transformation allows us to adapt the discrete protein input to a standard Gaussian diffusion framework.

Then we train a continuous denoising network, $\hat{z}_{\theta}()$, to recover $z_{0}$ from $z_{t}=\sqrt{\alpha_{t}}z_{0}+\sqrt{1-\alpha_{t}}\varepsilon$ using the standard objective:

$$ $\mathcal{L}(\theta)=\mathop{\mathbb{E}}_{\varepsilon\sim\mathcal{N}(0,\mathbf{ I}),t\sim U[0;1]}||z_{0}-\hat{z}_{\theta}(z_{t},t)||^{2}$ $$

### 3.3 Noise schedule

Figure: Figure 2: Left: the diffusion reconstruction loss of $z_{0}$ from $z_{t}$ with different noise schedules: $||z_{0}-\hat{z}_{\theta}(z_{t},t)||^{2}$. Right: $\alpha_{t}$ for different noise schedules.
Refer to caption: x1.png

We have found that while the linear noise scheduler used by Song et al. () andHo et al. (), and the cosine scheduler proposed by Nichol & Dhariwal () work well for the image domain, they are sub-optimal for the protein domain.
The reconstruction loss curves of diffusion models trained with such schedulers have large areas with small magnitudes and close to the minimum, as shown in Figure 2 (left).
We conjecture that this happens due to the high discreteness of the protein latent space.
Consequently, the reconstruction of $z_{0}$ from $z_{t}$ becomes quite trivial for the model for a long period of time, leading to inefficient training.
We employ the noise schedule proposed by Simple Diffusion (Hoogeboom et al., ) (SD):

$$ $\alpha_{t}=\frac{1}{1+d^{2}\tan^{2}(\frac{\pi t}{2})}$ $$

where $d$ is a hyperparameter that reflects the rate of the schedule.
The larger the value of $d$, the greater the data corruption rate.
Figure 2 (right) depicts classical and SD-$10$ ($d=10$) schedules in terms of $\alpha_{t}$.
We use the SD-$10$ schedule in our training setup.

### 3.4 Self-conditioning

We follow recent advances in sequence generation and apply the self-conditioning technique (Chen et al., ) in the implementation of our method.
Typically, the denoising network predicts $\hat{z}_{0}$ using the latent variable $z_{t}$ and timestep $t$ as an input.
Self-conditioning additionally proposes to utilize predicted $\hat{z}_{0,s}$ from the previous timestep $s$ for estimation $\hat{z}_{0,t}=\hat{z}_{\theta}(z_{t},t,\hat{z}_{0,s})$, $t<s$.
During iterative sampling at inference, we have already computed the prediction $\hat{z}_{0,s}$ from the previous timestep.
Consequently, there are no additional model launches at inference.
However, we need to modify the training process so that the diffusion model trains to exploit the additional input $\hat{z}_{0,s}$.

Just like during a standard training iteration, we sample the timestep $t\sim U[0;1]$.
Initially, in half of the cases, we provide no additional input to the model, setting $\hat{z}_{0,t}=\emptyset$. In the remaining cases, we estimate $\hat{z}_{0,t}=\hat{z}_{\theta}(z_{t},t,\emptyset)$, where $\emptyset$ is a zero vector in our implementation.
Subsequently, we compute $\hat{z}_{0}$ using the previously predicted estimation $\hat{z}_{0}=\hat{z}_{\theta}(z_{t},t,\text{SG}[\hat{z}_{0,t}])$, where $\text{SG}[\cdot]$ denotes the stop-gradient operation.
The last estimation is used to compute the loss:

$$ $\mathcal{L}(\theta)=\mathop{\mathbb{E}}_{\varepsilon\sim\mathcal{N}(0,\mathbf{ I}),t\sim U[0;1]}\left[||z_{0}-\hat{z}_{\theta}(z_{t},t,\text{SG}[\hat{z}_{0,t }])||^{2}\right]$ $$

This training procedure also allows sampling with zero self-condition, which is used at the first iteration of generation.
We implement additional conditioning on the previous estimate by using linear projection and summation with input at each transformer block as depicted in Figure 7.

### 3.5 Decoder

The proposed architecture allows us to use the decoder of ESM-2 pre-trained simultaneously with the encoder on masked language modeling objectives.
However, we found that additional training of the decoder results in a more precise reconstruction of amino acid sequences from the latents $x$ during inference.
The decoder architecture comprises two linear layers and an activation function.

### 3.6 Length sampling

A crucial aspect of the inference phase involves determining the length of the generated sequence.
While in training, the sequence length is straightforwardly determined by the corrupted latent, during inference, we choose to sample the length from the empirical distribution observed in the training dataset.
Length sampling is a crucial aspect of our model, as its absence leads to an inadequate length distribution (see Appendix E for more details).
We employ an attention mask to incorporate length information into the denoising network.
The generation process begins by sampling a random Gaussian embedding and length.
Using a fixed number of steps $T$, we iteratively generate the final $\hat{z}_{0}$.
Subsequently, we denormalize each embedding and utilize a decoder to map the sequence of vectors to the sequence of amino acids.

### 3.7 Model architecture

We use the 12-layer Transformer model with 16 attention heads and a hidden size of 320 as a backbone for our diffusion model. We modify the model to ensure the effective operation of denoising diffusion within the specific context of protein-related data (see Appendix C for more details).
One noteworthy modification involves incorporating the time embedding into each transformer block. To achieve this, we use a linear projection before summation prior to each transformer block. Our experiments consistently demonstrate the effectiveness of this approach for time conditioning (refer to Section 4.1 and Table 1).
Another modification involves incorporating long skip connections into the transformer model, inspired by the U-Vit architecture (Bao et al., ).

## 4 Experiments

We carry out experiments on two datasets, SwissProt and AFDBv4-90 (Section A), and compare our approach against a set of generative models operating in the amino acid sequence space (Section 4.2). Specifically, we compare our diffusion models with the AR transformer (NanoGPT), AR CNN (SeqDesign), GAN (ProteinGAN), and the categorical diffusion model EvoDiff. We train the baseline models from scratch on the same datasets.
We conduct a comprehensive ablation study to assess the impact of the architectural features and sampling techniques employed (Section 4.1).
We evaluate the quality and diversity of generated sequences as well as the ability of the models to capture the distribution of the training data (Section B). To assess quality generation, we employ sequence- and 3D structure-based metrics: pLDDT, pseudo-perplexity, scPerplexity, TM-score, and BLAST identity score. To examine the similarity between the distributions of the generated sequences and the datasets, we use Fréchet ProtT5 distance (FPD), maximum mean discrepancy (MMD), and 1-Wasserstein optimal transport (OT). To assess the diversity of generation, we measure average distances between sequences in batches and use joint clustering of training and generated data. Additionally, we employ InterPRO for functional annotation to compare the semantic distribution of sequences (Section 4.5).

**Table 1: Performance comparison between DiMA and baseline architectures of the same parameter count trained on SwissProt and AFDBv4-90 datasets.**
|  | Model | pLDDT ($\uparrow$) | ESM-2 pppl ($\downarrow$) | scPerplexity ($\downarrow$) | TM-score ($\uparrow$) | BLAST ($\uparrow$) | FPD ($\downarrow$) | MMD ($\downarrow$) | OT ($\downarrow$) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SwissProt | Dataset | 80.7 | 5.35 | 1.88 | 0.80 | 100 | 0.13 | 0.00 | 1.08 |
| Random sequences | 25.0 | 21.54 | 2.77 | 0.33 | 0 | 3.97 | 0.20 | 3.88 |  |
| nanoGPT | 61.0 | 8.18 | 2.04 | 0.63 | 43 | 1.24 | 0.03 | 2.53 |  |
| EvoDiff-OADM | 37.1 | 15.77 | 2.44 | 0.42 | 12 | 1.49 | 0.11 | 2.63 |  |
| SeqDesign | 43.1 | 11.89 | 2.35 | 0.41 | 17 | 3.53 | 0.19 | 5.12 |  |
| proteinGAN | 30.4 | 16.48 | 2.57 | 0.33 | 0 | 2.94 | 0.17 | 3.98 |  |
| DiMA | 80.8 | 5.20 | 1.80 | 0.85 | 68 | 0.41 | 0.01 | 1.41 |  |
| w/o skip connections | 77.3 | 5.84 | 1.87 | 0.82 | 61 | 0.48 | 0.02 | 1.51 |  |
| w/o time layers | 79.4 | 5.49 | 1.83 | 0.85 | 66 | 0.44 | 0.02 | 1.44 |  |
| w/o ESM encoder | 62.7 | 9.22 | 2.09 | 0.71 | 48 | 1.05 | 0.04 | 2.14 |  |
| w/o self-conditioning | 68.2 | 9.18 | 2.08 | 0.74 | 46 | 0.54 | 0.04 | 1.61 |  |
|  | w linear schedule | 77.0 | 6.29 | 1.89 | 0.82 | 58 | 0.50 | 0.02 | 1.51 |
|  | w cosine schedule | 54.1 | 10.86 | 2.16 | 0.60 | 34 | 0.97 | 0.06 | 2.02 |
| AFDB | Dataset | 83.9 | 5.79 | 1.75 | 0.92 | 100 | 0.18 | 0.00 | 1.57 |
| Random sequences | 26.2 | 21.67 | 2.75 | 0.35 | 0 | 3.02 | 0.18 | 4.15 |  |
| nanoGPT | 68.5 | 8.21 | 1.94 | 0.77 | 40 | 0.62 | 0.02 | 1.99 |  |
| DiMA | 73.9 | 8.50 | 1.90 | 0.85 | 48 | 0.69 | 0.03 | 1.86 |  |
| w/o self-conditioning | 56.3 | 12.08 | 2.18 | 0.69 | 31 | 0.96 | 0.05 | 2.29 |  |

### 4.1 Ablation study

To assess the contribution of the proposed design choices to the performance of DiMA, we train four models from scratch with the following modifications:

- •
Removing the long skip-connections between the shallow and the deep transformer blocks.
- •
Using time conditioning through admixing the time embeddings to the corrupted latent vectors of amino acids instead of employing a dedicated time layer before each transformer block.
- •
Omitting the transformer encoder (ESM-2), retaining only its initial projection layer.
- •
Training the model without self-conditioning.
- •
Using linear noise schedule.
- •
Using cosine noise schedule.

To evaluate the impact of each component, we measure quality and diversity metrics on 2048 samples of generated protein sequences for each ablated model.

Table 1 demonstrates that each proposed feature contributes significantly to the model’s performance individually. The most substantial decrease in both the quality and diversity of the generated sequences occurs in the ablated models without the ESM-2 encoder and when trained without self-conditioning. Disabling these features eliminates DiMA’s advantage over nanoGPT, particularly affecting the diversity of the generated proteins.
Removing skip-connections and time layers results in a less pronounced but still significant decrease in overall performance, as reflected in all metrics.

To ablate the impact of the SD-$10$ noise schedule, we train our diffusion model with standard linear and cosine schedules, leaving other parameters intact.
We find that SD-$10$ significantly outperforms the cosine schedule in both quality and diversity.
It also achieves less expressed but better results than the linear schedule.

### 4.2 Baseline models

We evaluate DiMA against various architectures for sequence generation. For each of the methods described below, we train from scratch a model with the same parameter count (33M) as DiMA on the same dataset(s). All the models use a standard character-level tokenization scheme, where each individual amino acid in the sequence corresponds to a single token (Lin et al., ; Madani et al., ). In this work, we consider only methods for sequence generation with published source code.

- •
EvoDiff-OADM (Alamdari et al., ), a recently developed order-agnostic autoregressive discrete diffusion method in which a forward process iteratively corrupts a protein sequence by converting single amino acid to a special mask token. In the reverse process, the trained dilated convolutional neural network (CNN) restores a fully masked sequence to generate a protein sequence.
- •
nanoGPT (Karpathy, ), a lean implementation of GPT-2 autoregressive language model.
- •
ProteinGAN (Repecka et al., ), a variant of the generative adversarial network in which both the discriminator and generator are CNNs based on ResNet blocks augmented with a self-attention layer. In the original study, ProteinGAN was trained on a dataset consisting of similar sequences from a single protein family. This approach focused on generalizing within the space of that particular protein family. Despite using diverse datasets in this work, we include ProteinGAN in the comparison list.
- •
SeqDesign (Shin et al., ), a residual causal dilated CNN that is shown to have strong generalization capabilities over protein sequence space.

We evaluate the performance of each model using batches of 2048 generated sequences. Additionally, to assess the features of the datasets and establish the lower expected quality of the generated sequences, we use 2048 sequences sampled from the datasets and randomly generated with the length distributions observed in each dataset, respectively.

### 4.3 Evaluation metrics

To address the challenging task of evaluating generative models for amino acid sequences, we deploy a diverse set of metrics, focusing on quality, distributional similarity, and diversity.
To assess sequence quality, we use established protein language and structure models, ensuring plausibility and coherence across different modalities (Table 2).
To measure distributional similarity, we compare generated batches with training data, affirming the model’s fidelity to the underlying data distribution.
To evaluate the diversity, we analyze distances within batches and perform co-clustering of generated and training data.
Finally, to assess the semantic relevance, we utilize protein annotation tools that provide insights into functional regions within the generated sequences.

Quality of generated sequences.
One of the foundational principles in protein science is that a protein’s structure is governed by its amino acid sequence. Therefore, ensuring foldability and structural plausibility
is crucial in generating protein sequences. We use ESMfold to predict 3D structures and pLDDT scores to assess confidence in amino acid positions. The average pLDDT score across all amino acids (Table 1) reflects overall confidence in the predicted protein structure.

Pseudoperplexity (Salazar et al., ) (ESM-2 pppl) reflects how well a given sequence aligns with the patterns the assessing model has learned from its training data. We calculate ESM-2 pppl using ESM-2 650M (Lin et al., ) pLM by masking each token (amino acid) in the sequence and, subsequently, predicting it considering all other tokens.

**Table 2: Review of the metrics across modalities for evaluating generation quality, diversity, and distribution similarilty**
|  | Quality | Diversity | Distribution |  |
| --- | --- | --- | --- | --- |
|  |  |  |  | similarity |
| Modality | Sequence | Structure | Sequence | ProtT5 features |
| Metrics | ESM-2 pppl | pLDDT | SeqD | FPD |
| BLAST | TM-score | Clustering | OT |  |
| scPerplexity | scPerplexity |  | MMD |  |

We use scPerplexity, introduced in (Alamdari et al., ), to estimate the quality and reliability of amino acid sequences.
For each sequence, we first predict its structure using ESMFold.
Then, we employ ProteinMPNN (Dauparas et al., ), which takes the bare backbone of the predicted structure with masked names of the amino acid residues and predicts a sequence that would naturally fold into that shape.
This process differs from ESM-2 pppl, which focuses on direct sequence comparison. Instead, with scPerplexity, we are checking how well the structure predicted by ESMFold aligns with a sequence generated by ProteinMPNN, providing a measure of the structural validity of our sequences.

To assess the structural relevance of generated sequences, we use the TM-score (Zhang & Skolnick, ), a widely recognized metric for evaluating structural similarity between protein pairs.
To compute TM-scores for each batch of sequences, we acquire their 3D structures through ESMFold. Subsequently, for each of these structures, we identify the closest natural protein in the SwissProt and AFDBv4-90 datasets from the AlphaFold Database (Tunyasuvunakool et al., ) using FoldSeek (Van Kempen et al., ).
TM-score does not depend on protein size, consistently ranging between 0 and 1. A TM-score above 0.5 indicates a similar fold in structure and values above 0.9 signal on nearly identical structures.

Another way to assess the model’s ability to generate natural-like sequences is the BLAST identity measure. BLAST (Altschul et al., ) is a standard bioinformatics tool for efficiently searching for similar sequences. This metric can also serve as an indicator of model quality (Repecka et al., ).

Distribution learning.
The desirable property of a generative model is to accurately capture a given data distribution. For this purpose, we employ three metrics: Fréchet distance, maximum mean discrepancy, and 1-Wasserstein distance. These metrics measure the distance between the distributions of ProtT5 (Elnaggar et al., ) feature representations for real and generated sequences, and they are zero when the training and generated distributions align perfectly and increase as the dissimilarity between them grows.
The features are obtained by embedding the sequences using the ProtT5 pLM, with global average pooling applied along the length dimension to ensure a fixed-size embedding. This approach has demonstrated robustness and efficacy in addressing a wide range of bioinformatic tasks (Elnaggar et al., ; Hou et al., ).

Inspired by the well-established Fréchet Inception distance (FID) (Heusel et al., ) used in image generation and applications of a similar metric to protein sequence data (Alamdari et al., ), we apply this approach to measure the extent to which the models have learned the training distribution. We use the embeddings from the ProtT5, hence the term Fréchet ProtT5 Distance (FPD).

Maximum mean discrepancy (MMD) is a kernel-based statistical test to determine whether two samples are drawn from different distributions, as proposed in (Gretton et al., ). This is achieved by measuring the difference in mean values resulting from applying a smooth kernel function to each sample. In practical applications, an empirical estimate of MMD is employed (Joshua Southern & Correia, ). We evaluate the MMD between batches of sequences sampled from the dataset and generated by the respective models, following the methodology proposed for 3D structures in (Joshua Southern & Correia, ).

To evaluate the similarity between batches of generated and natural sequences, we employ 1-Wasserstein optimal transport (OT). To establish matching pairs of sequences, we use the Earth Mover Distance (EMD) solver. The average distance between the optimal pairs 1 gives a measure of both the diversity and proximity of generated samples to the dataset.

Figure: Figure 3: Distributions of distances between optimal pairs of generated sequences and those sampled from the dataset.
Refer to caption: extracted/5452544/figures/both_ot_dists_hists.png

**Table 3: Diversity evaluation of DiMA and baseline architectures, trained on SwissProt and AFDBv4-90 datasets.**
|  | Model | SeqD ($\downarrow$) | Num. PS ($\uparrow$) | Frac. PS ($\uparrow$) | Num. PC ($\uparrow$) |
| --- | --- | --- | --- | --- | --- |
| SwissProt | Dataset | 0.00 | 1835 | 0.89 | 1228 |
| Random sequences | 1.44 | 0 | 0.00 | 0 |  |
| nanoGPT | 0.50 | 940 | 0.46 | 398 |  |
| EvoDiff-seq | 0.91 | 223 | 0.11 | 125 |  |
| SeqDesign | 0.26 | 44 | 0.02 | 23 |  |
| proteinGAN | 0.66 | 0 | 0.00 | 0 |  |
| DiMA | 0.72 | 1803 | 0.88 | 412 |  |
| w/o skip connections | 0.83 | 1413 | 0.69 | 322 |  |
| w/o time layers | 0.59 | 1504 | 0.73 | 260 |  |
| w/o ESM encoder | 2.52 | 1035 | 0.51 | 121 |  |
| w/o self-conditioning | 0.14 | 726 | 0.35 | 311 |  |
|  | w linear schedule | 0.74 | 1680 | 0.82 | 333 |
|  | w cosine schedule | 0.66 | 755 | 0.36 | 200 |
| AFDB | Dataset | 0.00 | 1720 | 0.83 | 1386 |
| Random sequences | 2.06 | 0 | 0.00 | 0 |  |
| nanoGPT | 0.88 | 1117 | 0.55 | 463 |  |
| DiMA | 0.93 | 1319 | 0.64 | 233 |  |
| w/o self-conditioning | 0.36 | 377 | 0.18 | 131 |  |

Diversity of Generated Sequences.
To assess the diversity of sequences, we used the metric dubbed Sequence Diversity (SeqD) (Wang et al., ). It reflects how distant the sequences are from each other in terms of Levenshtein distance. The desired SeqD value for a generative model is expected to be close to the training data. Therefore, we measure the absolute difference between the calculated SeqD values for each model and the dataset.

Sequence clustering
provides insights into the model’s generating ability, specifically, the inner diversity and similarity to the dataset. The desired behavior is to generate sequences similar to the dataset yet distinctly different from each other.
We use MMseqs2 (Steinegger & Söding, ) for clustering the generated sequences with the dataset to establish how they align with each other. The generated sequences that reside in clusters with the sequences from the dataset are positive sequences (PS) and positive clusters (PC). The number of PS determines the degree of resemblance between the generated sequences and the dataset, while the diversity can be assessed through PC.

### 4.4 Results

The evaluation of the quality of sequence generation on the SwissProt dataset (Table  1) shows that DiMA surpasses all baselines and displays metric values closely aligned with the dataset. NanoGPT, an autoregressive language model, achieves strong results but falls short of matching the metric levels of the dataset.
In comparison, other baselines exhibit notably poorer performance. SeqDesign and ProteinGAN, initially designed for narrow classes of proteins, may not be suitable for training on diverse datasets. EvoDiff, while outperforming SeqDesign and ProteinGAN, still demonstrates metric values closer to a random sample than to the dataset, consistent with observations in the original EvoDiff paper (Table S3 of (Alamdari et al., )).

On the larger and more diverse AFDBv4-90 dataset, the performance gap between DiMA and nanoGPT narrows. DiMA achieves higher values for pLDDT, TM-score, and BLAST, while nanoGPT shows better results in ESM-2 pseudoperplexity (8.21 against 8.50 for DiMA). Both models are close in terms of scPerplexity, with 1.94 for nanoGPT and 1.90 for DiMA. Despite these achievements, both models fall short of reaching the metric values of the dataset.

In distributional similarity between the generated sequences and the training data (refer to Table 1), DiMA outperforms all baselines when trained on the SwissProt dataset. In the case of the AFDBv4-90 dataset, both DiMA and nanoGPT demonstrate robust results, with minimal discrepancy between the dataset and the generated samples. However, nanoGPT slightly outperforms DiMA in terms of FPD and MMD. Despite this advantage, nanoGPT demonstrates a distribution of distances between optimal pairs that is less similar to the dataset distribution, as illustrated in Figure 3.

The evaluation of diversity (Table  3) reveals a tradeoff between quality and diversity. DiMA, when trained without self-conditioning, shows exceptional diversity on both datasets, as measured by the SeqD metric. However, the self-conditioned version, while achieving higher quality, generates slightly less diverse results.
All baselines exhibit SeqD values close to zero, indicating the absence of mode collapse.

Assessing the cluster structure of generated sequences, DiMA outperforms all baselines on SwissProt, with 88% of generated sequences having at least 30% similarity to the dataset sequences. This suggests that DiMA consistently produces sensible results resembling the training data. NanoGPT demonstrates a lower fraction of similar sequences (46%). Both DiMA and nanoGPT exhibit comparable numbers of clusters, 412 and 398, respectively, suggesting that while nanoGPT generates sequences resembling the dataset less often, those instances span the same area of protein space.
On AFDBv4-90, the difference between DiMA and nanoGPT is less pronounced. Both models generate fewer clusterable sequences than the dataset (83%), with DiMA at 64% and nanoGPT at 55%. The number of clusters differs, with 233 for DiMA and 463 for nanoGPT.

For both datasets, the quality of sequences (measured by pLDDT) generated by DiMA that did not cluster with the dataset (“negative” sequences) is significantly higher than the corresponding “negative” sequences generated by other models (refer to the Sequence clustering section in Appendix B). This suggests that the model is capable of generating plausible sequences that lie outside the distribution of the training data.

The observed disparity in performance on the two datasets may stem from their inherent structural characteristics. SwissProt, being smaller and less diverse, contains multiple similar proteins from different families, while AFDBv4-90 represents a subset of UniProt50 designed to maintain mutual similarity below 50%, resulting in a larger and more diverse dataset with 2.2 million entries.

### 4.5 Biological relevance

To evaluate the biological relevance of the generated sequences, we utilized the established protein annotation tool InterProScan (Paysan-Lafosse et al., ; Jones et al., ). InterProScan includes a set of pre-trained models based on hidden Markov models (HMMs), which allow for assigning potential folds and functions. This analysis involves annotating the generated protein sequences using the SUPERFAMILY HMM library (Oates et al., ), which provides sequence homology to SCOP structural domains (Murzin et al., ). For DiMA and DiMA without self-conditioning (DiMA_no_SC), 92% and 73% of the sequences are annotated as a SUPERFAMILY member (Figure 11A), which is close to 80% for the batch from the SwissProt dataset. In contrast, fewer than 60% of the sequences generated by nanoGPT show homology to SCOP structural domains, while the output of other models exhibits even lower homology. Additionally, DiMA and DiMA_no_SC reproduce the distribution of domain lengths from the natural set of proteins from SwissProt (Figures 9-10). Both DiMA and DiMA_no_SC can generate biologically relevant protein sequences. However, DiMA tends to produce proteins with higher similarity to known sequences, while DiMA_no_SC generates sequences with both lower homology and lower quality (Table  1).
Representative examples of ESMFold predictions of the 3D structures of generated sequences by DiMA and DiMA_no_SC are shown in Figure 4B.

Figure: Figure 4: Biological relevance of the generated sequences. A) Sequence annotation into known structural domains. Left: The fraction of protein sequences annotated per model. Right: The identity % with the closest hit from SwissProt. B) ESMFold predicted representative examples of proteins generated by DiMA (top) and the closest hit SwissProt (bottom) with UniProt IDs and the homology %. C) Left: MobiDBLite predictions of IDRs. Right: A representative example of a protein generated by DiMA, highlighting the predicted IDR.
Refer to caption: extracted/5452544/figures/seqdif_biol_cropped.png

Natural proteins encompass both structured regions and IDRs that lack regular structure but still play functional roles (Uversky, ) (Figure 12). To annotate these regions, we employ the MobiDB model within the InterProScan tool, which predicts IDRs in protein sequences using multiple classifiers (Piovesan et al., ). Sequences generated by DiMA exhibit a natural-like profile of IDR length distribution (Figures 4 and 12). Generation of both folded and unfolded structural regions provides a distinct advantage for sequence diffusion models over models exclusively trained on folded protein domains.

Finally, we calculate the frequency of secondary structure elements within the folded regions using the DSSP tool (Kabsch & Sander, ) against protein structures predicted via ESMFold. DiMA mirrors the amount of secondary structural elements of natural proteins, outperforming other models (Figure 13).

Overall, this data suggests that DiMA models are capable of generating diverse variants of natural-like proteins.

## 5 Conclusion

In this paper, we introduce DiMA, a continuous diffusion-based model for protein sequence generation that operates in the space of language model latent representations. Through extensive experiments, we evaluate the quality, diversity, distribution similarity, and biological relevance of the generated sequences. The results demonstrate that DiMA outperforms other methods for text generation. Additionally, a comprehensive ablation study quantitatively verifies the impact of DiMA’s architectural features and design choices on its performance.

## 6 Impact statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

## Appendix A Datasets

SwissProt is a dataset that contains a high-quality, manually annotated subset of the UniProt database. This dataset is small enough and good enough for proof-of-the-concept studies. After filtering out all sequences shorter than 128 and trimming all sequences longer than 254, we ended up with  470k sequences. MMseqs2 clustering of this dataset (>50% sequence identity and >80% sequence overlap) reveals the presence of clusters of similar sequences with the maximum number of sequences in a cluster equal to 1570. Each of those clusters comprises sequences that belong to a single protein family. For example, the most populous cluster is 1570 protein sequences of cytochrome b of different species, a very abundant protein involved in electron transport in eukaryotic cells. Around 120k sequences do not form clusters under the conditions used.

Another dataset that we used is AFDBv4-90 from Durairaj et al. (), which is a subset of the UniRef50 database. The sequences in this dataset obey two conditions: 1. The sequence identity between all members is no more than 50%, and 2. The average predicted pLDDT by AlphaFold is no less than 90. After filtering out all sequences shorter than 128 and longer than 254, we ended up with  2.2 million whole sequences of highly diverse proteins of high quality.

## Appendix B Metrics

pLDDT.
To assess the foldability of our generated sequences, we utilize ESMfold to predict the three-dimensional structure of the given protein sequence.
For each amino acid within the predicted structure, ESMfold provides a pLDDT score, which represents the confidence of the model in the predicted positions of amino acids in the 3D structure.
We took an average of these pLDDT scores for all amino acids in the sequence to gauge the overall confidence in the predicted protein structure.

ESM-2 pseudoperplexity.
To assess how probable the original sequence is under the model’s distribution, we used pseudoperplexity (Salazar et al., ) using ESM-2 650M encoder transformer-based language model (Lin et al., ). Each token (amino acid) in the sequence was masked and then predicted, considering all other tokens in the sequence.
The final pseudoperplexity value is aggregated using the following equation:

$$ $\mathcal{P}_{ESM-2}(S)=\exp\left\{-\frac{1}{|S|}\sum_{i=1}^{|S|}\log p(s_{i}|S _{\backslash i},\Theta_{ESM-2})\right\}$ $$

Here, $\mathcal{P}_{ESM-2}(S)$ represents the pseudoperplexity of sequence $S$, $|S|$ denotes the length of sequence $S$, $s_{i}$ is the $i$-th token in the sequence, $S_{\backslash i}$ represents the sequence without the $i$-th token, and $\Theta_{ESM-2}$ denotes the parameters of the ESM-2 model.

TM-score. To evaluate the structural relevance of the generated sequences, we turned to the TM-score (Zhang & Skolnick, ), a widely recognized metric for evaluating structural similarity between protein pairs.
The TM-score measures the similarity between two protein structures and helps distinguish proteins with a similar fold from those with different folds.
Unlike many other metrics for 3D-alignment, it does not depend on protein size and always ranges between 0 and 1, where a TM-score above 0.5 indicates a similar fold in structure.
The TM-score is given by:

$$ $\text{TM-score}=\frac{1}{L_{\text{target}}}\sum_{i=1}^{L_{\text{query}}}\frac{ 1}{1+\left(\frac{d_{i}}{d_{0}(L_{\text{target}})}\right)^{2}}$ $$

Here, $L_{\text{target}}$ is the length of the target protein, $L_{\text{query}}$ is the number of aligned residues between the two proteins, $d_{i}$ is the distance between the $i$-th aligned residue pairs, and $d_{0}$ is a scaling factor to normalize the length difference between query and target proteins.
To calculate TM-scores for each sample of generated sequences, we first obtained their 3D structures using ESMFold.
For each of these structures, we have found the closest natural protein in the SwissProt and AFDBv4-90 datasets from the AlphaFold Database (Tunyasuvunakool et al., ) using the FoldSeek (Van Kempen et al., ).

BLAST Identity.
For each sequence, we ran BLAST with specific parameters (e-value = 0.05 and BLOSUM62 substitution matrix) to identify similar sequences within the training dataset. The number of matching amino acids between the generated sequence and the most similar sequence found in training data was normalized by sequence length and multiplied by 100 to obtain percentages. The BLAST identity metric is the average over a batch of 2048 sequences.

Fréchet ProtT5 Distance (FPD).
The Fréchet distance, also known as the 2-Wasserstein distance, quantifies the dissimilarity between two samples drawn from multivariate Gaussian distributions, denoted as $X_{1}\sim\mathcal{N}(\mu_{1},\Sigma_{1})$ and $X_{2}\sim\mathcal{N}(\mu_{2},\Sigma_{2})$, and is defined as follows:

$$ $d(X_{1},X_{2})^{2}=||\mu_{1}-\mu_{2}||^{2}+\text{Tr}(\Sigma_{1}+\Sigma_{2}-2 \sqrt{\Sigma_{1}\Sigma_{2}})$ (1) $$

Maximum mean discrepancy (MMD).
The idea behind MMD involves assessing the distance between two samples by measuring the difference in mean values resulting from applying a smooth function to the samples. A biased empirical estimate of MMD between two samples $X=\{x_{1},...,x_{n}\}$ and $Y=\{y_{1},...,y_{n}\}$ using kernel $k$ is defined as follows:

$$ $\displaystyle MMD_{k}^{2}(X,Y)=\frac{1}{n^{2}}\sum_{i=1}^{n}\sum_{j=1}^{n} \left(k({x}_{i},{x}_{j})+k({y}_{i},{y}_{j})-2k({x}_{i},{y}_{j})\right)$ $$

As a kernel function, we used the radial basis function kernel (RBF). We evaluated the distance between batches of sequences, each of size $n$ equal to 2048, sampled from the dataset and generated by the respective models. Following the methodology proposed for 3D structures in (Joshua Southern & Correia, ), we utilized ProtT5 sequence representations to calculate MMD.

1-Wasserstein optimal transport (OT).
The BLAST identity metric effectively evaluates the similarity between generated sequences and natural ones. However, its limitation lies in assessing the model’s capability to produce diverse sequences, as it may identify the same dataset sequence as the closest match for every generated sequence. To overcome this limitation, we employ transportation theory to establish optimal pairs between generated sequences and the dataset.

Optimal transport theory, initially devised for solving economic problems, has found applications in various fields, including physics, biology, and tomography. To implement our approach, we calculate pairwise Levenshtein distances and use them as transportation costs. Subsequently, we determine optimal sequence pairs using the Earth Mover Distance (EMD) solver with a uniform distribution of the samples. We use the average distance between these optimal pairs, measuring both the diversity and proximity of generated samples to the dataset.

The inherent diversity of the dataset, i.e., when a sample from the dataset pairs with itself, gives zero distances $(OT(dataset)=0)$. In contrast, random sequences form optimal pairs with the highest mean distances, as illustrated in figures 3. Although our OT realization offers advantages over BLAST, it has a special feature: the EMD solver identifies an exact pair for each sequence. This poses a challenge when dealing with two query sequences that are similar to one dataset sequence but distant from others, resulting in one close pair and one distant pair. However, we employ EMD precisely to penalize such cases, reinforcing the generation of diverse rather than similar sequences.

Sequence Diversity (SeqD).
We calculate batch SeqD as pairwise normalized Levenshtein distances as follows:

$$ $SeqD_{\text{absolute}}=\frac{\sum_{i=1}^{n-1}\sum_{j=i+1}^{n}\frac{levenshtein (s_{\text{i}},s_{\text{j}})}{max(len(s_{\text{i}}),len(s_{\text{j}}))}(s_{ \text{i}},s_{\text{j}}\in G)}{|G|(|G|-1)}\cdot 100$ $$

Here, G represents the generated sequences, $|G|$ is the number of generated sequences, and Levenshtein ($\cdot$) is the Levenshtein distance between the $i$-th and $j$-th sequences. Unlike (Wang et al., ), we normalize each Levenshtein distance by the maximum sequence length to mitigate the impact of sequence length on this metric. The desired SeqD value for a generative model should be close to that of the training data, so we used the absolute difference between the dataset and model batch SeqD:

$$ $SeqD=|SeqD_{\text{batch}}(Dataset)-SeqD_{\text{batch}}(Query)|$ $$

All models exhibit values close to the dataset, indicating the absence of mode collapse. The SeqD values are documented in Table  3.

Sequence clustering.
To obtain clusters, we first removed duplicates from the dataset. For clustering we used MMseqs2 (Steinegger & Söding, ) with parameters: identity = 0.3, coverage = 0.8, cov-mode = 0, cluster-mode = 1. We focused on the number of query sequences that ended up in the same clusters with the dataset sequences (positive sequences or PS) and the number of such clusters (positive clusters or PC). It is notable that generating out of distribution is also very important, so we evaluated the quality of generated sequences from other (”negative”) clusters. We found that the average pLDDT of these sequences from DiMA (SwissProt and AFDBv4-90: 65 ±14 and 63 ±12) is significantly higher than from other models (nanoGPT: SwissProt and AFDBv4-90 43 ±12 and 52 ±16), which might indicate that the model generalizes beyond the training data.

UMAP.
To visually represent the distribution of generated sequences across PC, we trained UMAP on all sequences from PC for all models (parameters - n_neighbors - 25 and min_dict - 0.5).
The UMAP plots in Figures  6 and  6 show that despite the fact that the SeqD metric of the DIMA w/o self-conditioning is higher, DIMA with self-condition has the same coverage on the SwissProt (and even more coverage on AFDBv4-90). This and the fact that DIMA is closer to the dataset in terms of distribution learning metrics shows that DIMA w/o self condition achieved better diversity by generating sequences that greatly differ from those from the dataset.

Figure: Figure 5: UMAP projection of sequences from PC. Training dataset - SwissProt. Grey background points - dataset sequences from PC.
Refer to caption: extracted/5452544/figures/swissprot_final_25_0.5.png

## Appendix C Model architecture

We use the 12-layer Transformer model with 16 attention heads and a hidden size of 320 as a backbone for our diffusion model with several modifications tailored to ensure the effective operation of denoising diffusion within the specific context of protein-related data.
Firstly we add trainable positional encodings to noisy protein latents.
Then, the input for the current transformer block is a sum of the output of the previous block, time embeddings, and self-condition prediction projected with linear layers.
We also incorporate long skip connections because for time steps close to zero the output of the model is similar to the input.
Such modification helps to learn identical transformation.
Figure 7 depicts the whole architecture.

Figure: Figure 7: The architecture of the diffusion model.
Refer to caption: extracted/5452544/figures/protein_arch.png

## Appendix D Length sampling

During the inference phase, the model needs to define the length of the generated sequence.
We compare two approaches to tackle this problem: training diffusion models both with and without pad masking.
In the first case, we feed additionally to corrupted latents the attention mask of pad tokens during training and ignore pad tokens for computing diffusion loss.
During inference, we sample the length from the empirical distribution of lengths in the training set.
In the second case, we do not provide any information about pad tokens during training and compute loss using all tokens in sequence including pad.
Then, during generation, the model should define the length by itself.
Figure 8 depicts the distribution of lengths in the training and generated by second approach sets.
The distribution of generated sequences differs from the training set on both datasets.
To avoid this distribution mismatch, we choose to use an attention mask during training and sample length during inference.

## Appendix E Work limitations

While DiMA exhibits promising outcomes in unconditional protein sequence generation, it’s worth noting that our datasets were derived from a subset of available known protein sequences, and the scalability of our model size has yet to be explored. This includes investigating scalability concerning both model and dataset sizes. In this work we limit ourselves to the monomeric proteins with the range of sequences from 128 to 254. While this range covers a great deal of proteins, it leaves important classes of long proteins out of the scope. Additionally, the current approach to sampling sequence lengths may require refinement for seamless application in conditional generation. These aspects form part of our ongoing research, which also includes exploring the practical utility of conditional generation.

Figure: Figure 8: The distribution of lengths in the training and generated sets for models trained on SwissProt (top) and AFDB (bottom).
Refer to caption: x2.png

Figure: Figure 9: Histogram depicting the occurrence of the top 15 most frequent SUPERFAMILY domains in the dataset pool. (Oates et al., ; Jones et al., ). Percentage of annotated protein sequences per model.
Refer to caption: extracted/5452544/figures/superfamily_fig_swissprot.png

Figure: Figure 10: Sequence annotation into known structural domains using SUPERFAMILY tool within InterProScan (Oates et al., ; Jones et al., ). Histogram of domain lengths (Red - 2048 generated sequences; blue - 2048 dataset sequences). No hits were found for random sequences and ProteinGAN-generated sequences.
Refer to caption: extracted/5452544/figures/SUPERFAMILY_swissprot_coverage.png

Figure: Figure 11: Sequence annotation into known structural domains using SUPERFAMILY tool within InterProScan (Oates et al., ; Jones et al., ). ESMFold-predicted structures of representative SUPERFAMILY domains generated by DiMA.
Refer to caption: extracted/5452544/figures/superfamily_structures.png

Figure: Figure 12: Prediction of intrinsic disorder regions (IDR) using the MobiDBLite tool (Piovesan et al., ). (A) Histogram depicting the lengths of intrinsic disorder regions. The blue color represents the dataset, while the red color represents the generated sequences. No hits were found for random sequences. (B) Representative examples of proteins generated by DiMA, highlighting intrinsic disorder regions in red and folded structural domains in blue.
Refer to caption: extracted/5452544/figures/MobiDBLite_swissprot.png

Figure: Figure 13: The number of secondary structure elements calculated per residue from ESMFold predicted structures using DSSP (Kabsch & Sander, ) software. H = $\alpha$-helix; B = residue in isolated $\beta$-bridge; E = extended strand, participates in $\beta$ ladder; G = 3-helix (310 helix); I = 5 helix ($\pi$-helix); T = hydrogen bonded turn; S = bend.
Refer to caption: extracted/5452544/figures/figure_annotation_DSSP.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Why were the distributional metrics computed relative to the train set? How do the authors justify the use of these metrics as an appropriate means to benchmark the learning capacity and generalizability of the tested models?
2. What was the sequence-level diversity/novelty achieved, relative to natural sequences? The repeated subsequence metric does not reflect this, and the 95% identity threshold is extremely lenient.
3. What was the reasoning behind the use of the SwissProt dataset when it is relatively so small, almost certainly represented within the ESM pretraining set, and given its size very easy to memorize?
4. What are the relative strengths, limitations, and novelties of DiMA versus another latent diffusion approach, Pro-LDM? These are not discussed. 
5. The evaluations of “biological relevance” do not demonstrate evidence of real-world utility of the approach. It is not surprising that high homology is achieved given training on SwissProt. Could the authors comment on the ability of DiMA to generate new sequences relevant to an independent downstream evaluation or task?
soundness: 3
strengths: While many works have presented compelling approaches for unconditional protein sequence generation, including sequence-based generative models (ProGen, ProtGPT, EvoDiff, etc), and latent diffusion models (eg ProLDM), to my knowledge this is a work that leverages the pre-trained ESM to define a latent space over which to learn the diffusion process. This is a valuable idea to contribute to the field, but it must hold in terms of evaluations. It is also interesting to provide a head-to-head comparison of discrete diffusion over amino acid sequences versus continuous diffusion over the latent sequence embeddings provided by ESM.
summary: This work presents an approach for unconditional protein sequence generation that operates via continuous diffusion in the latent space of protein sequence embeddings derived from the encoding model ESM2. Generated latent vectors are decoded back to amino acid sequence space via a lightweight decoder. By leveraging a pre-trained encoder like ESM2, the size of the backbone network driving the latent space diffusion process is relatively small. The authors present a series of evaluations that attempt to quantify the distribution modeling capabilities of their approach and the quality of sampled generations, including relative to internal ablations and other approaches trained on the SwissProt dataset (evaluations on AFDB are not presented relative to other approaches).

Post rebuttal: The authors' rebuttal addresses my primary concerns, and the promised revision promises to correct the significant instances of imprecision in the manuscript. I will adjust my scores accordingly.
weaknesses: There are several weaknesses to this work:
- The claim that “the foundational task of unconditional generation remains underexplored and underappreciated” is simply incorrect — many works (e.g., ProGen, ProtGPT, EvoDiff) have explored and evaluated performance of sequence-based unconditional generation.
- Discussion of related works on latent diffusion for proteins, such as Pro-LDM and CHEAP, is missing from the Introduction and Related Works. In general, many related works on diffusion models for protein generation are not discussed.
- The claim that “The straightforward adaptation of image-based diffusion models has not yet yielded satisfactory results in this domain,” is also unfounded. In fact, diffusion-based models for protein generation have been proven state-of-the-art in structure space; in sequence space, discrete diffusion is not at all a straightforward adaptation of image-based diffusion. This is not directly applying existing image-based diffusion, as the authors imply. This imprecision signals an underlying weakness of the work and the authors’ claims.
- The claim “most comprehensive set of metrics” is a gross overclaim and significantly weakens the credibility of the paper.
- All the distributional metrics presented seem to be computed relative to the train set — this calls into question the strength and credibility of the evaluations, since such comparisons to train will only reflect the ability of the generative models to memorize the train set. Such distributional evaluations should be computed relative to an independent test set.
- The presence of repeated amino acid subsequences is not a metric of “diversity”; similarly, the 95% sequence identity threshold as the barrier to novelty is extremely lenient.
- The claims that DiMA performs as well or better than the benchmarked approaches is not upheld by the evaluations, where DiMA is often outperformed.
- The training and comparisons using SwissProt are not compelling, given the small (~200-300K sequences) size of SwissProt, making it easy to memorize, especially considering the use of the training set in the distributional metrics.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: * In the data section, could authors clarify how they'd used the AFDB dataset, since it is a structural dataset? Why not just train on UniRef or another similarly sized whole-protein dataset?

Minor:
* line 237: description of "ProteinMPNN structure representations" is a bit murky. Which layer was used? Which ProteinMPNN structure? Where did the structure come from given that the generation is for sequences?
soundness: 3
strengths: * Ablation results are thorough and provide insight for others venturing into generating proteins via their latent embeddings.
* Results are comprehensive. Retraining baselines on the same dataset & controlling for parameter size is a good and rigorous effort.
* On the whole, sequence generation results are quite strong against baselines.
* Focusing on distribution similarity is IMO a great standard for this community, and it's nice to see that many of these metrics perform well.
summary: This work introduces DiMA, a protein sequence generation model by performing latent diffusion on ESM2-650M embeddings. The model is evaluated by sequence diversity, BLAST identity to known sequences, perplexity under autoregressive models, and distribution matching to embeddings of protein families. Structure sensibility is also assessed by examining the pLDDT and self-consistency perplexity. Some emphasis is also given to methodological design, including ablating the noise schedule used, self-conditioning, ESM2 encoding, time embedding information insertion into the model. When compared to several baselines that include autoregressive, energy-based, and diffusion-based sequence generation methods, trained to have the same parameter count and trained on the same datasets,
weaknesses: * The encoder compression sections feels underdeveloped. At minimum, it should probably include a citation to CHEAP (https://www.biorxiv.org/content/10.1101/2024.08.06.606920v1.full.pdf). A stretch experiment would be to see how sequence diffusion performance changes when doing so in the CHEAP embedding space. ESM 3B yields better results in Table 3, but we know that larger embedding dimensions is generally harder for latent diffusion models, so all of this feels a bit counter intuitive. I think at this point it's not too much of a mystery that we can do sequence generation by diffusion on ESM embeddings, and ample literature exists on diffusion design choices, but the finer details of how information is captured in these latent embeddings is more under-explored, and including it in this paper would strengthen its utility.
* "Most comprehensive evaluation" in the introduction is a strong claim. The evaluations are indeed very impressive, but adding a qualifier like this is rather misleading, and does not promote the paper's longevity.
* I think the work has a lot of results, which is great, but they feel a bit under-analyzed. Reducing comparisons to single valued points without any examples of generated sequences (or their folded structures) feels short of convincing.  In future iterations, the work could be strengthened by more biological and qualitative analyses like those on pages 24-28, and moving that into the main text. In its current form, the captions are vague; for e.g. Figure 11, 12, 13 especially.
* I think it's productive to have figures like Appendix B.2 which demonstrate distribution comparisons in the main text. Though table space is limited, it would be nice to report the stds from the mean if possible. 
* In the main baseline tables, is it possible to also add the length distributions of baselines and generated sequences, given how important they are for many of the reported metrics?

Minor:
* Line 127: "High-generalizing" reads slightly awkwardly
* This is not critical for me to raise my score, but could I think the results section would be better organized with better subsections, such that a study is presented right before the results. Found it a bit hard to jump between the sections and find the correct \paragraph{} heading. 
* fontsize inconsistencies in tables could be polished
* noise schedule formula is duplicated between main text and Appendix D.1

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Several critical points require attention in the manuscript:

[Technical Motivation and Advantages]

While the manuscript identifies protein sequence design as an important problem, it lacks sufficient justification for using latent diffusion over established methods like discrete diffusion. The technical advantages and unique capabilities of latent diffusion in this specific context should be pointed out. 


[Evaluation Metrics]
1. A novelty assessment comparing generated sequences with native sequences (e.g., sequence identity analysis) should be provided. One major concern is that the high-performance metrics could result from training set memorization rather than genuine generative capabilities. It is suggested to calculate sequence identity between generated sequences and their nearest neighbors in the training set. 
2. The choice of Progen perplexity (PPL) as a quality metric requires justification. The authors should provide more discussion to clarify whether PPL is reliable for protein sequence design (ref: https://arxiv.org/abs/2210.05892). 
3. Sequence length: The evaluation is limited to a narrow sequence length window (128-254 amino acids). This restriction raises important questions about method generalizability: How does performance scale with longer sequences? Could discrete diffusion methods demonstrate superior quality and diversity for longer sequences?

[Model Scaling Effects]

The observed decrease in distribution fitness scores (FD-seq, MMD-seq) when using the ESM3b encoder appears to contradict previous findings about model scaling benefits. This unexpected result warrants deeper analysis and discussion.
soundness: 3
strengths: The authors present comprehensive ablation studies and evaluations using multiple metrics to demonstrate their method's effectiveness in generating protein sequences. Their results show the method's ability to generate diverse sequences while maintaining structural fidelity to the protein space. The inclusion of distributional similarity metrics is interesting. This evaluation framework offers valuable insights for the protein sequence design community and establishes useful benchmarking standards.
summary: In this manuscript, the authors introduce a latent diffusion method that leverages representations derived from the protein language model, ESM-2, to design proteins by generating amino acid sequences. Extensive ablation studies and evaluations based on diverse matrics were performed by the authors showing the method's capacity in generating protein sequences. This study shows the latent diffusion based method shows practical benefits for protein sequence design in terms of computational efficiency and sequence diversity.
weaknesses: The manuscript does not sufficiently justify why latent diffusion-based sequence design methods would provide advantages over alternative approaches, particularly discrete diffusion methods, in the context of protein sequence design. The authors may provide more theoretical or practical benefits of their latent diffusion approach. To strengthen the work, the author should provide additional evaluation metrics including sequence novelty, and provide more discussion about the underlying relationships between the metrics they used, clarify how they complement or potentially contradict each other.

---

confidence: 4
contribution: 2
initial_rating: 1
presentation: 2
questions: - Can you clarify how DiMA’s approach contributes uniquely to the generative model area compared to existing methodologies?
- What specific rationale guided your choice of Swissprot and AFDB for protein sequence generation, and how do these structure-based databases specifically benefit your model? If this is the common practice in previous work, do the authors think incorporating sequence-based database like UniProt/UniRef will enhance the performance of DiMA? For example, in one of the baselines, EvoDiff, the authors there use evolutionary-scale sequence data to train their model.
- Could you provide a more detailed breakdown of the metrics used in your evaluations and other studies (eg., baselines) to support your claim of comprehensiveness of evaluation? For example, you can make a table to justify and list the necessity of some newly introduced metrics. That would be helpful for new readers to understand, ok, this is important to evaluate the protein design models and bring new insights to the community. 
- Why the authors choose ESM2-8M - a very small representation learning model as the encoder across the study. As a usual practice, most of works use 650M model. Could the author reason the necessity of using 8M model? Moreover, as a control comparison, in Table 3, 10, 11, the authors combine “DiMA-8M” with other ESM2 encoders while in the main experiment, the authors claims using DiMA-33M. How do the authors explain the discrepancy between? Or is this a typo?
- Could the author provide a illustration figure in the main text? That would be helpful for readers to quickly grasp what you have done.
soundness: 2
strengths: - The paper claims extensive evaluations using multiple metrics, which basically supports the findings.
- The idea of doing diffusion or flow matching in the latent space makes sense and is a good direction to explore. 
- DiMA’s efficiency in utilizing ten times fewer parameters than leading models is a significant advantage, suggesting potential for broader applications in protein design.
summary: The paper presents DiMA, a model leveraging latent diffusion on protein language model encodings (in specific, ESM2) for generating protein sequences. The authors claim that DiMA achieves superior quality, diversity, and distribution matching compared to existing autoregressive and discrete diffusion protein sequence language models while using fewer parameters.
weaknesses: - Firstly, the approach presented in this paper does not significantly advance the field of **generative models**(the primary area selected). While the application of latent diffusion is noted, it does not bring substantial new insights or methodologies to the existing body of research on generative models (even in the field of protein design). I agree on the experimental contribution of this paper and think the protein design problem is important and relevant. However, it is under-qualified as a top-conference paper given its presentation and methodology contribution.
- Inconsistent results presentation. In Tables 4 and 5 of the main text, the bold digits intended to denote the best results (i guess) are inconsistently labeled. This lack of clarity can lead to confusion about which results are truly the best and detracts from the paper’s credibility.
- The use of Swissprot and AFDB as sequence-based protein databases is not adequately justified. The authors should provide a clearer rationale for selecting these databases and explain how they enhance the effectiveness of the sequence generation process. Why not use Uniprot or Uniref on which ESM2 is pre-trained?
- The authors assert that they use “the most comprehensive set of metrics” to evaluate their model. However, this claim lacks substantiation within the paper, raising questions about the thoroughness of their evaluation process.
- The sections/paragraphs are badly organized and hard to follow, please consider improve the manuscript during rebuttal.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: - Please review and correct the citation formatting for consistency.

- I suggest including the forward transition kernel in Section 3, clarifying how the noise schedule is used to determine the mean and variance.

- For the CD0.95 metric, does a higher value indicate greater diversity? In Table 2, 0.990 is bolded, suggesting it is considered the best. Please clarify this metric.

- I wonder why is the FD-seq metric not zero for the ground truth dataset, Is this due to sampling variability? Additionally, FD-STRUCT appears consistently lower than FD-seq—is this due to the lower dimensionality of the structure representation?

- What was the intuition behind choosing 320 as the hidden dimension for the diffusion backbone model? Also, have you considered using the fine-tuned ESM-2 model directly, as seen in DPLM?

- **Potential Bias in pLDDT Evaluation**: Since you use ESMFold to evaluate pLDDT and the ESM-2 encoder in DiMA, could this create a risk of biased evaluation? It may be worth considering alternative pLDDT metrics, such as those from AF2 or OmegaFold 2, to mitigate potential bias.
soundness: 3
strengths: - **Comprehensive Evaluation Metrics**: The authors use a robust set of evaluation metrics, covering quality, distribution similarity, diversity, biological relevance, and both sequence-based and structure-based evaluations.

- **Well-Designed Ablation Study**: The study effectively demonstrates the contribution of each component, particularly the ESM encoder and self-conditioning. The authors also investigate the impact of the noise schedule and model scaling, offering valuable insights into best practices for latent diffusion models in protein generation.

- **Empirical Performance**: The authors compare DiMA against a broad range of baseline models and pretrained models for the unconditional generation task, providing a thorough analysis. DiMA achieves comparable generation quality with larger pretrained models using significantly fewer parameters.
summary: The paper introduces DiMA, a latent diffusion model designed for generating protein sequences using embeddings from the pre-trained protein language model ESM-2. DiMA aims to address the limitations of current models, which often struggle to achieve both high quality and diversity in unconditional protein sequence generation.

This work provides an in-depth exploration of model components and performs extensive evaluation across multiple metrics to validate DiMA’s sequence quality and diversity.
weaknesses: - **Limited to Unconditional Generation**: DiMA’s design is currently limited to unconditional generation, which poses challenges for performing tasks such as sequence inpainting, motif scaffolding, or guided design. This is a key limitation of DiMA and diffusion models in continuous embedding space. I suggest that the authors discuss and explore the potential of extending DiMA to these tasks, which could greatly enhance the versatility of this work.

- The authors could also consider evaluating DiMA on protein representation learning tasks, similar to those presented in DPLM.

- In the main experiment, it is stated, “For a fair comparison, we train each method from scratch with the same parameter count (33M) as DiMA on the same dataset.” However, DiMA’s dependency on the ESM-2 encoder increases the effective parameter count, which could affect the reliability of these comparisons. This should be clarified to ensure fair evaluation.

- **Inconsistencies in Tables**: Some tables lack consistency, particularly Table 5, where the best-performing models are not consistently bolded. For example, in the Progen ppl column, DPLM-3B is bolded, even though DPLM-150M performs better. Correcting these inconsistencies would improve clarity and fairness.

- **Paper Structure**: Improving the structure of the paper would enhance readability. For example, Section 4 currently only has one subsection (4.1). Creating subsections for ablation studies and other experiments could provide better organization and make the results easier to follow.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 2
questions: The paper is a valuable contribution to protein sequence generation through latent diffusion modeling. However, the work could be strengthened by deeper motivation for the latent diffusion choice and comparisons with recent models like Chroma and MultiFlow. I really admire this work for its comprehensive evaluation, but I doubt it is good enough for this top AI conference.
soundness: 2
strengths: 1. The combination of latent diffusion modeling with protein language model encodings is interesting.
2. The manuscript is well-organized and clearly written.
3. The experiments are extensive, with a set of rigorous evaluation metrics.
summary: The paper presents DiMA (Diffusion on Language Model Encodings for Protein Sequence Generation), a framework for generating protein sequences. DiMA addresses the challenge of generating high-quality, diverse protein sequences by leveraging latent diffusion on protein language model (pLM) encodings from ESM-2. Existing protein generation models struggle to balance quality, diversity, and parameter efficiency, especially for unconditional generation, which is essential for creating novel protein structures without predefined constraints.
weaknesses: 1. While latent diffusion is novel in this context, the motivation for choosing this approach over more conventional methods like sequence-based discrete diffusion or 3D-structure-based continuous diffusion is not strongly justified. The authors could strengthen their argument by highlighting specific challenges with current discrete or continuous approaches in handling the protein data structure.
2. It omits comparisons with strong baseline models like Chroma [1] and MultiFlow [2].
3. Although DiMA is designed to generate diverse sequences, the paper does not sufficiently analyze sequence variability at the biochemical or functional level.
4. The current citation format, as mentioned, should follow the convention of \citep rather than \cite for clarity and consistency.


[1] Ingraham, John B., Max Baranov, Zak Costello, Karl W. Barber, Wujie Wang, Ahmed Ismail, Vincent Frappier et al. "Illuminating protein space with a programmable generative model." Nature 623, no. 7989 (2023): 1070-1078.

[2] Campbell, Andrew, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. "Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design." In Forty-first International Conference on Machine Learning.

Meta-review: After all discussion phases have concluded, three reviewers judge the paper to be marginally above the acceptance threshold, and three reviewers judge the paper to be marginally or clearly below the acceptance threshold. 

Among other things, as positives the reviewers highlight useful ablation studies of the proposed method.
As areas of improvement reviewers highlighted the need for more thorough evaluation, an improved structure of the paper, correction of statements in the paper that were not deemed correct by the reviewers, removing too strong claims and a better discussion of related work. The authors have put in a substantial effort in their rebuttal to alleviate concerns, which has led to several increased scores, moving from [1, 6, 5, 3, 5, 3] to [5, 6, 6, 5, 6, 3]. Unfortunately, despite the productive discussion between reviewers and authors, the paper still needs some further improvement in terms of the evaluation of the method and the presentation of the paper for acceptance. I therefore recommend to reject this paper for this conference. I hope the authors can build on the feedback of the reviewers to further improve their manuscript and resubmit to another venue in future. Thank you to all reviewers and authors for engaging during the reviewing period.

Decision: Reject