UniMoMo: Unified Generative Modeling of 3D Molecules for De Novo Binder Design

Thanks for your appreciation and the positive comments!

Q1: How did you balance data from different all three domains during training?

Thanks for the question! In the current implementation, we have not extensively explored domain-specific data balancing strategies. For our joint training approach, we utilized representative datasets from each domain: CrossDocked (~100k samples) for small molecules, PepBench (~80k samples) for peptides, and SAbDab (~10k samples) for antibodies. As an initial effort toward unified molecular modeling, we employed simple random sampling across these datasets during training, which has demonstrated promising results. Moving forward, particularly as we scale the framework to incorporate larger and more diverse datasets, we recognize the importance of investigating more sophisticated data balancing approaches to further enhance model performance.

Thanks for your appreciation and the constructive comments!

Q1: The paper only evaluates on in-silico metrics, which are known to be far from perfect. The proposed method do achieve good results on these metrics (compared to baselines benchmarked).

Thanks for the comments! While we are currently conducting wet-lab validations, the experimental timeline is extensive. Therefore, we plan to include these results in a future extension of this work.

Q2: During sampling time, how many "blocks" nodes are chosen before starting the (reverse) diffusion process? How does this choice affects experimental results (on all modalities)?

For peptides and antibodies, we follow the literature [a,b] and set the number of blocks equal to that of the reference sequence, as many metrics, including AAR and RMSD, assume generated sequences having the same lengths as the native binders.

For small molecules, the number of blocks is sampled from the statistical distribution based on the spatial size of the pocket [c]. We evaluated different settings on the GPCR case in Section 4.4, averaging results over 100 designs, where $n$ denotes the originally sampled number of blocks for small molecules:

Commonly, larger binders lead to lower energies, since they can form more interactions. Interestingly, for small molecules, the model adapts to the number of blocks. With fewer blocks, it generates larger fragments with more atoms to fill the pocket; otherwise, with more blocks, it generates smaller fragments, avoiding overcrowding the pocket.

Thanks for the insightful question again! We will include the discussion in the revision!

[a] Full-Atom Peptide Design based on Multi-modal Flow Matching. ICML 2024.

[b] Conditional Antibody Design as 3D Equivariant Graph Translation. ICLR 2023.

[c] CBGBench: Fill in the Blank of Protein-Molecule Complex Binding Graph. ICLR 2025.

Q3: I feel some information is missing on how the bonds are computed. Could the authors elaborate more on how the bonds between bonds and between blocks are computed? It seems that the latter depends on a NN prediction. How accurate is it?

Sorry for the confusion. For intra-block bonds, they are predetermined once the block type is assigned. For inter-block chemical bonds, we employ an MLP for prediction, using the hidden states of atom pairs as input (Eq. 6). Bond prediction is restricted to spatially neighboring atoms, excluding those that are too far apart. The predictions are dynamically changing during the flow matching process in the decoder (Figure 1B, bottom right). Ultimately, the reconstruction accuracy of the chemical bonds are around 97%, which we think is accurate enough.

Q4: How are the sequences of AAs (on both peptides and antibodies) extracted from the full atom point cloud?

Sorry for the confusion. In our block-level decomposition, each natural amino acid forms a block. Therefore, we can directly get the amino acid type of each block. We think this is also one merit of our block-based unified representation, without requirements to further derive an algorithm to extract the AAs.

Q5: What parts of the CDR are modelled on the Antibody experiments? More experimental details on this section is needed.

Sorry for the confusion. We follow the convention [d, e] and evaluate on CDR-H3, as it exhibits much higher irregularities than other CDRs and plays a crucial role in binding and interactions [f]. We appreciate the suggestion and will clarify this in the revision.

[d] End-to-End Full-Atom Antibody Design. ICML 2023.

[e] GeoAB: Towards Realistic Antibody Design and Reliable Affinity Maturation. ICML 2024.

[f] Antigen-Specific Antibody Design and Optimization with Diffusion-Based Generative Models for Protein Structures. NeurIPS 2022.

Q6: During sampling time, given a target pocket, how does the model decide when it generates blocks that belond to small molecules, or peptides or antibodies?

Thanks for the insightful question! We apologize for not making this clear in the paper. We use a binary prompt for each block, with 1 indicating the generation of amino acid (AA), and 0 indicating either AAs or molecular fragments. During benchmarking for peptides and antibodies, all blocks are assigned a prompt of 1 to ensure AA generation. For small molecules, we set the prompt to 0, allowing flexible generation of molecular fragments. We will add a detailed explanation in the revision to clarify this point.

Thanks for your appreciation and insightful feedback, which is very helpful in improving our paper!

W & Q2: If my understanding is correct, the method only works if the binding site on the target protein is known. What if we do not know the binding site?

Yes, your understanding is correct. Our method requires prior knowledge of the binding site, in line with the convention established in previous studies [a], since most public benchmarks are built upon this setting. This is also biologically reasonable. For example, to study a protein-protein interaction, a binder is often designed at the interface to inhibit the PPI. If truely no biological knowledge is available, pocket detection tools like MaSIF [b] should be able to identify potential sites for pocket-based design models.

[a] Pocket2mol: Efficient molecular sampling based on 3d protein pockets. ICML 2022

[b] Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods

C: Equation (13) is missing the square root over $1-\bar{\alpha}^t$.

Thanks for catching this! We apologize for the typo and will correct it in the revision

Q1: How exactly is the invariance of the latents guaranteed, as well as the equivariance of the coordinates?

For the diffusion part, GeoDiff [c] proves that the diffusion process maintain equivariance via an equivariant denoising kernel (Proposition 1), which in our work is the equivariant transformer[d]. The transformer is scalar-based[e], which maintain the equivariance with inner product and unbiased linear transformation on the velocities.

For the VAE part, the encoding equivariance is determined by the network used, which is also an equivariant transformer. The decoder employs a short flow matching, with a similar theoretical foundation as GeoDiff on equivariant flow matching [f] (Theorem 4.1).

Thanks for pointing this out! We will add a detailed discussion in the revision for clarity.

[c] GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation. ICLR 2022.

[d] An Equivariant Pretrained Transformer for Unified 3D Molecular Representation Learning. preprint.

[e] Scalars are universal: Equivariant machine learning, structured like classical physics. NeurIPS 2021.

[f] Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation. NeurIPS 2023.

Q3: Why not use the original building block coordinates $\vec{X}_i$, together with the other latents $z_i$ to encode all the additional information? Related, an ablation study over the KL weights $\lambda_1$ and $\lambda_2$ in equation (2) would be interesting.

Thanks for the valuable question. The block coordinates $\vec{X}_i\in\mathbb{R}^{n_i\times 3}$ vary in length due to different number of atoms per block (e.g. residue), which makes direct diffusion challenging, as it typically operates in a fixed-length space. A similar issue arises for $H_i$. Thus, we use an all-atom VAE to compress these irregular matrices into fixed-length latent vectors $z_i$ and $\vec{z}_i$, making diffusion feasible.

Regarding $\lambda_1$ and $\lambda_2$, this is a really insightful question. A higher $\lambda$ smooths the latent space, improving continuity for later generative modeling but also increasing compression, which may limit expressivity. Thus, the weight can neither be too low nor be too high, as shown in the validation loss curves. Ultimately, we select the combination with the lowest validation loss.

Q4: The authors manually add extra noise to $\vec{z}_i$ before feeding it to the decoder, to enforce robustness. A more principled way would be to increase the KL weight.

Thanks for the suggestion. The extra noises primarily benefits small molecules, which have more intricate inter-block connections than peptides and antibodies connected by peptide bonds, thus require higher robustness to coordinate errors introduced by diffusion. Given that the weight of KL loss is already high (0.8), further increasing it would overly constrain the latent space (as in Q3 above and the table below).

Q5: Are "motion vectors" the vector field in the flow matching framework?

Yes, we will replace them with "vector fields" in the revision to avoid ambiguity.

Q6：How exactly is the model told to either generate a peptide, a small molecule, or an antibody for the different applications?

Thanks for the insightful question! We apologize for not making this clear in the paper. We assign a binary prompt to each block, with 1 indicating amino acid (AA) and 0 without restriction. For peptide and antibody benchmarks, all blocks are assigned 1 to ensure AA generation. For small molecules, we use 0 to allow arbitrary fragment generation. We will clarify this in the revision.