Technical Summary: Adapting Evidential Neural Networks to Test-Time Neighbor Fusion

Problem Statement

Deep learning models for molecular property prediction often provide point estimates without sufficient uncertainty quantification (UQ) to guide experimental workflows. While Bayesian and uncertainty-aware methods exist, they typically focus on calibrating prediction intervals rather than actively improving the predictions themselves. Traditional active learning approaches can refine models by acquiring new data, but they require expensive and computationally intensive retraining cycles after each acquisition.

The core challenge addressed is how to refine a trained molecular property model at test time without retraining. Specifically, the authors investigate whether structurally similar training molecules with known labels can serve as corrective signals to sharpen predictions for a query molecule. A critical limitation of existing neighbor-based approaches (e.g., $k$-NN) is the reliance on structural similarity (e.g., Tanimoto similarity of fingerprints) as a proxy for property similarity. On "rough" Quantitative Structure-Activity Relationship (QSAR) landscapes characterized by "activity cliffs," structurally similar molecules may have vastly different properties, rendering structural similarity an unreliable guide for neighbor selection.

Methodology

The paper proposes PG-EVIKAL (Property-guided Evidential Kalman), a test-time refinement framework that integrates evidential neural networks with Gaussian process (GP) neighbor fusion. The methodology proceeds in three stages:

1. Evidential Neural Network Baseline

The foundation is an evidential regression model (using AttentiveFP as the encoder) that outputs parameters for a Normal-Inverse-Gamma distribution. This allows the model to decompose uncertainty into two components from a single forward pass:

• Aleatoric uncertainty ($u_a$): Irreducible noise inherent to the structure-property relationship (e.g., experimental measurement noise, conformational flexibility).
• Epistemic uncertainty ($u_e$): Model ignorance due to limited training data, which can be reduced by incorporating relevant observations.

2. Neighbor Fusion via Kalman Filtering

The framework treats the training set as a source of corrective signals. For a query molecule $q$, the evidential model provides a prior mean ($\mu_0$) and prior variance ($P_0 = u_e$). Neighboring training molecules are treated as noisy observations of the query's true property.

• Noise Model: The observation noise $R_k$ for a neighbor $k$ is modeled as $R_k = u_a + C(1 - \text{sim}_k)^2$, where $\text{sim}_k$ is the Tanimoto similarity. This ensures that even a structurally identical neighbor cannot reduce uncertainty below the aleatoric floor, while structurally distant neighbors contribute higher noise.
• Fusion Mechanism:
  • EVIKAL: A scalar Kalman filter that processes neighbors sequentially. It assumes neighbors are independent, which can lead to over-counting correlated evidence.
  • GP-EVIKAL: A batch Gaussian process approach that processes all $K$ neighbors simultaneously. It uses a Tanimoto covariance kernel to explicitly model correlations between neighbors, down-weighting redundant structural information.

3. Property-Guided Neighbor Selection (PG-EVIKAL)

The authors identify that structural similarity alone is insufficient for selecting informative neighbors. To address this, they introduce PropDist, a lightweight Multi-Layer Perceptron (MLP) trained to predict the absolute property distance between molecule pairs based on their fingerprint features (bitwise AND and XOR of ECFP4 fingerprints).

• Two-Stage Selection: At inference, the top 500 structurally similar neighbors are retrieved via Tanimoto similarity. PropDist then re-ranks these candidates by their predicted property relevance.
• Final Fusion: The top $K$ neighbors (selected by the re-ranked property relevance) are fused using the GP-EVIKAL posterior update. This ensures that the fusion algorithm receives neighbors whose labels are consistent with the assumed noise model.

Key Contributions

1. PG-EVIKAL Framework: The introduction of a property-guided neighbor selection mechanism that learns a property-distance metric to re-rank structurally similar candidates. This addresses the "activity cliff" problem where structural similarity does not guarantee property similarity.
2. Theoretical Unification: The work formalizes the connection between evidential uncertainty decomposition and test-time neighbor fusion. It demonstrates that epistemic uncertainty serves as a prior that can be refined by neighbors, while aleatoric uncertainty sets a hard floor on achievable precision.
3. Diagnostic Metric (QSAR Smoothness): The authors define a QSAR smoothness ratio ($s$) to diagnose when neighbor fusion is likely to succeed. They establish an empirical threshold ($s < 0.65$) where structural neighbors reliably provide corrective signal, and show that performance degrades on rough landscapes or when structural coverage is sparse.
4. Comparison of Fusion Algorithms: The paper systematically compares sequential Kalman filtering (EVIKAL), batch Gaussian processes with correlated neighbors (GP-EVIKAL), and property-guided selection (PG-EVIKAL), identifying neighbor selection quality as the primary bottleneck for refinement accuracy.

Results

The methods were evaluated on 16 molecular datasets spanning solubility, lipophilicity, kinase inhibition, receptor binding, toxicity, and quantum chemistry targets.

• Accuracy Improvements: PG-EVIKAL reduced the Root Mean Squared Error (RMSE) relative to the evidential baseline on 14 of 16 datasets. The median reduction was 19.4%.
  • Significant gains were observed on smooth datasets (e.g., CDK2: $-42.4\%$, BACE: $-38.9\%$, LD50: $-38.2\%$).
  • On "accuracy ceiling" datasets (QM8, QM9) where the base model already captures nearly all learnable information from its topological representation, gains were negligible or non-existent.
  • On the FreeSolv dataset, the method failed due to extreme structural sparsity (low Tanimoto similarity between query and nearest neighbors), highlighting that structural coverage is a prerequisite for success.
• Calibration: PG-EVIKAL improved calibration (measured by Prediction Interval Coverage Probability at 90%, PICP@90%) on 14 of 16 datasets. By selecting property-relevant neighbors, the method avoided the over-confidence often seen in simple neighbor averaging or Tanimoto-only selection.
• Sequential Assay Scenarios: In a simulated sequential assay setting (incorporating new data without retraining), PG-EVIKAL matched or exceeded the performance of full model retraining on smooth landscapes, particularly when seed datasets were small.
• Computational Cost: The method adds approximately $3\times$ overhead to a single model forward pass (dominated by neighbor lookup and GP solve), which is significantly cheaper than Monte Carlo Dropout or deep ensembles.

Significance and Claims

The paper argues that evidential uncertainty decomposition is not merely a calibration objective but an actionable inference resource. By treating the training set as a source of corrective signals at test time, practitioners can refine molecular property predictions without the computational cost or data requirements of retraining.

The authors emphasize that the success of this approach is contingent on two factors:

1. QSAR Smoothness: The underlying structure-property landscape must be smooth enough that structural neighbors are informative.
2. Neighbor Quality: Structural similarity must be supplemented by property relevance. The paper demonstrates that the fusion algorithm (GP vs. Kalman) is secondary to the quality of neighbor selection; without property-guided selection, even sophisticated fusion algorithms fail to improve accuracy on complex landscapes.

This work provides a practical mechanism for improving molecular property predictions and uncertainty calibration entirely at test time, offering a viable alternative to active learning for scenarios where retraining is infeasible or data is scarce.