Exploiting Latent Linearity in LLMs Improves Explainable Molecular Representation Learning

TL;DR

This paper introduces MoleX, a framework that leverages latent linearity in LLMs to enhance explainability and performance in molecular representation learning, with significant improvements in speed and efficiency.

Contribution

MoleX decomposes molecular embeddings into a concept-aligned space, revealing linear structures that align with chemical principles and improve downstream task performance.

Findings

MoleX outperforms existing methods in accuracy and explainability.

It achieves 300x faster CPU inference on large datasets.

Uses 100,000 fewer parameters than traditional LLMs.

Abstract

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.