# Distribution Learning for Molecular Regression

arXiv: https://arxiv.org/abs/2407.20475 · HTML: https://mlsanity.com/paper/2407.20475/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 26 of its 37 references that are in the map.
- First posted on arXiv 30 Jul 2024.

## Builds on (most cited first, 26 of 26)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast R-CNN](https://mlsanity.com/paper/1504.08083/) (2015)
- [Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://mlsanity.com/paper/1509.09292/) (2015)
- [DropEdge: Towards Deep Graph Convolutional Networks on Node Classification](https://mlsanity.com/paper/1907.10903/) (2019)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [Train faster, generalize better: Stability of stochastic gradient descent](https://mlsanity.com/paper/1509.01240/) (2015)
- [Regularizing Neural Networks by Penalizing Confident Output Distributions](https://mlsanity.com/paper/1701.06548/) (2017)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [PairNorm: Tackling Oversmoothing in GNNs](https://mlsanity.com/paper/1909.12223/) (2019)
- [Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules](https://mlsanity.com/paper/2011.14115/) (2020)
- [DeeperGCN: All You Need to Train Deeper GCNs](https://mlsanity.com/paper/2006.07739/) (2020)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Geometric and Physical Quantities Improve E(3) Equivariant Message Passing](https://mlsanity.com/paper/2110.02905/) (2021)
- [GraphNorm: A Principled Approach to Accelerating Graph Neural Network Training](https://mlsanity.com/paper/2009.03294/) (2020)
- [Towards Deeper Graph Neural Networks with Differentiable Group Normalization](https://mlsanity.com/paper/2006.06972/) (2020)
- [Improving Regression Performance with Distributional Losses](https://mlsanity.com/paper/1806.04613/) (2018)
- [Benchmarking Graphormer on Large-Scale Molecular Modeling Datasets](https://mlsanity.com/paper/2203.04810/) (2022)
- [Rotation Invariant Graph Neural Networks using Spin Convolutions](https://mlsanity.com/paper/2106.09575/) (2021)
- [Molecular Mechanics-Driven Graph Neural Network with Multiplex Graph for Molecular Structures](https://mlsanity.com/paper/2011.07457/) (2020)
- [Towards Training Billion Parameter Graph Neural Networks for Atomic Simulations](https://mlsanity.com/paper/2203.09697/) (2022)
- [Graph Convolutional Neural Networks with Node Transition Probability-based Message Passing and DropNode Regularization](https://mlsanity.com/paper/2008.12578/) (2020)

## Abstract (arXiv)

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.