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 ($\boldsymbol{\beta}=\log\boldsymbol{\tau}$) that is intended to remove router burn-in, and two-stage self-supervised graph pretraining. On scaffold-split BACE ($\approx1.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 $\approx0.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. These three modalities recur across various molecular property-prediction endpoints although their individual importance varies. TAME is a single interpretable recipe that delivers the flexibility for different endpoints, turning volatile low-data predictors into reliable, self-explaining surrogates, wherever labels are scarce and every failed prediction costs a real assay or simulation.
Chat is not available.
Successful Page Load