[Paper Review] Non-convex Distributionally Robust Optimization: Non-asymptotic Analysis

Jikai Jin, Bohang Zhang+2 more|arXiv (Cornell University)|Oct 24, 2021

Sparse and Compressive Sensing Techniques33 references4 citations

TL;DR

This paper proposes a non-convex distributionally robust optimization (DRO) framework with non-asymptotic convergence guarantees for general smooth non-convex losses. By leveraging a mini-batch normalized gradient descent with momentum on a penalized DRO objective, the method achieves an $\mathcal{O}(\epsilon^{-4})$ gradient complexity to find an $\epsilon$-first-order stationary point, even under unbounded gradient noise and non-convexity.

ABSTRACT

Distributionally robust optimization (DRO) is a widely-used approach to learn models that are robust against distribution shift. Compared with the standard optimization setting, the objective function in DRO is more difficult to optimize, and most of the existing theoretical results make strong assumptions on the loss function. In this work we bridge the gap by studying DRO algorithms for general smooth non-convex losses. By carefully exploiting the specific form of the DRO objective, we are able to provide non-asymptotic convergence guarantees even though the objective function is possibly non-convex, non-smooth and has unbounded gradient noise. In particular, we prove that a special algorithm called the mini-batch normalized gradient descent with momentum, can find an $ε$ first-order stationary point within $O( ε^{-4} )$ gradient complexity. We also discuss the conditional value-at-risk (CVaR) setting, where we propose a penalized DRO objective based on a smoothed version of the CVaR that allows us to obtain a similar convergence guarantee. We finally verify our theoretical results in a number of tasks and find that the proposed algorithm can consistently achieve prominent acceleration.

Motivation & Objective

• To close the gap in theoretical understanding of DRO for non-convex, smooth loss functions with unbounded gradient noise.
• To provide non-asymptotic convergence guarantees for DRO algorithms under general $\psi$-divergence uncertainty sets.
• To address the challenges of non-convexity and unbounded stochastic gradient variance in DRO optimization.
• To extend convergence analysis to the conditional value-at-risk (CVaR) setting using a smoothed DRO formulation.

Proposed method

• Proposes a penalized DRO objective using $\psi$-divergence to handle distributional uncertainty.
• Introduces a smoothed version of CVaR to enable non-asymptotic analysis in non-convex settings.
• Employs mini-batch normalized gradient descent with momentum to stabilize training and improve convergence.
• Exploits a generalized smoothness condition on the DRO objective to bound gradient variance.
• Uses a dual formulation via convex conjugate $\psi^*$ to derive gradient bounds and smoothness properties.
• Establishes convergence via bounding the variance of stochastic gradients using the true gradient norm.

Experimental results

Research questions

• RQ1 Can non-asymptotic convergence be established for DRO with general smooth non-convex losses and unbounded gradient noise?
• RQ2 Does the proposed normalized SGD with momentum achieve a provable gradient complexity for finding $\epsilon$-stationary points in non-convex DRO?
• RQ3 Can the conditional value-at-risk (CVaR) DRO formulation be smoothed to allow non-asymptotic analysis?
• RQ4 What is the gradient complexity of the proposed algorithm in terms of $\epsilon$, $\lambda$, $\alpha$, and problem parameters?
• RQ5 How does the method perform empirically on imbalanced and robust classification tasks?

Key findings

• The proposed mini-batch normalized gradient descent with momentum achieves an $\mathcal{O}(\epsilon^{-4})$ gradient complexity to find an $\epsilon$-first-order stationary point in non-convex DRO.
• The DRO objective satisfies a generalized smoothness condition, enabling convergence analysis despite non-convexity and non-smoothness.
• For the smoothed CVaR formulation, the gradient complexity is $\mathcal{O}(\alpha^{-3}\lambda^{-1}G^{2}(G^{2}+\lambda L)\Delta\epsilon^{-4})$, where $\Delta$ is the initial optimality gap.
• Empirical results on imbalanced CIFAR-10 show consistent acceleration, with smoothed CVaR achieving 90.2% test accuracy on the worst-performing class.
• The method outperforms standard SGD and baseline DRO methods, especially on minority classes, demonstrating robustness under distribution shift.
• Theoretical analysis confirms that gradient variance is bounded by the true gradient norm, enabling stable optimization under unbounded noise.