# Learning Invariant Molecular Representation in Latent Discrete Space

arXiv: https://arxiv.org/abs/2310.14170 · HTML: https://mlsanity.com/paper/2310.14170/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 70 references that are in the map.
- First posted on arXiv 22 Oct 2023.

## Builds on (most cited first, 30 of 44)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [SoundStream: An End-to-End Neural Audio Codec](https://mlsanity.com/paper/2107.03312/) (2021)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Towards Out-Of-Distribution Generalization: A Survey](https://mlsanity.com/paper/2108.13624/) (2021)
- [Explainability in Graph Neural Networks: A Taxonomic Survey](https://mlsanity.com/paper/2012.15445/) (2020)
- [Analyzing Learned Molecular Representations for Property Prediction](https://mlsanity.com/paper/1904.01561/) (2019)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Discovering Invariant Rationales for Graph Neural Networks](https://mlsanity.com/paper/2201.12872/) (2022)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism](https://mlsanity.com/paper/2201.12987/) (2022)

## Built on by (most cited first, 19 of 19)

- [UniMoT: Unified Molecule-Text Language Model with Discrete Token Representation](https://mlsanity.com/paper/2408.00863/) (2024)
- [Raising the Bar in Graph OOD Generalization: Invariant Learning Beyond Explicit Environment Modeling](https://mlsanity.com/paper/2502.10706/) (2025)
- [Discrete Tokenization for Multimodal LLMs: A Comprehensive Survey](https://mlsanity.com/paper/2507.22920/) (2025)
- [DIVE: Subgraph Disagreement for Graph Out-of-Distribution Generalization](https://mlsanity.com/paper/2408.04400/) (2024)
- [GraphTOP: Graph Topology-Oriented Prompting for Graph Neural Networks](https://mlsanity.com/paper/2510.22451/) (2025)
- [When Graph Neural Network Meets Causality: Opportunities, Methodologies and An Outlook](https://mlsanity.com/paper/2312.12477/) (2023)
- [A Survey of Out-of-distribution Generalization for Graph Machine Learning from a Causal View](https://mlsanity.com/paper/2409.09858/) (2024)
- [Pruning Spurious Subgraphs for Graph Out-of-Distribution Generalization](https://mlsanity.com/paper/2506.05957/) (2025)
- [Invariant Graph Learning Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2408.01697/) (2024)
- [Invariant Graph Transformer for Out-of-Distribution Generalization](https://mlsanity.com/paper/2508.00304/) (2025)
- [Quantifying Distributional Invariance in Causal Subgraph for IRM-Free Graph Generalization](https://mlsanity.com/paper/2510.20295/) (2025)
- [WED-Net: A Weather-Effect Disentanglement Network with Causal Augmentation for Urban Flow Prediction](https://mlsanity.com/paper/2601.22586/) (2026)
- [Improving Graph Out-of-distribution Generalization Beyond Causality](https://mlsanity.com/paper/2407.10204/) (2024)
- [Context-aware Graph Causality Inference for Few-Shot Molecular Property Prediction](https://mlsanity.com/paper/2601.11135/) (2026)
- [Distributional Priors Guided Diffusion for Generating 3D Molecules in Low Data Regimes](https://mlsanity.com/paper/2404.00962/) (2024)
- [Discrete Dictionary-based Decomposition Layer for Structured Representation Learning](https://mlsanity.com/paper/2406.06976/) (2024)
- [Teleporter Theory: A General and Simple Approach for Modeling Cross-World Counterfactual Causality](https://mlsanity.com/paper/2406.11501/) (2024)
- [Rethinking Graph Generalization through the Lens of Sharpness-Aware Minimization](https://mlsanity.com/paper/2602.08855/) (2026)
- [Molecular Property Prediction under Structural Shift with Tabular Foundation Models](https://mlsanity.com/paper/2609.38744/) (2026)

## Abstract (arXiv)

Molecular representation learning lays the foundation for drug discovery. However, existing methods suffer from poor out-of-distribution (OOD) generalization, particularly when data for training and testing originate from different environments. To address this issue, we propose a new framework for learning molecular representations that exhibit invariance and robustness against distribution shifts. Specifically, we propose a strategy called ``first-encoding-then-separation'' to identify invariant molecule features in the latent space, which deviates from conventional practices. Prior to the separation step, we introduce a residual vector quantization module that mitigates the over-fitting to training data distributions while preserving the expressivity of encoders. Furthermore, we design a task-agnostic self-supervised learning objective to encourage precise invariance identification, which enables our method widely applicable to a variety of tasks, such as regression and multi-label classification. Extensive experiments on 18 real-world molecular datasets demonstrate that our model achieves stronger generalization against state-of-the-art baselines in the presence of various distribution shifts. Our code is available at https://github.com/HICAI-ZJU/iMoLD.