Smoothing Dark Areas in Molecular Latent Diffusion

# Smoothing Dark Areas in Molecular Latent Diffusion
Source: [https://arxiv.org/html/2606.13955](https://arxiv.org/html/2606.13955)
Xi Wang1Jiahan Li1Yuxuan Xia1Yingcheng Wu2Shaoyi Zheng1Shengjie Wang1 1New York University 2Stanford University

###### Abstract

Latent diffusion is a promising framework for scalable 3D molecular generation, but it requires a latent space that remains smooth, valid, and navigable beyond posterior samples\. Existing molecular VAEs, however, are typically learned through reconstruction\-based objectives, which do not guarantee such a latent space\. We show that this leads todark areas: regions of latent space that are reachable during diffusion sampling but decode to disconnected or chemically invalid molecules\. Unlike in image generation, molecular decoding requires strict structural and chemical precision, so even small latent perturbations can produce catastrophic failures\. We therefore proposeTopVAE, a topology\-optimized VAE that reduces dark areas by making the decoder internalize structural and chemical constraints during training, eliminating the need for test\-time chemical correction\. TopVAE greatly improves off\-posterior robustness, and when paired with a standard DiT, achieves77%77\\%lower FCD3Don QM9, the highest V&C,52%52\\%lower FCD3Don GEOM\-Drugs, and1\.29×1\.29\{\\times\}more stable and connected molecules on zero\-shot scaffold inpainting\.

## 1Introduction

Latent diffusion models, originally developed for image generation\(Rombachet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib1); Peebles and Xie,[2023](https://arxiv.org/html/2606.13955#bib.bib2)\), have also been extended to 3D molecular generation\(Xuet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib14); Chenet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib18)\)\. This setting is more challenging because valid molecules must simultaneously satisfy structural, geometric, and chemical constraints\(Youet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib15); Luoet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib16)\)\. A latent\-variable formulation is therefore especially appealing: if these complexities can be absorbed by the decoder, diffusion can operate in a cleaner latent space, simplifying modeling and sampling\.

Molecular latent spaces are typically learned with reconstruction objectives plus regularization, as in VAE or VQ\-VAE frameworks\(Luoet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib16)\)\. However, strong reconstruction does not guarantee a latent space suitable for diffusion\. Recent work shows that reconstruction quality can poorly predict, and sometimes even conflict with, downstream generation performance\(Yaoet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib22); Skorokhodovet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib24); Xuet al\.,[2026](https://arxiv.org/html/2606.13955#bib.bib25)\), because diffusion requires a smooth, interpolatable latent manifold rather than correctness only at posterior samples\. Although latent regularization is intended to enlarge the region of valid latents, this is particularly difficult for molecules: molecular decoding demands strict structural and chemical precision, and unlike images, it lacks strong architectural biases that make nearby latent points decode naturally\. As a result, even small perturbations around posterior latents can produce severely corrupted molecules, including disconnected structures or graphs that fail chemical sanitization \(Fig\.[1](https://arxiv.org/html/2606.13955#S1.F1)\)\. We call these failure regionsdark areas: latent neighborhoods where the decoder breaks, rendering the surrounding manifold non\-interpolatable and difficult to navigate for diffusion\.

Dark areas arise mainly from two sources:*topological disconnection*, where the decoded graph fractures into isolated components, and*chemical invalidity*, where predicted bonds violate valence or compatibility rules\. Prior work often addresses these issues by enforcing chemical rules during generation\(Jinet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib29); Liuet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib30); Maet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib31); Krennet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib32)\)\. While such constraints can repair outputs, they do not remove the underlying latent\-space fragility and may also bias generation\. Instead, we aim to internalize these constraints during VAE training so that valid molecules occupy a more robust latent neighborhood, without requiring chemical constraint optimization \(ChemCO\) at inference time\. To this end, we proposeTopVAE\(Topology\-OPtimizedVAE\), which promotes connectivity through BFS\-based adjacency refinement, enforces valence and bond\-type constraints via unrolled primal–dual optimization during training, and selectively injects ChemCO’s corrections into the decoder\(Ohet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib40)\), enabling ChemCO\-free inference\.

In summary, our contributions are:

1. 1\.Dark areas in molecular latent space\.We identify and formalize dark areas: latent regions reachable by diffusion sampling but decoded as chemically invalid or disconnected molecules\.
2. 2\.TopVAEwith three components: a\) TopoBridge guarantees connectivity via adjacency refinement; b\) inherent ChemCO provides chemical constraints through unrolled primal–dual optimization; c\) Advantage\-Gated Constratint Learning \(AGCL\) that selectively feeds constraint\-based correction signals into the decoder training\.
3. 3\.State\-of\-the\-art generation with dark\-area closure\.TopVAE paired with a standard DiT achieves77%77\\%lower FCD3Don QM9, the highest V&C and52%52\\%lower FCD3Don GEOM\-Drugs, and1\.29×1\.29\{\\times\}more stable and connected molecules on zero\-shot scaffold inpainting\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/figures/main.png)Figure 1:Overview of TopVAE\.\(A\) Dark\-area diagnosis reveals that existing VAE decoders produce invalid molecules in off\-posterior latent regions, which TopVAE closes by internalizing chemical rules\. \(B\) TopVAE combines a topology\-first multi\-stage decoder with TopoBridge for guaranteed connectivity and ChemCO for chemical constraint learning during training\.
## 2Related Work

#### 3D molecular generation\.

Diffusion and flow models for de novo 3D molecules fall into two families\.*Data\-space*models generate directly in molecular data space: some diffuse atom types and coordinates and infer bonds post hoc\(Hoogeboomet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib19)\), while later graph\-aware or complete\-molecule variants reduce atom–bond mismatch by jointly modeling graph/topology and geometry, by explicitly predicting bond variables, or by introducing bond\-formation\-aware training objectives\(Penget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib9); Huanget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib10); Vignacet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib12); Leet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib13); Reidenbachet al\.,[2026](https://arxiv.org/html/2606.13955#bib.bib27); Xuet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib28)\)\.*Latent*models first compress molecules into continuous representations and then fit a diffusion model\(Xuet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib14); Youet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib15); Luoet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib16); Joshiet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib17)\), gaining scalability and controllability but requiring high quality latents for diffusion models\. Our dark\-area diagnostic makes this assumption explicit and testable, motivating a decoder\-centric complement to existing latent molecular diffusion models\.

#### Constraint\-aware molecular decoding\.

Chemical and structural validity has been enforced via grammar\-level constraints\(Kusneret al\.,[2017](https://arxiv.org/html/2606.13955#bib.bib35); Krennet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib32)\), structured graph decoders\(Jinet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib29); Liuet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib30)\), validity\-oriented regularization\(Maet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib31)\), and differentiable constraint\-satisfaction layers\(Wanget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib49); Zenget al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib50)\)\. TopVAE takes a different route: rather than retaining a permanent constraint layer at inference, we use constraint\-guided corrections from ChemCO as a selective training signal using AGCL, with the goal that the unconstrained decoder itself internalizes chemical rules and supports ChemCO\-free inference \(Sec\.[4\.5](https://arxiv.org/html/2606.13955#S4.SS5), Fig\.[2](https://arxiv.org/html/2606.13955#S4.F2)\)\. Extended related work is in Appendix[A](https://arxiv.org/html/2606.13955#A1)\.

## 3Prerequisites

#### 3D molecular VAE\.

A 3D molecule can represented asM=\(𝑨,𝑻,𝑩,𝑹\)M=\(\\bm\{A\},\\bm\{T\},\\bm\{B\},\\bm\{R\}\), where𝑨∈\{0,1\}N×N\\bm\{A\}\\in\\\{0,1\\\}^\{N\\times N\}is the binary adjacency matrix,𝑻\\bm\{T\}denotes atom types,𝑩\\bm\{B\}denotes bond types, and𝑹∈ℝN×3\\bm\{R\}\\in\\mathbb\{R\}^\{N\\times 3\}stores atomic coordinates\. A molecular VAE encodesMMinto a latent variable𝒛\\bm\{z\}and reconstructs the molecule through a decoderDθD\_\{\\theta\}:

qϕ​\(𝒛∣M\)=𝒩​\(𝝁ϕ​\(M\),diag⁡\(𝝈ϕ2​\(M\)\)\),M^=Dθ​\(𝒛\),𝒛=𝝁ϕ​\(M\)\+𝝈ϕ​\(M\)⊙ϵ\.q\_\{\\phi\}\(\\bm\{z\}\\mid M\)=\\mathcal\{N\}\\\!\\left\(\\bm\{\\mu\}\_\{\\phi\}\(M\),\\operatorname\{diag\}\(\\bm\{\\sigma\}\_\{\\phi\}^\{2\}\(M\)\)\\right\),\\qquad\\hat\{M\}=D\_\{\\theta\}\(\\bm\{z\}\),\\quad\\bm\{z\}=\\bm\{\\mu\}\_\{\\phi\}\(M\)\+\\bm\{\\sigma\}\_\{\\phi\}\(M\)\\odot\\bm\{\\epsilon\}\.\(1\)Hereqϕ​\(𝒛∣M\)q\_\{\\phi\}\(\\bm\{z\}\\mid M\)is the encoder posterior, parameterized by encoder parametersϕ\\phi, which maps an input molecule to the mean and variance of a Gaussian latent distribution\. The VAE is trained by balancing reconstruction and latent regularization:

ℒVAE=𝔼M∼𝒟​𝔼𝒛∼qϕ​\(𝒛∣M\)​\[−log⁡pθ​\(M∣𝒛\)\]\+β​KL​\(qϕ​\(𝒛∣M\)∥p​\(𝒛\)\)\.\\mathcal\{L\}\_\{\\mathrm\{VAE\}\}=\\mathbb\{E\}\_\{M\\sim\\mathcal\{D\}\}\\mathbb\{E\}\_\{\\bm\{z\}\\sim q\_\{\\phi\}\(\\bm\{z\}\\mid M\)\}\\big\[\-\\log p\_\{\\theta\}\(M\\mid\\bm\{z\}\)\\big\]\+\\beta\\,\\mathrm\{KL\}\\\!\\left\(q\_\{\\phi\}\(\\bm\{z\}\\mid M\)\\,\\\|\\,p\(\\bm\{z\}\)\\right\)\.\(2\)

#### Latent diffusion over molecular latents\.

A diffusion model is trained on the VAE latent space to generate new latent codes\. Starting from a clean latent𝒛0\\bm\{z\}\_\{0\}, the forward process adds Gaussian noise:

𝒛t=α¯t​𝒛0\+1−α¯t​ϵ,ϵ∼𝒩​\(𝟎,𝑰\)\.\\bm\{z\}\_\{t\}=\\sqrt\{\\bar\{\\alpha\}\_\{t\}\}\\,\\bm\{z\}\_\{0\}\+\\sqrt\{1\-\\bar\{\\alpha\}\_\{t\}\}\\,\\bm\{\\epsilon\},\\qquad\\bm\{\\epsilon\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\bm\{I\}\)\.\(3\)A denoiserϵψ​\(𝒛t,t\)\\bm\{\\epsilon\}\_\{\\psi\}\(\\bm\{z\}\_\{t\},t\)is trained to predictϵ\\bm\{\\epsilon\}, typically with an MSE loss\. At sampling time, iterative denoising produces a clean latent𝒛^0\\hat\{\\bm\{z\}\}\_\{0\}, which is decoded into a molecule:M^=Dθ​\(𝒛^0\)\\hat\{M\}=D\_\{\\theta\}\(\\hat\{\\bm\{z\}\}\_\{0\}\)\. Therefore, the decoder must work not only on posterior latents, but also on noisy and prior\-like latents visited during diffusion\.

#### Dark areas in molecular latent space\.

At inference time, a latent diffusion model may sample latents that lie between or outside the posterior codes seen during training\. We call a decoded graph*chemically valid*if it \(i\) passes RDKit sanitization—requiring legal valences, consistent aromaticity, and recognized bond types—and \(ii\) forms a single connected component\. Denoting the sampling\-reachable region by𝒜LDM\\mathcal\{A\}\_\{\\mathrm\{LDM\}\}and the valid molecule set byℳvalid\\mathcal\{M\}\_\{\\mathrm\{valid\}\}, we define thedark areasas

𝒟dark=\{𝒛∈𝒜LDM:Dθ​\(𝒛\)∉ℳvalid\}\.\\mathcal\{D\}\_\{\\mathrm\{dark\}\}=\\\{\\bm\{z\}\\in\\mathcal\{A\}\_\{\\mathrm\{LDM\}\}:D\_\{\\theta\}\(\\bm\{z\}\)\\notin\\mathcal\{M\}\_\{\\mathrm\{valid\}\}\\\}\.\(4\)

## 4TopVAE: Topology\-Optimized VAE for Chemical Constraint Learning

TopVAE targets dark areas in molecular latent space\. The goal is to make the VAE decoder learn chemical rules, closing dark\-areas\. TopVAE therefore combines a structured molecular decoder with two constraint\-aware modules\. TopoBridge refines the predicted adjacency into a guaranteed connected adjacency matrix\. ChemCO unrolls chemical constraint optimization and converts valence limits, bond exclusivity, atom\-pair legality, and degree rules into differentiable training signals\.

### 4\.1Encoder

The encoder maps the molecular graph into latent features with a Relational Transformer\(Diao and Loynd,[2022](https://arxiv.org/html/2606.13955#bib.bib46)\)\. We first embed atoms and atom pairs:

𝒉i\(0\)=fenc−node​\(Ti,𝑹i\),𝒑i​j=fenc−edge​\(Bi​j,‖𝑹i−𝑹j‖2\)\.\\bm\{h\}\_\{i\}^\{\(0\)\}=f\_\{\\mathrm\{enc\-node\}\}\(T\_\{i\},\\bm\{R\}\_\{i\}\),\\qquad\\bm\{p\}\_\{ij\}=f\_\{\\mathrm\{enc\-edge\}\}\(B\_\{ij\},\\\|\\bm\{R\}\_\{i\}\-\\bm\{R\}\_\{j\}\\\|\_\{2\}\)\.\(5\)The Relational Transformer updates atom features by edge\-aware attention, where pair features𝒑i​j\\bm\{p\}\_\{ij\}are injected into the attention between atomsiiandjj:

𝑯=RTransϕ⁡\(\{𝒉i\(0\)\}i=1N,\{𝒑i​j\}i,j=1N\)\.\\bm\{H\}=\\operatorname\{RTrans\}\_\{\\phi\}\\left\(\\\{\\bm\{h\}\_\{i\}^\{\(0\)\}\\\}\_\{i=1\}^\{N\},\\;\\\{\\bm\{p\}\_\{ij\}\\\}\_\{i,j=1\}^\{N\}\\right\)\.\(6\)The resulting atom\-level features𝑯=\(𝒉1,…,𝒉N\)\\bm\{H\}=\(\\bm\{h\}\_\{1\},\\ldots,\\bm\{h\}\_\{N\}\)are used to parameterize the latent distribution in Eq\. \([1](https://arxiv.org/html/2606.13955#S3.E1)\)\.

### 4\.2Topology\-aware multi\-stage decoder

Using the notation in Sec\.[3](https://arxiv.org/html/2606.13955#S3), TopVAE decodes each latent code𝒛\\bm\{z\}in a topology\-first order:

pθ​\(𝑨,𝑻,𝑩,𝑹∣𝒛\)=pθ​\(𝑨∣𝒛\)​pθ​\(𝑻∣𝑨,𝒛\)​pθ​\(𝑩∣𝑨,𝑻,𝒛\)​pθ​\(𝑹∣𝑨,𝑻,𝑩,𝒛\)\.p\_\{\\theta\}\(\\bm\{A\},\\bm\{T\},\\bm\{B\},\\bm\{R\}\\mid\\bm\{z\}\)=p\_\{\\theta\}\(\\bm\{A\}\\mid\\bm\{z\}\)\\,p\_\{\\theta\}\(\\bm\{T\}\\mid\\bm\{A\},\\bm\{z\}\)\\,p\_\{\\theta\}\(\\bm\{B\}\\mid\\bm\{A\},\\bm\{T\},\\bm\{z\}\)\\,p\_\{\\theta\}\(\\bm\{R\}\\mid\\bm\{A\},\\bm\{T\},\\bm\{B\},\\bm\{z\}\)\.\(7\)This factorization makes adjacency the first decoded object\. In particular, bond prediction is constrained by

Ai​j=0⇒Bi​j=0,Ai​j=1⇒Bi​j∈\{1,…,K\}\.A\_\{ij\}=0\\Rightarrow B\_\{ij\}=0,\\qquad A\_\{ij\}=1\\Rightarrow B\_\{ij\}\\in\\\{1,\\ldots,K\\\}\.\(8\)
Concretely, a node projection headfnodef\_\{\\mathrm\{node\}\}first maps each latent token𝒛i\\bm\{z\}\_\{i\}to an initial decoder feature𝒈i\(0\)\\bm\{g\}\_\{i\}^\{\(0\)\}\. The adjacency headfAf\_\{A\}then takes the pair feature\[𝒈i\(0\),𝒈j\(0\)\]\[\\bm\{g\}\_\{i\}^\{\(0\)\},\\bm\{g\}\_\{j\}^\{\(0\)\}\]and predicts the edge\-existence probability:

𝒈i\(0\)=fnode​\(𝒛i\),Pi​jA=σ​\(fA​\(\[𝒈i\(0\),𝒈j\(0\)\]\)\)\.\\bm\{g\}\_\{i\}^\{\(0\)\}=f\_\{\\mathrm\{node\}\}\(\\bm\{z\}\_\{i\}\),\\qquad P^\{A\}\_\{ij\}=\\sigma\\\!\\left\(f\_\{A\}\(\[\\bm\{g\}\_\{i\}^\{\(0\)\},\\bm\{g\}\_\{j\}^\{\(0\)\}\]\)\\right\)\.\(9\)
Because molecular bonds are undirected and self\-bonds are invalid, we symmetrizePAP^\{A\}and set its diagonal to zero; TopoBridge, denoted byΠTB\\Pi\_\{\\mathrm\{TB\}\}, then refines it into a connected adjacency matrix,

𝑨~=ΠTB​\(PA\)\.\\widetilde\{\\bm\{A\}\}=\\Pi\_\{\\mathrm\{TB\}\}\(P^\{A\}\)\.\(10\)
The refined adjacency is injected into TopoFormer as an attention bias:

αi​j=softmaxj⁡\(\(𝒒i\)⊤​𝒌jd\+𝒆A​\(A~i​j\)\),𝒈i←∑jαi​j​𝒗j,\\alpha\_\{ij\}=\\operatorname\{softmax\}\_\{j\}\\left\(\\frac\{\(\\bm\{q\}\_\{i\}\)^\{\\top\}\\bm\{k\}\_\{j\}\}\{\\sqrt\{d\}\}\+\\bm\{e\}\_\{A\}\(\\widetilde\{A\}\_\{ij\}\)\\right\),\\qquad\\bm\{g\}\_\{i\}\\leftarrow\\sum\_\{j\}\\alpha\_\{ij\}\\bm\{v\}\_\{j\},\(11\)where𝒒i,𝒌j,𝒗j\\bm\{q\}\_\{i\},\\bm\{k\}\_\{j\},\\bm\{v\}\_\{j\}are projections of the current decoder features and𝒆A​\(A~i​j\)\\bm\{e\}\_\{A\}\(\\widetilde\{A\}\_\{ij\}\)is a learned adjacency bias\.

Given the atom\-type distributionPiT=softmax⁡\(fT​\(𝒈i\)\)P\_\{i\}^\{T\}=\\operatorname\{softmax\}\(f\_\{T\}\(\\bm\{g\}\_\{i\}\)\), we form an atom\-conditioned feature𝒂iT\\bm\{a\}\_\{i\}^\{T\}fromPiTP\_\{i\}^\{T\}and use it in the bond head\. For each pair\(i,j\)\(i,j\), the bond head predicts logits over real bond types:

𝒐i​jB=fB​\(𝒈i,𝒈j,𝒂iT,𝒂jT,A~i​j\),𝒐i​jB∈ℝK\.\\bm\{o\}^\{B\}\_\{ij\}=f\_\{B\}\(\\bm\{g\}\_\{i\},\\bm\{g\}\_\{j\},\\bm\{a\}\_\{i\}^\{T\},\\bm\{a\}\_\{j\}^\{T\},\\widetilde\{A\}\_\{ij\}\),\\qquad\\bm\{o\}^\{B\}\_\{ij\}\\in\\mathbb\{R\}^\{K\}\.\(12\)The full bond distribution is then obtained by adjacency gating:

Pi​j,0B=1−A~i​j,Pi​j,cB=A~i​jsoftmax\(𝒐i​jB\)c,c=1,…,K\.P^\{B\}\_\{ij,0\}=1\-\\widetilde\{A\}\_\{ij\},\\qquad P^\{B\}\_\{ij,c\}=\\widetilde\{A\}\_\{ij\}\\operatorname\{softmax\}\(\\bm\{o\}^\{B\}\_\{ij\}\)\_\{c\},\\quad c=1,\\ldots,K\.\(13\)Thus bond prediction is both adjacency\-conditioned and atom\-conditioned: non\-adjacent pairs are assigned no bond, while adjacent pairs are classified among real bond types using the predicted atom information\.

For coordinate reconstruction, the coordinate head first predicts an initial coordinate matrix from the final decoder features,

𝑹\(0\)=fR​\(𝒈\)\.\\bm\{R\}^\{\(0\)\}=f\_\{R\}\(\\bm\{g\}\)\.\(14\)The bond distribution is first mapped to a continuous edge feature𝒆i​jB=∑c=0KPi​j,cB​𝒘c\\bm\{e\}\_\{ij\}^\{B\}=\\sum\_\{c=0\}^\{K\}P^\{B\}\_\{ij,c\}\\,\\bm\{w\}\_\{c\}via learned bond embeddings\{𝒘c\}\\\{\\bm\{w\}\_\{c\}\\\}\. The EGNN then takes𝑹\(0\)\\bm\{R\}^\{\(0\)\}, the decoder features𝒈\\bm\{g\}, and𝒆B\\bm\{e\}^\{B\}as inputs, and iteratively updates features and coordinates:

\(𝒈\(ℓ\+1\),𝑹\(ℓ\+1\)\)=EGNNℓ​\(𝒈\(ℓ\),𝑹\(ℓ\),𝒆B\),𝑹^=𝑹\(L\)\.\(\\bm\{g\}^\{\(\\ell\+1\)\},\\bm\{R\}^\{\(\\ell\+1\)\}\)=\\mathrm\{EGNN\}\_\{\\ell\}\(\\bm\{g\}^\{\(\\ell\)\},\\bm\{R\}^\{\(\\ell\)\},\\bm\{e\}^\{B\}\),\\qquad\\widehat\{\\bm\{R\}\}=\\bm\{R\}^\{\(L\)\}\.\(15\)

### 4\.3TopoBridge: connected adjacency refinement

TopoBridge converts the soft adjacency probabilitiesPAP^\{A\}into a binary connected adjacency matrix𝑨~\\widetilde\{\\bm\{A\}\}\. It first builds an initial undirected graph by keeping the top\-kkmost confident neighbours for each atom, then symmetrises and removes self\-loops\.

Disconnected graphs are repaired by greedily adding bridge edges\. Choose a root in the largest component and letSSbe the set of atoms reachable by BFS\. IfS≠VS\\neq V, add the most likely cross\-component edge

\(u⋆,v⋆\)=arg⁡maxu∈S,v∉S⁡Pu​vA,\(u^\{\\star\},v^\{\\star\}\)=\\arg\\max\_\{u\\in S,\\;v\\notin S\}P^\{A\}\_\{uv\},\(16\)and re\-run BFS\. Each insertion strictly grows\|S\|\|S\|, so the procedure terminates and returns𝑨~=ΠTB​\(PA\)\\widetilde\{\\bm\{A\}\}=\\Pi\_\{\\mathrm\{TB\}\}\(P^\{A\}\)with𝒢​\(𝑨~\)\\mathcal\{G\}\(\\widetilde\{\\bm\{A\}\}\)connected\.

BecauseΠTB\\Pi\_\{\\mathrm\{TB\}\}is discrete, TopVAE applies a straight\-through surrogate\(Bengioet al\.,[2013](https://arxiv.org/html/2606.13955#bib.bib47)\)for downstream decoder stages:

𝑨ST=PA\+stopgrad​\(𝑨~−PA\)\.\\bm\{A\}^\{\\mathrm\{ST\}\}=P^\{A\}\+\\mathrm\{stopgrad\}\\\!\\left\(\\widetilde\{\\bm\{A\}\}\-P^\{A\}\\right\)\.\(17\)The forward value is𝑨~\\widetilde\{\\bm\{A\}\}; gradients flow throughPAP^\{A\}\. The adjacency head is supervised from its logits against the ground\-truth adjacency\.

#### Over\-connection risk\.

Because TopoBridge forces edges, atoms connected by a bridge necessarily receive a nonzero bond type, which could bias the decoded graph toward over\-connectedness\. Empirically, however, atom\-degree and ring\-size distributions of TopVAE are comparable to UDM\-3D \(degree EMD:0\.5680\.568vs\.0\.5580\.558; see Appendix[J](https://arxiv.org/html/2606.13955#A10)and Tables[18](https://arxiv.org/html/2606.13955#A10.T18)–[19](https://arxiv.org/html/2606.13955#A10.T19)\)\. This is because bridge edges are selected by descending adjacency logit \(Eq\.[16](https://arxiv.org/html/2606.13955#S4.E16)\), inserting only edges the model already considers most likely, and the multi\-stage decoder conditions downstream bond\-type and coordinate heads on the refined adjacency, integrating forced edges into a globally consistent graph\.

#### Computational overhead\.

TopoBridge’s refinement adds only2\.82\.8ms per batch \(∼2\.8%\{\\sim\}2\.8\\%of forward time on an H100 NVL, batch size 64\) for a converged model\. Scaling is governed by the number of bridge insertions rather than atom countNN\. Full profiling and scaling curves are in Appendix[K](https://arxiv.org/html/2606.13955#A11)\.

### 4\.4ChemCO: Unrolled Chemical Constraint Optimization

TopoBridge guarantees connectivity but does not constrain bond*types*: atom pairs may still receive chemically illegal bonds or violate valence caps\. ChemCO closes this gap by solving a constrained optimization problem over bond\-type assignments, producing a chemically improved distribution that AGCL \(Sec\.[4\.5](https://arxiv.org/html/2606.13955#S4.SS5)\) selectively injects into training\.

Letk=1,…,Kk\\\!=\\\!1,\\ldots,Kindex bond types with bond ordersoko\_\{k\}, andk=0k\\\!=\\\!0denote no bond\. From the decoder’s raw logitsUi​j\(k\)U\_\{ij\}^\{\(k\)\}over allK\+1K\\\!\+\\\!1bond classes \(including no\-bondk=0k\\\!=\\\!0\), ChemCO takes the relative neural utility:

U¯i​j\(k\)=Ui​j\(k\)−Ui​j\(0\),k=1,…,K,\\bar\{U\}\_\{ij\}^\{\(k\)\}=U\_\{ij\}^\{\(k\)\}\-U\_\{ij\}^\{\(0\)\},\\qquad k=1,\\ldots,K,\(18\)as a fixed input encoding the decoder’s preference for each bond type\. ChemCO introduces free primal variablesΦ\(k\)\\Phi^\{\(k\)\}, initialized asΦ0\(k\)=U¯\(k\)\\Phi\_\{0\}^\{\(k\)\}\\\!=\\\!\\bar\{U\}^\{\(k\)\}, and maps them to a continuous assignment score via

Y\(k\)=Tk​\(Φ\):=12​\(g​\(Φ\(k\)\)\+g​\(Φ\(k\)\)⊤\)⊙A~⊙Ω\(k\),g​\(x\)=σ​\(x\)2,Y^\{\(k\)\}=T\_\{k\}\(\\Phi\):=\\tfrac\{1\}\{2\}\\\!\\bigl\(g\(\\Phi^\{\(k\)\}\)\+g\(\\Phi^\{\(k\)\}\)^\{\\top\}\\bigr\)\\odot\\widetilde\{A\}\\odot\\Omega^\{\(k\)\},\\quad g\(x\)=\\sigma\(x\)^\{2\},\(19\)whereA~\\widetilde\{A\}is the TopoBridge adjacency support restricting optimization to candidate edges, andΩ\(k\)\\Omega^\{\(k\)\}is a binary mask that zeros out chemically invalid atom–bond combinations \(e\.g\., a triple bond between two oxygen atoms\)\. During trainingΩ\(k\)\\Omega^\{\(k\)\}is constructed from predicted atom types\. FromYY, ChemCO derives

si​j=∑kYi​j\(k\),vali=∑j≠i∑kok​Yi​j\(k\),degi=∑j≠isi​j,s\_\{ij\}=\\textstyle\\sum\_\{k\}Y\_\{ij\}^\{\(k\)\},\\qquad\\operatorname\{val\}\_\{i\}=\\textstyle\\sum\_\{j\\neq i\}\\sum\_\{k\}o\_\{k\}Y\_\{ij\}^\{\(k\)\},\\qquad\\deg\_\{i\}=\\textstyle\\sum\_\{j\\neq i\}s\_\{ij\},\(20\)subject to three chemical constraints:si​j≤1s\_\{ij\}\\\!\\leq\\\!1\(each atom pair carries at most one bond type\),vali≤ci\\operatorname\{val\}\_\{i\}\\\!\\leq\\\!c\_\{i\}\(total bond order at each atom does not exceed its valence capacity\), anddegi≥dmin\\deg\_\{i\}\\\!\\geq\\\!d\_\{\\min\}\(every valid atom participates in at least one bond\)\.

#### Objective\.

ChemCO maximizes expected bond utility while penalizing chemical constraint violations through an adaptive\-penalty scheme\. Each violation is smoothly approximated byψγ​\(r\)=γ−1​log⁡\(1\+eγ​r\)\\psi\_\{\\gamma\}\(r\)=\\gamma^\{\-1\}\\log\(1\+e^\{\\gamma r\}\), a soft surrogate for\[r\]\+\[r\]\_\{\+\}\. The per\-step score is

𝒥t​\(Φ\)\\displaystyle\\mathcal\{J\}\_\{t\}\(\\Phi\)=∑i<j∑k=1KU¯i​j\(k\)​Yi​j\(k\)−∑i<jμi​j,t​ψγ​\(si​j−1\)\\displaystyle=\\textstyle\\sum\_\{i<j\}\\sum\_\{k=1\}^\{K\}\\bar\{U\}\_\{ij\}^\{\(k\)\}Y\_\{ij\}^\{\(k\)\}\-\\sum\_\{i<j\}\\mu\_\{ij,t\}\\psi\_\{\\gamma\}\(s\_\{ij\}\-1\)\(21\)−∑iλi,t​ψγ​\(vali−ci\)−∑iνi,t​ψγ​\(dmin−degi\),\\displaystyle\\quad\-\\textstyle\\sum\_\{i\}\\lambda\_\{i,t\}\\psi\_\{\\gamma\}\(\\operatorname\{val\}\_\{i\}\-c\_\{i\}\)\-\\sum\_\{i\}\\nu\_\{i,t\}\\psi\_\{\\gamma\}\(d\_\{\\min\}\-\\deg\_\{i\}\),whereY=T​\(Φ\)Y=T\(\\Phi\)\. The three penalty terms enforce bond exclusivity, valence limits, and minimum\-degree requirements respectively\. After each gradient step onΦ\\Phi, multipliers are raised for still\-violated constraints, automatically concentrating pressure where it is most needed\. Full update rules are in Appendix[M](https://arxiv.org/html/2606.13955#A13)\. We also provide analysis of the computational overhead; see Appendix[K](https://arxiv.org/html/2606.13955#A11)\.

### 4\.5Advantage\-Gated Constraint Learning

ChemCO provides a constraint\-optimization signal for learning chemical rules\. However, this signal should remain conservative: if constraint\-induced corrections are imposed indiscriminately, they may conflict with the reconstruction objective\. Following the positive\-advantage principleOhet al\.\([2018](https://arxiv.org/html/2606.13955#bib.bib40)\); Penget al\.\([2019](https://arxiv.org/html/2606.13955#bib.bib41)\), AGCL selectively injects the ChemCO signal only when it improves ground\-truth bond NLL on bonded pairs\.

LetPrawP\_\{\\mathrm\{raw\}\}andPchemP\_\{\\mathrm\{chem\}\}be the bond distributions before and after ChemCO\. For moleculebbwith bonded pair set𝒮b\\mathcal\{S\}\_\{b\}and ground\-truth labelsyi​j⋆y\_\{ij\}^\{\\star\}, the per\-molecule advantage is

ℰb​\(P\)=1\|𝒮b\|​∑\(i,j\)∈𝒮b−log⁡Pi​j,yi​j⋆,ab=\[ℰb​\(Praw\)−ℰb​\(Pchem\)\]\+\.\\mathcal\{E\}\_\{b\}\(P\)=\\frac\{1\}\{\|\\mathcal\{S\}\_\{b\}\|\}\\sum\_\{\(i,j\)\\in\\mathcal\{S\}\_\{b\}\}\-\\log P\_\{ij,y\_\{ij\}^\{\\star\}\},\\qquad a\_\{b\}=\\left\[\\mathcal\{E\}\_\{b\}\(P\_\{\\mathrm\{raw\}\}\)\-\\mathcal\{E\}\_\{b\}\(P\_\{\\mathrm\{chem\}\}\)\\right\]\_\{\+\}\.\(22\)Thusab\>0a\_\{b\}\\\!\>\\\!0only when ChemCO lowers the ground\-truth bond NLL\. AGCL uses this advantage to weight a consistency loss:

ℒAGCL=∑bab​∑\(i,j\)∈𝒮b‖Praw,i​j−sg⁡\(Pchem,i​j\)‖22\(K\+1\)​∑bab​\|𝒮b\|\+ϵ,\\mathcal\{L\}\_\{\\mathrm\{AGCL\}\}=\\frac\{\\sum\_\{b\}a\_\{b\}\\sum\_\{\(i,j\)\\in\\mathcal\{S\}\_\{b\}\}\\left\\\|P\_\{\\mathrm\{raw\},ij\}\-\\operatorname\{sg\}\(P\_\{\\mathrm\{chem\},ij\}\)\\right\\\|\_\{2\}^\{2\}\}\{\(K\+1\)\\sum\_\{b\}a\_\{b\}\|\\mathcal\{S\}\_\{b\}\|\+\\epsilon\},\(23\)wheresg⁡\(⋅\)\\operatorname\{sg\}\(\\cdot\)detaches the ChemCO output so thatℒAGCL\\mathcal\{L\}\_\{\\mathrm\{AGCL\}\}steers only the raw decoder without back\-propagating through the unrolled solver\. As training progresses, the raw decoder internalizes the chemical constraints: the fraction of molecules with positive advantageab\>0a\_\{b\}\\\!\>\\\!0approaches zero \(Figure[2](https://arxiv.org/html/2606.13955#S4.F2)\) and ChemCO can be removed at inference time \(Table[5](https://arxiv.org/html/2606.13955#S5.T5)\)\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x1.png)Figure 2:ChemCO and AGCL dynamics\.\(a\) Averaged ChemCO optimization curves: red reports over\-valent atoms, and blue reports mean pairwise bond\-type exclusivity excess\. \(b\) AGCL training curves: red reports the fraction of samples with positive advantage \(ab\>0a\_\{b\}\\\!\>\\\!0\), and blue reports molecule validity\.
### 4\.6Training Objective

TopVAE is trained end\-to\-end with reconstruction losses for each decoded modality, KL regularization, and the AGCL selective teacher:

ℒ=ℒtopo\+ℒatom\+λ3​D​ℒ3​D\+β​ℒKL\+λAGCL​ℒAGCL\.\\mathcal\{L\}=\\mathcal\{L\}\_\{\\mathrm\{topo\}\}\+\\mathcal\{L\}\_\{\\mathrm\{atom\}\}\+\\lambda\_\{3\\mathrm\{D\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{3D\}\}\+\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{KL\}\}\+\\lambda\_\{\\mathrm\{AGCL\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{AGCL\}\}\.\(24\)Hereℒtopo\\mathcal\{L\}\_\{\\mathrm\{topo\}\}groups focal binary cross\-entropy for the adjacency matrix and cross\-entropy for bond types,ℒatom\\mathcal\{L\}\_\{\\mathrm\{atom\}\}is cross\-entropy over atom types,ℒ3​D\\mathcal\{L\}\_\{\\mathrm\{3D\}\}combines coordinate regression and pairwise distance supervision, andℒKL\\mathcal\{L\}\_\{\\mathrm\{KL\}\}is the standard VAE KL divergence\. All individual loss terms, weights, architectural details, ChemCO hyperparameters, and warmup schedules are reported in Appendix[L](https://arxiv.org/html/2606.13955#A12)\.

## 5Experiments

### 5\.1Setup

Datasets\.QM9\(Ramakrishnanet al\.,[2014](https://arxiv.org/html/2606.13955#bib.bib37)\)contains∼\{\\sim\}134k small organic molecules with up to 9 heavy atoms \(C, N, O, F\) and their equilibrium geometries computed at the DFT level\. GEOM\-Drugs\(Axelrod and Gomez\-Bombarelli,[2022](https://arxiv.org/html/2606.13955#bib.bib38)\)contains∼\{\\sim\}304k drug\-like molecules with up to 90 heavy atoms and multiple conformers per molecule\.

Metrics\.Atom stability and molecular stability \(MolStab\) follow the definitions in EDM\(Hoogeboomet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib19)\): an atom is stable if its valence equals the reference value, and a molecule is stable if all its atoms are stable \(not necessarily connected\)\. Validity requires successful RDKit sanitization\. Connectivity requires the molecular graph to be a single connected component\. We report joint metricsV&C\(validity∧\\wedgeconnectivity\),V&U\(validity∧\\wedgeuniqueness\), andS∧\\wedgeC\(MolStable∧\\wedgeconnectivity\)\.FCDmeasures distributional similarity to the reference set in chemical descriptor space\(Preueret al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib39)\)\.iFIDadapts the interpolated FID to molecular VAE latent spaces\(Xuet al\.,[2026](https://arxiv.org/html/2606.13955#bib.bib25)\), decoding SLERP midpoints between posterior means and measuring the Fréchet distance of ChemNet activations \(details in Appendix[L\.7](https://arxiv.org/html/2606.13955#A12.SS7)\)\.

### 5\.2VAE Dark\-Area Diagnosis

To quantify dark areas, we evaluate each VAE decoder under latent perturbations of increasing magnitude\. For each validation moleculeMM, we encode its posterior mean𝒛=𝝁ϕ​\(M\)\\bm\{z\}=\\bm\{\\mu\}\_\{\\phi\}\(M\)and decode the perturbed latent

𝒛~=𝒛\+σ⋅semp⋅ϵ,ϵ∼𝒩​\(𝟎,𝑰\),\\tilde\{\\bm\{z\}\}=\\bm\{z\}\+\\sigma\\cdot s\_\{\\mathrm\{emp\}\}\\cdot\\bm\{\\epsilon\},\\qquad\\bm\{\\epsilon\}\\sim\\mathcal\{N\}\(\\bm\{0\},\\bm\{I\}\),\(25\)wheresemp=std​\(\{𝝁ϕ​\(M\)\}M∈𝒟\)s\_\{\\mathrm\{emp\}\}=\\mathrm\{std\}\(\\\{\\bm\{\\mu\}\_\{\\phi\}\(M\)\\\}\_\{M\\in\\mathcal\{D\}\}\)is the empirical standard deviation of the encoded latents andσ\\sigmacontrols the perturbation strength relative to the latent spread\.

Table[1](https://arxiv.org/html/2606.13955#S5.T1)reports results atσ=0\\sigma\{=\}0and0\.50\.5\. Existing VAEs collapse under perturbation: S∧\\wedgeC drops by over an order of magnitude on QM9 and falls to zero on GEOM\-Drugs, whereas TopVAE achieves higher S∧\\wedgeC\. PCA projection \(Figure[3](https://arxiv.org/html/2606.13955#A2.F3)\) confirms the gap—TopVAE decodes90\.5%90\.5\\%of perturbed latents as valid connected molecules versus44\.4%44\.4\\%\(UAE\) and32\.6%32\.6\\%\(ADiT\)\. This robustness extends to the diffusion trajectory: TopVAE\+DiT reaches\>\>90% validity by denoising progress 0\.25 while UDM\-3D does not saturate until progress∼\{\\sim\}0\.90 \(Figures[4](https://arxiv.org/html/2606.13955#A2.F4)–[5](https://arxiv.org/html/2606.13955#A2.F5), Appendix[B](https://arxiv.org/html/2606.13955#A2)\)\.

Table 1:VAE dark\-area diagnosison QM9 and GEOM\-Drugs \(n=10​kn\{=\}10\\text\{k\}\)\. TopVAE maintains substantially higher decodability under noise perturbation\.Post\.MolStab↑\\,\\uparrowPost\.S∧\\wedgeC↑\\,\\uparrowσ=0\.5\\sigma\{=\}0\.5S∧\\wedgeC↑\\,\\uparrowiFID↓\\downarrow*QM9*UAE1\.0000\.9980\.0232\.88ADiT\-VAE0\.9510\.9510\.0007\.08TopVAE1\.0001\.0000\.6620\.395*GEOM\-Drugs*UAE0\.9980\.0000\.00062\.00TopVAE0\.9600\.9600\.03342\.77
### 5\.3Latent Diffusion Generation on QM9 and GEOM\-Drugs

#### QM9\.

Table[2](https://arxiv.org/html/2606.13955#S5.T2)compares de novo generation on QM9\. TopVAE\+DiT achieves perfect V&C and the best FCD3D\(77%77\\%lower over previous SOTA\), indicating that dark\-area closure translates directly into higher\-quality 3D generation\. QM9 is near\-saturated for MolStab—all competitive methods exceed0\.970\.97—so the principal gains are in topology\-consistent generation and 3D distributional fidelity\.

#### GEOM\-Drugs\.

Table[3](https://arxiv.org/html/2606.13955#S5.T3)reports results on the larger GEOM\-Drugs benchmark, where dark areas are more severe\. TopVAE\+DiT achieves the highest V&C \(0\.9510\.951,\+8%\+8\\%\), V&U \(1\.0001\.000\), and FCD3D\(8\.258\.25,−52%\-52\\%\), while remaining competitive on 2D distributional metrics\.

#### Summary\.

The V&C improvements on both benchmarks are direct consequences of TopoBridge \(guaranteeing connectivity\) and ChemCO \(enforcing chemical validity\), which together mitigate the dark areas diagnosed in Section[5\.2](https://arxiv.org/html/2606.13955#S5.SS2)\. The FCD3Dgains further suggest that a smoother, more navigable latent space enables the diffusion prior to produce latent codes that decode into higher\-fidelity 3D geometries\. We provide additional analysis of 3D conformations, chemical properties, and diversity distributions \(Appendices[C](https://arxiv.org/html/2606.13955#A3)–[F](https://arxiv.org/html/2606.13955#A6)\)\.

Table 2:QM9 de novo 3D molecule generation\(n=10,000n\{=\}10\{,\}000\)\. Baseline results are from respective papers; TopVAE\+DiT is our evaluation\.2D Metrics3D MetricsMethodFCD↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowV&C↑\\uparrowV&U↑\\uparrowFCD3D↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowEDM1\.2850\.9860\.8170\.9340\.9071\.2850\.9860\.817CDGS0\.7980\.9970\.9510\.9360\.860–––GeoLDM1\.0300\.9890\.8940\.9510\.9271\.0300\.9890\.897LDM\-3DG0\.5590\.9760\.8691\.0000\.953–––MolFLAE–0\.9940\.920–0\.889–––GFMDiff–0\.9890\.877–0\.951–0\.9890\.877MiDi0\.1870\.9980\.9760\.9800\.9541\.1000\.9830\.842JODO0\.1380\.9990\.9880\.9900\.9600\.8850\.9920\.934EQGAT\-diff2\.0880\.9990\.9710\.9650\.9501\.5200\.9880\.888SemlaFlow0\.8630\.9950\.9490\.8570\.8211\.1270\.9710\.787UDM\-3D0\.1300\.9990\.9880\.9830\.9730\.8810\.9930\.935TopVAE\+DiT0\.1850\.9961\.0001\.0000\.9590\.207↓\\downarrow77%0\.9870\.925Table 3:GEOM\-Drugs de novo 3D molecule generation\(n=10,000n\{=\}10\{,\}000\)\. Baseline results are from respective papers; TopVAE\+DiT is our evaluation\.2D Metrics3D MetricsMethodFCD↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowV&C↑\\uparrowV&U↑\\uparrowFCD3D↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowEDM40\.140\.9910\.9140\.3590\.99131\.290\.8310\.002CDGS22\.050\.9910\.7060\.2850\.285–––GeoLDM39\.810\.9960\.9090\.4820\.99830\.680\.8430\.008MiDi7\.0540\.9680\.8220\.6330\.65423\.140\.7500\.003JODO2\.5231\.0000\.9810\.8740\.90219\.990\.8450\.010EQGAT\-diff5\.8981\.0000\.9890\.8450\.86326\.330\.8250\.007UDM\-3D0\.6921\.0000\.9250\.8790\.90717\.360\.8520\.014TopVAE\+DiT2\.681\.0000\.9810\.951↑\\uparrow8%1\.000↑\\uparrow10%8\.25↓\\downarrow52%0\.8320\.020↑\\uparrow43%

### 5\.4Zero\-Shot Scaffold Inpainting

Scaffold inpainting stress\-tests dark areas under OOD conditions: given a scaffold withkkatoms, we encode it intokklatent tokens and fill the remainingn−kn\{\-\}kslots \(heavy atoms\) via reverse diffusion\. As the expansion size grows, a larger fraction of the latent code is sampled off\-posterior, increasing exposure to𝒟dark\\mathcal\{D\}\_\{\\mathrm\{dark\}\}\.

Table[4](https://arxiv.org/html/2606.13955#S5.T4)reports results on GEOM\-Drugs across five scaffolds \(benzene, pyridine, naphthalene, indole, cyclohexane; 1k samples each; per\-scaffold breakdown in Appendix[I](https://arxiv.org/html/2606.13955#A9)\)\. TopVAE\+DiT consistently outperforms UDM\-3D, with the gap widening at larger expansions: at\+60\+60atoms, UAE collapses to5\.6%5\.6\\%S∧\\wedgeC while TopVAE maintains51\.4%51\.4\\%; at\+80\+80and beyond, UAE produces near\-zero valid molecules while TopVAE continues to generate scaffold\-preserving structures\.

Table 4:Zero\-shot scaffold inpainting on GEOM\-Drugs\(mean over scaffolds\)\. Per\-scaffold breakdown in Appendix Table[17](https://arxiv.org/html/2606.13955#A9.T17)\.Scaff\. Pres\.↑\\uparrowS∧\\wedgeC↑\\uparrowScaff\. Pres\.↑\\uparrowS∧\\wedgeC↑\\uparrow\+n\+nUAETopVAEUAETopVAE\+n\+nUAETopVAEUAETopVAE510\.011\.389\.895\.9605\.028\.95\.651\.41014\.815\.486\.191\.1800\.025\.20\.126\.42015\.215\.986\.984\.81000\.019\.80\.08\.64017\.927\.863\.969\.6Mean9\.020\.647\.561\.1
### 5\.5Ablation

#### Component ablation\.

Table[5](https://arxiv.org/html/2606.13955#S5.T5)isolates the contribution of each module\. Without ChemCO,σ=0\.5\\sigma\{=\}0\.5Stab∧\\wedgeConn drops from0\.6620\.662to0\.020\.02: the decoder loses chemical validity off\-manifold\. Without TopoBridge, posterior MolStab remains high \(0\.9510\.951\) but Stab∧\\wedgeConn falls to0\.3020\.302: molecules pass valence checks yet fragment into disconnected components\.

#### Inference\-time ablation\.

To verify that the decoder has learned chemical rules rather than relying on ChemCO as a post\-processing step, Table[5](https://arxiv.org/html/2606.13955#S5.T5)varies the number of ChemCO iterations at inference from0\(fully disabled\) to100100\. All metrics remain stable, confirming that AGCL has transferred the constraint knowledge into the decoder weights during training\.

#### Hyperparameter sensitivity\.

Appendix[G](https://arxiv.org/html/2606.13955#A7)reports sensitivity analysis for TopoBridge’s adjacency sparsity, ChemCO’s minimum\-degree penalty \(dmin∈\{0,1\}d\_\{\\min\}\\in\\\{0,1\\\}\), and theΩ\(k\)\\Omega^\{\(k\)\}chemical mask\. TopVAE is robust to adjacency sparsity and hyperparameters\.

Table 5:Ablation studies\.\(a\)Component ablation on QM9 \(n=10​kn\{=\}10k\)\.\(b\)Inference\-time ablation: varying ChemCO iterations at inference on GEOM\-Drugs \(n=10​kn\{=\}10k\)\. The decoder produces similar results with or without ChemCO\.\(a\) Component ablation Posterior \(σ=0\\sigma\{=\}0\)Noiseσ=0\.5\\sigma\{=\}0\.5ConfigurationMolStabConn\.S∧\\wedgeCΔ\\DeltaS∧\\wedgeCMolStabS∧\\wedgeCFull TopVAE1\.0001\.0001\.000–0\.6820\.662w/o ChemCO0\.9411\.0000\.941–5\.9%0\.4100\.020w/o TopoBridge0\.9510\.3180\.302–69\.8%0\.4050\.010

\(b\) Inference\-time ablation ChemCOStepsMolStab\.↑\\uparrowS∧\\wedgeC↑\\uparrowConn\.↑\\uparrowFCD↓\\downarrowFCD3D↓\\downarrow00\.9810\.9330\.9512\.688\.25100\.9800\.9310\.9512\.688\.251000\.9790\.9280\.9482\.688\.24

## 6Conclusion

We identified*dark areas*in molecular latent spaces, i\.e\., regions where existing VAE decoders produce chemically invalid or disconnected structures despite near\-perfect posterior reconstruction, and showed that per\-atom stability metrics mask this failure by accepting fragmented atom clouds as stable\. TopVAE closes dark areas through a selective\-teacher paradigm: TopoBridge guarantees connectivity, ChemCO formulates chemical rules as differentiable constrained optimization, and AGCL gates this signal by positive advantage so the decoder progressively internalizes the constraints and ChemCO can be removed at inference\. The resulting smoother latent manifold yields77%77\\%lower FCD3Don QM9, the highest V&C and52%52\\%lower FCD3Don GEOM\-Drugs, and1\.29×1\.29\{\\times\}higher molecular stability on zero\-shot scaffold inpainting, demonstrating that dark\-area closure translates directly into generation quality across both in\-distribution and out\-of\-distribution regimes\.

## 7Limitations

The ChemCO–AGCL framework currently enforces only graph\-level valence rules\. In principle, any differentiable constraint admitting a primal–dual formulation can serve as a selective training signal: 3D physical priors such as van der Waals clash penalties, force\-field energy bounds, and steric strain limits can be unrolled in the same manner, enabling decoders to internalize geometric feasibility alongside chemical validity and produce physically plausible molecules without post\-hoc relaxation; incorporating such priors remains future work\. Likewise, extending TopVAE to conditional generation targeting quantum properties, bioactivity profiles, or protein\-pocket constraints is a natural next step toward latent molecular design\.

## Acknowledgments and Disclosure of Funding

Funding disclosure will be added in the camera\-ready version\.

## References

- Optnet: differentiable optimization as a layer in neural networks\.InInternational conference on machine learning,pp\. 136–145\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p2.1)\.
- S\. Axelrod and R\. Gomez\-Bombarelli \(2022\)GEOM, energy\-annotated molecular conformations for property prediction and molecular generation\.Scientific data9\(1\),pp\. 185\.Cited by:[§5\.1](https://arxiv.org/html/2606.13955#S5.SS1.p1.2)\.
- Y\. Bengio, N\. Léonard, and A\. Courville \(2013\)Estimating or propagating gradients through stochastic neurons for conditional computation\.arXiv preprint arXiv:1308\.3432\.Cited by:[§4\.3](https://arxiv.org/html/2606.13955#S4.SS3.p3.1)\.
- Z\. Chen, Y\. Jia, Z\. Tian, W\. Ma, and Y\. Lan \(2025\)Manipulating 3d molecules in a fixed\-dimensional e \(3\)\-equivariant latent space\.arXiv preprint arXiv:2506\.00771\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p2.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1)\.
- H\. Dai, Y\. Tian, B\. Dai, S\. Skiena, and L\. Song \(2018\)Syntax\-directed variational autoencoder for structured data\.arXiv preprint arXiv:1802\.08786\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1)\.
- C\. Diao and R\. Loynd \(2022\)Relational attention: generalizing transformers for graph\-structured tasks\.arXiv preprint arXiv:2210\.05062\.Cited by:[§4\.1](https://arxiv.org/html/2606.13955#S4.SS1.p1.5)\.
- R\. Gómez\-Bombarelli, J\. N\. Wei, D\. Duvenaud, J\. M\. Hernández\-Lobato, B\. Sánchez\-Lengeling, D\. Sheberla, J\. Aguilera\-Iparraguirre, T\. D\. Hirzel, R\. P\. Adams, and A\. Aspuru\-Guzik \(2018\)Automatic chemical design using a data\-driven continuous representation of molecules\.ACS central science4\(2\),pp\. 268–276\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1)\.
- M\. D\. Hoffman, C\. Riquelme, and M\. J\. Johnson \(2017\)Theβ\\beta\-vae’s implicit prior\.InWorkshop on Bayesian Deep Learning, NIPS,pp\. 1–5\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p1.1)\.
- E\. Hoogeboom, V\. G\. Satorras, C\. Vignac, and M\. Welling \(2022\)Equivariant diffusion for molecule generation in 3d\.InInternational conference on machine learning,pp\. 8867–8887\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2606.13955#S5.SS1.p2.4)\.
- C\. Hua, S\. Luan, M\. Xu, Z\. Ying, J\. Fu, S\. Ermon, and D\. Precup \(2024\)Mudiff: unified diffusion for complete molecule generation\.InLearning on Graphs Conference,pp\. 33–1\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1)\.
- H\. Huang, L\. Sun, B\. Du, and W\. Lv \(2023\)Learning joint 2d & 3d diffusion models for complete molecule generation\.arXiv preprint arXiv:2305\.12347\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- W\. Jin, R\. Barzilay, and T\. Jaakkola \(2018\)Junction tree variational autoencoder for molecular graph generation\.InInternational conference on machine learning,pp\. 2323–2332\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1),[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p3.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- C\. K\. Joshi, X\. Fu, Y\. Liao, V\. Gharakhanyan, B\. K\. Miller, A\. Sriram, and Z\. W\. Ulissi \(2025\)All\-atom diffusion transformers: unified generative modelling of molecules and materials\.arXiv preprint arXiv:2503\.03965\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- M\. Krenn, Q\. Ai, S\. Barthel, N\. Carson, A\. Frei, N\. C\. Frey, P\. Friederich, T\. Gaudin, A\. A\. Gayle, K\. M\. Jablonka,et al\.\(2022\)SELFIES and the future of molecular string representations\.Patterns3\(10\)\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1),[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p3.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- M\. J\. Kusner, B\. Paige, and J\. M\. Hernández\-Lobato \(2017\)Grammar variational autoencoder\.InInternational conference on machine learning,pp\. 1945–1954\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1),[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p1.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- T\. Le, J\. Cremer, F\. Noe, D\. Clevert, and K\. Schütt \(2023\)Navigating the design space of equivariant diffusion\-based generative models for de novo 3d molecule generation\.arXiv preprint arXiv:2309\.17296\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- H\. Lee, M\. Kim, S\. Jang, J\. Jeong, and S\. J\. Hwang \(2025\)Enhancing variational autoencoders with smooth robust latent encoding\.arXiv preprint arXiv:2504\.17219\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1)\.
- R\. Li, X\. Peng, C\. Lin, W\. Rong, and Z\. Chen \(2021\)On the low\-density latent regions of vae\-based language models\.InNeurIPS 2020 Workshop on Pre\-registration in Machine Learning,pp\. 343–357\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p1.1)\.
- Q\. Liu, M\. Allamanis, M\. Brockschmidt, and A\. Gaunt \(2018\)Constrained graph variational autoencoders for molecule design\.Advances in neural information processing systems31\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p3.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- Y\. Luo, Z\. Liu, Y\. Zhao, S\. Li, H\. Cai, K\. Kawaguchi, T\. Chua, Y\. Zhang, and X\. Wang \(2025\)Towards unified and lossless latent space for 3d molecular latent diffusion modeling\.arXiv preprint arXiv:2503\.15567\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p2.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- T\. Ma, J\. Chen, and C\. Xiao \(2018\)Constrained generation of semantically valid graphs via regularizing variational autoencoders\.Advances in neural information processing systems31\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p3.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- M\. Madeira, C\. Vignac, D\. Thanou, and P\. Frossard \(2024\)Generative modelling of structurally constrained graphs\.Advances in Neural Information Processing Systems37,pp\. 137218–137262\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p3.1)\.
- T\. Ochiai, T\. Inukai, M\. Akiyama, K\. Furui, M\. Ohue, N\. Matsumori, S\. Inuki, M\. Uesugi, T\. Sunazuka, K\. Kikuchi,et al\.\(2023\)Variational autoencoder\-based chemical latent space for large molecular structures with 3d complexity\.Communications Chemistry6\(1\),pp\. 249\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px3.p2.1)\.
- J\. Oh, Y\. Guo, S\. Singh, and H\. Lee \(2018\)Self\-imitation learning\.InInternational conference on machine learning,pp\. 3878–3887\.Cited by:[§1](https://arxiv.org/html/2606.13955#S1.p3.1),[§4\.5](https://arxiv.org/html/2606.13955#S4.SS5.p1.1)\.
- W\. Peebles and S\. Xie \(2023\)Scalable diffusion models with transformers\.InProceedings of the IEEE/CVF international conference on computer vision,pp\. 4195–4205\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1)\.
- X\. Peng, J\. Guan, Q\. Liu, and J\. Ma \(2023\)Moldiff: addressing the atom\-bond inconsistency problem in 3d molecule diffusion generation\.arXiv preprint arXiv:2305\.07508\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p1.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- X\. B\. Peng, A\. Kumar, G\. Zhang, and S\. Levine \(2019\)Advantage\-weighted regression: simple and scalable off\-policy reinforcement learning\.arXiv preprint arXiv:1910\.00177\.Cited by:[§4\.5](https://arxiv.org/html/2606.13955#S4.SS5.p1.1)\.
- K\. Preuer, P\. Renz, T\. Unterthiner, S\. Hochreiter, and G\. Klambauer \(2018\)Fréchet chemnet distance: a metric for generative models for molecules in drug discovery\.Journal of chemical information and modeling58\(9\),pp\. 1736–1741\.Cited by:[§5\.1](https://arxiv.org/html/2606.13955#S5.SS1.p2.4)\.
- R\. Ramakrishnan, P\. O\. Dral, M\. Rupp, and O\. A\. Von Lilienfeld \(2014\)Quantum chemistry structures and properties of 134 kilo molecules\.Scientific data1\(1\),pp\. 1–7\.Cited by:[§5\.1](https://arxiv.org/html/2606.13955#S5.SS1.p1.2)\.
- D\. Reidenbach, F\. Nikitin, O\. Isayev, and S\. G\. Paliwal \(2026\)Applications of modular co\-design for de novo 3d molecule generation\.Digital Discovery5\(2\),pp\. 754–768\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- R\. Rombach, A\. Blattmann, D\. Lorenz, P\. Esser, and B\. Ommer \(2022\)High\-resolution image synthesis with latent diffusion models\.InProceedings of the IEEE/CVF conference on computer vision and pattern recognition,pp\. 10684–10695\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1)\.
- I\. Skorokhodov, S\. Girish, B\. Hu, W\. Menapace, Y\. Li, R\. Abdal, S\. Tulyakov, and A\. Siarohin \(2025\)Improving the diffusability of autoencoders\.arXiv preprint arXiv:2502\.14831\.Cited by:[§1](https://arxiv.org/html/2606.13955#S1.p2.1)\.
- V\. V\. Vazirani \(2001\)Approximation algorithms\.Vol\.1,Springer\.Cited by:[§M\.8](https://arxiv.org/html/2606.13955#A13.SS8.p1.2)\.
- C\. Vignac, N\. Osman, L\. Toni, and P\. Frossard \(2023\)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:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- R\. Wang, Y\. Zhang, Z\. Guo, T\. Chen, X\. Yang, and J\. Yan \(2023\)Linsatnet: the positive linear satisfiability neural networks\.InInternational Conference on Machine Learning,pp\. 36605–36625\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.
- C\. Xu, H\. Wang, W\. Wang, P\. Zheng, and H\. Chen \(2024\)Geometric\-facilitated denoising diffusion model for 3d molecule generation\.InProceedings of the AAAI Conference on Artificial Intelligence,Vol\.38,pp\. 338–346\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px1.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- M\. Xu, A\. S\. Powers, R\. O\. Dror, S\. Ermon, and J\. Leskovec \(2023\)Geometric latent diffusion models for 3d molecule generation\.InInternational Conference on Machine Learning,pp\. 38592–38610\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p2.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- T\. Xu, M\. He, S\. Abu\-Hussein, J\. M\. Hernandez\-Lobato, H\. Zhang, K\. Zhao, C\. Zhou, Y\. Zhang, and Y\. Wang \(2026\)Making reconstruction fid predictive of diffusion generation fid\.arXiv preprint arXiv:2603\.05630\.Cited by:[§L\.7](https://arxiv.org/html/2606.13955#A12.SS7.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2606.13955#S1.p2.1),[§5\.1](https://arxiv.org/html/2606.13955#S5.SS1.p2.4)\.
- J\. Yao, B\. Yang, and X\. Wang \(2025\)Reconstruction vs\. generation: taming optimization dilemma in latent diffusion models\.InProceedings of the Computer Vision and Pattern Recognition Conference,pp\. 15703–15712\.Cited by:[§1](https://arxiv.org/html/2606.13955#S1.p2.1)\.
- Y\. You, R\. Zhou, J\. Park, H\. Xu, C\. Tian, Z\. Wang, and Y\. Shen \(2024\)Latent 3d graph diffusion\.InThe Twelfth International Conference on Learning Representations,Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px2.p2.1),[§1](https://arxiv.org/html/2606.13955#S1.p1.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px1.p1.1)\.
- H\. Zeng, C\. Yang, Y\. Zhou, C\. Yang, and Q\. Guo \(2024\)GLinSAT: the general linear satisfiability neural network layer by accelerated gradient descent\.Advances in Neural Information Processing Systems37,pp\. 122584–122615\.Cited by:[Appendix A](https://arxiv.org/html/2606.13955#A1.SS0.SSS0.Px4.p2.1),[§2](https://arxiv.org/html/2606.13955#S2.SS0.SSS0.Px2.p1.1)\.

## Appendix AExtended Related Work

This appendix expands on the related work discussion in Sec\.[2](https://arxiv.org/html/2606.13955#S2), providing detailed comparisons with prior methods and clarifying where TopVAE fits relative to data\-space molecular generators, molecular latent diffusion models, VAE latent\-space analyses, and constraint\-aware decoding methods\.

#### Data\-space 3D molecular generation\.

Early 3D molecular diffusion models generate molecules directly in data space\. EDM introduced E\(3\)\-equivariant diffusion over atom coordinates and categorical atom features\[Hoogeboomet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib19)\]\. Because EDM does not generate bond variables explicitly, molecular bonds are typically inferred after generation through distance\-based or chemistry\-based post\-processing\. This post\-hoc bond recovery can create graph–geometry mismatch: the generated atom types and coordinates may not admit a chemically consistent bond graph\. MolDiff formalized this issue as atom–bond inconsistency and proposed to generate atom and bond information jointly\[Penget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib9)\]\.

Subsequent data\-space methods make molecular topology, bond variables, or bond\-formation information part of the generative process\. JODO jointly models atom types, formal charges, bond information, and 3D coordinates through a joint 2D–3D diffusion formulation\[Huanget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib10)\]\. MUDiff combines discrete graph diffusion with continuous coordinate diffusion for molecular graphs and conformations\[Huaet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib11)\]\. MiDi performs denoising over both molecular graphs and 3D atom arrangements\[Vignacet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib12)\]\. EQGAT\-diff studies equivariant diffusion with mixed categorical and continuous variables, including atom, bond, and coordinate channels\[Leet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib13)\]\. Recent scalable or graph\-aware molecular generators further improve continuous–discrete denoising architectures and bond\-formation\-aware training objectives\[Reidenbachet al\.,[2026](https://arxiv.org/html/2606.13955#bib.bib27), Xuet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib28)\]\.

These data\-space methods reduce the mismatch between generated geometry and molecular graph structure, but they run the generative dynamics over the full molecular representation\. They therefore do not isolate the molecular decoder as an object of study\. In contrast, TopVAE focuses on the decoder used by a latent generative pipeline: given a compact continuous latent code, the decoder must produce a connected and chemically valid molecule not only on posterior latents, but also on interpolated, prior\-like, and diffusion\-perturbed latents\.

#### Molecular latent diffusion and autoencoding\.

Latent diffusion separates representation learning from prior learning: a encoder first maps data into a lower\-dimensional latent space, and a diffusion model is then trained in that latent space\[Rombachet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib1)\]\. Transformer\-based diffusion priors such as DiT further improve the scalability of latent\-space generation\[Peebles and Xie,[2023](https://arxiv.org/html/2606.13955#bib.bib2)\]\. This paradigm is especially attractive for 3D molecules, where the raw data contain mixed discrete–continuous variables, permutation structure, chemical constraints, and geometric symmetries\.

Several recent methods adapt latent diffusion or latent autoencoding to 3D molecular generation\. GeoLDM constructs a point\-structured molecular latent space with invariant scalar and equivariant vector components, and trains a latent diffusion model over this representation\[Xuet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib14)\]\. Latent 3D Graph Diffusion provides theoretical motivation for molecular latent diffusion, emphasizing that useful latents should be low\-dimensional, reconstructive, and symmetry\-preserving\[Youet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib15)\]\. UAE\-3D compresses atom types, bonds, and coordinates into a unified latent sequence and applies a standard DiT prior without molecule\-specific inductive bias in the latent generator\[Luoet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib16)\]\. ADiT extends latent diffusion to a shared framework for all\-atom molecules and periodic materials\[Joshiet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib17)\]\. Related latent autoencoding work such as MolFLAE learns fixed\-dimensional E\(3\)\-equivariant molecular latents for zero\-shot molecular manipulation, without being limited to the standard “autoencoder plus diffusion prior” pipeline\[Chenet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib18)\]\.

These works primarily address how to construct expressive, compact, and symmetry\-aware molecular latent spaces\. TopVAE addresses a complementary failure mode: even when reconstruction quality is high, a decoder can fail on the latent regions actually queried by a diffusion prior\. Our dark\-area diagnostic evaluates whether prior samples, interpolation paths, and diffusion\-perturbed latents remain chemically decodable\. This shifts attention from reconstruction\-only autoencoding to the operational robustness required for latent molecular generation\.

#### Latent\-space quality in VAEs\.

A VAE decoder is trained on samples from the encoder posterior, but downstream latent generation may query regions closer to the aggregated posterior, interpolation paths between encoded molecules, or trajectories produced by a learned latent prior\. This creates a gap between reconstruction\-time decoding and generation\-time decoding\. Prior analyses relate this gap to aggregate\-posterior mismatch and to the hypothesized, empirically debated phenomenon of low\-density latent “holes”\[Hoffmanet al\.,[2017](https://arxiv.org/html/2606.13955#bib.bib20), Liet al\.,[2021](https://arxiv.org/html/2606.13955#bib.bib21)\]\. For molecular generation, such holes are especially consequential: an off\-posterior latent may decode not merely to a low\-quality sample, but to a disconnected graph or a molecule that violates chemical valence and sanitization constraints\.

Earlier molecular VAEs already showed that continuous molecular latent spaces are useful for optimization, while also exposing the importance of validity and representation bias in molecular decoding\. The original chemical VAE maps SMILES strings into a continuous latent space for molecular design\[Gómez\-Bombarelliet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib34)\]\. Grammar VAE uses grammar constraints to improve syntactic validity\[Kusneret al\.,[2017](https://arxiv.org/html/2606.13955#bib.bib35)\], and Syntax\-Directed VAE further incorporates syntax and semantic constraints into the decoder\[Daiet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib36)\]\. JT\-VAE generates a junction tree over chemical substructures before assembling the molecular graph, improving chemical validity through a substructure\-level generative process\[Jinet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib29)\]\. SELFIES provides a semantically robust molecular string representation designed to improve molecular validity at the representation level\[Krennet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib32)\]\. NP\-VAE constructs latent spaces for large natural products and complex molecular structures\[Ochiaiet al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib33)\]\. Beyond molecules, SRL\-VAE shows that adversarial smoothing can improve latent robustness in image VAEs\[Leeet al\.,[2025](https://arxiv.org/html/2606.13955#bib.bib26)\]\.

These methods motivate the importance of smooth and valid latent decoding, but they do not directly test a 3D molecular VAE under the latent distribution induced by modern latent diffusion\. TopVAE makes this test explicit through dark\-area analysis: we evaluate whether the decoder remains valid under posterior perturbations, interpolation, prior sampling, and diffusion sampling, and we identify topological disconnection and chemical invalidity as two major failure modes\.

#### Constraint\-aware graph and molecule decoding\.

Validity\-aware molecular generation has been addressed at multiple levels\. At the representation level, grammar\-based methods and SELFIES restrict the output language so that decoded strings better respect syntactic or semantic molecular validity\[Kusneret al\.,[2017](https://arxiv.org/html/2606.13955#bib.bib35), Krennet al\.,[2022](https://arxiv.org/html/2606.13955#bib.bib32)\]\. At the substructure or graph\-decoder level, JT\-VAE and constrained graph VAEs impose chemical validity through the generative procedure itself: JT\-VAE builds molecules from chemical substructure trees\[Jinet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib29)\], while CGVAE uses constrained graph extension to guide molecular graph generation\[Liuet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib30)\]\. At the graph\-regularization level, regularized graph VAEs penalize invalid matrix\-valued outputs to encourage properties such as connectivity and valence consistency\[Maet al\.,[2018](https://arxiv.org/html/2606.13955#bib.bib31)\]\.

Another line of work introduces differentiable optimization or satisfiability layers inside neural networks\. OptNet embeds a differentiable quadratic programming solver as a neural network layer\[Amos and Kolter,[2017](https://arxiv.org/html/2606.13955#bib.bib48)\], while LinSATNet and GLinSAT project neural outputs toward feasible sets defined by linear satisfiability constraints\[Wanget al\.,[2023](https://arxiv.org/html/2606.13955#bib.bib49), Zenget al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib50)\]\. These methods provide useful templates for differentiable constrained prediction, but molecular validity is not simply a generic linear feasibility problem\. It involves discrete support constraints, atom\-type\-dependent valence, bond\-type compatibility, global graph connectivity, aromaticity and charge consistency, and consistency between topology and 3D geometry\.

Constrained diffusion methods impose hard structural constraints along the generative trajectory\. For example, ConStruct maintains graph constraints such as planarity or acyclicity by combining edge\-absorbing noise with projection operators\[Madeiraet al\.,[2024](https://arxiv.org/html/2606.13955#bib.bib51)\]\. These structural graph constraints are related to molecular validity, but they are not identical to chemical sanitization constraints\. In particular, molecular decoding must handle both global graph structure and local chemistry\-specific rules\.

TopVAE takes a different route from permanent projection layers and post\-hoc sanitization\. A permanent projection layer can create decoder dependency: the raw neural decoder may never internalize the constraints and may require the projection module at inference time\. Post\-hoc sanitization has the opposite limitation: it can repair invalid outputs after decoding, but it does not provide a training signal that teaches the decoder why the output was invalid\. TopVAE uses ChemCO\+AGCL as a selective training\-time teacher\. ChemCO computes chemically constrained bond corrections through unrolled optimization, and AGCL injects these corrections only when they improve ground\-truth bond likelihood\. As a result, the default decoder progressively internalizes chemical rules and can be used without ChemCO during inference\.

#### Positioning\.

TopVAE sits between complete\-molecule data\-space generation and molecular latent diffusion\. Like complete\-molecule data\-space models, it treats molecular topology and bond types as first\-class objects rather than recovering bonds only after coordinate generation\. Like latent diffusion methods, it preserves the efficiency and controllability of a compact continuous latent space\. Its central contribution is decoder\-centric: instead of asking only whether the encoder reconstructs posterior latents, TopVAE asks whether the latent space is chemically navigable in the regions that a diffusion prior will actually visit\.

Concretely, TopVAE addresses the two dominant dark\-area failure modes\. TopoBridge guarantees connected candidate topology through adjacency refinement, preventing disconnected molecular fragments\. ChemCO enforces chemistry\-specific bond and valence constraints through unrolled primal–dual optimization\. AGCL then distills ChemCO’s corrections into the raw decoder during training, enabling ChemCO\-free inference\. This makes TopVAE a complement to prior molecular latent diffusion models: it improves the decodability and chemical robustness of the latent space on which such priors depend\.

## Appendix BExtended Dark\-Area Analysis

VAE decoders can produce invalid outputs when the latent code falls outside the support of the aggregate posterior—so\-called*dark areas*of the latent space\. The diffusion prior is designed to keep generated latent codes on the data manifold, but the degree to which it succeeds depends on both the prior’s expressiveness and the decoder’s robustness\. We analyse dark areas from two complementary perspectives: the static geometry of the VAE latent space \(Section[B\.1](https://arxiv.org/html/2606.13955#A2.SS1)\) and the dynamic denoising trajectory of the diffusion prior \(Section[B\.2](https://arxiv.org/html/2606.13955#A2.SS2)\)\.

### B\.1VAE Latent Space Geometry

To probe the extent of dark areas, we randomly sample 2,000 QM9 molecules and encode each through three VAEs \(UAE, TopVAE, ADiT\)\. For each latent codez=μϕ​\(M\)z=\\mu\_\{\\phi\}\(M\), we add Gaussian perturbationsz′=z\+σ⋅semp⋅εz^\{\\prime\}=z\+\\sigma\\cdot s\_\{\\mathrm\{emp\}\}\\cdot\\varepsilonat noise levelsσ∈\{0,0\.05,0\.1,0\.2,0\.3,0\.5\}\\sigma\\in\\\{0,0\.05,0\.1,0\.2,0\.3,0\.5\\\}\(50 samples per nonzeroσ\\sigma;semps\_\{\\mathrm\{emp\}\}is the empirical latent standard deviation\), decode eachz′z^\{\\prime\}, and check validity via RDKit sanitisation and single\-fragment connectivity\. Figure[3](https://arxiv.org/html/2606.13955#A2.F3)shows a 2\-D PCA projection of the pooled per\-molecule mean latents, coloured by decoded validity\. TopVAE retains 90\.5% validity across all perturbation levels, compared to 44\.4% for UAE and 32\.6% for ADiT, indicating that TopoBridge and ChemCO substantially shrink the dark areas of the latent space\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x2.png)Figure 3:VAE dark\-area embeddings\.2\-D PCA of per\-molecule mean latent vectors \(n=2,000n\{=\}2\{,\}000, QM9 test set\)\. Green points decode to valid molecules; red points decode to invalid ones\. Invalid molecules concentrate in low\-density peripheral regions, delineating the dark areas of each latent space\.
### B\.2Diffusion Denoising Trajectory

To understand how the diffusion prior navigates the latent space during generation, we track molecules through the full denoising process\. At each diffusion timestep we decode the intermediate latent code, check RDKit validity, and record the first principal component of the per\-molecule mean latent\.

Figure[4](https://arxiv.org/html/2606.13955#A2.F4)plots these trajectories for three models: UDM\-3D, TopVAE\+DiT, and ADiT\. The dashed grey curve shows the instantaneous validity ratio across the batch\. Two observations stand out:

1. 1\.Latent spread\.UDM\-3D’s intermediate latents span a PC 1 range of∼300\{\\sim\}300units early in denoising, reflecting high variance in the unconstrainted latent space\. TopVAE\+DiT operates within a∼60\{\\sim\}60\-unit range, and ADiT within∼1\{\\sim\}1unit, indicating progressively tighter concentration on the data manifold\.
2. 2\.Validity transition\.UDM\-3D’s validity ratio rises slowly and does not reach 80% until denoising progress∼\{\\sim\}0\.75\. TopVAE\+DiT crosses 90% validity by progress 0\.25, and its intermediate molecules are already chemically plausible at early denoising stages\. This early validity is a direct consequence of the TopoBridge connectivity guarantee and ChemCO valence constraints, which make the decoder robust to partially denoised \(i\.e\., off\-manifold\) latent codes\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x3.png)Figure 4:Latent\-space trajectories during diffusion denoising\.Each dot represents a molecule at a specific denoising step; colour indicates validity \(green = valid, red = invalid\)\. The dashed grey curve shows the batch validity ratio\.\(e\)UDM\-3D exhibits wide latent spread and delayed validity\.\(f\)TopVAE\+DiT shows compact trajectories with early validity onset\.\(g\)ADiT operates in a narrow latent range with a gradual validity ramp\.
### B\.3Molecule Trajectory Showcase

Figure[5](https://arxiv.org/html/2606.13955#A2.F5)visualises representative denoising trajectories for TopVAE\+DiT and UDM\-3D at five diffusion timesteps \(t=1\.0,0\.5,0\.25,0\.1t=1\.0,0\.5,0\.25,0\.1, and the final output\)\. Molecules are grouped into three categories based on their validity profile across the trajectory:

- •Always valid\(top rows\): molecules that decode to valid structures at every sampled timestep\. TopVAE\+DiT produces such trajectories more frequently, reflecting the decoder’s built\-in chemical constraints\.
- •Transition\(middle rows\): molecules that start invalid at high noise \(t=1\.0t\{=\}1\.0\) and become valid during denoising\. TopVAE\+DiT transitions earlier \(t∼0\.5t\{\\sim\}0\.5\) while UDM\-3D transitions later \(t∼0\.25t\{\\sim\}0\.25or0\.10\.1\), consistent with the validity curves in Figure[4](https://arxiv.org/html/2606.13955#A2.F4)\.
- •Hard case\(bottom rows\): molecules that remain invalid for most of the trajectory before recovering at the final step, or that oscillate between valid and invalid states\. Even in hard cases, TopVAE\+DiT recovers valid structures by the final step, whereas UDM\-3D sometimes produces final molecules with residual valence violations or disconnected fragments \(red borders\)\.

Green borders indicate valid molecules; red borders indicate invalid ones\. The visualisation confirms that TopVAE’s decoder\-level constraints \(TopoBridge \+ ChemCO\) provide a robustness buffer that allows the diffusion prior to produce valid molecules even from partially denoised, off\-manifold latent codes\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x4.png)Figure 5:Denoising trajectories of representative molecules\.Green/red borders indicate valid/invalid decoded molecules at each timestep\.Top:TopVAE\+DiT produces valid intermediates earlier and more consistently\.Bottom:UDM\-3D often remains invalid until late denoising stages, and hard cases may retain valence errors in the final output\.

## Appendix C3D Conformer Quality

Tables[6](https://arxiv.org/html/2606.13955#A3.T6)and[7](https://arxiv.org/html/2606.13955#A3.T7)report Maximum Mean Discrepancy \(MMD\) scores for bond lengths, bond angles, and dihedral angles between generated and reference molecules, following thegeom\_predictorevaluation protocol used by JODO and EDM \(bonds inferred from 3D coordinates\)\.

Table 6:3D geometry MMD on QM9\(nn=10,000\)\. Lower is better\. Baseline numbers are from respective papers\.MethodBond Len\.↓\\downarrowBond Ang\.↓\\downarrowDihedral↓\\downarrowEDM1\.30e\-11\.82e\-26\.64e\-4GeoLDM2\.40e\-11\.00e\-26\.59e\-4JODO1\.48e\-11\.21e\-26\.29e\-4ADiT9\.98e\-13\.38e\-21\.46e\-3UDM\-3D7\.04e\-29\.84e\-33\.47e\-4TopVAE\+DiT6\.65e\-25\.74e\-34\.83e\-3Table 7:3D geometry MMD on GEOM\-Drugs\(nn=10,000\)\. Lower is better\. Baseline numbers from respective papers\.MethodBond Len\.↓\\downarrowBond Ang\.↓\\downarrowDihedral↓\\downarrowEDM4\.29e\-14\.96e\-11\.46e\-2GeoLDM3\.91e\-14\.22e\-11\.69e\-2JODO8\.49e\-21\.15e\-26\.68e\-4UDM\-3D9\.89e\-35\.11e\-31\.74e\-4TopVAE\+DiT1\.78e\-21\.58e\-24\.86e\-3#### Analysis\.

On QM9, TopVAE\+DiT achieves the best bond\-length MMD \(6\.65e\-2 vs\. 7\.04e\-2 for the next\-best method UDM\-3D, a 6% reduction\) and the best bond\-angle MMD \(5\.74e\-3 vs\. 9\.84e\-3 for UDM\-3D, a 42% reduction\)\. The dihedral MMD \(4\.83e\-3\) is weaker than UDM\-3D’s 3\.47e\-4; we attribute this to the topology\-first decoder factorization, which prioritizes graph correctness over torsional accuracy—a trade\-off that future work on coordinate refinement may address\.

On GEOM\-Drugs, TopVAE\+DiT achieves bond\-length MMD of 1\.78e\-2, on the same order of magnitude as UDM\-3D \(9\.89e\-3\) and substantially better than EDM \(24×\\times\), GeoLDM \(22×\\times\), and JODO \(4\.8×\\times\)\. This demonstrates that the topology\-first decoder, despite lacking explicit E\(3\)\-equivariant denoising in data space, produces competitive 3D geometries on drug\-sized molecules\.

## Appendix DMolecular Property Distributions

To assess whether TopVAE\+DiT generates molecules that faithfully reproduce the training distribution beyond aggregate FCD scores, we compute 8 physicochemical descriptors on the generated set \(nn=10,000\) and compare them to the reference\.UDM\-3D’s checkpoint was not released when we completed the manuscript; the results in this section are from our own re\-training\. They are for reference only and may not reflect UDM\-3D’s true performance\.

Table 8:Property statistics on QM9\(nn=10,000\)\. Reference values computed on the QM9 test set\. Bold indicates closer to reference\.PropertyDataset Ref\.UDM\-3DTopVAE\+DiTMW \(Da\)122\.6122\.4122\.6logP0\.350\.430\.45QED0\.4660\.4670\.468HeavyAtoms8\.798\.788\.79Rings1\.741\.872\.07Fsp30\.7090\.7500\.778TPSA \(Å2\)35\.632\.931\.3RotBonds0\.940\.880\.70Valid molecules—9,98610,000Table 9:Property statistics on GEOM\-Drugs\(nn=10,000\)\. Reference values computed on the GEOM\-Drugs training set\. Bold indicates closer to reference\.PropertyDataset Ref\.UDM\-3DTopVAE\+DiTMW \(Da\)355\.5321\.5351\.5logP2\.862\.412\.91QED0\.6460\.5300\.668HeavyAtoms24\.922\.224\.9Rings3\.002\.223\.27Fsp30\.3050\.5370\.334TPSA \(Å2\)73\.967\.268\.2RotBonds5\.054\.884\.75Valid molecules—1,823 / 10,00010,000/ 10,000#### Analysis\.

On GEOM\-Drugs, TopVAE\+DiT matches the reference distribution more closely than UDM\-3D on six of eight properties \(all except rotatable bonds, where UDM\-3D is slightly closer: 4\.88 vs\. 4\.75 against a reference of 5\.05\)\. The improvement is particularly pronounced for molecular weight \(1\.1% deviation vs\. 9\.6%\), heavy\-atom count \(exact match vs\. 10\.8% undercount\), and ring count \(9% deviation vs\. 26%\), indicating that UDM\-3D’s UAE backbone tends to generate structurally simpler molecules on this larger\-molecule benchmark\. TopVAE\+DiT also achieves 100% validity \(10,000/10,000\), compared to only 1,823/10,000 for UDM\-3D—a direct consequence of dark\-area closure\.

On QM9, both models produce near\-identical property statistics, consistent with the benchmark being saturated for small molecules \(≤\\leq9 heavy atoms\)\. UDM\-3D is closer to the reference on four properties \(logP, Rings, Fsp3, TPSA, RotBonds\), while TopVAE\+DiT matches the reference more closely on MW, QED, and HeavyAtoms\. TopVAE\+DiT achieves perfect validity \(10,000 vs\. 9,986\)\. Figures[6](https://arxiv.org/html/2606.13955#A4.F6)and[7](https://arxiv.org/html/2606.13955#A4.F7)visualize the full property distributions underlying the summary statistics in Tables[9](https://arxiv.org/html/2606.13955#A4.T9)and[8](https://arxiv.org/html/2606.13955#A4.T8)\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x5.png)Figure 6:Property distributions on QM9\(nn=10,000\)\. Green: test\-set reference; blue: TopVAE\+DiT; gray: UDM\-3D \(re\-trained\); orange: ADiT\. All three models closely match the reference distribution, consistent with QM9 being near\-saturated for small molecules \(≤\\leq9 heavy atoms\)\. Minor differences are visible in ring count and rotatable bonds, where TopVAE\+DiT slightly overrepresents 2\-ring structures\.![Refer to caption](https://arxiv.org/html/2606.13955v1/x6.png)Figure 7:Property distributions on GEOM\-Drugs\(nn=10,000\)\. Green: test\-set reference; blue: TopVAE\+DiT; gray: UDM\-3D \(re\-trained\)\. TopVAE\+DiT closely tracks the reference across all eight properties\. UDM\-3D exhibits a systematic leftward shift in molecular weight and heavy\-atom count, underrepresentation of 3–4 ring systems, and a bimodal QED distribution—consistent with its tendency to decode structurally simpler molecules due to dark\-area failures\.

## Appendix EDiversity and Novelty

To verify that dark\-area closure does not reduce structural diversity, Tables[11](https://arxiv.org/html/2606.13955#A5.T11)and[10](https://arxiv.org/html/2606.13955#A5.T10)report diversity and novelty metrics on both datasets\. Internal diversity is computed as1−Tanimoto¯1\-\\overline\{\\mathrm\{Tanimoto\}\}over Morgan fingerprints \(radius 2, 2048 bits, up to 500k sampled pairs\)\. Novelty is the fraction of unique SMILES absent from the training set\. Unique scaffolds are Murcko scaffolds of unique molecules\. Since UDM\-3D has not released pretrained checkpoints,all UDM\-3D results in this section are re\-trained using the official UDM\-3D GitHub repository and default configurations\.The results of UDM\-3D are for reference only\.

Table 10:Diversity and novelty on QM9\(nn=10,000\)\.UDM\-3D is re\-trained using the official UDM\-3D repository\(for reference only\)\. QM9 contains≤\\leq9 heavy atoms; novelty and unique scaffold rates are naturally lower than on GEOM\-Drugs because the training set already covers most of the reachable chemical space at this size\.MetricTopVAE\+DiTUDM\-3DValid molecules10,000\(100%\)9,986 \(99\.9%\)FCD \(2D\)↓\\downarrow0\.1850\.207Internal diversity↑\\uparrow0\.9140\.916Novelty↑\\uparrow45\.1%42\.3%Unique Murcko scaffolds↑\\uparrow36\.3%29\.6%Uniqueness↑\\uparrow95\.7%96\.6%Unique SMILES9,5719,660Unique scaffold count3,4772,864Table 11:Diversity and novelty on GEOM\-Drugs\(nn=10,000 sampled\)\.UDM\-3D is re\-trained using the official UDM\-3D repository\(for reference only\)\. UDM\-3D produces only 1,823 valid molecules \(18\.2%\); its diversity and novelty metrics are computed on this smaller set, which naturally inflates per\-molecule diversity and novelty\.MetricTopVAE\+DiTUDM\-3DValid molecules10,000\(100%\)1,823 \(18\.2%\)FCD \(2D\)↓\\downarrow2\.6825\.73Internal diversity↑\\uparrow0\.8740\.898∗Novelty↑\\uparrow99\.95%100%∗Unique Murcko scaffolds↑\\uparrow88\.7%85\.4%∗Uniqueness↑\\uparrow99\.99%100%∗∗Computed on 1,823 valid molecules; small\-sample sizes inflate these metrics\.#### QM9 analysis\.

On QM9, internal diversity is effectively tied \(0\.914 vs\. 0\.916\), consistent with the∼\{\\sim\}0\.91 range reported by EDM and GeoLDM on this dataset\. TopVAE\+DiT explores a substantially broader scaffold space, producing 3,477 unique Murcko scaffolds versus 2,864 for UDM\-3D—a 21% increase—and achieves higher novelty \(45\.1% vs\. 42\.3%\)\. TopVAE\+DiT also attains a lower FCD \(0\.185 vs\. 0\.207\), indicating a generated distribution that better matches the training set\. UDM\-3D has slightly higher uniqueness \(96\.6% vs\. 95\.7%\), but TopVAE\+DiT’s advantage in scaffold coverage and novelty suggests it explores a broader region of chemical space\.

#### GEOM\-Drugs analysis\.

TopVAE\+DiT achieves high internal diversity \(0\.874\), near\-perfect novelty \(99\.95%\), and 88\.7% unique Murcko scaffolds across all 10,000 generated molecules, confirming the absence of mode collapse\. UDM\-3D shows slightly higher internal diversity \(0\.898\) and novelty \(100%\), but these metrics are computed on only 1,823 valid molecules—the remaining 81\.8% fail RDKit sanitization due to dark\-area decoding failures\. The 10×\\timesdifference in FCD \(2\.68 vs\. 25\.73\) reflects the severe distributional shift caused by dark areas in the UAE latent space: the few valid molecules that survive sanitization are not representative of the target distribution\. TopVAE\+DiT achieves a higher unique scaffold rate \(88\.7% vs\. 85\.4%\) despite evaluating over5×5\\timesmore molecules, indicating genuine structural diversity rather than small\-sample artifacts\.

## Appendix FEnergy Relaxation Evaluation

To assess whether TopVAE’s topological improvements translate into more physically plausible 3D geometries, we conduct an MMFF94 energy\-relaxation study on GEOM\-Drugs\. For each generated molecule, we compute the MMFF94 single\-point energy of the generated conformer, run a local geometry optimization to the nearest force\-field minimum, and report \(i\) the relaxation energyΔ​E=Egenerated−Erelaxed\\Delta E=E\_\{\\text\{generated\}\}\-E\_\{\\text\{relaxed\}\}\(lower indicates the generated geometry is closer to equilibrium\) and \(ii\) the heavy\-atom RMSD between the generated and relaxed conformers\.

#### Results\.

Table[12](https://arxiv.org/html/2606.13955#A6.T12)reports MMFF94 relaxation statistics for TopVAE\+DiT, UDM\-3D, and reference molecules from the GEOM\-Drugs test set \(n=2,000n\{=\}2\{,\}000per source\)\. TopVAE\+DiT generates 3D geometries substantially closer to force\-field minima than UDM\-3D: the median relaxation energy is3\.9×3\.9\{\\times\}lower \(19\.8 vs\. 76\.8 kcal/mol\), and the median RMSD is 12% lower \(0\.99 vs\. 1\.12 Å\)\. Both generative models are expectedly farther from equilibrium than the reference dataset \(medianΔ​E=13\.9\\Delta E=13\.9kcal/mol\), but TopVAE is considerably closer\.

These results suggest that enforcing topological correctness via TopoBridge and ChemCO provides an indirect benefit to 3D geometric quality: a chemically consistent bond graph constrains the coordinate head and EGNN to produce conformers in more physically reasonable regions of configuration space\.

Table 12:MMFF94 energy relaxation on GEOM\-Drugs \(n=2,000n\{=\}2\{,\}000\)\.Δ​E\\Delta Eis the energy difference between the generated and locally optimized conformer \(lower is better\)\. RMSD is the heavy\-atom root\-mean\-square deviation between the two conformers\.MetricReferenceTopVAE\+DiTUDM\-3DΔ​E\\Delta Emean \(kcal/mol\)↓\\downarrow16\.540\.6104\.4Δ​E\\Delta Emedian \(kcal/mol\)↓\\downarrow13\.919\.876\.8RMSD mean \(Å\)↓\\downarrow0\.651\.141\.24RMSD median \(Å\)↓\\downarrow0\.510\.991\.12

## Appendix GHyperparameter Sensitivity Analysis

We study the sensitivity of TopVAE to two key design choices: TopoBridge hyperparameters \(adjacency sparsity, minimum degree\) and theΩ\(k\)\\Omega^\{\(k\)\}chemical mask\. All runs use QM9 with 500 training epochs and posterior evaluation \(n=10​kn\{=\}10\\text\{k\}\); absolute numbers may shift with full convergence \(2000 epochs\), but cross\-run rankings are expected to hold\.

### G\.1TopoBridge Hyperparameters

Table[13](https://arxiv.org/html/2606.13955#A7.T13)varies the adjacency sparsity \(adj\_max\_deg, i\.e\., the top\-kkneighbor count\) and the minimum\-degree penalty \(dmind\_\{\\min\}\) in ChemCO\.

Table 13:TopoBridge hyperparameter sensitivityon QM9 \(500 epochs,n=10​kn\{=\}10\\text\{k\}posterior\)\. Default:adj\_max\_deg=4=4,dmin=1d\_\{\\min\}\{=\}1\.RunChange vs\. defaultValid\.↑\\uparrowUniq\.↑\\uparrow3DS∧\\wedgeC↑\\uparrow3D AtomStab\.↑\\uparrowtopk\_3adj\_max\_deg=3=3\(sparse\)1\.0000\.9160\.01460\.689topk\_8adj\_max\_deg=8=8\(dense\)1\.0000\.9250\.01600\.682dmin\_0dmin=0d\_\{\\min\}\{=\}0\(no anti\-degeneracy\)1\.0000\.9220\.04350\.722#### Findings\.

1. 1\.Robust to adjacency sparsity\.Varyingadj\_max\_degfrom 3 to 8 produces nearly identical metrics \(3D S∧\\wedgeC: 0\.0146 vs\. 0\.0160\)\. ChemCO compensates for sparser adjacency by reallocating bond mass under valence constraints\.
2. 2\.Per\-dataset tuning ofdmind\_\{\\min\}\.Removing the minimum\-degree penalty \(dmin=0d\_\{\\min\}\{=\}0\) yields the highest 3D S∧\\wedgeC on QM9 \(2\.7×2\.7\{\\times\}over denset​o​p​k​\_​8topk\\\_8setting\), suggesting that thedmind\_\{\\min\}term over\-regularizes on small molecules \(≤9\{\\leq\}9heavy atoms\)\. The term was designed for GEOM\-Drugs, where bond collapse is the dominant failure mode; this result justifies per\-dataset hyperparameter selection\.

### G\.2Ω\(k\)\\Omega^\{\(k\)\}Chemical Mask Sensitivity

Table[14](https://arxiv.org/html/2606.13955#A7.T14)compares the default QM9 configuration \(valence caps only,Ω\(k\)=𝟏\\Omega^\{\(k\)\}\\\!=\\\!\\bm\{1\}\) with the GEOM\-style explicit pair mask \(Appendix[M\.2](https://arxiv.org/html/2606.13955#A13.SS2)\)\.

Table 14:Ω\(k\)\\Omega^\{\(k\)\}mask sensitivityon QM9 \(500 epochs,n=10​kn\{=\}10\\text\{k\}posterior\)\.Ω\(k\)\\Omega^\{\(k\)\}configValid\.↑\\uparrowUniq\.↑\\uparrow3DS∧\\wedgeC↑\\uparrow3D AtomStab\.↑\\uparrowValence caps only \(default\)1\.0000\.9250\.01600\.682\+ Explicit pair restrictions1\.0000\.9200\.01910\.688Adding the explicit chemical pair mask yields a 19% relative gain in 3D S∧\\wedgeC \(0\.0160→\\to0\.0191\) with no loss in validity or uniqueness\. The gain is mild on QM9 \(only 4 atom types\) and is expected to be larger on datasets with more diverse atom\-type combinations\. This validates thatΩ\(k\)\\Omega^\{\(k\)\}captures pair\-specific chemical knowledge that valence caps alone cannot encode\.

## Appendix HVAE Posterior Reconstruction Quality

Tables[15](https://arxiv.org/html/2606.13955#A8.T15)and[16](https://arxiv.org/html/2606.13955#A8.T16)report VAE posterior reconstruction quality on QM9 and GEOM\-Drugs respectively\. These results complement the dark\-area diagnosis in Section[5\.2](https://arxiv.org/html/2606.13955#S5.SS2)by showing that high posterior MolStab does not guarantee high Stab∧\\wedgeConn: on GEOM\-Drugs, UAE achieves MolStab=0\.998\{=\}0\.998but Stab∧\\wedgeConn=0\.000\{=\}0\.000, as every reconstructed molecule fragments into disconnected atoms\.

Table 15:VAE posterior reconstruction on QM9\(n=10​kn\{=\}10\\text\{k\}validation set,†\\dagger\)\.MethodFCD↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowV&C↑\\uparrowFCD3D↓\\downarrowMolStab\.3D↑\\uparrowUAE0\.1001\.0001\.0000\.99924\.170\.000ADiT\-VAE0\.7951\.0000\.9510\.9040\.8980\.900TopVAE0\.0891\.0001\.0001\.00015\.390\.002Table 16:VAE posterior reconstruction on GEOM\-Drugs\(n=1997n\{=\}1997validation set,†\\dagger\)\. UAE achieves near\-perfect MolStab but zero connectivity—the atom\-cloud failure mode discussed in Section[5\.2](https://arxiv.org/html/2606.13955#S5.SS2)\. TopVAE numbers reported with the same greedy valence repair used by TopVAE\+DiT\.MethodFCD↓\\downarrowAtomStab\.↑\\uparrowMolStab\.↑\\uparrowV&C↑\\uparrowFCD3D↓\\downarrowMolStab\.3D↑\\uparrowUAE60\.671\.0000\.9980\.00044\.280\.002TopVAE46\.170\.9990\.9600\.99844\.280\.000
## Appendix IScaffold Inpainting: Per\-Scaffold Results

Table[17](https://arxiv.org/html/2606.13955#A9.T17)provides the per\-scaffold breakdown of the GEOM\-Drugs inpainting results summarized in Table[4](https://arxiv.org/html/2606.13955#S5.T4)\. All experiments use differential inpainting with noise scale=0\.3\{=\}0\.3, 1k samples per cell\. The advantage of TopVAE is consistent across all five scaffolds, with cyclohexane \(saturated, non\-aromatic\) being the easiest and naphthalene/indole \(fused aromatic rings\) the most challenging for scaffold preservation \(Figure[9](https://arxiv.org/html/2606.13955#A9.F9),[11](https://arxiv.org/html/2606.13955#A9.F11)\)\.

Table 17:GEOM\-Drugs scaffold inpainting: per\-scaffold breakdown\.1k samples per cell, differential inpainting \(noise scale=0\.3\{=\}0\.3\)\. Values in %\. Bold = better per cell\.Scaff\. Pres\.↑\\uparrowS∧\\wedgeC↑\\uparrowScaffold\+n\+nNNUAETopVAEUAETopVAE![[Uncaptioned image]](https://arxiv.org/html/2606.13955v1/scaffold_imgs/benzene.png)5117\.04\.095\.697\.210168\.49\.289\.094\.820267\.421\.286\.687\.6404628\.460\.862\.270\.260660\.650\.21\.651\.280860\.034\.80\.025\.21001060\.016\.40\.09\.6![[Uncaptioned image]](https://arxiv.org/html/2606.13955v1/scaffold_imgs/pyridine.png)5113\.43\.495\.697\.610164\.05\.290\.293\.620262\.67\.887\.887\.240462\.811\.069\.476\.060660\.011\.22\.053\.680860\.06\.80\.029\.41001060\.02\.80\.08\.8![[Uncaptioned image]](https://arxiv.org/html/2606.13955v1/scaffold_imgs/naphthalene.png)5151\.81\.681\.292\.010201\.01\.279\.689\.420300\.01\.685\.085\.840500\.03\.046\.268\.260700\.03\.00\.452\.680900\.02\.40\.031\.01001100\.00\.40\.08\.2![[Uncaptioned image]](https://arxiv.org/html/2606.13955v1/scaffold_imgs/indole.png)5142\.42\.079\.695\.010193\.23\.276\.485\.420290\.22\.484\.684\.440490\.06\.466\.669\.060690\.09\.01\.848\.480890\.06\.40\.025\.01001090\.02\.80\.09\.4![[Uncaptioned image]](https://arxiv.org/html/2606.13955v1/scaffold_imgs/cyclohexane.png)51135\.245\.697\.097\.8101657\.258\.095\.492\.2202665\.646\.690\.479\.2404658\.457\.875\.064\.8606624\.271\.222\.251\.280860\.075\.80\.421\.21001060\.076\.60\.07\.2Mean9\.020\.647\.561\.1![Refer to caption](https://arxiv.org/html/2606.13955v1/x7.png)Figure 8:3D Generation Molecules\.![Refer to caption](https://arxiv.org/html/2606.13955v1/x8.png)Figure 9:Scaffold Inpainting \(Part 1\)\.![Refer to caption](https://arxiv.org/html/2606.13955v1/x9.png)Figure 10:Scaffold Inpainting \(Part 2\)\.![Refer to caption](https://arxiv.org/html/2606.13955v1/x10.png)Figure 11:Scaffold Inpainting \(Part 3\)\.
## Appendix JTopoBridge Chemical Plausibility Analysis

TopoBridge guarantees graph connectivity by adding repair edges between disconnected components through greedy BFS refinement \(Eq\.[16](https://arxiv.org/html/2606.13955#S4.E16)–[17](https://arxiv.org/html/2606.13955#S4.E17)\)\. To isolate the effect of BFS connectivity repair from the decoder’s own predictions, we instrument the inference pipeline to track*BFS\-inserted edges*: edges that did not exist after initial thresholding/top\-kkof the adjacency logits and were explicitly added by the BFS repair procedure \(isolated\-node fix or component\-merging\) to ensure connectivity\. Concretely, we snapshot the adjacency matrix𝐀init\\mathbf\{A\}\_\{\\mathrm\{init\}\}immediately after thresholding and before any repair, then identify inserted edges via𝐌ins=clamp​\(𝐀final−𝐀init,min=0\)\\mathbf\{M\}\_\{\\mathrm\{ins\}\}=\\mathrm\{clamp\}\(\\mathbf\{A\}\_\{\\mathrm\{final\}\}\-\\mathbf\{A\}\_\{\\mathrm\{init\}\},\\,\\min\{=\}0\)\. This differs from the broader set of graph\-theoretic bridge bonds \(cut edges found by Tarjan’s algorithm\), most of which originate from the decoder’s adjacency head and are not artifacts of repair\.

#### BFS\-inserted edge chemistry\.

Table[18](https://arxiv.org/html/2606.13955#A10.T18)summarizes the chemical properties of BFS\-inserted edges in TopVAE\+DiT generations on GEOM\-Drugs\. Among 10,000 generated molecules, only 2\.1% \(210 molecules\) receive any BFS\-inserted edge, totaling 211 inserted edges out of approximately 270,000 total edges — i\.e\., 0\.078% of all bonds\. This shows that BFS repair is nearly idle at inference time: the decoder has learned to produce connected graphs during training via the BFS\-STE straight\-through estimator, and the repair mechanism serves as a safety net that rarely fires\.

Among the 211 BFS\-inserted edges, 98\.1% connect atom pairs for which the assigned bond type is chemically allowed under theΩ\(k\)\\Omega^\{\(k\)\}atom\-pair mask\. This is expected because the repair selects the highest\-probability edge predicted by the adjacency head, so even forced connections involve plausible atom pairs\. The mean bond length of inserted edges is 2\.55 Å \(σ=0\.41\\sigma\{=\}0\.41\), longer than typical covalent bonds, with 40\.8% falling within the standard covalent bonding range of 0\.8–2\.5 Å\. This is consistent with their role as last\-resort cross\-fragment connections between components that were nearly but not quite linked by the decoder\.

Table 18:BFS\-inserted edge audit on GEOM\-Drugs \(n=10,000n\{=\}10\{,\}000generated molecules\)\. A BFS\-inserted edge is an edge absent after initial adjacency thresholding that was explicitly added by BFS connectivity repair\. For comparison, the corresponding statistics for all graph\-theoretic bridge bonds \(cut edges\) are shown in the right column; note that the vast majority of bridge bonds are predicted by the decoder, not inserted by repair\.StatisticBFS\-inserted edgesAll bridge bondsMolecules affected2\.1% \(210/10,000\)99\.9% \(9,989/10,000\)Total edges21196,791Mean per molecule0\.029\.68Mean fraction of all edges0\.078%35\.8%Chemically allowed atom pairs98\.1%99\.97%Bond length in covalent range \[0\.8, 2\.5\] Å40\.8%99\.87%Mean bond length2\.55 Å \(σ=0\.41\\sigma\{=\}0\.41\)1\.43 Å \(σ=0\.15\\sigma\{=\}0\.15\)Bond type: Single / Double / Triple98\.1% / 1\.5% /<1\{<\}1%81% / 18% /<1\{<\}1%
#### Degree and ring distributions\.

To assess whether TopoBridge systematically increases graph density, we compare atom\-degree, ring\-size, and bridge\-bond distributions against the reference test set and UDM\-3D, which shares the same diffusion backbone but has no TopoBridge module \(Table[19](https://arxiv.org/html/2606.13955#A10.T19)\)\.

Table 19:Degree, ring, and graph\-theoretic bridge\-bond statistics on GEOM\-Drugs \(n=10,000n\{=\}10\{,\}000\)\. Lower EMD↓\\downarrowindicates closer match to the reference distribution\.MetricReferenceTopVAE\+DiTUDM\-3DMean atom degree2\.0892\.1852\.146Graph\-theoretic bridge\-bond fraction0\.6430\.3560\.384Degree EMD vs\. ref\.↓\\downarrow—0\.5680\.558Ring\-size EMD vs\. ref\.↓\\downarrow—0\.5610\.706The mean degree of TopVAE\-generated molecules \(2\.185\) is within 5% of the reference value \(2\.089\) and comparable to UDM\-3D \(2\.146\), which uses no TopoBridge connectivity refinement\. The two models trade wins across the EMD metrics: TopVAE\+DiT better reproduces ring sizes \(0\.5610\.561vs\.0\.7060\.706\), while the degree EMDs are nearly tied \(0\.5680\.568vs\.0\.5580\.558\)\. The graph\-theoretic bridge\-bond fraction of TopVAE\+DiT \(0\.356\) is lower than the reference value \(0\.643\) and close to UDM\-3D \(0\.384\), showing that TopoBridge does not inflate the fraction of final bonds occupying graph\-theoretic bridge positions\.

#### Summary\.

Taken together, these results support three conclusions\. First, TopoBridge’s primary contribution is at*training*time: the BFS\-STE gradient signal teaches the decoder to predict connected graphs, so that BFS repair at inference is nearly idle \(2\.1% of molecules, 0\.078% of edges\)\. Second, when repair does fire, it inserts a single edge connecting chemically valid atom pairs \(98\.1% allowed\), though with atypically long bond lengths \(mean 2\.55 Å\) reflecting their last\-resort nature\. This affects a negligible fraction of all edges and does not measurably perturb overall geometric quality\. Third, TopoBridge does not systematically increase graph density: generated molecules exhibit near\-reference mean degree and competitive degree EMD\. Since BFS repair fires on only 2\.1% of generated molecules, the remaining connectivity gain is attributable to learned decoder behavior rather than chemically implausible forced edges\.

### J\.1Design Justification: Why Not Soft Connectivity Penalties?

A natural alternative to TopoBridge is a differentiable soft connectivity penalty—for example, maximising the Fiedler value \(second\-smallest eigenvalue of the graph Laplacian\) or penalising the number of connected components via a smooth surrogate\. We explored soft penalty formulations and encountered several practical difficulties based on our experiments:

1. 1\.No guarantee\.Soft penalties reduce but do not eliminate disconnected outputs; molecules can still fragment when the penalty weight is insufficient or when the penalty landscape has local minima\.
2. 2\.Loss balancing\.The appropriate penalty weight varies with molecule size, batch composition, and training stage, requiring extensive hyperparameter search that hard TopoBridge avoids entirely\.
3. 3\.Gradient quality\.Computing eigenvalues of theN×NN\{\\times\}NLaplacian at each training step isO​\(N3\)O\(N^\{3\}\)and produces gradients that are numerically unstable near degenerate eigenvalues \(exactly the regime of near\-disconnected graphs\)\. TopoBridge’s BFS procedure isO​\(N\+E\)O\(N\{\+\}E\)and uses a straight\-through estimator that avoids these issues\.

TopoBridge provides a hard connectivity guarantee at modest computational cost that self\-anneals during training \(Appendix[K\.1](https://arxiv.org/html/2606.13955#A11.SS1)\), and the chemical plausibility audit \(Appendix[J](https://arxiv.org/html/2606.13955#A10)\) confirms that this guarantee does not come at the expense of chemical validity\. We therefore adopt TopoBridge as the preferred connectivity mechanism\.

## Appendix KComputational Cost Analysis

We profile TopVAE’s training cost on an NVIDIA H100 NVL GPU with batch size 64 on GEOM\-Drugs\. Component timings are forward\-only and averaged over 50 forward passes after 10 warmup iterations\. Training\-step timings are averaged over 50 forward\+backward passes after the same warmup\.

#### Component\-level breakdown\.

Table[20](https://arxiv.org/html/2606.13955#A11.T20)reports the per\-component cost of a single TopVAE forward pass\. The encoder accounts for 36\.7% of forward time\. ChemCO contributes 24\.0% but isremoved entirely at inference, making it a training\-only investment\. TopoBridge’s BFS\-STE adds only 2\.8 ms per batch \(2\.8% of forward time\) when measured at convergence \(α=1\.0\\alpha\{=\}1\.0\); as we show in Section[K\.1](https://arxiv.org/html/2606.13955#A11.SS1), this cost is self\-annealing and remains negligible throughout the majority of training\. Note that Tarjan bridge detection is a standardO​\(V\+E\)O\(V\{\+\}E\)graph algorithm independent of TopoBridge and is folded into the “Other” category\.

Table 20:Per\-component forward\-pass cost of TopVAE on GEOM\-Drugs \(H100 NVL, batch size 64\)\. TopoBridge timing uses the BFS\-only cost from a converged model \(α=1\.0\\alpha\{=\}1\.0\), which is representative of the vast majority of training\.ComponentTime \(ms\)% of ForwardEncoder36\.836\.7%AdjHead \+ BFS\-STE \(converged\)2\.82\.8%TopoFormer2\.82\.8%AtomHead0\.20\.2%BondTypeHead1\.11\.1%Bond Distribution Assembly0\.50\.5%ChemCO \(40 Steps\)24\.124\.0%CoordHead \+ EGNN7\.87\.8%Loss computation2\.72\.7%Other \(Tarjan, data transfer, etc\.\)21\.821\.7%Total forward100\.3100%
#### Comparison with UAE\.

Table[21](https://arxiv.org/html/2606.13955#A11.T21)compares the total training\-step cost between TopVAE and UAE\. TopVAE is4\.3×4\.3\{\\times\}slower per forward\+backward step and uses5\.1×5\.1\{\\times\}more peak GPU memory\. Under 4\-GPU DDP \(H100 NVL\), communication partially overlaps with computation and the effective per\-step ratio drops to3\.11×3\.11\{\\times\}\(295295ms/step vs\.9595ms/step from training logs\)\. This overhead is a training\-time investment: at inference, ChemCO is disabled \(saving 24% of forward time\), and the remaining modules produce guaranteed\-connected, chemically valid molecules\-\.

Table 21:Training cost comparison between TopVAE and UAE on GEOM\-Drugs \(H100 NVL, batch size 64\)\.MetricTopVAEUAERatioForward only \(ms\)100\.341\.12\.4×2\.4\{\\times\}Forward \+ backward \(ms\)180\.941\.44\.3×4\.3\{\\times\}Peak GPU memory \(MB\)4,5468955\.1×5\.1\{\\times\}
### K\.1TopoBridge Cost is Self\-Annealing

The TopoBridge’s BFS repair cost depends on the number of disconnected components in the initial top\-kkgraph, which in turn depends on the quality of the adjacency predictor\. Early in training, the predictor outputs near\-random probabilities, producing highly fragmented graphs that require many bridge insertions; as training progresses, the predictor sharpens and the top\-kkgraph is already connected for most molecules, so the BFS loop fires rarely\.

Figure[12](https://arxiv.org/html/2606.13955#A11.F12)quantifies this relationship by interpolating between a random adjacency matrix and the converged predictor output via a mixing coefficientα∈\[0,1\]\\alpha\\in\[0,1\]:PαA=α⋅PtrainedA\+\(1−α\)⋅PrandomAP^\{A\}\_\{\\alpha\}=\\alpha\\cdot P^\{A\}\_\{\\text\{trained\}\}\+\(1\{\-\}\\alpha\)\\cdot P^\{A\}\_\{\\text\{random\}\}\.The data is simulated which might different from real training data\.Atα=0\\alpha\{=\}0\(random\), the BFS repair takes11\.811\.8ms/batch and requires∼20\{\\sim\}20bridge insertions per molecule; byα=0\.25\\alpha\{=\}0\.25the cost has already collapsed to∼3\{\\sim\}3ms, and atα=1\.0\\alpha\{=\}1\.0\(converged\) it stabilises at2\.82\.8ms \(∼0\{\\sim\}0repair rounds for most molecules\)\. Panel \(b\) shows that the average node degree \(pre\-prune\) drops from∼25\{\\sim\}25to themax\_deg==6 ceiling as the predictor improves, explaining why the initial graph becomes well\-connected and repair becomes unnecessary\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x11.png)Figure 12:BFS repair cost as a function of adjacency predictor qualityα\\alpha\(GEOM\-Drugs, H100 NVL, batch 64\)\.\(a\)Wall\-clock time drops4\.2×4\.2\{\\times\}from random \(11\.811\.8ms\) to converged \(2\.82\.8ms\)\.\(b\)The number of BFS repair rounds \(blue\) and the pre\-prune average node degree \(orange\) both decrease as the predictor sharpens\.Figure[13](https://arxiv.org/html/2606.13955#A11.F13)shows the same effect measured at three discrete training stages \(random initialisation, early training, converged\)\. The BFS repair overhead drops from7\.57\.5ms \(2\.5% of step\) at random initialisation to1\.21\.2ms \(0\.4%\) during early training and2\.12\.1ms \(0\.7%\) at convergence, confirming that the cost self\-anneals and remains negligible throughout the majority of training\. The data is collected during real training process\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x12.png)Figure 13:TopoBridge BFS repair overhead at three training stages\. The cost drops from 2\.5% to<1\{<\}1% of total step time as the adjacency predictor improves\.
### K\.2Scaling with Molecule Size and Batch Size

A natural concern is whether TopoBridge and ChemCO scale to molecules larger than those in GEOM\-Drugs \(≤90\{\\leq\}90heavy atoms\)\. We profile both components on synthetic inputs withNNranging from 25 to 200 heavy atoms \(Figure[14](https://arxiv.org/html/2606.13955#A11.F14)\)\.

![Refer to caption](https://arxiv.org/html/2606.13955v1/x13.png)Figure 14:Per\-molecule wall\-clock time for TopoBridge refinement and ChemCO \(40 steps\)\.\(a\)Scaling with heavy\-atom countNN\(batch 64\)\. TopoBridge grows slowly \(0\.05→0\.130\.05\\to 0\.13ms/mol fromN=25N\{=\}25to200200\), while ChemCO remains the dominant cost at∼1\.1\{\\sim\}1\.1ms/mol across all sizes\.\(b\)Batch\-size effect atN=100N\{=\}100\. TopoBridge is batch\-size\-insensitive; ChemCO benefits from GPU parallelism at larger batches\.#### TopoBridge refinement

scales nearly flat with molecule size: per\-molecule cost increases only from∼0\.05\{\\sim\}0\.05ms atN=25N\{=\}25to∼0\.13\{\\sim\}0\.13ms atN=200N\{=\}200\(Figure[14](https://arxiv.org/html/2606.13955#A11.F14)a\)\. This is expected because the BFS repair cost is governed by the number of disconnected components \(and hence bridge insertions\), not directly by atom count\. For a molecule withcccomponents the greedy loop executes exactlyc−1c\{\-\}1insertions, each requiring anO​\(N\)O\(N\)BFS\. The worst case isO​\(N2\)O\(N^\{2\}\)\(all atoms isolated\), but the converged adjacency predictor produces less components compared to early stage, keeping repair cost near\-constant\. TopoBridge is also insensitive to batch size \(Figure[14](https://arxiv.org/html/2606.13955#A11.F14)b\), remaining at∼0\.1\{\\sim\}0\.1ms/mol across batch sizes 16–128\.

#### ChemCO

is the true scalability bottleneck\. Its per\-molecule cost \(∼1\.1\{\\sim\}1\.1ms, 40 steps\) is roughly10×10\{\\times\}that of TopoBridge across all molecule sizes and dominates at small batch sizes due to lower GPU utilisation \(Figure[14](https://arxiv.org/html/2606.13955#A11.F14)b\)\. Because ChemCO is removed at inference with no quality loss \(Table 6\), the scalability concern applies only to training\. For future work on molecules with\>\>100 heavy atoms, ChemCO’s iterative projection is the natural target for amortised or approximate alternatives; TopoBridge itself is expected to remain negligible\.

Table[22](https://arxiv.org/html/2606.13955#A11.T22)summarises the end\-to\-end forward\-pass scaling on real GEOM\-Drugs molecules, confirming modest overhead growth\.

Table 22:Forward\-pass time by molecule size on GEOM\-Drugs \(single GPU, batch 64\)\.Atom countTopVAE \(ms\)UAE \(ms\)RatioN∈\[15,25\)N\\in\[15,25\)88\.741\.12\.2×2\.2\{\\times\}N∈\[25,35\)N\\in\[25,35\)101\.141\.42\.4×2\.4\{\\times\}

## Appendix LTraining Protocol

This section provides complete training details for reproducing TopVAE\.

### L\.1Data Preprocessing

#### Kekulization\.

All molecules are kekulized viaChem\.Kekulize\(mol, clearAromaticFlags=True\), converting aromatic bonds to alternating single/double bonds and reducing the bond vocabulary from\{single, double, triple, aromatic\}\\\{\\text\{single, double, triple, aromatic\}\\\}toK=3K\\\!=\\\!3kekulized types\{single, double, triple\}\\\{\\text\{single, double, triple\}\\\}\. Molecules in datasets that fail kekulization are discarded\.

#### Atom and bond vocabularies\.

QM9 \(heavy\-only\): 4 atom types\{C, N, O, F\}\\\{\\text\{C, N, O, F\}\\\}, max 9 heavy atoms\. GEOM\-Drugs \(heavy\-only\): 12 atom types\{C, N, O, F, S, Cl, Br, P, I, B, Si, Bi\}\\\{\\text\{C, N, O, F, S, Cl, Br, P, I, B, Si, Bi\}\\\}, max 90 heavy atoms\.

#### Input features\.

The encoder receives a fully\-connected edge graph \(N×NN\\\!\\times\\\!N\)\. Bond edges carry one\-hot bond type; non\-bond edges carry zeros; self\-loops receive a dedicated flag\. Node features concatenate atom\-type one\-hot with 39 GeoMol\-style features \(atomic number, aromaticity, degree, hybridization, implicit valence, formal charge, ring membership\), plus 3D coordinates\.

#### Augmentation\.

All coordinates are mean\-centered\. GEOM\-Drugs training applies random SO\(3\) rotation and Gaussian translation \(scale0\.10\.1\); QM9 does not\.

### L\.2Model Architectures

#### Encoder\.

Both UAE and TopVAE share the same encoder based on DMTBlock transformer layers, with Gaussian Basis Function distance embedding and message\-passing attention\. The encoder projects per\-atom hidden states to𝝁\\bm\{\\mu\}andlog⁡𝝈2\\log\\bm\{\\sigma\}^\{2\}via linear heads, with posterior log\-variance clamped to\[−10,4\]\[\-10,4\]on GEOM\-Drugs\.

#### TopoFormer blocks\.

The topology\-conditioned transformer \(Sec\.[4\.2](https://arxiv.org/html/2606.13955#S4.SS2)\) uses Graphormer\-style additive pair bias from the binary adjacency:bias=nn\.Embedding​\(2,nheads\)​\[A~\]\\text\{bias\}\\\!=\\\!\\texttt\{nn\.Embedding\}\(2,n\_\{\\mathrm\{heads\}\}\)\[\\widetilde\{A\}\]\. Each block applies pre\-LayerNorm, multi\-head self\-attention with additive bias, and a GELU FFN \(4×4\\timesexpansion\), with residual connections\.

#### EGNN coordinate head\.

The E\(nn\)\-equivariant coordinate prediction uses normalized directions with tanh\-bounded magnitude:

Δ​𝒓i=∑j𝒓i−𝒓j‖𝒓i−𝒓j‖\+ϵ​tanh⁡\(ϕx​\(mi​j\)\),𝒓i←𝒓i\+Δ​𝒓i,\\Delta\\bm\{r\}\_\{i\}=\\sum\_\{j\}\\frac\{\\bm\{r\}\_\{i\}\-\\bm\{r\}\_\{j\}\}\{\\\|\\bm\{r\}\_\{i\}\-\\bm\{r\}\_\{j\}\\\|\+\\epsilon\}\\;\\tanh\\\!\\bigl\(\\phi\_\{x\}\(m\_\{ij\}\)\\bigr\),\\qquad\\bm\{r\}\_\{i\}\\leftarrow\\bm\{r\}\_\{i\}\+\\Delta\\bm\{r\}\_\{i\},\(26\)whereϵ=10−6\\epsilon\\\!=\\\!10^\{\-6\}andϕx\\phi\_\{x\}is a learned scalar MLP\. Coordinates are re\-centered after each EGNN layer\.

#### Bond type head\.

For each pair\(i,j\)\(i,j\)on the supportA~\\widetilde\{A\}, the bond head concatenates node pair features, learned atom\-type embeddings, adjacency embedding, and optionally RBF distance features \(16 Gaussian centers spanning\[0\.8,3\.0\]\[0\.8,3\.0\]Å on GEOM\-Drugs\)\.

#### Adjacency head\.

Edge\-existence logits are produced by a 2\-layer MLP on concatenated node pairs, symmetrized and zero\-diagonal before sigmoid\.

#### Architecture summary\.

Table 23:Architecture configurations\.ParameterQM9GEOM\-DrugsLatent dimdd1632Encoder hidden / heads / blocks64 / 4 / 3128 / 8 / 4Decoder \(Topo / Geo layers\)3 / 34 / 4Decoder hidden / heads64 / 4128 / 8Edge dim64128Dropout0\.10\.1RBF distance featuresNoYes \(16 centers\)

### L\.3Training Hyperparameters

Table 24:Training hyperparameters for TopVAE \(Stage 1: VAE\)\.ParameterQM9GEOM\-DrugsOptimizerAdamWAdamWLearning rate10−310^\{\-3\}3×10−43\\times 10^\{\-4\}Weight decay10−510^\{\-5\}10−510^\{\-5\}LR schedulerConstantCosine \(Tmax=1200T\_\{\\mathrm\{max\}\}\\\!=\\\!1200\)LR warmup–Linear, 50 epochsMax epochs20001200Batch size×\\timesGPUs256×4256\\times 4128×4128\\times 4Precisionfp32fp32Gradient clipping–1\.0*Loss weights \(Eq\.[24](https://arxiv.org/html/2606.13955#S4.E24)\)*wadjw\_\{\\mathrm\{adj\}\}1\.010\.0watomw\_\{\\mathrm\{atom\}\},wbondw\_\{\\mathrm\{bond\}\},λR\\lambda\_\{R\}1, 1, 11, 1, 1λD\\lambda\_\{D\}\(all\-pairs \+ bonded\-pair distance\)01\.0λAGCL\\lambda\_\{\\mathrm\{AGCL\}\}1\.00\.1β\\beta\(KL weight\)10−610^\{\-6\}5×10−65\\times 10^\{\-6\}*TopoBridge*Thresholdτ\\tau0\.50\.3Max degree cap46 \(Considering charges\)Adjacency lossBCEFocal \(α=0\.75\\alpha\\\!=\\\!0\.75,γ=1\.0\\gamma\\\!=\\\!1\.0\)*ChemCO \(Sec\.[4\.4](https://arxiv.org/html/2606.13955#S4.SS4), Appendix[M](https://arxiv.org/html/2606.13955#A13)\)*TCOT\_\{\\mathrm\{CO\}\}\(train / eval\)20 / 4040 / 100Primal base rateα\\alpha/ decayγα\\gamma\_\{\\alpha\}0\.05 / 0\.990\.05 / 0\.99Dual base ratesημ\\eta\_\{\\mu\},η\+\\eta\_\{\+\},η−\\eta\_\{\-\}0\.1, 0\.2, 0\.20\.1, 0\.2, 0\.2Active\-set sharpnessγ\\gamma10\.010\.0Chemical maskΩ\\OmegaNoYes
### L\.4Loss Function Details

The full training objective \(Eq\.[24](https://arxiv.org/html/2606.13955#S4.E24)\) consists of the terms defined below\. Let𝑴node∈\{0,1\}B×N\\bm\{M\}\_\{\\mathrm\{node\}\}\\\!\\in\\\!\\\{0,1\\\}^\{B\\times N\}denote the valid\-atom mask and𝑴pair=𝑴node\(i\)∧𝑴node\(j\)∧𝟙​\[i<j\]\\bm\{M\}\_\{\\mathrm\{pair\}\}=\\bm\{M\}\_\{\\mathrm\{node\}\}^\{\(i\)\}\\wedge\\bm\{M\}\_\{\\mathrm\{node\}\}^\{\(j\)\}\\wedge\\mathbb\{1\}\[i\\\!<\\\!j\]the upper\-triangular valid\-pair mask\.

#### \(i\) Adjacency lossℒadj\\mathcal\{L\}\_\{\\mathrm\{adj\}\}\.

On GEOM\-Drugs we use focal binary cross\-entropy:

ℒadj=1\|𝑴pair\|​∑\(i,j\)∈𝑴pairαt​\(1−pt\)γ​BCE​\(si​j,ti​j\),\\mathcal\{L\}\_\{\\mathrm\{adj\}\}=\\frac\{1\}\{\|\\bm\{M\}\_\{\\mathrm\{pair\}\}\|\}\\sum\_\{\(i,j\)\\in\\bm\{M\}\_\{\\mathrm\{pair\}\}\}\\alpha\_\{t\}\\,\(1\-p\_\{t\}\)^\{\\gamma\}\\,\\mathrm\{BCE\}\(s\_\{ij\},t\_\{ij\}\),\(27\)withα=0\.75\\alpha\\\!=\\\!0\.75\(up\-weight bond edges\) andγ=1\.0\\gamma\\\!=\\\!1\.0\. QM9 uses standard BCE \(γ=0\\gamma\\\!=\\\!0\)\. We foundγ=2\.0\\gamma\\\!=\\\!2\.0caused negative\-logit drift; reducing toγ=1\.0\\gamma\\\!=\\\!1\.0eliminated this without additional regularization\.

#### \(ii\) Atom type lossℒatom\\mathcal\{L\}\_\{\\mathrm\{atom\}\}\.

Cross\-entropy over predicted atom types, averaged over valid atoms:

ℒatom=1\|𝑴node\|​∑i∈𝑴nodeCE​\(Alogits,i,aigt\)\.\\mathcal\{L\}\_\{\\mathrm\{atom\}\}=\\frac\{1\}\{\|\\bm\{M\}\_\{\\mathrm\{node\}\}\|\}\\sum\_\{i\\in\\bm\{M\}\_\{\\mathrm\{node\}\}\}\\mathrm\{CE\}\\\!\\bigl\(A\_\{\\mathrm\{logits\},i\},\\;a\_\{i\}^\{\\mathrm\{gt\}\}\\bigr\)\.\(28\)

#### \(iii\) Bond type lossℒbond\\mathcal\{L\}\_\{\\mathrm\{bond\}\}\.

Conditional cross\-entropy computed only on ground\-truth bond edges within the TopoBridge supportA~\\widetilde\{A\}:

ℒbond=1\|𝑴bond\|​∑\(i,j\)∈𝑴bondCE​\(qi​j,ti​jbond−1\),\\mathcal\{L\}\_\{\\mathrm\{bond\}\}=\\frac\{1\}\{\|\\bm\{M\}\_\{\\mathrm\{bond\}\}\|\}\\sum\_\{\(i,j\)\\in\\bm\{M\}\_\{\\mathrm\{bond\}\}\}\\mathrm\{CE\}\\\!\\bigl\(q\_\{ij\},\\;t\_\{ij\}^\{\\mathrm\{bond\}\}\\\!\-\\\!1\\bigr\),\(29\)where𝑴bond=𝑴pair∧\(ti​jbond\>0\)∧A~\\bm\{M\}\_\{\\mathrm\{bond\}\}=\\bm\{M\}\_\{\\mathrm\{pair\}\}\\wedge\(t\_\{ij\}^\{\\mathrm\{bond\}\}\\\!\>\\\!0\)\\wedge\\widetilde\{A\}, andqi​j∈ℝKq\_\{ij\}\\\!\\in\\\!\\mathbb\{R\}^\{K\}are the bond\-type logits\. The target is shifted by−1\-1because the bond head predicts over\{1,…,K\}\\\{1,\\ldots,K\\\}; bond absence is handled by the adjacency gate\.

#### \(iv\) Coordinate lossℒcoord\\mathcal\{L\}\_\{\\mathrm\{coord\}\}\.

MSE between predicted and ground\-truth coordinates, both zero\-centered per molecule:

ℒcoord=13​\|𝑴node\|​∑i∈𝑴node‖𝒓^i−𝒓igt‖22\.\\mathcal\{L\}\_\{\\mathrm\{coord\}\}=\\frac\{1\}\{3\|\\bm\{M\}\_\{\\mathrm\{node\}\}\|\}\\sum\_\{i\\in\\bm\{M\}\_\{\\mathrm\{node\}\}\}\\\|\\hat\{\\bm\{r\}\}\_\{i\}\-\\bm\{r\}\_\{i\}^\{\\mathrm\{gt\}\}\\\|\_\{2\}^\{2\}\.\(30\)

#### \(v\) Distance lossℒdist\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\.

This term provides rotation\-invariant geometric supervision and is the sum of an all\-pairs distance MSE and a bonded\-pair distance MSE:

ℒdist=1\|𝑴full\|​∑\(i,j\)∈𝑴full\(D^i​j−Di​jgt\)2\+1\|𝑴bonded\|​∑\(i,j\)∈𝑴bonded\(D^i​j−Di​jgt\)2,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}=\\frac\{1\}\{\|\\bm\{M\}\_\{\\mathrm\{full\}\}\|\}\\sum\_\{\(i,j\)\\in\\bm\{M\}\_\{\\mathrm\{full\}\}\}\(\\hat\{D\}\_\{ij\}\-D\_\{ij\}^\{\\mathrm\{gt\}\}\)^\{2\}\\;\+\\;\\frac\{1\}\{\|\\bm\{M\}\_\{\\mathrm\{bonded\}\}\|\}\\sum\_\{\(i,j\)\\in\\bm\{M\}\_\{\\mathrm\{bonded\}\}\}\(\\hat\{D\}\_\{ij\}\-D\_\{ij\}^\{\\mathrm\{gt\}\}\)^\{2\},\(31\)whereDi​j=‖𝒓i−𝒓j‖2D\_\{ij\}\\\!=\\\!\\\|\\bm\{r\}\_\{i\}\-\\bm\{r\}\_\{j\}\\\|\_\{2\}\. The bonded\-pair term focuses geometric learning on bond lengths, where distance errors have the largest chemical impact\. Active on GEOM\-Drugs \(λD=1\\lambda\_\{D\}\\\!=\\\!1\); inactive on QM9 \(λD=0\\lambda\_\{D\}\\\!=\\\!0\)\.

#### \(vi\) KL divergenceℒKL\\mathcal\{L\}\_\{\\mathrm\{KL\}\}\.

The encoder produces per\-atom posteriorsq​\(zi∣M\)=𝒩​\(μi,diag​\(exp⁡\(ℓi\)\)\)q\(z\_\{i\}\\mid M\)=\\mathcal\{N\}\(\\mu\_\{i\},\\mathrm\{diag\}\(\\exp\(\\ell\_\{i\}\)\)\)withℓi=clamp​\(log⁡σi2,−10,4\)\\ell\_\{i\}=\\mathrm\{clamp\}\(\\log\\sigma\_\{i\}^\{2\},\-10,4\):

ℒKL=1B​N​∑b,i\[−12​∑d=1D\(1\+ℓi,d−μi,d2−eℓi,d\)\]\.\\mathcal\{L\}\_\{\\mathrm\{KL\}\}=\\frac\{1\}\{BN\}\\sum\_\{b,i\}\\Bigl\[\-\\tfrac\{1\}\{2\}\\sum\_\{d=1\}^\{D\}\\bigl\(1\+\\ell\_\{i,d\}\-\\mu\_\{i,d\}^\{2\}\-e^\{\\ell\_\{i,d\}\}\\bigr\)\\Bigr\]\.\(32\)The near\-zero weight \(β=5×10−6\\beta\\\!=\\\!5\\\!\\times\\\!10^\{\-6\}on GEOM,10−610^\{\-6\}on QM9\) yields a near\-deterministic autoencoder whose latent space remains structured enough for downstream diffusion\.

#### \(vii\) AGCL lossℒAGCL\\mathcal\{L\}\_\{\\mathrm\{AGCL\}\}\.

This is the advantage\-gated selective teacher loss defined in Sec\.[4\.5](https://arxiv.org/html/2606.13955#S4.SS5), Eqs\. \([22](https://arxiv.org/html/2606.13955#S4.E22)\)–\([23](https://arxiv.org/html/2606.13955#S4.E23)\)\. LetPrawP\_\{\\mathrm\{raw\}\}andPchemP\_\{\\mathrm\{chem\}\}denote the raw decoder and ChemCO\-projected bond distributions, respectively\. For each moleculebbin the batch, the advantageab=\[ℰb​\(Praw\)−ℰb​\(Pchem\)\]\+a\_\{b\}=\[\\mathcal\{E\}\_\{b\}\(P\_\{\\mathrm\{raw\}\}\)\-\\mathcal\{E\}\_\{b\}\(P\_\{\\mathrm\{chem\}\}\)\]\_\{\+\}gates anL2L\_\{2\}consistency loss that steersPrawP\_\{\\mathrm\{raw\}\}towardsg⁡\(Pchem\)\\operatorname\{sg\}\(P\_\{\\mathrm\{chem\}\}\)\. The stop\-gradient detaches the ChemCO output so that only the raw decoder is updated\. When∑bab<10−8\\sum\_\{b\}a\_\{b\}<10^\{\-8\}\(i\.e\., ChemCO provides no advantage\), the loss returns zero\. Weight:λAGCL=0\.1\\lambda\_\{\\mathrm\{AGCL\}\}\\\!=\\\!0\.1\(GEOM\),1\.01\.0\(QM9\)\.

#### Raw bond distribution\.

The decoder constructsPrawP\_\{\\mathrm\{raw\}\}by gating the conditional bond\-type softmax with the STE adjacency:

Praw,i​j=\[1−A~i​jste⏟P​\(no bond\),A~i​jste⋅softmax​\(qi​j\)⏟P​\(type​1​…​K\)\],P\_\{\\mathrm\{raw\},ij\}=\\bigl\[\\underbrace\{1\-\\widetilde\{A\}\_\{ij\}^\{\\mathrm\{ste\}\}\}\_\{P\(\\text\{no bond\}\)\},\\;\\underbrace\{\\widetilde\{A\}\_\{ij\}^\{\\mathrm\{ste\}\}\\cdot\\mathrm\{softmax\}\(q\_\{ij\}\)\}\_\{P\(\\text\{type \}1\\ldots K\)\}\\bigr\],\(33\)whereA~ste\\widetilde\{A\}^\{\\mathrm\{ste\}\}is the straight\-through adjacency from TopoBridge\. Pairs outside the support are forced to\[1,0,…,0\]\[1,0,\\ldots,0\]\(no bond\)\.

#### Loss weight summary\.

Table 25:Active loss weights per dataset\. Zero entries are inactive\.LossQM9GEOM\-Drugsℒadj\\mathcal\{L\}\_\{\\mathrm\{adj\}\}\(adjacency\)1\.010\.0ℒatom\\mathcal\{L\}\_\{\\mathrm\{atom\}\}\(atom type\)1\.01\.0ℒbond\\mathcal\{L\}\_\{\\mathrm\{bond\}\}\(bond type\)1\.01\.0λR​ℒcoord\\lambda\_\{R\}\\,\\mathcal\{L\}\_\{\\mathrm\{coord\}\}\(coordinates\)1\.01\.0λD​ℒdist\\lambda\_\{D\}\\,\\mathcal\{L\}\_\{\\mathrm\{dist\}\}\(distances\)01\.0β​ℒKL\\beta\\,\\mathcal\{L\}\_\{\\mathrm\{KL\}\}\(KL\)1×10−61\\\!\\times\\\!10^\{\-6\}5×10−65\\\!\\times\\\!10^\{\-6\}λAGCL​ℒAGCL\\lambda\_\{\\mathrm\{AGCL\}\}\\,\\mathcal\{L\}\_\{\\mathrm\{AGCL\}\}\(AGCL\)1\.00\.1

### L\.5GT\-Support Warmup Schedule

#### Motivation\.

Early in training, the adjacency predictor produces near\-random graphs\. If the bond\-type and coordinate heads must condition on these noisy adjacencies, they receive inconsistent inputs that slow convergence\. GT\-support warmup decouples this by providing ground\-truth topology initially, then transitioning to predicted TopoBridge supports\.

#### Three\-phase schedule\.

LetWgtW\_\{\\mathrm\{gt\}\}andWtrW\_\{\\mathrm\{tr\}\}be the warmup and transition durations \(in epochs\)\. The probability of using ground\-truth support is:

pgt​\(epoch\)=\{1\.0epoch<Wgt​\(Phase A: pure GT\),1−epoch−WgtWtrWgt≤epoch<Wgt\+Wtr​\(Phase B: linear transition\),0\.0epoch≥Wgt\+Wtr​\(Phase C: fully predicted\)\.p\_\{\\mathrm\{gt\}\}\(\\text\{epoch\}\)=\\begin\{cases\}1\.0&\\text\{epoch\}<W\_\{\\mathrm\{gt\}\}\\;\\;\\text\{\(Phase A: pure GT\)\},\\\\ 1\-\\frac\{\\text\{epoch\}\-W\_\{\\mathrm\{gt\}\}\}\{W\_\{\\mathrm\{tr\}\}\}&W\_\{\\mathrm\{gt\}\}\\leq\\text\{epoch\}<W\_\{\\mathrm\{gt\}\}\+W\_\{\\mathrm\{tr\}\}\\;\\;\\text\{\(Phase B: linear transition\)\},\\\\ 0\.0&\\text\{epoch\}\\geq W\_\{\\mathrm\{gt\}\}\+W\_\{\\mathrm\{tr\}\}\\;\\;\\text\{\(Phase C: fully predicted\)\}\.\\end\{cases\}\(34\)During Phase B, gating is a per\-batch Bernoulli coin flip: with probabilitypgtp\_\{\\mathrm\{gt\}\}, the entire batch receives ground\-truth adjacency and bond types; otherwise the batch uses predicted TopoBridge supports\. Settings:Wgt=100W\_\{\\mathrm\{gt\}\}\\\!=\\\!100,Wtr=100W\_\{\\mathrm\{tr\}\}\\\!=\\\!100on GEOM\-Drugs; QM9 usesWgt=0W\_\{\\mathrm\{gt\}\}\\\!=\\\!0\.

### L\.6Latent Diffusion Model \(Stage 2\)

The diffusion prior is trained on frozen TopVAE latents\. We describe the backbone, noise schedule, and sampling procedure\.

#### DiT backbone\.

A standard TransformerEncoder with sinusoidal timestep embeddings \(base 10000\)\. Timestep embedding is projected through a 2\-layer MLP and added to input token embeddings\.

#### Architecture configurations\.

Table 26:DiT configurations for Stage 2 latent diffusion\.ParameterQM9GEOM DiT\-BHidden dim512768Heads812Layers812MLP ratio4×4\\times4×4\\timesOptimizerAdamW, lr=10−4=10^\{\-4\}, weight decay=0\.05=0\.05LR scheduleLinear warmup \(1k steps\)\+\+cosine decayMax epochs500010000Batch size×\\timesGPUs1024×41024\\\!\\times\\\!4768×4768\\\!\\times\\\!4Precisionfp16\-mixed
#### Noise schedule\.

VP\-SDE with the cosine schedule:α¯\(t\)=cos\(t\+s1\+s⋅π2\)2\\bar\{\\alpha\}\(t\)=\\cos\\\!\\bigl\(\\frac\{t\+s\}\{1\+s\}\\cdot\\frac\{\\pi\}\{2\}\\bigr\)^\{2\}with offsets=0\.008s\\\!=\\\!0\.008\.

#### Latent whitening\.

Before diffusion training, latent codes from the frozen VAE encoder are whitened:znorm=\(z−μz\)/σzz\_\{\\mathrm\{norm\}\}=\(z\-\\mu\_\{z\}\)/\\sigma\_\{z\}, whereμz,σz\\mu\_\{z\},\\sigma\_\{z\}are computed over all training latents\. The inverse transform is applied before decoding at generation time\.

#### Sampling\.

We use ancestral DDPM reverse sampling withT=100T\\\!=\\\!100steps ontarray=linspace​\(1−ϵ,ϵ,T\)t\_\{\\mathrm\{array\}\}=\\mathrm\{linspace\}\(1\\\!\-\\\!\\epsilon,\\epsilon,T\),ϵ=10−3\\epsilon\\\!=\\\!10^\{\-3\}\. Noise temperatureτ=1\.0\\tau\\\!=\\\!1\.0\. No exponential moving average is used\.

#### VAE freezing\.

The VAE is fully frozen during Stage 2: all parameters are detached, and online encoding runs undertorch\.no\_grad\(\)\.

### L\.7Metrics

#### iFID Computation Details\.

Our iFID metric adapts the interpolated FID from pixel space to the latent space of molecular VAEs\[Xuet al\.,[2026](https://arxiv.org/html/2606.13955#bib.bib25)\]\. The pipeline proceeds in four stages: encoding, interpolation, decoding, and scoring\.

Encoding\.Every moleculexxin the training set𝒟train\\mathcal\{D\}\_\{\\text\{train\}\}or validation set𝒟val\\mathcal\{D\}\_\{\\text\{val\}\}is mapped to its deterministic posterior meanμϕ​\(x\)\\mu\_\{\\phi\}\(x\)by a frozen encoder with all augmentations disabled \(no rotation, translation, or noise\)\. Using the mean rather than a reparameterised sample ensures the metric measures the geometry of the learned manifold, not sampling noise\.

Interpolation\.Each dense latentz∈ℝNmax×dz\\in\\mathbb\{R\}^\{N\_\{\\max\}\\times d\}is mean\-pooled along the node dimension toz¯\\bar\{z\}\. For each validation moleculexvx\_\{v\}, we retrieve the top\-K=10K\{=\}10neighbours by cosine similarity from training molecules with the*same*atom count\|x\|=\|xv\|\|x\|\{=\}\|x\_\{v\}\|, avoiding size\-mismatch artefacts\. One neighbourztz\_\{t\}is drawn categorically with weights

πk=softmax​\(−‖zv−zt\(k\)‖22\)\.\\pi\_\{k\}=\\mathrm\{softmax\}\\\!\\bigl\(\-\\\|z\_\{v\}\-z\_\{t\}^\{\(k\)\}\\\|\_\{2\}^\{2\}\\bigr\)\.The midpoint latent is then constructed via spherical linear interpolation:

zα=SLERP​\(zv,zt,α\),α=0\.5,z\_\{\\alpha\}=\\mathrm\{SLERP\}\(z\_\{v\},\\,z\_\{t\},\\,\\alpha\),\\qquad\\alpha=0\.5,computed on the node\-level flattened vectors, falling back to LERP whensin⁡θ<10−6\\sin\\theta<10^\{\-6\}\.

Decoding\.Bothzαz\_\{\\alpha\}and the originalzvz\_\{v\}are decoded by a frozen decoder following each model’s standard protocol, assembled into RDKit molecules, and converted to canonical SMILES \(RemoveAllHs,MolToSmiles\(isomericSmiles=False\)\)\.

Scoring and invalid\-molecule handling\.FID features are 512\-d ChemNet activations \(fcd\_torch\), from which we estimate means and covariances for the reference, interpolated, and reconstructed sets and compute the Fréchet distance:

FID=‖μ1−μ2‖22\+Tr​\(Σ1\+Σ2−2​\(Σ1​Σ2\)1/2\)\.\\mathrm\{FID\}=\\\|\\mu\_\{1\}\-\\mu\_\{2\}\\\|\_\{2\}^\{2\}\+\\mathrm\{Tr\}\\\!\\bigl\(\\Sigma\_\{1\}\+\\Sigma\_\{2\}\-2\(\\Sigma\_\{1\}\\Sigma\_\{2\}\)^\{1/2\}\\bigr\)\.We deliberately do*not*restrict evaluation to valid\-and\-connected subsets\. Decoding failures yield empty strings; a subsequentMolFromSmiles\(sanitize=True\)filter further removes unsanitisable or overly long \(\>\>350 character\) SMILES\. Only surviving molecules enter ChemNet, so invalid samples do not pollute the FID value\.

## Appendix MDetails of ChemCO

This appendix provides the complete optimization details for ChemCO\. We use the same notation as Sec\.[4\.4](https://arxiv.org/html/2606.13955#S4.SS4):A~\\widetilde\{A\}is the TopoBridge adjacency support,Ω\(k\)\\Omega^\{\(k\)\}is the chemical mask for bond typekk,U¯\\bar\{U\}is the relative neural utility \(Eq\.[18](https://arxiv.org/html/2606.13955#S4.E18)\),Φ\\Phidenotes the free primal logits, andY\(k\)=Tk​\(Φ\)Y^\{\(k\)\}\\\!=\\\!T\_\{k\}\(\\Phi\)is the induced continuous bond\-type assignment \(Eq\.[19](https://arxiv.org/html/2606.13955#S4.E19)\)\.

### M\.1Initialization

ChemCO optimizes only real bond typesk=1,…,Kk=1,\\ldots,K; the no\-bond classk=0k\\\!=\\\!0is represented by the residual mass\. The primal logits are initialized from the relative neural utility:

Φ0,i​j\(k\)=U¯i​j\(k\),k=1,…,K\.\\Phi\_\{0,ij\}^\{\(k\)\}=\\bar\{U\}\_\{ij\}^\{\(k\)\},\\qquad k=1,\\ldots,K\.\(35\)This initialization seeds the optimizer at the decoder’s own belief, so that the subsequent unroll only needs to*correct*constraint violations rather than discover the bond structure from scratch\.

#### Continuous statistics\.

At iterationtt, ChemCO computes the soft bond\-type assignment

Yt\(k\)=Tk​\(Φt\),k=1,…,K,Y\_\{t\}^\{\(k\)\}=T\_\{k\}\(\\Phi\_\{t\}\),\\qquad k=1,\\ldots,K,\(36\)using the transform defined in Eq\. \([19](https://arxiv.org/html/2606.13955#S4.E19)\)\. Recall thatTkT\_\{k\}includes masking by the TopoBridge supportA~\\widetilde\{A\}and the chemical maskΩ\(k\)\\Omega^\{\(k\)\}, so thatYt\(k\)Y\_\{t\}^\{\(k\)\}is nonzero only on candidate edges with chemically allowed bond types\. The three continuous statistics introduced in Eq\. \([20](https://arxiv.org/html/2606.13955#S4.E20)\) are then evaluated at iterationtt:

si​j,t=∑k=1KYi​j,t\(k\),vali,t=∑j∈𝒱,j≠i∑k=1Kok​Yi​j,t\(k\),degi,t=∑j∈𝒱,j≠isi​j,t,s\_\{ij,t\}=\\sum\_\{k=1\}^\{K\}Y\_\{ij,t\}^\{\(k\)\},\\qquad\\operatorname\{val\}\_\{i,t\}=\\sum\_\{j\\in\\mathcal\{V\},\\,j\\neq i\}\\sum\_\{k=1\}^\{K\}o\_\{k\}Y\_\{ij,t\}^\{\(k\)\},\\qquad\\deg\_\{i,t\}=\\sum\_\{j\\in\\mathcal\{V\},\\,j\\neq i\}s\_\{ij,t\},\(37\)where the subscriptttmakes the iteration explicit \(the main text omitsttfor conciseness\)\. The constraint residuals \(denotedhhto avoid confusion withg​\(x\)=σ​\(x\)2g\(x\)=\\sigma\(x\)^\{2\}in Eq\.[19](https://arxiv.org/html/2606.13955#S4.E19)\) are

hi​j,tpair=si​j,t−1,hi,tval=vali,t−ci,hi,tdeg=dmin−degi,t\.h\_\{ij,t\}^\{\\mathrm\{pair\}\}=s\_\{ij,t\}\-1,\\qquad h\_\{i,t\}^\{\\mathrm\{val\}\}=\\operatorname\{val\}\_\{i,t\}\-c\_\{i\},\\qquad h\_\{i,t\}^\{\\mathrm\{deg\}\}=d\_\{\\min\}\-\\deg\_\{i,t\}\.\(38\)By convention, a*positive*residual always indicates a violated constraint\.

#### Dual warm\-start\.

Rather than initializing all dual variables at zero, ChemCO warm\-starts them from the initial constraint violations:

μi​j,0=\[si​j,0−1\]\+,λi,0\+=\[vali,0−ci\]\+,λi,0−=\[dmin−degi,0\]\+\.\\mu\_\{ij,0\}=\\bigl\[s\_\{ij,0\}\-1\\bigr\]\_\{\+\},\\qquad\\lambda\_\{i,0\}^\{\+\}=\\bigl\[\\operatorname\{val\}\_\{i,0\}\-c\_\{i\}\\bigr\]\_\{\+\},\\qquad\\lambda\_\{i,0\}^\{\-\}=\\bigl\[d\_\{\\min\}\-\\deg\_\{i,0\}\\bigr\]\_\{\+\}\.\(39\)This gives the solver a nonzero correction signal from the first iteration, improving convergence speed compared to a cold start\. Nonnegativity is guaranteed because\[⋅\]\+≥0\[\\cdot\]\_\{\+\}\\\!\\geq\\\!0\.

### M\.2Chemical MaskΩ\(k\)\\Omega^\{\(k\)\}Construction

The chemical maskΩ\(k\)∈\{0,1\}N×N\\Omega^\{\(k\)\}\\\!\\in\\\!\\\{0,1\\\}^\{N\\times N\}encodes which bond types are chemically allowed for each atom pair\. It is*static*and derived from standard valence\-bond chemistry, not learned or fitted to any dataset\. During training,Ω\(k\)\\Omega^\{\(k\)\}is constructed from the predicted atom\-typearg⁡max\\arg\\max\. SinceΩ\(k\)\\Omega^\{\(k\)\}enters only as a binary mask in the forward pass and is not differentiated through, this discrete operation does not affect gradient computation\. At inference time, ChemCO is removed entirely \(Table[5](https://arxiv.org/html/2606.13955#S5.T5)\), so the mask is not needed\. The mask is constructed bybuild\_atom\_pair\_type\_mask\(\)using the following rules:

Table 27:Allowed bond types per atom\-pair combination inΩ\(k\)\\Omega^\{\(k\)\}\. S = single, D = double, T = triple\.Atom PairAllowed Bond TypesC–CS, D, TC–NS, D, TC–OS, DC–SS, DC–F / C–Cl / C–Br / C–IS onlyN–NS, D, TN–OS, DO–OS onlyF/Cl/Br/I – anyS onlyS–SS, DS–NS, DS–OS, DP–C / P–N / P–O / P–SS, DB–C / B–N / B–OS, DSi–C / Si–N / Si–OS, DBi – anyS onlyFor atom\-pair combinations not listed above, the default fallback allows single bonds only \(Ωi​j\(single\)=1\\Omega^\{\(\\mathrm\{single\}\)\}\_\{ij\}\\\!=\\\!1,Ωi​j\(k\)=0\\Omega^\{\(k\)\}\_\{ij\}\\\!=\\\!0fork≥2k\\\!\\geq\\\!2\)\.

On QM9 \(4 atom types: C, N, O, F\), the default configuration uses only implicit valence caps without explicit pair restrictions \(Ω\(k\)=𝟏\\Omega^\{\(k\)\}\\\!=\\\!\\bm\{1\}\); on GEOM\-Drugs \(12 atom types\), the full mask is active\. The sensitivity ofΩ\(k\)\\Omega^\{\(k\)\}is studied in Appendix[G](https://arxiv.org/html/2606.13955#A7)\(Table[14](https://arxiv.org/html/2606.13955#A7.T14)\): adding the explicit mask to QM9 yields a 19% relative gain in 3D Stab∧\\wedgeConn, confirming that pair\-level restrictions capture chemical knowledge beyond what valence caps alone provide\.

### M\.3Adaptive penalty objective and smoothed active\-set direction

We introduce nonnegative multipliersμi​j≥0\\mu\_\{ij\}\\geq 0for pair exclusivity,λi\+≥0\\lambda\_\{i\}^\{\+\}\\geq 0for valence upper bounds, andλi−≥0\\lambda\_\{i\}^\{\-\}\\geq 0for minimum\-degree constraints\. These correspond to the main\-text multipliersμ\\mu,λ\\lambda, andν\\nu, respectively; the split notationλ\+\\lambda^\{\+\}/λ−\\lambda^\{\-\}makes the sign convention explicit\.

At iterationtt, ChemCO uses a fixed\-dual adaptive penalty score

𝒥t​\(Φ\)\\displaystyle\\mathcal\{J\}\_\{t\}\(\\Phi\)=∑i<ji,j∈𝒱∑k=1KU¯i​j\(k\)​Yi​j\(k\)−ρpair​∑i<ji,j∈𝒱μi​j,t​ψγ​\(hi​j,tpair\)\\displaystyle=\\sum\_\{\\begin\{subarray\}\{c\}i<j\\\\ i,j\\in\\mathcal\{V\}\\end\{subarray\}\}\\sum\_\{k=1\}^\{K\}\\bar\{U\}\_\{ij\}^\{\(k\)\}Y\_\{ij\}^\{\(k\)\}\-\\rho\_\{\\mathrm\{pair\}\}\\sum\_\{\\begin\{subarray\}\{c\}i<j\\\\ i,j\\in\\mathcal\{V\}\\end\{subarray\}\}\\mu\_\{ij,t\}\\psi\_\{\\gamma\}\(h\_\{ij,t\}^\{\\mathrm\{pair\}\}\)\(40\)−∑i∈𝒱λi,t\+​ψγ​\(hi,tval\)−∑i∈𝒱λi,t−​ψγ​\(hi,tdeg\),\\displaystyle\\quad\-\\sum\_\{i\\in\\mathcal\{V\}\}\\lambda\_\{i,t\}^\{\+\}\\psi\_\{\\gamma\}\(h\_\{i,t\}^\{\\mathrm\{val\}\}\)\-\\sum\_\{i\\in\\mathcal\{V\}\}\\lambda\_\{i,t\}^\{\-\}\\psi\_\{\\gamma\}\(h\_\{i,t\}^\{\\mathrm\{deg\}\}\),where

hi​j,tpair=si​j,t−1,hi,tval=vali,t−ci,hi,tdeg=dmin−degi,t\.h\_\{ij,t\}^\{\\mathrm\{pair\}\}=s\_\{ij,t\}\-1,\\qquad h\_\{i,t\}^\{\\mathrm\{val\}\}=\\operatorname\{val\}\_\{i,t\}\-c\_\{i\},\\qquad h\_\{i,t\}^\{\\mathrm\{deg\}\}=d\_\{\\min\}\-\\deg\_\{i,t\}\.\(41\)Hereψγ​\(r\)=γ−1​log⁡\(1\+exp⁡\(γ​r\)\)\\psi\_\{\\gamma\}\(r\)=\\gamma^\{\-1\}\\log\(1\+\\exp\(\\gamma r\)\)is a smooth approximation to\[r\]\+\[r\]\_\{\+\}, and

ψγ′​\(r\)=σ​\(γ​r\)\.\\psi\_\{\\gamma\}^\{\\prime\}\(r\)=\\sigma\(\\gamma r\)\.\(42\)Thus the smoothed active\-set gates are

ωi​j,tpair=σ​\(γ​hi​j,tpair\),ωi,tval=σ​\(γ​hi,tval\),ωi,tdeg=σ​\(γ​hi,tdeg\)\.\\omega\_\{ij,t\}^\{\\mathrm\{pair\}\}=\\sigma\(\\gamma h\_\{ij,t\}^\{\\mathrm\{pair\}\}\),\\qquad\\omega\_\{i,t\}^\{\\mathrm\{val\}\}=\\sigma\(\\gamma h\_\{i,t\}^\{\\mathrm\{val\}\}\),\\qquad\\omega\_\{i,t\}^\{\\mathrm\{deg\}\}=\\sigma\(\\gamma h\_\{i,t\}^\{\\mathrm\{deg\}\}\)\.\(43\)Whenγ\\gammais large, these gates approach binary indicators of violated constraints; for finiteγ\\gamma, they provide a smooth active\-set approximation\.

The induced smoothed direction with respect to the continuous assignment massYi​j\(k\)Y\_\{ij\}^\{\(k\)\}is

G~i​j,t\(k\)\\displaystyle\\widetilde\{G\}\_\{ij,t\}^\{\(k\)\}=U¯i​j\(k\)−ρpair​ωi​j,tpair​μi​j,t\\displaystyle=\\bar\{U\}\_\{ij\}^\{\(k\)\}\-\\rho\_\{\\mathrm\{pair\}\}\\omega\_\{ij,t\}^\{\\mathrm\{pair\}\}\\mu\_\{ij,t\}\(44\)−ok​\(ωi,tval​λi,t\+\+ωj,tval​λj,t\+\)\+ωi,tdeg​λi,t−\+ωj,tdeg​λj,t−\.\\displaystyle\\quad\-o\_\{k\}\\\!\\left\(\\omega\_\{i,t\}^\{\\mathrm\{val\}\}\\lambda\_\{i,t\}^\{\+\}\+\\omega\_\{j,t\}^\{\\mathrm\{val\}\}\\lambda\_\{j,t\}^\{\+\}\\right\)\+\\omega\_\{i,t\}^\{\\mathrm\{deg\}\}\\lambda\_\{i,t\}^\{\-\}\+\\omega\_\{j,t\}^\{\\mathrm\{deg\}\}\\lambda\_\{j,t\}^\{\-\}\.The implementation usesρpair=2\\rho\_\{\\mathrm\{pair\}\}=2, matching the row\-wise pair penalty followed by explicit symmetrization\. Equivalently, this constant can be absorbed into the pair multiplier scale\.

### M\.4Row\-centric direction and direct logit\-space update

For efficient batched implementation, ChemCO constructs the direction in two stages\. First, it forms a row\-centric direction that contains atomii’s node\-level multiplier contributions:

G^i​j,t\(k\)=12​U¯i​j\(k\)−ρpair2​ωi​j,tpair​μi​j,t−ok​ωi,tval​λi,t\+\+ωi,tdeg​λi,t−\.\\widehat\{G\}\_\{ij,t\}^\{\(k\)\}=\\tfrac\{1\}\{2\}\\bar\{U\}\_\{ij\}^\{\(k\)\}\-\\tfrac\{\\rho\_\{\\mathrm\{pair\}\}\}\{2\}\\omega\_\{ij,t\}^\{\\mathrm\{pair\}\}\\mu\_\{ij,t\}\-o\_\{k\}\\omega\_\{i,t\}^\{\\mathrm\{val\}\}\\lambda\_\{i,t\}^\{\+\}\+\\omega\_\{i,t\}^\{\\mathrm\{deg\}\}\\lambda\_\{i,t\}^\{\-\}\.\(45\)The symmetrized direction is then

G~i​j,t\(k\)=G^i​j,t\(k\)\+G^j​i,t\(k\)\.\\widetilde\{G\}\_\{ij,t\}^\{\(k\)\}=\\widehat\{G\}\_\{ij,t\}^\{\(k\)\}\+\\widehat\{G\}\_\{ji,t\}^\{\(k\)\}\.\(46\)SinceU¯\\bar\{U\}andμ\\muare symmetric, the utility terms sum toU¯i​j\(k\)\\bar\{U\}\_\{ij\}^\{\(k\)\}, and the pair\-exclusivity terms sum to−ρpair​ωi​j,tpair​μi​j,t\-\\rho\_\{\\mathrm\{pair\}\}\\omega\_\{ij,t\}^\{\\mathrm\{pair\}\}\\mu\_\{ij,t\}\. The node\-level terms collect the valence and degree contributions from both endpoints, yielding Eq\. \([44](https://arxiv.org/html/2606.13955#A13.E44)\)\.

#### Direct logit\-space heuristic update\.

ChemCO applies this direction directly to the free logits:

Φt\+1\(k\)=Φt\(k\)\+ηΦ\(t\)​G~t\(k\)⊙Ω\(k\),k=1,…,K,\\Phi\_\{t\+1\}^\{\(k\)\}=\\Phi\_\{t\}^\{\(k\)\}\+\\eta\_\{\\Phi\}^\{\(t\)\}\\widetilde\{G\}\_\{t\}^\{\(k\)\}\\odot\\Omega^\{\(k\)\},\\qquad k=1,\\ldots,K,\(47\)whereηΦ\(t\)=α​γαt\\eta\_\{\\Phi\}^\{\(t\)\}=\\alpha\\gamma\_\{\\alpha\}^\{t\}\. The TopoBridge supportA~\\widetilde\{A\}is enforced by the transformTk​\(Φ\)T\_\{k\}\(\\Phi\), and the chemical maskΩ\(k\)\\Omega^\{\(k\)\}zeros out forbidden atom–bond combinations\.

This update is not the exact chain\-rule gradient of Eq\. \([40](https://arxiv.org/html/2606.13955#A13.E40)\) with respect toΦ\\Phi\. A true chain\-rule gradient would include the derivative ofg​\(Φ\)=σ​\(Φ\)2g\(\\Phi\)=\\sigma\(\\Phi\)^\{2\}, which contains a factor proportional toσ​\(Φ\)2​\(1−σ​\(Φ\)\)\\sigma\(\\Phi\)^\{2\}\(1\-\\sigma\(\\Phi\)\)and can vanish when a currently absent bond hasΦ≪0\\Phi\\ll 0\. ChemCO instead uses the smoothedYY\-space direction directly in logit space, allowing absent bonds to be activated by the unrolled correction\.

#### Direct logit\-space heuristic update\.

The primal update addsG~t\\widetilde\{G\}\_\{t\}directly to the logits:

Φt\+1\(k\)=Φt\(k\)\+ηΦ\(t\)​G~t\(k\)⊙Ω\(k\),k=1,…,K,\\Phi\_\{t\+1\}^\{\(k\)\}=\\Phi\_\{t\}^\{\(k\)\}\+\\eta\_\{\\Phi\}^\{\(t\)\}\\;\\widetilde\{G\}\_\{t\}^\{\(k\)\}\\odot\\Omega^\{\(k\)\},\\qquad k=1,\\ldots,K,\(48\)where the masking byΩ\(k\)\\Omega^\{\(k\)\}zeros out updates on chemically forbidden pairs, andηΦ\(t\)=α​γαt\\eta\_\{\\Phi\}^\{\(t\)\}=\\alpha\\,\\gamma\_\{\\alpha\}^\{t\}is an exponentially decayed step size with base rateα\\alphaand decay factorγα∈\(0,1\]\\gamma\_\{\\alpha\}\\\!\\in\\\!\(0,1\]\.

#### Why a direct logit\-space update instead of a chain\-rule gradient?

A standard approach would chain\-ruleG~\\widetilde\{G\}through the transformTkT\_\{k\}\(Eq\.[19](https://arxiv.org/html/2606.13955#S4.E19)\), yielding a logit\-space gradient proportional toσ​\(Φ\)2​\(1−σ​\(Φ\)\)⊙G~\\sigma\(\\Phi\)^\{2\}\(1\\\!\-\\\!\\sigma\(\\Phi\)\)\\odot\\widetilde\{G\}\. However, theσ2​\(1−σ\)\\sigma^\{2\}\(1\\\!\-\\\!\\sigma\)factor suppresses the gradient whenσ​\(Φ\)→0\\sigma\(\\Phi\)\\\!\\to\\\!0, making it difficult for currently\-absent bonds \(Φ≪0\\Phi\\\!\\ll\\\!0\) to become active\. The direct logit\-space heuristic bypasses this suppression, providing uniform update magnitude across all logit values and allowing bonds to transition smoothly from absent to present\. In practice this improves convergence and is essential for effective end\-to\-end differentiable training through the unrolled solver\. We note that this update is not the true gradient of the penalty objective with respect toΦ\\Phi; it is a heuristic logit\-space update that uses theYY\-space direction directly in logit space\.

After the logit update, the assignment is recomputed viaYt\+1\(k\)=Tk​\(Φt\+1\)Y\_\{t\+1\}^\{\(k\)\}=T\_\{k\}\(\\Phi\_\{t\+1\}\), and the continuous statisticsss,val\\operatorname\{val\},deg\\degare refreshed\.

### M\.5Dual updates

After the logit update, ChemCO recomputesYt\+1\(k\)=Tk​\(Φt\+1\)Y\_\{t\+1\}^\{\(k\)\}=T\_\{k\}\(\\Phi\_\{t\+1\}\)and refreshesst\+1s\_\{t\+1\},valt\+1\\operatorname\{val\}\_\{t\+1\}, anddegt\+1\\deg\_\{t\+1\}\. The multipliers are then updated by accumulating positive residuals:

μi​j,t\+1=μi​j,t\+ημ\(t\)​\[si​j,t\+1−1\]\+,μt\+1←12​\(μt\+1\+μt\+1⊤\),\\mu\_\{ij,t\+1\}=\\mu\_\{ij,t\}\+\\eta\_\{\\mu\}^\{\(t\)\}\\bigl\[s\_\{ij,t\+1\}\-1\\bigr\]\_\{\+\},\\qquad\\mu\_\{t\+1\}\\leftarrow\\tfrac\{1\}\{2\}\(\\mu\_\{t\+1\}\+\\mu\_\{t\+1\}^\{\\top\}\),\(49\)λi,t\+1\+=λi,t\+\+η\+\(t\)​\[vali,t\+1−ci\]\+,λi,t\+1−=λi,t−\+η−\(t\)​\[dmin−degi,t\+1\]\+\.\\lambda\_\{i,t\+1\}^\{\+\}=\\lambda\_\{i,t\}^\{\+\}\+\\eta\_\{\+\}^\{\(t\)\}\\bigl\[\\operatorname\{val\}\_\{i,t\+1\}\-c\_\{i\}\\bigr\]\_\{\+\},\\qquad\\lambda\_\{i,t\+1\}^\{\-\}=\\lambda\_\{i,t\}^\{\-\}\+\\eta\_\{\-\}^\{\(t\)\}\\bigl\[d\_\{\\min\}\-\\deg\_\{i,t\+1\}\\bigr\]\_\{\+\}\.\(50\)Hereημ\(t\)=ημ​γμt\\eta\_\{\\mu\}^\{\(t\)\}=\\eta\_\{\\mu\}\\gamma\_\{\\mu\}^\{t\},η\+\(t\)=η\+​γ\+t\\eta\_\{\+\}^\{\(t\)\}=\\eta\_\{\+\}\\gamma\_\{\+\}^\{t\}, andη−\(t\)=η−​γ−t\\eta\_\{\-\}^\{\(t\)\}=\\eta\_\{\-\}\\gamma\_\{\-\}^\{t\}are exponentially decayed dual step sizes\.

#### Nonnegativity invariant\.

The warm\-start initialization is nonnegative, and every update adds a nonnegative positive\-residual term\. Thereforeμi​j,t≥0\\mu\_\{ij,t\}\\geq 0,λi,t\+≥0\\lambda\_\{i,t\}^\{\+\}\\geq 0, andλi,t−≥0\\lambda\_\{i,t\}^\{\-\}\\geq 0for all iterationstt\.

### M\.6Complete algorithm

Algorithm[1](https://arxiv.org/html/2606.13955#alg1)summarizes the implemented ChemCO core\. It contains the continuous logit\-space unroll and positive\-residual dual accumulation\.

Algorithm 1ChemCO: implemented unrolled logit\-space correction0:Raw bond logits

U∈ℝN×N×\(1\+K\)U\\in\\mathbb\{R\}^\{N\\times N\\times\(1\+K\)\}; node mask

mm; TopoBridge support

A~\\widetilde\{A\}; chemical masks

\{Ω\(k\)\}k=1K\\\{\\Omega^\{\(k\)\}\\\}\_\{k=1\}^\{K\}; valence caps

\{ci\}\\\{c\_\{i\}\\\}; minimum degree

dmind\_\{\\min\}; iterations

TCOT\_\{\\mathrm\{CO\}\}\.

0:Projected bond distribution

P∈\[0,1\]N×N×\(1\+K\)P\\in\[0,1\]^\{N\\times N\\times\(1\+K\)\}\.

1:Build valid\-pair mask

MMfrom

mmand effective support

Me←M⊙A~M\_\{e\}\\leftarrow M\\odot\\widetilde\{A\}\.

2:Symmetrize logits:

U←12​\(U\+U⊤\)U\\leftarrow\\tfrac\{1\}\{2\}\(U\+U^\{\\top\}\)\.

3:Compute relative utility:

U¯\(k\)←U\(k\)−U\(0\)\\bar\{U\}^\{\(k\)\}\\leftarrow U^\{\(k\)\}\-U^\{\(0\)\},

k=1,…,Kk=1,\\ldots,K\.

4:Initialize logits:

Φ0\(k\)←U¯\(k\)\\Phi\_\{0\}^\{\(k\)\}\\leftarrow\\bar\{U\}^\{\(k\)\}\.

5:Set

a0\(k\)←σ​\(Φ0\(k\)\)⊙Me⊙Ω\(k\)a\_\{0\}^\{\(k\)\}\\leftarrow\\sigma\(\\Phi\_\{0\}^\{\(k\)\}\)\\odot M\_\{e\}\\odot\\Omega^\{\(k\)\}\.

6:Compute

Y0\(k\)←Tk​\(a0\)Y\_\{0\}^\{\(k\)\}\\leftarrow T\_\{k\}\(a\_\{0\}\), then

s0s\_\{0\},

val0\\operatorname\{val\}\_\{0\}, and

deg0\\deg\_\{0\}\.

7:Warm\-start duals:

μ0←\[s0−1\]\+\\mu\_\{0\}\\leftarrow\[s\_\{0\}\-1\]\_\{\+\},

λ0\+←\[val0−c\]\+\\lambda\_\{0\}^\{\+\}\\leftarrow\[\\operatorname\{val\}\_\{0\}\-c\]\_\{\+\}, and

λ0−←\[dmin−deg0\]\+\\lambda\_\{0\}^\{\-\}\\leftarrow\[d\_\{\\min\}\-\\deg\_\{0\}\]\_\{\+\}\.

8:for

t=0,…,TCO−1t=0,\\ldots,T\_\{\\mathrm\{CO\}\}\-1do

9:Compute gates

ωtpair\\omega\_\{t\}^\{\\mathrm\{pair\}\},

ωtval\\omega\_\{t\}^\{\\mathrm\{val\}\},

ωtdeg\\omega\_\{t\}^\{\\mathrm\{deg\}\}using Eq\. \([43](https://arxiv.org/html/2606.13955#A13.E43)\)\.

10:Compute row\-centric direction

G^t\(k\)\\widehat\{G\}\_\{t\}^\{\(k\)\}using Eq\. \([45](https://arxiv.org/html/2606.13955#A13.E45)\)\.

11:Symmetrize:

G~t\(k\)←G^t\(k\)\+\(G^t\(k\)\)⊤\\widetilde\{G\}\_\{t\}^\{\(k\)\}\\leftarrow\\widehat\{G\}\_\{t\}^\{\(k\)\}\+\(\\widehat\{G\}\_\{t\}^\{\(k\)\}\)^\{\\top\}\.

12:Update logits:

Φt\+1\(k\)←Φt\(k\)\+α​γαt​G~t\(k\)⊙Ω\(k\)\\Phi\_\{t\+1\}^\{\(k\)\}\\leftarrow\\Phi\_\{t\}^\{\(k\)\}\+\\alpha\\gamma\_\{\\alpha\}^\{t\}\\widetilde\{G\}\_\{t\}^\{\(k\)\}\\odot\\Omega^\{\(k\)\}\.

13:Recompute

at\+1\(k\)←σ​\(Φt\+1\(k\)\)⊙Me⊙Ω\(k\)a\_\{t\+1\}^\{\(k\)\}\\leftarrow\\sigma\(\\Phi\_\{t\+1\}^\{\(k\)\}\)\\odot M\_\{e\}\\odot\\Omega^\{\(k\)\},

Yt\+1\(k\)←Tk​\(at\+1\)Y\_\{t\+1\}^\{\(k\)\}\\leftarrow T\_\{k\}\(a\_\{t\+1\}\), and refresh

st\+1s\_\{t\+1\},

valt\+1\\operatorname\{val\}\_\{t\+1\},

degt\+1\\deg\_\{t\+1\}\.

14:Update

μt\+1\\mu\_\{t\+1\}using Eq\. \([49](https://arxiv.org/html/2606.13955#A13.E49)\)\.

15:Update

λt\+1\+\\lambda\_\{t\+1\}^\{\+\}and

λt\+1−\\lambda\_\{t\+1\}^\{\-\}using Eq\. \([50](https://arxiv.org/html/2606.13955#A13.E50)\)\.

16:endfor

17:Assemble soft distribution:

P\(k\)←YTCO\(k\)P^\{\(k\)\}\\leftarrow Y\_\{T\_\{\\mathrm\{CO\}\}\}^\{\(k\)\}for

k=1,…,Kk=1,\\ldots,K\.

18:Set no\-bond mass:

P\(0\)←clamp⁡\(1−∑k=1KP\(k\),0,1\)P^\{\(0\)\}\\leftarrow\\operatorname\{clamp\}\(1\-\\sum\_\{k=1\}^\{K\}P^\{\(k\)\},0,1\)\.

19:For invalid pairs, set

Pi​j=\[1,0,…,0\]P\_\{ij\}=\[1,0,\\ldots,0\]\.

20:return

PP\.

### M\.7Full bond distribution and hard decoding

AfterTCOT\_\{\\text\{CO\}\}iterations, ChemCO returns the optimized assignmentY^\(k\)=Tk​\(Φ^\)\\widehat\{Y\}^\{\(k\)\}\\\!=\\\!T\_\{k\}\(\\widehat\{\\Phi\}\)\. The full soft bond distribution over\{0,1,…,K\}\\\{0,1,\\ldots,K\\\}assigns no\-bond mass as the clamped residual:

Pi​j\(0\)=clamp⁡\(1−∑k=1KY^i​j\(k\),0,1\),Pi​j\(k\)=Y^i​j\(k\),k=1,…,K\.P\_\{ij\}^\{\(0\)\}=\\operatorname\{clamp\}\\\!\\Bigl\(1\-\\sum\_\{k=1\}^\{K\}\\widehat\{Y\}\_\{ij\}^\{\(k\)\},\\;0,\\;1\\Bigr\),\\qquad P\_\{ij\}^\{\(k\)\}=\\widehat\{Y\}\_\{ij\}^\{\(k\)\},\\quad k=1,\\ldots,K\.\(51\)The clamping ensures numerical stability when the pair exclusivity constraintsi​j≤1s\_\{ij\}\\\!\\leq\\\!1is not yet fully satisfied after a finite unroll\.

The discrete bond type is selected by

B^i​j=arg​maxk∈\{0,1,…,K\}⁡Pi​j\(k\)\.\\widehat\{B\}\_\{ij\}=\\operatorname\*\{arg\\,max\}\_\{k\\in\\\{0,1,\\ldots,K\\\}\}P\_\{ij\}^\{\(k\)\}\.\(52\)During training, gradients are propagated through the full unrolled soft projectionΦ0↦ΦTCO↦Y^\\Phi\_\{0\}\\\!\\mapsto\\\!\\Phi\_\{T\_\{\\mathrm\{CO\}\}\}\\\!\\mapsto\\\!\\widehat\{Y\}\. Hard decoding \(Eq\.[52](https://arxiv.org/html/2606.13955#A13.E52)\) is used only for discrete molecule construction or straight\-through variants, and is not on the gradient path\.

### M\.8Discrete valence refinement

ChemCO’s continuous solver enforces one\-sided valence inequalities \(vali≤ci\\operatorname\{val\}\_\{i\}\\\!\\leq\\\!c\_\{i\}\), which is the natural formulation in continuous space\. However, thearg⁡max\\arg\\maxhard decoding \(Eq\.[52](https://arxiv.org/html/2606.13955#A13.E52)\) can introduce discrete valence violations that the continuous solver cannot resolve, because its soft assignment landscape may contain local minima where the inequality is satisfied but the rounded solution is infeasible\[Vazirani,[2001](https://arxiv.org/html/2606.13955#bib.bib42)\]\. To escape such local minima, ChemCO incorporates a discrete valence refinement step as part of its optimization loop\.

#### Hybrid continuous–discrete loop\.

As shown in Algorithm[1](https://arxiv.org/html/2606.13955#alg1), after each continuous primal–dual unroll producesB^\\widehat\{B\}via hard decoding, ChemCO checks whether any atomiihas discrete valencevi=∑joB^i​jv\_\{i\}=\\sum\_\{j\}o\_\{\\widehat\{B\}\_\{ij\}\}exceeding its capacitycic\_\{i\}\. If violations exist, a greedy discrete repair is applied toB^\\widehat\{B\}, and the repaired bond graph is re\-injected as initialization for a subsequent continuous unroll\. This continuous→\\todiscrete→\\tocontinuous cycle repeats until no discrete violations remain or the maximum number of outer iterationsTouterT\_\{\\mathrm\{outer\}\}is reached\. During training, gradients flow only through the final continuous unroll; the discrete repair steps are non\-differentiable\.

#### Greedy bond\-order reduction\.

The discrete repair operates on the rounded bond graphB^\\widehat\{B\}by iteratively resolving overvalent atoms:

1. 1\.Scan\.Identify all atoms withvi\>civ\_\{i\}\>c\_\{i\}and sort by decreasing violationvi−civ\_\{i\}\-c\_\{i\}\.
2. 2\.Prioritize\.For the most\-violated atom, rank its incident bonds by a priority score that favors: \(a\) bonds to other overvalent atoms \(resolving mutual violations\), \(b\) higher\-order bonds \(downgrading triple→\\todouble or double→\\tosingle removes more excess valence per edit\), while protecting terminal atoms \(degree\-11neighbors are deprioritized\)\.
3. 3\.Repair\.Apply the highest\-priority feasible action: reduce bond order by one, or remove a single bond if the neighbor retains at least one other bond\.
4. 4\.Iterate\.Repeat until no violations remain or a maximum iteration count is reached\.

The discrete repair acts as a structured perturbation that moves the solution out of a continuous local minimum\. In practice, the continuous solver already produces near\-feasible solutions, so the discrete phase typically requires only a few bond\-order edits per molecule\.

### M\.9Bridge\-edge protection during optimization

ChemCO optimizes bond\-type assignments subject to valence and degree constraints, which may drive the soft assignment mass on certain edges toward zero—effectively removing them from the molecular graph\. If such an edge happens to be a*bridge*\(an edge whose removal disconnects the graph\), the connectivity established by TopoBridge would be destroyed\.

#### Bridge\-edge set\.

Letℬ⊆ℰ\\mathcal\{B\}\\\!\\subseteq\\\!\\mathcal\{E\}denote the set of*graph\-theoretic*bridge edges in the current molecular graph, identified by running Tarjan’s bridge\-finding algorithm on the TopoBridge output\. Note that these are distinct from the edges*inserted*by TopoBridge’s BFS reconnection: an inserted edge may or may not be a graph\-theoretic bridge, and a bridge may be an original predicted edge rather than an inserted one\.

#### Post\-processing mask\.

After computing the bond distributionPPvia Eq\. \([51](https://arxiv.org/html/2606.13955#A13.E51)\), bridge\-edge protection is applied as a post\-processing step: for each bridge edge\(i,j\)∈ℬ\(i,j\)\\\!\\in\\\!\\mathcal\{B\}, the no\-bond probability is set to zero and the remaining bond\-type probabilities are renormalized:

Pi​j\(0\)←0,Pi​j\(k\)←Y^i​j\(k\)∑k′=1KY^i​j\(k′\)\+ϵ,\(i,j\)∈ℬ,k=1,…,K\.P\_\{ij\}^\{\(0\)\}\\;\\leftarrow\\;0,\\qquad P\_\{ij\}^\{\(k\)\}\\;\\leftarrow\\;\\frac\{\\widehat\{Y\}\_\{ij\}^\{\(k\)\}\}\{\\sum\_\{k^\{\\prime\}=1\}^\{K\}\\widehat\{Y\}\_\{ij\}^\{\(k^\{\\prime\}\)\}\+\\epsilon\},\\qquad\(i,j\)\\in\\mathcal\{B\},\\;\\;k=1,\\ldots,K\.\(53\)This guaranteesPi​j\(0\)\+∑kPi​j\(k\)=1P\_\{ij\}^\{\(0\)\}\+\\sum\_\{k\}P\_\{ij\}^\{\(k\)\}=1for bridge edges and ensures that they always retain a real bond type \(k≥1k\\\!\\geq\\\!1\) throughout ChemCO’s optimization\. For non\-bridge edges the distribution is unchanged\.

This mechanism complements the valence and degree constraints: while ChemCO freely adjusts bond types and removes redundant edges to satisfy chemical constraints, it is prevented from disrupting the connected topology guaranteed by TopoBridge\. The bridge mask thus propagates Stage 2’s connectivity guarantee through Stage 3’s bond\-type optimization, ensuring the final discrete molecular graph remains connected\.

#### Relationship to adjacency support\.

The TopoBridge adjacency supportA~\\widetilde\{A\}defines the set of*candidate*edges on which ChemCO may place bonds\. It does not guarantee that every candidate edge will carry a real bond in the final output: ChemCO may assignPi​j\(0\)≈1P\_\{ij\}^\{\(0\)\}\\\!\\approx\\\!1\(no bond\) to non\-bridge candidate edges when doing so satisfies valence constraints\. Only bridge edges inℬ\\mathcal\{B\}are protected from removal\.

### M\.10Notation correspondence

For clarity we summarize the correspondence between the compact main\-text notation and the iteration\-explicit appendix notation:

Main textAppendixDescriptionsi​js\_\{ij\}si​j,ts\_\{ij,t\}pair\-level total assignmentvali\\operatorname\{val\}\_\{i\}vali,t\\operatorname\{val\}\_\{i,t\}soft valence at atomiidegi\\deg\_\{i\}degi,t\\deg\_\{i,t\}soft degree at atomiiλi\\lambda\_\{i\}\(valence\)λi,t\+\\lambda\_\{i,t\}^\{\+\}valence cap dual variableνi\\nu\_\{i\}\(degree\)λi,t−\\lambda\_\{i,t\}^\{\-\}minimum\-degree dual variableμi​j\\mu\_\{ij\}μi​j,t\\mu\_\{ij,t\}pair exclusivity dual variableηΦ\\eta\_\{\\Phi\}α​γαt\\alpha\\gamma\_\{\\alpha\}^\{t\}decayed primal step sizeημ\\eta\_\{\\mu\},η\+\\eta\_\{\+\},η−\\eta\_\{\-\}η∙​γ∙t\\eta\_\{\\bullet\}\\gamma\_\{\\bullet\}^\{t\}decayed dual step sizes
The main text omits the iteration subscriptttand the step\-size decay for conciseness, and usesνi\\nu\_\{i\}in place ofλi−\\lambda\_\{i\}^\{\-\}to distinguish it from the valence dual; the two notations are interchangeable\.

## NeurIPS Paper Checklist

1. 1\.Claims
2. Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
3. Answer:\[Yes\]
4. Justification: The abstract and introduction state the main contributions: dark\-area diagnosis, TopVAE with TopoBridge/ChemCO/AGCL, and improved generation and scaffold\-inpainting performance\.
5. Guidelines: - •The answer\[N/A\]means that the abstract and introduction do not include the claims made in the paper\. - •The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations\. A\[No\]or\[N/A\]answer to this question will not be perceived well by the reviewers\. - •The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings\. - •It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper\.
6. 2\.Limitations
7. Question: Does the paper discuss the limitations of the work performed by the authors?
8. Answer:\[Yes\]
9. Justification: Paper discusses limitations that future extensions to physical constraints and conditional molecular generation\.
10. Guidelines: - •The answer\[N/A\]means that the paper has no limitation while the answer\[No\]means that the paper has limitations, but those are not discussed in the paper\. - •The authors are encouraged to create a separate “Limitations” section in their paper\. - •The paper should point out any strong assumptions and how robust the results are to violations of these assumptions \(e\.g\., independence assumptions, noiseless settings, model well\-specification, asymptotic approximations only holding locally\)\. The authors should reflect on how these assumptions might be violated in practice and what the implications would be\. - •The authors should reflect on the scope of the claims made, e\.g\., if the approach was only tested on a few datasets or with a few runs\. In general, empirical results often depend on implicit assumptions, which should be articulated\. - •The authors should reflect on the factors that influence the performance of the approach\. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting\. Or a speech\-to\-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon\. - •The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size\. - •If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness\. - •While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper\. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community\. Reviewers will be specifically instructed to not penalize honesty concerning limitations\.
11. 3\.Theory assumptions and proofs
12. Question: For each theoretical result, does the paper provide the full set of assumptions and a complete \(and correct\) proof?
13. Answer:\[N/A\]
14. Justification: The paper does not present standalone theorem/lemma statements\.
15. Guidelines: - •The answer\[N/A\]means that the paper does not include theoretical results\. - •All the theorems, formulas, and proofs in the paper should be numbered and cross\-referenced\. - •All assumptions should be clearly stated or referenced in the statement of any theorems\. - •The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition\. - •Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material\. - •Theorems and Lemmas that the proof relies upon should be properly referenced\.
16. 4\.Experimental result reproducibility
17. Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper \(regardless of whether the code and data are provided or not\)?
18. Answer:\[Yes\]
19. Justification: The paper discloses the datasets, metrics, perturbation protocol, generation and inpainting settings, model architecture, training hyperparameters, loss terms, diffusion setup, and ChemCO update rules\. These details provide a reproducible specification of the main experimental pipeline independent of code release\.
20. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •If the paper includes experiments, a\[No\]answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not\. - •If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable\. - •Depending on the contribution, reproducibility can be accomplished in various ways\. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model\. In general\. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model \(e\.g\., in the case of a large language model\), releasing of a model checkpoint, or other means that are appropriate to the research performed\. - •While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution\. For example 1. \(a\)If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm\. 2. \(b\)If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully\. 3. \(c\)If the contribution is a new model \(e\.g\., a large language model\), then there should either be a way to access this model for reproducing the results or a way to reproduce the model \(e\.g\., with an open\-source dataset or instructions for how to construct the dataset\)\. 4. \(d\)We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility\. In the case of closed\-source models, it may be that access to the model is limited in some way \(e\.g\., to registered users\), but it should be possible for other researchers to have some path to reproducing or verifying the results\.
21. 5\.Open access to data and code
22. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
23. Answer:\[No\]
24. Justification: The experiments use public datasets and the paper provides detailed implementation and training information, but the current submission does not provide an anonymized code release, execution commands, or a documented artifact for reproducing the main results\. The authors plan to release the code and related artifacts after paper acceptance\.
25. Guidelines: - •The answer\[N/A\]means that paper does not include experiments requiring code\. - • - •While we encourage the release of code and data, we understand that this might not be possible, so\[No\]is an acceptable answer\. Papers cannot be rejected simply for not including code, unless this is central to the contribution \(e\.g\., for a new open\-source benchmark\)\. - •The instructions should contain the exact command and environment needed to run to reproduce the results\. See the NeurIPS code and data submission guidelines \([https://neurips\.cc/public/guides/CodeSubmissionPolicy](https://neurips.cc/public/guides/CodeSubmissionPolicy)\) for more details\. - •The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc\. - •The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines\. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why\. - •At submission time, to preserve anonymity, the authors should release anonymized versions \(if applicable\)\. - •Providing as much information as possible in supplemental material \(appended to the paper\) is recommended, but including URLs to data and code is permitted\.
26. 6\.Experimental setting/details
27. Question: Does the paper specify all the training and test details \(e\.g\., data splits, hyperparameters, how they were chosen, type of optimizer\) necessary to understand the results?
28. Answer:\[Yes\]
29. Justification: Method part defines the datasets and evaluation metrics, while Appendix specifies preprocessing, model architecture, optimizer, learning\-rate schedule, epochs, batch sizes, loss weights, warmup schedules, and diffusion training/sampling details\.
30. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them\. - •The full details can be provided either with the code, in appendix, or as supplemental material\.
31. 7\.Experiment statistical significance
32. Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
33. Answer:\[No\]
34. Justification: The paper reports results over fixed evaluation sample sizes, such as 10,000 generated molecules, but does not report error bars, confidence intervals, statistical significance tests, or variability across independent training runs\. The reported results should therefore be interpreted as point estimates under the stated experimental settings\.
35. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The authors should answer\[Yes\]if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper\. - •The factors of variability that the error bars are capturing should be clearly stated \(for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions\)\. - •The method for calculating the error bars should be explained \(closed form formula, call to a library function, bootstrap, etc\.\) - •The assumptions made should be given \(e\.g\., Normally distributed errors\)\. - •It should be clear whether the error bar is the standard deviation or the standard error of the mean\. - •It is OK to report 1\-sigma error bars, but one should state it\. The authors should preferably report a 2\-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified\. - •For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range \(e\.g\., negative error rates\)\. - •If error bars are reported in tables or plots, the authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text\.
36. 8\.Experiments compute resources
37. Question: For each experiment, does the paper provide sufficient information on the computer resources \(type of compute workers, memory, time of execution\) needed to reproduce the experiments?
38. Answer:\[Yes\]
39. Justification: Appendix reports the compute setup and profiling results, including NVIDIA H100 NVL GPUs, batch sizes, per\-component timing, peak memory, forward/backward cost, DDP step\-time comparisons, and scaling curves\. The full set of reported experiments used approximately 384 H100 GPU\-hours in total\.
40. Guidelines: - •The answer\[N/A\]means that the paper does not include experiments\. - •The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage\. - •The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute\. - •The paper should disclose whether the full research project required more compute than the experiments reported in the paper \(e\.g\., preliminary or failed experiments that didn’t make it into the paper\)\.
41. 9\.Code of ethics
43. Answer:\[Yes\]
44. Justification: The research uses public molecular datasets and standard machine\-learning experimentation without human subjects or private data\. The authors have reviewed the NeurIPS Code of Ethics and believe the work conforms to it\.
45. Guidelines: - •The answer\[N/A\]means that the authors have not reviewed the NeurIPS Code of Ethics\. - •If the authors answer\[No\], they should explain the special circumstances that require a deviation from the Code of Ethics\. - •The authors should make sure to preserve anonymity \(e\.g\., if there is a special consideration due to laws or regulations in their jurisdiction\)\.
46. 10\.Broader impacts
47. Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
48. Answer:\[No\]
49. Justification: The current manuscript focuses on methodological and empirical contributions for 3D molecular generation and does not include a dedicated discussion of both positive and negative societal impacts\. Potential impacts include beneficial use in molecular design as well as possible dual\-use concerns if generative molecular models are used to propose harmful compounds\.
50. Guidelines: - •The answer\[N/A\]means that there is no societal impact of the work performed\. - •If the authors answer\[N/A\]or\[No\], they should explain why their work has no societal impact or why the paper does not address societal impact\. - •Examples of negative societal impacts include potential malicious or unintended uses \(e\.g\., disinformation, generating fake profiles, surveillance\), fairness considerations \(e\.g\., deployment of technologies that could make decisions that unfairly impact specific groups\), privacy considerations, and security considerations\. - •The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments\. However, if there is a direct path to any negative applications, the authors should point it out\. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate Deepfakes for disinformation\. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster\. - •The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from \(intentional or unintentional\) misuse of the technology\. - •If there are negative societal impacts, the authors could also discuss possible mitigation strategies \(e\.g\., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML\)\.
51. 11\.Safeguards
52. Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse \(e\.g\., pre\-trained language models, image generators, or scraped datasets\)?
53. Answer:\[N/A\]
54. Justification: The current submission does not release a scraped dataset, a deployed generation service, or a high\-risk pretrained language/image model\. The planned post\-acceptance release is research code and benchmark\-trained molecular generation artifacts; if released checkpoints are included, they will be accompanied by usage guidelines in the repository\.
55. Guidelines: - •The answer\[N/A\]means that the paper poses no such risks\. - •Released models that have a high risk for misuse or dual\-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters\. - •Datasets that have been scraped from the Internet could pose safety risks\. The authors should describe how they avoided releasing unsafe images\. - •We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort\.
56. 12\.Licenses for existing assets
57. Question: Are the creators or original owners of assets \(e\.g\., code, data, models\), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
58. Answer:\[No\]
59. Justification: The paper cites the original datasets, baselines, and software resources used in the experiments, including QM9, GEOM\-Drugs, and prior molecular generation methods\. However, the current submission does not explicitly enumerate the licenses, versions, or terms of use for all existing assets; these will be listed in the post\-acceptance repository\.
60. Guidelines: - •The answer\[N/A\]means that the paper does not use existing assets\. - •The authors should cite the original paper that produced the code package or dataset\. - •The authors should state which version of the asset is used and, if possible, include a URL\. - •The name of the license \(e\.g\., CC\-BY 4\.0\) should be included for each asset\. - •For scraped data from a particular source \(e\.g\., website\), the copyright and terms of service of that source should be provided\. - •If assets are released, the license, copyright information, and terms of use in the package should be provided\. For popular datasets,[paperswithcode\.com/datasets](https://arxiv.org/html/2606.13955v1/paperswithcode.com/datasets)has curated licenses for some datasets\. Their licensing guide can help determine the license of a dataset\. - •For existing datasets that are re\-packaged, both the original license and the license of the derived asset \(if it has changed\) should be provided\. - •If this information is not available online, the authors are encouraged to reach out to the asset’s creators\.
61. 13\.New assets
62. Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
63. Answer:\[N/A\]
64. Justification: The current submission does not release new standalone assets such as a dataset, benchmark, codebase, or model checkpoint\. The authors plan to release code and related artifacts after acceptance, with documentation, licensing information, and reproduction instructions provided alongside the release\.
65. Guidelines: - •The answer\[N/A\]means that the paper does not release new assets\. - •Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates\. This includes details about training, license, limitations, etc\. - •The paper should discuss whether and how consent was obtained from people whose asset is used\. - •At submission time, remember to anonymize your assets \(if applicable\)\. You can either create an anonymized URL or include an anonymized zip file\.
66. 14\.Crowdsourcing and research with human subjects
67. Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation \(if any\)?
68. Answer:\[N/A\]
69. Justification: The work does not involve crowdsourcing, human\-subject experiments, user studies, human annotation, or human evaluation\.
70. Guidelines: - •The answer\[N/A\]means that the paper does not involve crowdsourcing nor research with human subjects\. - •Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper\. - •According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector\.
71. 15\.Institutional review board \(IRB\) approvals or equivalent for research with human subjects
72. Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board \(IRB\) approvals \(or an equivalent approval/review based on the requirements of your country or institution\) were obtained?
73. Answer:\[N/A\]
74. Justification: The work does not involve crowdsourcing or human\-subject research, so IRB or equivalent review is not applicable\.
75. Guidelines: - •The answer\[N/A\]means that the paper does not involve crowdsourcing nor research with human subjects\. - •Depending on the country in which research is conducted, IRB approval \(or equivalent\) may be required for any human subjects research\. If you obtained IRB approval, you should clearly state this in the paper\. - •We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution\. - •For initial submissions, do not include any information that would break anonymity \(if applicable\), such as the institution conducting the review\.
76. 16\.Declaration of LLM usage
77. Question: Does the paper describe the usage of LLMs if it is an important, original, or non\-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does*not*impact the core methodology, scientific rigor, or originality of the research, declaration is not required\.
78. Answer:\[N/A\]\.
79. Justification: LLMs were used only for writing, editing, or formatting assistance and did not affect the core methodology, experiments, scientific rigor, or originality of the research\. The core method development does not involve LLMs as an important, original, or non\-standard component\.
80. Guidelines: - •The answer\[N/A\]means that the core method development in this research does not involve LLMs as any important, original, or non\-standard components\. - •Please refer to our LLM policy in the NeurIPS handbook for what should or should not be described\.