Technical Summary: Distributionally Robust Multi-Objective Optimization

1. Problem Formulation

The paper addresses Multi-Objective Optimization (MOO) in settings where data distributions are subject to shifts, a common issue in modern multi-task learning (MTL) where shared representations are vulnerable to domain mismatches, class imbalances, or adversarial perturbations. Standard MOO formulations optimize objectives under fixed nominal distributions $\{P_i\}_{i=1}^m$, which may fail to remain optimal when evaluated under shifted distributions.

To address this, the authors introduce Distributionally Robust Multi-Objective Optimization (DR-MOO). The problem is formulated as minimizing a vector of objectives, where each objective $i$ is an $f$-divergence regularized Distributionally Robust Optimization (DRO) risk:
$$\min_{\theta \in \mathbb{R}^n} \sup_{Q_1, \dots, Q_m} \left[ H_1(\theta, Q_1), \dots, H_m(\theta, Q_m) \right]$$
where $H_i(\theta, Q_i) = \mathbb{E}_{\xi \sim Q_i}[\ell_i(\theta, \xi)] - \lambda D_f(Q_i \| P_i)$. Here, $D_f$ denotes an $f$-divergence measuring the discrepancy between the nominal distribution $P_i$ and the worst-case distribution $Q_i$.

Key Challenges:

1. Solution Concepts: Standard Pareto optimality does not directly apply due to distributional uncertainty. The paper defines Distributionally Robust Pareto Dominance and Optimality, comparing parameters after adversarial perturbations of each objective's distribution.
2. Algorithmic Design: Existing MOO algorithms (e.g., MGDA) rely on fixed nominal distributions. Applying them directly to DR-MOO is non-trivial because the dual reformulation of DRO objectives alters the geometry, potentially violating assumptions like uniformly bounded gradients. Furthermore, estimating the dual variables required for the gradient introduces bias and computational overhead.

2. Methodology

The authors propose a dual-based approach to transform the DR-MOO problem into a standard MOO form amenable to first-order methods.

Dual Reformulation

Using Lagrangian duality, the primal DR-MOO problem is rewritten as an equivalent dual MOO problem:
$$\min_{\theta} \Phi(\theta) := [\phi_1(\theta), \dots, \phi_m(\theta)]$$
where each dual objective $\phi_i(\theta) = \min_{\eta_i \in \mathbb{R}} \left( \lambda \mathbb{E}_{\xi \sim P_i}[f^*(\frac{\ell_i(\theta, \xi) - \eta_i}{\lambda})] + \eta_i \right)$.
The gradient of the dual objective admits a closed-form expression involving the optimal dual variable $\eta_i^*$, which is the solution to a one-dimensional convex minimization problem.

Algorithm 1: Double-Loop MGDA

To solve the dual problem, the authors first propose a Double-Loop Multi-Gradient Descent Algorithm (MGDA):

• Inner Loop: Uses Stochastic Gradient Descent (SGD) to approximately solve for the dual variables $\eta_i^*$ for each objective.
• Outer Loop: Updates the model parameters $\theta$ and the preference vector $w$ using a stochastic MGDA update.
• Bias Control: To handle the bias introduced by stochastic estimation of gradients and the inexactness of the inner loop, the outer loop employs a double-sampling strategy (drawing three independent samples per iteration) to construct unbiased estimators for the preference vector update.
• Complexity: This method achieves convergence to an $\epsilon$-distributionally robust Pareto-stationary point with a total sample complexity of $O(\epsilon^{-12})$ (specifically $O(\epsilon^{-4})$ outer iterations and $O(\epsilon^{-8})$ inner iterations), without assuming bounded gradients on the dual objectives.

Algorithm 2: Single-Loop Double-Clip MGDA

To improve efficiency and remove the high cost of the inner loop, the authors derive a Single-Loop Double-Clip MGDA:

• Surrogate Stationarity: They establish an equivalent stationarity criterion that avoids the need for accurate tracking of $\eta_i^*$.
• Gradient Clipping: The algorithm employs double gradient clipping on both the parameter update direction ($X_t w_t$) and the dual-variable update direction ($Z_t w_t$). This controls the ill-conditioned geometry and the bias induced by affine-bounded stochastic noise.
• Single-Loop Structure: By clipping gradients, the method eliminates the need for the double-sampling strategy and the inner loop, allowing for a single-loop update of $\theta$, $\eta$, and $w$.
• Complexity: This method achieves a significantly improved sample complexity of $O(\epsilon^{-4})$, matching the best-known rates for nonconvex stochastic optimization, while relaxing the requirement for bounded gradients.

3. Key Contributions

1. Formulation of DR-MOO: The paper systematically formulates DR-MOO as a multi-objective problem with $f$-divergence regularized risks. It introduces distributionally robust Pareto dominance and optimality, and defines distributionally robust Pareto stationarity via an equivalent dual formulation.
2. Theoretical Analysis of Double-Loop MGDA: The authors establish an $L$-smoothness property for the dual objectives along update trajectories. They prove that the Double-Loop MGDA converges to an $\epsilon$-stationary point with $O(\epsilon^{-12})$ sample complexity, without requiring bounded gradient assumptions on the dual objectives.
3. Efficient Single-Loop Algorithm: They propose the Double-Clip MGDA, which utilizes gradient clipping to handle generalized smoothness and biased estimates. This method removes the need for double sampling and inner-loop tracking, achieving $O(\epsilon^{-4})$ sample complexity.
4. Empirical Validation: Experiments on Multi-MNIST and CelebA datasets demonstrate that the proposed methods are competitive with state-of-the-art MGDA baselines (e.g., MoCo, SDMGrad, NashMTL) in terms of test accuracy, particularly under adversarial attacks (FGSM) and label imbalance.

4. Results

• Adversarial Robustness: On Multi-MNIST datasets, the proposed Double-Clip MGDA and Double-Loop MGDA maintained higher test accuracy under FGSM attacks compared to standard MGDA baselines. For instance, under a 0.08 attack level on 2-digit Multi-MNIST, Double-Clip MGDA achieved 57.13% accuracy, outperforming NashMTL (47.09%) and standard MGDA (48.30%).
• Label Imbalance: On the CelebA dataset with imbalanced labels, the proposed methods showed improved performance in averaged accuracy, balanced accuracy, and AUC compared to baselines, validating the robustness of the DR-MOO formulation.
• Ablation Studies: The authors analyzed the impact of hyperparameters ($\lambda$, learning rates, batch sizes, and inner-loop iterations). They found that the Double-Clip MGDA is robust to batch sizes as low as 256 and that the Double-Loop MGDA is relatively insensitive to the number of inner iterations (5 iterations were sufficient for deep learning tasks).

5. Significance and Claims

The paper claims to be the first to systematically formulate DR-MOO and provide non-asymptotic convergence guarantees for MGDA-type algorithms in a nonconvex setting with distributional uncertainty.

• Theoretical Advancement: The work extends the understanding of MOO to distributionally robust settings, showing that Pareto-stationary optimization under distributional uncertainty can be addressed through principled stochastic approximation. It challenges the necessity of bounded gradient assumptions often found in existing MOO literature by leveraging gradient-function value relationships under generalized smoothness.
• Practical Efficiency: The transition from a double-loop to a single-loop algorithm via gradient clipping offers a practical solution to the computational bottleneck of dual variable tracking, making DR-MOO more scalable for deep learning applications.
• Broader Impact: The authors suggest that their framework and the Double-Clip analysis could inspire further studies on tractable formulations for multi-objective alignment problems and extend to other normalized optimization methods (e.g., Adam, Muon) in MOO settings.

The paper maintains a modest tone regarding its contributions, acknowledging theoretical limitations such as the large-batch requirement ($O(\epsilon^{-2})$) for the single-loop method to achieve optimal complexity under affine-bounded noise, and noting that while the inner-loop length has negligible impact in practice, the theoretical bound remains high.