Title: Sampling 3D Molecular Conformers with Diffusion Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: 3D Molecular Graphs

Molecular Modeling

Generative Modeling

Diffusion Transformers

Geometric Deep Learning

Flow Matching
Abstract: Diffusion Transformers (DiTs) have demonstrated strong performance in generative modeling, particularly in image synthesis, making them a compelling choice for molecular conformer generation. However, applying DiTs to molecules introduces novel challenges, such as integrating discrete molecular graph information with continuous 3D geometry, handling Euclidean symmetries, and designing conditioning mechanisms that generalize across molecules of varying sizes and structures. We propose DiTMC, a framework that adapts DiTs to address these challenges through a modular architecture that separates the processing of 3D coordinates from conditioning on atomic connectivity. To this end, we introduce two complementary graph-based conditioning strategies that integrate seamlessly with the DiT architecture. These are combined with different attention mechanisms, including both standard non-equivariant and SO(3)-equivariant formulations, enabling flexible control over the trade-off between between accuracy and computational efficiency.
Experiments on standard conformer generation benchmarks (GEOM-QM9, -DRUGS, -XL) demonstrate that DiTMC achieves state-of-the-art precision and physical validity. Our results highlight how architectural choices and symmetry priors affect sample quality and efficiency, suggesting promising directions for large-scale generative modeling of molecular structures. Code is available at [https://github.com/ML4MolSim/dit_mc](https://github.com/ML4MolSim/dit_mc).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work modifies and applies the diffusion transformers architecture to the task of molecular conformer generation, i.e., generating low-energy geometric configurations for a fixed molecular structure formula. In order to translate the DiT architecture to this task, the authors introduce a molecular graph-based conditioning scheme as well as three positional encodings that are ablated. The authors claim state-of-the-art accuracy for their resulting DiTMC model on GEOMQM9 and GEOMDRUGS.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n* The authors did a great job of communicating their contributions and documenting their methods and experiments.\n* The downstream application is straightforward and a valuable one to solve. \n* The proposed DiTMC seems to perform well on the selected tasks.\n\nWeaknesses:\n* The main weakness of this work is its limited novelty. Transferring a model with minimal adjustments from one modality to another does not, in my humble opinion, clear the bar for NeurIPS. Though I reckon that this point is subjective, and acknowledge if other reviewers disagree.\n* The experiment sections are inconsistent in that while rPE in all experiments within the main body and the appendix is the best performing method, it does not occur in most tables.\n* The empirical gains compared to previous works are present but limited.\n* The related work section is rudimentary and misses the field of Boltzmann generators. \n\nThis work touches many interesting discussion points within the AI4Science community, but leaves them without clear conclusions. I highly encourage the authors to take the opportunity and reframe their work as a thorough investigation of these topics, such as equivariance, positional encoding, cutoffs, etc., by performing detailed empirical ablation studies.\n\nQUESTIONS:\nWhy is rPE frequently omitted?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents DiTMC, a novel framework for adapting the powerful Diffusion Transformer (DiT) architecture to the task of 3D molecular conformer generation. The authors identify and address key challenges in this adaptation, namely the integration of discrete molecular graph information with continuous 3D geometry and the handling of Euclidean symmetries. The core of their contribution lies in a modular DiT architecture combined with two new graph-based conditioning strategies (bond-pair and all-pair) that effectively inject atomic connectivity information into the self-attention mechanism. The framework is designed to systematically explore the impact of different symmetry priors by comparing standard non-equivariant attention mechanisms with a fully SO(3)-equivariant formulation. Experiments on the GEOM-QM9 and GEOM-DRUGS benchmarks show that DiTMC variants achieve state-of-the-art results in both precision and physical validity, providing valuable insights into the trade-off between architectural choices, computational efficiency, and sample quality.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1/ Systematic and Rigorous Investigation: A key strength of this paper is its systematic and rigorous approach. Instead of merely presenting a single new model, the authors design a modular framework that allows them to ablate and compare different architectural choices, particularly regarding symmetry handling (aPE vs. rPE vs. PE(3)). This comparative study provides a clear and insightful analysis of the trade-offs between model accuracy, computational cost, and the explicit enforcement of physical priors.\n\n2/ Novel and Effective Conditioning: The proposed conditioning strategies are well-designed for the DiT architecture. The \"all-pair\" conditioning, which uses graph geodesic distances to inform the attention mechanism about global connectivity, is particularly smart and is shown empirically to be the most effective strategy. This represents a non-trivial and successful adaptation for incorporating graph structure into a Transformer.\n\n3/ Strong Empirical Performance: The paper demonstrates state-of-the-art results on established and challenging benchmarks.\n\nWeaknesses\n\n1/ Computational Cost of Equivariance: While the analysis of the cost-fidelity trade-off is a strength, the high computational cost of the fully equivariant model is a practical weakness. The paper reports that the PE(3) model is approximately 3-3.5x slower than its non-equivariant counterparts, even with a smaller model size. The authors give a dicussion about this in the manuscript, and it sounds reasonable.\n\nQUESTIONS:\nThe paper uses two variants of model size, \"Base\" and \"Large\", demonstrating a clear trade-off: the non-equivariant aPE model scales efficiently in terms of computation and achieves state-of-the-art precision, while the fully-equivariant PE(3) model offers higher fidelity at a significant computational cost. Which direction is more promising when scaling this framework up to much larger and more complex systems? Please comment.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors extend the DiT architecture to irregular molecular geometries. Key moves are (i) two graph-based conditioning schemes (bond-pair vs. all-pair geodesic tokens) and (ii) interchangeable self-attention/positional-embedding blocks that range from standard (non-equivariant) to full SO(3)-equivariant forms. On GEOM-QM9, -DRUGS and -XL, DiTMC obtain state-of-the-art precision and comparable recall results while remaining computationally tractable, especially when using simple absolute PEs. Ablations quantify the fidelity/cost trade-off of enforcing Euclidean symmetries. Limitations include slower training for equivariant variants and restricted assessment to small/medium molecules.\n\nSTRENGTHS AND WEAKNESSES:\nStrength\n1. **Clear, modular architecture**\n    \n    By cleanly separating coordinate processing from graph conditioning, the authors can tease apart how positional encodings and symmetry priors influence results, turning these architectural knobs into concrete design guidelines.\n    \n2. **High accuracy with compact models**\n    \n    A 9.6 M-parameter DiTMC reaches 0.080 Å AMR-P and 95.7 % COV-P on GEOM-QM9, while a 77 M-parameter aPE-L variant secures state-of-the-art precision on GEOM-DRUGS—matching or outpacing far larger baselines.\n    \n3. **Comprehensive ablation suite**\n    \n    Systematic experiments on conditioning schemes, positional-encoding types, and model capacity lend strong empirical backing to the paper’s claims.\n    \n4. **Pragmatic symmetry–cost trade-offs**\n    \n    The study shows full equivariance incurs a 3–3.5× compute overhead yet yields only modest accuracy gains beyond ~0.5 Å, arming practitioners with data to balance fidelity against resources.\n\nWeakness\n\n1. **Incremental novelty**\n    \n    The core idea—plugging flow-matching/diffusion into a transformer with graph conditioning—extends recent DiT/Flow-Matching work\n    \n2. **Dependence on Training Data Quality**:\n    \n    DITMC’s performance relies heavily on the quality and diversity of training datasets (e.g., GEOM-QM9, GEOM-DRUGS). Limited or biased datasets may lead to poor generalization, especially for rare or complex molecular structures.\n    \n3. **Generalization to Large and Complex Molecules**:\n    \n    While effective on benchmark datasets, DITMC’s performance on larger, more complex molecules (e.g., proteins or macromolecules) is not fully validated due to dataset constraints. The modular architecture may struggle with the increased structural diversity and conformational flexibility of such systems.\n\nQUESTIONS:\n1. **Data Quality Sensitivity**: The paper demonstrates strong performance on high-quality datasets like GEOM. Could the authors discuss how the model's performance might degrade when trained on lower-quality or noisy conformer data? Are there any built-in robustness mechanisms?\n2. **Computational Cost Comparison**: While the inference times for DiTMC variants are provided, how do these compare to other state-of-the-art methods (e.g., MCF, ET-Flow) in terms of FLOPs, memory, or wall-clock time? A direct comparison would help assess practical trade-offs.\n3. **Scalability to Large Systems**: The experiments focus on small to medium-sized molecules. How does DiTMC scale to very large systems (e.g., >400 atoms), such as proteins or polymers? Are there architectural or sampling adjustments needed to maintain efficiency?\n\nLIMITATIONS:\n1. **Incremental Conceptual Advance**\n    \n    The work mainly repackages existing diffusion/flow-matching ideas within a transformer framework. While the engineering is solid, it introduces no fundamentally new generative principle, limiting its conceptual impact.\n    \n2. **Heavy Reliance on High-Quality Conformer Labels**\n    \n    Training hinges on well-curated datasets such as GEOM-QM9 and GEOM-DRUGS. In domains where reference conformers are scarce, noisy, or biased, model fidelity and transferability can degrade sharply.\n    \n3. **Restricted Molecular Scope**\n    \n    All experiments target small- to medium-sized, neutral organic molecules (≤ ~100 atoms). Performance on larger, highly flexible, charged, or metal-containing species—including peptides, macrocycles, and organometallics—remains untested.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a diffusion transformer (DiT) learning system for generating ensembles of 3D conformers for small molecules. Algorithmic components are introduced to embed a full suite of molecular features within transformer architectures, including atomic and pairwise graph embeddings. Additionally, multiple 3D positional embedding schemes are tested, including absolute, relative, and equivariant-pairwise, which are incorporated into self-attention-based updates to the molecular structure along diffusion trajectories. Altogether, the method introduces graph-conditioned molecular structure optimization under a DiT framework. Results are presented for low-energy conformer generation in two standard geometry datasets and are compared with a handful of SOTA methods from the literature.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. Generating ensembles with high realism is significant \n2. The formulation of the transformer architecture is elegant and extensively described (Figure 1 and Sec 4), w/ conditioning signals from full-graph embedding\n3. Including all-pair representations using the graph geodesic is nice, and seeing that this conditioning performs best is significant\n4. It is satisfying that many positional embedding schemes were tested\n5. Results for ensemble properties (Table 3) are compelling, though in E.3 it is described that this is after conformer relaxation with GFN2-xTB (this may be standard, but is still a limitation for practical use)\n\nWeaknesses:\n1. Motivation: I'm not sure that the “encod[ing of] molecular connectivity and incorporat[ing] Euclidean symmetries, such as translational and rotational invariance/equivariance” is such a major challenge as posed, though, again, the formulation of the architectural components is extensive. Equivariant transformer architectures, e.g., Equiformer have existed for several years specifically for molecular generation, which incorporate many of the components necessary. Equivariant diffusion mechanisms for 3D molecule generation, e.g., EDM, GeoLDM, GeoBFN and the like have also been well demonstrated for conformer generation, though not necessarily conditioned on graph structure as proposed or with transformers (that I know of), which is nice and, yes, challenging. Equivariant, graph-conditioned conformer generation architectures do exist, e.g., MoleCLUEs, though again not with DiTs.\n2. The formulation modeling a velocity, i.e., vector field looks very similar to neural force fields, which is not super novel, though again the utilization of DiTs for this task is interesting\n3. More details in the main text on how ensembles are generated, e.g., via rollout of different trajectories, would be appreciated, as would discussion or experimentation on the effects of noise scheduling in the diffusion process, as this is expected to play a significant role in model fidelity\n\nQUESTIONS:\n1. Was a pairwise-$\\textit{distance}$ positional embedding considered? \n2. I must have missed somewhere, what is used as the prior ($t=0$) for each sample, i.e., where is the initial conformer obtained from?\n3. How are ensembles obtained? I.e., what does the rollout of parallel trajectories entail?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to adapt the diffusion transformer, which has shown promise in general domains, to the task of molecular conformer prediction. To this end, several modifications are proposed within a new framework called DiTMC, including two types of molecular graph conditioning strategies. Moreover, the paper provides a detailed and informative discussion on the use of geometric priors (equivariant attention) versus standard self-attention for molecular conformation prediction. By combining conditioning strategies, attention mechanisms, and positional embedding, the proposed DiTMC has demonstrated effectiveness on conformation prediction tasks.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. Studying the inductive bias for molecular data is a very important topic. The paper aims to investigate the challenges of applying a widely adopted and proven framework, DiT, to molecular data, which should be encouraged.\n2. The paper provides a detailed comparison between equivariant and non-equivariant approaches. I appreciate the authors' effort to keep the methods simple and effective. I believe this discussion could inspire further exploration in model architecture design.\n3. Overall, the paper is clearly written, making the key elements easy to understand.\n\n**Weaknesses**\n\n1. The paper could be further improved by making the study more systematic. The scope of experiments and tasks is somewhat limited. As a new framework, many of the proposed elements—such as conditioning strategies, positional embedding, and attention—are not necessarily restricted to the task of conformation prediction. As the focus is on model architecture, I believe a more comprehensive evaluation would be beneficial, for example, on 3D molecule generation or molecular property prediction tasks.\n2. A key concern lies in the experimental results. While DiTMC-L, as claimed, achieves competitive performance with ~30M parameters (I highly recommend mentioning the parameter size in the main text instead of the appendix for clearer comparison), there is not a significant performance boost over the ET-Flow - SS baseline (8.3M) on Geom-Drugs. This raises questions about the effectiveness of the proposed approaches.\n3. The scalability of the proposed architecture is not clear. According to Fig. 3, increasing the number of parameters by four times does not yield the expected performance improvement, which may raise concerns about the model’s scalability.\n4. The technical novelty is limited. While I appreciate the efforts to adapt DiT to molecular graphs, from a methodological perspective, the contribution is relatively minor. This sets a high bar for assessing the completeness and effectiveness/performance of the method.\n\nQUESTIONS:\n1. Could the authors provide a more complete study over the scaling of the proposed approach by considering more model size? \n\n2. Could the authors provide a comparison between the base model and ET-Flow - SS baseline (8.3M)  over Geom-Drugs?\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}