Technical Summary: MAdam: Metric-Aware Multi-Objective Adam

1. Problem Statement

Multi-objective optimization (MOO) is fundamental to many machine learning tasks, including multi-task learning (MTL), physics-informed neural networks (PINNs), and medical imaging. While diverse MOO solvers exist—categorized into loss-balancing, gradient-balancing, and Pareto-based families—they almost universally delegate the execution of the update step to the Adam optimizer.

The authors identify a systematic disconnect between the intent of these MOO solvers and the execution of Adam, resulting in two specific mismatches:

1. Weighting Mismatch: In methods where the preference vector $\lambda(t)$ varies over time (e.g., loss balancing or Pareto front sweeping), Adam's second-moment exponential moving average (EMA) entangles the time-varying preferences with gradient statistics. Instead of applying the current preference $\lambda(t)$ to the gradient statistics, Adam absorbs the preference into a history-averaged quantity. This marginalizes the intended preference strategy, collapsing distinct Pareto trade-offs into a near-uniform mixture and causing the optimizer to lag behind preference switches.
2. Geometric Mismatch: MOO solvers typically assume a raw Euclidean geometry when constructing reconciled descent directions (e.g., via linear scalarization). However, Adam implicitly imposes its own adaptive diagonal metric (derived from the second-moment denominator) on every update. This metric distorts the Euclidean geometry assumed by the solver, potentially turning aligned objectives into apparent conflicts or vice versa, thereby altering the realized per-objective loss changes.

2. Methodology: MAdam

To resolve these issues without modifying the underlying MOO solver or the Adam optimizer, the authors propose MAdam (Metric-Aware Multi-Objective Adam), a "drop-in" wrapper that preconditions the solver's output.

Theoretical Foundation

The authors derive the correct curvature for MOO updates. They argue that the appropriate metric is the preference-conditioned diagonal Fisher information matrix of the scalarized objective at the active preference $\lambda(t)$.

• Let $g_i$ be the gradient of the $i$-th objective.
• The scalarized gradient is $g_\lambda = \sum \lambda_i g_i$.
• The correct curvature estimate is $C_\lambda = \mathbb{E}[g_\lambda \odot g_\lambda] = \sum_{i,j} \lambda_i \lambda_j F_{ij}$, where $F_{ij}$ represents the cross-objective gradient correlations (Fisher blocks).
• Unlike Adam, which computes $\mathbb{E}[(\sum \lambda_i g_i) \odot (\sum \lambda_j g_j)]$ and averages the product of preferences and gradients jointly, MAdam keeps the current preference $\lambda(t)$ fixed outside the expectation, allowing it to control the contribution of each cross-moment dynamically.

Algorithmic Implementation

MAdam operates as a wrapper in the solver-then-Adam pipeline:

1. Solver Step: The MOO solver produces a reconciled direction $d(t) = \sum \lambda_i(t) g_i(t)$.
2. Metric Estimation: MAdam maintains online estimates of the cross-objective Fisher blocks $F_{ij}$ using EMAs. To reduce computational cost from $O(C)$ to $O(1)$ backward passes per step, it employs stochastic pair sampling, updating only the Fisher entries corresponding to a randomly sampled pair of tasks $(i, j)$ at each iteration.
3. Preconditioning: The reconciled direction is preconditioned by the inverse of the estimated metric: $\tilde{d}(t) = M_\lambda(t)^{-1} d(t)$, where $M_\lambda(t) = \text{Diag}(\sqrt{C_\lambda + \epsilon})$.
4. Optimization: This preconditioned direction $\tilde{d}(t)$ is fed into Adam.
   • Because the input is "whitened" by the preference-conditioned metric, Adam's second-moment EMA on this input collapses to the identity matrix ($M_{Adam} \approx I$).
   • Consequently, the realized update becomes $\theta(t+1) = \theta(t) - \eta M_\lambda(t)^{-1} d(t)$, governed entirely by the preference-conditioned metric rather than Adam's history-biased metric.
5. Warmup: A ramp-up coefficient $\alpha(t)$ is used during initialization to blend the preconditioner between the identity (standard Adam) and the full MAdam metric, preventing instability from noisy early Fisher estimates.

3. Key Contributions

The paper makes three primary contributions:

1. Diagnosis of Solver-Adam Mismatch: The authors formally prove two failure modes in the standard MOO pipeline: a weighting mismatch that marginalizes time-varying preferences (Proposition 1) and a geometric mismatch that distorts the Euclidean geometry assumed by solvers (Proposition 2).
2. Metric-Aware Gradient Preconditioning: They derive the preference-conditioned diagonal Fisher of the scalarized objective and implement it as a plug-and-play preconditioner. This mechanism resolves both mismatches simultaneously without altering the solver logic or the optimizer's internal state.
3. Empirical Validation: The method is validated across diverse domains, demonstrating consistent improvements over standard Adam when paired with various MOO solvers.

4. Experimental Results

The authors evaluated MAdam across four regimes: Multi-Task Learning (MTL), Pareto-front recovery, Physics-Informed Neural Networks (PINNs), and Medical Image Analysis (MIA).

• Multi-Task Learning (MTL): On benchmarks including MultiMNIST, SARCOS, UTKFace, Cityscapes, and NYUv2, MAdam consistently improved performance over base solvers (Linear Scalarization, DWA, UW, IMTL, CAGrad). For instance, on Cityscapes, it reduced the relative depth error of IMTL from 74.2 to 61.8.
• Pareto-front Recovery: When paired with Pareto-based solvers (PaMaL, PaLoRA) on MultiMNIST and SARCOS, MAdam improved the hypervolume and reduced per-ray errors, effectively tracing the Pareto front with higher diversity and accuracy.
• PINNs: On the PINNacle benchmark (20 PDEs), MAdam improved performance across Vanilla, Loss Reweighting, and adaptive activation baselines, with significant gains on geometrically complex domains and chaotic systems where loss scales vary widely.
• Medical Imaging: In skin-lesion segmentation (ISIC2018) and brain MRI super-resolution (OASIS3), MAdam achieved better perceptual fidelity (LPIPS, DISTS) with minimal degradation in pixel-wise metrics, and produced more balanced Dice/HD95 profiles.
• Ablation Studies: Removing the ramp-up mechanism caused performance to drop below the baseline, confirming the need for stable Fisher estimates. Removing off-diagonal Fisher terms ($F_{ij}$) also reduced gains, highlighting the importance of modeling cross-objective correlations.

5. Significance and Limitations

Significance:
The paper claims that MAdam provides a universal, drop-in solution to the solver-optimizer mismatch in MOO. By decoupling the preference vector from the optimizer's history and correcting the geometric distortion, it allows existing MOO solvers to execute their intended updates faithfully. It bridges the gap between the theoretical assumptions of MOO solvers (Euclidean geometry, explicit preferences) and the practical reality of adaptive optimizers.

Limitations and Future Work:
The authors acknowledge two limitations:

1. Heterogeneous Objectives: The estimation of off-diagonal Fisher blocks ($F_{ij}$) for heterogeneous objectives (e.g., classification paired with regression) is currently sub-optimal, often yielding small and noisy entries. Future work aims to develop more principled estimators for these cross-objective curvatures.
2. Scalarization Assumption: The current derivation assumes linear scalarization ($\sum \lambda_i L_i$). The authors plan to extend MAdam to handle nonlinear scalarizations, such as Tchebycheff scalarization.