Source-linked AI summary
All SMILES Variational Autoencoder
Zaccary Alperstein, Artem Cherkasov, Jason Tyler Rolfe
TL;DR
Molecular VAEs must handle discrete molecular structure and the non-unique, geometrically misleading nature of SMILES strings. The All SMILES VAE encodes multiple traversals with pooled recurrent networks, reconstructs disjoint strings, and uses attention for a fixed-length molecular latent representation. It reports state-of-the-art property prediction and molecular optimization while constraining search to the prior’s high-probability region.
Problem
Molecular optimization lacks a straightforward continuous, smooth search space, and non-unique SMILES strings can distort distances between molecules.
Method
The model processes multiple SMILES strings per molecule with stacked pooled RNNs, uses attention to form a fixed-length latent representation, and decodes disjoint strings.
Results
The All SMILES VAE surpasses state-of-the-art results across fully and semi-supervised property prediction and molecular property optimization tasks.
Takeaways & Limitations
The learned molecule-based latent space supports gradient-based optimization constrained near the prior region containing almost all probability mass.
Takeaways & Limitations
Unconstrained and MAP optimization leave the practical support of the trained prior, producing large prediction errors and poor optimization.
Abstract
from arXiv · showhide
Variational autoencoders (VAEs) defined over SMILES string and graph-based representations of molecules promise to improve the optimization of molecular properties, thereby revolutionizing the pharmaceuticals and materials industries. However, these VAEs are hindered by the non-unique nature of SMILES strings and the computational cost of graph convolutions. To efficiently pass messages along all paths through the molecular graph, we encode multiple SMILES strings of a single molecule using a set of stacked recurrent neural networks, pooling hidden representations of each atom between SMILES representations, and use attentional pooling to build a final fixed-length latent representation. By then decoding to a disjoint set of SMILES strings of the molecule, our All SMILES VAE learns an almost bijective mapping between molecules and latent representations near the high-probability-mass subspace of the prior. Our SMILES-derived but molecule-based latent representations significantly surpass the state-of-the-art in a variety of fully- and semi-supervised property regression and molecular property optimization tasks.
1 Introduction
Molecular optimization is difficult because molecules occupy a discrete, sparse space and SMILES strings provide non-unique representations with misleading geometry. The All SMILES VAE addresses this by encoding multiple SMILES strings and decoding disjoint representations of the same molecule.
- 1 Introduction: Molecular optimization spans a discrete, sparse space containing many chemically invalid combinations despite an enormous number of possible small molecules.Molecules are represented as graphs with labeled atoms and bonds, while most combinations do not form stable molecules.
- 1 Introduction: Existing SMILES models constrain decoding with grammar mechanisms because most strings over the character set do not represent valid molecules.SMILES strings flatten spanning-tree traversals of molecular graphs and encode atoms, bonds, branching, and ring closures.
- 1 Introduction: Different valid SMILES strings for one molecule can be farther apart than strings from dissimilar molecules, creating unnecessary local optima for property optimization.This makes models of individual SMILES strings poorly aligned with molecular identity and property geometry.
- 1 Introduction: Reinforcement-learning molecule generators can optimize simple in-silico properties but generally scale poorly when properties require time-consuming wet-lab experiments.This limitation motivates methods that optimize properties from a fixed training dataset.
- 1 Introduction: The All SMILES VAE encodes multiple SMILES strings with recurrent networks, pools information across molecular paths, and decodes disjoint strings representing the same molecule.Its latent representation is intended to capture molecular features rather than a particular SMILES realization.
2 Efficient molecular encoding with multiple SMILES strings
Graph convolutions propagate information locally and often require multiple rounds, whereas SMILES RNNs can transmit information along flattened molecular paths in a single sequential pass. Multiple traversals therefore provide efficient access to spatially distributed molecular information.
- 2 Efficient molecular encoding with multiple SMILES strings: A VAE optimizes the ELBO, balancing reconstruction through p(x|z) against divergence between q(z|x) and the prior p(z).This objective defines the stochastic encoder-decoder framework used by the model.
- 2 Efficient molecular encoding with multiple SMILES strings: Molecule encoders commonly use only three to seven graph-message-passing rounds, limiting the propagation of long-range information through molecular graphs.Non-local properties such as binding-related behavior can depend on spatially distributed molecular features.
- 2 Efficient molecular encoding with multiple SMILES strings: SMILES RNNs pass hidden-state messages sequentially along a depth-first spanning-tree traversal, allowing earlier information to influence all later branches.With a suitable spanning tree, information can cross the molecular graph’s full width in one RNN update.
- 2 Efficient molecular encoding with multiple SMILES strings: The All SMILES encoder combines parallel RNNs over distinct SMILES strings with graph-like pooling to capture multiple molecular message-passing pathways efficiently.The approach aims to combine graph-convolution latent geometry with the information-transfer efficiency of SMILES RNNs.
3 Model architecture
The model combines stacked pooled RNNs over multiple SMILES strings with attentional latent-variable construction and disjoint-string reconstruction. Latent optimization is constrained to the prior’s high-probability region because unconstrained and MAP strategies leave the trained support and perform poorly.
- 3 Model architecture: The All SMILES encoder applies stacked RNN modules to parallel SMILES strings and pools homologous atom representations between layers.This lets information flow along the union of implicit SMILES pathways while preserving separate sequence processing.
- 3 Model architecture: Homologous atom characters can be aligned across different-length SMILES strings even when branching and ring-closure symbols do not match.The pooled atom representation replaces the original atom representation in each stream, while syntactic characters remain unchanged.
- 3 Model architecture: The posterior forms a fixed-length autoregressive Gaussian hierarchy: max-pooled terminal states parameterize z1, while later latent subsets use attention over pooled atom outputs.Queries for later subsets are computed from preceding latent layers, enabling aggregation of spatially diffuse molecular features.
- 3 Model architecture: The decoder is a single-layer LSTM trained to reconstruct SMILES strings disjoint from the encoder inputs but representing the same molecule.This reconstruction setup pressures the latent representation to encode molecular identity rather than a particular string realization.
- 3.1 Latent space optimization: Unconstrained and MAP optimization leave the practical support of the hierarchical prior, causing large prediction errors and poor optimization.Unconstrained optimization moves away from the origin, whereas Gaussian-prior MAP optimization is pulled toward it.
- 3.1 Latent space optimization: The method instead reparameterizes the autoregressive prior to a standard Gaussian and optimizes on a spherical shell containing almost all prior probability mass.The radius constraint is applied separately to each hierarchical layer while optimizing angular coordinates.
4 Results
The All SMILES VAE improves property prediction and enables gradient-based molecular optimization through a molecule-consistent latent representation. It achieves state-of-the-art results across supervised, semi-supervised, and optimization tasks while requiring only a fixed training dataset.
- Generative performance: 87.4% ± 1% of held-out ZINC250k molecules are reconstructed accurately, while 98.5% ± 0.1% of prior samples decode to valid SMILES strings.Among 50,000 prior samples, all decoded molecules were unique and 99.958% were novel relative to training data.
- Property prediction: The model significantly improves semi-supervised prediction of logP, molecular weight, and QED, and surpasses prior fully supervised and toxicity-prediction results.The semi-supervised comparison uses ZINC310k; fully supervised regression uses ZINC250k, while toxicity prediction uses Tox21.
- Molecular optimization: The latent space is evaluated for property smoothness by comparing true logP across a densely sampled two-dimensional slice aligned with the logP regressor.Figure 6b tracks predicted and true penalized logP along an optimization trajectory.
- Molecular optimization: Gradient-based optimization maximizes learned property regressors with a log-prior regularizer and hierarchical radius constraint before decoding optimized latent points into molecules.The evaluated objectives include QED and penalized logP, with penalized logP subtracting synthetic accessibility and large-ring penalties [9] [27] [35] [56] [70] [72].
- Molecular optimization: Unlike reinforcement-learning methods that evaluate generated molecules during training, the All SMILES VAE optimizes properties using only a fixed training dataset.This is presented as important for properties such as pharmacological efficacy, toxicity, and OLED efficiency when accurate in-silico approximations are unavailable.
5 Conclusion
The All SMILES VAE builds molecule-consistent latent representations by processing multiple SMILES strings with pooled recurrent networks and attentional aggregation. These representations support state-of-the-art property prediction and constrained gradient-based molecular optimization.
- 5 Conclusion: Multiple SMILES strings are encoded with stacked, pooled RNNs, while decoding targets disjoint strings of the same molecule to enforce consistent molecular representations.Attentional pooling summarizes spatially diffuse features into a fixed-length, non-factorial posterior.
- 5 Conclusion: Linear and logistic regressors on the latent representation achieve state-of-the-art semi- and fully supervised property prediction.Gradient optimization of these regressors within a high-probability prior subspace produces state-of-the-art optimized molecules with an RNN decoder.
A Datasets
The experiments use ZINC and Tox21 datasets spanning small organic molecules, molecular properties, and toxicity assays. Dataset scales and molecular graph sizes vary substantially across tasks.
- SMILES representation: Different SMILES strings for one molecule can be less similar than strings representing radically different molecules.The highlighted example compares 30% similarity between strings for one molecule with 60% similarity to a different molecule.
- ZINC datasets: ZINC250k contains 250,000 organic molecules with 6–38 heavy atoms and penalized logP values from -13 to 5.Penalized logP subtracts synthetic accessibility and the number of rings with more than six atoms.
- ZINC datasets: ZINC310k contains 310,000 organic molecules with 6–38 heavy atoms for semi-supervised prediction of logP, molecular weight, and QED.These properties are computed using RDKit.
- Molecular graph structure: ZINC250k molecules have a mean graph diameter of 11.1 and a maximum diameter of 24, exceeding the message-passing range of typical three-to-seven-round graph convolutions.The diameter is the maximum eccentricity over all atoms in a molecular graph.
- Tox21 dataset: Tox21 contains binarized binding affinities against up to 12 proteins, with 6264 training, 783 validation, and 784 test molecules.Its molecules include structures with up to 140 atoms, ranging from large peptides to metal-containing compounds.
B Extended model architecture
The All SMILES VAE combines multiple SMILES encodings of one molecule with recurrent message passing, atom-level pooling, hierarchical latent inference, and constrained optimization. Its design aims to propagate information across molecular graph pathways efficiently while learning molecule-based latent representations.
- Training objective: The decoder uses LSTMs to generate a disjoint set of SMILES strings, while the ELBO combines their conditional log-likelihoods with the posterior-prior KL divergence.For molecular optimization, the KL term is usually scaled by the number of decoder SMILES strings, but it remains unscaled for property prediction.
- Latent inference: The hierarchical approximating posterior samples latent variables from conditional Gaussian distributions whose parameters are produced by attention over final atom representations.The attention mechanism uses encoder hidden vectors as keys and a query computed from previous latent layers; its output determines conditional means and log-variances.
- Latent space optimization: Latent optimization adds λ log(p(z)) and applies hard tanh before regression to constrain searches to well-trained regions and moderate strictly monotonic regressors.The hard-tanh linear region encompasses all values observed in the training set.
- Architecture: Multiple SMILES strings are encoded in parallel, with recurrent modules pooling homologous atom representations so information can flow across all molecular graph pathways efficiently.The approach is intended to propagate information across the molecule in a single layer, avoiding the depth and computational expense of graph convolutions.
- Contributions: The method extends SMILES transcoders into a variational generative model over molecules, using a learned hierarchical prior to regularize property prediction and molecular optimization.A hierarchical radius constraint is introduced so optimization remains within the practical support containing almost all prior probability mass.
C Extended results
The paper evaluates the All SMILES VAE against several state-of-the-art molecular property prediction and optimization methods on standard benchmark tasks.
C.1 Reconstruction accuracy and validity
Reconstruction and validity evaluation is complicated by the model’s deliberately harder cross-encoding reconstruction task and deterministic beam-search decoding. These choices make results less directly comparable with prior stochastic VAE evaluations.
- Reconstruction accuracy: Different input and target SMILES encodings force the model to capture the molecule rather than a particular string, but make decoding a complex multimodal task.This increases the chance of syntactic or semantic errors relative to reconstructing the input encoding.
- Decoding procedure: Beam search selects the most probable decoded SMILES string under p(x|z), producing an unambiguous decoding for every latent point.The full decoder distribution is still used during training.
- Comparability: Beam-search decoding makes the reported reconstruction and validity results incommensurable with much prior work using stochastic encoders and decoders.The limitation concerns direct comparison of these evaluation metrics rather than the training procedure itself.
C.2 Property prediction
The All SMILES VAE produces strong property-prediction results across supervised and semi-supervised settings. Jointly learned latent representations outperform unsupervised-VAE Gaussian-process baselines and comparison algorithms on key molecular properties.
- Fully supervised prediction: Property prediction is two orders of magnitude more accurate than sparse Gaussian-process regression on an unsupervised VAE latent representation.The paper attributes smooth molecular-property structure to jointly training the property regressor with the generative model.
- Semi-supervised prediction: Mean absolute error is at least three times smaller than comparison algorithms for logP and molecular weight in semi-supervised prediction.The reported semi-supervised evaluation also includes QED on ZINC310k.
- Latent-space visualization: A latent-space slice compares predicted and true logP after applying the hierarchical radius constraint to decoded molecules.The visualization uses a 2D slice with its y axis aligned to the trained logP regressor.
- Tox21 prediction: The evaluation includes per-assay AUC-ROC measurements on the Tox21 dataset.The reported assays cover nuclear-receptor signaling and stress-response pathways.
C.3 Molecular optimization
The All SMILES VAE is evaluated through molecular-property optimization trajectories, including QED and penalized LogP objectives. The reported QED trajectory reaches a maximum observed true QED of 0.948, while penalized LogP optimization examines top candidates under a hierarchical radius constraint.
- Optimization setup: The QED experiment scales the ELBO KL term by the number of SMILES strings in the decoder, making the loss analogous to parallel single-SMILES VAEs.Message passing between encoders still produces a shared latent representation.
- QED optimization: 0.948 maximum observed true QED is reached along the reported drug-likeness optimization trajectory.Figure 13 compares predicted and true QED throughout the trajectory.
- Penalized LogP optimization: Top three true penalized LogP molecules are produced by gradient-based optimization with a hierarchical radius constraint.The experiment uses an All SMILES VAE variant whose KL term is not scaled by the number of decoder SMILES strings.
C.4 Ablation of model components
Ablations show that message flow between multiple SMILES representations and the hierarchical posterior contribute to model performance, while the radius constraint improves penalized LogP optimization. The section also specifies the grammar and automata-based constraints used to generate valid SMILES strings.
- Message-passing ablations: Significant performance decay follows ablations that disrupt message flow between flattened spanning trees, while single-SMILES encoding and decoding causes a further reduction.The single-SMILES variant lets the latent representation encode a particular SMILES string rather than only the underlying molecule.
- Posterior hierarchy: The hierarchical approximating posterior with attentional pooling makes a meaningful contribution to model performance.Removing posterior hierarchy moves all latent variables to the first layer, leaves a standard Gaussian prior, and removes attentional pooling over atom representations.
- Radius-constraint ablation: The hierarchical radius constraint significantly improves molecular optimization under penalized LogP-only optimization without a log-prior regularizer.Table 8 reports the top three true penalized LogPs from 1000 trajectories evaluated on the 100 best predicted trajectories, with each optimization repeated five times.
- SMILES grammar: The SMILES grammar represents atoms, bonds, branches, ringbonds, charges, stereochemistry, and optional bracketed atom attributes through context-free productions.The supplied grammar includes productions for atom types, bond symbols, isotope, chiral, hydrogen count, charge, and class.
- Grammar enforcement: A pushdown automaton enforces SMILES syntactic restrictions during decoding by tracking atom and bond sequences and matching parentheses.Ringbond matching additionally uses a small random-access memory, with paired bonds required to carry matching labels under the stated grammar extension.
- Grammar enforcement: The grammar’s nested atomic contexts can be arbitrarily deep, but branch traversal returns to the parent atom only after the branch is completed.This structure motivates stack-based tracking of branching and valence-related context during SMILES generation.