Adapting Evidential Neural Networks to Test-Time Neighbor Fusion Improves Molecular Property Prediction

How can the uncertainty estimates produced by evidential neural networks be leveraged to improve molecular property predictions at test time without the computational cost of retraining the model? The authors investigate whether structurally similar training molecules can act as a source of corrective information to refine a query molecule's property prediction.

The researchers introduce PG-EVIKAL, a framework that performs test-time neighbor fusion. First, an evidential neural network provides a property prediction along with aleatoric and epistemic uncertainty estimates from a single forward pass. Second, the system retrieves structurally similar neighbors from the training set using Tanimoto similarity. Third, a learned property-distance metric (PropDist) re-ranks these neighbors to prioritize those most relevant to the target property. Finally, these neighbors are fused with the initial prediction using a Gaussian process posterior, where the evidential uncertainty parameters define the noise model for the fusion process.

PG-EVIKAL demonstrates significant improvements across 16 diverse molecular datasets. It achieved a median reduction in root-mean-square error (RMSE) of 19.4% compared to the baseline evidential model on 14 of the 16 datasets. The study highlights that the bottleneck for test-time refinement is not the fusion algorithm itself, but the quality of neighbor selection. By incorporating a property-guided selection step, PG-EVIKAL outperforms simpler variants (EVIKAL and GP-EVIKAL) that rely solely on structural similarity, while also improving the calibration of the resulting uncertainty intervals.

This work shifts the role of the training set from a static resource used only during model development to an active, on-demand inference resource. For researchers working with expensive experimental assays or models that cannot be frequently retrained, this approach offers a practical, computationally efficient way to refine predictions and improve uncertainty quantification in real-time deployment scenarios.