AIM: Adaptive Intervention for Deep Multi-task Learning of Molecular Properties
2509.25955v1
cs.LG, cs.AI, physics.chem-ph
2025-10-02
Авторы:
Mason Minot, Gisbert Schneider
Abstract
Simultaneously optimizing multiple, frequently conflicting, molecular
properties is a key bottleneck in the development of novel therapeutics.
Although a promising approach, the efficacy of multi-task learning is often
compromised by destructive gradient interference, especially in the data-scarce
regimes common to drug discovery. To address this, we propose AIM, an
optimization framework that learns a dynamic policy to mediate gradient
conflicts. The policy is trained jointly with the main network using a novel
augmented objective composed of dense, differentiable regularizers. This
objective guides the policy to produce updates that are geometrically stable
and dynamically efficient, prioritizing progress on the most challenging tasks.
We demonstrate that AIM achieves statistically significant improvements over
multi-task baselines on subsets of the QM9 and targeted protein degraders
benchmarks, with its advantage being most pronounced in data-scarce regimes.
Beyond performance, AIM's key contribution is its interpretability; the learned
policy matrix serves as a diagnostic tool for analyzing inter-task
relationships. This combination of data-efficient performance and diagnostic
insight highlights the potential of adaptive optimizers to accelerate
scientific discovery by creating more robust and insightful models for
multi-property molecular design.