Structure Language Models for Protein Conformation Generation

## Abstract

Abstract Proteins adopt multiple structural conformations to perform their diverse biological functions, and understanding these conformations is crucial for advancing drug discovery.
Traditional physics-based simulation methods often struggle with sampling equilibrium conformations and are computationally expensive.
Recently, deep generative models have shown promise in generating protein conformations as a more efficient alternative.
However, these methods predominantly rely on the diffusion process within a 3D geometric space, which typically centers around the vicinity of metastable states and is often inefficient in terms of runtime.
In this paper, we introduce Structure Language Modeling (SLM) as a novel framework for efficient protein conformation generation.
Specifically, the protein structures are first encoded into a compact latent space using a discrete variational auto-encoder, followed by conditional language modeling that effectively captures sequence-specific conformation distributions.
This enables a more efficient and interpretable exploration of diverse ensemble modes compared to existing methods.
Based on this general framework, we instantiate SLM with various popular LM architectures as well as proposing the
ESMDiff, a novel BERT-like structure language model fine-tuned from ESM3 with masked diffusion.
We verify our approach in various scenarios, including the equilibrium dynamics of BPTI, conformational change pairs, and intrinsically disordered proteins. SLM provides a highly efficient solution, offering a 20-100x speedup than existing methods in generating diverse conformations, shedding light on promising avenues for future research.

## 1 Introduction

Protein structure dynamics are fundamental to understanding the biological functions of proteins. The ability of proteins to adopt multiple conformations is crucial for their function in influencing interactions with other biomolecules and the environment. Traditional computational methods, such as molecular dynamics (MD) simulations, have long been used to explore these dynamics. However, these methods are computationally expensive and time-consuming. Structure prediction models, such as AlphaFold 2 (Jumper et al., ) and RosettaFold (Baek et al., ), have made significant strides in predicting static protein structures, yet often fail to accurately capture the dynamic nature of proteins and their multiple conformations (Chakravarty & Porter, ).
(^0^00 ^†Equal contribution. Code available at [https://github.com/lujiarui/esmdiff](https://github.com/lujiarui/esmdiff).)
(^0^00 Correspondence to: [jiarui.lu@mila.quebec](mailto:jiarui.lu@mila.quebec))

Recently, significant progress has been made by adopting deep generative models as conformation samplers to efficiently explore the complicated protein conformational space. For example, Noé et al. () adopts normalizing flow to match the underlying Boltzmann distribution by learning from simulation data.
Despite their potential, normalizing flow-based methods (Noé et al., ; Klein et al., ) face significant challenges in modeling large protein systems with hundreds of amino acids, as the invertibility constraint becomes a major obstacle when scaling up model parameters.
As a remedy, denoising diffusion approaches (Jing et al., ; Lu et al., ; Wang et al., ; Zheng et al., ) can efficiently learn from structural data, achieve good generalization, and perform amortized inference.
However, modeling high-dimensional protein structures explicitly in their 3D Euclidean space can demand intensive computation and usually requires accounting for special equivariant properties (Köhler et al., ).
Furthermore, L2-based training objectives such as denoising score matching (Song et al., ) tend to predict local perturbations rather than capturing remote modes of alternative conformations (Wang et al., ). Consequently, these models may overallocate their capacity to learn structural noises in the training data instead of focusing on low-frequency structural changes (Chou, ).

In complement with existing approaches, we present Structure Language Modeling (SLM), a novel framework for protein conformation generation that performs generative modeling in the latent space of protein structures.
Inspired by the recent progress in developing structural vocabularies for protein representation learning (Su et al., ; Hayes et al., ), our approach first encodes structural flexibility into a distribution over latent tokens using a discrete variational autoencoder, as illustrated in Fig. 1.
The discrete latent encoding removes high-frequency details of protein structures, forming “structure languages” that effectively capture the uncertainty of complex protein conformations (Fig. 2a);
Conditional language modeling is then applied to these latent structure tokens, using amino acid types as context to capture sequence-specific conformation distributions (Fig. 2b);
Protein conformations can finally be reconstructed by mapping structure tokens into 3D space with a learned decoder (Fig. 2c).
By leveraging generative language modeling in the discrete latent space, SLM bypasses the complexity of equivariant constraints associated with geometric symmetries and benefits from enhanced model capacity.
As a general framework, SLM is fully compatible with any existing language model (LM) architectures and shows promising scalability.
To further demonstrate the versatility of our approach, we introduce ESMDiff, a novel BERT-like structure language model instantiation fine-tuned from ESM3 (Hayes et al., ) with masked discrete diffusion (Austin et al., ; Zhao et al., ) grounded in the SLM framework.
Experimental results across various conformation generation scenarios demonstrate the state-of-the-art performance of SLM including the representative ESMDiff model, achieving orders of magnitude faster speeds compared to existing generative methods.
The proposed framework paves the way for new research avenues in addressing the protein conformation sampling challenge.

We summarize our key contributions as follows.

- •
We comprehensively explore an innovative conformation generation framework based on language modeling in the latent space, which opens up potential research avenues.
- •
We introduce ESMDiff, a novel fine-tuned variant of a state-of-the-art protein language model, built on masked discrete diffusion.
- •
We demonstrate the superior capability of structure language models by evaluating them on various conformation generation settings and comparing them with existing methods.

Figure: Figure 1: Residue flexibility (BPTI clusters, Shaw et al. ()) reflected by the categorical distribution over latent structure tokens. Different tokens (colored in different shades) are used to encode varying local structural patterns.
Refer to caption: x1.png

## 2 Related Work

Protein language models.
In recent years, several language models of protein sequence have been built. Among these, the ESM-series (Rives et al., ; Lin et al., ; Hayes et al., ) and other similar models (Elnaggar et al., ; Alley et al., ) have garnered great attention because of their wide range of downstream applications such as protein engineering (Meier et al., ). On the other hand, auto-regressive protein language models, based on either recurrent neural networks (Alley et al., ), or Transformer including ProGen (Madani et al., ) and ProtGPT2 (Ferruz et al., ), are able to generate de novo sequences with input controlling tokens. Specially, inverse folding models (Ingraham et al., ; Jing et al., ; Hsu et al., ; Dauparas et al., ; Gao et al., ) learn to perform structure-based protein design with geometric-aware encoders.

Generative conformation sampling. Given the intensive computation of traditional MD simulations, generative models have been used to learn conformation distributions in a data-driven fashion. The Boltzmann generator (Noé et al., ) uses normalizing flow to fit the Boltzmann distribution from target-specific simulation data. Arts et al. () extends this by using denoising diffusion models for coarse-grained protein conformations. Furthermore, EigenFold (Jing et al., ), Str2Str (Lu et al., ), AlphaFlow (Jing et al., ), ConfDiff (Wang et al., ), and DiG (Zheng et al., ) leverage diffusion or flow matching to conditionally sample protein conformations by learning from PDB data. Recently, AlphaFold3 (Abramson et al., ) revised the structure decoder of AlphaFold2 to a diffusion-based module for diversified structure prediction.

Quantized representation for protein structures.
Beside the prevailing diffusion models for protein structure, representation learning of protein structures using discrete variational autoencoders (dVAE) has gained increasing attention in recent years. FoldSeek (van Kempen et al., ) is one of the earliest attempt to build dVAE for fast structure search and alignment. Based on this, SaProt (Su et al., ) constructs learned representations with both sequence and structure tokens as input, while ProtT5 (Heinzinger et al., ) fine-tuned an existing language model to accept structure tokens as input. PVQD (Haiyan et al., ) applied latent diffusion in the embedding space of dVAE for conditional protein structure generation. ProSST (Li et al., ) trained an autoencoder with K-means clustering applied in the latent space. Gaujac et al. () and Gao et al. () respectively build dVAE with large vocabularies for learning protein structure representations.

Remarks: Our work is closely related to these concurrent research directions by leveraging LMs to model and efficiently perform conformation generation over the quantized representation of protein structures. We refer to this framework as “structure language models” and describe them in detail.

## 3 Protein Conformation Generation with Language modeling

Notation.
A protein with $L$ residues is identified by its sequence of amino acid types ${\bm{c}}\in|{\mathcal{S}}|^{L}$ where ${\mathcal{S}}$ is the vocabulary of 20 standard amino acids. The protein (backbone) structure is represented by its composing 3D atom positions ${\bm{x}}\in{\mathcal{X}}\equiv\mathbb{R}^{L\times 4\times 3}$ including all backbone heavy atoms. Through an encoder $q({\bm{z}}|{\bm{x}})$, the structure ${\bm{x}}$ is encoded to a sequence of latent codes ${\bm{z}}\sim q({\bm{z}}|{\bm{x}})$ where ${\bm{z}}\equiv({\bm{z}}_{1},\dots,{\bm{z}}_{L})\in|V|^{L}$ and $V$ is the pre-specified vocabulary of latent codes; the structure tokens ${\bm{z}}$ are decoded by first mapping to embedding vectors and then to the 3D structure ${\bm{x}}$.

Figure: Figure 2: An illustration of the proposed SLM framework.
Refer to caption: x2.png

### 3.1 Learning the sequence-structure distribution

To address the conformation generation problem, we start with modeling the sequence-to-structure translation distribution $p({\bm{x}}|{\bm{c}})$ of interest and derive the learning objective in this section.
To circumvent explicitly learning in the structure space, the roto-translation invariant(^1^11For example, features like distance and angle are roto-translation invariant. This relationship can be formally written as $q({\bm{z}}|T\circ{\bm{x}})\triangleq q({\bm{z}}|R\circ{\bm{x}}+t)=q({\bm{z}}|{ \bm{x}}),\forall T$.) latent representation ${\bm{z}}$ is introduced to encode 3D atomic protein structure.
Given this, the target distribution $p({\bm{x}}|{\bm{c}})$ can be derived by marginalizing the joint distribution $p({\bm{x}}|{\bm{c}})=\int_{\bm{z}}p({\bm{x}},{\bm{z}}|{\bm{c}})$.
We further factorize this joint distribution according to the Bayes’ rule by isolating the latent variable ${\bm{z}}$: $p_{\theta,\phi}({\bm{x}},{\bm{z}}|{\bm{c}})=p_{\phi}({\bm{x}}|{\bm{c}},{\bm{z}
})p_{\theta}({\bm{z}}|{\bm{c}})$, where $p_{\phi}$ denotes the (decoding) distribution over the 3D protein structures given the structure token and sequence, and $p_{\theta}$ denotes the conditional distribution over the structure tokens, respectively modeled by neural networks with parameter set $\phi,\theta$.
This gives rise to the evidence lower bound on the likelihood of model distribution over protein structures conditioned on sequence:

$$ $\log p_{\theta,\psi}({\bm{x}}|{\bm{c}})\geq\mathbb{E}_{q_{\psi}({\bm{z}}|{\bm{ x}})}\left[\log p_{\phi}({\bm{x}}|{\bm{c}},{\bm{z}})\right]-D_{\text{KL}}(q_{ \psi}({\bm{z}}|{\bm{x}})\|p_{\theta}({\bm{z}}|{\bm{c}}))\triangleq\mathcal{L}( {\phi,\theta}),$ (1) $$

where $\psi$ is introduced to parameterized the posterior distribution over latent representation ${\bm{z}}$. Please refer to the Appendix LABEL:app:elbo for the full derivation of Eq. 1. Directly optimizing the right-hand side of Eq. 1 can be intractable and difficult since we have unknown posterior $q_{\psi}$. As a result, we adopt an one-step expectation–maximization (EM) approach (Dempster et al., ) by first jointly learning $p_{\phi}$ and $q_{\psi}$ with a simple and parameter-free prior distribution $p({\bm{z}}|{\bm{c}})$, followed by optimization on $p_{\theta}$ with the learned $p_{\phi}^{*}$ and $q_{\psi}^{*}$.
This yields the overall two-stage and separable training pipeline:

Learning quantized representation for structure.
With the prior $p({\bm{z}}|{\bm{c}})$ fixed, we begin by maximizing the ELBO $\mathcal{L}({\phi,\theta})$ with respect to the encoder $\psi$ and decoder $\phi$, using protein structure samples $\mathcal{D}={{\bm{c}},{\bm{x}}}$. In the context of discrete latent spaces, this process is analogous to training a discrete VAE (dVAE) (Van Den Oord et al., ) to learn quantized representations for protein structures. Here, the encoder $q_{\psi}({\bm{z}}|{\bm{x}})$ maps structures to latent tokens, while the decoder $p_{\phi}({\bm{x}}|{\bm{z}},{\bm{c}})$ reconstructs structures from these tokens(^2^22We assume that ${\bm{x}}$ is conditionally independent of ${\bm{c}}$ given latent variable ${\bm{z}}$.). The prior $p({\bm{z}}|{\bm{c}})$ is fixed to be uniform during this stage.

Learning the prior over latent tokens.
In this stage, we fix the learned parameters $\phi^{*}$ and $\psi^{*}$, and train the prior $p_{\theta}$ by maximizing the ELBO: $\arg\max_{\theta}\mathcal{L}({\phi^{*},\theta})$. Since both $\phi^{*}$ and $\psi^{*}$ are fixed, the reconstruction term in the ELBO cancels out, and training reduces to minimizing the KL divergence $D_{\text{KL}}(q_{\psi^{*}}\|p_{\theta})$. This is equivalent to performing maximum likelihood estimation, as $\mathbb{E}_{({\bm{c}},{\bm{x}})\sim\mathcal{D}}\mathbb{E}_{{\bm{z}}\sim q_{
\psi}({\bm{z}}|{\bm{x}})}p_{\theta}({\bm{z}}|{\bm{c}})$ with respect to $p_{\theta}$. Given that both ${\bm{z}}$ and ${\bm{c}}$ are categorical variables, this formulation resembles a translation task, allowing $p_{\theta}$ to be parameterized by language models.

### 3.2 structure language modeling

The prior learned in the previous stage is now applied to conformation generation, which can be framed as a conditional generative modeling problem for sequence-to-structure (seq2str) translation. Given an input condition ${\bm{c}}\in|{\mathcal{S}}|^{L}$ which determines the molecular topology, the goal is to sample a conformation ensemble from $p({\bm{x}}|{\bm{c}})$. To do this, we first sample a set of latent variables from the prior distribution learned earlier, ${\bm{z}}\sim p_{\theta}({\bm{z}}|{\bm{c}})$, and then decode these latents using the decoder $p_{\phi}({\bm{x}}|{\bm{c}},{\bm{z}})$. The decoder is jointly trained with the encoder $q_{\psi}({\bm{z}}|{\bm{x}})$ in the first stage, ensuring that the sampled latents align with the reconstruction. This framework supports roto-translation invariant inference and is described in Algorithm 1.
Next, we illustrate this approach with two straightforward examples of structure language models (SLM): the encoder-decoder and decoder-only architectures.

Encoder-decoder.
Given the conditional nature of translation, the prior $p_{\theta}({\bm{z}}|{\bm{c}})$ can be explicitly modeled by an encoder-decoder architecture like T5 (Raffel et al., ). The decoder conditions on the context ${\bm{c}}$ and factorizes the structure tokens sequentially:
$p({\bm{z}}|{\bm{c}})=\prod_{l=1}^{L}p({\bm{z}}_{l}|{\bm{z}}_{<l},{\bm{c}}),$
where ${\bm{z}}\in{\mathcal{Z}}$ represents the quantized structure tokens. The training objective is the negative log-likelihood (NLL) loss conditioned on ${\bm{c}}$:
${\mathcal{L}}(\theta)=-\mathbb{E}_{({\bm{c}},{\bm{x}})\sim{\mathcal{D}}}
\mathbb{E}_{{\bm{z}}\sim q({\bm{z}}|{\bm{x}})}\sum_{l=1}^{L}\log p_{\theta}({
\bm{z}}_{l}|{\bm{z}}_{<l},{\bm{c}}),{\bm{z}}_{<1}\equiv\varnothing.$

Decoder-only. Alternatively, the latent prior $p_{\theta}({\bm{z}}|{\bm{c}})\propto p_{\theta}({\bm{c}},{\bm{z}})$ can be modeled autoregressively using a decoder-only architecture, such as GPT (Radford et al., ), where ${\bm{c}}$ serves as the “prompt”. We define ${\bm{y}}\triangleq[{\bm{c}},{\bm{z}}]=[{\bm{c}}^{1},\dots,{\bm{c}}^{L},{\bm{z}
}^{1},\dots,{\bm{z}}^{L}]$, and the training involves maximizing the likelihood over ${\bm{y}}$ via the NLL minimization:
${\mathcal{L}}(\theta)=-\mathbb{E}_{({\bm{c}},{\bm{x}})\sim{\mathcal{D}}}
\mathbb{E}_{{\bm{z}}\sim q({\bm{z}}|{\bm{x}})}\sum_{l=1}^{2L}\log p_{\theta}({
\bm{y}}_{l}|{\bm{y}}_{<l}),$
where ${\bm{c}},{\bm{x}}\sim{\mathcal{D}}$ is the i.i.d. samples from the data distribution over structures each with the associated amino-acid sequence as condition ${\bm{c}}$. In practice, we add an additional special token [sep] to differentiate between these two modalities.

Inference involves sampling with a left-to-right decoding order, as defined by the autoregressive factorization of both language models. Figure 3 briefly illustrates these two modeling strategies.

Figure: Figure 3: Autoregressive prior modeling the for latent structure tokens discussed in Section 3.2.
Refer to caption: x3.png

## 4 ESMDiff: a Masked Diffusion instantiation

Building on the foundation of SLM, we here propose ESMDiff as an instantiation based on discrete diffusion models (Austin et al., ). ESMDiff incorporates the inductive bias of seq2str translation and leverages the protein foundation model ESM3 (Hayes et al., ) through masked diffusion fine-tuning. The effectively fine-tuning of ESMDiff also exemplifies how a large pretrained BERT-like masked language model can be adapted to acquire additional generative capabilities, making it well-suited for broader downstream tasks such as conformation generation.

### 4.1 Revisiting Discrete diffusion as distribution interpolation

The discrete diffusion models (Austin et al., ; Lou et al., ; Sun et al., ; Campbell et al., ; Zheng et al., ) can be generally defined by a sequential process of progressive noisy variables $z_{t}\in V$ from the categorical variable $z_{0}\in V$. Denote the one-hot (row) vector of $z_{t}$ as ${\bm{z}}_{t}\in\{0,1\}^{|V|}$, in the discrete-time case (Austin et al., ), the forward marginal probability of ${\bm{z}}_{t}$ at time $t$ has the following form as a composition of Markov kernel defined by ${\bm{Q}}_{t}~{}(t=1,2\dots,T)$:

$$ $q({\bm{z}}_{t}|{\bm{z}}_{0})=\mathrm{Cat}\left({\bm{z}}_{t};{\bm{z}}_{0}\bar{{ \bm{Q}}}_{t}\right)\triangleq\mathrm{Cat}\left({\bm{z}}_{t};{\bm{z}}_{0}{\bm{Q }}_{1}\cdot\dots\cdot{\bm{Q}}_{t}\right),$ (2) $$

where ${\bm{Q}}_{t}$ indicates the transition probability matrix for time $t$ represented by $[{\bm{Q}}_{t}]_{ij}=q(z_{t}=j|z_{t-1}=i)$, and $\mathrm{Cat}(\cdot;{\bm{p}}),{\bm{p}}\in\Delta^{|V|}$ indicates the categorical distribution with probability and $\Delta^{|V|}$ is the $|V|$-simplex. Eq. 2 also induce the form of the marginal distribution for $\forall t>s$ is $q({\bm{z}}_{t}|{\bm{z}}_{s})=\mathrm{Cat}\left({\bm{z}}_{t};{\bm{z}}_{s}\bar{{
\bm{Q}}}_{t|s}\right)\triangleq\mathrm{Cat}\left({\bm{z}}_{t};{\bm{z}}_{s}{\bm
{Q}}_{s+1}\cdot\dots\cdot{\bm{Q}}_{t}\right)$. Correspondingly, the posterior $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})$ can be obtained by the reverse process (Austin et al., ):

$$ $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})=\frac{q({\bm{z}}_{t}|{\bm{z}}_{s},{ \bm{z}}_{0})q({\bm{z}}_{s}|{\bm{z}}_{0})}{q({\bm{z}}_{t}|{\bm{z}}_{0})}= \mathrm{Cat}\left({\bm{z}}_{s};\frac{{\bm{z}}_{t}{{\bm{Q}}}_{t|s}^{\top}\odot{ \bm{z}}_{0}\bar{{\bm{Q}}}_{s}}{{\bm{z}}_{0}\bar{{\bm{Q}}}_{t}{\bm{z}}_{t}^{ \top}}\right),\forall~{}s<t.$ (3) $$

Both Zhao et al. () and Shi et al. () discuss how the discrete-time diffusion process can be generalized to the time domain $t\in[0,1]$, akin to the diffusion over continuous space (Song et al., ), by demonstrating the continuous-time limit as $T\to\infty$. Notably, when the stationary distribution is explicitly specified (denoted as ${\bm{p}}\in\Delta^{|V|}$), we can choose a state-independent transition kernel in the simple form:
${{\bm{Q}}}_{t|s}\triangleq\left[\alpha(s)^{-1}\alpha(t){\bm{I}}+(1-\alpha(s)^{
-1}\alpha(t)){\bm{1}}^{\top}{\bm{p}}\right],$
thus simplifying the continuous-time forward marginal to:

$$ $q({\bm{z}}_{t}|{\bm{z}}_{s})=\mathrm{Cat}\left({\bm{z}}_{t};\frac{\alpha(t)}{ \alpha(s)}{\bm{z}}_{s}+(1-\frac{\alpha(t)}{\alpha(s)}){\bm{p}}\right),\forall~ {}0\leq s<t<1,$ (4) $$

where $\alpha(t)\in[0,1)$ is a strictly monotone decreasing function with $\alpha_{0}=1$ and $\alpha_{1}\to 0$. The equation above demonstrates that the discrete diffusion, when defined with an explicit stationary distribution, can be viewed as an interpolation between two categorical distributions controlled by $\alpha(t)$.
According to Eq. 3, the reverse process of diffusion defined in Eq. 4 takes the following form for the posterior distribution, where $0\leq s<t<1$:

$$ $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})=\mathrm{Cat}\left({\bm{z}}_{s};\frac {\left[\mu(t,s){\bm{z}}_{t}+(1-\mu(t,s))\lambda_{{\bm{z}}_{t}}({\bm{p}}){\bm{1 }}\right]\odot\left[\alpha(s){\bm{z}}_{0}+(1-\alpha(s)){\bm{p}}\right]}{\alpha (t)\lambda_{{\bm{z}}_{t}}({\bm{z}}_{0})+(1-\alpha(t))\lambda_{{\bm{z}}_{t}}({ \bm{p}})}\right),$ (5) $$

where $\mu(t,s)\triangleq\alpha(s)^{-1}\alpha(t)>0$ and indicator function $\lambda_{{\bm{z}}_{t}}(\cdot)\triangleq\langle{{\bm{z}}_{t}},\cdot\rangle$ for concision. Readers are referred to Appendix LABEL:app:dd for more details on deriving the Eq. 4 and 5.

Unlike open-ended text generation, protein conformation generation is well-defined within discrete diffusion models, as it conditions on the input amino acid sequence, allowing each output token to correspond uniquely to a position in the input and thus enjoy a fixed-length context window(^3^33This indicates that ${\bm{c}}_{[i]}$ and ${\bm{z}}_{[i]}$ are aligned at the same position index $i$.).

### 4.2 Conditional Masked Diffusion Language Model

Masked diffusion (Austin et al., ; Lou et al., ; Shi et al., ; Sahoo et al., ) represents a special case in which the transition includes an “absorbing state”, denoted as [MASK] . In this formulation, the stationary distribution in Eq. 2 assigns all probability mass to the unique special token [MASK] , such that $P(z=\textrm{[MASK] })=1$ and $P(z\neq\textrm{[MASK] })=0$. For convenience, we define ${\bm{p}}_{M}\in\{0,1\}^{|\bar{V}|}$ ($\bar{V}\triangleq V\cup\{\textrm{[MASK] }\}$) as the one-hot vector representing [MASK] . In masked diffusion, the stochastic forward process maps $z_{0}\to\textrm{[MASK] }$ and remains in this state thereafter (i.e., “absorbing”). Conversely, the reverse process gradually unmasks (denoises) the [MASK] token to produce the data sample ${\bm{z}}_{0}$, where $s<t$ (see Appendix LABEL:app:proof):

$$ $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})=\mathrm{Cat}\left({\bm{z}}_{s};[ \beta(s,t)+(1-\lambda_{M}({\bm{z}}_{t}))(1-\beta(s,t))]{\bm{z}}_{t}+\lambda_{M }({\bm{z}}_{t})(1-\beta(s,t)){\bm{z}}_{0}\right),$ (6) $$

where $\beta(s,t)=\frac{1-\alpha(s)}{1-\alpha(t)}$ and $\lambda_{M}({\bm{z}}_{t})=\langle{\bm{p}}_{M},{\bm{z}}_{t}\rangle$. Eq. 6 implies when ${\bm{z}}_{t}\neq{\textrm{[MASK] }}$, the backward process simply copies the unmasked token by ${\bm{z}}_{s}\leftarrow{\bm{z}}_{t}$, i.e. $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})=\mathrm{Cat}\left({\bm{z}}_{s};{\bm{
z}}_{t}\right)$; otherwise the probability mass interpolates between ${\bm{p}}_{M}$ and ${\bm{z}}_{0}$. The posterior $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0})$ can be approximated by $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t})$ using re-parameterization: $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t})=q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{u}}_{
\theta}(t,{\bm{z}}_{t}))$, where the neural net ${\bm{u}}_{\theta}\in\Delta^{|\bar{V}|}$ is a neural network that outputs a probability vector that remains in $\Delta^{|\bar{V}|}$.

For conformation generation, we now consider the conditional case of masked diffusion. Given the amino acid types ${\bm{c}}$, our goal is to sample structure tokens through Eq. 6, utilizing a conditional posterior $q({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{z}}_{0};{\bm{c}})$. This posterior can be re-parameterized similarly by incorporating the condition into the backbone model, resulting in $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t};{\bm{c}})=q({\bm{z}}_{s}|{\bm{z}}_{t},{
\bm{u}}_{\theta}(t,{\bm{z}}_{t},{\bm{c}}))$.
To achieve this goal, the reverse process simulated by $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t};{\bm{c}})$ must effectively approximate the data distribution $p({\bm{z}}|{\bm{c}})$. A feasible training objective is to optimize the estimation of the conditional ELBO within the continuous-time integral, resulting in the following loss (see Appendix LABEL:app:loss for details):

$$ $\mathcal{L}(\theta)=\mathbb{E}_{{\bm{c}},{\bm{z}}_{0}}\left\{\int_{t\in[0,1)} \mathbb{E}_{{\bm{z}}_{t}\sim q({\bm{z}}_{t}|{\bm{z}}_{0})}\left[\frac{1}{1- \alpha(t)}\frac{\partial\alpha(t)}{\partial t}\lambda_{M}({\bm{z}}_{t})\log \langle{\bm{u}}_{\theta}(t,{\bm{z}}_{t},{\bm{c}}),{\bm{z}}_{0}\rangle\right] \text{dt}\right\},$ (7) $$

where ${\bm{z}}_{0}$ is sampled from the learned encoder $q_{\phi}({\bm{z}}|{\bm{x}})$ with the corresponding amino acid condition ${\bm{c}}$ from the data distribution $p({\bm{x}},{\bm{c}})$, and $\lambda_{M}({\bm{z}}_{t})$ implies the loss is only applied for the latents $\forall t,\textrm{s.t. }{\bm{z}}_{t}=\textrm{[MASK] }$. In practice, we can employ Monte Carlo estimation to compute the integral.

Figure: Figure 4: Illustration for conditional denoising network ${\bm{u}}_{\theta}({\bm{z}}_{t},{\bm{c}})$ where the masked (structure) tokens are colored grey. The unmasked tokens are carried over to the output without update, while the masked tokens are under random transition into either unmasked (✓) or still masked (✗) state.
Refer to caption: x4.png

Figure: Algorithm 1 Inference: Conformation Generation of SLM

### 4.3 Bidirectional Encoder as denoising network

We now discuss the implementation of the conditional denoising network using bidirectional encoder language models, such as BERT (Devlin, ). First, consider the sequential generalization of masked diffusion with a sequence of categorical variables. Let ${\bm{z}}_{t}$ now be a sequence of discrete structure tokens $\left[{\bm{z}}_{t,[1]},{\bm{z}}_{t,[2]},\dots,{\bm{z}}_{t,[L]}\right]$ where ${\bm{z}}_{t,[i]}\in\bar{V},\forall i=1,\dots,L$. Due to the interpolation scheme of Eq. 4 and Eq. 6, we assume conditional independence and factorize the posterior distribution $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t};{\bm{c}})$ across the $L$ output tokens, such that $p_{\theta}({\bm{z}}_{s}|{\bm{z}}_{t},{\bm{c}})=\prod_{i=1}^{L}p_{\theta}({\bm{
z}}_{s,[i]}|{\bm{z}}_{t},{\bm{c}})=\prod_{i=1}^{L}q\left({\bm{z}}_{s,[i]}|{\bm
{z}}_{t,[i]},{\bm{u}}_{\theta,[i]}(t,{\bm{z}}_{t},{\bm{c}})\right)$ where ${\bm{u}}_{\theta,[i]}(t,{\bm{z}}_{t},{\bm{c}})$ represents the $i$-th output channel of neural network.
This implement coincides with the BERT-style transformer architecture and allow us to take advantage of existing protein foundation model, for example ESM3 (Hayes et al., ). For a sequence of tokens, the masked $\log$-term in the training objective from Eq. 7 is replaced by the summation: $\sum_{i=1}^{L}\lambda_{M}({\bm{z}}_{t,[i]})\log\langle{\bm{u}}_{\theta,[i]}(t,
{\bm{z}}_{t},{\bm{c}}),{\bm{z}}_{0,[i]}\rangle$, with notations the same as defined above.

Modifications.
The following are special considerations for the network:
(1) Position-coupled encoding. Unlike general translation problem, SLMs maintain strict position-to-position correspondence between amino acid types and the latent tokens(^4^44The underlying co-evolutionary relationships between residues are fully shared across both amino acid types and its spatial patterns.)
. This inductive bias enables us to construct the input embedding for all position $i$ as follows:
${\bm{e}}_{[i]}=f_{\theta}[e_{z}({\bm{z}}_{t,[i]})+e_{c}({\bm{c}}_{[i]})+e_{t}(
t)]\in\mathbb{R}^{D},$
where $e_{z}:|\bar{V}|\mapsto\mathbb{R}^{D},e_{c}:|{\mathcal{S}}|\mapsto\mathbb{R}^{D
},e_{t}:\mathbb{R}\mapsto\mathbb{R}^{D}$ are the embedding functions and $f_{\theta}$ is a linear transformation. (2) Copying. The unmasked tokens ${\bm{z}}_{t,[i]}\neq\textrm{[MASK] }$ remain the same in spite of the model output. (3) Zero-out [MASK] . Since ${\bm{u}}_{\theta}$ parameterize the approximated clean data ${\bm{z}}_{0}$ (fully unmasked), the [MASK] token cannot present in the output and its probability should be zero-out. This is equivalent to adding $-\infty$ to the logit. In our study, the pre-trained LM head of ESM3 is replaced with a randomly initialized head with augmented vocabulary ($\bar{V}$) during fine-tuning.

## 5 Experiments

Base settings. We start with the pre-trained dVAE established in Hayes et al. () as the structure tokenizer (frozen). The structure quantization is perform residue-level with a receptive field over local geometric neighborhoods of protein structure. The structure language models as described in Section 3.2 are based off state-of-the-art language models as follows: (1) S-T5 (384M) adopts the T5 (Raffel et al., ) architecture with a bidirectional encoder and an autoregressive decoder.
(2) S-GPT (961M) models the joint distribution with an uni-directional decoder like GPT2 (Radford et al., ). (3) ESM3 (zero shot) is a pre-trained BERT model over multi-modality data and perform zero-shot inference by Gibbs sampling.
(4) ESMDiff is a fine-tuned variant of (3) on the PDB data using the masked diffusion objective in Eq. 7. For S-T5 and S-GPT, we embed the sequence tokens with the pre-trained ESM3-1.4B encoder to provide model condition. For ESMDiff, two different sampling paradigm: Gibbs and DDPM are considered (see Appendix B for details).

**Table 1: Evaluation results of different conformation generation methods on generating the BPTI conformations. JS-* represents the Jensen-Shannon (JS) divergence between the sampled ensemble and long MD ensemble on pairwise distance (PwD), time-lagged independent components (TIC), and radius of gyration (Rg). The validity as the frequency of clash-free samples in the ensemble is also included. Moreover, the ensemble TM-score and ensemble RMSD are reported with respect to the five kinetic clusters in Shaw et al. (). The best results among the SLM family are bold.**
$$ $f(i)=-\sum_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime})\log p_{[i]}({ \bm{z}}_{[i]}=z^{\prime}),$ (8) $$

### 5.1 Structural dynamics of BPTI

In the first experiment, we evaluate the models by generating the conformations of protein bovine pancreatic trypsin inhibitor (BPTI). The structural dynamic patterns of BPTI are well acknowledged in Shaw et al. () with 1ms-long MD simulations, based on which five kinetic clusters have been revealed. Similar to Lu et al. (), we report the Jensen Shannon (JS) divergence for distributions of pairwise distance (PwD), time-lagged independent components (TIC), and radius of gyration (Rg); the clash-free validity and the ensemble TM-score and root-mean-square-deviation (RMSD) w.r.t. the kinetic clusters. The benchmark results are shown in Table 5.
Following Wang et al. (), we also evaluate the best distance of the generated samples to each cluster, as shown in Table 5.1. The RMSD and TM-score are both calculated using the TM-score binary (Zhang & Skolnick, ) with structural alignment. Note that the Cluster 3 is a difficult remote folding mode (Wang et al., ), yet SLMs achieve a significant improvement by modeling with a smaller matching RMSD.

**Table 2: Evaluating on the best matching RMSD with respect to each of the kinetic clusters reported in Shaw et al. () with an non-exhaustive ensemble size of $N=100$. $\textrm{RMSD}_{\textrm{Ci}}(i=1,2,3,4,5)$ are the lower the better ($\downarrow$). Among them, the Cluster 3 is the most challenging case to model (Wang et al., ) and highlighted in red, of which the best sample from ESMDiff is visualized.**
$$ $f(i)=-\sum_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime})\log p_{[i]}({ \bm{z}}_{[i]}=z^{\prime}),$ (8) $$

### 5.2 Conformation changing pairs

We continue the study on the task of modeling and predicting conformational changes in structural proteins. The authors of Jing et al. () have curated two benchmarking sets of pairing data including (1) 77 pairs of fold-switching proteins (Chakravarty & Porter, ); and (2) 90 apo/holo pairs with ligand-induced conformational change (Saldaño et al., ) to evaluate the modeling capacity of conformation diversity. Following the setting and evaluation metrics of Jing et al. (), we randomly sample a few-shot ensemble with five structures per target and evaluate them based on the correlation metrics of residue flexibility and the ensemble TM-score (Zhang & Skolnick, ). The evaluation results for both test sets are shown in Table 5.2.

**Table 3: Evaluation on the conformation changing pairs (Saldaño et al., ; Chakravarty & Porter, ): (1) Pearson correlation $r$ between sampled diversity and ground-truth diversity measured by the residue flexibility (ResFlex, absolute deviation after alignment), and (2) the ensemble TM-score (TM-ens). For the residue flexibility, both global (gl.) and per-target (pt.) mean/median correlations are reported; for the TM-ens, both mean/median are reported. Metrics are the higher the better ($\uparrow$).**
$$ $f(i)=-\sum_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime})\log p_{[i]}({ \bm{z}}_{[i]}=z^{\prime}),$ (8) $$

### 5.3 Intrinsically disordered proteins

Different from structural proteins, intrinsically disordered proteins (IDPs) do not have a fixed or stable tertiary structure under normal conditions. IDPs possess inherent flexibility and usually exist as dynamic ensembles of conformations, allowing them to adapt to different binding partners or cellular environments. We have curated in total 114 entries from the protein ensemble database (PED) (Lazar et al., ) as benchmarking set. In specific, we only select the experimentally validated (eg. NMR spectroscopy) structure ensembles and excluding similar protein records with training set to avoid data leakage (see Appendix A.4).
Due to the disordered structural characterization for IDPs, alignment-based metrics such as TM-score is not applicable. Because different targets have different sizes of ensemble (from ten to thousand), we follow the metrics used in Janson et al. () to evaluate the mean absolute error (MAE), specifically the pairwise distance, radius of gyration, and contact map between the predicted ensemble and the ground-truth ensemble.

**Table 4: Mean absolute error (MAE) for the IDP test set measured in different characteristics for each method, where both mean / median are reported over all targets. All metrics are the lower the better ($\downarrow$). The contact map of PED00247 (PDB: 2MTF) is shown to the right as an example.**
$$ $f(i)=-\sum_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime})\log p_{[i]}({ \bm{z}}_{[i]}=z^{\prime}),$ (8) $$

### 5.4 Runtime analysis

To demonstrate the efficiency of SLM, we benchmark the runtime of each SLM and compare them with diffusion-based baselines. The measurement is based on the elapsed wall clock time for sampling an ensemble of $N=100$ across different protein lengths (see Appendix A.2). As shown in Fig. 5, SLMs exhibit superior scalability with respect to protein size and are $\textrm{20-100}\times$ faster than diffusion models like AlphaFlow, highlighting their potential for real-world applications.

Figure: Figure 5: Runtime profiling for SLMs and baseline methods.
Refer to caption: x7.png

## 6 Conclusion and Limitations

In this work, we propose a novel conformation generation framework of structure language models (SLMs). The overall task is divided into two stages: conditional sampling of latent structure tokens and roto-translation invariant structure decoding. We develop and train a variety of conditional language models, especially the masked diffusion-based ESMDiff, which is a fine-tuned variant to enhance the capabilities of ESM3 adapting for conformation generation.
Unlike existing methods, SLMs perform amortized distribution learning within an invariant latent space, leading to more efficient inference. By alleviating the need for geometric modeling, SLMs can fully exploit the scalability of modern language model architectures and take advantage of advanced hardware optimizations.
Benchmarking results across various conformation generation tasks demonstrate the compelling performance and application potential of SLMs. In summary, the proposed method opens up an intriguing and novel research direction for related communities to explore.

Limitations.
The current study presents several limitations worth exploring for future works. Firstly, one can design more advanced dVAE architecture to balance between structure disentanglement and reconstruction fidelity.
In addition to the discrete latent space, continuous latent space can also be considered, for example using the latent diffusion models (Rombach et al., ). Secondly, it is worthwhile exploring alternative SLM instances specially tailored for the sequence-to-structure translation in consideration of proper inductive biases.
Lastly, the outstanding performance of MSA-based methods also indicates the potential to build MSA-conditioned structure language models.

## Appendix A Implementation details

### A.1 Baseline evaluations

For baselines, we compare our structure language models against: EigenFold (Jing et al., ), Str2Str (Lu et al., ), MSA subsampling (Del Alamo et al., ), AlphaFlow & ESMFlow (Jing et al., ), and idpGAN (Janson et al., ).
For MSA subsampling, we leverage the official repository of Del Alamo et al. () under AlphaFold v2.3.2 with an MSA depth of 32, i.e. 32 maximum extra MSAs and 16 maximum MSA clusters (half the depth value (Del Alamo et al., )).
For AlphaFlow and ESMFlow, we use the base PDB model (no distillation) with no resampling. The input MSAs for test proteins are queried via the ColabFold server (Mirdita et al., ). All other parameters are set to the recommended values according to the original implementation by the authors (Jing et al., ).
For EigenFold, we follow the pipeline in Jing et al. () by using OmegaFold (Wu et al., ) to make embeddings for the test set and sampling via the pre-trained diffusion model with default parameters, where $\alpha=1$ and $\beta=3$.
For Str2Str, we use their public codebase to with initial structure guess generated by the ESMFold (Lin et al., ) using the script provided in the repository. We use a $T=1000$ diffusion sampling steps for both probability flow and SDE (noise scale is set to $1.0$ for SDE runs) and other parameters as default.
For idpGAN, we run the inference pipeline following the Jupyter notebook
in their open-source repository with the default inference settings for custom proteins.

### A.2 Runtime Profiling

To profile the inference time for each method, we set a ladder of tasks of different protein lengths including 50, 100, 200, and 500 residues, with an ensemble size of $N=100$ for each benchmark task. The profiling is carried out on a single NVIDIA A100 SXM4 GPU with 40GB memory, we measure the wall clock time elapsed for generating the target ensembles for each model. Note that for the MSA-based methods, the time-consuming MSA search is conducted only once and is amortized in the overall computational cost for multiple batches. For autoregressive SLMs namely S-T5 and S-GPT, the past computed key-value pairs for each layer are cached for fast inference.

### A.3 Training details

#### S-T5 training.

For S-T5, we build upon the implementation from the Transformers package (Wolf et al., ). Table S1 shows the hyperparameter settings for training. The total number of trainable parameters is 384M. The model is trained without learning rate scheduler for up to 30 epochs. The final tested model and the reported epoch number in Table S1 come from the best checkpoint, selected according to the NLL of structure tokens on a hold-out validation set. To leverage the ESM3 sequence embeddings, we bypass the encoder’s input token embedding layer and feed the precomputed embeddings directly into the transformer encoder.

#### S-GPT training.

For S-GPT, we also use the implementation from the Transformers package. The total number of trainable parameters is 961M. We insert a special [sep]  token in between the sequence and structure tokens when they are fed into the model, with its embedding being randomly initialized and learnable. In practice, we find it useful to add a hyperparameter $\lambda$ to balance between the NLL of sequence and structure prediction, i.e.,
${\mathcal{L}}(\theta)=-\mathbb{E}_{({\bm{c}},{\bm{x}})\sim{\mathcal{D}}}
\mathbb{E}_{{\bm{z}}\sim q({\bm{z}}|{\bm{x}})}\left[\lambda\sum_{l=1}^{L}\log p
_{\theta}({\bm{y}}_{l}|{\bm{y}}_{<l})+\sum_{l=L}^{2L}\log p_{\theta}({\bm{y}}_
{l}|{\bm{y}}_{<l})\right]$, where the first half of ${\bm{y}}$ consists of the sequence tokens and the second half consists of the structure tokens, as defined in Sec. 3.2.
Table S2 shows the hyperparameter settings for training. The model is trained and up to 30 epochs but the best checkpoint is evaluated and reported.

Table: Table S1: Hyperparameters of S-T5

Table: Table S2: Hyperparameters of S-GPT

#### Masked Diffusion fine-tuning of ESM3.

The masked diffusion is scheduled with a log linear noise schedule (Austin et al., ) for clipped time domain $t\in[0,1-\epsilon]$ with infinite small time $\epsilon=0.001$, such that $\sigma(t)\triangleq-\log(1-t)$ while the coefficient is parameterized as $\alpha(t)=\exp[-\sigma(t)]\in(0,1]$. We leverage this schedule for both training and inference. When fine-tuning, we replace the pre-trained LM head in the structure track of the ESM3 with a randomly initialized head for the augmented vocabulary $\bar{V}=V\cup\textrm{[MASK] }$ by setting output dimension to be 4101. We use a learned embedding module to encode timestep $t$ via the sinusoidal time embeddings followed a simple 3-layer MLP coupled with SiLU activation function.
We use the AdamW optimizer with betas=(0.9, 0.999) and weight decay=0.01. The training is scheduled with the constant scheduler with warm up steps of 2,500.

Figure: Algorithm 2 Masked diffusion fine-tuning of ESM3

### A.4 Curation of IDP test set

We curated the test set for intrinsically disordered proteins (IDPs) by downloading data from the Protein Ensemble Database (PED) (Lazar et al., ) on August 10, 2024, using the official API, which provided a total of 481 raw entries. To prepare the evaluation set, we filtered the data to include proteins with sequence lengths between 20 and 500 and ensemble sizes of $N\geq 10$. To ensure balance in the dataset, we performed clustering using MMseqs 2 (Steinegger & Söding, ) with the flags -s 7.5 --min-seq-id 9.0. We further excluded the records with more than 90% sequence identity by conducting another search against the PDB training data to avoid data leakage. Finally, we select the experimentally validated (eg., NMR) ensembles by filtering with respect to the measurement method, which excludes computationally generated ensemble by models such as idpGAN (Janson et al., ).

Figure: Algorithm 3 Minibatch Gibbs Sampling with Positional Ranking

## Appendix B Inference of Bidirectional SLMs

### B.1 Zero-shot conformation generation of ESM3

The sampling method for the BERT model was first introduced by Wang & Cho (), which treats the bidirectional encoder as a Markov random field graphical model and employs Gibbs sampling to generate new sequences. In Hayes et al. (), this Gibbs sampling approach is extended to enable simultaneous decoding of multiple positions rather than processing them one at a time. We refer to this latter method as the minibatch Gibbs sampling and take advantages of this sampling strategy for zero-shot conformation generation.

In specific, we leveraged the open-source 1.4B ESM3 protein language model (Hayes et al., ) for zero-shot conformation sampling. ESM3 is pre-trained on multi-modal data related to protein and contains several tracks of language model heads for prediction as well as input embeddings, including sequence (amino acid types), structure, secondary structure, SASA and functions. In our main experiment, only the sequence and structure tracks are turned on and the input tokens for other tracks are set to the default value (i.e. [MASK] in each vocabulary). For sampling configurations, we use a sampling schedule of 25 steps and set the sampling temperature to be 1.0 without specially pointing out. During sampling, we also adopted the nucleus (top-p) sampling strategy (Holtzman et al., ) to improve the quality of samples with a probability threshold of 0.95. Throughout the experimental Section 5, we use the entropy-ranked minibatch Gibbs sampling to obtain the results for ESM3 (described below).

### B.2 Position-Ranked Gibbs Sampling

Unlike one-shot decoding, which maps all masked structure tokens [MASK] to their predictions in a single feed-forward pass, the minibatch Gibbs sampling can adopt various strategies to iteratively select where to unmask at each sampling step from the masked structure tokens of the target protein. These strategies are governed by different ranking functions $f(\cdot)$ applied over the masked positions, as detailed below. The overall inference pipeline (with ranking) is outlined in Algorithm 3. We investigate and compare these ranking strategies in Table B.2, using ESM3 as the base model.

Entropy ranking. The entropy for categorical distribution quantifies the predicted uncertainty for the current evaluation of probability. Formally, the ranking score of the masked position indexed by $i~{}(1<i<N)$ is the position-specific entropy:

$$ $f(i)=-\sum_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime})\log p_{[i]}({ \bm{z}}_{[i]}=z^{\prime}),$ (8) $$

where $p_{[i]}=\mathrm{SoftMax(logits_{[i]}))}$ is the predicted probability mass function of position $i$ over the pre-defined vocabulary of valid structure tokens $V$. The entropy of each candidate position is ranked ascending, i.e. positions with smaller entropy are firstly decoded.
We also adopt the adaptive entropy calculation, which means the logits prediction in the current decoding step is based on the unmasked structure tokens from the last step and calculated on-the-fly.

Maxlogit ranking. Beside the uncertainty-based ranking, we similarly consider the maximum of logits (maxlogit) from prediction as the ranking score. This indicates we choose the positions with the top (highest) logit score as candidates. The maxlogit reflects the maximal confidence of the model among its choice as well as the probability of the selected token during greedy decoding. Correspondingly, the ranking score of maxlogit for index $i$ is defined as:

$$ $f(i)=\max_{z^{\prime}\in V}p_{[i]}({\bm{z}}_{[i]}=z^{\prime}).$ (9) $$

Secondary structure ranking. The ranking scores discussed above can be applied to arbitrary data for masked decoding. Here, we introduce a novel decoding schedule designed for proteins, ordered by predicted secondary structure:

$$ $f(i)=-g(s^{*}_{[i]}),\mathrm{where\;}s^{*}_{[i]}=\operatorname*{arg\,max}_{{s^ {*}}\in{\mathcal{S}}}p_{ss}(s_{*}|{\bm{c}},{\bm{z}},i),$ (10) $$

where $p_{ss}$ is a predicted categorical distribution over all eight secondary structures (SS8) defined by the Dictionary of Protein Secondary Structure (DSSP), including helix of 3-turn, 4-turn, 5-turn, and hydrogen bonded turn; extended strand of $\beta$-sheet and isolated $\beta$-bridge; bend and coil, and $g(\cdot)\in\{0,\dots,7\}$ maps the SS8 type to the finite ordering score. Given the predicted distribution over SS8, we use greedy selection to assign the SS8 label (with the max logit) for all masked positions. The positions are further ranked from the most structural to the most disordered, or the same order as above. The intuition behind is that we want to predict the structural region in the first place, followed by modeling the disordered region (loop). In practice, we adopt the secondary structure prediction heads of pre-trained ESM3 (Hayes et al., ) for the prediction of SS8 types for each masked candidate positions.

Table: Table S3: The performance of different ranking functions on the three benchmarking tasks described in Appendix B.2.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: * [major] While the experiments compare a diverse set of baseline models, they lack a comparison with pure language models. Previous work [a] demonstrated that language models can generate molecular and protein structures with atomic coordinates represented as regular text, and other research [b] has shown that this approach is viable for molecular conformation space modeling. Given the context sizes of modern language models, it seems feasible to model protein backbone coordinates directly as text. Including a comparison with pure language models could provide a strong justification for using dVAE in this framework.
* [minor] In reference to lines 175–186, the training of the dVAE and the language model prior happens in two separate stages. Previous research on learnable priors in VAEs [c, d] showed that the encoder, decoder, and prior can be trained simultaneously. It would be helpful to clarify, either theoretically or through an ablation study, why this work uses a multi-stage training approach rather than a joint training method, as this separation introduces additional complexity to the training process.
* [minor] The scope of this work is limited to modeling the coordinates of the protein backbone (lines 145–150). However, accurate spatial positioning of side chains is essential for many practical applications, especially when modeling interactions between the protein and other proteins or molecular structures is required [e]. Including a discussion of this limitation in the Limitations section would strengthen the paper by addressing the potential impact of excluding side chains on the model’s practical applications.

a. Language models can generate molecules, materials, and protein binding sites directly in three dimensions as XYZ, CIF, and PDB files, 2023

b. BindGPT: A Scalable Framework for 3D Molecular Design via Language Modeling and Reinforcement Learning, 2024

c. VAE with a VampPrior, 2017

d. A Prior of a Googol Gaussians: a Tensor Ring Induced Prior for Generative Models, 2019

e. Rotamer Density Estimator is an Unsupervised Learner of the Effect of Mutations on Protein-Protein Interaction, 2023
soundness: 3
strengths: The paper is well-structured, with a clear and comprehensive presentation of the theoretical background and framework. Although flexible learnable priors in VAEs have a substantial history of research, the idea of coupling the dVAE and language models as flexible conditional prior is fresh and novel in the area of protein conformation modeling. The selected case studies are relevant and show the potential for wide practical application.
summary: In this paper, the authors present a novel Structure Language Modeling (SLM) framework that integrates recent advances in discrete variational autoencoders (dVAEs) to quantise the latent space and language models to learn the conditional distribution on this latent space with masked diffusion models framework. This combined approach is applied to model the conformation space of proteins. The proposed framework offers flexibility in selecting the dVAE and language model components; however, the authors introduce a specific setup, named ESMDiff, which utilizes ESM3 and BERT as its primary components. To demonstrate the effectiveness of this approach, the authors evaluate ESMDiff on several case studies, including the structural dynamics of bovine pancreatic trypsin inhibitor (BPTI), modeling conformational changes in structural proteins, and exploring the conformation space of intrinsically disordered proteins (IDPs).
weaknesses: While the overall quality of the paper is high, there are a few weaknesses and areas that could benefit from further clarification. Details on these points are provided in the Questions section.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See above
soundness: 3
strengths: * The problem addressed by the authors is important and timely. 
* The approach, although based on accessible and popular ideas, has not been explored in the literature and is likely to be of great interest to the community.
* The paper is thorough and has detailed appendices describing training, inference, and evaluation.
* The paper has a diversity of experimental results, including several that appear only in the appendices (?)
summary: This paper proposes to use protein language models, operating on structure tokens, to model protein conformation ensembles. Several variations on this idea are tried: (1) an encoder-decoder model that translates sequence-tokens to structure-tokens (2) a decoder only model supplied with sequence-token prompts (3) Gibbs sampling with the bidirectional unmasking model ESM3 (4) Gibbs sampling and ancestral sampling with ESMDiff, a version of ESM3 fine-tuned under masked diffusion. These models are evaluated on the BPTI benchmark and conformation pair datasets and compared against several existing methods.
weaknesses: **Novelty**
* The paper's approach is relatively straightforward and is almost an off-the-shelf application of ESM3 or any structural-token protein language model. With that said, I don't think this is a big drawback because the focus of the paper is on the novelty of application and empirical benchmarks.

**Quality**
The paper's evaluations are, however, relatively limited and could be improved in quality and rigor.
* The evaluations lean heavily on the BPTI benchmark, and do not make use of MD datasets like ATLAS and MD-CATH. It would be much more interesting to see how well the model can learn from those datasets and generalize to a larger number of new proteins.
* ESM3, out-of-the-box, can already be used as a sequence->structure generative model with iterative decoding. However, the authors seem to only evaluate ESM3 under Gibbs sampling. Thus a proper comparison with pretrained ESM3 is missing.
* The results themselves are also relatively mixed, with zero-shot ESM3 outperforming ESMDiff in several instances.

**Clarity**
* In all tables, the authors have bolded the best method _among the several methods they propose_, not the best method overall. This gives a very misleading impression and is a major issue that must be fixed.
* At times, the paper has a forced and unnecessary level of formalism. At best, this obscures what the model is actually doing, and at worst leads to misleading or incorrect statements.
    * The authors spend quite a bit of time building up a probabilistic framework for structure based on a standard variational lower bound. However, the authors universally use the ESM3 tokenizer, which is not trained with maximum-likelihood reconstruction. Indeed, "we begin by maximizing the ELBO $\mathcal{L}(\phi, \theta)$ with respect to the encoder $\psi$ and decoder $\phi$" as this is soon undercut by "We start with the pre-trained dVAE established in Hayes et al. (2024) as the structure tokenizer (frozen)." So are $psi$, $\phi$, ever trained, or not? Furthermore, the ESM3 decoder is not conditioned on sequence tokens, unlike the formally defined decoder $p_\phi(x \mid c, z)$.
    * The exposition on general discrete diffusion takes up quite a lot of space, but ultimately the authors use masked diffusion, which could be presented much more clearly and succinctly on its own.
    * Since ESM3 is also trained at all mask levels, it is not immediately clear what is different about the fine-tuning model. As I can currently tell, it is only that the model is explicitly time-conditioned and the unmasking is stochastic instead of a fixed number of tokens unmasked per step. But is this really a significant enough difference to merit a completely different presentation?

To summarize, I like to paper's core idea and would recommend acceptance based on that promise. However, the paper's presentation obfuscates the simplicity of the method, and the evaluations are incompletely and misleadingly presented. It would be preferable for the authors to de-emphasize methodological novelty but focus on a thorough and convincing empirical study of the ability of structural PLMs to sample conformational ensembles.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: In Table 2, the best SLM variant in each cluster is all different. Can authors provide more insights about these results?
soundness: 3
strengths: -  The paper is well motivated.
-  The framework is a meaningful combination of VAE and LM that take the benefits of both.
-  The 20x-100x speed up seems to be promising.
summary: This paper introduces the Structure Language Modeling (SLM) framework for protein conformation generation. The frameworks consist of an auto-encoder and a conditional language model, which can be BERT-like or GPT-like. This method gets a good speedup than previous methods.
weaknesses: In results tables, seems all systems (including baselines and SLMs) have different model sizes and pre-training data. It would be better if these details are clearly described in the table to better understand the significance of SLM.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see "Weaknesses".
soundness: 3
strengths: - The proposed method offers an interesting alternative to the task of protein conformation sampling, a task that has so far mostly been tackled with 3D-aware generative models that may suffer from poor computational efficiency. The proposed method naturally leverages pre trained ESM3 model and extends its capablities.

- The experiments validate the efficiency of the proposed method. The performance seems competitive with SOTA structure-based generative models.
summary: This paper aims to generate protein conformation ensembles with Structure Language Modeling (SLM), a method that first tokenizes protein 3D structure through learned autoencoding with a VQVAE, and then generates the 3D structure through language modeling and decoding. Compared to alternative methods such as equivariant diffusion models, the proposed method enjoys a significant efficiency advantage. The proposed method is applied to protein conformation ensemble generation tasks with a diverse set of evaluation metrics.
weaknesses: - There are many metrics being benchmarked and it is hard to evaluate the model performance. It seems that the proposed method is quite competitive with SOTA alternatives. However, this reviewer is not fully convinced about that because the results are so scattered around in different tasks and different evaluation metrics, and the proposed method sometimes works better and sometimes works worse compared to baseline methods. Can the authors explain what metrics are the most practically relevant, or propose consolidated metrics for comparing different models? Can the authors include benchmarks that were originally used in the AlphaFlow/ESMFLow paper (Jing et al, 2024), if they are solving the same task?

Meta-review: The submission presents a protein conformation generative model based on the VAE framework, where the latent space is discrete and sequential, on which the prior is flexibly modeled by a language model that can properly handle a given protein sequence. The encoder and decoder have properly taken geometric invariance into consideration. Although using a language model, the proposed formulation can still capture the spatial interactions mutually between every pair of residues, and the discretization does not lose the spatial resolution.

I agree with the reviewers' appreciation on the soundness, good presentation, and the promising future potential. Reviewers also pointed out a few insufficiencies, including lack of comparison with a pure language model, limited evaluation cases, and some mismatches between the introduction of a general framework and the actually employed method (I don't think the former is unnecessary, though). In the rebuttal, the authors updated the presentation, and presented some results on nanobody and ATLAS, which have addressed some of the limitations, while some remains (e.g. limited advantageous results on ATLAS, alignment and presentation in metric, model, and dataset settings for comparisons with other methods). Overall, I agree with the reviewers' opinion to appreciate the contributions over the remaining limitations, hence recommend an accept, and hope the authors improve the manuscript further according to the feedback.

Decision: Accept (Poster)