Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Exploiting Latent Linearity in LLMs Improves Explainable Molecular Representation Learning

Zhuoran Li, Xu Sun, Wanyu Lin, Jiannong Cao

TL;DR

This work tackles the explainability gap in LLM-driven molecular property prediction by revealing latent linearity in LLM representations. It introduces MoleX, which decomposes embeddings into a concept space rooted in functional groups and learns a linear predictor augmented by residual calibration, yielding accurate and chemically faithful explanations. The approach achieves state-of-the-art performance across multiple benchmarks while enabling fast CPU inference and dramatically reducing parameter counts. These findings demonstrate that aligning foundation models with domain concepts can enhance both predictive power and mechanistic interpretability for scientific applications.

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.

Exploiting Latent Linearity in LLMs Improves Explainable Molecular Representation Learning

TL;DR

This work tackles the explainability gap in LLM-driven molecular property prediction by revealing latent linearity in LLM representations. It introduces MoleX, which decomposes embeddings into a concept space rooted in functional groups and learns a linear predictor augmented by residual calibration, yielding accurate and chemically faithful explanations. The approach achieves state-of-the-art performance across multiple benchmarks while enabling fast CPU inference and dramatically reducing parameter counts. These findings demonstrate that aligning foundation models with domain concepts can enhance both predictive power and mechanistic interpretability for scientific applications.

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.

Paper Structure

This paper contains 12 sections, 15 equations, 7 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Our Framework: MoleX
  5. Constructing and Decomposing LLM Representations
  6. Learning a Linear Mapping with Residual Calibration
  7. Experiments
  8. Experimental Setup
  9. Results and Analysis
  10. Ablation Studies
  11. Scaling the Linear Predictor
  12. Conclusion and Discussion

Figures (7)

  • Figure 1: The overview of MoleX: (1) employ ChemBERTa-2 as the foundation model, (2) fine-tune it using a tailored objective to retain maximal task-relevant semantics, (3) decompose and project latent representations onto chemically meaningful concepts, (4) train a linear predictor on the derived concepts, (5) incorporate residual calibration to iteratively refine model predictions.
  • Figure 2: The molecular representations are decomposed into a linear mapping of functional groups.
  • Figure 3: Illustration of revealed functional groups. (Left) Explanation visualization across different methods, with functional group attribution offered by MoleX. (Right) The 3D UMAP displays well-separated clusters for decomposed functional groups, indicating intrinsic linear structure in the LLM representations with concept-aligned separability. Rare groups form isolated clusters, while chemically related groups remain proximal yet distinguishable, reflecting linearly decodable chemical semantics.
  • Figure 4: Execution time comparison on seven datasets. MoleX achieves the best efficiency.
  • Figure 5: Effect of LLM model size on downstream performance.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Remark 4.1
  • Remark 4.2