Unveiling Molecular Secrets: An LLM-Augmented Linear Model for Explainable and Calibratable Molecular Property Prediction

We sincerely appreciate your constructive feedback and your insightful concerns. We are glad that you find our study "state-of-the-art" and "exhaustive experiments". Below we will address your questions in detail.

W1, Q1: Add an algorithm to explain the workflow

We appreciate your suggestion to use an algorithm to explain the workflow. We add pseudo code describing the procedure of the LLM fine-tuning and the training and inference of MoleX in the Appendix. Due to the limited space and the length of the algorithm, we apologize that it is difficult to include them in the main part of the paper. Considering its importance in facilitating readers' understanding of MoleX's mechanism, we have included it as the first part of the Appendix. The manuscript has been updated accordingly.

W2: Revision to math descriptors.

We sincerely apologize for the misused or unclear math descriptors as we explained some of them in the Appendix. Thank you for pointing them out. We have fixed them with detailed explanations and checked our manuscript over and over again. We now updated our manuscript.

Q2: The reason why residual calibration improves the explanation accuracy.

This is a very good question touching the core of our design. We introduce the residual calibration to compensates for the expressiveness gap between linear models and complex LLM embeddings. Here, the calibrator is designed to fix the prediction errors stemming from the linear model's underfitting limitation.

Note that the Group SELFIES, represented by functional groups, serves as the input data. Residual calibration enables MoleX to iteratively correct mispredicted functional groups and their interactions, aligning explanations with chemical expertise. By progressively calibrating prediction errors toward the ground truth, the method ensures explanations use chemically accurate functional groups and interactions. This process enhances both the accuracy (in terms of AUC) and the faithfulness (in terms of chemical expertise) of explanations.

Thank you very much for the response. After reading the response and the manuscript again, I could not understand the definitions of $f_H$ and $f_R$.

• In the pseudo code added by the authors, these are given without any explanations.
• Line 288 says $f_H$ is defined in Equation 3.1, but it defines the objective function to learn $r$. I would like to ask the authors to clarify the definitions.

For the pseudo code, since it is the only explanation on the overall process of the proposed method, I consider it is an essential part of the paper that should be included in the main body. "Call for papers" says that,

Authors may use as many pages of appendices (after the bibliography) as they wish, but reviewers are not required to read the appendix.

and I understand it is mandatory to put all the essential materials in the main body, not in Appendix (otherwise, the page limit is meaningless). Therefore, I would like to request the authors to revise the paper so that the main body is self-contained.

We thank your prompt feedback. Your questions are addressed as follows.

1. $f_H(x)$ and $f_R(x)$ are dimensionality-reduced features used to train the explainable model $h$ and residual calibrator $r$, respectively. $f_H(x)$ captures explainable features related to $y$, while $f_R(x)$ captures the remaining features with residual variance. As the orthogonality condition given, they satisfy $\langle f_H(x), f_R(x) \rangle = 0$. In the molecular setup, $f_H(x)$ represents functional group features critical to $y$, fitting by $h$ to explain their relationship, while $f_R(x)$ captures residual variance not explained by $h$. We have added explanations into our pseudo code.

2. Apologies for the lack of clarity. We have corrected the $r$ notation in Equation 3.1, as both $h$ and $r$ are learned using this objective. In our implementation, the model first learns $h$, followed by $r$ (which learns residuals of $h$). The learning of $r$ involves residual calibration, updating the parameters of both $h$ and $r$ to improve overall model performance.

We have revised the paper and included the pseudo code in the main body. Modifications are highlighted in blue in the updated manuscript. We appreciate your invaluable suggestions again.

We sincerely appreciate your constructive feedback and your insightful concerns. We are glad that you find our study "outperforms existing baselines" and "holistic experiment evaluation". Below we will address your questions in detail.

Since the mentioned W1, W2, and Q pose the same concern, we would clarify it here. In summary, our technical innovation and contributions involve:

1. Innovative adaptation.
   We effectively adapt/refine and integrate the established approaches to achieve state-of-the-art explanation and predictive accuracy, even compared with advanced neural nets and LLMs.
   (as commented by reviewer P2P1)

2. Novel design of residual calibration.
   We introduce a novel residual calibration module that compensates for the expressiveness gap between linear models and complex LLM embeddings, effectively restoring predictive performance while preserving explainability through iterative prediction error recapturing.
   (as commented by reviewers 3Cvr, P2P1, and CoqY)

3. A leap forward in explainable molecular property prediction.
   Linear models offer high explainability for their simple architecture but struggle with challenging tasks. However, we design novel approaches to augment them with LLM knowledge and residual calibration, enabling them to handle complex molecular data. Our strong performance marks a significant advancement in explainable molecular property prediction.
   (as commented by reviewers 3Cvr and P2P1)

4. Chemical expertise-aligned explanations.
   We explain molecular properties with chemically meaningful functional groups and their interactions, highly aligning with chemical domain knowledge.
   (as commented by reviewers wPWb and CoqY)

5. Holistic analysis.
   We provide exhaustive theoretical and empirical justification to demonstrate our framework.
   (as commented by reviewers 3Cvr and CoqY)

While individual components such as linear models and IB are established, our key innovation consists of adapting and integrating them to create an explainable linear model enhanced by LLM embeddings for molecular property prediction. For instance, our adaptation of the variational information bottleneck is uniquely tailored to extract task-relevant knowledge during LLM embeddings in molecular contexts, which has not been previously explored. Additionally, we propose EFPCA to effectively extract the most relevant information from the complex LLM embedding space.

The novelty of this work has also been acknowledged by other reviewers. We received comments such as "stands out for its innovative approach/innovative" (reviewers P2P1 and wPWb), and "creative solution" (reviewer P2P1).

continue

• Line 308

  Are $\boldsymbol{f_R}$ and $\boldsymbol{f_H}$ outputs of EFPCA and why are they different?

  Thanks for your question. Yes, the $f_H$ and $f_R$ are outputs of EFPCA (as definitions are given in the preliminaries section), and they are different because EFPCA applies dimensionality reduction to the feature vector $f(x)$, transforming it into a set of orthogonal principal components (PCs). We construct $f_H(x)$ by selecting the top PCs that are explainable with respect to the molecular property $y$, while $f_R(x)$ consists of the remaining principal components capturing residual variance. Since these PCs are orthogonal, $f_H(x)$ and $f_R(x)$ are distinct and satisfy $\langle f_H(x), f_R(x) \rangle = 0$, resulting in EFPCA producing different outputs from the same input $f(x)$. Specifically, the vector $f_H(x)$ contains features corresponding to functional groups that contribute significantly to $y$. These features are used by the explainable model $h$ to interpret the relationship between functional groups and $y$. The vector $f_R(x)$ comprises the remaining features, representing additional variance in $y$ not captured by $h$.

  How does EFPCA produce distinct outputs from the same input?

  This is an inspiring question. Note that EFPCA ensures that all principal components are orthogonal ($\langle \text{PC}_i, \text{PC}_j \rangle = 0$ for $i \ne j$), so $f_H(x)$ and $f_R(x)$ are orthogonal, satisfying $\langle f_H(x), f_R(x) \rangle = 0$. This orthogonality allows $h$ and $r$ to operate on different subsets of the transformed functional group features without interfering with each other. Thus, $f_H(x)$ and $f_R(x)$ are distinct both in terms of their feature content and their roles in modeling $y$, enhancing model performance without compromising explainability.

  Since $f_H(x)$ and $f_R(x)$ are different components of $f(x)$, EFPCA produces distinct outputs from the same input $f(x)$ by decomposing the feature space into these orthogonal subspaces. By projecting $f(x)$ onto the subspace spanned by the selected PCs, we obtain $f_H(x)$, and projecting onto the remaining orthogonal subspace yields $f_R(x)$. Therefore, EFPCA generates $f_H(x)$ and $f_R(x)$ as distinct outputs corresponding to explainable and residual features, respectively.

  • Inconsistent Notation

    We appreciate your advice. We have offered more details and explanations to unclear or non-explained math notations in our updated manuscript.

• W3: Omission of LLM-Based Method Comparisons in Table 2

  We sincerely thank your reminders and will clarify this in our updated manuscript. We did not offer consistent evaluation metrics in Table 2 for LLMs because AUC, a widely-used metric for explanation accuracy, is inappropriate for LLMs' continuous, structured text outputs: LLMs' probabilistic distributions over large vocabularies are incompatible with the binary classification framework required for AUC computation.

We sincerely appreciate your constructive feedback and your insightful concerns. We are glad that you find our study "high interpretability", "innovative", and "efficient, reduced complexity". Below we will address your questions in detail.

• W1: Lack of Comparison with Relevant Explainability Studies.

  Thanks for your suggestions. However, our objective is to maximize task-relevant information during LLM fine-tuning, which is realized with information bottleneck. In contrast, SGIB and VGIB are proposed to identify explainable subgraphs in GNNs via information bottleneck. The purpose of using IB in MoleX differs from the mentioned papers, and comparing MoleX with SGIB and VGIB might deviate from the main idea of this study.

• W2, Q1-4: Mathematical Explanation.

  We apologize for missing these details. Since Q1-4 specify the confusion in W2, we would elaborate our explanations to these questions here. Your questions are addressed as follows.

  • Line 261. More Explanations

    To explain how the $\ell_0$ penalty, the coefficients $a_{kj}$, and the basis functions $\phi_j(t)$ are derived in the EFPCA, we explain each component respectively:

    1. The basis functions $\phi_j(t)$

       The basis functions $\phi_j(t)$ are derived by selecting a set of functions that have local support on subintervals $S_j \subset [a, b]$. Typically, we use B-spline basis functions because they are piecewise polynomial functions that are nonzero only within specific intervals. This local support property enables the principal components $\xi_k(t)$ to capture localized features in the data, which enhances explainability.

    2. The coefficients $a_{kj}$

       The coefficients $a_{kj}$ are derived through the expansion of the principal components $\xi_k(t)$ in terms of the chosen basis functions. We express $\xi_k(t)$ as a linear combination $\xi_k(t) = \sum_{j=1}^p a_{kj} \phi_j(t)$ where $p$ is the number of basis functions. The coefficients $a_{kj}$ are determined by solving the optimization problem posed in the EFPCA, which aims to maximize the variance explained by $\xi_k(t)$ while inducing sparsity.

    3. The $\ell_0$ penalty $\rho_k \|| a_k \||_0$

       The $\ell_0$ penalty $\rho_k \|| a_k \||_0$ is derived from the need to encourage sparsity in the coefficients $a_{kj}$. Here, $\|| a_k \||_0$ counts the number of nonzero coefficients in principal components By including the $\ell_0$ penalty into our designed optimization objective, we can impose a penalty proportional to the number of nonzero coefficients. This promotes solutions where many $a_{kj}$ are exactly zero, resulting in principal components that are sparse and focused on the most significant features of the data.

    Consequently, the derivation involves:

    • Selecting basis functions $\phi_j(t)$ with local support to enable explainability.
    • Expressing $\xi_k(t)$ as a linear combination of these basis functions and deriving the coefficients $a_{kj}$ through optimization.
    • Introducing the $\ell_0$ penalty to induce sparsity in $a_{kj}$, which enhances the explainability of the resulting principal components.

  • Input Consistency

    Great question. The explainable model and residual calibrator share the same input: dimensionality-reduced LLM embeddings. MoleX's workflow involves three steps:

    1. Feeding the input into the explainable linear model and collecting prediction error samples.
    2. Inputting the prediction errors into the residual calibrator (also a linear model).
    3. Aggregating the calibrator's output with the explainable linear model.

    During its workflow, the calibrator recaptures prediction errors and simultaneously updates the entire model's parameters, ensuring consistent input throughout the process.

The definition of $\|a_k\|$ is clear to me. However, every symbol appearing in the main text must be explicitly defined within the main text itself, rather than relying on definitions provided in the appendix. Additionally, the definition for EFPCA is directly adopted from [1], but it is inconsistent with the notational framework used in this paper. Ensuring consistent notation throughout the paper is essential for clarity and readability. For instance, the relationship between $(f_H, f_R)$, and $\xi$ in EFPCA should be explicitly clarified when $f_H$ and $f_R$ first appear at Line 172. Alternatively, you could specify that these elements will be discussed in detail in specific sections (e.g., Section XX and XX).

The definition of $t$ also presents issues of clarity and consistency. In EFPCA, $t$ represents the independent variable within a compact interval $[a, b]$. However, in Section 4.1 of this paper, $t$ is firstly defined as LLM embeddings, while in the algorithm section, these embeddings are instead denoted as $x^{\text{emb}}$. Moreover, the notation for the $i$-th group SELFIES is inconsistent: it is defined as $x^{(i)}$ at Line 168 but appears as $x_i$ in Algorithm 1.

Such inconsistencies in notation and definitions detract from the overall coherence and readability of the paper. The paper currently feels fragmented, resembling a collection of disjointed sections rather than a cohesive, self-contained work. A major revision is still needed to improve the overall presentation, ensure notational consistency, and enhance the clarity of the paper.

[1] Zhenhua Lin, Liangliang Wang, and Jiguo Cao. Interpretable functional principal component analysis. Biometrics, 72(3):846–854, 2016.

BTW, I am satisfied with the feedback for W1 and W3.

continue

• W1 and W2: More experimental evaluations on the effectiveness of MoleX.

  We appreciate your suggestions on our empirical results. In response to your feedback, we conducted additional evaluations of MoleX on three widely used datasets spanning diverse domains and tasks in molecular property prediction. The results of classification and explanation accuracy are shown as follows and have been incorporated into our updated manuscript.

Classification accuracy over three datasets (%)

Explanation accuracy over three datasets (%)

• W3: Suggestion of future works.

  We apologize for the confusion. We have discussed potential future works at the end of the Appendix. This section includes the exploration of framework generalizability, the diversity of LLMs, and the trade-off between complexity and performance. We hope this discussion will shed light on further research based on our study.

We sincerely appreciate your constructive feedback and your insightful concerns. We are glad that you find our study "well-written, well-organized", "clear and concise overview of the framework, its theoretical background", "a clear motivation", and "better interpret the result". Below we will address your questions in detail.

• Q1: Limitation of linearity assumption.

  This is an insightful question closely tied to the core of our design. Indeed, the linearity assumption has two inherent limitations, which we address with proposed solutions:

  • Linear models' lack of expressiveness can lead to underfitting on molecular data. We propose residual calibration to iteratively calibrate prediction errors made by the explainable model. Experiments show that the calibrator improves model performance by more than 12%.
  • MoleX's success depends on the LLM's ability to capture significant features during linear model fitting. We thus propose information bottleneck-inspired fine-tuning and sparsity-inducing dimensionality reduction to ensure task-relevant information is maximally maintained for linear model fitting.

  Our techniques resolve limitations and augment the linear model, with extensive experiments validating its effectiveness and robustness.

• Q2: Interpretation of n-grams.

  This is a good question. N-gram features represent n interconnected functional groups and their interactions, providing chemical contextual semantics about molecules. For example, in the Mutag dataset, 3-grams include functional groups such as benzene, nitro, and interactions such as pop and branch. This substructure has been validated to contribute to a molecule's mutagenicity.

• Q3: Robustness of MoleX to noises.

  We appreciate your inspiring question. Molecular data potentially contain noises [1-3], which we accounted for when developing MoleX. Additionally, noise can also arise from LLM embeddings. To mitigate this, we propose EFPCA and IB-inspired fine-tuning to eliminate noisy information, enabling the linear model to focus only on informative features and ensuring effective feature contribution computation. Experiments across ten datasets—spanning mutagenicity, carcinogenicity, drug toxicity, and permeability—consistently demonstrate MoleX's robustness.

• Q4: Theoretical foundations.

  We thank your suggestion to clarify. The term "theoretically" refers to the theoretical analysis of MoleX's explainability in the Appendix. This analysis is grounded in the mathematical foundations of linear models, convex optimization, and information bottleneck theories. These principles, while theoretical in nature, are adapted to our context through validated corollaries and lemmas, supported by mathematical proofs.

References

[1] Hillis, D. M., and J. P. Huelsenbeck. "Signal, noise, and reliability in molecular phylogenetic analyses." Journal of Heredity 83.3 (1992): 189-195.

[2] Plesa, Tomislav, et al. "Noise control for molecular computing." Journal of the Royal Society Interface 15.144 (2018): 20180199.

[3] Ioannidis, John PA. "Microarrays and molecular research: noise discovery?" The Lancet 365.9458 (2005): 454-455.

We sincerely appreciate your constructive feedback and your insightful concerns. We are glad that you find our study "meaningful contribution", "innovative", and "technical rigor, clarity, and significant practical implications". Below we will address your questions in detail.

• W, Q2: Difference of MoleX between prototype-based approaches and our advantages.

  We appreciate your concerns about the difference among explainability approaches. MoleX differs from prototype-based approaches, such as ProtoPNet and xDNN, in several aspects and holds some advantages over them.

  • Global vs. local explanations. MoleX provides global explanations by modeling the data distribution over the entire feature space, while prototype-based approaches offer case-based, post-hoc explanations. Global explainability better reveals the model's overall behavior, making MoleX more capable of interpreting the entire variability in complex molecular structures.

  • Explicit explainability. MoleX's explainability arises from the linearity of the model's coefficients, where individual features' contributions provide explanations. In contrast, prototype-based methods rely on similarity to real examples without explicitly articulating the key contributing features, making the precise reasoning harder for humans to understand. Thus, MoleX offers more explicit explanations behind predictions.

  • Chemical expertise-aligned explanations. In the context of our study, we explain the molecular properties with chemically meaningful functional groups and their interactions, highly aligning with chemical expertise.

  • Innovative design to enhance explainability. We introduce a residual calibration strategy that improves both the faithfulness and accuracy of explainability via iterative prediction error fixing.

• Q1: The evaluation of faithfulness from a chemical standpoint.

  This is a good question. We evaluate how faithful the explanations offered by MoleX are through chemical expertise to highlight its explainability (i.e., how correct the explanations are from a chemical perspective). As mentioned in the Appendix, domain experts marked correct functional groups that influence molecular properties as ground truth. Empirically, MoleX offers explanations highly aligned with these ground truths.

• Q3: More details about residual calibration strategy.

  These are very great questions. We are glad to summarize more details about residual calibration to resolve your confusion.

  • High-level ideas. We introduce residual calibration to address the linear model's limited expressiveness with high-dimensional LLM embeddings. It refits prediction errors caused by underfitting, with a training objective designed to improve overall performance by iteratively calibrating the explainable model's residuals (i.e., update both models' parameters during prediction error calibration).

  • How residual calibration improves explainability. Note that the Group SELFIES, represented by functional groups, serves as the input data. Residual calibration enables MoleX to iteratively correct mispredicted functional groups and their interactions, aligning explanations with chemical expertise. By progressively calibrating prediction errors toward the ground truth, the method ensures explanations use chemically accurate functional groups and interactions. This process enhances both the accuracy (in terms of AUC) and the faithfulness (in terms of chemical expertise) of the explanations.

  • Implementation details. The workflow mainly involves three steps:

    1. Feeding the input into the explainable linear model and collecting residual samples.
    2. Feeding residual samples into the residual calibrator (also a linear model).
    3. Aggregating the output of both models. The training objective updates both models' parameters simultaneously while fixing mispredicted samples.