Distribution Learning for Molecular Regression

TL;DR

This paper introduces DMoE, a distributional regression method for molecular property prediction that predicts probability distributions of targets, addressing biases in existing soft target methods and improving performance across multiple datasets and models.

Contribution

The paper proposes Distributional Mixture of Experts (DMoE), a novel, model- and data-independent regression approach that predicts target distributions and mitigates biases in soft target methods.

Findings

DMoE outperforms baseline methods on all evaluated datasets.

The proposed loss function enhances robustness to biases.

Evaluation across multiple architectures demonstrates versatility.

Abstract

Using "soft" targets to improve model performance has been shown to be effective in classification settings, but the usage of soft targets for regression is a much less studied topic in machine learning. The existing literature on the usage of soft targets for regression fails to properly assess the method's limitations, and empirical evaluation is quite limited. In this work, we assess the strengths and drawbacks of existing methods when applied to molecular property regression tasks. Our assessment outlines key biases present in existing methods and proposes methods to address them, evaluated through careful ablation studies. We leverage these insights to propose Distributional Mixture of Experts (DMoE): A model-independent, and data-independent method for regression which trains a model to predict probability distributions of its targets. Our proposed loss function combines the cross entropy between predicted and target distributions and the L1 distance between their expected values to produce a loss function that is robust to the outlined biases. We evaluate the performance of DMoE on different molecular property prediction datasets -- Open Catalyst (OC20), MD17, and QM9 -- across different backbone model architectures -- SchNet, GemNet, and Graphormer. Our results demonstrate that the proposed method is a promising alternative to classical regression for molecular property prediction tasks, showing improvements over baselines on all datasets and architectures.