# Exploiting Latent Linearity in LLMs Improves Explainable Molecular Representation Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 77 references that are in the map.
- First posted on arXiv 11 Oct 2024.

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

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Language Models as Knowledge Bases?](https://mlsanity.com/paper/1909.01066/) (2019)
- [Sanity Checks for Saliency Maps](https://mlsanity.com/paper/1810.03292/) (2018)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [Deep Variational Information Bottleneck](https://mlsanity.com/paper/1612.00410/) (2016)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [The Linear Representation Hypothesis and the Geometry of Large Language Models](https://mlsanity.com/paper/2311.03658/) (2023)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [Is Attention Interpretable?](https://mlsanity.com/paper/1906.03731/) (2019)
- [Analyzing Learned Molecular Representations for Property Prediction](https://mlsanity.com/paper/1904.01561/) (2019)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [The (Un)reliability of saliency methods](https://mlsanity.com/paper/1711.00867/) (2017)
- [Pre-training Molecular Graph Representation with 3D Geometry](https://mlsanity.com/paper/2110.07728/) (2021)
- [Large-Scale Chemical Language Representations Capture Molecular Structure and Properties](https://mlsanity.com/paper/2106.09553/) (2021)
- [ChemRL-GEM: Geometry Enhanced Molecular Representation Learning for Property Prediction](https://mlsanity.com/paper/2106.06130/) (2021)
- [Evaluating Explainable AI: Which Algorithmic Explanations Help Users Predict Model Behavior?](https://mlsanity.com/paper/2005.01831/) (2020)
- [Representation Learning with Large Language Models for Recommendation](https://mlsanity.com/paper/2310.15950/) (2023)
- [Harnessing Explanations: LLM-to-LM Interpreter for Enhanced Text-Attributed Graph Representation Learning](https://mlsanity.com/paper/2305.19523/) (2023)
- [ChemBERTa-2: Towards Chemical Foundation Models](https://mlsanity.com/paper/2209.01712/) (2022)
- [Generative Causal Explanations for Graph Neural Networks](https://mlsanity.com/paper/2104.06643/) (2021)
- [Towards 3D Molecule-Text Interpretation in Language Models](https://mlsanity.com/paper/2401.13923/) (2024)
- [A Quantitative and Qualitative Evaluation of LLM-Based Explainable Fault Localization](https://mlsanity.com/paper/2308.05487/) (2023)

## Abstract (arXiv)

Large language models (LLMs) have demonstrated broad utility across molecular domains, spanning drug discovery and materials design. Analyzing LLMs' latent representations is crucial for elucidating their underlying mechanisms, improving explainability, and ultimately advancing downstream performance. We propose MoleX, a simple yet effective framework that decomposes molecular embeddings within LLM representations into a concept-aligned space for explainable molecular representation learning. We further show that these high-dimensional embeddings admit a linear mapping onto chemically consistent concepts. Our analysis suggests that the uncovered linearity aligns with established chemical principles, indicating a mechanistically explainable latent structure in LLM representations for scientific applications. When applied to downstream tasks, this latent linearity improves both predictive and explanatory performance. Extensive experiments demonstrate that MoleX outperforms existing approaches in accuracy, explainability, and efficiency, achieving CPU inference on large-scale datasets 300 times faster with 100,000 fewer parameters than LLMs.