Kullback–Leibler (KL) divergence measures how much one probability distribution differs from another. Imagine using a coin known to be fair to predict an unknown coin: the divergence is zero if their probabilities match and larger as they differ. In a variational graph autoencoder (VGAE), KL divergence compares the learned latent distribution with a prior distribution. Adding it to the training loss helps shape a latent space from which the model can sample, rather than relying on fixed-point encodings. section 5.5.2 section 5.6
The practical goal is a latent space that supports generation. A regular space is continuous—nearby points decode to approximately similar outputs—and compact—sampled points can yield meaningful outputs. KL regularization helps structure that space, but it must be balanced against reconstruction: a model can reconstruct training inputs well yet interpolate poorly between them, or gain smoother sampling at the expense of reconstruction. section 5.3
To use KL divergence when training a VGAE:
- Have the graph encoder return a mean and log variance for its latent representation, rather than only a fixed encoding. section 5.3.1
- Combine reconstruction loss with a KL loss. PyTorch Geometric’s VGAE handles reparameterization and exposes `model.kl_loss()`; the edge-prediction training loop adds that term, scaled by the number of nodes. section 5.3.1
- Validate the result against a graph autoencoder without KL regularization. A VGAE need not achieve higher edge-prediction accuracy. section 5.3.1
Molecule generation
For molecular graphs, KL loss can sit alongside adjacency-matrix reconstruction loss and a property-prediction loss. One implementation computes binary cross-entropy for reconstructed edges, mean squared error for predicted QED, and KL loss from the encoder’s mean and log variance, then adds the three terms. The resulting latent space can be inspected for regions associated with higher QED before sampling candidate molecules; proposed candidates still need verification. section 5.4.2 section 5.4.4
Listing 5.18 Loss for molecule graph generation
def calculate_loss(self, pred_adj, \
true_adj, qed_pred, qed_true, mu, logvar):
adj_loss = F.binary_cross_entropy\
(pred_adj, true_adj) #1
qed_loss = F.mse_loss\
(qed_pred.view(-1), qed_true.view(-1)) #2
kl_loss = -0.5 * torch.mean\
(torch.sum(1 + logvar - mu.pow(2)\ #3
- logvar.exp(), dim=1))
return adj_loss + qed_loss + kl_loss
The loss listing shows the KL calculation beside adjacency reconstruction and QED prediction, making its role in the optimization concrete. section 5.4.2
The same loss pattern also appears in neural relational inference for temporal graph predictions: reconstruction and KL terms are both needed when learning latent relationships between nodes. KL divergence is therefore a training objective, not an edge-prediction or generation score by itself. section 6.4 section 5.3.1
