HemePLM-Diffuse: Protein–Ligand Generative Dynamics

• The paper introduces HemePLM-Diffuse, which generates plausible all-atom protein–ligand trajectories via a transformer-based, flow matching diffusion approach.
• It employs SE(3)-invariant tokenization for proteins and ligand fragments, ensuring scalability and precise modeling of complex, large systems.
• Empirical results reveal lower RMSD values and over 100× speedup compared to classical MD, highlighting its potential in rapid drug discovery.

Searching arXiv for the specified paper to ground the article in the published record. HemePLM-Diffuse is a generative, transformer-based diffusion framework for simulating long-timescale, all-atom protein–ligand dynamics, with explicit support for large systems and ligand-centric tasks such as inpainting and transition path sampling. It is designed for accurate simulation of protein-ligand trajectories, inpaints the missing ligand fragments, and samples transition paths in systems with more than 10,000 atoms. The framework is positioned against the core bottleneck of conventional molecular dynamics, where integrating Newton’s equations at femtosecond resolution becomes prohibitively expensive for large protein–ligand systems over microsecond–millisecond timescales and for rare events such as binding/unbinding and large conformational changes (Thakur et al., 7 Aug 2025).

1. Problem setting and model scope

The framework addresses long-timescale dynamics of protein-ligand complexes in drug discovery and structural biology. In the formulation described for HemePLM-Diffuse, the central objective is not force prediction alone, but direct generation of plausible trajectories in a latent, SE(3)-invariant token space using a diffusion-style, velocity-based generative process described as flow matching. This design is intended for full all-atom protein–ligand systems, explicit representation of ligand fragments, and atom counts characteristic of realistic drug-discovery targets (Thakur et al., 7 Aug 2025).

Its positioning relative to earlier machine-learning approaches is explicit. Force-field models such as TorchMD-Net and SchNet still require numerical integration and are typically limited to small systems without ligands. Generative dynamics models such as MDGEN operate on coarse-grained backbone traces, for example C$\alpha$ only, and omit explicit all-atom ligand dynamics. Uni-Mol and related 3D pre-training models provide strong SE(3)-aware structure encoders, but are static and not designed for trajectory generation. HemePLM-Diffuse is therefore framed as a model that unifies full trajectories, explicit ligand modeling, scalability beyond 1k atoms, all-atom protein–ligand complexes, ligand inpainting, trajectory upsampling, transition path sampling, and SE(3)-invariant dynamics modeling within a single generative backbone (Thakur et al., 7 Aug 2025).

A recurrent point of clarification is the distinction between invariance and equivariance. HemePLM-Diffuse uses SE(3)-invariant token features rather than fully SE(3)-equivariant vector features. The stated rationale is to preserve the geometry of 3D space while allowing transformer-style attention to scale. This contrasts with fully SE(3)-equivariant models, where outputs transform covariantly under rotations. The framework sacrifices strict equivariant vector outputs in favor of invariant feature tokens plus a simpler, more scalable transformer backbone; because the primary tasks are generative trajectory prediction and structural completion rather than direct force prediction, the paper states that invariance is sufficient and enables high scalability (Thakur et al., 7 Aug 2025).

2. SE(3)-invariant tokenization of proteins and ligands

A central design choice is tokenization at the residue/backbone level for proteins and in a fragment-aware manner for ligands. For proteins, each residue token encodes the 3D coordinates of the C$\alpha$ atom, backbone torsion angles $\Phi$, $\Psi$, and $\Omega$, and scalar auxiliary features such as B-factors, solvent accessible surface area (SASA), and binding-site scores, including examples from AlphaFold annotations. These features are aggregated into a tensor

$P \in \mathbb{R}^{L_p \times d_p},$

where $L_p$ is the number of protein tokens and $d_p$ is the embedding dimension (Thakur et al., 7 Aug 2025).

For ligands, tokenization uses RDKit and encodes fragment identifiers and atom types, internal dihedrals, relative 3D geometric features with respect to the protein, top-$k$ distances from each ligand atom or fragment to nearby residues, “relative 3D positions” to residues expressed through distance-based and other SE(3)-invariant descriptors, and atom-wise physicochemical descriptors. These are stored as

$L \in \mathbb{R}^{L_l \times d_l}.$

Although the paper does not spell out explicit distance formulas, the SE(3)-invariant nature is stated to depend on quantities such as pairwise distances and angles, for example

$\alpha$0

and internal torsions, all of which are unchanged under rigid-body rotations and translations. Absolute positions $\alpha$1 are not used directly in attention; instead, the model uses invariant encodings derived from them (Thakur et al., 7 Aug 2025).

This tokenization strategy has architectural and physical consequences. Because the entire system can be rotated or translated without changing the token features, the output distribution is E(3)-/SE(3)-invariant. At the same time, tokenization at the residue and fragment level reduces sequence length relative to atom-level modeling. This suggests a deliberate exchange: strict equivariant vector outputs are not retained, but scalability and compatibility with a standard transformer backbone are improved (Thakur et al., 7 Aug 2025).

3. Transformer backbone and diffusion formulation

The core model is a transformer that operates over protein and ligand token sequences with intra-molecule self-attention, cross-attention between protein and ligand tokens, and explicit time conditioning for diffusion dynamics. At a given diffusion time step $\alpha$2, the model takes noised protein tokens $\alpha$3 and ligand tokens $\alpha$4, a scalar time $\alpha$5 encoded via sinusoidal time embeddings, and optional conditioning information such as start and end conformations for transition path sampling, partial ligands for inpainting, and contact maps or other structural constraints. The time scalar is mapped to a sinusoidal embedding $\alpha$6 and added or concatenated into token embeddings, followed by an MLP projection to a shared hidden dimension (Thakur et al., 7 Aug 2025).

The encoder stack consists of $\alpha$7 transformer layers with interleaved self-attention over protein tokens, self-attention over ligand tokens, cross-attention in both directions, and SE(3)-invariant MLPs operating on token features. The high-level layer structure is described by

$\alpha$8

followed by

$\alpha$9

The decoder or prediction head maps final hidden states back to velocities or denoised token states for each atom, residue, or ligand fragment. Cross-attention is explicitly intended to capture binding-coupled motions such as induced fit, allosteric shifts, and ligand rearrangements within the pocket (Thakur et al., 7 Aug 2025).

The diffusion formulation is described as a stochastic velocity-based denoising approach, namely flow matching, rather than a standard discrete-time DDPM-style diffusion model. Perturbed states are written as

$\Phi$0

where $\Phi$1 is the clean latent representation and $\Phi$2 is the noisy version at time $\Phi$3. The model $\Phi$4 is trained to predict the “trajectory velocity” or score-like vector field guiding the reverse denoising process, with a core diffusion-related loss given as

$\Phi$5

In practice, this is presented as a flow-matching objective that learns the instantaneous velocity field transporting noisy data toward the data manifold (Thakur et al., 7 Aug 2025).

Training uses a weighted multi-objective loss. The trajectory MSE loss is

$\Phi$6

The free energy surface loss matches Markov state models derived from generated and MD trajectories through KL divergence,

$\Phi$7

The ligand inpainting loss is a cross-entropy loss over masked positions,

$\Phi$8

The transition path sampling loss is

$\Phi$9

These components are combined as

$\Psi$0

with coefficients not specified explicitly (Thakur et al., 7 Aug 2025).

4. Conditional generative tasks

Because the same backbone uses time-embedded tokens with optional conditioning, HemePLM-Diffuse supports several distinct tasks. In trajectory generation for long-timescale dynamics, an initial frame or a small prefix of MD is used to generate full trajectories by iterative denoising in the SE(3)-invariant latent space. The outputs are sequences of protein and ligand token states over time that can be decoded back into Cartesian coordinates and torsions. Since no force computation or numerical integration is required after training, the paper states that trajectories can be sampled orders of magnitude faster than MD (Thakur et al., 7 Aug 2025).

Trajectory upsampling is formulated as filling in missing intermediates when MD snapshots are sparse, for example every 10 frames. Conditioning on the observed frames, the model generates intermediate frames that are temporally smooth and consistent with both endpoints. The paper attributes this behavior to the combination of the TPS-style loss and flow matching, which are said to ensure that interpolated frames lie on physically plausible paths (Thakur et al., 7 Aug 2025).

Ligand inpainting is defined for incomplete ligands with missing atoms or fragments. The model conditions on protein context tokens $\Psi$1 and the observed ligand tokens $\Psi$2 and generates the missing $\Psi$3. The bidirectional cross-attention between ligand and protein is described as the mechanism that ensures inpainting respects the shape and chemistry of the binding pocket. The stated use cases include repairing incomplete structures, suggesting alternative ring closures, and exploring local modifications (Thakur et al., 7 Aug 2025).

Transition path sampling is performed by conditioning on initial and final states $\Psi$4 and $\Psi$5, endpoint tokens that may be clamped, the time embedding $\Psi$6, and the interpolant $\Psi$7 as a guide. The TPS loss is introduced to encourage intermediate states to lie on plausible low-energy paths connecting the endpoints. The paper also notes ligand optimization and dynamic edits: although not detailed with separate equations, the model is said to propose dynamic “edits” that promote stable binding interactions, effectively suggesting ligand modifications and testing them within the learned dynamical environment (Thakur et al., 7 Aug 2025).

5. Data, system setup, and scalability profile

The principal experimental system is the PDB 3CQV complex, namely the Rev-erb$\Psi$8 nuclear receptor bound to a heme ligand. The system is solvated and equilibrated with GROMACS using the Amber99SB-ILDN force field, and atomistic MD trajectories are generated. In the experiments, the solvated protein–ligand complex has over 1,700 atoms, while the broader design target is systems with more than 10,000 atoms. Preprocessing extracts C$\Psi$9 coordinates and torsions for each residue, ligand fragment positions, internal dihedrals, distances to nearby residues, and derived Markov state models for free energy surface alignment (Thakur et al., 7 Aug 2025).

The stated scalability rests on three architectural choices. First, token-level modeling replaces per-atom attention with residue tokens for proteins and fragment or atom tokens for ligands, reducing the effective sequence length and giving attention complexity roughly

$\Omega$0

with $\Omega$1. Second, SE(3)-invariance replaces full SE(3)-equivariance; invariant MLPs and scalar features are described as cheaper than tensor field networks because they avoid expensive spherical harmonics and higher-order irreducible representations. Third, numerical integration is not required at sampling time; trajectory generation consists of repeated forward passes of the transformer with noise (Thakur et al., 7 Aug 2025).

The computational profile reported for 1 ns-equivalent dynamics on the 3CQV+HEME system is concrete. GROMACS on a single GPU requires approximately 20 hours. TorchMD-Net, on simpler systems and still requiring integration, requires approximately 4 hours. HemePLM-Diffuse requires approximately 12 minutes, corresponding to a speedup greater than $\Omega$2 over classical MD. A wall-clock time versus atom count plot is described as showing much gentler scaling than both GROMACS and MDGEN as system size approaches 10k+ atoms (Thakur et al., 7 Aug 2025).

6. Empirical performance, implications, and limitations

The benchmarked comparisons against TorchMD-Net, MDGEN, and Uni-Mol cover ligand inpainting, trajectory upsampling, transition path sampling, and physical plausibility. In a 10-atom masking experiment on the HEME ligand, HemePLM-Diffuse achieves RMSD $\Omega$3 Å, MDGEN $\Omega$4 Å, and TorchMD-Net $\Omega$5 Å on simpler systems; Uni-Mol is reported as worse than HemePLM-Diffuse, though no precise number is given. The interpretation offered in the paper is that the lower RMSD reflects more effective use of protein context and binding-site geometry in reconstructing missing ligand atoms, particularly in macrocyclic systems such as heme (Thakur et al., 7 Aug 2025).

For trajectory upsampling and dynamics reconstruction from trajectories subsampled every 10 frames, HemePLM-Diffuse achieves average per-frame RMSD $\Omega$6 Å versus MD ground truth, while MDGEN is $\Omega$7 Å and Uni-Mol is $\Omega$8 Å. A “Trajectory RMSD vs Time” plot is described as keeping HemePLM-Diffuse consistently below 1.5 Å, whereas baselines remain above 2 Å and fluctuate more; a histogram of $\Omega$9RMSD/frame is reported to indicate smooth transitions between consecutive frames. This is presented as evidence for temporally coherent dynamics (Thakur et al., 7 Aug 2025).

For transition path sampling on the 3CQV+HEME system, the paper defines a TPS quality score in $P \in \mathbb{R}^{L_p \times d_p},$0 combining contact map overlap with reference, backbone RMSD smoothness, and continuity of a low-energy pathway. The reported values are approximately 0.95 for HemePLM-Diffuse, 0.75 for MDGEN, and 0.40 for Uni-Mol. Visualizations are described as showing the ligand smoothly docking into the pocket with contact heatmaps closely matching native contacts at corresponding stages (Thakur et al., 7 Aug 2025).

A compact summary of the main reported metrics is useful:

Evaluation	HemePLM-Diffuse	Baselines
10-atom HEME masking RMSD	0.91 Å	MDGEN 1.67 Å; TorchMD-Net 1.82 Å
Upsampling per-frame RMSD	$P \in \mathbb{R}^{L_p \times d_p},$1 Å	MDGEN $P \in \mathbb{R}^{L_p \times d_p},$2 Å; Uni-Mol $P \in \mathbb{R}^{L_p \times d_p},$3 Å
TPS quality score	$P \in \mathbb{R}^{L_p \times d_p},$4	MDGEN 0.75; Uni-Mol 0.40
1 ns-equivalent runtime	$P \in \mathbb{R}^{L_p \times d_p},$5 minutes	GROMACS $P \in \mathbb{R}^{L_p \times d_p},$6 hours; TorchMD-Net $P \in \mathbb{R}^{L_p \times d_p},$7 hours

The paper also reports post hoc physical plausibility checks using OpenMM. Generated binding trajectories remain within approximately 2 Å structural RMSD of MD throughout the binding trajectory, show contact map overlap greater than 92% between generated and native binding-mode contacts, and remain within approximately 2 kcal/mol of MD-derived energies. This suggests that the generated trajectories are not only geometrically similar but also consistent with the underlying free energy surface as enforced by the MSM-based FES loss (Thakur et al., 7 Aug 2025).

The stated implications are in binding pose refinement, binding and unbinding pathway exploration, ligand optimization through perturbing ligand fragments and sampling effects on binding stability and dynamics, and use as a fast surrogate sampler in structure-based drug design. The limitations are also explicit. Solvent is not modeled explicitly as separate dynamic tokens, which can limit solvent-mediated interactions and hydration effects. Generalization across diverse systems is not yet fully explored because current training is centered on 3CQV+HEME. No explicit force field or DFT-level physics is integrated as a hard prior, although future integration of DFT-based active learning is mentioned. Finally, while MSM-based FES matching suggests reasonable equilibrium ensemble reproduction, detailed kinetic fidelity is not deeply analyzed. Proposed future directions include explicit or coarse-grained solvent modeling, multi-ligand and multi-protein systems, diverse target classes such as membrane receptors, more direct integration of physics through DFT-informed corrections or energy guidance during diffusion, hybrid MD–ML schemes, and incorporation of experimental constraints such as cryo-EM maps and NMR restraints into conditioning (Thakur et al., 7 Aug 2025).