Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AIM: Adaptive Intervention for Deep Multi-task Learning of Molecular Properties

Mason Minot, Gisbert Schneider

TL;DR

AIM tackles the problem of efficiently optimizing multiple molecular-property objectives that often conflict, especially when data are scarce. It introduces a learnable adaptive gradient-intervention policy that converts raw task gradients into an intervened update, guided by a jointly optimized augmented objective and differentiable regularizers. The approach yields data-efficient improvements over static multi-task baselines on QM9 and a frontier TPD-ADME benchmark and provides an interpretable diagnostic view of task relationships via the learned policy matrix. This combination of improved data efficiency and mechanistic interpretability offers a practical path toward robust, multi-property molecular design in scientific discovery.

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.

AIM: Adaptive Intervention for Deep Multi-task Learning of Molecular Properties

TL;DR

AIM tackles the problem of efficiently optimizing multiple molecular-property objectives that often conflict, especially when data are scarce. It introduces a learnable adaptive gradient-intervention policy that converts raw task gradients into an intervened update, guided by a jointly optimized augmented objective and differentiable regularizers. The approach yields data-efficient improvements over static multi-task baselines on QM9 and a frontier TPD-ADME benchmark and provides an interpretable diagnostic view of task relationships via the learned policy matrix. This combination of improved data efficiency and mechanistic interpretability offers a practical path toward robust, multi-property molecular design in scientific discovery.

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.

Paper Structure

This paper contains 23 sections, 3 equations, 3 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. The challenge of multi-task optimization
  4. Heuristics for mitigating gradient interference
  5. Navigating the front: from heuristics to a data-driven approach
  6. Method: Adaptive Intervention for deep Multi-task learning (AIM)
  7. Differentiable gradient intervention
  8. Guiding the policy
  9. Experiments and results
  10. Visualizing learning trajectories on a toy problem
  11. QM9 benchmark
  12. The learned policy as a diagnostic tool
  13. Quantitative performance
  14. Targeted protein degraders (TPD) ADME benchmark
  15. Discussion
  16. ...and 8 more sections

Figures (3)

  • Figure 1: (Top) Loss landscape of the 2D toy problem, showing conflicting task minima (stars) and the resulting Pareto front. (Bottom) Optimization trajectories for linear scalarization (LS), PCGrad, and AIM on the conflicted landscape.
  • Figure 2: Evolution of the learned AIM policy on QM9. The heatmaps show four diagnostic metrics for the policy early in training (Epoch 10, top) and later in training (Epoch 100, bottom). From left to right: learned conflict thresholds ($\tau$), conflict rate, raw gradient cosine similarity, and final projection weights.
  • Figure 3: Evolution of the learned AIM policy on TPD. The heatmaps show the policy's diagnostic metrics early (Epoch 10, top) and late (Epoch 300, bottom) in training. From left to right: learned conflict thresholds ($\tau$), conflict rate, raw gradient similarity, and final projection weights.