Sesame: Structure-Aware Molecular Generation via Spatial Density-Map Conditioning

# Sesame: Structure-Aware Molecular Generation via Spatial Density-Map Conditioning
Source: [https://arxiv.org/html/2606.23856](https://arxiv.org/html/2606.23856)
Konstantin Yatsenko, Arvind Thiagarajan Tessel Biosciences, Inc\., 750 Main Street, Cambridge, MA 02139

###### Abstract

Generative molecular models for drug design are a promising direction with much active research\. In the next phase of computational drug design, such models will need to understand small molecule structure and protein\-ligand interactions, and they will need to possess the machinery to generate moleculesde novo\. Incorporating each feature poses a critical challenge\. Equally important, yet often treated as secondary, is the ability to grow a molecule from a partial starting point – a scaffold or fragment supplied by a chemist – which is the central operation of lead optimization\. We present Sesame \(Spatial Evoformer for a Structure\-Aware Molecular Engine\), a diffusion\-based molecular generation model that leverages a novel spatial pairformer module to condition on partial molecular structure and the surrounding protein pocket, both expressed as continuous spatial density maps\. This single conditioning mechanism supports bothde novogeneration and fragment\-conditioned lead optimization, letting a medicinal chemist prune a hit to a scaffold and have Sesame grow it in productive ways\. In addition to this module, we also introduce a diffusion framework for joint denoising of atom types, bond types, and positions, along with a trajectory finetuning scheme that trains on the model’s own sampling rollouts to improve generation quality\. Sesame is trained on a large corpus of ligand\-only and protein\-ligand datasets\.

Keywords:Molecular generation, Diffusion models, Structure\-aware design, Density maps, Drug discovery

![Refer to caption](https://arxiv.org/html/2606.23856v1/figures/diffusion_series.png)Figure 1:Snapshot of the diffusion process, left\-to\-right\. Here, we denoise a ligand with fragment conditioning, allowing the model to fill in the gaps and extend the existing core\.## 1Introduction

### 1\.1Problem Statement

Modern protein\-folding models possess capabilities to reason about protein\-ligand interactions in a meaningful way, with good accuracy across a number of tasks\. However, these models are fundamentally constrained by knowledge about a specific ligand, often passed as hard constraints for the generation algorithm\. These architectures are incompatible with the generation ofde novomolecules, which means that any system using these models for drug discovery must carefully balance exploration of chemical space with computational costs and constraints\. Furthermore, these protein\-ligand models are not perfect, and can often generate structures with small inconsistencies, e\.g\. incorrect enantiomers\. Our approach aims to learn from successes in the protein\-folding space, extending them to perform full molecular generation\.

Beyondde novogeneration, a second capability is equally central to real drug\-discovery campaigns and is often underserved by generative models: lead optimization, in which a known hit is refined rather than designed from scratch\. In practice this is where human expertise is most valuable – a medicinal chemist can identify the substructure of a hit worth preserving and prune the molecule down to that scaffold\. A useful generative model should treat such a fragment as a soft prior to be grown and elaborated, not as a rigid atom\-level constraint, so that human insight and generative chemistry compose rather than compete\. We therefore treat fragment\-conditioned generation as a primary objective alongsidede novodesign, and, as described below, realize both through a single conditioning mechanism\.

### 1\.2Motivation

Our approach is motivated by three key insights\. First, encoding the local structure of a protein by estimating various forces on a dense grid is a common technique, and largely used in various docking algorithms\. These so\-called density maps provide a natural representation for 3D structural information, encoding multiple physical properties \(charge, hydrophobicity, hydrogen bonding, van der Waals interactions\) in a unified grid\-based format\. This allows for modeling diverse sets of protein pockets with identically\-sized inputs, which confers a number of computational efficiencies\. Second, Pairformer architectures perform admirably for a variety of tasks that require complex multi\-way interactions between components\. An extension to this architecture to allow for information to pass from the density map to the prediction is an essential component to a model whose inputs are density maps\. Third, diffusion models whose goal is to produce molecules require a diffusion approach that handles both discrete atom types and continuous atomic coordinates\. Many naive approaches fail to take into account biases that trained models produce, and a careful examination of the reverse diffusion process is necessary to have a robust generative model\.

### 1\.3Contributions

Our main contributions are:

1. 1\.Novel density map conditioning architecture: An attention\-based mechanism that adaptively samples from density maps to guide molecular generation, enabling structure\-aware generation\.
2. 2\.Unified fragment and pocket conditioning for lead optimization: Because both the protein pocket and a partial ligand are supplied to the model as the same continuous density field, a single mechanism supportsde novogeneration and fragment\-conditioned growth with no architectural changes\. This makes lead optimization a first\-class capability: a medicinal chemist can prune a hit molecule down to a scaffold, and Sesame grows chemically sensible, pocket\-compatible completions around it – a direct integration of human insight with generative chemistry\. Across fragment\-conditioned generation, 94\.8% of generated molecules retain the seeding fragment as a substructure, confirming that the conditioning is honored rather than ignored\.
3. 3\.Hybrid discrete\-continuous diffusion: A unified diffusion process that handles discrete atom types and bond types via categorical diffusion and continuous coordinates via Gaussian diffusion\.
4. 4\.Trajectory finetuning: A self\-distillation finetuning scheme that rolls out the model’s own reverse\-diffusion trajectories, re\-matches the predicted atoms to ground truth, and trains on the resulting \(noisy state, aligned target\) pairs, improving the quality of fully generated molecules\.

## 2Related Work

##### Generating 3D molecules\.

Denoising generative models operating directly on 3D molecular structure have become the dominant paradigm for de novo design\. Equivariant diffusion models established the template by jointly denoising atomic coordinates and types under E\(3\)/SE\(3\) symmetry\(Hoogeboomet al\.,[2022](https://arxiv.org/html/2606.23856#bib.bib8)\), with subsequent work unifying the diffusion of continuous coordinates and discrete atom and bond features\(Vignacet al\.,[2023b](https://arxiv.org/html/2606.23856#bib.bib15)\)\. A complementary line replaces denoising diffusion with flow matching, trading stochastic reverse\-time dynamics for deterministic transport and typically improving sampling efficiency and physical quality\(Irwinet al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib10); Dunn and Koes,[2026](https://arxiv.org/html/2606.23856#bib.bib5)\)\. Most of these methods target generic small molecules – frequently benchmarked on QM9\-scale data – rather than drug\-like chemistry per se\(Hoogeboomet al\.,[2022](https://arxiv.org/html/2606.23856#bib.bib8); Vignacet al\.,[2023b](https://arxiv.org/html/2606.23856#bib.bib15); Irwinet al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib10)\), and they share a structural limitation relevant to our setting: the number of atoms is fixed prior to sampling, so the model cannot grow or trim a molecule to fit a target\. A distinct architectural lineage comes not from molecule generation but from biomolecular structure prediction: AlphaFold3 couples a Pairformer trunk to a diffusion module that denoises raw atomic coordinates\(Abramsonet al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib17)\), and the open Boltz models adopt and extend this co\-folding architecture\(Wohlwendet al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib18); Passaroet al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib19)\)\. While these systems predict the structure of a given complex rather than generating novel ligands, their pairwise\-representation\-plus\-diffusion\-decoder design directly informs the architecture of Sesame\.

##### Spatial\-field and voxel representations\.

A smaller body of work represents molecules not as point sets but as continuous spatial maps – most directly, voxel grids denoised in a learned density space, for both unconditional 3D generation\(Pinheiroet al\.,[2024b](https://arxiv.org/html/2606.23856#bib.bib12)\)and structure\-based design\(Pinheiroet al\.,[2024a](https://arxiv.org/html/2606.23856#bib.bib13)\)\. This framing is the closest precedent for our conditioning scheme, in which both the pocket and the \(partial\) ligand are expressed as maps of physical potentials\. ShEPhERD similarly operates over interaction fields – shape, electrostatic potential, and pharmacophores – generating molecules matched to a target profile\(Adamset al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib1)\), a representation that is conceptually the inverse of our pocket\-conditioned setup\.

##### Conditioning on scaffolds and fragments\.

Several methods generate molecules from a fixed structural prior – decorating a scaffold, linking fragments, or growing from a seed\. DiffLinker designs linkers conditioned on disconnected 3D fragments\(Igashovet al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib9)\); DiffDec performs structure\-aware scaffold decoration end\-to\-end\(Xieet al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib16)\); and 3D\-Scaffold generates 3D coordinates of drug\-like molecules around a specified scaffold\(Joshiet al\.,[2021](https://arxiv.org/html/2606.23856#bib.bib11)\)\. These correspond to the second conditioning mode in Sesame, in which a full or partial map of a scaffold or fragment seeds generation; unlike methods that fix the substructure as discrete atoms, Sesame supplies it as the same continuous field used for the pocket, allowing the prior to be partial and the surrounding chemistry to be grown to fit\.

##### Generative models for drug discovery\.

The methods above are largely concerned with generation in broad chemical space; a parallel body of work shares our specific goal of discovering drug\-like hits against a target\. For target\-aware generation, diffusion models conditioned on the protein binding site generate ligands directly within the pocket: TargetDiff jointly models ligand atoms and their geometry conditioned on the protein context and additionally estimates binding affinity\(Guanet al\.,[2023](https://arxiv.org/html/2606.23856#bib.bib7)\), while DiffSBDD applies equivariant diffusion to structure\-based design across pocket\-conditioned generation and inpainting tasks\(Schneuinget al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib14)\)\. More recent systems extend pocket conditioning to flow matching with multi\-objective guidance\(Cremeret al\.,[2024](https://arxiv.org/html/2606.23856#bib.bib2)\)and to unified interaction\- and fragment\-based generation\(Cremeret al\.,[2026](https://arxiv.org/html/2606.23856#bib.bib3)\)\. Where an explicit pocket is unavailable, ShEPhERD pursues the same target\-driven goal by conditioning on the interaction profile of a known ligand and generating bioisosteric replacements\(Adamset al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib1)\)\. A second, persistent gap between these generative models and practical hit finding is synthesizability: structurally valid molecules are not necessarily makeable\. A growing line of work addresses this by generating directly in synthesizable chemical space – composing molecules from purchasable building blocks and reaction templates rather than from atoms\(Gaoet al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib6)\), or by amortizing the search over synthetic routes with generative flow networks\(Cretuet al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib4)\)\. Sesame differs from the pocket\-conditioned methods in that the pocket enters as a spatial potential map rather than as an explicit atomistic context, and pocket conditioning is optional rather than required; synthesizability\-aware generation we view as complementary to spatial\-field conditioning rather than competing with it, composable as a downstream filter or alternative decoder\. Across these lines of work, conditioning on a spatial field, growing from a partial fragment, and targeting a protein pocket have been pursued largely in isolation; Sesame is designed to support all three within a single generative process\.

## 3Methods

### 3\.1Overall Architecture

Our training pipeline consists of several key components: data generation from multiple sources, density map computation, model forward pass, diffusion process, and loss computation\.

The training process begins with data generation from two sources: ZINC \(ligand\-only\) and SAIR \(protein\-ligand pairs\)\. Each sample undergoes preprocessing to generate density maps and extract molecular graphs\. During training, the model receives noisy atom types and positions from the forward diffusion process, along with density maps and conditioning signals\. The model predicts denoised atom types, positions, and bonds, which are compared to ground truth through a multi\-component loss function\.

### 3\.2MoleculePairformer Architecture

Table 1:Model dimensions and hyperparameters#### 3\.2\.1Core Components

TheMoleculePairformer111TheMoleculePairformerfollows the Evoformer/Pairformer lineage of AlphaFold; architecturally it is a pairformer – operating on single and pair representations with no MSA track – and we refer to it as such throughout\. The “Evoformer” in Sesame’s expansion \(Spatial Evoformer for a Structure\-Aware Molecular Engine\) reflects this lineage\.operates on two main representations \(Figure[2](https://arxiv.org/html/2606.23856#S3.F2)A\):

- •Single representation𝐬∈ℝN×ds\\mathbf\{s\}\\in\\mathbb\{R\}^\{N\\times d\_\{s\}\}: Per\-atom features of dimensionds=384d\_\{s\}=384, whereNNis the total potential number of atoms\.
- •Pair representation𝐩∈ℝN×N×dp\\mathbf\{p\}\\in\\mathbb\{R\}^\{N\\times N\\times d\_\{p\}\}: Atom\-pair features of dimensiondp=128d\_\{p\}=128\.

Atom types are represented as 16 discrete classes corresponding to \(element, number of implicit hydrogens\) pairs that occur in drug\-like molecules, plus one additional pseudo\-None class, givingA​T=17AT=17total atom types\. The 16 drug\-like classes are: \(C,0\), \(C,1\), \(C,2\), \(C,3\), \(N,0\), \(N,1\), \(N,2\), \(O,0\), \(O,1\), \(F,0\), \(P,0\), \(S,0\), \(S,1\), \(Cl,0\), \(Br,0\), \(I,0\)\. These cover the vast majority of heavy atoms in drug\-like molecules, with implicit hydrogen count encoding local valence context without requiring explicit hydrogen atoms in the model\. One key feature of the molecular generation process is that the model needs to predict not only which atom types exist where in space, but also how many atoms are needed for a given molecule\. Put differently, the model may output a maximum ofNNatoms, but must predict the pseudo\-None type for all atoms that do not exist in the final molecule\.

The model uses multiple embedding layers:

1. 1\.Atom type embedding: Linear layer mapping fromA​TATnoised atom types todsd\_\{s\}dimensions\.
2. 2\.Atom position embedding: Fourier embedding of atomic positions, followed by a linear projection todsd\_\{s\}dimensions\.
3. 3\.Bond type embedding: Linear layer mapping fromB​TBTnoised bond types todpd\_\{p\}dimensions, added to the position pair embedding\.
4. 4\.Position pair embedding: Fourier embedding of relative atomic positions, followed by a linear projection todpd\_\{p\}dimensions\. Diagonal elements receive the atom position instead of the \(zero\) relative position\.
5. 5\.Time embeddings: Three separate Fourier embeddings of the diffusion timestepτ∈\[0,1\]3\\tau\\in\[0,1\]^\{3\}, one per noise channel \(coordinate, atom type, bond type\), each projected todtd\_\{t\}dimensions\.

The Fourier embedding for a scalarttis defined as:

Fourier​\(t\)=2​cos⁡\(2​π​\(t​𝐰\+𝐛\)\)\\text\{Fourier\}\(t\)=\\sqrt\{2\}\\cos\(2\\pi\(t\\mathbf\{w\}\+\\mathbf\{b\}\)\)\(1\)where𝐰\\mathbf\{w\}and𝐛\\mathbf\{b\}are learnable parameters\.

Notably, for both the position and position pair embeddings, all 3 spatial dimensions are embedded with the same𝐰\\mathbf\{w\}and𝐛\\mathbf\{b\}\. The linear projection that follows does distinguish between the 3 dimensions, though, so the symmetry between dimensions ends there\.

The atom type and atom position embeddings are summed to form the initial single representation𝐬\\mathbf\{s\}\. The position pair embedding and bond type embedding are summed to form the initial pair representation𝐩\\mathbf\{p\}\. The coordinate and atom type time embeddings are each incorporated into the single representation via separateGatedResiduallayers, while the bond type time embedding is incorporated into the pair representation via aGatedResiduallayer\.

Finally, each initial6×32×32×326\\times 32\\times 32\\times 32density map \(protein and ligand\) is processed into add×12×12×12d\_\{d\}\\times 12\\times 12\\times 12volume via a series of 3D convolutions and MaxPool layers\.

#### 3\.2\.2Pairformer Layers

EachMoleculePairformerlayer consists of several components applied in sequence:

Density Map Conditioning Operations\(Figure[2](https://arxiv.org/html/2606.23856#S3.F2)C\): First, we generate sampling points via attention on the single representation\. Rather than self\-attention, we use a per\-layer learned key vector of shapeH×O×dhH\\times O\\times d\_\{h\}, whereOOis the number of sampling points emitted per attention head\. These then attend to the query vectors derived from the single representation\. The value vectors \(of shapeH×3H\\times 3, whereHHis the number of heads, typically 4\) are also derived from the single representation\. WithO=256O=256sampling points per head andH=4H=4heads, this yields10241024sampling points in total\.

The sampling points are used to extract features from the density map via trilinear interpolation:

𝐅density=GridSample​\(𝐃protein,𝐏sample\)\\mathbf\{F\}\_\{\\text\{density\}\}=\\text\{GridSample\}\(\\mathbf\{D\}\_\{\\text\{protein\}\},\\mathbf\{P\}\_\{\\text\{sample\}\}\)\(2\)where𝐃protein∈ℝdd×12×12×12\\mathbf\{D\}\_\{\\text\{protein\}\}\\in\\mathbb\{R\}^\{d\_\{d\}\\times 12\\times 12\\times 12\}is the post\-3D convnet density map\.

Theddd\_\{d\}density features are concatenated with sampling point coordinates and used in cross\-attention to update single and pair representations\.

Queries are single\- and pair\-derived vectors, while key and value vectors are derived from the sampled points\. Key vectors are shared between both representations, while value vectors are computed separately\.

Finally, these attention terms are added back to the model via gated residual connections\. Details of this module can be found in Appendix[A](https://arxiv.org/html/2606.23856#A1)\.

Triangle Operations: Following AlphaFold’s Evoformer architecture, we use triangle multiplication and triangle attention operations\. Triangle multiplication computes:

𝐳i​j=∑k𝐚i​k​𝐛j​k\\mathbf\{z\}\_\{ij\}=\\sum\_\{k\}\\mathbf\{a\}\_\{ik\}\\mathbf\{b\}\_\{jk\}\(3\)for outgoing triangles, and similarly for incoming triangles\.

Triangle Attention: Applies attention along one dimension of the pair representation:

Attention​\(𝐐,𝐊,𝐕\)=softmax⁡\(𝐐𝐊Tdk\+𝐁\)​𝐕\\text\{Attention\}\(\\mathbf\{Q\},\\mathbf\{K\},\\mathbf\{V\}\)=\\operatorname\{softmax\}\\left\(\\frac\{\\mathbf\{Q\}\\mathbf\{K\}^\{T\}\}\{\\sqrt\{d\_\{k\}\}\}\+\\mathbf\{B\}\\right\)\\mathbf\{V\}\(4\)where𝐁\\mathbf\{B\}is a pair bias learned from the pair representation\.

Attention with Pair Bias: The single representation is updated via attention, with bias terms from the pair representation:

𝐬′=𝐬\+Attention​\(𝐬,𝐬,𝐬;bias=𝐩\)\\mathbf\{s\}^\{\\prime\}=\\mathbf\{s\}\+\\text\{Attention\}\(\\mathbf\{s\},\\mathbf\{s\},\\mathbf\{s\};\\text\{bias\}=\\mathbf\{p\}\)\(5\)
Feed\-Forward Networks: We use SwiGLU\(Shazeer,[2020](https://arxiv.org/html/2606.23856#bib.bib21)\)activation:

SwiGLU​\(𝐱\)=SiLU​\(𝐖1​𝐱\)⊙\(𝐖2​𝐱\)\\text\{SwiGLU\}\(\\mathbf\{x\}\)=\\text\{SiLU\}\(\\mathbf\{W\}\_\{1\}\\mathbf\{x\}\)\\odot\(\\mathbf\{W\}\_\{2\}\\mathbf\{x\}\)\(6\)
All operations are wrapped with residual connections and pre\-layer normalization\. Figure[2](https://arxiv.org/html/2606.23856#S3.F2)B illustrates the layer structure\.

![Refer to caption](https://arxiv.org/html/2606.23856v1/x1.png)Figure 2:TheMoleculePairformerarchitecture\.\(A\)Noisy atoms, positions, and bonds are embedded into single and pair representations, conditioned on noise level and encoded protein/ligand density maps, and passed throughL=24L=24Pairformer layers before the atom\-type, coordinate, and bond output heads\.\(B\)A single Pairformer layer applies density\-map conditioning, then triangle and feed\-forward updates on the pair track and pair\-biased attention on the single track\.\(C\)Density\-map conditioning: the single representation predicts query points that sample the density volume, whose features then update both tracks via cross\-attention\.
#### 3\.2\.3Output Heads

The model produces three outputs:

Atom Type Head: Predicts logits overA​TATatom types:

𝐀pred=MLP​\(𝐬\)\\mathbf\{A\}\_\{\\text\{pred\}\}=\\text\{MLP\}\(\\mathbf\{s\}\)\(7\)
Coordinate Head: Predicts 3D coordinates:

𝐂pred=MLP​\(𝐬\)\\mathbf\{C\}\_\{\\text\{pred\}\}=\\text\{MLP\}\(\\mathbf\{s\}\)\(8\)
Bond Head: Predicts bond types between atom pairs:

𝐁pred=Symmetric​\(MLP​\(𝐩\)\)\\mathbf\{B\}\_\{\\text\{pred\}\}=\\text\{Symmetric\}\(\\text\{MLP\}\(\\mathbf\{p\}\)\)\(9\)whereSymmetricensures the bond matrix is symmetric \(bond between atomsiiandjjequals bond betweenjjandii\)\.

### 3\.3Diffusion Process

Molecular generation requires handling both discrete atom types and continuous atomic coordinates\. We introduce a novel diffusion process that jointly models these two modalities within a unified framework\.

The forward diffusion process adds noise independently to each modality, while the reverse process uses joint model predictions to denoise both simultaneously\.

#### 3\.3\.1Forward Process

The forward process defines a distributionq​\(𝐂τ,𝐀τ,𝐁τ\|𝐂0,𝐀0,𝐁0,τ\)q\(\\mathbf\{C\}\_\{\\tau\},\\mathbf\{A\}\_\{\\tau\},\\mathbf\{B\}\_\{\\tau\}\|\\mathbf\{C\}\_\{0\},\\mathbf\{A\}\_\{0\},\\mathbf\{B\}\_\{0\},\\tau\)over noisy coordinates𝐂τ\\mathbf\{C\}\_\{\\tau\}, noisy atom types𝐀τ\\mathbf\{A\}\_\{\\tau\}, and noisy bond types𝐁τ\\mathbf\{B\}\_\{\\tau\}, given true atom coordinates𝐂0\\mathbf\{C\}\_\{0\}, true atom types𝐀0∈\{0,…,A​T−1\}N\\mathbf\{A\}\_\{0\}\\in\\\{0,\\ldots,AT\-1\\\}^\{N\}, true bond types𝐁0∈\{0,…,B​T−1\}N×N\\mathbf\{B\}\_\{0\}\\in\\\{0,\\ldots,BT\-1\\\}^\{N\\times N\}, at timestepτ∈\[0,1\]3\\tau\\in\[0,1\]^\{3\}, whereτ​\[0\]\\tau\[0\]controls coordinate noise,τ​\[1\]\\tau\[1\]controls atom type noise, andτ​\[2\]\\tau\[2\]controls bond type noise, each sampled independently during training\. Atτ=\(1,1,1\)\\tau=\(1,1,1\), all modalities are considered fully noised\.

##### Coordinate diffusion

For coordinates, we perform a relatively simple set of operations\. First, we add noise sampled from𝒩​\(0,σ2​\(τ​\[0\]\)​𝐈\),\\mathcal\{N\}\\left\(0,\\sigma^\{2\}\(\\tau\[0\]\)\\mathbf\{I\}\\right\),whereσ:\[0,1\]→\[smin,smax\]\\sigma:\[0,1\]\\to\[s\_\{\\min\},s\_\{\\max\}\]is a power\-law function, withp=7,p=7,smin=0\.01​Å,s\_\{\\min\}=0\.01\\text\{\\AA \},andsmax=8​Ås\_\{\\max\}=8\\text\{\\AA \}\.

σ​\(τ​\[0\]\)=\(smax1/p\+\(1−τ​\[0\]\)​\(smin1/p−smax1/p\)\)p\\sigma\(\\tau\[0\]\)=\\left\(s\_\{\\max\}^\{1/p\}\+\(1\-\\tau\[0\]\)\(s\_\{\\min\}^\{1/p\}\-s\_\{\\max\}^\{1/p\}\)\\right\)^\{p\}\(10\)Because coordinate noise is added independently to each atom, a freshly noised cloud is not arranged the way the model expects to denoise it: the point that lands near a given ground\-truth atom is not necessarily the one sharing its index\. We therefore re\-pair the noised points to the ground\-truth atoms by minimum total displacement, using a modified Jonker\-Volgenant assignment, so that the noisy state is already in correspondence with the target the model predicts\. Importantly, this re\-pairing does not move any points; as an unordered cloud the noised positions are unchanged, and the assignment only selects which ground\-truth atom each point denoises toward\.

Concretely, let𝐂~τ=𝐂0\+ϵ\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}=\\mathbf\{C\}\_\{0\}\+\\boldsymbol\{\\epsilon\}withϵ∼𝒩​\(0,σ2​\(τ​\[0\]\)​𝐈\)\\boldsymbol\{\\epsilon\}\\sim\\mathcal\{N\}\(0,\\sigma^\{2\}\(\\tau\[0\]\)\\mathbf\{I\}\)be the noised positions\. These are sampled independently, giving the product form

q​\(𝐂~τ∣𝐂0,τ\)=∏i=1N𝒩​\(𝐂~τ​\[i\];𝐂0​\[i\],σ2​\(τ​\[0\]\)​𝐈\)\.q\(\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}\\mid\\mathbf\{C\}\_\{0\},\\tau\)=\\prod\_\{i=1\}^\{N\}\\mathcal\{N\}\\\!\\left\(\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}\[i\];\\,\\mathbf\{C\}\_\{0\}\[i\],\\,\\sigma^\{2\}\(\\tau\[0\]\)\\mathbf\{I\}\\right\)\.\(11\)We then compute the pairwise distance cost matrixCosti​j=‖𝐂~τ​\[i\]−𝐂0​\[j\]‖2\\text\{Cost\}\_\{ij\}=\\\|\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}\[i\]\-\\mathbf\{C\}\_\{0\}\[j\]\\\|\_\{2\}and find the assignment minimizing total displacement:

π∗=arg​minπ∈SN​∑i=1N‖𝐂~τ​\[π​\(i\)\]−𝐂0​\[i\]‖2,\\pi^\{\*\}=\\operatorname\*\{arg\\,min\}\_\{\\pi\\in S\_\{N\}\}\\sum\_\{i=1\}^\{N\}\\\|\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}\[\\pi\(i\)\]\-\\mathbf\{C\}\_\{0\}\[i\]\\\|\_\{2\},\(12\)which pairs each ground\-truth atomiiwith the noised pointπ∗​\(i\)\\pi^\{\*\}\(i\)it should denoise toward\. The matched coordinates are𝐂τ​\[i\]=𝐂~τ​\[π∗​\(i\)\]\\mathbf\{C\}\_\{\\tau\}\[i\]=\\tilde\{\\mathbf\{C\}\}\_\{\\tau\}\[\\pi^\{\*\}\(i\)\], and the discrete targets are permuted by the same assignment \(𝐀0\\mathbf\{A\}\_\{0\}reindexed byπ∗\\pi^\{\*\}, and the rows and columns of𝐁0\\mathbf\{B\}\_\{0\}permuted identically\), so that type and bond identity remain consistent with the re\-paired positions\. The discrete corruptions are then applied to these permuted targets:

q​\(𝐀τ∣𝐀0,π∗,τ\)\\displaystyle q\(\\mathbf\{A\}\_\{\\tau\}\\mid\\mathbf\{A\}\_\{0\},\\pi^\{\*\},\\tau\)=∏i=1Nq​\(aτ​\[i\]∣a0​\[i\],τ​\[1\]\),\\displaystyle=\\prod\_\{i=1\}^\{N\}q\(a\_\{\\tau\}\[i\]\\mid a\_\{0\}\[i\],\\tau\[1\]\),\(13\)q​\(𝐁τ∣𝐁0,π∗,τ\)\\displaystyle q\(\\mathbf\{B\}\_\{\\tau\}\\mid\\mathbf\{B\}\_\{0\},\\pi^\{\*\},\\tau\)=∏i=1N∏j=1Nq​\(bτ​\[i,j\]∣b0​\[i,j\],τ​\[2\]\),\\displaystyle=\\prod\_\{i=1\}^\{N\}\\prod\_\{j=1\}^\{N\}q\(b\_\{\\tau\}\[i,j\]\\mid b\_\{0\}\[i,j\],\\tau\[2\]\),\(14\)wherea0​\[i\]a\_\{0\}\[i\]andb0​\[i,j\]b\_\{0\}\[i,j\]denote the permuted ground\-truth types \(𝐀0​\[π∗​\(i\)\]\\mathbf\{A\}\_\{0\}\[\\pi^\{\*\}\(i\)\]and𝐁0​\[π∗​\(i\),π∗​\(j\)\]\\mathbf\{B\}\_\{0\}\[\\pi^\{\*\}\(i\),\\pi^\{\*\}\(j\)\]\), and the per\-atom and per\-pair discrete marginals are as defined in the paragraphs below\.

##### Atom Type Diffusion

We use a marginal diffusion scheme, following the framework of DiGress\(Vignacet al\.,[2023a](https://arxiv.org/html/2606.23856#bib.bib22)\)\. The atom type noise level is parameterized byτ​\[1\]∈\[0,1\]\\tau\[1\]\\in\[0,1\]via a cosine schedule:

α¯​\(τ​\[1\]\)=cos2⁡\(π​τ​\[1\]2\)\\bar\{\\alpha\}\(\\tau\[1\]\)=\\cos^\{2\}\\\!\\left\(\\frac\{\\pi\\,\\tau\[1\]\}\{2\}\\right\)\(15\)so thatα¯​\(0\)=1\\bar\{\\alpha\}\(0\)=1\(no noise\) andα¯​\(1\)=0\\bar\{\\alpha\}\(1\)=0\(fully noised\)\. The forward marginal for a single atom type is then:

q​\(aτ​\[1\]=k∣a0=i\)=α¯​\(τ​\[1\]\)​1​\[i=k\]\+\(1−α¯​\(τ​\[1\]\)\)​mk\(A\)q\(a\_\{\\tau\[1\]\}=k\\mid a\_\{0\}=i\)=\\bar\{\\alpha\}\(\\tau\[1\]\)\\,\\mathbf\{1\}\[i=k\]\+\(1\-\\bar\{\\alpha\}\(\\tau\[1\]\)\)\\,m\_\{k\}^\{\(A\)\}\(16\)wheremk\(A\)=1/A​Tm^\{\(A\)\}\_\{k\}=1/ATis the uniform distribution over atom types\. A uniform prior means that at full noise \(α¯=0\\bar\{\\alpha\}=0\) every atom type is equally likely, avoiding any systematic bias toward chemically common types during the noising process\.

For a pair of consecutive noise levelsτi​\[1\]\>τi−1​\[1\]\\tau\_\{i\}\[1\]\>\\tau\_\{i\-1\}\[1\], the effective one\-step corruption probability is:

βτi−1→τi=1−α¯​\(τi​\[1\]\)α¯​\(τi−1​\[1\]\)\\beta\_\{\\tau\_\{i\-1\}\\to\\tau\_\{i\}\}=1\-\\frac\{\\bar\{\\alpha\}\(\\tau\_\{i\}\[1\]\)\}\{\\bar\{\\alpha\}\(\\tau\_\{i\-1\}\[1\]\)\}\(17\)This ratio\-based formulation means no fixed discrete schedule is required: during training,τ​\[1\]\\tau\[1\]is sampled uniformly from\[0,1\]\[0,1\], and during generation a trajectory of sorted noise levels is used, withβ\\betaderived from consecutiveα¯\\bar\{\\alpha\}ratios at each step\.

##### Bond Type Diffusion

Bond types undergo the same marginal diffusion scheme as atom types, but use an independent noise coordinateτ​\[2\]∈\[0,1\]\\tau\[2\]\\in\[0,1\]with its own cosine scheduleα¯​\(τ​\[2\]\)=cos2⁡\(π​τ​\[2\]/2\)\\bar\{\\alpha\}\(\\tau\[2\]\)=\\cos^\{2\}\(\\pi\\tau\[2\]/2\)\. Each pairwise bond type is independently noised as:

q​\(bτ​\[2\]=k∣b0=i\)=α¯​\(τ​\[2\]\)​1​\[i=k\]\+\(1−α¯​\(τ​\[2\]\)\)​mk\(B\)q\(b\_\{\\tau\[2\]\}=k\\mid b\_\{0\}=i\)=\\bar\{\\alpha\}\(\\tau\[2\]\)\\,\\mathbf\{1\}\[i=k\]\+\(1\-\\bar\{\\alpha\}\(\\tau\[2\]\)\)\\,m\_\{k\}^\{\(B\)\}\(18\)wherem\(B\)m^\{\(B\)\}is a fixed prior over bond types, chosen to reflect the strong sparsity of molecular graphs:

m\(B\)=\[0,0\.0025,0\.0025,0\.0025,0\.0025,0\.99\]m^\{\(B\)\}=\[0,\\;0\.0025,\\;0\.0025,\\;0\.0025,\\;0\.0025,\\;0\.99\]\(19\)corresponding to types 0 \(unused\), 1 \(single\), 2 \(double\), 3 \(triple\), 4 \(aromatic\), and 5 \(None\) respectively\. The dominant weight on type 5 reflects the fact that the vast majority of atom pairs in a molecule are unbonded\. The small equal weights on types 1–4 provide a weak non\-zero prior over all bonded types, ensuring the posterior never assigns exactly zero probability to any valid bond during the reverse process\. Type 0 receives zero prior weight, as it is unused\.

The noised bond types𝐁τ\\mathbf\{B\}\_\{\\tau\}are provided as input to the model alongside the noised atom types𝐀τ\\mathbf\{A\}\_\{\\tau\}, encoded via the pair representation\.

#### 3\.3\.2Reverse Process

The reverse process defines a distributionp​\(𝐀τi−1,𝐁τi−1,𝐂τi−1\|𝐀τi,𝐁τi,𝐂τi,θ,d\)p\(\\mathbf\{A\}\_\{\\tau\_\{i\-1\}\},\\mathbf\{B\}\_\{\\tau\_\{i\-1\}\},\\mathbf\{C\}\_\{\\tau\_\{i\-1\}\}\|\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\theta,d\)over denoised atom type logits, bond type logits, and coordinates, given noisy inputs, model parametersθ\\theta, and density mapsdd\. The model predicts all three jointly, but the denoising processes are conditionally independent given these predictions:

p​\(𝐂τi−1,𝐀τi−1,𝐁τi−1∣𝐀τi,𝐁τi,𝐂τi,θ,𝐝\)\\displaystyle p\(\\mathbf\{C\}\_\{\\tau\_\{i\-1\}\},\\mathbf\{A\}\_\{\\tau\_\{i\-1\}\},\\mathbf\{B\}\_\{\\tau\_\{i\-1\}\}\\mid\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\theta,\\mathbf\{d\}\)=p​\(𝐀τi−1∣𝐀τi,𝐁τi,𝐂τi,θ,𝐝\)⋅p​\(𝐁τi−1∣𝐀τi,𝐁τi,𝐂τi,θ,𝐝\)⋅p​\(𝐂τi−1∣𝐀τi,𝐁τi,𝐂τi,θ,𝐝\)\\displaystyle\\quad=p\(\\mathbf\{A\}\_\{\\tau\_\{i\-1\}\}\\mid\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\theta,\\mathbf\{d\}\)\\cdot p\(\\mathbf\{B\}\_\{\\tau\_\{i\-1\}\}\\mid\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\theta,\\mathbf\{d\}\)\\cdot p\(\\mathbf\{C\}\_\{\\tau\_\{i\-1\}\}\\mid\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\theta,\\mathbf\{d\}\)\(20\)
##### Coordinates

The model predicts denoised positions𝐂pred=gθ​\(𝐀τi,𝐁τi,𝐂τi,τi,𝐝\)\\mathbf\{C\}\_\{\\text\{pred\}\}=g\_\{\\theta\}\(\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\tau\_\{i\},\\mathbf\{d\}\)\. Coordinates are updated deterministically, while the discrete atom and bond types are sampled from their reverse posteriors \(below\); the coordinate update follows the DDIM\-style probability\-flow rule\(Songet al\.,[2021](https://arxiv.org/html/2606.23856#bib.bib20)\)for a variance\-exploding schedule:

𝐂τi−1=𝐂pred\+σ​\(τi−1​\[0\]\)σ​\(τi​\[0\]\)​\(𝐂τi−𝐂pred\)\\mathbf\{C\}\_\{\\tau\_\{i\-1\}\}=\\mathbf\{C\}\_\{\\text\{pred\}\}\+\\frac\{\\sigma\(\\tau\_\{i\-1\}\[0\]\)\}\{\\sigma\(\\tau\_\{i\}\[0\]\)\}\(\\mathbf\{C\}\_\{\\tau\_\{i\}\}\-\\mathbf\{C\}\_\{\\text\{pred\}\}\)\(21\)
Note that since the model is trained to predict positions post\-matching, we assume that the model predictions already correspond to the correct input atoms, and do not perform any matching algorithm during a denoising step\.

##### Types

The model predictsx0x\_\{0\}logits𝐀pred=hθ​\(𝐀τi,𝐁τi,𝐂τi,τi,𝐝\)\\mathbf\{A\}\_\{\\text\{pred\}\}=h\_\{\\theta\}\(\\mathbf\{A\}\_\{\\tau\_\{i\}\},\\mathbf\{B\}\_\{\\tau\_\{i\}\},\\mathbf\{C\}\_\{\\tau\_\{i\}\},\\tau\_\{i\},\\mathbf\{d\}\)for atom types\. The reverse posterior is derived by Bayes’ rule and marginalizing overx0x\_\{0\}:

pθ​\(aτi−1=j∣aτi=k\)∝q​\(aτi=k∣aτi−1=j\)⋅∑i′pθ​\(a0=i′∣aτi\)q​\(aτi=k∣a0=i′\)⋅q​\(aτi−1=j∣a0=i′\)p\_\{\\theta\}\(a\_\{\\tau\_\{i\-1\}\}=j\\mid a\_\{\\tau\_\{i\}\}=k\)\\;\\propto\\;q\(a\_\{\\tau\_\{i\}\}=k\\mid a\_\{\\tau\_\{i\-1\}\}=j\)\\cdot\\sum\_\{i^\{\\prime\}\}\\frac\{p\_\{\\theta\}\(a\_\{0\}=i^\{\\prime\}\\mid a\_\{\\tau\_\{i\}\}\)\}\{q\(a\_\{\\tau\_\{i\}\}=k\\mid a\_\{0\}=i^\{\\prime\}\)\}\\cdot q\(a\_\{\\tau\_\{i\-1\}\}=j\\mid a\_\{0\}=i^\{\\prime\}\)\(22\)where all three terms have closed forms under the marginal diffusion structure \(Section[3\.3\.1](https://arxiv.org/html/2606.23856#S3.SS3.SSS1)\); in particular, the one\-step transitionq​\(aτi∣aτi−1\)q\(a\_\{\\tau\_\{i\}\}\\mid a\_\{\\tau\_\{i\-1\}\}\)is the marginal\-scheme transition with corruption probabilityβτi−1→τi\\beta\_\{\\tau\_\{i\-1\}\\to\\tau\_\{i\}\}defined there\. Critically, theq​\(aτi=k∣a0=i′\)q\(a\_\{\\tau\_\{i\}\}=k\\mid a\_\{0\}=i^\{\\prime\}\)denominator must remain inside the sum — factoring it out drops ani′i^\{\\prime\}\-dependent correction that systematically over\-concentrates probability mass on the most likely class, producing anO​\(K\)O\(K\)error on minority classes at low noise levels\.

In practice this reduces to aK×KK\\times Kmatrix multiply per atom: let𝐰=pθ​\(a0∣aτi\)/q​\(aτi∣a0\)\\mathbf\{w\}=p\_\{\\theta\}\(a\_\{0\}\\mid a\_\{\\tau\_\{i\}\}\)/q\(a\_\{\\tau\_\{i\}\}\\mid a\_\{0\}\)\(element\-wise, clamped for numerical stability\), and\[𝐐τi−1\]i′​j=q​\(aτi−1=j∣a0=i′\)\[\\mathbf\{Q\}\_\{\\tau\_\{i\-1\}\}\]\_\{i^\{\\prime\}j\}=q\(a\_\{\\tau\_\{i\-1\}\}=j\\mid a\_\{0\}=i^\{\\prime\}\); the unnormalized posterior is:

p~j∝q​\(aτi=k∣aτi−1=j\)⋅\(𝐰⊤​𝐐τi−1\)j\\tilde\{p\}\_\{j\}\\propto q\(a\_\{\\tau\_\{i\}\}=k\\mid a\_\{\\tau\_\{i\-1\}\}=j\)\\cdot\(\\mathbf\{w\}^\{\\top\}\\mathbf\{Q\}\_\{\\tau\_\{i\-1\}\}\)\_\{j\}\(23\)normalized overj∈\{0,…,A​T−1\}j\\in\\\{0,\\ldots,AT\-1\\\}, from whichaτi−1a\_\{\\tau\_\{i\-1\}\}is sampled\.

##### Bond Types

Bond types are denoised using the identical posterior formula, withm\(B\)m^\{\(B\)\},B​TBT,τ​\[2\]\\tau\[2\], and𝐁pred\\mathbf\{B\}\_\{\\text\{pred\}\}replacing their atom\-type counterparts\.

### 3\.4Training Methodology

#### 3\.4\.1Diffusion Training Objective

During initial training, the model learns to predict both atom types and positions by maximizing the likelihood of the data under the forward diffusion process\. The model predicts𝐂pred\\mathbf\{C\}\_\{\\text\{pred\}\},𝐀pred\\mathbf\{A\}\_\{\\text\{pred\}\}, and𝐁pred\\mathbf\{B\}\_\{\\text\{pred\}\}from noisy inputs\(𝐀τ,𝐁τ,𝐂τ\)\(\\mathbf\{A\}\_\{\\tau\},\\mathbf\{B\}\_\{\\tau\},\\mathbf\{C\}\_\{\\tau\}\)at a randomly sampled timestepτ\\tau\.

#### 3\.4\.2Loss Functions

Our loss function combines multiple components:

Atom Type Loss: Cross\-entropy with label smoothing \(smoothing factor 0\.02\):

ℒtype=CrossEntropy​\(𝐀pred,𝐀true;smoothing=0\.02\)\\mathcal\{L\}\_\{\\text\{type\}\}=\\text\{CrossEntropy\}\(\\mathbf\{A\}\_\{\\text\{pred\}\},\\mathbf\{A\}\_\{\\text\{true\}\};\\text\{smoothing\}=0\.02\)\(24\)
Bond Type Loss: Cross\-entropy on bonds:

ℒbond=CrossEntropy​\(𝐁pred,𝐁true;smoothing=0\.001\)\\mathcal\{L\}\_\{\\text\{bond\}\}=\\text\{CrossEntropy\}\(\\mathbf\{B\}\_\{\\text\{pred\}\},\\mathbf\{B\}\_\{\\text\{true\}\};\\text\{smoothing\}=0\.001\)\(25\)
Position Loss: Mean squared error:

ℒpos=1N​∑i=1N‖𝐂pred\(i\)−𝐂true\(i\)‖2\\mathcal\{L\}\_\{\\text\{pos\}\}=\\frac\{1\}\{N\}\\sum\_\{i=1\}^\{N\}\\\|\\mathbf\{C\}\_\{\\text\{pred\}\}^\{\(i\)\}\-\\mathbf\{C\}\_\{\\text\{true\}\}^\{\(i\)\}\\\|^\{2\}\(26\)
LDDT Loss: We use the same loss as many other protein\-folding models, in order to ensure good bond distances, bond angles, dihedral angles, and other multi\-atom distances\.

ℒlddt=1−1N​\(N−1\)​∑i≠j1\|𝒯\|​∑τ∈𝒯σ​\(τ−\|‖𝐂pred\(i\)−𝐂pred\(j\)‖−‖𝐂true\(i\)−𝐂true\(j\)‖\|\)\\mathcal\{L\}\_\{\\text\{lddt\}\}=1\-\\frac\{1\}\{N\(N\-1\)\}\\sum\_\{i\\neq j\}\\frac\{1\}\{\|\\mathcal\{T\}\|\}\\sum\_\{\\tau\\in\\mathcal\{T\}\}\\sigma\\\!\\left\(\\tau\-\\left\|\\\|\\mathbf\{C\}\_\{\\text\{pred\}\}^\{\(i\)\}\-\\mathbf\{C\}\_\{\\text\{pred\}\}^\{\(j\)\}\\\|\-\\\|\\mathbf\{C\}\_\{\\text\{true\}\}^\{\(i\)\}\-\\mathbf\{C\}\_\{\\text\{true\}\}^\{\(j\)\}\\\|\\right\|\\right\)\(27\)where𝒯=\{0\.5,1\.0,2\.0,4\.0\}\\mathcal\{T\}=\\\{0\.5,1\.0,2\.0,4\.0\\\}Å are distance thresholds andσ​\(x\)=1/\(1\+e−x\)\\sigma\(x\)=1/\(1\+e^\{\-x\}\)is the logistic sigmoid\.

The total loss is a simple unweighted combination:

ℒtotal=ℒtype\+ℒbond\+ℒpos\+ℒlddt\\mathcal\{L\}\_\{\\text\{total\}\}=\\mathcal\{L\}\_\{\\text\{type\}\}\+\\mathcal\{L\}\_\{\\text\{bond\}\}\+\\mathcal\{L\}\_\{\\text\{pos\}\}\+\\mathcal\{L\}\_\{\\text\{lddt\}\}\(28\)

#### 3\.4\.3Pretraining Procedure

Density Map Augmentation: During training, we apply augmentation to density maps:

- •50% dropout: Zero out the density map \(unconditioned generation\)
- •50% unchanged: Use original density map

This augmentation encourages the model to be robust to missing or noisy conditioning information\. It also ensures that the model does not over\-index on the existence of protein or ligand maps, as the information contained therein is non\-trivial\. Datapoints with no protein map necessarily come from the ligand\-only data source, while datapoints with no ligand map come from the protein\-ligand data source, but had an error during ligand map generation, which is strongly correlated with the presence of certain atoms, such as Boron or Silicon, for which the forcefield generation can fail\.

Optimization: We use AdamW optimizer with learning rate2⋅10−32\\cdot 10^\{\-3\}andϵ=10−6\\epsilon=10^\{\-6\}\. Notably, we use a higher\-than\-default weight decay of 0\.1\. Learning rate scheduling:

- •Linear warmup: 500 steps, start factor10−310^\{\-3\}
- •Cosine decay:Tm​a​x=T\_\{max\}=100k,ηm​i​n=10−3\\eta\_\{min\}=10^\{\-3\}

We used a batch size of 128, and an epoch size of 8192, leading to a set of validation metrics every 64 optimizer steps\.

Table 2:Training hyperparameters
#### 3\.4\.4Trajectory finetuning

To aid the model in creating high quality molecules during generation, we applied a fine\-tuning pass during which we perform the following procedure:

- •Sample a datapoint identically to that in pretraining\.
- •Noise that datapoint to a random noise level in the full noise space\.
- •Perform a 100\-step diffusion process with that starting noised datapoint, recording each intermediate state as a separate datapoint\.
- •Perform a matching algorithm from the ground truth to the final prediction, and use that permutation to re\-order the ground truth\.
- •Emit each pair of \(input, permuted ground truth\) as a datapoint for finetuning\.

The matching step is critical, because although the individual atom slots have fixed identity, as they were noised from a given input, often a pretrained model will choose to permute several slots early in the trajectory, and not recomputing the permutation leads to a misalignment between the training\-time prediction and the ground truth\.

For finetuning, we use a learning rate of2⋅10−6\.2\\cdot 10^\{\-6\}\.Other hyperparameters are the same as pretraining\.

### 3\.5Data Pipeline

#### 3\.5\.1Ligand\-Only Datasets

ZINC Database: We use the ZINC22 database\(Tingleet al\.,[2023](https://arxiv.org/html/2606.23856#bib.bib23)\), which contains over 37 billion ready\-to\-make commercially available compounds\. We have downloaded approximately 15 billion compounds for training, filtered to those with 3–50 heavy atoms\. For each compound, we generate a 3D conformer using RDKit, apply MMFF94 force field initialization, and compute a density map from the resulting forces and RDKit\-derived properties\.

ZINC provides diverse molecular structures without protein context, enabling the model to learn general molecular geometry and chemistry\.

#### 3\.5\.2Ligand\-Protein Datasets

SAIR Dataset: The Structurally\-Augmented IC50 Repository\(Lemoset al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib24)\)contains protein\-ligand complexes, predominantly co\-folded \(synthetic\) structures rather than experimental ones\. We process SAIR by:

- •Extracting binding pockets around ligands \(radius8​38\\sqrt\{3\}Å\)
- •Computing protein properties \(charge, hydrophobicity, H\-bond donors/acceptors, aromaticity, VdW radii\)
- •Generating density maps for both protein and ligand as above

#### 3\.5\.3Density Map Generation

Density maps are generated with a grid resolution of 0\.5 Å over a fixed cube size of32×32×3232\\times 32\\times 32voxels, or a cube with side length 16Åcentered at the molecule’s centroid\. We used 6 channels encoding charge, hydrophobicity, H\-bond donor, H\-bond acceptor, aromatic interactions, and van der Waals potential\.

For ligand density maps, an MMFF94 forcefield was generated using RDKit, and the parameters of that forcefield were used to determine the parameters for the density map generation below\. Hydrophobicity channel parameters were determined by using RDKit’s Crippen contributions per\-atom\. Aromaticity was determined from automatic recognition by RDKit, and H\-bond acceptor and donors were detected with SMARTS patterns\.

For protein density maps, each protein heavy atom had a hardcoded list of the parameters, with no adjustments made for their environment\.

The density map generator places atomic contributions onto the grid using Gaussian kernels for the hydrophobicity,π\\pi\-stacking, and H\-bond channels, with channel\-specific kernel widths for the different interaction types\. The electrostatic and van der Waals potentials were computed using their standard analytic formulas\. For the van der Waals channel, a post\-processing transform of a scaled tanh was applied for numerical stability\.

### 3\.6Data Augmentation

#### 3\.6\.1Rotational Augmentation

Random 3D rotations are applied to both ligands and proteins, ensuring rotational invariance\. Rotations are generated using QR decomposition of random matrices, ensuring uniform sampling from SO\(3\)\. This augmentation prevents overfitting to specific orientations\.

#### 3\.6\.2Fragment augmentation

In order to obtain a model capable of denoising a molecule with a soft constraint of a molecular fragment, we perform the following data augmentation when generating the ligand data on\-the\-fly during training\. First, we use RDKit’s Murcko scaffold implementation to select the canonical scaffold for the target molecule\. Then, we sample uniformly the desired fragment size\. If that size is smaller than the scaffold, then we remove atoms one at a time at random, restricting to those that would not create disconnected fragments\. If the desired size is larger than the scaffold, we add atoms from the molecule one at a time at random, restricting to those that directly connect to atoms already within the fragment\. In this way, we obtain an augmented set of fragments that are uniformly size\-distributed, and heavily biased towards scaffold\-like fragments\. This fragment is then used to create the ligand density map in 50% of ligand\-conditioned cases\. Note that the direct inputs and outputs of the model are unchanged; only the conditioning itself is affected\.

Table 3:Dataset statistics

## 4Experiments

### 4\.1Training Setup

Training was performed on an A100 GPU\. We used PyTorch 2\.7\.1 with CUDA 12\.6 support and PyTorch Lightning 2\.5\.0\. Pretraining ran for approximately a week over 100k optimizer steps, which corresponds to 12M training datapoints\. Finetuning ran for approximately 24 hours, over 1M datapoints\.

### 4\.2Training Validation

We validate the model across multiple conditioning modes to assess its ability to use structural information:

- •Fully conditioned \(P\+L\): Both protein and ligand density maps provided
- •Protein and fragment conditioned \(P\+F\): Both protein density maps and a ligand density map from a fragment provided
- •Protein\-only \(P\): Only protein density map provided
- •Full ligand conditioning \(L\): Only ligand density map provided
- •Fragment conditioning only \(F\): Only a ligand density map from a fragment is provided
- •Unconditioned \(U\): No density maps provided

Note that LDDT here is a sigmoid\-normalized fraction of atom\-pairs that correctly lie within \[0\.5, 1, 2, 4\]Åof the ground truth; a surrogate to the loss used in training\. Single\-fragment is the fraction of molecules that are fully\-connected, which proves to be a surprising challenge for variable\-atom count models\. Atom validity is the proportion of atoms that have correct valence structure\. Note that Sesame splits out different atom types according to their hydrogen count, so this metric is stricter than the same metric for many other similar models\. Molecule validity is the fraction of generated molecules that satisfy both of the above, for all the constituent atoms\. Table[4](https://arxiv.org/html/2606.23856#S4.T4)shows the results for pretraining, and table[5](https://arxiv.org/html/2606.23856#S4.T5)shows the results for finetuning\. Notably, trajectory\-level metrics improve with finetuning, while single\-step metrics stay roughly equivalent or degrade slightly\.

Table 4:Validation metrics by conditioning mode \(pretrained\)Table 5:Validation metrics by conditioning mode \(finetuned\)The main type of error the model makes is creating spurious atoms\. This is most noticeable in the fragment\-conditioning regime, but occurs across all conditioning environments\. These excess atoms are typically singletons that do not have any bonds, and the issue can likely be solved with more finetuning or cleverer noise schedules\. However, we note that with simple post\-processing to remove these atoms, the validity metrics for the two most salient regimes for drug discovery – protein \+ fragment for lead optimization and protein\-only forde novogeneration – jump to 92\.4% and 88\.7% respectively on the finetuned model\.

Beyond producing valid molecules, fragment conditioning must actually steer generation toward the supplied seed\. Across all fragment\-conditioned generation \(both protein \+ fragment and fragment\-only\), 94\.8% of generated molecules contain the seeding fragment as a substructure, confirming that the soft density\-map conditioning reliably preserves the intended scaffold rather than discarding it\.

Validity alone, however, does not guarantee that generated molecules resemble viable drug candidates\. To assess this, we compare the property distributions of Sesame\-generated molecules against ligands drawn from SandboxAQ’s SAIR dataset\(Lemoset al\.,[2025](https://arxiv.org/html/2606.23856#bib.bib24)\), a large synthetic structural dataset that has been widely used to train recent structure\-based models\. As shown in Figure[3](https://arxiv.org/html/2606.23856#S4.F3), Sesame produces molecules whose property distributions closely track those of real drug\-like ligands and largely respect the Lipinski Rule of 5 thresholds\. Notably, seeding generation with an existing scaffold yields slightly more drug\-like distributions thande novogeneration\.

![Refer to caption](https://arxiv.org/html/2606.23856v1/x2.png)Figure 3:Sesame produces drug\-like molecules\.We compare Sesame’s molecule generation to ligands drawn from SandboxAQ’s SAIR dataset, which has been widely used to train recent models\. Sesame produces molecules with a similar distribution of properties as real drugs; notably, starting with an existing scaffold \(blue distributions\) performs slightly better thande novogeneration \(green distributions\)\. Red dashed lines indicate Lipinski Rule of 5 thresholds, where applicable\.
### 4\.3Sampling

When sampling with a Sesame model, there are several distinct choices that can be made\. The most salient of these is to choose the full 3\-dimensional noise schedule, as the model is trained to denoise in a schedule\-agnostic manner, but not all paths through noise space are equal\. While a full optimization of this path is certainly an option, we chose to adopt a simple strategy: Given that each noise coordinate is normalized to\[0,1\],\[0,1\],choose a functionσ​\(t\)=\(tγposition,tγatom,tγbond\)\.\\sigma\(t\)=\(t^\{\\gamma\_\{\\text\{position\}\}\},t^\{\\gamma\_\{\\text\{atom\}\}\},t^\{\\gamma\_\{\\text\{bond\}\}\}\)\.We can then sweep through a small set of choices for each of theγ\\gammaby running a 20\-step denoising trajectory with a linearttfrom 1 to 0\. For our sweep, we chose 4 values ofγ\\gamma\(0\.5, 1, 2, 4\), to look at 64 schedules with 2048 randomly sampled initial states, protein conditioning, and no ligand conditioning\. We scored the generated molecules with a custom score that penalizes the above metrics in a smooth manner\. Results of this sweep are in Figure[4](https://arxiv.org/html/2606.23856#S4.F4)\. As expected, the idealγposition\\gamma\_\{\\text\{position\}\}was larger thanγatom\\gamma\_\{\\text\{atom\}\}\. This essentially states that fixing positions first is preferable to fixing atom types, and then trying to match them to the correct point in space\. Surprisingly, a largeγ\\gammafor bonds was strongly favored, one that implies that bonds are fixed before positions become well\-determined\. Our leading hypothesis is that this is due to two main factors: first, the pairwise matrix of bonds has a magnitude of order more micro\-decisions to make, but most of these are relatively trivial\. As a result, fixing most bonds to be the non\-bonding class is easier to set early and correct at low noise levels, rather than one\-shotting the prediction near the end of the diffusion process\. The second factor that we believe to be important is that the bond matrix at low noise provides a strong bias for the general structure of the molecule, and make both position and type denoising easier\.

![Refer to caption](https://arxiv.org/html/2606.23856v1/x3.png)Figure 4:Results from the schedule sweep\. Top: gradient\-smoothed average score from the 64 schedules\. Bottom: 3d path of the schedule in normalized noise space\. Paths are colored according to their average final score\.

## 5Discussion

##### Human\-in\-the\-loop lead optimization\.

A central design goal of Sesame is to make lead optimization a first\-class mode of use rather than a special case ofde novogeneration\. The most common real\-world drug\-discovery task is not designing a molecule from nothing, but improving a known hit\. Because a partial ligand enters the model as the same continuous density field used for the protein pocket \(Section[3\.6\.2](https://arxiv.org/html/2606.23856#S3.SS6.SSS2)\), a chemist can prune a hit to the substructure worth keeping and hand Sesame that fragment as a soft prior; the model then grows the remaining chemistry to fit both the fragment and the pocket\. Unlike methods that fix a scaffold as discrete atoms, the field\-based prior can be partial and approximate, and the surrounding atoms, bonds, and geometry are all free to adapt around it\.

This dual capability is reflected in our validation\. The protein\-plus\-fragment \(P\+F\) regime is among the strongest across the trajectory\-level metrics \(Table[5](https://arxiv.org/html/2606.23856#S4.T5)\), and after simple post\-processing the two regimes we consider most salient for drug discovery – protein\-plus\-fragment for lead optimization and protein\-only forde novogeneration – reach 92\.4% and 88\.7% molecule validity respectively\. We view this combination, exposed through a single conditioning interface, as the practical core of the system: human structural insight and generative chemistry are composed rather than traded off\.

## 6Conclusion

We presented Sesame, a structure\-aware molecular generation system that uses density map conditioning and hybrid diffusion\. Our key innovation is an attention\-based density map conditioning architecture that adaptively extracts structural information to guide generation\. The hybrid diffusion process handles both discrete atom types and continuous coordinates within a unified framework\.

Experiments demonstrate that the model effectively uses conditioning information, with performance improving when both protein and ligand density maps are provided\. Multi\-source training enables learning from both general molecular chemistry and specific binding interactions\.

Critically, because the pocket and a partial ligand are presented to the model through the same density\-map interface, Sesame supports bothde novodesign and fragment\-conditioned lead optimization within one framework\. Our validation across protein\-only and protein\-plus\-fragment regimes shows the model makes productive use of a supplied scaffold, pointing toward a workflow in which a medicinal chemist’s structural insight directly seeds generation and the model elaborates it into pocket\-compatible molecules\.

## References

- J\. Abramson, J\. Adler, J\. Dunger, R\. Evans, T\. Green, A\. Pritzel, O\. Ronneberger, L\. Willmore, A\. J\. Ballard, J\. Bambrick, S\. W\. Bodenstein, D\. A\. Evans, C\. Hung, M\. O’Neill, D\. Reiman, K\. Tunyasuvunakool, Z\. Wu, A\. Žemgulytė, E\. Arvaniti, C\. Beattie, O\. Bertolli, A\. Bridgland, A\. Cherepanov, M\. Congreve, A\. I\. Cowen\-Rivers, A\. Cowie, M\. Figurnov, F\. B\. Fuchs, H\. Gladman, R\. Jain, Y\. A\. Khan, C\. M\. R\. Low, K\. Perlin, A\. Potapenko, P\. Savy, S\. Singh, A\. Stecula, A\. Thillaisundaram, C\. Tong, S\. Yakneen, E\. D\. Zhong, M\. Zielinski, A\. Žídek, V\. Bapst, P\. Kohli, M\. Jaderberg, D\. Hassabis, and J\. M\. Jumper \(2024\)Accurate structure prediction of biomolecular interactions with AlphaFold 3\.Nature630\(8016\),pp\. 493–500\.External Links:[Document](https://dx.doi.org/10.1038/s41586-024-07487-w),[Link](https://doi.org/10.1038/s41586-024-07487-w)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- K\. Adams, K\. Abeywardane, J\. Fromer, and C\. W\. Coley \(2025\)ShEPhERD: diffusing shape, electrostatics, and pharmacophores for bioisosteric drug design\.InThe Thirteenth International Conference on Learning Representations \(ICLR 2025\),External Links:[Link](https://openreview.net/forum?id=KSLkFYHlYg)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px2.p1.1),[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- J\. Cremer, R\. Irwin, A\. Tibo, J\. P\. Janet, S\. Olsson, and D\. Clevert \(2026\)FLOWR: flow matching for structure\-aware de novo, interaction\- and fragment\-based ligand generation\.Nature Computational Science\.External Links:[Document](https://dx.doi.org/10.1038/s43588-026-00998-8),[Link](https://doi.org/10.1038/s43588-026-00998-8)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- J\. Cremer, T\. Le, F\. Noé, D\. Clevert, and K\. T\. Schütt \(2024\)PILOT: equivariant diffusion for pocket\-conditioned de novo ligand generation with multi\-objective guidance via importance sampling\.Chemical Science15,pp\. 14954–14967\.External Links:[Document](https://dx.doi.org/10.1039/D4SC03523B),[Link](https://doi.org/10.1039/D4SC03523B)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- M\. Cretu, C\. Harris, I\. Igashov, A\. Schneuing, M\. Segler, B\. Correia, J\. Roy, E\. Bengio, and P\. Liò \(2025\)SynFlowNet: towards molecule design with guaranteed synthesis pathways\.InThe Thirteenth International Conference on Learning Representations \(ICLR 2025\),External Links:[Link](https://openreview.net/forum?id=uvHmnahyp1)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- I\. Dunn and D\. R\. Koes \(2026\)FlowMol3: flow matching for 3d de novo small\-molecule generation\.Digital Discovery5\(5\),pp\. 2052–2066\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- W\. Gao, S\. Luo, and C\. W\. Coley \(2025\)Generative AI for navigating synthesizable chemical space\.Proceedings of the National Academy of Sciences122\(41\),pp\. e2415665122\.External Links:[Document](https://dx.doi.org/10.1073/pnas.2415665122),[Link](https://doi.org/10.1073/pnas.2415665122)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- J\. Guan, W\. W\. Qian, X\. Peng, Y\. Su, J\. Peng, and J\. Ma \(2023\)3D equivariant diffusion for target\-aware molecule generation and affinity prediction\.InThe Eleventh International Conference on Learning Representations \(ICLR 2023\),External Links:[Link](https://openreview.net/forum?id=kJqXEPXMsE0)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- E\. Hoogeboom, V\. G\. Satorras, C\. Vignac, and M\. Welling \(2022\)Equivariant diffusion for molecule generation in 3d\.InProceedings of the 39th International Conference on Machine Learning,pp\. 8867–8887\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- I\. Igashov, H\. Stärk, C\. Vignac, A\. Schneuing, V\. G\. Satorras, P\. Frossard, M\. Welling, M\. Bronstein, and B\. Correia \(2024\)Equivariant 3d\-conditional diffusion model for molecular linker design\.Nature Machine Intelligence6\(4\),pp\. 417–427\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px3.p1.1)\.
- R\. Irwin, A\. Tibo, J\. P\. Janet, and S\. Olsson \(2025\)SemlaFlow – efficient 3d molecular generation with latent attention and equivariant flow matching\.arXiv preprint arXiv:2406\.07266\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- R\. P\. Joshi, N\. W\. A\. Gebauer, M\. Bontha, M\. Khazaieli, R\. M\. James, J\. B\. Brown, and N\. Kumar \(2021\)3D\-Scaffold: a deep learning framework to generate 3d coordinates of drug\-like molecules with desired scaffolds\.The Journal of Physical Chemistry B125\(44\),pp\. 12166–12176\.External Links:[Document](https://dx.doi.org/10.1021/acs.jpcb.1c06437),[Link](https://doi.org/10.1021/acs.jpcb.1c06437)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px3.p1.1)\.
- P\. Lemos, Z\. Beckwith, S\. Bandi, M\. van Damme, J\. Crivelli\-Decker, B\. J\. Shields, T\. Merth, P\. K\. Jha, N\. De Mitri, T\. J\. Callahan, A\. J\. Nish, P\. Abruzzo, R\. Salomon\-Ferrer, and M\. Ganahl \(2025\)SAIR: enabling deep learning for protein–ligand interactions with a synthetic structural dataset\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2025.06.17.660168),[Link](https://doi.org/10.1101/2025.06.17.660168)Cited by:[§3\.5\.2](https://arxiv.org/html/2606.23856#S3.SS5.SSS2.p1.1),[§4\.2](https://arxiv.org/html/2606.23856#S4.SS2.p5.1)\.
- S\. Passaro, G\. Corso, J\. Wohlwend, M\. Reveiz, S\. Thaler, V\. R\. Somnath, N\. Getz, T\. Portnoi, J\. Roy, H\. Stark, D\. Kwabi\-Addo, D\. Beaini, T\. Jaakkola, and R\. Barzilay \(2025\)Boltz\-2: towards accurate and efficient binding affinity prediction\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2025.06.14.659707),[Link](https://doi.org/10.1101/2025.06.14.659707)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- P\. O\. Pinheiro, A\. Jamasb, O\. Mahmood, V\. Sresht, and S\. Saremi \(2024a\)Structure\-based drug design by denoising voxel grids\.arXiv preprint arXiv:2405\.03961\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px2.p1.1)\.
- P\. O\. Pinheiro, J\. Rackers, J\. Kleinhenz, M\. Maser, O\. Mahmood, A\. M\. Watkins, S\. Ra, V\. Sresht, and S\. Saremi \(2024b\)3D molecule generation by denoising voxel grids\.arXiv preprint arXiv:2306\.07473\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px2.p1.1)\.
- A\. Schneuing, C\. Harris, Y\. Du, K\. Didi, A\. Jamasb, I\. Igashov, W\. Du, C\. Gomes, T\. L\. Blundell, P\. Liò, M\. Welling, M\. Bronstein, and B\. Correia \(2024\)Structure\-based drug design with equivariant diffusion models\.Nature Computational Science4\(12\),pp\. 899–909\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px4.p1.1)\.
- N\. Shazeer \(2020\)GLU variants improve transformer\.arXiv preprint arXiv:2002\.05202\.External Links:[Link](https://arxiv.org/abs/2002.05202)Cited by:[§3\.2\.2](https://arxiv.org/html/2606.23856#S3.SS2.SSS2.p10.1)\.
- J\. Song, C\. Meng, and S\. Ermon \(2021\)Denoising diffusion implicit models\.InInternational Conference on Learning Representations \(ICLR\),External Links:[Link](https://arxiv.org/abs/2010.02502)Cited by:[§3\.3\.2](https://arxiv.org/html/2606.23856#S3.SS3.SSS2.Px1.p1.1)\.
- B\. I\. Tingle, K\. G\. Tang, M\. Castanon, J\. J\. Gutierrez, M\. Khurelbaatar, C\. Dandarchuluun, Y\. S\. Moroz, and J\. J\. Irwin \(2023\)ZINC\-22—a free multi\-billion\-scale database of tangible compounds for ligand discovery\.Journal of Chemical Information and Modeling63\(4\),pp\. 1166–1176\.Note:PMID: 36790087External Links:[Document](https://dx.doi.org/10.1021/acs.jcim.2c01253),[Link](https://doi.org/10.1021/acs.jcim.2c01253)Cited by:[§3\.5\.1](https://arxiv.org/html/2606.23856#S3.SS5.SSS1.p1.1)\.
- C\. Vignac, I\. Krawczuk, A\. Siraudin, B\. Wang, V\. Cevher, and P\. Frossard \(2023a\)DiGress: discrete denoising diffusion for graph generation\.InThe Eleventh International Conference on Learning Representations \(ICLR 2023\),External Links:[Link](https://openreview.net/forum?id=UaAD-Nu86WX)Cited by:[§3\.3\.1](https://arxiv.org/html/2606.23856#S3.SS3.SSS1.Px2.p1.1)\.
- C\. Vignac, N\. Osman, L\. Toni, and P\. Frossard \(2023b\)MiDi: mixed graph and 3d denoising diffusion for molecule generation\.InJoint European Conference on Machine Learning and Knowledge Discovery in Databases,pp\. 560–576\.Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Wohlwend, G\. Corso, S\. Passaro, N\. Getz, M\. Reveiz, K\. Leidal, W\. Swiderski, L\. Atkinson, T\. Portnoi, I\. Chinn, J\. Silterra, T\. Jaakkola, and R\. Barzilay \(2024\)Boltz\-1: democratizing biomolecular interaction modeling\.bioRxiv\.External Links:[Document](https://dx.doi.org/10.1101/2024.11.19.624167),[Link](https://doi.org/10.1101/2024.11.19.624167)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Xie, S\. Chen, J\. Lei, and Y\. Yang \(2024\)DiffDec: structure\-aware scaffold decoration with an end\-to\-end diffusion model\.Journal of Chemical Information and Modeling64\(7\),pp\. 2554–2564\.External Links:[Document](https://dx.doi.org/10.1021/acs.jcim.3c01466),[Link](https://doi.org/10.1021/acs.jcim.3c01466)Cited by:[§2](https://arxiv.org/html/2606.23856#S2.SS0.SSS0.Px3.p1.1)\.

## Appendix ADensity Map conditioning Operations

Given single representation𝐬∈ℝB×N×ds\\mathbf\{s\}\\in\\mathbb\{R\}^\{B\\times N\\times d\_\{s\}\}whereBBis batch size,NNis number of atoms, andds=384d\_\{s\}=384is the single dimension, we compute:

𝐐sample\\displaystyle\\mathbf\{Q\}\_\{\\text\{sample\}\}=Linear​\(𝐬\)∈ℝB×N×H×dh\\displaystyle=\\text\{Linear\}\(\\mathbf\{s\}\)\\in\\mathbb\{R\}^\{B\\times N\\times H\\times d\_\{h\}\}\(29\)𝐊sample\\displaystyle\\mathbf\{K\}\_\{\\text\{sample\}\}=𝐊learned∈ℝH×O×dh\\displaystyle=\\mathbf\{K\}\_\{\\text\{learned\}\}\\in\\mathbb\{R\}^\{H\\times O\\times d\_\{h\}\}\(30\)𝐕sample\\displaystyle\\mathbf\{V\}\_\{\\text\{sample\}\}=Linear​\(𝐬\)∈ℝB×N×H×3\\displaystyle=\\text\{Linear\}\(\\mathbf\{s\}\)\\in\\mathbb\{R\}^\{B\\times N\\times H\\times 3\}\(31\)whereH=4H=4is the number of attention heads,O=1024/H=256O=1024/H=256is the number of sampling points per head \(for a total of10241024sampling points\), anddh=ds/H=96d\_\{h\}=d\_\{s\}/H=96is the dimension per head\. The attention mechanism computes:

𝐀=softmax⁡\(𝐊sample​𝐐sampleTdh\)∈ℝB×H×O×N\\mathbf\{A\}=\\operatorname\{softmax\}\\left\(\\frac\{\\mathbf\{K\}\_\{\\text\{sample\}\}\\mathbf\{Q\}\_\{\\text\{sample\}\}^\{T\}\}\{\\sqrt\{d\_\{h\}\}\}\\right\)\\in\\mathbb\{R\}^\{B\\times H\\times O\\times N\}\(32\)and the sampling points are computed as:

𝐏sample=tanh⁡\(reshape​\(𝐀𝐕sample\)\)∈ℝB×O⋅H×3\\mathbf\{P\}\_\{\\text\{sample\}\}=\\tanh\\left\(\\text\{reshape\}\(\\mathbf\{A\}\\mathbf\{V\}\_\{\\text\{sample\}\}\)\\right\)\\in\\mathbb\{R\}^\{B\\times O\\cdot H\\times 3\}\(33\)where the reshape combines the head and sampling point dimensions\. The tanh ensures the range in\[−1,1\]\[\-1,1\]for grid sampling\.

Grid Sampling: The sampling points are used to extract features from the density map via trilinear interpolation:

𝐅density=GridSample​\(𝐃protein,𝐏sample\)\\mathbf\{F\}\_\{\\text\{density\}\}=\\text\{GridSample\}\(\\mathbf\{D\}\_\{\\text\{protein\}\},\\mathbf\{P\}\_\{\\text\{sample\}\}\)\(34\)where𝐃protein∈ℝdd×12×12×12\\mathbf\{D\}\_\{\\text\{protein\}\}\\in\\mathbb\{R\}^\{d\_\{d\}\\times 12\\times 12\\times 12\}is the post\-3D convnet density map\.

Cross\-Attention: Theddd\_\{d\}density features are concatenated with sampling point coordinates and used in cross\-attention to update single and pair representations:

𝐅combined=\[𝐅density,𝐏sample\]∈ℝB×1024×\(dd\+3\)\\mathbf\{F\}\_\{\\text\{combined\}\}=\[\\mathbf\{F\}\_\{\\text\{density\}\},\\mathbf\{P\}\_\{\\text\{sample\}\}\]\\in\\mathbb\{R\}^\{B\\times 1024\\times\(d\_\{d\}\+3\)\}\(35\)wheredd=384d\_\{d\}=384is the processed density map channel dimension\. The cross\-attention computes:

𝐐single\\displaystyle\\mathbf\{Q\}\_\{\\text\{single\}\}=Linear​\(𝐬\)∈ℝB×N×H×dh\\displaystyle=\\text\{Linear\}\(\\mathbf\{s\}\)\\in\\mathbb\{R\}^\{B\\times N\\times H\\times d\_\{h\}\}\(36\)𝐐pair\\displaystyle\\mathbf\{Q\}\_\{\\text\{pair\}\}=Linear​\(𝐩\)∈ℝB×N×N×H×dh\\displaystyle=\\text\{Linear\}\(\\mathbf\{p\}\)\\in\\mathbb\{R\}^\{B\\times N\\times N\\times H\\times d\_\{h\}\}\(37\)𝐊density\\displaystyle\\mathbf\{K\}\_\{\\text\{density\}\}=Linear​\(𝐅combined\)∈ℝB×1024×H×dh\\displaystyle=\\text\{Linear\}\(\\mathbf\{F\}\_\{\\text\{combined\}\}\)\\in\\mathbb\{R\}^\{B\\times 1024\\times H\\times d\_\{h\}\}\(38\)𝐕single\\displaystyle\\mathbf\{V\}\_\{\\text\{single\}\}=Linear​\(𝐅combined\)∈ℝB×1024×ds\\displaystyle=\\text\{Linear\}\(\\mathbf\{F\}\_\{\\text\{combined\}\}\)\\in\\mathbb\{R\}^\{B\\times 1024\\times d\_\{s\}\}\(39\)𝐕pair\\displaystyle\\mathbf\{V\}\_\{\\text\{pair\}\}=Linear​\(𝐅combined\)∈ℝB×1024×dp\\displaystyle=\\text\{Linear\}\(\\mathbf\{F\}\_\{\\text\{combined\}\}\)\\in\\mathbb\{R\}^\{B\\times 1024\\times d\_\{p\}\}\(40\)where𝐩∈ℝB×N×N×dp\\mathbf\{p\}\\in\\mathbb\{R\}^\{B\\times N\\times N\\times d\_\{p\}\}is the pair representation withdp=128d\_\{p\}=128\. The attention outputs are:

𝐬conditioned\\displaystyle\\mathbf\{s\}\_\{\\text\{conditioned\}\}=reshape​\(softmax⁡\(𝐐single​𝐊densityT/dh\)​𝐕single\)∈ℝB×N×ds\\displaystyle=\\text\{reshape\}\(\\operatorname\{softmax\}\(\\mathbf\{Q\}\_\{\\text\{single\}\}\\mathbf\{K\}\_\{\\text\{density\}\}^\{T\}/\\sqrt\{d\_\{h\}\}\)\\mathbf\{V\}\_\{\\text\{single\}\}\)\\in\\mathbb\{R\}^\{B\\times N\\times d\_\{s\}\}\(41\)𝐩conditioned\\displaystyle\\mathbf\{p\}\_\{\\text\{conditioned\}\}=reshape​\(softmax⁡\(𝐐pair​𝐊densityT/dh\)​𝐕pair\)∈ℝB×N×N×dp\\displaystyle=\\text\{reshape\}\(\\operatorname\{softmax\}\(\\mathbf\{Q\}\_\{\\text\{pair\}\}\\mathbf\{K\}\_\{\\text\{density\}\}^\{T\}/\\sqrt\{d\_\{h\}\}\)\\mathbf\{V\}\_\{\\text\{pair\}\}\)\\in\\mathbb\{R\}^\{B\\times N\\times N\\times d\_\{p\}\}\(42\)The conditioning is applied to the single and pair representations, respectively, through gated residual connections\.