What Does a Chemical Language Model Know About Molecules?

Technical Summary: What Does a Chemical Language Model Know About Molecules?

Problem Statement
Chemical Language Models (cLMs) have demonstrated strong performance in molecular property prediction and generative drug design. However, a significant skepticism remains within the community regarding whether these models learn meaningful molecular semantics or merely exploit surface-level syntactic patterns of molecular string representations (e.g., SMILES). Given that molecules can be represented as strings in multiple equivalent ways, it is unclear if cLMs internalize chemical structure or simply parse string grammar. Addressing this requires a mechanistic understanding of how cLMs internally represent molecules.

Methodology
The authors apply Sparse Autoencoders (SAEs), a technique for mechanistic interpretability, to MolFormer-XL (specifically the open-weights MolFormer-10 variant), an encoder-only cLM pretrained on 1.1 billion molecular SMILES strings.

• SAE Training: They train k-SAEs on the residual stream of the model to decompose activations into sparse, human-interpretable features.
• Feature Classification: SAE latents are classified based on dominant token types (atoms, ring indices, branches, bonds) and mapped to atom-related concepts (atom, atom-in-substructure, substructure) using SMARTS notation.
• Perturbation Analysis: The study investigates representation shifts caused by non-canonical SMILES (valid variations in string parsing) and invalid SMILES (syntax or valence errors) to test the model's invariance to string notation.
• Downstream Probing: Linear probes are trained on 21 ADMET (Absorption, Distribution, Metabolism, Excretion, Toxicity) tasks to identify which SAE latents correlate with pharmacologically relevant properties.
• Tooling: The authors developed InterMol, an interactive web visualizer for exploring SAE activations on molecular strings and structures.

Key Contributions and Results

1. Hierarchical Representation Building:
   The analysis reveals a hierarchical progression in how molecular representations are constructed:

   • Early Layers: Rely heavily on position-tracking latents to parse molecular grammar (e.g., tracking ring indices, branch closings, and token positions). These layers show a high density of "position latents" that correlate with relative token positions.
   • Later Layers: Shift toward encoding atom-in-substructure and substructure-level features. The model appears to learn shared atom-pair latents for common motifs (e.g., C=O, C-N) rather than distinct latents for every specific functional group, allowing for reuse across complex structures.

2. Syntax Parsing Features:
   Beyond atom-related concepts, the authors identify two distinct types of syntax parsing features:

   • Position Latents: Prominent in early layers, these track specific token positions or the completeness of syntactic substructures (e.g., tracking a double-bonded oxygen within a branch from opening to closing).
   • Anything-But-X Latents: These activate on all tokens except specific types (e.g., "anything but carbons"). While their exact function is unclear, they may assist in discriminating heteroatoms, though the authors note they could be artifacts of SAE size.

3. Impact of SMILES Variants:

   • Non-canonical SMILES: These produce more disruptive representation shifts than invalid SMILES. The disruption is driven by position-latent disruption in early layers, which propagates to deeper layers. This suggests that models pretrained only on canonical SMILES may struggle with invariance to string variations due to their reliance on position-tracking mechanisms.
   • Invalid SMILES: Syntax errors (e.g., truncation) cause wide similarity distributions, while valence errors significantly affect specific non-positional latents. Notably, a specific latent in Layer 12 (L12/f/956) acts as a detector for both canonicality and validity, showing strong activation when inputs are augmented or corrupted.

4. Pharmacologically Relevant Latents:
   Linear probing on ADMET tasks demonstrates that SAE embeddings outperform baseline molecular fingerprints and descriptors on 16 of 21 tasks. The authors identified specific latents corresponding to known pharmacophoric features:

   • Human Intestinal Absorption (HIA): Latents activating on halogen substituents (Cl, F) and double-bonded oxygens.
   • CYP2C9 Inhibition: A latent activating on heteroatoms within heterocycles.
   • Ames Mutagenicity: Latents activating on nitrogen-nitrogen bonds (nitrosamines) and nitro/azide groups.
   • Lipophilicity: Latents activating on aromatic carbons (increasing lipophilicity) and carboxylic acid motifs (decreasing it).

Significance and Claims
The paper claims to provide a grounded, mechanistic answer to what chemical language models know about molecules. It demonstrates that cLMs do not merely learn surface patterns but build hierarchical representations that progress from syntactic parsing to chemically meaningful substructure concepts. The findings suggest that while cLMs encode pharmacologically relevant features, their reliance on position-tracking latents creates a vulnerability to non-canonical string variations.

The authors modestly frame their work as a foundation for future cLM design, suggesting that training strategies involving SMILES augmentation or contrastive learning could mitigate positional bias. They also propose that future work could utilize cross-layer transcoders to identify causal molecular circuits. The paper explicitly avoids claiming that cLMs have achieved human-level chemical reasoning, instead focusing on the identification of interpretable features that align with known chemical principles. The introduction of InterMol is presented as a tool to facilitate community exploration of these uninterpreted latents.