PaperlayerView exact revision on arXiv

Source-linked AI summary

NEAT-POCKET: Pocket-Conditioned Autoregressive 3D Molecular Generation with a Neighborhood-Guided Set Transformer

Roxane Axel Jacob, Daniel Rose, Thierry Langer, Johannes Kirchmair

arXiv:2609.05097v1cs.LGcs.AI

TL;DR

Early-stage drug discovery needs efficient ways to generate 3D ligands tailored to protein binding pockets, where existing screening and transport-based generation can be computationally demanding. NEAT-POCKET extends autoregressive NEAT with pocket conditioning and achieves competitive benchmark performance with substantially faster sampling, while also supporting fragment completion. Its scope remains bounded by autoregressive error propagation, dataset artifacts, rigid-pocket assumptions, and omitted water molecules.

  • Problem

    Large-library screening and transport-based 3D generation can be computationally expensive, motivating efficient de novo design of ligands tailored to protein pockets.

  • Method

    NEAT-POCKET extends autoregressive NEAT to generate molecules atom by atom in protein pockets while preserving permutation-invariant set modeling and using pocket conditioning.

  • Results

    NEAT-POCKET achieved competitive structure-based generation performance on CrossDocked and SPINDR while sampling faster than existing baselines.

  • Takeaways & Limitations

    The framework supports pocket-conditioned fragment completion without architectural changes, retraining, or task-specific modifications, including workflows for fragment growing, scaffold decoration, and lead optimization.

  • Takeaways & Limitations

    Performance is constrained by autoregressive error propagation, biases and artifacts in training data, rigid-pocket modeling, and the omission of explicit water molecules.

Abstract

from arXiv · show

AI-driven de novo molecular design offers a promising route to accelerate early-stage drug discovery by generating novel ligands directly within target protein binding pockets. We present NEAT-POCKET, a pocket-conditioned extension of the autoregressive NEAT model for 3D molecular generation. NEAT-POCKET generates molecules atom by atom in protein pocket environments while preserving atom permutation invariance and explicitly modeling hydrogen atoms. Benchmarks on the CrossDocked and SPINDR datasets show that NEAT-POCKET achieves competitive structure-based generation performance while sampling substantially faster than existing baselines. Beyond full-molecule generation, NEAT-POCKET naturally enables pocket-conditioned fragment completion, a task directly relevant to lead optimization and scaffold elaboration. These results position NEAT-POCKET as a fast, flexible, and practical framework for structure-based drug design.

1 Introduction

Structure-based molecular generation aims to design 3D ligands that fit protein pockets and form favorable interactions, but transport-based methods can be computationally expensive. NEAT-POCKET extends an efficient autoregressive approach to pocket-conditioned generation while supporting broad evaluation and fragment-based workflows.

  • Ultra-large-library screening is computationally expensive and constrained by the size and diversity of available compounds or fragments.
  • 3D generation is valuable because protein–ligand binding depends on molecular conformation, motivating pocket-specific ligand design.
  • Transport-based generators repeatedly evaluate equivariant architectures over many atom pairs, creating substantial sampling latency.
  • NEAT constructs molecules atom by atom, encoding the partial molecule once per step before predicting atom types and coordinates with lightweight heads.
  • NEAT-POCKET extends this autoregressive framework to protein-binding pockets and is reported as competitive with diffusion- and flow-based methods while providing a substantial inference-time speed-up.
  • The work benchmarks pocket-conditioned methods across molecular validity, target compatibility, training-data similarity, and sampling efficiency, while supporting fragment growing, scaffold decoration, and lead optimization.

2 Related works

NEAT-POCKET extends autoregressive NEAT for pocket-conditioned 3D molecular generation, combining hierarchical pocket encoding with cross-attention, adaptive layer normalization, classifier-free guidance, and clash control. Its autoregressive design also supports fragment completion for lead optimization, addressing limitations of empty-pocket generation and scaffold-specific or hydrogen-implicit conditioning approaches.

  • 2 Related works: NEAT-POCKET targets pocket-conditioned 3D molecular generation while supporting fragment completion for lead optimization.The model generates ligands within protein pockets and completes molecules from partial molecular prefixes.
  • 2 Related works: Autoregressive prefix conditioning is suited to lead optimization, whereas scaffold-trained and inpainting-based alternatives impose training-distribution or inference constraints.FLOWR is identified as the closest existing competitor for fragment conditioning because it includes fragment-conditioned training examples and supports explicit hydrogen atoms.
  • 2 Related works: NEAT-POCKET encodes pocket information with a fine-coarse-fine transformer and integrates it into molecular generation through cross-attention and adaptive layer normalization.Cross-attention supplies atom-level conditioning, while adaptive layer normalization supplies global pocket-level information.
  • 2 Related works: The NEAT base model uses neighborhood guidance to map permutation-invariant connected subgraphs to next-atom predictions without canonical atom orderings.During inference, the representation generates the next atom type and 3D position autoregressively.
  • 2 Related works: Classifier-free guidance controls adherence to pocket conditioning by combining conditional and unconditional model outputs during inference.Conditional inputs are randomly dropped for 10–20% of training samples, and the guidance factor controls exploration versus conditional exploitation.
  • 2 Related works: The pocket-conditioned encoder can produce a representation from an empty source set using a learned start token, while later model components remain unchanged.The start token initializes pocket-conditioned autoregressive inference, and the representation is passed to subsequent layers.

4 Methods

The study pre-trains NEAT on GEOM-Drugs, fine-tunes separate pocket-conditioned models on CrossDocked and SPINDR, and evaluates generated ligands across validity, physical quality, docking, distributional similarity, and sampling time. Comparisons use established baselines with consistent hydrogen preprocessing where needed.

  • 4 Methods: NEAT-POCKET is initialized from NEAT pre-trained on GEOM-Drugs and fine-tuned separately on CrossDocked and SPINDR.The reported fine-tuning schedules are 2,000 epochs on CrossDocked and 5,000 epochs on SPINDR.
  • 4 Methods: CrossDocked uses 100,000 protein–ligand pairs for training and 100 proteins from remaining sequence clusters for testing.The split is derived from a refined 184,057-point subset clustered at 30% sequence identity.
  • 4 Methods: SPINDR is used as a refined alternative to CrossDocked because rigid-pocket cross-docking creates unrealistic ligand–pocket interactions.SPINDR derives from PLINDER with filtering intended to improve structure quality and reduce information leakage.
  • 4 Methods: Evaluation measures PoseBusters validity, protein–ligand clashes, ligand strain energy, docking scores before and after minimization, physicochemical similarity, and sampling time.Sampling time is measured as the average time to generate 100 molecules for one protein pocket.
  • 4 Methods: CrossDocked models are compared with Pocket2Mol, TargetDiff, DiffSBDD, and DrugFlow, while SPINDR models are compared with FLOWR.Hydrogens are added to baselines without explicit hydrogen generation using a consistent preprocessing protocol, and baseline models are not retrained.

5 Results

NEAT-POCKET delivers competitive pocket-conditioned molecular generation with strong validity, low clash counts, close physicochemical matching, and substantially faster sampling. It also supports pocket-aware guidance and fragment completion, while revealing trade-offs between pocket adaptation, molecular quality, and diversity.

  • CrossDocked: 69.2% PB validity and 6.6 average protein–ligand clashes make NEAT-POCKET competitive while generating compact, low-clash ligands.It ranked second in PB validity and best in clash avoidance among generative models.
  • CrossDocked: 2.48 descriptor rank indicates that NEAT-POCKET most closely reproduces CrossDocked physicochemical properties among evaluated baselines.It ranked first for several atom and stereochemical composition descriptors, although its molecules contained fewer rings than the dataset.
  • CrossDocked: 4 seconds per 100 molecules makes NEAT-POCKET approximately 20 times faster than DiffSBDD and more than an order of magnitude faster than the remaining baselines.DiffSBDD required 88 seconds on average for the same generation workload.
  • SPINDR: 80.3% versus 75.8% PB validity and 3.7 versus 4.6 clashes favor NEAT-POCKET over FLOWR on SPINDR.FLOWR nevertheless achieved lower strain energy and more favorable docking scores, while NEAT-POCKET generated smaller molecules and was nine times faster.
  • Representative generations: NEAT-POCKET generates pocket-compatible ligands that recover many reference interactions while sampling distinct new chemotypes.Examples in the DNA photolyase pocket reproduce hydrogen-bond, ionic, and hydrophobic contacts observed for the crystallographic ligand.
  • Classifier-free guidance: 8.3 to 4.5 CrossDocked clashes and 4.8 to 2.3 SPINDR clashes decrease as classifier-free guidance increases.Raw Vina scores also improve with stronger guidance, but PB validity falls to 47.9% and 55.2% at ω = 2.5 while strain energy rises.
  • Prefix-conditioned generation: Fragment completion is supported directly, but larger prefixes trade higher-quality constrained completions for reduced structural diversity.This capability is relevant to fragment growing, scaffold decoration, and lead optimization workflows.

6 Conclusions

NEAT-POCKET is a fast, flexible framework for pocket-conditioned 3D molecular generation, combining competitive benchmark performance with fragment completion and controllable conditioning. Its efficiency and flexibility support rapid exploration in hit-generation and lead-optimization workflows.

  • 6 Conclusions: NEAT-POCKET combines pocket-conditioned generation, explicit hydrogen modeling, controllable conditioning, and fragment-based completion.Pocket information is incorporated through cross-attention and global adaptive layer normalization, while classifier-free guidance controls conditioning strength.
  • 6 Conclusions: NEAT-POCKET achieved competitive performance across CrossDocked and SPINDR, including strong validity, low clash counts, and physicochemical similarity to training distributions.On CrossDocked it achieved the second-highest PoseBusters validity and lowest average protein–ligand clashes; on SPINDR it outperformed FLOWR on validity and clashes.
  • 6 Conclusions: Larger prefixes generally improved validity, strain energy, and docking performance but reduced structural diversity by restricting accessible chemical space.The model supports arbitrary fragment conditioning without architectural changes, retraining, or task-specific modifications.
  • 6 Conclusions: 4 seconds were required to generate 100 molecules for a pocket, approximately nine times faster than FLOWR and more than an order of magnitude faster than the remaining baselines.This efficiency matters when exploring many pockets, molecular prefixes, and optimization hypotheses.
  • 6 Conclusions: Overall, NEAT-POCKET provides a fast and practical approach for early-stage hit generation and lead optimization workflows.The authors position rapid exploration of chemically plausible design hypotheses as the relevant use case.

Conflicts of interest

The authors report no conflicts of interest.

  • Conflicts of interest: No conflicts of interest were declared.

Notes and references

The supplied section contains bibliographic references supporting the paper's methodological and related-work context.

  • Notes and references: The section lists references cited throughout the paper.

S1 Implementation details

The implementation and evaluation pipeline combines specified hardware and software, standardized sampling and hydrogenation procedures, and method-specific handling of molecular connectivity and filtering. Baseline inference protocols were modified to enable fixed-budget comparisons across methods.

  • S1 Implementation details: Training and inference used separate Linux systems with NVIDIA GPUs, AMD CPUs, CUDA 13.0, Python 3.11.15, PyTorch 2.11.0, and PyTorch Geometric 2.8.0.The training system used an RTX Pro 6000 Blackwell GPU, while inference used an RTX 4090 GPU.
  • S1 Implementation details: The best-performing NEAT-POCKET model has 242M parameters and was optimized with AdamW, gradient clipping, cosine annealing, and linear warm-up.The reported optimizer parameters are β1 = 0.9 and β2 = 0.95, with a maximum gradient norm of 1.0.
  • S1 Implementation details: The bond-prediction model infers molecular connectivity from generated atom types and positions using a lightweight E(3)-invariant gated point network.It uses three convolutional layers and hidden dimension 64, incorporating distances and angles through an E(3)-invariant architecture.
  • S1 Implementation details: Evaluation was standardized because pocket-conditioned generation protocols vary in sampling, downstream processing, docking software, hydrogen treatment, and connectivity handling.The study used Gnina 1.1 for docking and generated exactly 100 molecules per pocket for each method.
  • S1 Implementation details: Hydrogen-implicit baselines received hydrogen atoms with RDKit before downstream evaluation, whereas each method's original bond-inference strategy was retained.The hydrogenation procedure used Chem.AddHs(addCoords=True), and connectivity handling was not otherwise altered.
  • S1 Implementation details: Baseline inference procedures were modified to impose fixed generation budgets rather than selectively filtering or extending sampling until successful molecules were obtained.Pocket2Mol was stopped after 100 completion events without duplicate removal, while DiffSBDD used one batch of exactly 100 molecules.
  • S1 Implementation details: Inpainting evaluation used the authors' dedicated script with BRICS-derived test-set fragments, hydrogen removal for DiffSBDD, sanitization, and filtering of fragmented structures.

S3 Post hoc hydrogenation of CrossDocked data

CrossDocked largely omits explicit hydrogens, making post hoc reconstruction chemically difficult and motivating direct hydrogen modeling. Among tested approaches, RDKit coordinate addition was the most practical compromise despite known failures and dataset limitations.

  • Only 4,320 hydrogen atoms remain among 100,100 combined CrossDocked ligands, approximately 0.2% of all ligand atoms.This hydrogen depletion affects the datasets used by many pocket-conditioned generation models.
  • Hydrogens influence ligand shape, hydrogen bonding, protonation, tautomerism, formal charge, and donor or acceptor assignments relevant to molecular recognition.
  • Post hoc hydrogenation depends on chemically consistent bond orders, aromaticity, tautomeric states, valences, and charges, which CrossDocked may not reliably provide.Only 12 CrossDocked ligands contained formal charges, all corresponding to zwitterionic N-oxide annotations.
  • RDKit Chem.AddHs(addCoords=True) can fail or place hydrogens implausibly when heavy-atom representations have incorrect valences, aromaticity, or protonation states.N-containing heteroaromatic rings are a common source of such errors.
  • Open Babel exhibits similar limitations on chemically inconsistent inputs, while constrained force-field embedding can perturb original heavy-atom coordinates.
  • RDKit Chem.AddHs(addCoords=True) provided the most practical tested compromise, but post hoc hydrogenation remains a workaround rather than a principled solution.The authors identify explicit-hydrogen datasets with curated bond orders, charges, and protonation states as the appropriate long-term direction.

S4 Evaluation metrics

The evaluation aggregates molecule-level measurements within pockets, reports means with Student’s t confidence intervals, and combines validity, interaction, energy, docking, and physicochemical metrics. These measures assess both chemical plausibility and pocket-conditioned pose quality.

  • Per-pocket metric means are computed first, followed by aggregation across pockets and reporting with 95% Student’s t confidence intervals.The interval uses the sample standard deviation and the 97.5th percentile of the t distribution.
  • A molecule is PB valid when RDKit parses it successfully and it passes all 22 conditional PoseBusters checks.
  • PoseCheck measures steric clashes, hydrogen-bond and hydrophobic interactions, van der Waals contacts, and ligand strain energy.Ligand strain energy compares the generated conformation with its UFF-optimized geometry.
  • Gnina provides Vina and minimized Vina scores, alongside average molecular weight because docking scores depend strongly on ligand size.
  • The physchem rank compares generated molecules with CrossDocked across 31 physicochemical descriptors, with lower ranks indicating better performance.The descriptors include general, ring-size-related, and element-type-related properties.

S5 Additional results

Additional analyses examine interaction quality and strain-energy robustness across pockets. Strain-energy distributions reveal greater extreme-value sensitivity for several transport-based models on CrossDocked, while NEAT-POCKET shows a broader upper tail on SPINDR.

  • Table S2 compares favorable protein–ligand interactions, SA score, and related quality metrics across NEAT-POCKET and five pocket-conditioned generation baselines.The interaction categories include hydrogen-bond acceptors, hydrogen-bond donors, hydrophobic interactions, and van der Waals contacts.
  • Strain-energy summaries report means, medians, maxima, outlier counts, and values after removing extreme or detected outlier pockets.
  • Outliers are identified using modified z-scores on log10-transformed per-pocket average strain energies, with |z_i| > 3.5 as the threshold.
  • 1.5 × 10^8 versus 255, DrugFlow’s mean and median strain energies on CrossDocked show strong sensitivity to extreme pockets; removing the largest value reduces the mean to 14,422.
  • On CrossDocked, diffusion- and flow-based models are more prone to extremely large average strain energies, whereas Pocket2Mol and NEAT-POCKET have more stable profiles.
  • On SPINDR, NEAT-POCKET’s median strain energy is 62 kcal/mol versus FLOWR’s 35, while the means are 848 versus 48 kcal/mol.NEAT-POCKET has no flagged outlier pockets, and removing its largest pocket changes the mean only from 848 to 825 kcal/mol.

S7 Fragment generation

NEAT-POCKET is evaluated for completing molecules from fragments of varying sizes derived from reference ligands. BRICS fragments serve as molecular prefixes for autoregressive completion.

  • Reference ligands from the CrossDocked and SPINDR test sets are decomposed with BRICS, and the resulting fragments are used as molecular prefixes.
  • Three prefix sizes are tested for each ligand: the largest fragment, the second-largest fragment, and the smallest fragment.NEAT-POCKET is then tasked with completing the molecule from each prefix.

S8 Examples of generated molecules

This section presents examples of molecules generated by the evaluated models, showing each model in both 3D and 2D representations. The examples include baseline methods and NEAT-POCKET trained on CrossDocked or SPINDR.

  • Models evaluated: Examples are provided for Pocket2Mol, TargetDiff, DiffSBDD, DrugFlow, FLOWR, and NEAT-POCKET.The section includes generated-molecule examples for each evaluated model.
  • Visualization: Each model’s generated molecules are plotted in both 3D and 2D.The figures pair 3D and 2D visualizations for the listed models.
  • NEAT-POCKET training datasets: NEAT-POCKET examples are shown separately for models trained on CrossDocked and SPINDR.CrossDocked examples appear in Figures S15–S16, while SPINDR examples appear in Figures S19–S20.
  • Reference-ligand context: The supplementary figures also report absolute and relative fragment-size distributions for CrossDocked and SPINDR reference ligands.These distributions are shown separately for the two datasets.
Loading 2609.05097v1…