CODLAD

A backmapped protein can have every atom within two angstroms of the truth and still be unusable: one over-long bond, one side chain threaded through a ring, and the coordinates will not survive the first step of an all-atom simulation. Diffusion models that denoise directly in Cartesian all-atom space have to police that chemistry inside the sampling loop, and that is where the field's standing trade-off comes from. Denoising residue by residue keeps local geometry clean but accumulates error along the chain and multiplies cost; denoising the whole structure at once is faster but drifts away from correct bond topology.

CODLAD (COnstraint Decoupled LAtent Diffusion), from the Beijing Academy of Artificial Intelligence with the Institute of Automation, CAS and the University of Chinese Academy of Sciences, moves the constraints out of that loop. A vector-quantized autoencoder is trained first to compress all-atom structures into discrete residue-level latents under explicit bond-length, bond-angle, torsion, bond-topology and steric-clash penalties, so that any point on the learned latent manifold decodes to a chemically plausible structure. A conditional diffusion model then operates only in that latent space, conditioned on the coarse-grained Cα trace and the amino acid sequence, and never reasons about individual atoms. Geometric validity comes from the decoder rather than from constrained sampling.

The work first appeared on arXiv in October 2024 under a different name, was substantially revised, and was published in the Journal of Chemical Theory and Computation in December 2025. Code, dataset splits and pretrained checkpoints are released under an MIT license. It joins CGBack and StruCloze as a backmapping model from an unrelated group.

#Key Features

• Constraints learned, not enforced: bond, angle, torsion, topology and clash terms enter as reconstruction losses during compression, so the generative stage needs no coordinate-level constraint handling and no post-hoc energy minimization.
• Discrete latent space: a 4096-entry codebook with an embedding dimension of 3 outperforms a continuous VAE on structural fidelity, which the authors attribute to discretization resisting the mode collapse that blurs distinct conformational states.
• Internal-coordinate decoder: structures are represented as bond lengths, bond angles and dihedrals over up to 13 heavy atoms per residue, then mapped deterministically back to Cartesian coordinates.
• Non-autoregressive denoising: a ProteinMPNN-based network updates every residue latent in parallel, conditioned on the whole coarse-grained graph, rather than walking the chain residue by residue.
• Checkpoints per training corpus: the release ships separate autoencoder variants for PED, PDB and ATLAS, and a reconstruction and sampling entry point that runs against them without retraining.

#Technical Details

Compression uses hierarchical message passing at both atom level and residue level with cross-level exchange, producing residue-level latents with a node feature dimension of 36. The denoiser is a three-layer encoder-decoder with hidden size 128 built on ProteinMPNN, taking noisy latents, Cα coordinates, sequence and diffusion timestep, the last injected through adaptive layer normalization. Sampling uses 100 steps of a 1000-step linear variance schedule. Training used a single 40 GB NVIDIA A100.

Evaluation spans the Protein Ensemble Database (85 intrinsically disordered proteins at about 100 conformations each), ATLAS (1,297 training proteins after filtering above 512 residues, 70 held out), a SidechainNet-derived PDB set (62,105 training structures), and the DES fast-folding trajectories held out entirely as an out-of-distribution test. Against GenZProt, DiAMoNDBack and FlowBack, CODLAD reaches the lowest RMSD on all three in-domain sets (1.710 Å on PED, 1.236 Å on PDB, 1.539 Å on ATLAS) and cuts graph edit distance by 41.1% on PED and 45.3% on ATLAS. On DES, after pretraining on PDB only, it improves RMSD by 12.6% and graph edit distance by 56.9% over the strongest baseline. Latent denoising also cuts total inference time by roughly 70% — 9.9 hours versus 33.6 and 37.3 for the two all-atom diffusion baselines on the DES test set.

#Applications

Backmapping is the bridge between the timescales coarse-grained simulation can reach and the atomic detail that ligand docking, protein-protein interface analysis and all-atom production runs require. CODLAD suits groups that generate long coarse-grained trajectories and need ensembles rather than single structures: disordered-protein modeling, and conformational analysis where the cost advantage compounds over thousands of frames.

#Impact

The result that matters most is the ablation: swapping a continuous VAE for a vector-quantized one, and moving constraint handling from sampling into compression, each improve topological fidelity independently. That is a design argument other structure generators can adopt, separate from where CODLAD sits on any current leaderboard. The limitations are stated plainly by the authors. The formulation is tied to Cα-only coarse-graining, so MARTINI or mixed-resolution mappings would require retraining, and the released checkpoints are per corpus rather than one general model. Weights are distributed as a Google Drive archive rather than through a model hub, and the repository is a research codebase, not a packaged tool.