TAME: Element-wise Mixture-of-Experts Fusion for Reliable and Interpretable Molecular Property Prediction
Robert Schiller ⋅ Christoph Weisser ⋅ Klaus-Robert Müller ⋅ ⋅ Parastoo Semnani
Abstract
Property prediction for molecules and materials is bottlenecked by label scarcity, and in this regime the operative failure mode is not a low mean but an unpredictable tail: individual training runs collapse. The three dominant molecular representations fail in complementary ways — physicochemical descriptors are exact but fixed, graph neural networks learn topology but degenerate under sparse supervision, and language models supply semantic context but cannot compute exact quantities — yet combining them is itself the learning problem, because modality-level scalar gating commits the whole model to one trust weight per source. We introduce **TAME**, a tri-modal encoder whose *element-wise* Mixture-of-Experts router assigns an independent modality mixture to every hidden coordinate, stabilized by a balance–entropy regularizer, a closed-form gate initialization ($\beta = \log \tau$) that is intended to remove router burn-in, and two-stage self-supervised graph pretraining. On scaffold-split BACE ($\approx 1.5$k molecules, 100 seeds) we evaluate two fusion topologies, each against its own graph-only control. In both, adding the text and descriptor experts contracts the seed-to-seed distribution of the threshold-free metrics — ROC-AUC $\sigma$ falls five-fold in the flat topology and 1.6-fold in the hierarchical one, whose control was already the tighter of the two — and the collapse tail reaching ROC-AUC $\approx 0.45$ that the weaker single-modality configurations carry is absent from the fused models while accuracy is retained. Each comparison is matched on everything but the fusion stage, isolating its effect. The router allocates a distinct mixture to each coordinate, which no scalar gate can represent, and the entropy weight moves that allocation continuously between graded mixing and a near-binary regime in which each coordinate is claimed outright by one modality. Fused-representation alignment tracks encoder quality without supervision, shifting from text to topology once pretraining lifts the graph embedding out of rank-1 collapse. The same three experts map onto crystal graphs, composition–structure descriptors and literature-mined text. These three modalities recur unchanged across molecular and inorganic chemistry, so TAME is a single interpretable recipe for turning volatile low-data predictors into reliable, self-explaining ones, wherever data are scarce and every failed prediction costs a real experiment.
Chat is not available.
Successful Page Load