Appendix C. Knowledge Distillation: Shrinking Models for Efficient, Hierarchical Molecular Generation

The chapter explains why compressing a strong molecular generator can be more valuable than simply maximizing its raw performance. Large hierarchical models are powerful but costly to train and deploy, especially in workflows like virtual screening, reinforcement learning, and active learning where generation happens repeatedly. Knowledge distillation is presented as the solution: a smaller student model learns from a larger teacher model, inheriting much of its chemical competence while becoming faster, lighter, and more practical for real-world use.

A major focus is the hierarchical nature of molecular generation. The text contrasts early SMILES-based approaches and atom-by-atom graph generation with motif-based methods that preserve chemical validity more naturally. It then describes hierarchical representations that decompose molecules into chemically meaningful substructures, enabling multi-scale reasoning from atoms to motifs to whole molecules. Hierarchical encoders and decoders, as used in HierVAE, exploit this structure to produce valid molecules more efficiently, and that same structure creates natural points for distillation at the latent, feature, and output levels.

The implementation strategy centers on an offline teacher-student setup with a frozen pre-trained teacher, a smaller student, and a multi-component loss that combines reconstruction, KL regularization, latent matching, and feature matching. Training stability is addressed with KL annealing, gradient clipping, and careful handling of dimension mismatches, while evaluation emphasizes validity, uniqueness, reconstruction, novelty, speed, and model size. The results show the expected trade-off: the student is much smaller and faster, with strong validity but some loss in diversity and reconstruction quality, illustrating how distillation can make molecular generation much more deployable without discarding most of the teacher’s knowledge.

Figure C.1 Atom-based methods require many generation steps, which provoke vanishing gradients and error accumulation over long horizons. Substructure-based methods, like JT-VAE, only considered single rings and bonds as building blocks. Whereas JT-VAE assembled each neighborhood consisting of multiple motifs in one shot, HierVAE decomposes the assembly process into multiple steps: predict attaching points, match atoms, then one motif at a time and not an entire neighborhood. For a detailed, guided view into hgraph2graph implementation, we recommend the author’s ICML presentation [3].

Figure C.2 The knowledge distillation paradigm. We train a teacher model (top), comparing its predictions to the true labels, y, and using the difference between the prediction and the ground truth to backpropagate into the weights of the teacher and adjust them so that on the next forward pass its predictions are a little bit better over time. We also train a student model in parallel (bottom) by using a combined loss: one comparing its predictions to ground truth (y) and a distillation loss comparing its predictions to the teacher’s “soft labels”.

Figure C.3 The power of dark knowledge. The top shows the typical error computation when training with hard labels, whereas the bottom shows how incorporating soft labels into knowledge distillation provides an additional, valuable signal as to the student model’s capabilities.

Figure C.4 Think of offline knowledge distillation as spreading the knowledge distillation process across discrete time steps. Once we've trained the teacher model, we use the teacher model to label x during inference time. This results in teacher model predictions denoted as y'. The student is now trained as before with x and y, but we have additionally new labels, y'. It’s the same input data that we have labels y for, but we have the soft labels (y') as well.

Figure C.5 Since the "Inference Teacher" step is not dependent on y at all, we can use additional x (marked x2) that we have even if it is not labeled. Even though we can't provide hard labels for these examples, we can still derive pseudo labels on additional data to feed into the distillation loss, allowing us to leverage a much larger pool of unlabeled data!

Figure C.6 Our implementation consists of four main components that work together to transfer knowledge from the teacher to the student: training (or, in our case, loading) a teacher model, producing soft labels, training a student model, and evaluating the model. Ideally, our student model sacrifices marginal performance for notable gains in efficiency metrics such as molecule generation speed and memory overhead.

• Knowledge distillation for hierarchical molecular generation provides a practical solution to the deployment challenges posed by large neural models. By training compact student models to mimic capable teacher models, we can achieve most of the teacher's performance with a fraction of the computational cost. This chapter has explored the complete pipeline from hierarchical molecular representation through distillation training to comprehensive evaluation.
• The hierarchical representation employed by HierVAE captures the modular nature of chemical structure, decomposing molecules into trees of motifs rather than flat sequences of atoms. This representation provides inductive biases that improve generation quality and efficiency. The HierVAE architecture combines variational autoencoders, message-passing neural networks, and hierarchical processing to enable both molecular reconstruction and novel molecule generation. Understanding this architecture was essential for effective distillation, as we needed to consider which components to compress and which knowledge to transfer.
• The multi-component distillation loss transfers knowledge through four complementary channels: reconstruction quality, latent space organization, and intermediate feature representations. Balancing these components through hyperparameters allows us to control which aspects of the teacher's knowledge to prioritize. The reconstruction loss ensures basic functionality, the KL loss maintains latent space structure, the latent distillation loss transfers semantic organization, and the feature matching loss transfers learned chemical patterns.
• Training distilled models requires careful attention to optimization dynamics. KL annealing prevents posterior collapse by initially focusing on reconstruction before gradually enforcing latent space constraints. The dual forward pass architecture enables efficient distillation by freezing the teacher and updating only the student. Gradient clipping, learning rate scheduling, and appropriate checkpointing ensure stable, reliable training. Monitoring multiple metrics throughout training helps diagnose problems and guide hyperparameter tuning.
• Student model design balances compression and capability through architectural choices. Reducing hidden dimensions and latent dimensions cuts parameter count while preserving structural inductive biases. Keeping message-passing depths constant maintains the receptive field needed for long-range dependencies. Empirical evaluation of different compression ratios reveals the performance-efficiency trade-off and guides architecture selection for specific deployment scenarios.
• Comprehensive evaluation characterizes both quality and efficiency of compressed models. Validity, uniqueness, reconstruction accuracy, property distributions, generation speed, and model size together capture the multifaceted nature of generative model quality. Statistical analysis provides confidence in observed differences and enables principled comparison of compression strategies. Ablation studies isolate the contribution of each loss component and identify opportunities for further improvement.
• The techniques presented here generalize beyond molecular generation to other structured generation tasks. The core principles of knowledge distillation (learning from teacher outputs, intermediate representations, and latent organization) apply to any scenario where a large model must be compressed for deployment. The specific architectural choices (hierarchical decomposition, message passing, variational sampling) apply to generation of other graph-structured objects like proteins.

FAQ