Mirror Descent Under Generalized Smoothness

TL;DR

This paper extends the concept of smoothness in optimization to general norms and dual norms, enabling convergence guarantees for mirror descent algorithms in broader settings including non-Euclidean geometries and stochastic, non-convex, or composite problems.

Contribution

It introduces a new ll*-smoothness framework and a generalized self-bounding property, broadening the applicability of mirror descent beyond Euclidean spaces with theoretical guarantees.

Findings

Established convergence for mirror descent under ll*-smoothness.

Extended results to stochastic, non-convex, and composite optimization.

Matched state-of-the-art rates under classic smoothness conditions.

Abstract

Smoothness is crucial for attaining fast rates in first-order optimization. However, many optimization problems in modern machine learning involve non-smooth objectives. Recent studies relax the smoothness assumption by allowing the Lipschitz constant of the gradient to grow with respect to the gradient norm, which accommodates a broad range of objectives in practice. Despite this progress, existing generalizations of smoothness are restricted to Euclidean geometry with $\ell_2$-norm and only have theoretical guarantees for optimization in the Euclidean space. In this paper, we address this limitation by introducing a new $\ell*$-smoothness concept that measures the norm of Hessians in terms of a general norm and its dual, and establish convergence for mirror-descent-type algorithms, matching the rates under the classic smoothness. Notably, we propose a generalized self-bounding property that facilitates bounding the gradients via controlling suboptimality gaps, serving as a principal component for convergence analysis. Beyond deterministic optimization, we establish sharp convergence for stochastic mirror descent, matching state-of-the-art under classic smoothness. Our theory also extends to non-convex and composite optimization, which may shed light on practical usages of mirror descent, including pre-training and post-training of LLMs.