# Molecules Meet Language: Confound-Aware Representation Learning and Chemical Property Steering in Transformer-VAE Latent Spaces

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Interpreting biomedical and science models, one of 13 directions in Mechanistic interpretability.
- Builds on 21 of its 32 references that are in the map.
- First posted on arXiv 7 May 2026.

## Placed in

- Mechanistic interpretability / [Interpreting biomedical and science models](https://mlsanity.com/mechinterp/interpreting-biomedical-and-science-models/)

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Understanding intermediate layers using linear classifier probes](https://mlsanity.com/paper/1610.01644/) (2016)
- [Probing Classifiers: Promises, Shortcomings, and Advances](https://mlsanity.com/paper/2102.12452/) (2021)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Designing and Interpreting Probes with Control Tasks](https://mlsanity.com/paper/1909.03368/) (2019)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [MolGAN: An implicit generative model for small molecular graphs](https://mlsanity.com/paper/1805.11973/) (2018)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [Optimization of Molecules via Deep Reinforcement Learning](https://mlsanity.com/paper/1810.08678/) (2018)
- [MoFlow: An Invertible Flow Model for Generating Molecular Graphs](https://mlsanity.com/paper/2006.10137/) (2020)
- [SMILES Transformer: Pre-trained Molecular Fingerprint for Low Data Drug Discovery](https://mlsanity.com/paper/1911.04738/) (2019)
- [MolGenSurvey: A Systematic Survey in Machine Learning Models for Molecule Design](https://mlsanity.com/paper/2203.14500/) (2022)
- [SMILES Enumeration as Data Augmentation for Neural Network Modeling of Molecules](https://mlsanity.com/paper/1703.07076/) (2017)
- [Reinforced Genetic Algorithm for Structure-based Drug Design](https://mlsanity.com/paper/2211.16508/) (2022)
- [LIMO: Latent Inceptionism for Targeted Molecule Generation](https://mlsanity.com/paper/2206.09010/) (2022)
- [SELFormer: Molecular Representation Learning via SELFIES Language Models](https://mlsanity.com/paper/2304.04662/) (2023)
- [Unraveling Molecular Structure: A Multimodal Spectroscopic Dataset for Chemistry](https://mlsanity.com/paper/2407.17492/) (2024)
- [Navigating Chemical Space with Latent Flows](https://mlsanity.com/paper/2405.03987/) (2024)

## Abstract (arXiv)

Molecular generative models often assume meaningful latent geometry, but apparent property predictability can reflect sequence-level shortcuts rather than chemical organization. We study this issue in an unsupervised autoregressive Transformer-VAE trained on SELFIES. After training, we freeze the model, fit linear probes to RDKit descriptors, and use the probe weights as candidate global steering directions. To separate chemical signal from SELFIES artifacts, we introduce a confound-aware evaluation based on residualization, confound-direction alignment analysis, and decoded-molecule traversal. This is necessary because SELFIES length, branch tokens, ring tokens, and token entropy are strongly encoded in the latent space. Under this confound-aware evaluation, we find robust monotonic steering for cLogP, FractionCSP3, HeavyAtomCount, TPSA, BertzCT, and HBA. Nonlinear probes further show that some properties admit stable global directions, while others are better described by local latent gradients. Overall, our results show that chemically meaningful steering can emerge in entangled molecular latent spaces, but only when validated through decoded molecules and controlled for representation-level confounds.