AIM: Adaptive Intervention for Deep Multi-task Learning of Molecular Properties
This work addresses a bottleneck in drug discovery by improving multi-task learning for molecular design, offering both performance gains and interpretability, though it is incremental as it builds on existing optimization methods.
The paper tackled the problem of destructive gradient interference in multi-task learning for molecular property optimization, particularly in data-scarce regimes, by proposing AIM, an adaptive optimization framework that achieved statistically significant improvements over baselines on QM9 and targeted protein degraders benchmarks.
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.