Equivariant Shape-Conditioned Generation of 3D Molecules for Ligand-Based Drug Design

Shape-based virtual screening is widely employed in ligand-based drug design to search chemical libraries for molecules with similar 3D shapes yet novel 2D chemical structures compared to known ligands. 3D deep generative models have the potential to automate this exploration of shape-conditioned 3D chemical space; however, no existing models can reliably generate valid drug-like molecules in conformations that adopt a specific shape such as a known binding pose. We introduce a new multimodal 3D generative model that enables shape-conditioned 3D molecular design by equivariantly encoding molecular shape and variationally encoding chemical identity. We ensure local geometric and chemical validity of generated molecules by using autoregressive fragment-based generation with heuristic bonding geometries, allowing the model to prioritize the scoring of rotatable bonds to best align the growing conformational structure to the target shape. We evaluate our 3D generative model in tasks relevant to drug design including shape-conditioned generation of chemically diverse molecular structures and shape-constrained molecular property optimization, demonstrating its utility over virtual screening of enumerated libraries.

Ligand-based drug design (LBDD) does not assume knowledge of protein structure. Instead, molecules are compared against previously identified "actives" on the basis of 3D pharmacophore or 3D shape similarity under the principle that molecules with similar structures should share similar activity (Vázquez et al., 2020;Cleves & Jain, 2020). In particular, ROCS (Rapid Overlay of Chemical Structures) is commonly used as a shape-based virtual screening tool to identify molecules with similar shapes to a reference inhibitor and has shown promising results for scaffold-hopping tasks (Rush et al., 2005;Hawkins et al., 2007;Nicholls et al., 2010). However, virtual screening relies on enumeration of chemical libraries, fundamentally restricting its ability to probe new chemical space.
Here, we consider the novel task of generating chemically diverse 3D molecular structures conditioned on a molecular shape, thereby facilitating the shape-conditioned exploration of chemical To surmount these challenges, we design a new generative model, SQUID 1 , to enable the shapeconditioned generation of chemically diverse molecules in 3D. Our contributions are as follows: • Given a 3D molecule with a target shape, we use equivariant point cloud networks to encode the shape into (rotationally) equivariant features. We then use graph neural networks (GNNs) to variationally encode chemical identity into invariant features. By mixing chemical features with equivariant shape features, we can generate diverse molecules in aligned poses that fit the shape. • We develop a sequential fragment-based 3D generation procedure that fixes local bond lengths and angles to prioritize the scoring of rotatable bonds. By massively simplifying 3D coordinate generation, we generate drug-like molecules while maintaining chemical and geometric validity. • We design a rotatable bond scoring network that learns how local bond rotations affect global shape, enabling our decoder to generate 3D conformations that best fit the target shape.
We evaluate the utility of SQUID over virtual screening in shape-conditioned 3D molecular design tasks that mimic ligand-based drug design objectives, including shape-conditioned generation of diverse 3D structures and shape-constrained molecular optimization. To inspire further research, we note that our tasks could also be approached with a hypothetical 3D generative model that disentangles latent variables controlling 2D chemical identity and 3D shape, thus enabling zero-shot generation of topologically distinct molecules with similar shapes to any encoded target.

RELATED WORK
Fragment-based molecular generation. Seminal works in autoregressive molecular generation applied language models to generate 1D SMILES string representations of molecules character-bycharacter (Gómez-Bombarelli et al., 2018;Segler et al., 2018), or GNNs to generate 2D molecular graphs atom-by-atom (Liu et al., 2018;Simonovsky & Komodakis, 2018;Li et al., 2018). Recent works construct molecules fragment-by-fragment to improve the chemical validity of intermediate graphs and to scale generation to larger organic molecules (Podda et al., 2020;Jin et al., 2019;. Our fragment-based decoder is most related to MoLeR (Maziarz et al., 2022), which iteratively generates molecules by selecting a new fragment (or atom) to add to the partial graph, choosing attachment sites in the new fragment, and predicting new bonds to the partial graph. Yet, MoLeR only generates 2D graphs; we generate 3D molecular structures. Beyond 2D generation, Flam-Shepherd et al. (2022) train an RL agent to generate 3D molecules by sampling and connecting molecular fragments. However, their agent samples from a small multiset of fragment types, severely restricting the accessible chemical space. Powers et al. (2022) also use fragments to generate 3D molecules inside a protein pocket, but employ a limited library containing just 7 distinct rings.
Generation of drug-like molecules in 3D. In this work, we generate novel drug-like 3D molecular structures in free space, e.g., not conformers given a known molecular graph (Ganea et al., 2021;Jing et al., 2022). Myriad models have been proposed to generate small 3D molecules such as E(3)-equivariant normalizing flows and diffusion models (Satorras et al., 2022a;Hoogeboom et al., 2022), RL agents with an SE(3)-covariant action space (Simm et al., 2020b), and autoregressive generators that build molecules atom-by-atom with SE(3)-invariant internal coordinates (Luo & Ji, 2022;Gebauer et al., 2022). However, fewer 3D generative models can generate larger drug-like molecules for realistic chemical design tasks. Of these, Hoogeboom et al. (2022) and Arcidiacono & Koes (2021) fail to generate chemically valid molecules, while Ragoza et al. (2020) rely on postprocessing and geometry relaxation to extract stable molecules from their generated atom density grids. Only Roney et al. (2021) and Li et al. (2021), who develop autoregressive generators that simultaneously predict graph structure and internal coordinates, have shown to reliably generate valid drug-like molecules. We also couple graph generation with 3D coordinate prediction; however, we employ fragment-based generation with fixed local geometries to ensure local chemical and geometric validity. Futher, we focus on shape-conditioned molecular design; none of these works can natively address the aforementioned challenges posed by shape-conditioned molecular generation.
Shape-conditioned molecular generation. Other works have only partially addressed shapeconditioned 3D molecular generation. Skalic et al. (2019) and Imrie et al. (2021) REINVENT (Olivecrona et al., 2017) on SMILES strings to propose molecules whose conformers have high shape similarity to a target, but they must re-optimize the agent for each target shape. Roney et al. (2021) fine-tune a 3D generative model on the hits of a ROCS virtual screen of > 10 10 drug-like molecules to shift the learned distribution towards a particular shape. Yet, this expensive screening approach must be repeated for each new shape. Instead, we seek to achieve zero-shot generation of 3D molecules with similar shapes to any encoded target shape, without requiring fine-tuning or costly post facto optimization strategies.
Equivariant geometric deep learning on point clouds. Various equivariant networks have been designed to encode point clouds for updating coordinates in R 3 (Satorras et al., 2022b), predicting tensorial properties (Thomas et al., 2018), or modeling 3D structures natively in Cartesian space (Fuchs et al., 2020). Especially noteworthy are architectures which lift scalar neuron features to vector features in R 3 and employ simple operations to mix invariant and equivariant features without relying on expensive higher-order tensor products or Clebsh-Gordan coefficients Jing et al., 2021). In this work, we employ 's Vector Neurons (VN)-based equivariant point cloud encoder VN-DGCNN to encode molecules into equivariant latent representations in order to generate molecules which are natively aligned to the target shape. Two recent works also employ VN operations for structure-based drug design and linker design (Peng et al., 2022;Huang et al., 2022). Huang et al. (2022) also build molecules in free space; however, they generate just a few atoms to connect existing fragments and do not condition on molecular shape.

METHODOLOGY
Problem definition. We model a conditional distribution P (M |S) over 3D molecules M = (G, G) with graph G and atomic coordinates G = {r a ∈ R 3 } given a 3D molecular shape S. Specifically, we aim to sample molecules M ∼ P (M |S) with high shape similarity (sim S (M , M S ) ≈ 1) and low graph (chemical) similarity (sim G (M , M S ) < 1) to a target molecule M S with shape S. This scheme differs from 1) typical 3D generative models that learn P (M ) without modeling P (M |S), and from 2) shape-conditioned 1D/2D generators that attempt to model P (G|S), the distribution of molecular graphs that could adopt shape S, but do not actually generate specific 3D conformations.
We define graph (chemical) similarity sim G ∈ [0, 1] between two molecules as the Tanimoto similarity computed by RDKit with default settings (2048-bit fingerprints). We define shape similarity sim * S ∈ [0, 1] using Gaussian descriptions of molecular shape, modeling atoms a ∈ M A and b ∈ M B from molecules M A and M B as isotropic Gaussians in R 3 (Grant & Pickup, 1995;Grant et al., 1996). We compute sim * S using (2-body) volume overlaps between atom-centered Gaussians: where α controls the Gaussian width. Setting α = 0.81 approximates the shape similarity function used by the ROCS program (App. A.6). sim * S is sensitive to SE as the shape similarity when M A is optimally aligned to M B . We perform such alignments with ROCS. Approach. At a high level, we model P (M |S) with an encoder-decoder architecture. Given a molecule M S = (G S , G S ) with shape S, we encode S (a point cloud) into equivariant features. We then variationally encode G S into atomic features, conditioned on the shape features. We then mix these shape and atom features to pass global SE(3) {in,equi}variant latent codes to the decoder, which samples new molecules from P (M |S). We autoregressively generate molecules by factoring P (M |S) = P (M 0 |S)P (M 1 |M 0 , S)...P (M |M n−1 , S), where each M l = (G l , G l ) are partial molecules defined by a BFS traversal of a tree-representation of the molecular graph (Fig. 2). Tree-nodes denote either non-ring atoms or rigid (ring-containing) fragments, and tree-links denote acyclic (rotatable, double, or triple) bonds. We generate M l+1 by growing the graph G l+1 around a focus atom/fragment, and then predict G l+1 by scoring a query rotatable bond to best fit shape S.
Simplifying assumptions. (1) We ignore hydrogens and only consider heavy atoms, as is common in molecular generation. (2) We only consider molecules with fragments present in our fragment library to ensure that graph generation can be expressed as tree generation. (3) Rather than generating all coordinates, we use rigid fragments, fix bond distances, and set bond angles according to hybridization heuristics (App. A.8); this lets the model focus on scoring rotatable bonds to best fit the growing conformer to the encoded shape. (4) We seed generation with M 0 (the root tree-node), restricted to be a small (3-6 atoms) substructure from M S ; hence, we only model P (M |S, M 0 ).

ENCODER
Featurization. We construct a molecular graph G using atoms as nodes and bonds as edges. We featurize each node with the atomic mass; one-hot codes of atomic number, charge, and aromaticity; and one-hot codes of the number of single, double, aromatic, and triple bonds the atom forms (including bonds to implicit hydrogens). This helps us fix bond angles during generation (App. A.8). We featurize each edge with one-hot codes of bond order. We represent a shape S as a point cloud built by sampling n p points from each of n h atom-centered Gaussians with (adjustable) variance σ 2 p . Fragment encoder. We also featurize each node with a learned embedding f i ∈ R d f of the atom/fragment type to which that atom belongs, making each node "fragment-aware" (similar to MoLeR). In principle, fragments could be any rigid molecular substructure containing ≥ 2 atoms. Figure 2: Encoder-decod er architecture of SQUID, which equivariantly encodes molecular shape and variationally encodes chemical identity to generate chemically diverse 3D molecules that fit the shape. SQUID generates molecules atom-by-atom and fragment-by-fragment, iteratively growing the molecular graph in a tree-expansion and generating 3D coordinates by scoring rotatable bonds.
Here, we specify fragments as ring-containing substructures without acyclic single bonds (benzene, naphthalene, pyridazinone, etc.). We construct a library L f of atom/fragment types by extracting the top-k (k = 100) most frequent fragments from the dataset and adding these, along with each distinct atom type, to L f (App. A.13). We then encode each atom/fragment in L f with a simple GNN (App. A.12) to yield the global atom/fragment embeddings: fi are per-atom features. Shape encoder. Given M S with n h heavy atoms, we use VN-DGCNN (App. A.11) to encode the molecular point cloud P S ∈ R (n h np)×3 into a set of equivariant per-point vector featuresX p ∈ R (n h np)×q×3 . We then locally mean-pool the n p equivariant features per atom: whereX ∈ R n h ×q×3 are per-atom equivariant representations of the molecular shape. Because VN operations are SO(3)-equivariant, rotating the point cloud will rotateX:XR = LocalPool(VN-DGCNN(P S R)). Although VN operations are strictly SO(3)-equivariant, we subtract the molecule's centroid from the atomic coordinates prior to encoding, makingX effectively SE(3)-equivariant. Throughout this work, we denote SO(3)-equivariant vector features with tildes.
Variational graph encoder. To model P (M |S), we first use a GNN (App. A.12) to encode G S into learned atom embeddings H = {h (a) ∀ a ∈ G S }. We condition the GNN on per-atom invariant shape features X = {x (a) } ∈ R n h ×6q , which we form by passingX through a VN-Inv (App. A.11): where H 0 ∈ R n h ×(da+d f ) are the set of initial atom features concatenated with the learned fragment embeddings, H ∈ R n h ×d h , and (·, ·) denotes concatenation in the feature dimension. For each atom in M S , we then encode h where (a) ∼ N (0, 1) ∈ R d h , H var ∈ R n h ×d h , and denotes elementwise multiplication. Here, the second argument of N (·, ·) is the standard deviation vector of the diagonal covariance matrix.
Mixing shape and variational features. The variational atom features H var are insensitive to rotations of S. However, we desire the decoder to construct molecules in poses that are natively aligned to S (Challenge 2). We achieve this by conditioning the decoder on an equivariant latent representation of P (M |S) that mixes both shape and chemical information. Specifically, we mix H var with X by encoding each h (a) var ∈ H var into linear transformations, which are applied atom-wise toX. We then pass the mixed equivariant features through a separate VN-MLP (App. A.11): H ∈ R q ×q ,X (a) ∈ R q×3 , andX Hvar ∈ R n h ×dz×3 . This maintains equivariance since W  var for each atom in M S , mixing H var withX, and passing the resultant (Z, z) to the decoder. We seed generation with a small structure M 0 (extracted from M S ), and build M by sequentially generating larger structures M l+1 in a tree-like manner (Fig.  2). Specifically, we grow new atoms/fragments around a "focus" atom/fragment in M l , which is popped from a BFS queue. To generate M l+1 from M l (e.g., grow the tree from the focus), we factor P (M l+1 |M l , S) = P (G l+1 |M l , S)P (G l+1 |G l+1 , M l , S). Given (Z, z), we sample the new graph G l+1 by iteratively attaching (a variable) C new atoms/fragments (children tree-nodes) around the focus, yielding G (c) l for c = 1, ..., C, where G (C) l = G l+1 and G (0) l = G l . We then generate coordinates G l+1 by scoring the (rotatable) bond between the focus and its parent tree-node. New bonds from the focus to its children are left unscored in M l+1 until the children become "in focus".
Partial molecule encoder. Before bonding each new atom/fragment to the focus (or scoring bonds), we encode the partial molecule M (c−1) l with the same scheme as for M S (using a parallel encoder; Fig. 2), except we do not variationally embed H . 2 Instead, we process H analogously to H var . Further, in addition to globally pooling the per-atom embeddings to obtainZ = Note thatZ dec ∈ R q×3 is equivariant to rotations of the overall system (M (c−1) l , M S ). Finally, we form a global invariant feature z dec ∈ R ddec to condition graph (or coordinate) generation: Graph generation. We factor P (G l+1 |M l , S) into a sequence of generation steps by which we iteratively connect children atoms/fragments to the focus until the network generates a (local) stop token. Fig. 2  . Given z dec , the model first predicts whether to stop (local) generation 2 We have dropped the (c) notation for clarity. However, each Zdec is specific to each (M (c−1) l , MS) system. via p ∅ = sigmoid(MLP(z dec )) ∈ (0, 1). If p ∅ ≥ τ ∅ (a threshold, App. A.16), we stop and proceed to bond scoring. Otherwise, we select which atom a foc on the focus (if multiple) to grow from: The decoder then predicts which atom/fragment f next ∈ L f to connect to the focus next: If the selected f next is a fragment, we predict the attachment site a site on the fragment G fnext : where h (a) fnext are the encoded atom features for G fnext . Lastly, we predict the bond order (1 • , 2 • , fnext )). We repeat this sequence of steps until p ∅ ≥ τ ∅ , yielding G l+1 . At each step, we greedily select the action after masking actions that violate known chemical valence rules. After each sequence, we bond a new atom or fragment to the focus, giving G (c) l . If an atom, the atom's position relative to the focus is fixed by heuristic bonding geometries (App. A.8). If a fragment, the position of the attachment site is fixed, but the dihedral of the new bond is yet unknown. Thus, in subsequent generation steps we only encode the attachment site and mask the remaining atoms in the new fragment until that fragment is "in focus" (Fig. 2). This means that prior to bond scoring, the rotation angle of the focus is random. To account for this when training (with teacher forcing), we randomize the focal dihedral when encoding each M Scoring rotatable bonds. After sampling G l+1 ∼ P (G l+1 |M l , S), we generate G l+1 by scoring the rotation angle ψ l+1 of the bond connecting the focus to its parent node in the generation tree ( Fig. 2). Since we ultimately seek to maximize sim S (M , M S ), we exploit the fact that our model generates shape-aligned structures to predict max for various query dihedrals ψ l+1 = ψ foc of the focus rotatable bond in a supervised regression setting. Intuitively, the scorer is trained to predict how the choice of ψ foc affects the maximum possible shape similarity of the final molecule M to the target M S under an optimal policy. App. A.2 details how regression targets are computed. During generation, we sweep over each query ψ foc ∈ [−π, π), encode each resultant structure M (ψfoc) l+1 into z (ψfoc) dec, scorer 3 , and select the ψ foc that maximizes the predicted score: At generation time, we also score chirality by enumerating stereoisomers G χ foc ∈ G foc of the focus and selecting the (G χ foc , ψ foc ) that maximizes Eq. 12 (App. A.2). Training. We supervise each step of graph generation with a multi-component loss function: L graph-gen = L ∅ +L focus +L next +L site +L bond +β KL L KL +β next-shape L next-shape +β ∅-shape L ∅-shape (13) L ∅ , L focus , L next , and L bond are standard cross-entropy losses.
) is a modified cross-entropy loss that accounts for symmetric attachment sites in the fragments site are the predicted attachment-site probabilities and c a are multi-hot class probabilities. L KL is the KL-divergence between the learned N (h µ , h σ ) and the prior N (0, 1). We also employ two auxiliary losses L next-shape and L ∅-shape in order to 1) help the generator distinguish between incorrect shape-similar (near-miss) vs. shape-dissimilar fragments, and 2) encourage the generator to generate structures that fill the entire target shape (App. A.10). We train the rotatable bond scorer separately from the generator with an MSE regression loss. See App. A.15 for training details.

EXPERIMENTS
Dataset. We train SQUID with drug-like molecules (up to n h = 27) from MOSES (Polykovskiy et al., 2020) using their train/test sets. L f includes 100 fragments extracted from the dataset and 24 atom types. We remove molecules that contain excluded fragments. For remaining molecules, we generate a 3D conformer with RDKit, set acyclic bond distances to their empirical means, and fix acyclic bond angles using heuristic rules. While this 3D manipulation neglects distorted bonding geometries in real molecules, the global shapes are marginally impacted, and we may recover refined geometries without seriously altering the shape (App. A.8). The final dataset contains 1.3M 3D molecules, partitioned into 80/20 train/validation splits. The test set contains 147K 3D molecules.
In the following experiments, we only consider molecules M S for which we can extract a small (3-6 atoms) 3D substructure M 0 containing a terminal atom, which we use to seed generation. In principle, M 0 could include larger structures from M S , e.g., for scaffold-constrained tasks. Here, we use the smallest substructures to ensure that the shape-conditioned generation tasks are not trivial.
Shape-conditioned generation of chemically diverse molecules. "Scaffold-hopping"-designing molecules with high 3D shape similarity but novel 2D graph topology compared to known inhibitors-is pursued in LBDD to develop chemical lead series, optimize drug activity, or evade intellectual property restrictions (Hu et al., 2017). We imitate this task by evaluating SQUID's ability to generate molecules M with high sim S (M , M S ) but low sim G (M , M S ). Specifically, for 1000 molecules M S with target shapes S in the test set, we use SQUID to generate 50 molecules per M S . To generate chemically diverse species, we linearly interpolate between the posterior N (h µ , h σ ) and the prior N (0, 1), sampling each h var ∼ N ((1 − λ)h µ , (1 − λ)h σ + λ1) using either λ = 0.3 or λ = 1.0 (prior). We then filter the generated molecules to have sim G (M , M S ) < 0.7, or < 0.3 to only evaluate molecules with substantial chemical differences compared to M S . Of the filtered molecules, we randomly choose N max samples and select the sample with highest sim S (M , M S ). Figure 3A plots distributions of sim S (M , M S ) between the selected molecules and their respective target shapes, using different sampling (N max = 1, 20) and filtering (sim G (M , M S ) < 0.7, 0.3) schemes. We compare against analogously sampling random 3D molecules from the training set. Overall, SQUID generates diverse 3D molecules that are quantitatively enriched in shape similarity compared to molecules sampled from the dataset, particularly for N max = 20. Qualitatively, the molecules generated by SQUID have significantly more atoms which directly overlap with the atoms of M S , even in cases where the computed shape similarity is comparable between SQUIDgenerated molecules and molecules sampled from the dataset (Fig. 3C). We quantitatively explore this observation in App. A.7. We also find that using λ = 0.3 yields greater sim S (M , M S ) than λ = 1.0, in part because using λ = 0.3 yields less chemically diverse molecules ( Fig. 3B; Challenge 3). Even so, sampling N max = 20 molecules from the prior with sim G (M , M S ) < 0.3 still yields more shape-similar molecules than sampling N max = 500 molecules from the dataset. We emphasize that 99% of samples from the prior are novel, 95% are unique, and 100% are chemically valid (App. A.4). Moreover, 87% of generated structures do not have any steric clashes (App. A.4), indicating that SQUID generates realistic 3D geometries of the flexible drug-like molecules. Ablating equivariance. SQUID's success in 3D shape-conditioned molecular generation is partly attributable to SQUID aligning the generated structures to the target shape in equivariant feature space (Eq. 7), which enables SQUID to generate 3D structures that fit the target shape without having to implicitly learn how to align two structures in R 3 (Challenge 2). We explicitly validate this design choice by settingZ = 0 in Eq. 7, which prevents the decoder from accessing the 3D orientation of M S during training/generation. As expected, ablating SQUID's equivariance reduces the enrichment in shape similarity (relative to the dataset baseline) by as much as 33% (App. A.9).
Shape-constrained molecular optimization. Scaffold-hopping is often goal-directed; e.g., aiming to reduce toxicity or improve bioactivity of a hit compound without altering its 3D shape. We mimic this shape-constrained MO setting by applying SQUID to optimize objectives from Gau-caMol (Brown et al., 2019) while preserving high shape similarity (sim S (M, M S ) ≥ 0.85) to various "hit" 3D molecules M S from the test set. This task considerably differs from typical MO tasks, which optimize objectives without constraining 3D shape and without generating 3D structures.
To adapt SQUID to shape-constrained MO, we implement a genetic algorithm (App. A.5) that iteratively mutates the variational atom embeddings H var of encoded seed molecules ("hits") M S in order to generate 3D molecules M * with improved objective scores, but which still fit the shape of M S . Table 1 reports the optimized top-1 scores across 6 objectives and 8 seed molecules M S (per objective, sampled from the test set), constrained such that sim S (M * , M S ) ≥ 0.85. We compare against the score of M S , as well as the (shape-constrained) top-1 score obtained by virtual screening (VS) our training dataset (>1M 3D molecules). Of the 8 seeds M S per objective, 3 were selected from top-scoring molecules to serve as hypothetical "hits", 3 were selected from top-scoring large molecules (≥ 26 heavy atoms), and 2 were randomly selected from all large molecules.
In 40/48 tasks, SQUID improves the objective score of the seed M S while maintaining sim S (M * , M S ) ≥ 0.85. Qualitatively, SQUID optimizes the objectives through chemical alterations such as adding/deleting individual atoms, switching bonding patterns, or replacing entire substructures -all while generating 3D structures that fit the target shape (App. A.5). In 29/40 of successful cases, SQUID (limited to 31K samples) surpasses the baseline of virtual screening 1M molecules, demonstrating the ability to efficiently explore new shape-constrained chemical space.

CONCLUSION
We designed a novel 3D generative model, SQUID, to enable shape-conditioned exploration of chemically diverse molecular space. SQUID generates realistic 3D geometries of larger molecules that are chemically valid, and uniquely exploits equivariant operations to construct conformations that fit a target 3D shape. We envision our model, alongside future work, will advance creative shape-based drug design tasks such as 3D scaffold hopping and shape-constrained 3D ligand design.

REPRODUCIBILITY STATEMENT
We have taken care to facilitate the reproduciblility of this work by detailing the precise architecture of SQUID throughout the main text; we also provide extensive details on training protocols, model parameters, and further evaluations in the Appendices. Additionally, our code is available at https://github.com/keiradams/SQUID. Beyond the model implementation, our code includes links to access our datasets, as well as scripts to process the training dataset, train the model, and evaluate our trained models across the shape-conditioned generation and shape-constrained optimization tasks described in this paper.

ETHICS STATEMENT
Advancing the shape-conditioned 3D generative modeling of drug-like molecules has the potential to accelerate pharmaceutical drug design, showing particular promise for drug discovery campaigns involving scaffold hopping, hit expansion, or the discovery of novel ligand analogues. However, such advancements could also be exploited for nefarious pharmaceutical research and harmful biological applications.   conditional distribution of 3D molecules given a molecular shape S M l partial 3D molecule (e.g., during generation) G l partial molecular graph (e.g., during generation) M 0 small 3D substructure used to seed generation M (c) l partial 3D molecule after bonding c atoms/fragments to the focus G (c) l partial molecular graph after bonding c atoms/fragments to the focus n h number of heavy atoms n p number of sampled points per heavy atom σ 2 p variance of isotropic atom-centered Gaussians in R 3 L f the atom/fragment library f i an atom/fragment in L f f i the learned embedding of atom/fragment f i ∈ L f h (a) learned embedding of atom a P S point cloud representation of shape S R an arbitrary rotation matrix in R 3×3 SO (3) the special orthogonal group (rotations) in R 3 SE (3) the special Euclidean group (rotations and translations) in R 3 X p tensor containing the learned equivariant features of each point in a point cloud;X p ∈ R (n h np)×q×3 X tensor containing the learned equivariant features of each atom in the encoded 3D molecule;X ∈ R n h ×q×3 q (1st) dimensionality of vector featuresX,X p H matrix containing the learned atom embeddings; H = {h (a) ∀a ∈ G} x (a) invariant shape features of atom a; x (a) ∈ R 6q X matrix containing the invariant shape features of each atom in the encoded 3D molecule; X ∈ R n h ×6q d a dimension of input atom features d f dimension of learned fragment embeddings d h dimension of learned atom embeddings h µ vector of means of the posterior distribution N (h µ , h σ ) h σ vector of standard deviations of the posterior distribution N (h µ , h σ ) h var sampled variational atom features λ interpolation factor between N (h µ , h σ ) and N (0, 1) N (0, 1) multidimensional standard normal distribution with µ = 0, σ = 1 sampled Gaussian noise vector for atom a; (a) ∼ N (0, 1) W (a) H learned linear transformation for atom a, applied toX (a) X Hvar tensor containing the per-atom equivariant shape features, after mixing with H var ;X Hvar ∈ R n h ×dz×3 X the predicted probability of stopping local generation τ ∅ probability threshold to stop local generation p focus predicted probabilities over the attachment sites on the focus p next predicted probabilities over the atom/fragment types in L f p site predicted probabilities over the attachment sites on the next fragment a foc the atom of attachment on the focus f next the atom/fragment to be added next G fnext the graph of the fragment to be added next a site the atom of attachment on the fragment to be added next ψ l+1 rotation angle (dihedral) of the bond connecting the focus to its parent (tree) node in the partially generated molecule ψ foc query dihedral angle when scoring the rotation angle of the focal rotatable bond (ψ l+1 = ψ foc ) M (ψfoc) l+1 the 3D molecule M l+1 after setting the query dihedral ψ l+1 = ψ foc z (ψfoc) dec, scorer analogue to z dec , but for the rotatable bond scorer when encoding the system (M S , M (ψfoc) l+1 ) G χ foc a stereoisomer of the focus G foc L graph-gen total loss for training graph generator L ∅ binary cross entropy loss for predicting to stop local generation L focus cross entropy loss for selecting the attachment site on the focus L next cross entropy loss for selecting which atom/fragment to add next L site modified cross entropy loss for selecting the attachment site on the next fragment L bond cross entropy loss for predicting the bond order L KL KL-divergence loss between N (h µ , h σ ) and N (0, 1) L next-shape auxiliary loss used when predicting the atom/fragment to add next L ∅-shape auxiliary loss used when predicting whether to stop local generation β KL weighting of L KL in L graph-gen β next-shape weighting of L next-shape in L graph-gen β ∅-shape weighting of L ∅-shape in L graph-gen Figure 4: Demonstration of conformer sampling when scoring rotatable bonds. For each query dihedral ψ foc = ψ 1 , ψ 2 , ... ∈ [−π, π), we sample N ψ conformations of the subgraph G Tfoc (keeping ψ foc held fixed) induced by the subtree T foc whose root node is the current focus. For each sampled (subgraph) conformation G Tfoc , we compute the shape similarity of between the sampled conformation and the ground truth (subgraph) conformation G S T foc . We select the maximum computed (nonaligned) shape similarity sim * S amongst the N ψ sampled conformations as the regression target for that ψ foc .

A.2 SCORING ROTATABLE BONDS AND STEREOCHEMISTRY
Recall that our goal is to train the scorer to predict max ψ l+2 ,ψ l+3 ,... sim * S (G (ψfoc) , G S ) for various query dihedrals ψ l+1 = ψ foc . That is, we wish to predict the maximum possible shape similarity of the final molecule M to M S when fixing ψ l+1 = ψ foc and optimally rotating all the yet-to-be-scored (or generated) rotatable bond dihedrals ψ l+2 , ψ l+3 , ... so as to maximize sim * S (G (ψfoc) , G S ). Training. We train the scorer independently from the graph generator (with a parallel architecture) using a mean squared error loss between the predicted scoresŝ (ψfoc) dec, scorer = sigmoid(MLP(z (ψfoc) dec, scorer )) and the regression targets s (ψfoc) for N s different query dihedrals ψ foc ∈ [−π, π): Computing regression targets. When training with teacher forcing (M l = M S l , G = G S ), we compute regression targets s ψfoc ≈ max ψ l+2 ,ψ l+3 ,... sim * S (G (ψfoc) , G S ) by setting the focal dihedral ψ l+1 = ψ foc , sampling N ψ conformations of the "future" graph G Tfoc induced by the subtree T foc whose root (sub)tree-node is the focus, and computing s ψfoc = max Tfoc , G S T foc ; α = 2.0). Since we fix bonding geometries, we need only sample N ψ sets of dihedrals of the rotatable bonds in G S T foc to sample N ψ conformers, making this conformer enumeration very fast. Note that rather than using α = 0.81 in these regression targets, we use α = 2.0 to make the scorer more sensitive to shape differences (App. A.7). When computing regression targets, we use N ψ < 1800 and select 36 (evenly spaced) ψ focus ∈ [−π, π) per rotatable bond. Figure 4 visualizes how regression targets are computed. App. A.15 contains further training specifics.
Scoring stereochemistry. At generation time, we also enumerate all possible stereoisomers of the focus (except cis/trans bonds) and score each stereoisomer separately, ultimately selecting the (stereoisomer, ψ foc ) pair that maximizes the predicted score. Figure 5 illustrates how we enumerate stereoisomers. Note that although we use the learned scoring function to score stereoisomerism at generation time, we do not explicitly train the scorer to score different stereoisomers.
Masking severe steric clashes. At generation time, we do not score any query dihedral ψ foc that causes a severe steric clash (< 1Å) with the existing partially generated structure (unless all query dihedrals cause a severe clash). Figure 5: In addition to scoring the dihedral of the focal bond, we also enumerate all stereoisomers around the focus and score each stereoisomer's focal dihedral separately. We then select the best scoring (steroisomer, dihedral) pair. We enumerate stereoisomers by combinatorially swapping chiral atoms, swapping axial/equatorial attachments to the focus (if relevant), and swapping the axial/equatorial attachment of the focus to the rest of the partial molecule (if relevant). Although swapping cis/trans bonds is feasible within SQUID's framework, we currently do not enumerate cis/trans isomers because of their minor presence in the dataset. Hence, any generated double bonds will have random stereochemistry. Figures 6 and 7 show additional random examples of molecules generated by SQUID when sampling N max = 1, 20 molecules with sim G (M , M S ) < 0.7 from the prior (λ = 1.0) or λ = 0.3 and selecting the sample with the highest sim S (M , M S ). Note that the visualized poses of the generated conformers are those which are directly generated by SQUID; the generated conformers have not been explicitly aligned to M S (e.g., using ROCS). Even so, the conformers are (for the most part) aligned to M S since SQUID's equivariance enables the model to generate natively aligned structures.

A.3 RANDOM EXAMPLES OF GENERATED 3D MOLECULES
It is apparent in these examples that using larger N max yields molecules with significantly improved shape similarity to M S , both qualitatively and quantitatively. This is in part caused by: 1) stochasticity in the variationally sampled atom embeddings H var ; 2) stochasticity in the input molecular point clouds, which are sampled from atom-centered isotropic Gaussians in R 3 ; 3) sampling sets of variational atom embeddings that may not be entirely self-consistent (e.g., for instance, if we sample only 1 atom embedding that implicitly encodes a ring structure); and 4) the choice of τ ∅ , the threshold for stopping local generation. While a small τ ∅ (we use τ ∅ = 0.01) helps prevent the model from adding too many atoms or fragments around a single focus, a small τ ∅ can also lead to early (local) stoppage, yielding molecules that do not completely fill the target shape. By sampling more molecules (using larger N max ), we have more chances to avoid these adverse random effects. Further work will attempt to improve the robustness of the encoding scheme and generation procedure in order to increase SQUID's overall sample efficiency. Figure 6: Random examples of molecules generated by SQUID using different sampling strategies. The generated 3D molecules are shown overlaid on the target molecule M S (blue). In these examples, the generated molecules have not been explicitly aligned to M S , and their displayed poses are those directly generated by SQUID.  Table 4 reports the percentage of molecules that are chemically valid, novel, and unique when sampling 50 molecules from the prior (λ = 1.0) for 1000 encoded molecules M S (e.g., target shapes) from the test set, yielding a total of 50K generated molecules. We define chemical validity as passing RDKit sanitization. Since we directly generate the molecular graph and mask actions which violate chemical valency, 100% of generated molecules are valid. We define novelty as the percentage of generated molecules whose molecular graphs are not present in the training data. We define uniqueness as the percentage of generated molecular graphs (of the 50K total) that are only generated once. For novelty and uniqueness calculations, we consider different stereoisomers to have the same molecular graph. We also report the percentage of generated 3D structures that have an apparent steric clash, defined to be a non-bonded interatomic distance below 2Å. Table 5 reports the graph reconstruction accuracy when sampling 3D molecules from the posterior (λ = 0.0), for 1000 target molecules M S from the test set. We report the top-k graph reconstruction accuracy (ignoring stereochemical differences) when sampling k = 1 molecule per encoded M S , and when sampling k = 20 molecules per encoded M S . Since we have intentionally trained SQUID inside a shape-conditioned variational autoencoder framework in order to generate chemically diverse molecules with similar 3D shapes, the significance of graph reconstruction accuracy is debatable in our setting. However, it is worth noting that the top-1 reconstruction accuracy is 16.3%, while the top-20 reconstruction accuracy is much higher (57.2%). This large difference is likely attributable to both stochasticity in the variational atom embeddings and stochasticity in the input 3D point clouds. Table 4: Generation statistics computed by sampling 50 molecules from the prior (λ = 1.0) per encoded target molecule M S for 1000 random targets from the test set.

A.5.1 GENETIC ALGORITHM
We adapt SQUID to shape-constrained molecular optimization by implementing a genetic algorithm on the variational atom embeddings H var . Algorithm 1 details the exact optimization procedure. In summary, given the seed molecule M S with a target 3D shape and an initial substructure M 0 (which is contained by all generated molecules for a given M S ), we first generate an initial population of generated molecules M by repeatedly sampling H var for various interpolation factors λ, mixing these H var with the encoded shape features of M S , and decoding new 3D molecules. We only add a generated molecule to the population if sim S (M , M S ) ≥ τ S (we use τ S = 0.75), so that the GA does not overly explore regions of chemical space that have no chance of satisfying the ultimate constraint sim S (M , M S ) ≥ 0.85. After generating the initial population, we iteratively 1) select the top-scoring samples in the population, 2) cross the top-scoring H var in crossover events, 3) mutate the top and crossed H var via adding random noise, and 4) generate new molecules M for each mutated H var . The final optimized molecule M * is the top-scoring generated molecule that satisfies the shape-similarity constraint sim S (M , M S ) ≥ 0.85. Figure 8 visualizes the structures of the SQUID-optimized molecules M * and their respective seed molecules M S (e.g., the starting "hit" molecules with target shapes) for each of the optimization tasks which led to an improvement in the objective score. We also overlay the generated 3D conformations of M * on those of M S , and report the objective scores for each M * and M S .  A.6 COMPARING SIM S TO ROCS SCORING FUNCTION Our shape similarity function described in Equation 1 closely approximates the shape (only) scoring function employed by ROCS, when α = 0.81. Figure 9 demonstrates the near-perfect correlation between our computed shape scores and those computed by ROCS for 50,000 shape comparisons, with a mean absolute error of 0.0016. Note that Equation 1 computes non-aligned shape similarity. We still employ ROCS to align the generated molecules M to the target molecule M S before computing their (aligned) shape similarity in our experiments. However, we do not require explicit alignment when training SQUID; we do not use the commercial ROCS program during training.

A.7 EXPLORING DIFFERENT VALUES OF α IN SIM S
Our analysis of shape similarity thus far has used Equation 1 with α = 0.81 in order to recapitulate the shape similarity function used by ROCS, which is widely used in drug discovery. However, compared to randomly sampled molecules in the dataset, the molecules generated by SQUID qualitatively appear to do a significantly better job at fitting the target shape S on an atom-by-atom basis, even if the computed shape similarities (with α = 0.81) are comparable (see examples in Figure 3). We quantify this observation by increasing the value of α when computing sim S (M , M S ; α) for generated molecules M , as α is inversely related to the width of the isotropic 3D Gaussians used in the volume overlap calculations in Equation 1. Intuitively, increasing α will greater penalize sim S if the atoms of M and M S do not perfectly align. Figure 10 plots the mean sim S (M, M S ; α) for the most shape-similar molecule M of N max sampled molecules M for increasing values of α. Averages are calculated over 1000 target molecules M S from the test set, and we only consider generated molecules for which sim G (M , M S ) < 0.7. Crucially, the gap between the mean sim S (M, M S ; α) obtained by generating molecules with SQUID vs. randomly sampling molecules from the dataset significantly widens with increasing α. This effect is especially apparent when using SQUID with λ = 0.3 and N max = 20, although can be observed with other generation strategies as well. Hence, SQUID does a much better job at generating (still chemically diverse) molecules that have significant atom-to-atom overlap with M S .

A.8 HEURISTIC BONDING GEOMETRIES AND THEIR IMPACT ON GLOBAL SHAPE
In all molecules (dataset and generated) considered in this work, we fix acyclic bond distances to their empirical averages and set acyclic bond angles to heuristic values based on hybridization rules in order to reduce the degrees of freedom in 3D coordinate generation. Here, we describe how we fix these bonding geometries and explore whether this local 3D structure manipulation significantly alters the global molecular shape.
Fixing bonding geometries. We fix acyclic bond distances by computing the mean bond distance between pairs of atom types across all the RDKit-generated conformers in our training set. After collecting these empirical mean values, we manually set each acyclic bond distance to its respective mean value for each conformer in our datasets.
We set acyclic bond angles using simple hybridization rules. Specifically, sp3-hybridized atoms will have bond angles of 109.5 • , sp2-hybridized atoms will have bond angles of 120 • , and sp-hybridized atoms will have bond angles of 180 • . We manually fix the acyclic bond angles to these heuristic values for all conformers in our datasets. We use RDKit to determine the hybridization states of each atom. During generation, occasionally the hybridization of certain atoms (N, O) may change once they are bonded to new neighbors. For instance, an sp3 nitrogen can become sp2 once bonded to an aromatic ring. We adjust bond angles on-the-fly in these edge cases.
Impact on global shape. Figure 11 plots the histogram of sim S (M fixed , M relaxed ) for 1000 test set conformers M fixed whose bonding geometries have been fixed, and the original RDKitgenerated conformers M relaxed with relaxed (true) bonding geometries. In the vast majority of cases, fixing the bonding geometries negligibly impacts the global shape of the 3D molecule (sim S (M fixed , M relaxed ) ≈ 1). This is because the main factor influencing global molecular shape is rotatable bonds (e.g., flexible dihedrals), which are not altered by fixing bond distances and angles.
Recovering refined bonding geometries. Even though fixing bond distances and angles only marginally impacts molecular shape, we still may wish to recover refined bonding geometries of the generated 3D molecules without altering the generated 3D shape. We can accomplish this (to a first approximation) for generated molecules by creating a geometrically relaxed conformation of the generated molecular graph with RDKit, and then manually setting the dihedrals of the rotatable bonds in the relaxed conformer to match the corresponding dihedrals in the generated conformers. Importantly, if we perform this relaxation procedure for both the dataset molecules and the SQUID-generated molecules, the (relaxed) generated molecules still have significantly enriched shape-similarity to the (relaxed) target shape compared to (relaxed) random molecules from the dataset (Fig. 12).
Figure 11: Histogram of shape similarity sim S (M fixed , M relaxed ) between 1000 test molecules whose acylic bonding geometries are fixed by heuristics and their geometrically relaxed counterparts.  Overall, ablating SQUID's equivariance signficantly reduces the enrichment in shape similarity of generated molecules to the target shape.
A.9 ABLATING EQUIVARIANCE SQUID aligns the equivariant representations of the encoded target shape and the partially generated structures in order to generate 3D conformations that natively fit the target shape, without having to implicitly learn SE(3)-alignments (Challenge 2). We achieve this in Equation 7, where we mix the equivariant representations of M S and the partially generated structure M (c−1) l . To empirically motivate this design choice, we ablate the equivariant alignment by settingZ = 0 in Eq. 7. We denote this ablated model as SQUID-NoEqui. Note that because we still pass the unablated invariant features z to the decoder (Eq. 8), SQUID-NoEqui is still conditioned on the shape of M S -the model simply no longer has access to any explicit information about the relative spatial orientation of M (c−1) l to M S (and thus must learn this spatial relationship from scratch). As expected, ablating SQUID's equivariance significantly reduces SQUID's ability to generate chemically diverse molecules that fit the target shape. Figure 13 plots the distributions of sim S (M , M S ) for the best of N max generated molecules with sim G (M , M S ) < 0.7 or 0.3 when using SQUID or SQUID-NoEqui. Crucially, the mean shape similarity when sampling with (λ = 1.0, N max = 20, sim G (M , M S ) < 0.7) decreases from 0.828 (SQUID) to 0.805 (SQUID-NoEqui). When sampling with (λ = 0.3, N max = 20, sim G (M , M S ) < 0.7), the mean shape similarity also decreases from 0.879 (SQUID) to 0.839 (SQUID-NoEqui). Relative to the mean shape similarity of 0.758 achieved by sampling random molecules from the dataset (N max = 20, sim G (M , M S ) < 0.7), this corresponds to a substantial 33% reduction in the shape-enrichment of SQUID-generated molecules.
Interestingly, sampling (λ = 1.0, N max = 20, sim G (M , M S ) < 0.7) with SQUID-NoEqui still yields shape-enriched molecules compared to analogously sampling random molecules from the dataset (mean shape similarity of 0.805 vs. 0.758). This is because even without the equivariant feature alignment, SQUID-NoEqui still conditions molecular generation on the (invariant) encoding of the target shape S, and hence biases generation towards molecules which better fit the target shape (after alignment with ROCS).

A.10 AUXILIARY TRAINING LOSSES
We employ two auxiliary losses when training the graph generator in order to encourage the generated graphs to better fit the encoded target shape.
The first auxiliary loss penalizes the graph generator if it adds an incorrect atom/fragment to the focus that is of significantly different size than the correct (ground truth) atom/fragment. We first compute a matrix ∆V f ∈ R (|L f |×|L f |) + containing the (pairwise) volume difference between all atoms/fragments in the library L f where v fi is the volume of atom/fragment f i ∈ L f (computed with RDKit).
We then compute the auxiliary loss L next-shape as: where g is the index of the correct (ground truth) next atom/fragment f next, true , ∆V (g) f is the gth row of ∆V f , and p next are the predicted probabilities over the next atom/fragment types to be connected to the focus (see Eq. 10).
The second auxiliary loss penalizes the graph generator if it prematurely stops (local) generation, with larger penalties if the premature stop would result in larger portions of the (ground truth) graph not being generated. When predicting (local) stop tokens during graph generation (with teacher forcing), we compute the number of atoms in the subgraph induced by the subtree whose root treenode is the next atom/fragment to be added to the focus (in the current generation sequence). We then multiply the predicted probability for the local stop token by this number of "future" atoms that would not be generated if a premature stop token were generated. Hence, if the correct action is to indeed stop generation around the focus, the penalty will be zero. However, if the correct action is to add a large fragment to the current focus but the generator predicts a stop token, the penalty will be large. Formally, we compute: where p ∅, true is the ground truth action for local stopping (p ∅, true = 0 indicates that the correct action is to not stop local generation), and G S T next is the subgraph induced by the subtree whose root node is the next atom/fragment (to be generated) in the ground-truth molecular graph. where wherek = UX for a learnable weight matrix U ∈ R 1×q , and wherex ∈ R 3 .
where KNN f (n) are the k-nearest neighbors of point m in feature space, Φ (t) and Θ (t) are weight matrices, andX n ∈ R q×3 are the per-point equivariant features.

A.12 GRAPH NEURAL NETWORKS
In this work, we employ graph neural networks (GNNs) to encode: • each atom/fragment in the library L f  (Satorras et al., 2022b). Given a molecular graph G with atoms as nodes and bonds as edges, we use graph convolutional layers defined by the following: where h t i are the learned atom embeddings at each GNN-layer, m t ij are learned (directed) messages, r i ∈ R 3 are the coordinates of atom i, N (i) is the set of 1-hop bonded neighbors of atom i, and each φ t m , φ t h are MLPs. Note that h 0 i are the initial atom features, and m 0 ij are the initial bond features for the bond between atoms i and j. In general, m t ij = m t ji for t > 0, but here m 0 ij = m 0 ji . Note that since we only aggregate messages from directly bonded neighbors, ||r i − r j || only encodes bond distances, and does not encode any information about specific 3D conformations. Hence, our GNNs effectively only encode 2D chemical identity, as opposed to 3D shape. Our atom/fragment library L f includes 100 distinct fragments (Fig. 14) and 24 unique atom types. The 100 fragments were selected based on the top-100 most frequently occurring fragments in our training set. In this work, we specify fragments as ring-containing substructures that do not contain any acyclic single bonds. However, in principle fragments could be any (valid) chemical substructure. Note that we only use 1 (geometrically optimized) conformation per fragment, which is assumed to be rigid. Hence, in its current implementation, SQUID does not consider different ring conformations (e.g., boat vs. chair conformations of cyclohexane).
A.14 MODEL PARAMETERS Parameter sharing. For both the graph generator and the rotatable bond score, the (variational) molecule encoder (in the Encoder, Fig. 2) and the partial molecule encoder (in the Decoder, Fig. 2) share the same fragment encoder (L f -GNN), which is trained end-to-end with the rest of the model. Apart from L f -GNN, these encoders do not share any learnable parameters, despite having parallel architectures. The graph generator and the rotatable bond scorer are completely independent, and are trained separately.
Hyperparameters. Tables 6 and 7 tabulate the set of hyperparameters used for SQUID across all the experiments conducted in this paper. Table 8 summarizes training and generation parameters, but we refer the reader to App. A.15 and A.16 for more detailed discussion of training and generation protocols.
Because of the large hyperparameter search space and long training times, we did not perform extensive hyperparameter optimizations. We manually tuned the learning rates and schedulers to maintain training stability, and we maxed-out batch sizes given memory constraints. We set β ∅-shape = 10 and β next-shape = 10 to make the magnitudes of L ∅-shape and L next-shape comparable to the other loss components for graph-generation. We slowly increase β KL over the course of training from 10 −5 to a maximum of 10 −1 , which we found to provide a reasonable balance between L KL and graph reconstruction. Table 6: Architecture hyperparameters for SQUID's graph generator. Many of the GNN, VN-DGCNN, Vn-Inv, VN-MLP, and MLP modules use the same set of hyperparameters (not learnable parameters), but some do not. Where it appears, (→) indicates the outputs of the module(s) to which we refer (see Fig. 2). The graph generator has a total of 1468103 learnable parameters.

Symbol Parameter Description
Value σ 2 p variance of atom-centered Gaussians 0.049Å 2 n p number of sampled points per atom 5 d f dimensionality of atom/fragment embeddings Table 7: Architecture hyperparameters for SQUID's rotatable bond scorer. Many of the GNN, VN-DGCNN, Vn-Inv, VN-MLP, and MLP modules use the same set of hyperparameters (not learnable parameters), but some do not. Where it appears, (→) indicates the outputs of the module(s) to which we refer (see Fig. 2). The rotatable bond scorer has a total of 1270080 learnable parameters.
Training times. Training the graph generator on 1 GPU takes ∼ 5-7 days. Training the rotatable bond scorer on 1 GPU takes ∼ 5 days.
Dataset processing times. Pre-computing the regression targets for the rotatable bond scorer is the most compute-intensive part of dataset pre-processing. We required ∼ 2-3 days when parallelizing across 384 cpu cores. Figure 15: Examples of starting seed structures M 0 highlighting where new atoms/fragments might be attached (orange).

A.16 ADDITIONAL GENERATION DETAILS
Seeding generation. When generating new molecules M from M S , we start by extracting a small structure M 0 from M S to seed generation. In our experiments, M 0 is determined by the following criteria: • M 0 has ≤ 6 heavy atoms • M 0 contains at least 1 terminal atom (an atom with only 1 bonded neighbor) • Any new atom/fragment that is bonded to M 0 forms a valid dihedral. For this to be satisfied, M 0 must contain at least 3 atoms.
• M 0 does not contain a sequence of 4 atoms that forms a valid dihedral of a flexible rotatable bond. Hence, all rotatable bonds are scored during generation. We note that the above seeding procedure could be replaced by a suitable model that builds and (equivariantly) predicts the coordinates of the starting structure M 0 . In this work, we extract seeds for simplicity, but use the above criteria to ensure that the downstream shape-conditioned 3D generation and design tasks are non-trivial. We also emphasize that our choice of seeding procedure is fairly arbitrary, and does not necessarily restrict how the model is used in other generation tasks. For instance, larger seed structures could be used.
Masking invalid actions. When generating the molecular graph, we mask actions which would violate chemical valency. Because we define atom types in part by the number of single, double, and triple bonds the atom forms (including bonds to implicit hydrogens), we create valency masks by ensuring that each atom forms the correct number and types of bonds. When scoring rotatable bonds, we also mask query dihedrals that would lead to a severe steric clash (<1Å) with the existing partially generated molecular structure (App. A.2).
Bond scoring. When scoring a rotable bond at generation time, we sample 36 dihedral angles separated by 10 • . In principle, one could use finer discretization schemes to generate more refined 3D conformations, at the cost of speed and/or memory.
Threshold for stopping local generation. We manually tuned τ ∅ , the threshold for local stopping, to τ ∅ = 0.01 in order to balance tendencies of the model to 1) prematurely stop local generation (thus leading to molecules that do not completely fill the target shape), and 2) generate extra atoms/fragments around the focus (leading to steric crowding).
Generation time. Typically, generating (in serial, on a cpu) one 3D molecule with 20-30 atoms takes 2-3 seconds. Overall generation time (on a cpu) and memory cost significantly increase if many stereoisomers are enumerated when scoring any particular rotatable bond (App. A.2). We thus cap the number of enumerated stereoisomers to 32 (per focus).