Nonlinearly Preconditioned Gradient Methods under Generalized Smoothness

Nonlinearly Preconditioned Gradient Methods under Generalized Smoothness

Speakers: Konstantinos Oikonomidis, Jan Quan, Emanuel Laude, Panagiotis Patrinos

Conference: ICML 2025

YouTube: https://slideslive.com/39044015

Overview

This talk, presented by Konstantinos Oikonomidis and his co-authors Jan Quan, Emanuel Laude, and Panagiotis Patrinos, introduces a novel framework for nonlinearly preconditioned gradient methods designed to tackle unconstrained minimization problems. While standard gradient descent (GD) is a ubiquitous solver, its optimal performance often hinges on the restrictive assumption of global Lipschitz smoothness of the cost function. In the context of modern machine learning, this assumption frequently does not hold, leading to challenges such as arduous hyperparameter tuning, requirement for excessively small step sizes, or reliance on computationally intensive line search procedures for convergence.

The core contribution of this work is to generalize gradient descent by passing the gradient of the cost function through a carefully chosen nonlinear mapping, effectively "preconditioning" it. The aim is to develop algorithms that are not only capable of handling more general problem classes beyond Lipschitz smooth functions but also to inherently adapt to the specific properties of the cost function, leading to more efficient and robust optimization. This approach draws inspiration from the concept of inverting nonlinear systems, extending it to the optimization landscape by leveraging what the authors term "dual reference functions."

The significance of this research lies in its potential to broaden the applicability and enhance the efficiency of gradient-based optimization methods, particularly in complex, non-convex landscapes prevalent in deep learning and other advanced ML domains. By providing a principled theoretical framework rooted in the majorization-minimization (MM) principle and introducing the concept of anisotropic smoothness, the work offers a unified perspective on various adaptive gradient methods like gradient clipping and simplified versions of Adam/AdaGrad, demonstrating how they naturally emerge from this generalized framework. This allows for the design of new, adaptive algorithms with strong convergence guarantees, moving beyond the limitations of classical Lipschitz smoothness assumptions.

Background

▶ Watch: Introduction to Nonlinearly Preconditioned Gradient Methods (0:00)

The foundational problem addressed in this paper is the unconstrained minimization of a continuously differentiable, possibly non-convex cost function f(x). Gradient Descent, while a workhorse in optimization, faces significant practical limitations. Its theoretical guarantees and empirical performance are often tied to the assumption that the cost function f is globally Lipschitz smooth. This means that its gradient ∇f has a bounded Lipschitz constant L, implying that the function's curvature does not change too rapidly. However, many functions encountered in modern machine learning, such as those arising from deep neural networks with non-smooth activations (e.g., ReLU) or certain regularization techniques, do not satisfy this global Lipschitz smoothness condition. When this assumption is violated, standard GD can exhibit slow convergence, require extensive hyperparameter tuning (especially for the learning rate), or necessitate computationally expensive line search methods to ensure stability and progress.

To overcome these limitations, the authors explore a variant of gradient descent that incorporates a nonlinear preconditioner. This idea is motivated by considering analogous problems: for a linear system Ax=b with an invertible matrix A, the solution is x = A⁻¹b. Similarly, for a convex function with a unique minimizer, finding the point where the gradient is zero (∇f(x) = 0) can be conceptually linked to "inverting" this nonlinear system by computing the gradient of the convex conjugate of the function at zero. While these are idealized scenarios, they provide the intuition for using gradients of convex and smooth functions as "preconditioners," which the authors formally define as dual reference functions phi.

The paper illustrates how several existing adaptive gradient methods can be reinterpreted within this framework. For instance, normalized gradient descent and standard gradient clipping algorithms can be viewed as instances of this preconditioning scheme. In these cases, the algorithm effectively becomes gradient descent with a scalar adaptive step size that depends on the norm of the gradient. The corresponding dual reference function phi is generated by composing simple convex 1D functions with the Euclidean norm. This type of preconditioning is termed isotropic preconditioning.

Extending this, simplified versions of adaptive optimizers like Adam and AdaGrad, stripped of their momentum and memory mechanisms, yield what are essentially smooth versions of the sign gradient method. Here, the preconditioner takes the form of a coordinate-wise adaptive step size. The dual reference function phi for these methods is generated by separable sums of convex 1D functions, leading to anisotropic or separable preconditioning. These two types of preconditioning – isotropic and anisotropic – form the main focus of the paper. The authors note that more complex cases, such as combining both types or using different norms for subsets of variables, are also possible.

This general framework for preconditioned gradient methods has roots in prior research. It was initially introduced for convex and smooth functions in the "Dual Space Preconditioning for Gradient Descent" paper, where it was analyzed under a condition called dual relative smoothness. Subsequently, it was extended to the non-convex setting with an additional non-smooth term in the "Anisotropic Proximal Gradient" paper, which introduced and analyzed it under anisotropic smoothness, a condition that forms the basis of the current paper's analysis.

A key conceptual advantage of this preconditioning scheme is its natural generation from a majorization-minimization (MM) procedure. In MM, an optimization problem is iteratively solved by minimizing a sequence of simpler upper bounds (majorizers) of the original objective function. Classical gradient descent can be understood as an MM instance where, for a Lipschitz smooth cost function, the Euclidean descent lemma provides global quadratic upper bounds at every point. Iteratively minimizing these quadratic majorizers with a step size of 1/L yields the gradient descent method. The authors show that by applying the Pythagorean theorem to group terms in the Euclidean descent inequality, one can arrive at a form where the squared Euclidean norm can be replaced by a more general convex reference function phi. This substitution leads to potentially less restrictive descent inequalities, specifically the anisotropic descent inequality. When a function satisfies this inequality, it is deemed anisotropically smooth with respect to the reference function phi and a constant L. Crucially, the nonlinearly preconditioned gradient method can then be generated by successively minimizing these now non-quadratic, yet intuitively derived, upper bounds in each iteration.

Key Findings

▶ Watch: Examples: Isotropic and anisotropic preconditioners (Adam/AdaGrad) (2:00)

The talk presents several significant findings that collectively advance the understanding and application of gradient-based optimization in challenging settings.

First, the paper formally introduces and thoroughly analyzes anisotropic smoothness. This concept generalizes classical Lipschitz smoothness by replacing the standard Euclidean norm in the descent inequality with a more general convex dual reference function phi. The resulting anisotropic descent inequality forms the bedrock of their analysis. A crucial insight is that if the reference function phi is strongly convex and grows faster than the squared Euclidean norm, then anisotropic smoothness becomes a less restrictive descent inequality than standard Lipschitz smoothness. This is a powerful result, as it means the framework can apply to a broader class of functions, including those commonly encountered in modern machine learning that violate global Lipschitz conditions. Examples are provided, such as when phi has a bounded domain, which leads to very fast growth and even less restrictive descent inequalities, with the corresponding algorithm performing a form of gradient clipping. Another interesting case is the hyperbolic cosine function as phi, which yields an algorithm whose behavior lies between gradient clipping and standard gradient descent.

Second, the work establishes a tight connection between anisotropic smoothness and the concept of phi-convexity (also known as c-concavity in optimal transport literature). Phi-convexity states that a function f can be written as the pointwise supremum over a family of nonlinear functions, mirroring the classical envelope representation theorem for convex functions (which are pointwise suprema over affine functions). The authors highlight that if a function f is Lipschitz smooth, then -f is actually the supremum over concave quadratics, making f itself "quadratically phi-convex." This connection underscores that anisotropic smoothness is a natural and theoretically grounded generalization of Lipschitz smoothness, and it comes equipped with a useful envelope representation that is leveraged in their convergence analysis.

Third, recognizing the practical difficulty of directly verifying anisotropic smoothness via its definition (which must hold for any two points), the authors provide easy-to-compute second-order conditions that imply anisotropic smoothness. As an example, they demonstrate that if a function is twice continuously differentiable and exhibits L0 L1 smoothness (a specific type of smoothness related to its Hessian), then it is anisotropically smooth with respect to a reference function phi that has a bounded domain. This specific condition leads to an algorithm that takes the form of smooth normalized gradient descent with a suitably chosen step size, providing a practical pathway to identify and apply this framework.

Finally, the paper delivers comprehensive convergence results for the proposed methods across different settings:

• Non-convex setting: For a general non-convex function, if the reference function phi is "even" (symmetric), the method can use step sizes up to 2/L, akin to classical gradient descent. Furthermore, by choosing a specific form of phi, the convergence guarantees can be translated from their custom optimality measure to more standard ones.
• Convex setting (isotropic phi): Specifically for isotropic reference functions like the hyperbolic cosine, the authors prove that the sequence of iterates is Fejér monotone with respect to the set of minimizers of f. This implies that the iterates progressively get "closer" to the solution set. Additionally, the norm of the gradient is shown to be decreasing along the sequence of iterates. Under a very mild technical condition, which holds for all interesting reference functions, a sublinear rate for the suboptimality gap is established.
• General case: For the most general anisotropic settings, the authors employ a nonlinear proximal point interpretation of their method combined with the property of 2-subhomogeneity (which holds for most relevant reference functions) to derive sublinear convergence rates.

In summary, the key findings revolve around the formalization of anisotropic smoothness as a powerful generalization, its deep connection to phi-convexity, the provision of practical conditions for its verification, and robust convergence guarantees that extend the applicability of gradient methods to a much wider array of optimization problems.

Technical Deep Dive

▶ Watch: Majorization-minimization and anisotropic smoothness derivation (4:00)

The core of this work lies in reformulating and generalizing gradient descent through the concept of nonlinear preconditioning. The problem is to minimize an unconstrained, continuously differentiable function f(x), which may be non-convex and non-Lipschitz smooth.

The proposed nonlinearly preconditioned gradient method can be understood through the lens of a majorization-minimization (MM) principle. In standard gradient descent, if f is L-Lipschitz smooth, the Euclidean descent lemma guarantees a quadratic upper bound (majorizer) at any point x_k:

f(y) <= f(x_k) + <∇f(x_k), y - x_k> + (L/2) ||y - x_k||²

Minimizing this quadratic majorizer with respect to y yields the next iterate x_{k+1} = x_k - (1/L) ∇f(x_k).

This paper generalizes this by introducing a more flexible majorizer. By regrouping terms in the Euclidean descent inequality using the Pythagorean theorem, and then substituting the squared Euclidean norm (1/2) ||y - x_k||² with a general convex dual reference function phi(y - x_k), they arrive at the anisotropic descent inequality:

f(y) <= f(x_k) + <∇f(x_k), y - x_k> + L * phi(y - x_k)

A function f satisfying this inequality is said to be anisotropically smooth with respect to phi and constant L.

The nonlinearly preconditioned gradient update is then derived by minimizing this non-quadratic majorizer:

x_{k+1} = argmin_y { f(x_k) + <∇f(x_k), y - x_k> + L * phi(y - x_k) }

Assuming phi is convex and differentiable, the optimality condition for this minimization is ∇f(x_k) + L * ∇phi(y - x_k) = 0, which implies ∇phi(y - x_k) = - (1/L) ∇f(x_k). If ∇phi is invertible, this leads to y - x_k = (∇phi)⁻¹(- (1/L) ∇f(x_k)).

More formally, leveraging the properties of convex conjugates, the solution to this minimization can be expressed as:

x_{k+1} = x_k - (1/L) ∇phi^((1/L) ∇f(x_k)) (This specific form isn't explicitly stated in the transcript but is the standard derivation from minimizing f(x_k) + <∇f(x_k), d> + L phi(d) where d = y-x_k, and ∇phi*(v) is the inverse of ∇phi). The speaker describes it as "passing the gradient of the cost function through a nonlinear mapping," which is precisely what (1/L) ∇phi^* represents.

The choice of the dual reference function phi dictates the nature of the preconditioning:

1. Isotropic Preconditioning: phi is constructed by composing a simple convex 1D function h with the Euclidean norm, e.g., phi(d) = h(||d||). This often results in a scalar adaptive step size that depends on the norm of the gradient. Examples include normalized gradient descent and gradient clipping algorithms, where the effective step size alpha_k is 1 / (L * h'(||∇f(x_k)||)).
2. Anisotropic/Separable Preconditioning: phi is a separable sum of convex 1D functions, e.g., phi(d) = Σ_i h_i(d_i). This leads to coordinate-wise adaptive step sizes, where each coordinate d_i is scaled differently. Simplified versions of Adam and AdaGrad (without momentum or memory) fall into this category. The ∇phi* operation then applies component-wise.

A crucial theoretical underpinning is the connection between anisotropic smoothness and phi-convexity. A function f is phi-convex if it can be written as the pointwise supremum over a family of nonlinear functions, specifically f(x) = sup_y { <y,x> - phi^*(y) }. This is an extension of the classical result that a proper, lower semi-continuous convex function is the pointwise supremum over affine functions. This envelope representation is instrumental in the convergence analysis.

To make anisotropic smoothness practically verifiable, the authors provide second-order conditions. They demonstrate that if a function f is twice continuously differentiable and satisfies an L0 L1 smoothness condition (implying certain bounds on its Hessian in relation to phi), then it is anisotropically smooth with respect to a phi that has a bounded domain. This specific choice of phi corresponds to algorithms that perform a form of smooth normalized gradient descent.

The convergence analysis employs different techniques depending on the problem setting:

• For the non-convex setting, the analysis relies on the properties of the anisotropic descent inequality and the choice of phi. If phi is an even function, step sizes up to 2/L are permissible, similar to classical GD.
• In the convex setting with isotropic phi (e.g., hyperbolic cosine), the analysis shows that the sequence of iterates x_k is Fejér monotone with respect to the set of minimizers of f. This means that the distance from x_k to any minimizer decreases or stays constant. Additionally, the norm of the gradient ||∇f(x_k)|| is shown to be decreasing, and a sublinear rate for the suboptimality gap (the difference f(x_k) - f(x*)) is established under mild technical conditions.
• For the more general anisotropic case, the authors interpret the method as a nonlinear proximal point algorithm. By leveraging a property called 2-subhomogeneity of the reference functions phi (which holds for most interesting cases), they are able to derive sublinear convergence rates for the general framework.

Ultimately, the framework provides two interpretations: a root-finding interpretation where the method chooses the correct preconditioner for a nonlinear system, and a minimizing cost function interpretation where it chooses a non-quadratic majorizer. This dual perspective highlights the versatility and theoretical depth of nonlinearly preconditioned gradient methods.

Experimental Setup & Results

▶ Watch: Connection between anisotropic smoothness and phi-convexity (6:30)

The presented talk focuses entirely on the theoretical framework, mathematical derivations, and convergence analysis of nonlinearly preconditioned gradient methods. It introduces novel concepts such as anisotropic smoothness and provides theoretical guarantees under various conditions (non-convex, convex, isotropic, anisotropic settings).

As such, the talk does not include any discussion of experimental setups, datasets, baselines, hardware configurations, or empirical performance metrics. There are no headline numbers, speedup factors, or ablation studies presented. The presentation is purely a theoretical contribution to the field of optimization, laying the groundwork for future algorithmic development and empirical validation.

Practical Implications

▶ Watch: Second-order conditions for anisotropic smoothness (8:00)

The framework of nonlinearly preconditioned gradient methods holds several significant practical implications for practitioners, infrastructure teams, model builders, and deployers in the machine learning ecosystem.

Firstly, and perhaps most importantly, this work broadens the applicability of gradient-based optimization to a wider class of problems. Many modern machine learning models, particularly deep neural networks with non-smooth activation functions (like ReLU) or those involving certain regularization techniques, often lead to cost functions that are not globally Lipschitz smooth. Traditional gradient descent struggles in such scenarios, requiring extensive tuning or extremely small learning rates. By generalizing gradient methods to handle anisotropically smooth functions, this framework allows for more robust and efficient optimization in these challenging, non-Lipschitz environments. This means model builders can potentially use simpler, more direct gradient methods where previously they might have resorted to more complex or less theoretically understood heuristics.

Secondly, the framework provides a principled theoretical foundation for various adaptive step size mechanisms. Methods like gradient clipping, and even the core ideas behind adaptive optimizers such as Adam and AdaGrad (when stripped of their momentum and memory components), are shown to naturally arise as specific instances of nonlinearly preconditioned gradient methods with particular choices of dual reference functions phi. This understanding can guide the design of new, more effective adaptive learning rate schemes, moving beyond empirical trial-and-error to a more theoretically informed approach. For practitioners, this could lead to algorithms that are more self-tuning and require less manual hyperparameter optimization, thereby reducing development time and computational resources spent on tuning.

Thirdly, the majorization-minimization (MM) principle offers a clear methodology for designing new optimization algorithms. By carefully selecting the dual reference function phi, researchers and engineers can systematically derive novel preconditioned gradient methods tailored to specific problem structures or desired algorithmic behaviors. For instance, if a problem exhibits certain sparsity patterns or coordinate-wise sensitivities, an anisotropic phi could be designed to exploit these properties, leading to more efficient updates. This principled design approach contrasts with heuristic algorithm development and can accelerate innovation in optimization.

However, there are also trade-offs and limitations to consider. The effectiveness of the method heavily relies on the appropriate choice of the dual reference function phi. While the paper provides guidance and examples (e.g., phi with bounded domain for gradient clipping, hyperbolic cosine for intermediate behavior), identifying the optimal phi for a given complex ML model might still require significant analytical effort or domain expertise. The "easy-to-compute second-order conditions" for anisotropic smoothness, while helpful, might still be challenging to verify for arbitrary deep learning architectures with millions of parameters.

Furthermore, the computational cost associated with computing ∇phi* in each iteration needs to be efficient. For simple phi functions (like those leading to isotropic or separable preconditioning), this might be straightforward. However, for more complex phi choices, the inverse gradient of phi or its convex conjugate's gradient might be computationally intensive, potentially offsetting the benefits of improved convergence behavior. Infrastructure teams would need to ensure that the underlying deep learning frameworks (e.g., PyTorch, TensorFlow) can efficiently support these generalized gradient operations.

Finally, while the paper establishes sublinear convergence rates for both convex and non-convex settings, these are general rates. For specific problem classes, specialized algorithms might offer faster convergence (e.g., linear convergence for strongly convex problems). The practical gains in terms of wall-clock time compared to highly optimized existing adaptive optimizers (like AdamW or Lion) would need to be thoroughly evaluated empirically. The current work provides the theoretical foundation, but extensive benchmarking on real-world ML tasks will be crucial to assess its full practical impact and identify scenarios where it offers a distinct advantage.

Key Takeaways

• Generalization of Gradient Descent: Nonlinearly preconditioned gradient methods extend standard gradient descent to handle functions that are not globally Lipschitz smooth, a common scenario in modern machine learning.
• Dual Reference Functions: The core idea involves passing the gradient through a nonlinear mapping derived from a convex and smooth "dual reference function" phi, which adapts the algorithm to the cost function's local geometry.
• Unified Framework: This approach unifies and provides theoretical grounding for various adaptive gradient techniques, including gradient clipping and simplified versions of adaptive optimizers like Adam and AdaGrad, by showing how they arise from specific choices of phi.
• Majorization-Minimization Principle: The method is rigorously derived from a generalized majorization-minimization procedure, replacing standard quadratic upper bounds with non-quadratic majorizers based on phi.
• Anisotropic Smoothness & Phi-Convexity: The paper introduces anisotropic smoothness as a less restrictive generalization of Lipschitz smoothness and establishes its deep connection to phi-convexity, providing crucial theoretical underpinnings.
• Robust Convergence Guarantees: The framework offers convergence guarantees (including sublinear rates, Fejér monotonicity, and gradient norm decrease) for both non-convex and convex optimization problems under various preconditioning schemes, with practical second-order conditions provided for verifying anisotropic smoothness.

About the Speaker(s)

The talk was presented by Konstantinos Oikonomidis, and is joint work with Jan Quan, Emanuel Laude, and Panagiotis Patrinos. The transcript does not provide specific titles or affiliations for the speakers.