General framework for online-to-nonconvex conversion: Schedule-free SGD is also effective for nonconvex optimization

General framework for online-to-nonconvex conversion: Schedule-free SGD is also effective for nonconvex optimization

Speakers: Gagik Magakyan, Kwangjun Ahn, Ashok Cutkosky

Conference: ICML 2025

YouTube: https://slideslive.com/39044018

Overview

This article delves into a significant theoretical advancement in the field of optimization for machine learning, presented at ICML 2025. The talk, titled "General framework for online-to-nonconvex conversion: Schedule-free SGD is also effective for nonconvex optimization," by Gagik Magakyan, Kwangjun Ahn, and Ashok Cutkosky, introduces a novel theoretical framework that extends the applicability and understanding of schedule-free Stochastic Gradient Descent (SGD). Traditionally, neural network training heavily relies on carefully tuned learning rate schedules, which are notoriously difficult to optimize and often require pre-specification for the entire training budget. Schedule-free methods aim to alleviate this burden by eliminating the need for such schedules, making optimization more robust and user-friendly.

The core contribution of this work is a rigorous theoretical analysis demonstrating that schedule-free SGD is not only empirically effective but also theoretically optimal for nonconvex and nonsmooth problems, a common characteristic of modern deep learning landscapes. This represents a crucial generalization over prior work, which primarily offered theoretical guarantees only for convex problems. The authors achieve this by developing a generalized online-to-nonconvex conversion framework, which provides a powerful lens through which to analyze and derive convergence bounds for a broad class of optimization algorithms. Furthermore, their analysis provides the first theoretical explanation for a critical practical observation: the interpolation parameter kappa in schedule-free SGD must be chosen close to one for optimal performance.

This research is particularly pertinent in an era dominated by large-scale models like Large Language Models (LLMs), where the computational cost and expertise required for hyperparameter tuning, especially learning rate schedules, are immense. By offering a robust, theoretically grounded, and empirically validated alternative to scheduled optimization, this work paves the way for more efficient and accessible training of complex neural networks, thereby accelerating research and deployment in AI/ML.

Background

▶ Watch: Introduction: Schedule-free SGD benefits and challenges (0:00)

The training of modern machine learning models, particularly deep neural networks, is fundamentally an iterative optimization process. At its heart lies Stochastic Gradient Descent (SGD) and its variants, where model parameters xt are updated at each step t using a gradient estimate delta_t and a learning rate eta_t: xt+1 = xt - eta_t * delta_t. The learning rate eta_t, often referred to as a schedule, is crucial for effective training. It dictates the step size and typically decays over time to ensure convergence and prevent oscillations.

Despite their necessity, learning rate schedules present significant practical challenges. They demand extensive manual tuning, a time-consuming and resource-intensive process that often requires specialized expertise. Moreover, schedules typically need to be pre-specified for a fixed optimization budget, meaning the total number of training iterations must be known in advance. This rigidity can be problematic in dynamic research environments or when training budgets are uncertain. The complexity of these schedules can vary widely; while some models like Llama 3 might employ well-known cosine decay schedules, others, such as those from DeepSeek-AI, might utilize much more intricate piecewise linear schedules, further exacerbating the tuning burden.

Recognizing these challenges, Aaron Defazio and collaborators introduced schedule-free methods in their NeurIPS work. Their approach aimed to eliminate the need for explicit learning rate schedules, demonstrating impressive empirical performance. Notably, their method won the AlgoPerf Challenge self-tuning track, showcasing its practical efficacy. Empirical results, such as training a GPT-2 small network on OpenWebText or a ResNet-50 on ImageNet, showed that schedule-free methods could track the Pareto frontier of optimization curves, often outperforming heavily tuned cosine schedules across different optimization budgets.

However, the initial theoretical foundations for these schedule-free methods had limitations. The theory provided by Defazio et al. was primarily applicable to convex problems, which do not accurately represent the highly nonconvex loss landscapes of deep neural networks. Furthermore, the hyperparameter selections suggested by their convex theory often did not align with the values found to be effective in practical applications, leaving a gap between theory and practice. This highlighted the need for a more comprehensive theoretical framework that could address nonconvexity and explain empirical observations, which is precisely what the presented ICML 2025 work aims to provide.

Key Findings

▶ Watch: Our work: Schedule-free SGD for nonconvex and nonsmooth problems (2:20)

The ICML 2025 talk by Magakyan, Ahn, and Cutkosky unveils several pivotal findings that significantly advance the understanding and application of schedule-free SGD:

Firstly, and most critically, the research demonstrates that schedule-free SGD is effective for nonconvex problems. This extends the theoretical guarantees beyond the previously established convex regime, directly addressing the complexities of real-world deep learning optimization. The authors rigorously show that schedule-free SGD achieves optimal rates for nonconvex and nonsmooth problems, a significant theoretical achievement given the challenging nature of such optimization landscapes.

Secondly, the work introduces a novel analysis technique rooted in a generalized online-to-nonconvex conversion framework. This framework builds upon existing literature that connects online learning algorithms to nonconvex optimizers, allowing for a more flexible and powerful way to derive convergence bounds. The authors prove that schedule-free SGD is a specific instance of this generalized conversion, solidifying its theoretical foundation within a broader optimization paradigm.

A particularly impactful finding is the theoretical explanation for a crucial practical hyperparameter selection. The interpolation parameter kappa, which dictates the blend between the current iterate zt and its weighted average xt in the schedule-free update rule, was empirically observed to perform best when chosen close to one (e.g., 0.95). Previous convex theories could not explain this observation, suggesting optimal convergence for any kappa between zero and one. This new nonconvex analysis provides the first theoretical justification for why kappa needs to be close to one, bridging a significant gap between theory and practical heuristics.

Finally, the work establishes that the schedule-free SGD algorithm, when analyzed through their generalized framework, achieves an optimal rate of convergence in a specific function class under a generalized notion of stationarity called lambda-epsilon stationarity. This measure, which penalizes variance rather than strictly constraining a delta-ball, is more appropriate for the challenging nonconvex and nonsmooth regime, where classical stationary points are not always achievable. These findings collectively provide a robust theoretical underpinning for schedule-free optimization, enhancing its credibility and guiding its practical application.

Technical Deep Dive

▶ Watch: Dissecting the Schedule-Free SGD algorithm (3:10)

The core of this work lies in its detailed technical exposition of schedule-free SGD and its theoretical analysis via a generalized online-to-nonconvex conversion framework.

The schedule-free SGD algorithm itself is defined by a pair of iterates: zt representing the primary optimization variable, and xt as a weighted average of past zt values. The update rule is given by:

1. xt = (1/Ct) sum_{i=1 to t} (ci zi) where Ct = sum_{i=1 to t} ci. If ct = 1/t, xt becomes an exact average; if ct is a constant, xt approximates an Exponential Moving Average (EMA). Crucially, ct is a fixed sequence, not a tuned schedule.
2. The gradient is sampled at an interpolated point yt: yt = (1 - kappa) xt + kappa zt.
3. The primary iterate zt is updated using a fixed learning rate eta (no schedule) and the gradient at yt: zt+1 = zt - eta * grad(yt).

A critical aspect observed in practice is that the interpolation parameter kappa must be chosen close to one (e.g., kappa ≈ 0.95) for good practical performance. The previous convex theory offered no explanation for this, suggesting optimality across the entire [0,1] interval. This new work provides the theoretical justification for this empirical observation.

The problem setting considered by the authors deviates from standard smooth nonconvex optimization. They analyze functions f: R^d -> R that are:

• Differentiable.
• Bounded below.
• Lipschitz.
• Satisfy a "well-behavedness" condition, which is a weaker assumption than typical smoothness, allowing for nonsmooth problems.

Given this setting, aiming for classical stationary points (where the gradient is zero) is often not feasible. Therefore, the authors consider a generalization of Goldstein stationarity, which typically requires finding a random vector in a delta-ball with a small average gradient. Instead, they introduce lambda-epsilon stationarity, which penalizes the variance of the gradient, effectively seeking a point where the gradient is small in expectation while controlling its variability. This notion is more robust for the challenging nonsmooth regime.

The analytical backbone of the paper is the online-to-nonconvex conversion framework, initially introduced by co-author Ashok Cutkosky and his collaborators. This framework establishes a profound connection between online learning algorithms and nonconvex optimization.

• In this framework, an online learner chooses an optimization update delta_t.
• The optimization iterate moves from xt-1 to xt via xt = xt-1 + delta_t.
• A gradient is sampled uniformly randomly in the interval between xt-1 and xt, and a loss (related to the gradient) is provided back to the online learner.
• Remarkably, famous online learning algorithms map to practical optimization algorithms: Optimistic Mirror Descent (OMD) corresponds to Stochastic Gradient Descent with Momentum, and Follow-The-Regularized-Leader (FTRL) corresponds to Adam.
• A key insight is that a sublinear regret bound for the online learner directly translates to an optimal convergence bound for the corresponding optimization algorithm.

Magakyan and his co-authors further generalize this framework. The previous conversion restricted the gradient sampling to the interval defined by xt-1 and xt. The generalized framework provides the algorithm with the flexibility of choosing xt arbitrarily (subject to certain regularity conditions) and, crucially, allows for freedom in choosing where the gradient is sampled. This means the gradient is no longer strictly sampled between the optimization iterates.

The main theorem, presented informally, connects the optimization algorithm's primary quantity of interest – the measure of stationarity (left-hand side) – to the online learner's regret, plus a term related to the function iterate difference, and other smaller order terms (right-hand side). As a sanity check, if xt is chosen to be wt-1 (where w is an auxiliary variable), the generalized framework recovers the previous online-to-nonconvex conversion.

The pivotal result is that with a smart choice of xt and by employing Optimistic Mirror Descent (OMD) as the online learner, the resulting optimization algorithm is precisely schedule-free SGD. This theoretical construction rigorously proves that schedule-free SGD is optimal for nonconvex optimization in this generalized setting. The analysis also directly explains why kappa must be close to one: the theoretical bounds derived through this generalized conversion framework reveal that values of kappa near one are essential for achieving optimal convergence rates under nonconvexity.

Experimental Setup & Results

▶ Watch: Understanding the online-to-nonconvex conversion framework (5:40)

The presented talk primarily focuses on the theoretical advancements and a novel analytical framework for schedule-free SGD, rather than introducing new empirical results from the authors' own experiments. The talk builds upon the strong empirical foundations established by previous work, particularly the NeurIPS paper by Aaron Defazio and collaborators, which introduced schedule-free methods.

The speaker references these prior empirical successes to motivate the theoretical inquiry. Specifically, the previous work demonstrated impressive performance on standard neural network training tasks:

• Training a GPT-2 small network on the OpenWebText dataset.
• Training a ResNet-50 on the ImageNet dataset.

In these experiments, the schedule-free method (represented by a black curve in the original work) was shown to effectively track the Pareto frontier of the optimization curve. It consistently outperformed heavily tuned cosine schedules across various optimization budgets. This empirical evidence highlighted the practical viability and superiority of schedule-free methods over traditional scheduled approaches, even without a complete theoretical understanding for nonconvex problems.

While this ICML 2025 talk does not present new datasets, baselines, hardware configurations, or headline numbers from the authors' own studies, it is crucial to understand that its contribution is fundamentally theoretical. The purpose is to provide the missing theoretical guarantees and explanations for the observed empirical performance, particularly for nonconvex and nonsmooth optimization, and to shed light on hyperparameter choices like kappa that were previously only empirically driven. The existing empirical results serve as a testament to the practical relevance of the theoretical framework developed in this paper.

Practical Implications

▶ Watch: Generalized online-to-nonconvex framework and main theorem (8:00)

The theoretical breakthroughs presented in this work have profound practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML space.

The most immediate benefit is the reduction in hyperparameter tuning burden. Learning rate schedules are notoriously difficult and time-consuming to tune, often requiring extensive grid searches or expert intuition. By validating schedule-free SGD for nonconvex problems, this research offers a path towards more robust and "self-tuning" optimization algorithms. This translates directly to reduced computational costs and faster iteration cycles in model development, as less time is spent on trial-and-error hyperparameter optimization. For large-scale models, such as Large Language Models (LLMs), where training runs can span weeks or months and involve thousands of GPUs, even minor reductions in tuning effort can lead to substantial savings in compute and engineering time.

For infrastructure teams, the adoption of schedule-free methods could simplify the deployment and management of training pipelines. Algorithms that are less sensitive to hyperparameter choices are inherently more robust to variations in hardware, data distributions, or training budget changes. This can lead to more stable and predictable training processes, reducing the need for constant monitoring and intervention.

Model builders and deployers will benefit from more reliable training outcomes. The theoretical optimality guarantees for nonconvex and nonsmooth problems instill greater confidence in schedule-free SGD's ability to converge to good solutions. Furthermore, the theoretical explanation for the kappa hyperparameter (that it should be close to one) provides clear guidance for practitioners. Instead of blindly searching for an optimal kappa value, they now have a theoretically justified starting point (e.g., 0.95), which can significantly narrow the search space and improve initial model performance.

However, it's important to acknowledge practical tradeoffs and limitations. While schedule-free SGD eliminates the learning rate schedule, it still requires the setting of a fixed learning rate eta and the interpolation parameter kappa. While the guidance for kappa is a significant step, eta may still require some degree of tuning, although likely less sensitive than a full schedule. The theoretical nature of this work means that empirical validation on a broader range of architectures, datasets, and hardware configurations will still be necessary for widespread adoption. Additionally, the definition of lambda-epsilon stationarity provides a theoretical convergence guarantee, but its direct correlation to desired generalization performance in all practical scenarios may warrant further investigation. Nevertheless, this work represents a crucial step towards making deep learning optimization more accessible, efficient, and robust.

Key Takeaways

• Schedule-free SGD eliminates the need for manual learning rate schedule tuning, addressing a major pain point in deep learning optimization.
• This research extends the theoretical understanding of schedule-free SGD, proving its optimality for nonconvex and nonsmooth problems, which are characteristic of modern neural networks.
• The work introduces a generalized online-to-nonconvex conversion framework, providing a powerful new tool for analyzing and deriving convergence bounds for complex optimization algorithms.
• A key practical insight is theoretically justified: the interpolation parameter kappa in schedule-free SGD must be chosen close to one (e.g., 0.95) for optimal performance, explaining previous empirical observations.
• By leveraging Optimistic Mirror Descent (OMD) within the generalized framework, the authors demonstrate that schedule-free SGD is an instance of this optimal conversion.
• These advancements pave the way for more robust, efficient, and less-tuned training of large-scale machine learning models, significantly reducing computational costs and development time.

About the Speaker(s)

The talk "General framework for online-to-nonconvex conversion: Schedule-free SGD is also effective for nonconvex optimization" was presented by Gagik Magakyan, and is a collaborative effort with Kwangjun Ahn and Ashok Cutkosky.

Ashok Cutkosky is a distinguished researcher in the field of optimization and machine learning, and a co-author of this work. He is particularly recognized for his foundational contributions to the online-to-nonconvex conversion framework, which serves as the theoretical bedrock for the analysis presented in this talk. His prior work in this area has established critical connections between online learning algorithms and practical optimization methods, significantly influencing the theoretical understanding of deep learning optimizers.

Kwangjun Ahn is also a co-author of this research and has been actively involved in the line of work connecting online learning to optimization algorithms, as highlighted by the speaker when referencing prior papers in this domain. His expertise contributes to the rigorous theoretical development and generalization of the conversion framework.

The presentation by Gagik Magakyan effectively conveyed the intricate theoretical details and the practical significance of their findings, guiding the audience through the complexities of schedule-free SGD and its optimal performance in nonconvex settings.