Semi-Supervised JT-VAE SeMole

• The paper introduces SeMole, an extension of JT-VAE that employs a semi-supervised objective to improve molecular property prediction under limited labeled data.
• The model leverages multi-latent variational inference with tree and graph encoders/decoders to jointly capture molecular structures and associated properties.
• Empirical results show that pretraining and progressive supervision in SeMole dramatically reduce MAE and achieve high chemical validity and property accuracy in generated molecules.

The Semi-Supervised Junction Tree Variational Autoencoder (SeMole) is a generative model that augments the Junction Tree VAE (JT-VAE) paradigm for molecular graphs with a semi-supervised learning framework. SeMole is specifically designed for molecular property prediction under scarcity of labeled data, exploiting both labeled and unlabeled molecules to improve sample efficiency and support property-conditioned molecular graph generation. The underlying approach integrates property prediction with conditional structure generation via multi-latent-variable variational inference and progressive pretraining (Hamidizadeh et al., 2022).

1. Model Architecture

SeMole builds directly on the JT-VAE architecture, which factorizes each molecule $G$ into two levels: (i) the junction tree $T$ over chemically valid substructures ("scaffolds"), and (ii) the molecular graph $G$ assembled from these substructures. The autoencoder is composed of four main subnetworks:

• Tree encoder $q_{\phi}(z_T \mid T)$: Implements a two-layer GRU over the node sequence of the junction tree, with learned substructure embeddings. It outputs a Gaussian posterior for the latent scaffold code $z_T \in \mathbb{R}^d$.
• Graph encoder $q_{\phi}(z_G \mid G)$: Utilizes a Message-Passing Neural Network (MPNN) to aggregate node-level representations, yielding a Gaussian posterior for the latent code $z_G \in \mathbb{R}^d$.
• Tree decoder $p_{\theta}(T \mid z_T, y)$: Predicts the topology and labeling of the junction tree, conditioned on both $z_T$ and the molecular property variable $y$.
• Graph decoder $T$0: Predicts the assembly of molecular substructures defined by the tree and the detailed attachment patterns and bond types, taking as input $T$1, $T$2, and $T$3. This subnetwork is implemented via another MPNN.

SeMole treats $T$4 and $T$5 as independent Gaussian priors: $T$6, $T$7, with a third latent variable $T$8 for the target molecular property, which is only partially observed.

The generative process is factorized as: $T$9

Recognition models are formulated differently for labeled ($G$0) and unlabeled ($G$1) samples:

• For labeled: $G$2
• For unlabeled: $G$3

2. Semi-Supervised Learning Objective

The SeMole framework employs a semi-supervised variational inference scheme in the "M2" style of Kingma et al. (2014). Two evidence lower bound (ELBO) objectives are defined:

• Labeled ELBO (for $G$4), denoted $G$5: $G$6
• Unlabeled ELBO (for $G$7), denoted $G$8: $G$9 An additional entropy term $q_{\phi}(z_T \mid T)$0 appears due to the marginalization over $q_{\phi}(z_T \mid T)$1.

The final objective combines labeled and unlabeled ELBOs, augmented with a weighted property prediction term: $q_{\phi}(z_T \mid T)$2 where $q_{\phi}(z_T \mid T)$3 regulates the strength of the supervised property-prediction loss.

The reconstruction losses decompose into:

• Tree topology and label cross-entropy (for $q_{\phi}(z_T \mid T)$4)
• Graph attachment cross-entropy (for $q_{\phi}(z_T \mid T)$5)
• Two KL divergences: $q_{\phi}(z_T \mid T)$6 and $q_{\phi}(z_T \mid T)$7

Classification/regression is enforced by the negative log-likelihood term $q_{\phi}(z_T \mid T)$8, which is mean squared error for regression or cross-entropy for classification.

3. Pretraining and Training Stabilization

Training multi-latent-variable VAEs is empirically unstable. SeMole introduces a pre-training schedule ("SeMoleₚₑₜᵣₐᵢₙₑd", Editor's term), where the supervised loss term ($q_{\phi}(z_T \mid T)$9) is set to zero for an initial warm-up period (epochs 1–10), meaning only the reconstruction ELBOs are active and the property prediction auxiliary network is not updated. During ramp-up (epochs 11–N), $z_T \in \mathbb{R}^d$0 is increased linearly up to a maximum hyperparameter $z_T \in \mathbb{R}^d$1.

This staged optimization allows the model to first learn a robust reconstruction manifold before encouraging the disentanglement of molecular properties in latent space. This curriculum, based on Kingma et al. (2014) and Maaløe et al. (2016), improves both final accuracy and training stability.

4. Property-Conditioned Generation

Unlike vanilla JT-VAE, SeMole injects the molecular property code $z_T \in \mathbb{R}^d$2 into both the tree and graph decoders:

• The tree decoder conditions every structure and label prediction on $z_T \in \mathbb{R}^d$3.
• The graph decoder's MPNN uses $z_T \in \mathbb{R}^d$4, in combination with the tree embedding, to inform chemical graph assembly.

The inference network $z_T \in \mathbb{R}^d$5 is a compact MLP over concatenated encoder embeddings, producing a Gaussian mean for property regression.

Through this property injection in both generative steps, the latent representations for structure ($z_T \in \mathbb{R}^d$6, $z_T \in \mathbb{R}^d$7) are tightly coupled with the molecular property variable $z_T \in \mathbb{R}^d$8, enabling both property-conditional graph generation and property inference from molecular structure.

5. Experimental Protocol and Baselines

The experimental evaluation uses 310,000 molecules from ZINC 15, with three scalar properties computed via RDKit: Molecular weight (MolWt), LogP (Wildman–Crippen partition coefficient), and QED (Quantitative Estimation of Drug-likeness). Each property is standardized. Experiments simulate label scarcity by varying the labeled fraction between 5% and 50%.

A summary of dataset and protocol:

Component	Details	Notes
Dataset	310k from ZINC 15	Drug-like molecules
Properties	MolWt, LogP, QED	Standardized
Test set	10,000 held-out molecules	Fixed
Labeled splits	5%, 10%, 20%, 50% labeled	Rest are unlabeled
Validation set	5% of labeled data
Hyperparameters	Batch: 16, Adam lr: 1e-3, $z_T \in \mathbb{R}^d$9	2-layer GRU (tree), 3-layer MPNN (graph decoder)
$q_{\phi}(z_G \mid G)$0	Tuned in $q_{\phi}(z_G \mid G)$1	Used for property-signal ramp-up

Baseline and ablation comparisons:

1. SSVAE (Kang & Cho 2018) operating on SMILES strings.
2. SeMoleₛᵤₚₑᵣᵥᵢₛₑd: graph JT-VAE without unlabeled data ($q_{\phi}(z_G \mid G)$2).
3. SeMole: full semi-supervised, constant $q_{\phi}(z_G \mid G)$3.
4. SeMoleₚₑₜᵣₐᵢₙₑd: pretraining with warm-up and ramp-up of $q_{\phi}(z_G \mid G)$4.

6. Empirical Performance and Analysis

SeMole is evaluated using Mean Absolute Error (MAE) on the test set for property prediction. Results synthesized from Table 1 of (Hamidizadeh et al., 2022):

• With 5% labeled data:
  • LogP: SSVAE 0.133±0.002, SeMole 0.120±0.006, SeMoleₚₑₜᵣₐᵢₙₑd 0.117±0.001
  • QED: SSVAE 0.038±0.000, SeMole 0.028±0.001, SeMoleₚₑₜᵣₐᵢₙₑd 0.021±0.000
• As the labeled fraction increases, SeMoleₚₑₜᵣₐᵢₙₑd maintains lower MAE than both SSVAE and the supervised-only SeMole variant. At 50% labels, MAE for LogP is ≈ 0.043 and for QED ≈ 0.009, compared to SSVAE’s 0.047/0.011.

Ablation findings indicate that removing unlabeled data (SeMoleₛᵤₚₑᵣᵥᵢₛₑd) significantly degrades performance under sparse supervision, while semi-supervised training with pretraining (SeMoleₚₑₜᵣₐᵢₙₑd) achieves both optimal accuracy and robust convergence. This suggests that the model’s capacity to learn effective representations for property-conditioned generation relies on first establishing a solid generative manifold.

Conditional molecule generation experiments show SeMoleₚₑₜᵣₐᵢₙₑd achieves 100% chemical validity, approximately 90% uniqueness/novelty, and approximately 80% of generated samples within ±5% of the target property value. SSVAE only attains ≈70% validity and ≈60% property accuracy.

7. Related Approaches and Significance

By marrying graph-structured molecular generation (JT-VAE backbone) with the probabilistic semi-supervised inference of Kingma et al. (2014), SeMole enables accurate molecular property prediction and controlled molecule synthesis in data regimes where labeled annotations are scarce. The model demonstrates empirically that reconstructive pretraining followed by progressive supervision yields more stable and performant semi-supervised VAEs for structured domains. The property-conditional decoding strategy notably enhances sample efficiency and controllability compared to previous graph and SMILES-based SSVAEs (Hamidizadeh et al., 2022).