ProteinAE: Protein Diffusion Autoencoders for Structure Encoding
Authors: Shaoning Li
Le Zhuo
Yusong Wang
Mingyu Li
Xinheng He
Fandi Wu
Hongsheng Li
Pheng-Ann Heng
Venue: ICLR 2026 Poster (2026)
Keywords: Protein Auto-encoder; Protein Structure modeling
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Developing effective representations of protein structures is essential for advancing protein science, particularly for protein generative modeling. Current approaches often grapple with the complexities of the $\operatorname{SE}(3)$ manifold, rely on discrete tokenization, or the need for multiple training objectives, all of which can hinder the model optimization and generalization. We introduce ProteinAE, a novel and streamlined protein diffusion autoencoder designed to overcome these challenges by directly mapping protein backbone coordinates from $\operatorname{E}(3)$ into a continuous, compact latent space. ProteinAE employs a non-equivariant Diffusion Transformer with a bottleneck design for efficient compression and is trained end-to-end with a single flow matching objective, substantially simplifying the optimization pipeline. We demonstrate that ProteinAE achieves state-of-the-art reconstruction quality, outperforming existing autoencoders. The resulting latent space serves as a powerful foundation for a latent diffusion model that bypasses the need for explicit equivariance. This enables efficient, high-quality structure generation that is competitive with leading structure-based approaches and significantly outperforms prior latent-based methods. Code is available at https://github.com/OnlyLoveKFC/ProteinAE_v1.

Published as a conference paper at ICLR 2026
PROTEINAE: PROTEIN DIFFUSION AUTOENCODERS FOR STRUCTURE ENCODING
Shaoning L Le Zhuol.3, Yusong Wang?* Mingyu Li', Xinheng He,, Fandi Wu?, Hongsheng Li', Pheng-Ann Heng'
CUHK ZGC Academy Krea AI 4 XITU SITU AstraZeneca 7 Tencent
ABSTRACT
Developing effective representations protein structures is essential for advanc- ing protein science, particularly for protein generative modeling. Current ap- proaches often grapple with the complexities of the SE(3) manifold, rely on dis- crete tokenization, the need for multiple training objectives, all of which hinder the model optimization and generalization. We introduce PROTEINAF, a novel and streamlined protein diffusion autoencoder designed overcome these challenges by directly mapping protein backbone coordinates from E(3) into a continuous, compact latent space. PROTEINAE employs non-equivariant Diffu- sion Transformer with a bottleneck design for efficient compression and is trained end-to-end with a single flow matching objective, substantially simplifying the optimization pipeline. We demonstrate that PROTEINAE achieves state-of-the- art reconstruction quality, outperforming existing autoencoders. The resulting latent space serves as powerful foundation for a latent diffusion model that bypasses the need for explicit equivariance. This enables efficient, high-quality structure generation that is competitive with leading structure- based approaches and significantly outperforms prior latent- methods. Code available at https///ghhnb...ommmmmyyy/xxtttt
INTRODUCTION
Proteins serve as the foundation of life. Protein structures, determined by the folding of various amino acid sequences, dictate their biological functions and behaviors. Understanding and repre- senting these structures is important for various protein tasks, typically protein generative modeling. A dominant paradigm for visual generative models involves a two-step process (Esser et al., |021 Rombach et al.) 2022): initially compressing the pixel input into a compact latent representation with visual autoencoders (tokenizers), and then performing generative modeling within this latent space. This paradigm significantly enhances both efficiency and performance of modeling complex visual distributions. In contrast to the visual domain, which are typically represented at the pixel level and are high-dimensional, protein structures are represented by continuous atom coordinates in a lower dimensional space, i.c., 3D Euclidean space (E(3)) and differ greatly in size due to their diverse amino acid sequence lengths (Jumper et al. 202 Abramson et al. 2024). These character- istics require new recipes for creating protein structure autoencoders and necessitate the exploration of more efficient protein generative modeling methods.
Researchers have made several attempts to encode protein structures with autoencoders. Pioneer- ing efforts in this area include the ESM3 VQ- VAE tokenizer (Hayes et al. 2025) and the DPLM-2 lookup-free quantization (LFQ) tokenizer (Wang et al. 2024 2025). which were among the first to convert continuous 3D coordinates into discrete tokens to jointly model sequences and structures used in generative masked language models. (Yuan et al.) (025) proposes AminoAseed with further codebook improvements. While achieving initial success, these autoencoders are not ideal in the following aspects: (i) operating on the intricate SE(3) manifold (involving translations and frame rotations) necessitates incorporating equivariance and physical constraints, thus introducing com- plexity to the latent space and the design of model architecture; (ii) discretizing continuous atom coordinates to tokens leads to a loss of reconstruction accuracy; (iii) requiring the combination of
"Equal contribution.

---

Published as a conference paper at ICLR 2026
complex training objectives (e.g.. FAPE loss, distance loss, violation loss, KL loss, etc.), which de- mands tuning individual weights for each term; (iv) being limited to fixed input size (sequence length) and lacking compact bottleneck latent space for efficient generative modeling. These raise the question: Could we design a simpler, more accurate, and effective protein autoencoder in continuous, compact latent space?
In this work, we propose PROTEINAE, a pro- tein diffusion autoencoder designed for effec- tive and efficient structure encoding and gener- ation. Specifically, PROTEINAE operates in non-equivariant manner, and conducts autoen- coding protein backbone atoms (Ca: N, O) directly on E(3), avoiding the discretization. Inspired by recent works on denoising autoen- coders (Chen al. 2025). PROTEINAE em- ploys a simple diffusion loss for training, which makes the learned structural representations maximize the ELBO of the likelihood of the in- put protein structures, which is similarly vali- dated by AlphaFold3 (Abramson et al., (024). An architecture comparison between the tradi- tional protein autoencoder ESM3 structure au- toencoder and PROTEINAE is shown Fig. ESM3 structure autoencoder employs stan- dard VQ- VAE framework (Esser et al.] 2021 with modified geometry-based transformer encoder on SE(3) and focuses more on the local structural pattern. In contrast, PRO- TEINAE's encoder and decoder are designed with scalable Diffusion Transformers (DiT) (Peebles Xie. 2023). The encoder maps the protein backbone atom coordinates to latent representation We also incorporate bottleneck design by downsampling the protein in length and channel dimensions, which is common in vision for efficiency (Ronneberger al. 2013). The decoder takes condition and predicts the clean structure or velocity from noisy structures. At inference time, protein structures are reconstructed from the noisy coordinates with latent representation 2 via a diffusion sampler.
fraite
distance
citteiion loss
Quant
noise/x;
Bottleneck
Geom. Encoder
DiT Encader
ånn
Trans
Rots
Atom coords
Prutein Frame
Protein Backhome
(a) ESM3 VQ-VAE
(b) ProteinAE
Operate SE(3)
Operate B(3)
Discrete encoding
Continuous encoding
Multiple training Joss
Simple training loss
Figure I: Comparison of ESM3 VQ-VAI and our PROTEINAE.
Building on this protein autoencoder, we can further perform downstream structural latent analysis, such as protein latent diffusion modeling (PLDM) and physicochemical prediction. Previous gener- ative methods are mostly structure-bar based diffusion models over SE(3) (Yim et al. 2223b a: Watson (et al.) 2023] Bose et al. 2023), but they face challenges with equivariance and physical constraints such bond lengths and angles (Ym et al. 2025). Protein multi-modal models like ESM3 use iter- ative decoding to generate discrete structure tokens, their generation quality not as expected.
Yim et al.) 2025) proposes a hierarchical structure latent diffusion model, but the latent diffusion is conducted on the contact map and still relies on FrameFlow (Yim et al. 2023a) at the second fine-grain stage. Based the lightweight PROTEINAE, standard DiT architecture, bypassing the need for explicit equivariance or physical constraints, and eliminates the computationally expensive triangle attention module for better efficiency.
develop scalable PLDM. It employs
We rigorously evaluate PROTEINAE on the CASP14 and CASPI5 benchmarks, demonstrating state- of-the-art reconstruction quality. We further show that the learned latent space is highly effective for physicochemical property prediction. Finally, our PLDM is competitive with leading structure- based generative models and substantially outperforms existing latent-based methods in both sample quality and efficiency. Our primary contributions are:
PROTEINAE: A Simple and Effective Protein Diffusion Autoencoder. We introduce a non-equivariant autoencoder based on Diffusion Transformers that operates directly on backbone atom coordinates in E(3). It lears continuous, compact latent representation using a single flow-matching loss, avoiding the complexities of SE(3) manifold and dis- crete tokenization.
2

---

Published as a conference paper at ICLR 2026
(a) ProteinAE Encoder
(c) Protein Latent Diftusion Modeling
Xi Refesece feturrss
Soy, ldx
Random Gaussan
All Encoder
Seq. Module
Cloud.
Regasers
Imput
DIT w attention bias
Deneiiina
18, 85,8/1,33 D)
68 8//6, 4)
Libear
LayerNoras
Rege
Lhoemampll
Dmm Divmmample
(b) ProteinAE Decoder
(PP)
All-Atom Antn. Trecnder
Pair Cond. Modille
Seg Cond, Mediule
Interpolate
Len. Upsanple
Recons tructed
Seq. Repe
Regiaes
Pair Repe p
Seq Cond, e
Lizear
Upeample
(8.N.D) D)
DIT w attention bias
Ali-Atom Alter
Decoder
Velocity of
Flow Loes
Sampling
Figure 2: Overall architecture of PROTEINAE. (a) The encoder maps a protein structure to latent representation 2; (b) The flow decoder predicts the velocity field of for structure reconstruction conditioned on z; (c) Downstream tasks like PLDM operating over the learned latent space.
High-Fidelity Encoding for Downstream Generative and Predictive Tasks. PRO- TEINAE achieves state-of-the-art protein structure reconstruction. Its learned latent space enables accurate physicochemical property prediction and serves as the basis for protein latent diffusion model (PLDM) that significantly outperforms latent-based methods and is competitive with leading structure-based generative approaches with improved efficiency.
PROTEINAE: PROTEIN DIFFUSION AUTOENCODERS
As illustrated in Fig. PROTEINAE employs an encoder-decoder architecture for learning latent protein structure representations. We explain abbreviations of some components in Fig. in italics.
Feature Preparation Both the encoder and decoder take a protein backbone structure x e N4x3 as input (z, for the encoder, representing the clean structure, and for the decoder, rep- resenting the noisy structure at timestep f; N denotes the protein length). Given the input structure, initial residue-level features condition feature e (Seq. Cond.) are generated through the Pair Condition Module (Pair Cond. Module) and Sequence Condition Module (Seq. Cond. Module), respectively. For the encoder, in- puts to these modules include xn and sequence index features. For the decoder (inputs shown in gray below), the diffusion timestep and self-condition f are added to the Pair Condition Module inputs, and the latent representation is added to the output of the Sequence Condition Module. Reference amino acid features are used by both modules. The feature preparation process can be summarized:
constructed. The pair representation p (Pair Repr.) and sequence
p PairConnMModull(a""., poon_idxx (2f
(1)
ScqCondModulle(f) +2.
(2)
Note that inputs shown in gray specific the decoder. These initial token-level features (p and c) are used alongside the atom-level input æ. lightweight All-Atom Attention Encoder (All- Atom Atin. Encoder) processes the atom-level features from generate token-level sequence
3

---

Published as a conference paper at ICLR 2026
representation (Seq. Repr.). This module also outputs skip connection features (g*kin, cskin prkin) used later in the decoder pipeline. The output of the All-Atom Attention Encoder is:
8, kip Chiip p Al
(3)
More details regarding the All-Atom Attention mechanism are provided in the next paragraph.
Diffusion Transformers and All-Atom Attention Unlike many previous structure-based diffu- sion or representation learning models that rely on geometric equivariance, PROTEINAE adopts non-equivariant architecture, forming the core of our feature processing. encoder, and decoder. This choice aligns with recent works such as AlphaFold3 (Abramson et al., 2024) and Proteina (Geffnnr er al. 2025), which feature stacks of conditioned and biased multi-head self-attention layers com- bined with transition blocks and residual connections. Operating on sequence and pair representa- tions, these layers are optionally enhanced with attention biases derived from geometric relationships in the input structures. The input tokens are further augmented with additional registers (Darcet et al. 2023). To effectively handle variable- lenghh protein inputs, employ Rotary Positional En- codings (RoPE) (Su et al. 2024) instead of traditional absolute positional encodings. The Diffusion Transformer (DiT) architecture utilized in three distinct parts of our model: (i) the All-Atom At- tention modules during feature preparation and decoding, (ii) the core token-level Attention stacks in the PROTEINAE Encoder and Decoder, and (iii) the token-level Attention in the subsequent PLDM stage. We note that pair bias is incorporated in the DiT stacks of the PROTEINAE Encoder and Decoder, as well as within the All-Atom Attention modules, but it omitted in the PLDM stage. A DIT block operating on sequence (s) and condition (c) representations, potentially incorporating pair bias (p. Biy). can be summarized as:
s' DiT bias( .p,
PpC, 98y);
(4)
s TransitionBlock(s'.o
(5)
where 1 denotes the 1-th layer in a stack of L blocks. B, here represents the attention mask. A comprehensive algorithm is provided in the Appendix A.
To explicitly model and generate atom-level details for the backbone structure, we incorporate All- Atom Attention Encoder and Decoder modules, inspired by AlphaFold3 but implemented with fewer parameters. These modules utilize sequence-local atom attention, allowing interactions among all backbone atoms within defined sequence neighborhood. This approach offers richer local inter- action modeling compared to methods relying solely on KNN graphs, such as those adopted in the ESM3 VQ- VAE pipeline (see Fig. D. The All-Atom Attention Encoder, described previously as part of feature preparation, aggregates atom-level features from the input structure into token-level rep- resentations. For consistency within the All-Atom Attention mechanisms, we fix reference features f (including positions, charge, mask, elements, and atom name characters) based on the standard Glycine (GLY) residue from the Chemical Components Dictionary (CCD). The All-Atom Attention Decoder operates after the DiT stack in the decoder pipeline. It takes the final sequence representa- tion s from the DiT stack and broadcasts loken-level features back to atom-level. After computing atom attention, it projects these features to predict the noise velocity of R x4x3 in coordinate space. The output prediction of the All-Atom Attention Decoder is summarized as:
of AllAtomAttnDecodee
gakip, skip kkip, 3.y))
Note that while the DiT stack and All-Atom Attention modules are distinct components operating in sequence, their underlying attention mechanisms share principles of conditioned and biased self- attention.
Autoencoder Bottleneck To obtain a compact latent representation and improve the efficiency of subsequent generative modeling, we introduce two types of bottleneck compression within the encoder pipeline (Fig. 2b): length bottleneck and a dimension bottleneck. Following the stack of DiT layers, we obtain the sequence representation s". For the length bottleneck, we apply one or more ID convolution layers with kernel size of 3 and a stride of 2. This downsamples the protein length from N to Naw = N/r, where r is the total downsampling ratio. Subsequently, for the dimension bottleneck, linear layer projects the sequence representation from its token dimension
4

---

Published as a conference paper at ICLR 2026
D to reduced bottleneck dimension d. This compression process can be summarized as:
LincarNoBias
Convid(stride = 2, kernel = 3) . (transpose(s*))
(7)
Dimension downsample
Length downsample
Here, Iranspose(-) denotes the batch transpose, transforming a tensor of shape N, D) to (B, DN for convolution. The resulting latent representation z has shape (B, NVown. d).
In the decoder pipeline, the latent representation needs to be upsampled and expanded to match the target protein length Naaret and channel size of conditioning features before being added (Eq 2
This upsampling process involves expanding the dimension of and then interpolating its length:
Zupsampled = transpose
Interpolate
iinaaaaaaarraaapeeoeaiaaiaaaaaan =)), Nurget
(8)
Dimension upsample
Replace KL regularization with LayerNorm Unlike traditional VAEs that employ KL regu- larization loss on the latent feature we apply LayerNorm without learnable scales following DiTo (Chen et al. 2025). This operation is applied to the output of the bottleneck (Eq. 7 and yields the final latent representation used for both the PROTEINAE decoder and PLDM training. This approach eliminates the need for KL loss weight tuning and empirically demonstrates better recon- struction performance. Furthermore, by normalizing the latent space in this manner, train the PLDM on without requiring additional normalization within the diffusion process (Rom- bach al. 2022. In conclusion, this design choice simplifies the overall training procedure for both PROTHINAE E and PLDM.
directly
Autoencoder Pipelines Having introduced the core architectural components, we now detail their arrangement within the PROTEINAE encoder and decoder pipelines (Fig.
Encoder (Fig. 2h) takes a native protein structure (x1) as input. As described "Feature prepara- tion", this structure is used to derive initial sequence (s), pair (p). and condition (c) features. These features are processed through a stack of DiT blocks incorporating pair bias. The resulting repre- sentation then undergoes bottleneck stage, which compresses both the length and dimension of the feature maps (as detailed in Autoencoder Bottleneck). Finally, the compressed latent representation is normalized via LayerNorm before being passed to the decoder used for downstream tasks.
Decoder (Fig. 2p) operates within the flow matching framework, taking a noisy protein structure (12) at timestep input. Similar to the encoder, de is used for feature preparation to obtain initial 8, p. features. key distinction is that the latent representation z (from the encoder during training, sampled during inference) is upsampled and incorporated into the sequence condition feature c. These conditioned features (s, p, and the modified c) are then processed by stack of DiT blocks with pair bias, mirroring the encoder's architecture. The output from the DiT stack is fed into an All-Atom Attention Decoder, which predicts the velocity vector of required for structure reconstruction via ODE integration (Appendix A.22)
conditioning information
Reconstruction Loss PROTEINAE is trained using simple flow-matching loss, with no other auxiliary losses. The objective to train the model to predict the velocity vector field of at noisy structure I, and timestep t, conditioned flow matching is defined as v(t) predicted velocity of are all represented as tensors in g x3, The PROTEINAE's training objective
the latent representation 2. The target velocity field in Zg. The input structures I1 Tr, the noise xo, and the
is:
1 4N
min
| (I,, (x1 xo0
where Pas (x) is the data distribution, p(t) is the timestep sampling distribution, and ef (2,, 2) is the velocity predicted by the decoder network. Following common practice in maneture-based flow matching models, we use the i-sampling distribution p(t) 0.022/(0, 1) + 0.983(1.9, 1.0).
5

---

Published as a conference paper at ICLR 2026
3 EXPERIMENT
3.1 EXPERIMENTAL. SETUP
Dataset and Model Configuration We use the AFDB-FS (Jumper et al., 2021 Lin et al. dataset for training PROTEINAE and PLDM. This is large-scale dataset of single-chain protein structures derived from the AlphaFold Protein Structure Database (AFDB) through sequential fil- tering and clustering steps utilizing sequence-based MMseqs2 (Steinegger & Söding 2017) and structure-base based Foldseek (Van Kempen et al., 2024). It contains 588,318 structures with lengths ranging from 32 256 residues. For PROTEINAE training, we apply random global rotation the protein structures as data augmentation. For structure reconstruction evaluation, we use the bench- mark sets from CASP 14 and CASP 15. We use ATLAS Vander Meersche et al. (2024), protein molecular dynamies (MD) dataset for downstream flexibility prediction. For our reconstruction and downstream results, we primarily use a default PROTEINAE configuration. The encoder and decoder DiT stacks have L = 5 layers and token dimension D 256. The bottleneck is configured with a length downsampling ratio r = 1 (i.e., no length downsampling) and compresses the dimension from D 256 to latent dimension d 8. The PLDM is also based
2024)
DiT architecture with 200M parameters, featuring I. = 15 and D = 768. Different configurations and their impact on performance explored in the ablation studies. The overall details can be found in Appendix A..4
3.2 STRUCTURE RECONSTRUCTION
Table I: Different protein autoencoders' structure reconstruction quality measured by RMSD (L).
Lower is better. Bold indicates the best performance. *Note that ProeT and DPLM-2 can only process proteins under 2,048 residues.
CASP14
CASP15
Methods
T
T-dom
oligo
TS-domains
oligo
CHEAP (u et
2024
11.15
8,9911.18
9.93111055
10.22:111.23
9.22-112.57
ESM3 VA- (Haees et al.
2025)
1.02+11822
0.66+0042
3.08-7.39
1.23+
1.94+2243
C. RMSD
ProTokken*
Lin et al. 20258
20258
0.99=2069
0,96::058
1.15=1.12
1.15#1.000
1.18:10.89
DPLM-2
Wang et al. 2024
2024
1.99+22.03
1.87+1.78
2.70=5.05
3.31+55.69
3.50=+6.23
PROTEINAE
0.23:10.15
0.22:10.11
0.21-0.22
0,28:10.20
0.37
A comprehensive benchmark analysis in Table D reveals that our proposed model exhibits clear and consistent superiority in the task of protein structure reconstruction. We give a comprehen- sive reviews of these baselines in Appendix B PROTEINAE systematically outperforms all baseline methods, including prominent vector-quantized autoencoders, across the full suite of challenging targets from the CASP14 and CASP15 assessments. This robust outperformance holds true irre- spective of the protein's structural class or complexity. The performance margin of PROTEINAE particularly pronounced when reconstructing targets of high structural complexity. For challeng- ing oligomeric assemblies, where many baseline models show a marked degradation quality, our approach consistently maintains a high degree of fidelity. This capability underscores the robust- ness our framework in capturing the intricate spatial arrangements and fine- fine-graing grained geometric details that are often lost by competing methods. Collectively, these results strongly suggest that the diffusion autoencoder framework is more effective at modeling the manifold of protein structures than discrete, vector- vector-quantized approaches. By circumventing the information bottlenecks inherent to tokenization, our model learns richer and more expressive latent representation. This leads to significantly more accurate and reliable reconstructions, establishing a new state of the art in high-fidelity protein structure encoding. It is noting that CHEAP consistently records the poorest performance; this is an anticipated outcome, it relies on encoding features from ESM2, thereby inheriting its limitations and being fundamentally capped by ESMFold's prediction accuracy.
3.3 DOWNSTREAM STRUCTURAL LATENT ANALYSIS
Unconditional Backbone Generation We assess the performance of PROTEINAE-PLDM on unconditional protein backbone generation against state-of-the-art Structure Diffusion Models (SDMs), Mudti-Modal Language Models (MLLMss, and other Latent Diffusion Models (LDMs),

---

Published as a conference paper at ICLR 2026
with results presented in Table 2] The details of PROTEINAE-PLDY, methods are illustrated in Appendix A.3] including the flow training objectives and sampling within the PROTEINAE'S latent space. The evaluation demonstrates that PROTEINAE-PLDM achieves sfate-of-the performance among latent-space generative models. Our approach outperforms those MILM and LDM base- lines, achieving high designability and diversity, where these prior latent-space methods have tra- ditionally struggled. Furthermore, PROTEINAE-PLDMM' performance extends beyond dominating its own category rival that of classical SDMs, effectively closing the significant performance gap that has long separated latent- and structure- -baee diffusion models. PROTHINAE-PLDM can also provide a controllable trade-off with sampling temperature y, where a higher temperature setting modestly decreases designability in exchange for notable increase in structural diversity. Visual case studies (Figure 3D, purple indicates helix and yellow indicates sheet) further reveal that sam- ples with high designability converge to common folds, while those with lower designability often represent novel yet physically plausible conformations. We also observe that the model's tendency to generate varied loop regions, reminiscent of AFDB structures (according to Huguet et al. (2024); Gefiner et al. (2025), , AFDB's s overall designability is 33%, far from 100% designable), may con- tribute to these novel topologies and explain the small remaining designability gap SDMs focused on canonical folds.
Table 2: Unconditional protein backbone generation performance among SDM, MLLM and LDM approaches. Baseline results are adopted from Yim al. (2025). Appendix[C]
Metric details are illustrated in
Type
Methods
Des (1)
Div (t)
DPT (4)
Nov (L)
RFdiffusion Watson et al. 2023
96%
247
0.43
0.71
ProteinSGM LLe cl al. 20033
49%
122
0.37
0.51
SDM
FrameFlow PDB (Ym et al.
2023a
91%
278
0,48
0.65
FrameFlow AFDB
23%
54
0.42
0.70
ESM3 (Hayes al.
2025
61%
127
0.37
0.84
MLLM
DPLM-2 650M Wang et al.
2024
63%
130
0.37
0.72
semi-LDM
LSD
Yim el al.
2025
69%
203
0.46
0.74
LatentDiff (Fu et 2024
17%
34
0.51
0.73
LDM
PROTEINAE-PLDM 7 = 0.35
93%
204
0.36
0.70
PROTEINAE-PLLMM 7 0.5
86%
228
0.35
0.66
Generation Efficiency We evaluate the generation efficiency of PROTEINAE-PLDD against prominent structure diffusion models, RFDiffusion (Watson et al.) 2023) and multi-modal protein language model DPLM-2 650M Wang al. (024). The comparison is based on the average sam- pling time and GPU memory required to generate backbones 200 residues with a batch size of 5 on single 80G A100 GPU. As illustrated in Fig. BB, PROTEINAE-PLDA demonstrates substantially higher efficiency than both baselines. It achieves the lowest sampling time (~1.6 seconds) while consuming the least GPU memory (~0.3 GB). RFDiffusion is the most computationally demand- ing, requiring 15 seconds and ~5 GB of memory. DPLM-2 exhibits intermediate performance, with a sampling time of ~3 seconds and memory usage of ~1 GB. The remarkable efficiency of AATTEINAE-PLDM. is primarily attributed to its dimension bottleneck design and the elimination of triangular attention. This allows the PLDM to operate entirely within compact, low-dimensional latent space, bypassing the complex geometric or physical constraints inherent to direct structure generation and thereby drastically reducing the computational burden.
Effectiveness Analysis To evaluate the downstream effectiveness of the latent space learned by PROTHINAF, we benchmarked its performance on physicochemical property prediction. Following the setup from StrucTTokenBench (Yuan et al., 2025), we predicted residue-level structural flexibility (RMSF and B-factor) on the ATLAS dataset. The training details are illustrated in Appendix [A.4| As shown in Table 3 PROTEINAE achieves the highest scores across all evaluated tasks and splits, demonstrating a clear advantage over the baseline models. the task of predicting FlexRMSF, PROTEINAE attains the highest Spearman's 's correlation
both fold (close level) and superfam- ily (more remote) splits, improving upon ESM3 by over 10%. A similar trend is observed in the BlextF prediction, where PROTEINAE again leads in performance across both splits. This
7

---

Published as a conference paper at ICLR 2026
Table 3: Physicochemical Property Prediction (Spearman's p%) ATLAS. Baseline results adopted from StructTokenBench (Yuan et al. 2025).
Model
Task
Split
FoldSeck
ProTokens
ESM3
Vanilla VQ
AminoAsced
PROTEINAF
Fold SupFam
15.35
13.81
44.53
44.22
44.63
45.36
FlxxRMSF
11.99
7.62
39.68
39.08
40,99
44.71
Fold
4.17
6.67
23.60
22.32
21.30
34.87
FlexBFactor
SupPam
6.97
5.47
25.80
23.73
21.76
26.54
indicates that the continuous representations learned by PROTEINAE have effectively captured the generalizable principles governing protein geometry and dynamics.
A
B
Dimcosion Bottleneck
Length
2o
DQ VAE
1.3321.50
12te11
2.321
ESM
03225333
0.3c0033
0222:013
0.5
03
RMSD: 2,43A
RMSD: L.18A
d-32
0-156
(-2
Scalahility
Registers compression
20
ProteinAL
0.5000.56
9-11
15.00..00
c
10
Under 2584A
Over 256M
RMSD:
D
Hiod Desimbbliyy
Samping time (203
3,50
GPU Hemery (00)
3163
Low Designaaiity
IPON 073R22 hoterRE--LLDM
Figure 3: Summary PROTEINAE reconstruction, generation, and model architecture anal- ysis. (A) Visual comparison of protein structure reconstruction quality between PROTEINAE and ESM3 VQ-VAE. (B) Ablation studies on key architectural components: the impact of dimension and length bottlenecks, evaluation of model scalability with increased parameters, and analysis of using DiT registers for latent compression. (C) Generation efficiency comparison PROTEINAE.
PLDM, RFDiffusion, and DPLM-2, including average sampling time and GPU memory usage. (D) Visual examples of protein backbone structures generated unconditionally by PROTEINAE-PLDMM,
3.4 ABLATION STUDY
In this section, we explore the influences of different designs of PROTEINAE. Specifically, we con- duct experiments investigate the impact of the protein length and dimension bottlenecks, evaluate
8

---

Published as a conference paper at ICLR 2026
the model scalability, and analyze the effectiveness of using registers alternative strategy for latent representation. All the results are conducted on CASPI5 TS-domains
Protein Length & Dimension Bottleneck We depict the reconstruction quality of protein length
and dimension bottleneck in Fig. B We choose 3 downsampling ratios r 1, 2, 4, where r = 8, 32, 256,
signifies no length downsampling. We also investigate 3 bottleneck dimensions where d = 256 corresponds to no dimension reduction relative to the full feature dimension. As shown in Fig. EB, both length and dimension compression impact reconstruction accuracy. Increas- ing the dimension bottleneck (decreasing d) generally leads to moderate increase in RMSD, as less information is preserved in the latent space. However, increasing the length downsampling ra- tio (increasing r) results in significantly larger degradation in reconstruction quality, with RMSD increasing substantially as r goes from I to 4. This indicates that preserving the sequential length dimension is more critical for accurate protein backbone reconstruction than maintaining a large feature dimension.
Scalability We study the scalability of PROTHINAE on two variants: PROTEINAF-Base (~20M parameters) and PROTEINAE-Large (~100M parameters). The results, showing reconstruction quality (RMSD) for these variants under different bottleneck configurations (r d 32 and 2, = 32), are presented in the "Scalability" plot in Fig. BB. As the model size increases from Base to Large, we observe slight improvement reconstruction quality (decrease RMSD) for the configuration with length downsampling (r 2, d = 32). For the configuration without length downsampling (r = 1, = 32), the RMSD remains low and stable across both model sizes. This indicates that PROTEINAE exhibits positive scalability, where increasing model capacity leads to comparable or slightly improved performance, particularly noticeable when the task is more chal- lenging due to bottlenecking (e.g., with r = 2). The plot also reinforces that maintaining the original protein length (r = 1) is beneficial for reconstruction accuracy regardless of the model size.
Conditional Decoding using Registers Beyond the bottleneck compression strategy for obtaining the latent representation 2, we also explore using the learnable register tokens to compress the struc- tures to fixed collection. The registers are concatenated participate in attention computations (as mentioned in Section 2 Recent works, such as FlexTok (Bachmann et al. 2025), have demonstrated that these tokens can effectively serve as compact latent representations for images. Inspired by this, we investigate variant, termed AROTEINAE-Regian where the registers from the encoder (z**) are directly used as the registers in the decoder, enabling conditional generation based on these tokens instead of the bottlenecked latent z. We evaluate the reconstruction quality of PROTEINAE A-RRgister, with results depicted in Fig. 3B. The plot shows that PROTEINAE-Regin Register achieves good reconstruction quality for protein lengths up to 256 residues, which corresponds to the maximum length in our training dataset. However, performance degrades dramatically for structures exceeding this training length limit. Fig. 3 also visualizes an example of this failure mode, showing the reconstruction breaking down abruptly around residue 231 for a longer protein. These results highlight a limitation of directly applying a register-based compression strategy to variable-length protein sequences, The success of registers as latent representations in vision tasks may be partly attributable to the fixed size of image tokens after patchification, a char- acteristic not present in native protein sequences of arbitrary length. This suggests that for protein structures, alternative compression mechanisms like length and dimension bottlenecks are more ro- bust to variable input sizes, particularly for handling lengths beyond the training distribution. We will also explore better way to encode protein structures into fixed-length vectors in the future.
the input sequence representation and
CONCLUSION AND LIMITATIONS
We introduced PROTEINAE, novel protein diffusion autoencoder that maps protein backbone structures into a continuous, compact latent space using a non-equivariant Diffusion Transformer ar- chitecture and a simple flow matching objective. avoids the operations on intricate SE(3) and dis- crete representations inherent in existing methods. Building on this learned latent space, we further developed PLDM for structure generation. Our experiments demonstrate that PROTEINAE achieves high reconstruction quality and enables efficient protein structure generation. Despite its capabili- ties, PROTEINAE still has several limitations that warrant future investigation. For instance, PRO- TEINAF is currently limited modeling protein monomers and cannot model other biomolecules

---

Published as a conference paper at ICLR 2026
such as ligands, DNA, RNA. Besides, the generative performance of PLDM is not yet signif- icantly outperforming state-of-ther art structure- generative models. Finally, PLDM exhibits challenges related to sequence length handling, which can lead to structural collapse unrealistic geometries for certain residues in the generated outputs. We plan to address these limitations and explore extensions PROTRINAE in future work.
REPRODUCIBILITY STATEMENT
We demonstrate the dataset and model configuration in Sec. 3.1 the model architecture details in Sec. 2 and Appendix A.1 the training details of different tasks in Appendix A.4
THE USE OF LARGE LANGUAGE MODELS (LLMS)
In the course writing this paper. a large language model (LLM) was used as a tool to refine the language and grammar of the background section. We meticulously reviewed and revised all content to take full responsibility for the final text. Additionally, we consulted an LLM for the implementation of certain metries. The underlying code for these metries was carefully examined and validated to ensure its correctness and appropriateness for our study.
ACKNOWLEDGMENTS
The work described in this paper was supported in part by the Research Grants Council the Hong Kong Special Administrative Region, China, under Project T45-401/22-N. We thank Prof. Yaosen Min for valuable discussion.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 63((8016):4935500 2024.
Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijndent Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303.0879) 2023.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 14:117981-177 2021.
Roman Bachmann, Jesse Allardice, David Mizrahi, Enrico Fini, Oguzhan Fatih Kar, Elmira Amir- loo, Alaaeldin El-Nouby, Amir Zamir, and Afshin Dehghan. Flextok: Resampling images into Id token sequences of flexible length. arXiv preprint arXiv:2502.1 13677, 2025.
Avishek Joey Bose, Tara Akhound-Sadegh, Guillaume Huguet, Kilian Fatras, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexan- der Tong. Se (3)-stocasttic flow matching for protein backbone generation. arXiv preprint arXiv:2310.0239// 2023.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.0499)) 2024.
Bo Chen, Zhilei Bei, Xingyi Cheng, Pan Li, Jie Tang, and Le Song. Msagpt: Neural prompting protein structure prediction via msa generative pre-training. Advances in Neural Information Processing Systems, 37:37504-3777 37534, 2024.
Ricky TQ Chen and Yaron Lipman. Riemannian flow matching general geometries. arXiv e- prints, pp. XiX2222302, 2023.
Yinbo Chen, Rohit Girdhar, Xiaolong Wang, Sai Saketh Rambhatla, and Ishan Misra. Diffusion autoencoders scalable image tokenizers. arXiv preprint arXiv:2501. 18593, 2025.
10

---

Published as a conference paper at ICLR 2026
Chaoran Cheng, Jiahan Li, Jian Peng, and Ge Liu. Categorical flow matching statistical mani- folds. arXiv preprint arXiv:240S. 16441, 2024.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. arXiv preprint 0rXiv:2309.16558 2023.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning- based protein sequence design using proteinmpnn. Science, 378(6615):49-56 2022.
Nicki Skafte Detlefsen, Søren Hauberg, and Wouter Boomsma. Learning meaningful representations of protein sequences. Nature communications, 13(1):1914, 2022.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference computer vision and pattern recogni- tion, pp. 12873- 12883, 2021.
Cong Fu, Keqiang Yan, Limei Wang, Wing Yee Au, Michael Curtis McThrow, Tao Komikado, Koji Maruhashi, Kanji Uchino, Xiaoning Qian, and Shuiwang Ji. latent diffusion model for protein structure generation. In Learning on Graphs Conference, pp. 29-1. PMIR, 2024.
Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, et al. Proteina: Scaling flow-based protein structure generative models. arXiv preprint arXiv:2503.00711 2025.
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q Tran, Jonathan Deaton, Marius Wiggert, et al. Simulating 500 million years of evolution with language model. Science, pp. eads0018, 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-683 2020.
Guillaume Huguet, James Vuckovic, Kilian Fatras, Eric Thibodeau-Laufe Pablo Lemos, Riashat Tara Akhound-Sadepe Michael Bronstein, et al. Sequence-augmeented se (3)-flow matching for conditional protein backhone generation. arXiv preprint arXiv:2405.20311 2024.
Islam, Cheng-Hao Liu, Jarrid
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anma Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596 2021.
Jin Sub Lee, Jisun Kim, and Philip Kim. bcsee- generative modeling for novo protein design. Nature Computational Science, 3(5):382 2023.
Jiahan Li, Chaoran Cheng, Zuofan Wu, Ruihan Guo, Shitong Luo, Zhizhou Ren, Jian Peng, and Jianzhu Ma. Full-atom peptide design based on multi-modal flow matching. arXiv preprint arXiv:2406.00735, 2024.
Xiaohan Lin, Zhenyu Chen, Yanheng Li, Zicheng Ma, Chuanliu Fan, Ziqiang Cao, Shihao Feng, Yi Qin Gao, and Jun Zhang. Tokenizing foldable protein structures with machine-learme artificial amino-acid vocabulary. bioRxiv, pp. 2023-11. 2023a.
Yeqing Lin, Minji Lee, Zhao Zhang, and Mohammed AlQuraishi. Out of many, one: Design- ing and scaffolding proteins the scale of the structural universe with genie 2. arXiv preprint arXiv:2405. 15489, 2024.
Zeming Lin, Halil Akin, Roshan Rao. Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 1200066667)1111111133333 2023b.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hammn Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02777. 2022.

---

Published as a conference paper at ICLR 2026
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0m 03000, 2022.
Amy X Lu, Wilson Yan, Kevin K Yang, Vladimir Gligorijevic, Kyunghyun Cho, Pieter Abbeel, Richard Bonneau, and Nathan Frey. Tokenized and continuous embedding compressions of pro- tein sequence and structure. bioRxiv, pp. 2024-08, 2024.
Amy X Lu, Wilson Yan, Sarah A Robinson, Simon Kelow, Kevin K Yang, Vladimir Gligorijevic, Kyunghyun Cho, Richard Bonneau, Pieter Abbeel, and Nathan C Frey, All-atom protein genera- tion with latent diffusion. In ICLR 2025 Workshop on Generative and Experimental Perspectives for Biomolecular Design, 2025.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp.
2023.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Medical image computing and computer-assisted intervention- MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceed- ings, part III 18, pp. 234-241. Springer, 2015.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2265. pmlr, 2015.
Martin Steinegger and Johannes Söding, Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology. 35(11):0026-102 2017.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Michel Van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Jeongjae Lee, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Fast and accurate protein struc- ture search with foldseek. Nature biotechnology. 42(2):243-244 2024.
Yann Vander Meersche, Gabriel Cretin, Aria Gheeraert, ChnnCChrstoophe Gelly, and Tatiana Ga- lochkina. Atlas: protein ficxibility description from atomistic molccular dynamics simulations. Nucleic acids research, (DDD)):D344-D330 2024.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Dplm-2: A multimodal diffusion protein language model. arXiv preprint arXiv:24100 13782, 2024.
Xinyou Wang, Daiheng Zhang, Dongyu Xue, Fei Ye, Shujian Huang, Zaixiang Zheng, Quanquan Gu, al. Elucidating the design space of multimodal protein language models. arXiv preprint 2rXiv:2504.11455. 2025.
Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, 62097976): 1100, 2023.
Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Vietor Garcia Satorras, Bastiaan Veeling, Regina Barzilay, Tommi Jaakkola, al. Fast protein backbone generation with se (3) flow matching, arXiv preprint arXiv:2310.05297, 2023a.
Jason Yim, Brian I. Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint ar Xiv:2302.00277 2023b.
Jason Yim, Marouane Jaakik, Ge Liu, Jacob Gershon, Karsten Kreis, David Baker, Regina Barzi- lay, and Tommi Jaakkola. Hierarchical protein backbone generation with latent and structure diffusion. arXiv preprint 2rXiv:2504.09374, 2025.
12

---

Published as a conference paper at ICLR 2026
Xinyu Yuan, Zichen Wang, Marcus Collins, and Huzefa Rangwala. Protein structure tokenization: Benchmarking and new recipe. arXiv preprint arXiv:2503. 00089, 2025.
Yang Zhang and Jeffrey Skolnick. Scoring function for automated assessment of protein structure template quality. Proteins: Structure, Function, and Bioinformatics, 57(4):702-710 2004.
13

---

Published as a conference paper at ICLR 2026
A
ARCHITECTURE DETAILS
A. DIFFUSION TRANSFORMERS AND ALL-ATOM ATTENTION
The yellow background indicates the differences between PROTEINAE and AlphaFold3.
Algorithm Attention with pair bias and mask
1:
function ATTENTONPARBBLAS((S})) {csh. (Pal. (Big). Nhcud)
2:
a, f- AdaLN(s;, c,)
3: q Linear(s;)
4:
k. v LinearNoBias(s))
5:
k,, V LayNorm(k)) LayNom((v)))
6:
K v RoPEEK, v())
7:
D Linear(p:;) Big
8:
g <- sigmoid(Linear((a.
Attention
9:
4 + softmax,
(ta (q + bi)
10: S; LinearNoBias concata (g. E, A1, )
11: return {s:}
12: end function
Algorithm 2 Diffusion Transformer
function DIFRUSIONTRASSPORMEE((SS. {c,), {pt). (Big]. Notock Nuena)
2: for all ( Natua: do
3: {s.) AttentionnPiirBias(( (s.), (c.), (pa), (B,y), Nhead)
4: a, f- + ContiitonedIhannntiiinllooo,,,,
5: end for
6: return {s:}
7:
end function
Algorithm 3 Atom Transformer
1: function ATOMTRANSFORMERK ). {ci), [pim). Notock = 3, Noadd N Nqueries 32, Nkeys {15.5, 47.5, 79.5....
128, Saubset centres
0
if - cll Nqueries Im cll N Ve Suubset centres else
2: Bim -0
3: {qr} f- IiffusionTransfrrmme ((r:)) (ct), (Pim); (Bmm). Notock, Nheus)
4: return (q))
5: end function
14

---

Published as a conference paper at ICLR 2026
Algorithm 4 All-Atom Attention Encoder
function ALLATTOMAATNEENCODDE(((E)) (x), (s). {c), (Pij). Catoen 64, Catormpair 16,
Ctoken 256)
2: Ci <- LinearNoBiaasconccallec pos fqat charge charg, poot mask fiet element ffee atom name chars)
3: dim pfotpos pfot.pos
Vim - pref_suceeiis dim
4:
Im.
5:
Pim LinearNoBias( dim) Vim
6:
Pim LinearNoBias
d+ Vim
7:
Pim += LinearNoBias(w..)
8:
q + C
9:
C += LineaNoBBissLLyyeeNornniCC. idx(
10:
Pim + Lanerrt/ooiiisseeeeeeeooooooooaa
11:
q += LinearNoBiassrr))
12:
Pim + LinearNoBBiss(reut(())) + LinearNoBBiastrllleee
13:
Pim +=
))))))
14:
{qi} Atom Troomfraaformmr ((q:) {(t) (Pim), ). Notuck = 3, Nueat 4)
15:
S: f- me.. -. idxx) (inneaainanNNoas(a)))))
<kip qutip, pxiip qı Ct Pim
17: return {si), (q:
skin) ) (c (pim)
skip
end function
Algorithm All-Atom Attention Decoder
1: function ALL.ATOOMATNDDEEOORRR(SSS (qit)),, (cthip), (pime)
2: q LincarNoBiass(saaaac)) + skip
3: {q:} } AttonTTransfommer(( (q:), {c,i)), (pmn)) Notook 3, Nhead 4)
4: f- Limear\ooBis(LyverNomm(((())
5: return {u}
6: end function
A.2 PROTEINAE FLOW DECODING
PROTHINAE decodes reconstructed protein structures by simulating the learned ODE velocity field, conditioned on the latent representation
dz, of (x,, t,
(10)
dt
Starting from a noise sample ZT, the ODE integrated from t = 1 to t = 0 to obtain the denoised structure Zg. In practice, integrate the ODE using numerical solver (e.g., Euler method) with a small number discretization steps to reduce computational cost. The latent representation acts powerful condition during decoding, guiding the structure generation process. It serves a role analogous to sequence features (like MSAs) in models such as AlphaFold3, enabling more deterministic structure generation conditioned on the learned latent.
15

---

Published as a conference paper at ICLR 2026
A.3 PROTEIN LATENT DIFFUSION MODELING (PLDM)
Protein Latent Diffusion Modeling (PLDM) is our framework for generating novel latent represen- tations. We employ flow matching the generative approach in the latent space, which learns continuous transformation to map samples from a simple distribution (e.g.. Gaussian) to the target latent distribution. Unlike structure-basee generative models that operate on the SE(3) manifold of protein structures, our PLDM operates purely on the compact latent representation learned by the PROTHINAE, as illustrated in Fig. 2
PLDM Training Objective PLDM is trained using flow-min matching loss applied the latent representation 2 R No The model v trained to predict the velocity field the latent space. Given clean structure 21 from the data distribution (a) its latent representation is E(x)). We define a linear path in the latent space between a noise sample 2g ~ NN T) and E(r)): 2; = ttt The target velocity for this path is dz. = -(x1)= -2-- The PLDM is parameterized by Φ, and its training objective is:
dt
min B
(2) (E(x)) 20)1
(11)
Naown
where v (2t, t) is the velocity predicted by the PLDM network given the noisy latent sample ž and timestep t. The expectation is taken over data samples TT, noise samples zg. and timesteps sampled from same p(t).
PLDM Sampling Following observations from Proteina (Geffiner al. 2025) regarding the limi- tations of directly sampling the full learned distribution, we perform sampling from the latent space using an SDE-based schedule. The SDE for generating latent samples é is defined as:
dz = u t) dt + g(t)s (z,, t)dt + 2g((t) dW,,
(12)
where (ze 1) is the velocity predicted by the trained PLDDM model, g(t) is the diffusion coeffi- cient, (21, 1) is the score function derived from the predicted velocity (specifically, (=,, t) Rs
V.. log Pr (2x)), 7 is noise scale parameter, and dVW, is standard Wiener process. Starting from sample from the Gaussian distribution 2T ~ (O, 1). we integrate this SDE numerically from f
to 0 to obtain the generated latent E. Finally, this generated latent is decoded into protein
structure using the trained PROTEINAE flow decoder: = D(2) (Eq. 10
16

---

Published as a conference paper at ICLR 2026
A.4 TRAINING DETAILS
Table 4: Training and model hyperparameters for PROTEINAE, PLDM, and Flexibility Prediction.
Parameter
PROTEINAE
PLDM
Flexibility Prediction
Training Details GPUs
4 x 80G A100
8 x 80G A100
1 x 80G A100
Batch size per GPU
64
Global batch size
32
512
8 300 epochs
Training duration
10 epochs Adam
200,000 steps
Optimizer
Adam
Adam
Learning rate
x 10
x 10
x 10 MSE
Loss function
Flow-matching
Model Architecture
Token dimension
256
768
Layers
15
Attention heads
8
12
Condition dimension
256
512
Timestep embedding size
256
Sequence index embedding
128
Downsampling factor (r)
Dimension bottleneck
8
All-Atom Attention Module (PROTEINAE only)
Atom dimension
64
Atom pair dimension
6
Attention mask (query/key)
32 / 128
Flexibility Prediction Model
Architecture
MLP
Dimension
512
Activation function
Relu
17

---

Published as a conference paper at ICLR 2026
B BACKGROUND
The foundational diffusion model concepts introduced by (Sohl-Dickstein et al. 2015) and significantly advanced with denoising diffusion probabilistic models (DDPMs) (Ho al.
Diffusion models
2020 whiic iteratively add noise to data and then learn to reverse the process. More recently, flow
matching has emerged as powerful and efficient alternative (Lipman et al. 2022 Liu et al. 2022 Albergo et al. 2023). These flow-based methods directly learn a vector field, often an Ordinary Differential Equation (ODE), to transform a simple noise distribution into a complex data distri- bution, frequently enabling simulation-free training and faster sampling. Both diffusion and flow matching paradigms can be further extends to general manifolds (Chen & Lipman] 2023) as well as discrete data (Cheng et al., 2024: Austin al. 2021). They have found extensive applications in de novo protein design, with models like RFDiftusion, FrameDiff (initially iittusioonbbseed and their successor RFdiffusion2, FrameFlow (which incorporates flow matching for training). Some work (Campbell et al.] 2024 Li et al.] (2024) has also attempted the co-design of sequences and structures. More recently, Proteina (Geffner al.] (025) attempted to model coordinates directly using a non-equívariant architecture on E(3), but they only work with C..
Protein Autoencoders Autoencoders play significant role in learning compact and meaningful representations for both protein sequences and structures, which are vital for efficient generative modeling and other downstream tasks. While protein sequence autoencoders are widely used, often leveraging features like MSAs or ESM embeddings (Detlefsen et al. 2022 Lu et al. 2024 2025
Chen et al.) 2024). the development of autoencoders for protein structures is a more recent area of research. A notable example that bridges these areas is CHEAP (Lu et al. 2024). This model operates on protein language model features from ESM2, which are then used to predict structure via ESMFold's structure module. While this method can yield valuable structural information, its encoding capability is inherently dependent on, and thus limited by, ESMFold's ' prediction accuracy.
In parallel, pioneering efforts have focused on autoencoding the protein structure directly. These models tackle the challenge of converting continuous 3D atomic coordinates into discrete tokens. Key examples include the ProToken Lin al.. (2023a), ESM3 VQ-VAE tokenizer (Hayes et al.) 2025) and the DPLM-2 lookup-free quantization (LFQ) tokenizer (Wang et al. 2024 2025). By en- abling this conversion, they allow for the joint modeling of protein sequences and structures within large generative frameworks. Subsequent work, such as AminoAseed (Yuan et al. 2025), has further explored improvements in the codebook design for these discrete structure representations. How- ever, a common challenge for these tokenization- based autoencoders is the inherent information loss that occurs during the discretization process. This makes it difficult to fully capture the complexity and subtle geometric details present in the continuous 3D space of a protein structure.
Latent Space Generative Models for Proteins While diffusion models operating directly on atomic coordinates or torsion angles have achieved remarkable success, they often suffer from high computational costs due to the vast exploration space of 3D protein structures. To address this, re- cent works have investigated difffusing proposed LafentDiff, which employs an SE(3) -E3iivaauaaa graph autoencoder to map protein back- bones into a condensed latent representation. By performing the diffusion process within this lower- dimensional latent space rather than the original coordinate space, LatentDiff significantly reduces computational complexity, achieving sampling speeds orders of magnitude faster than frame-based approaches like FrameDiff Yim et al. (2023b) RFdiffusion Watson et al. (2023).
compressed latent space. For instance, Fu al. (2024)
Table 5: Comparison of different protein autoencoder architectures, highlighting their input modal- ities, representation types, and operational scopes.
Methods
Input Modality
Representation Type
Operational Scope
Compression Strategy
ProteinAF
Atom courdinates
(Structure)
Continuous
All proteins
structural compression
/
ESM2 features (Sequence)
CHEAP
Continuous
Proteins Eld predicts well
Soquence feature compression
Pro ESM3 VQ-VAE LFQ
Backbone Frame (Structure)
All proteins (with 2,048 length coustraints for ProToken and DPLM-2 LFQ)
Discrete
18

---

Published as a conference paper at ICLR 2026
B.1 BACKBONE-LEVEL RECONSTRUCTION QUALITY
Table 6: Different protein autoencoders' structure reconstruction quality measured by RMSD (1).
Lower better. Bold indicates the best performance. Proteins missing backbone atoms are ignored. "Note that ProTokee and DPLM-2 can only process proteins under 2,048 residues.
CASP14
CASPI5
Methods
T
T-dom
oligo
TS-dmmanns
oligo
CHEAP
11.16+10.09
4.71+5.21
11.10+11.95
10.98-12.44
8.24+11.12
ESM3 VQ-VAE
1.28++232
0.66++0.42
3.1147.41
1.25+1.28
2.47+2.277
Backbone
ProToken*
1.14+0.66
1.09++0.16
1.55+1.46
1.29-0067
1.33+00.00
RMSD
DPLM-2*
1.94.11.99
1.47:10.22
3.8117.20
4.58-6.62
3.8311.90
PROTEINAE
0.51
0.23.10.11
0.31110.27
0.31
0.4310.51
19

---

Published as a conference paper at ICLR 2026
C METRIC DETAILS AND ADDITIONAL EXPERIMENTS
We complete the details of protein unconditional generation Sec.
Designability Protein designability is assessed based on whether protein backbone structure can be generated from a specific amino acid sequence that folds into that structure. Following the method of FrfmeDiff (Yim et al. 2023a), eight sequences are generated per backbone using MPNN (Dauparas al.] (022) with sampling temperature of 0.1. Structures are predicted using ESMFold (Lin et al. 2023b). The Root Mean Square Deviation (RMSD) between the predicted and original structures is calculated. A sample is deemed designable if its lowest RMSD, termed self-consistency RMSD (scRMSD), is v 2Å. The overall designability score the fraction of designable samples.
Designable Pairwise TM-score (DPT) This is the first of methods for evaluating diversity, based on the methodology by (Bose et al. 2023). It involves calculating pairwise TM-scores (Zhang & Skolnick 2004) among all designable samples for each specified protein length, and then ag- gregating the averages across different lengths. TM-scores range from 0 to 1, where higher scores indicate greater structural similarity. Therefore, lower TM-scores are preferred for this metric as they suggest greater diversity.
The second diversity measure follows the approach by FrameDiff (Yim et al, 2023b). Designable backbones clustered using Foldseek (Van Kempen et al. 2024) based TM-score threshold of 0.5. Diversity is calculated as the ratio the total number of clusters to the number of designable samples:
Diversity (Cluster)
Number of designable clusters
Diversity (Cluster)
(13)
Number of designable samples
In this case, higher scores are preferable, indicating that the designable samples form a larger number of distinct clusters, thus representing greater diversity.
Novelty Novelty quantifies how distinct model's generated structures are compared to structures present in predefined reference databases (AFDB and PDB in our work). For each designable struc- ture generated by the model, its TM-score is computed against every structure in the reference set, and the maximum TM-score recorded. The novelty score is the average of these maximum TM- scores over all designable samples. Lower scores are considered better, as they suggest the generated structures are less similar to known structures in the reference sets, implying higher novelty.
Evaluation Details For protein unconditional generation, we sample 10 proteins of each length between 60-128. resulting in 690 samples in total. is noting that the designability is not a accurate indicator of how well generative model matches the training distribution since the training dataset far from 100% designable (Huguet et al. 2024).
20

---

Published as a conference paper at ICLR 2026
D IMPACT OF LENGTH BOTTLENECK ON GENERATION
We further investigate the impact of imposing a length bottleneck to PROTEINAE on the PLDM generation process.
Figure 4: Illustration of structural collapse observed during protein generation by PLDM under length bottleneck conditions.
As depicted in Fig. 4 observe that specific regions of the generated protein structures collapse towards the center, indicating a failure to maintain realistic local or global geometry under these conditions. We hypothesize that this phenomenon is attributed the use of nearst interpolation dur- ing the upsampling process within PROTEINAE or the padding value in the downsampling training. We plan to investigate this issue detail in future work.
21

---

Published as a conference paper at ICLR 2026
E PROTEINAE v2 PREVIEW
6ARE
4AWS
2MNE
Figure 5: PROTEINAE v2 achieves high-fidelity all-atom reconstruction multi-modal com- plexes. The reconstructed structures are depicted in color, while the ground truth complexes shown in grey for comparison.
22


Reviews:
[{"review_id": "YExOw0oSBV", "reviewer": "Reviewer_FiwP", "summary": "In this work, the authors propose using a non-equivariant diffusion autoencoder called ProteinAE to encode protein structure by directly mapping backbone coordinates into a continuous latent space. And they show the performance of their model compared to recent generative frameworks along with its performance as a structure embedder for a downstream prediction task. Overall, the field has largely treated equivariance as gospel—driven by the desire to model structure from first principles and inspired by the success of AlphaFold2. However, AlphaFold3 recently shifted this paradigm. This work presents an interesting application of non-equivariant modeling for generative purposes, while also offering a fast structural encoder, though it still requires further validation and benchmarking to fully establish its advantages.", "questions": "- In the introduction, authors mention that other models use SE(3) representations and are therefore slow. However, SE(3) representations are also essential in protein structure modeling specifically in generative protein design, so this point needs to be explained more carefully — perhaps discussing when SE(3) equivariance is beneficial versus when it introduces unnecessary computational cost. \n- Authors include CASP14 and CASP15 in their benchmark for structure reconstruction quality and compare against ESM3 etc, It would be good to see the performance on CAMEO as well for the same period as ESM3.  \n- Since the model is positioned also as a faster alternative for encoding structural information, it would strengthen the paper to demonstrate its applicability beyond the current benchmark on physiochemical property prediction. For example, evaluating its performance on other downstream structure-based prediction tasks, such as GO-term prediction, could better illustrate the practical impact of the approach.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Not imposing equivariant restriction on the protein structure representation enables a much more computational efficient framework for modeling protein structure. \n- Superior benchmark performance of the trained representation layers for physiochemical property prediction. \n- Authors provided benchmarks against some of the newer models for structure generation. Compared to RFDiffusion model which has been extensively validated in the lab, authors report ProteinAE while performing slightly worse than RFDDiffusion for designability and diversity, the sampling time is much more efficient. This can be useful for some protein design pipelines. \n- The model also seems to perform well at reconstructing structure based on the CASP14/15 benchmarks.", "weaknesses": "- The paper would benefit from clarifying some sections [see the question below]. \n- While the model shows promise, the included benchmarks are across a limited range of tasks, specifically for protein structure encoding and its performance in downstream tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "inwxfx0YW8", "reviewer": "Reviewer_MnGk", "summary": "This paper proposes PROTEINAE, a protein diffusion autoencoder mapping protein backbone coordinates from E(3) into a continuous, compact latent space for protein modeling and generation. PROTEINAE uses a non-equivariant Diffusion Transformer and is trained with a single flow matching objective to simplify the training objective. PROTEINAE achieves better reconstruction quality and high-quality structure generation that significantly outperforms prior latent-based methods.", "questions": "Please refer to the weakness section", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- PROTEINAE uses a non-equivariant autoencoder based on Diffusion Transformers that operates directly on atom coordinate without the need to considering complicated equivarience.\n- Performance of PROTEINAE is better than latent-based methods and competitive with structure-based methods.\n- The model is trained using a simple flow matching loss, which makes the training objective easy to design.", "weaknesses": "- Structure-based method is still better than the latent-based model, it’s not clear if the latent model has potential to achieve better performance than structure-based method. Could authors elaborate more on that?\n- PROTEINAE doesn’t consider equivariance to simplify the design but would that lead to generating non-physical structures?\n- The comparison missed some important baselines, La-proteina and Proteina. Proteina also uses non-equivariant model design and La-proteina is a latent model based on Proteina.\n\n[1] Geffner, Tomas, et al. \"Proteina: Scaling flow-based protein structure generative models.\" arXiv preprint arXiv:2503.00710 (2025)\\\n[2] Geffner, Tomas, et al. \"La-proteina: Atomistic protein generation via partially latent flow matching.\" arXiv preprint arXiv:2507.09466 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KlhJ79VGN7", "reviewer": "Reviewer_dnyu", "summary": "This paper proposes a new generative model for protein tertiary structure generation and reconstruction. Most existing approaches rely on coordinate prediction models or diffusion models based on stochastic noise denoising, which are computationally expensive and often struggle to maintain structural consistency and reconstruction accuracy. In contrast, this study combines the Flow Matching framework with a Transformer-based network, achieving continuous and stable structure generation. Unlike conventional SE(3)-equivariant models, the proposed method intentionally adopts a non-equivariant architecture on E(3) space. This design eliminates the need for explicit and computationally costly handling of rotations and translations, while still maintaining effective geometric robustness through structure alignment, pairwise bias based on relative distances and angles, and data augmentation.\n\nThe proposed framework consists of three stages. In the first stage, the model takes the backbone atoms of each residue as input and compresses them into residue-level representations using an All-Atom Attention Encoder, which is designed to capture both spatial dependencies among residue pairs and primary sequence information. In the second stage, these features are processed by an Encoder–Decoder structure based on a Diffusion Transformer (DiT), which learns the global 3D correlations of the entire protein. Pairwise representations are incorporated as attention biases, enabling the model to encode global dependencies that reflect inter-residue distances and angles. The third stage introduces the Protein Latent Diffusion Model (PLDM), which treats the latent representations of residue sequences obtained by the encoder as latent token sequences. After compressing both the sequence length and feature dimension, PLDM learns the diffusion process directly in the latent space. Rather than performing stochastic noise denoising, PLDM is designed as a Transformer that predicts the velocity field of latent vectors based on Flow Matching. To preserve the sequential order while integrating temporal embeddings, positional encodings are applied to the latent tokens.\n\nExperimental evaluation on protein structures from the AlphaFold Database demonstrates that the proposed method achieves smoother and more stable structure generation than conventional diffusion models, with higher reconstruction accuracy as measured by RMSD and FAPE. The generated structures also preserve physical consistency, confirming that the proposed approach enables efficient and high-precision protein structure generation.", "questions": "Do the authors have any ideas or future plans for incorporating additional constraints or introducing some form of equivariant mechanism to extend the applicability of the proposed model to larger and more complex protein structures?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The main strength of this paper lies in its novel design choice to intentionally train the model in a non-equivariant manner on E(3) space, rather than relying on conventional SE(3)-equivariant frameworks. By doing so, the authors successfully avoid the computational cost of explicit rotation and translation handling while maintaining effective geometric robustness through structure alignment, pairwise bias based on relative distances and angles, and appropriate data augmentation. This represents a meaningful and practical contribution to the field of protein structure generation.\n\nAnother notable strength is the use of a simple Flow Matching loss, which removes the need for complex KL or reconstruction losses commonly required in traditional diffusion models. This design enables a smoother and more stable generation process, allowing structures to evolve naturally from noise with higher accuracy and computational efficiency than probabilistic baselines. By integrating Flow Matching with a Transformer architecture, the paper reformulates protein structure generation as learning a continuous flow rather than stochastic denoising, effectively preserving structural continuity and physical consistency while enabling efficient generation in the latent space.", "weaknesses": "The authors clearly state this limitation, but the proposed model still has difficulty handling very large proteins and multi-chain complexes. Although using a non-equivariant design on E(3) makes the computation much faster, it also creates concerns about whether the model can keep geometric consistency and produce physically reliable results, especially for rotation and translation. Because of this, the method may not yet be suitable for very large or complex protein structures. As the authors mention, adding more geometric constraints or partially equivariant mechanisms in the future could help overcome this problem and further improve the model’s applicability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Flv2PTwJUR", "reviewer": "Reviewer_qmAJ", "summary": "This paper proposed PROTEINAE, a protein diffusion autoencoder that learns compact, continuous representations of protein structures directly from 3D backbone coordinates in Euclidean space (E(3)). \n\nUnlike prior approaches that depend on SE(3)-equivariant models, discrete tokenization, or multiple loss functions, PROTEINAE simplifies training with a single flow-matching objective and a non-equivariant Diffusion Transformer (DiT) architecture. The model efficiently encodes protein backbones into a latent space through a bottleneck design, enabling high-fidelity reconstruction and scalable downstream modeling. On benchmarks such as CASP14 and CASP15, PROTEINAE achieves state-of-the-art reconstruction accuracy, substantially outperforming discrete and equivariant autoencoders.\n\nThe authors argue that strict equivariance may not be essential for effective protein representation learning, citing recent advances such as AlphaFold3 and Proteina, which also employ non-equivariant architectures. However, the paper could benefit from a more comprehensive discussion of related work on 3D protein autoencoder design, as well as a deeper theoretical justification for the choice of a non-equivariant model.", "questions": "As listed above in weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "## Strengths\n\n- This paper introduces a non-equivariant diffusion autoencoder that learns continuous protein structure representations directly from 3D coordinates, simplifying model design and training compared to SE(3)-equivariant or discrete approaches. But as mentioned before, the paper could benefit from a deeper theoretical justification for the choice of a non-equivariant model.\n\n- The proposed method achieves state-of-the-art reconstruction accuracy on CASP14/15 and competitive structure generation quality relative to more complex diffusion models.\n\n- Computational efficiency: The latent diffusion framework offers good speed and memory advantages over previous diffusion methods.", "weaknesses": "## Weaknesses:\n\n- This paper lacks sufficient discussion of prior work on 3D protein backbone structure autoencoder design. For example, the model LatentDiff is compared in Table 2, but it is not discussed in the related-work section. It shares similar design ideas with the proposed method, including an autoencoder that reduces the latent diffusion modelling burden, shrinks the modelling space in terms of length, and offers computational efficiency.\n\n- The authors argue that strict equivariance may not be essential for effective protein representation learning by citing recent advances such as AlphaFold3 and Proteina, which also employ non-equivariant architectures. However, the paper would benefit from a deeper theoretical justification for choosing a non-equivariant model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)