TAME: Element-wise Mixture-of-Experts Fusion for Reliable and Interpretable Molecular Property Prediction
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 (β = log τ ) that is intended to remove router burn-in, and two-stage self-supervised graph pretraining. On scaffold-split BACE (≈ 1.5k 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 σ 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 ≈ 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.
