arXiv subjects

Zaiwen Wen

Publications and source records attributed to Zaiwen Wen.

At least 55 records · Page 3Linked to original sources

arXiv2025

Riemannian EXTRA: Communication-efficient decentralized optimization over compact submanifolds with data heterogeneity

We consider decentralized optimization over a compact Riemannian submanifold in a network of $n$ agents, where each agent holds a smooth, nonconvex local objective defined by its private data. The goal is to collaboratively minimize the sum of these local objective functions. In the presence of data heterogeneity across nodes, existing algorithms typically require communicating both local gradients and iterates to ensure exact convergence with constant step sizes. In this work, we propose REXTRA, a Riemannian extension of the EXTRA algorithm [Shi et al., SIOPT, 2015], to address this limitation. On the theoretical side, we leverage proximal smoothness to overcome the challenges of manifold nonconvexity and establish a global sublinear convergence rate of $\mathcal{O}(1/k)$, matching the best-known results. To our knowledge, REXTRA is the first algorithm to achieve a global sublinear convergence rate under a constant step size while requiring only a single round of local iterate communication per iteration. Numerical experiments show that REXTRA achieves superior performance compared to state-of-the-art methods, while supporting larger step sizes and reducing total communication by over 50\%.

math.OC↗

arXiv2025

An efficient primal dual semismooth Newton method for semidefinite programming

In this paper, we present an efficient semismooth Newton method, named SSNCP, for solving a class of semidefinite programming problems. Our approach is rooted in an equivalent semismooth system derived from the saddle point problem induced by the augmented Lagrangian duality. An additional correction step is incorporated after the semismooth Newton step to ensure that the iterates eventually reside on a manifold where the semismooth system is locally smooth. Global convergence is achieved by carefully designing inexact criteria and leveraging the $α$-averaged property to analyze the error. The correction steps address challenges related to the lack of smoothness in local convergence analysis. Leveraging the smoothness established by the correction steps and assuming a local error bound condition, we establish the local superlinear convergence rate without requiring the stringent assumptions of nonsingularity or strict complementarity. Furthermore, we prove that SSNCP converges to an $\varepsilon$-stationary point with an iteration complexity of $\widetilde{\mathcal{O}}(\varepsilon^{-3/2})$. Numerical experiments on various datasets, especially the Mittelmann benchmark, demonstrate the high efficiency and robustness of SSNCP compared to state-of-the-art solvers.

math.OC↗

arXiv2025

Accelerated Natural Gradient Method for Parametric Manifold Optimization

Parametric manifold optimization problems frequently arise in various machine learning tasks, where state functions are defined on infinite-dimensional manifolds. We propose a unified accelerated natural gradient descent (ANGD) framework to address these problems. By incorporating a Hessian-driven damping term into the manifold update, we derive an accelerated Riemannian gradient (ARG) flow that mitigates oscillations. An equivalent first-order system is further presented for the ARG flow, enabling a unified discretization scheme that leads to the ANGD method. In our discrete update, our framework considers various advanced techniques, including least squares approximation of the update direction, projected momentum to accelerate convergence, and efficient approximation methods through the Kronecker product. It accommodates various metrics, including $H^s$, Fisher-Rao, and Wasserstein-2 metrics, providing a computationally efficient solution for large-scale parameter spaces. We establish a convergence rate for the ARG flow under geodesic convexity assumptions. Numerical experiments demonstrate that ANGD outperforms standard NGD, underscoring its effectiveness across diverse deep learning tasks.

math.OC↗

arXiv2025

Formalization of Algorithms for Optimization with Block Structures

Block-structured problems are central to advances in numerical optimization and machine learning. This paper provides the formalization of convergence analysis for two pivotal algorithms in such settings: the block coordinate descent (BCD) method and the alternating direction method of multipliers (ADMM). Utilizing the type-theory-based proof assistant Lean4, we develop a rigorous framework to formally represent these algorithms. Essential concepts in nonsmooth and nonconvex optimization are formalized, notably subdifferentials, which extend the classical differentiability to handle nonsmooth scenarios, and the Kurdyka-Lojasiewicz (KL) property, which provides essential tools to analyze convergence in nonconvex settings. Such definitions and properties are crucial for the corresponding convergence analyses. We formalize the convergence proofs of these algorithms, demonstrating that our definitions and structures are coherent and robust. These formalizations lay a basis for analyzing the convergence of more general optimization algorithms.

math.OC↗

arXiv2025

Formalization of Optimality Conditions for Smooth Constrained Optimization Problems

Optimality conditions are central to analysis of optimization problems, characterizing necessary criteria for local minima. Formalizing the optimality conditions within the type-theory-based proof assistant Lean4 provides a precise, robust, and reusable framework essential for rigorous verification in optimization theory. In this paper, we introduce a formalization of the first-order optimality conditions (also known as the Karush-Kuhn-Tucker (KKT) conditions) for smooth constrained optimization problems by beginning with concepts such as the Lagrangian function and constraint qualifications. The geometric optimality conditions are then formalized, offering insights into local minima through tangent cones. We also establish the critical equivalence between the tangent cone and linearized feasible directions under appropriate constraint qualifications. Building on these key elements, the formalization concludes the KKT conditions through the proof of the Farkas lemma. Additionally, this study provides a formalization of the dual problem and the weak duality property.

math.OC↗

arXiv2025

OptMATH: A Scalable Bidirectional Data Synthesis Framework for Optimization Modeling

Despite the rapid development of large language models (LLMs), a fundamental challenge persists: the lack of high-quality optimization modeling datasets hampers LLMs' robust modeling of practical optimization problems from natural language descriptions (NL). This data scarcity also contributes to the generalization difficulties experienced by learning-based methods. To address these challenges, we propose a scalable framework for synthesizing a high-quality dataset, named OptMATH. Starting from curated seed data with mathematical formulations (MF), this framework automatically generates problem data (PD) with controllable complexity. Then, a back-translation step is employed to obtain NL. To verify the correspondence between the NL and the PD, a forward modeling step followed by rejection sampling is used. The accepted pairs constitute the training part of OptMATH. Then a collection of rejected pairs is identified and further filtered. This collection serves as a new benchmark for optimization modeling, containing difficult instances whose lengths are much longer than these of NL4OPT and MAMO. Through extensive experiments, we demonstrate that models of various sizes (0.5B-32B parameters) trained on OptMATH achieve superior results on multiple modeling benchmarks, thereby validating the effectiveness and scalability of our approach. Our dataset is publicly available at https://github.com/AuroraLHL/OptMATH.

cs.AI↗

arXiv2025

A Memory Efficient Randomized Subspace Optimization Method for Training Large Language Models

The memory challenges associated with training Large Language Models (LLMs) have become a critical concern, particularly when using the Adam optimizer. To address this issue, numerous memory-efficient techniques have been proposed, with GaLore standing out as a notable example designed to reduce the memory footprint of optimizer states. However, these approaches do not alleviate the memory burden imposed by activations, rendering them unsuitable for scenarios involving long context sequences or large mini-batches. Moreover, their convergence properties are still not well-understood in the literature. In this work, we introduce a Randomized Subspace Optimization framework for pre-training and fine-tuning LLMs. Our approach decomposes the high-dimensional training problem into a series of lower-dimensional subproblems. At each iteration, a random subspace is selected, and the parameters within that subspace are optimized. This structured reduction in dimensionality allows our method to simultaneously reduce memory usage for both activations and optimizer states. We establish comprehensive convergence guarantees and derive rates for various scenarios, accommodating different optimization strategies to solve the subproblems. Extensive experiments validate the superior memory and communication efficiency of our method, achieving performance comparable to GaLore and Adam.

cs.LG↗

arXiv2024

Enhancing Zeroth-order Fine-tuning for Language Models with Low-rank Structures

Parameter-efficient fine-tuning (PEFT) significantly reduces memory costs when adapting large language models (LLMs) for downstream applications. However, traditional first-order (FO) fine-tuning algorithms incur substantial memory overhead due to the need to store activation values for back-propagation during gradient computation, particularly in long-context fine-tuning tasks. Zeroth-order (ZO) algorithms offer a promising alternative by approximating gradients using finite differences of function values, thus eliminating the need for activation storage. Nevertheless, existing ZO methods struggle to capture the low-rank gradient structure common in LLM fine-tuning, leading to suboptimal performance. This paper proposes a low-rank ZO gradient estimator and introduces a novel low-rank ZO algorithm (LOZO) that effectively captures this structure in LLMs. We provide convergence guarantees for LOZO by framing it as a subspace optimization method. Additionally, its low-rank nature enables LOZO to integrate with momentum techniques while incurring negligible extra memory costs. Extensive experiments across various model sizes and downstream tasks demonstrate that LOZO and its momentum-based variant outperform existing ZO methods and closely approach the performance of FO algorithms.

cs.LG↗

arXiv2024

Formalization of Complexity Analysis of the First-order Algorithms for Convex Optimization

The convergence rate of various first-order optimization algorithms is a pivotal concern within the numerical optimization community, as it directly reflects the efficiency of these algorithms across different optimization problems. Our goal is making a significant step forward in the formal mathematical representation of optimization techniques using the Lean4 theorem prover. We first formalize the gradient for smooth functions and the subgradient for convex functions on a Hilbert space, laying the groundwork for the accurate formalization of algorithmic structures. Then, we extend our contribution by proving several properties of differentiable convex functions that have not yet been formalized in Mathlib. Finally, a comprehensive formalization of these algorithms is presented. These developments are not only noteworthy on their own but also serve as essential precursors to the formalization of a broader spectrum of numerical algorithms and their applications in machine learning as well as many other areas.

math.OC↗

arXiv2024

A Customized Augmented Lagrangian Method for Block-Structured Integer Programming

Integer programming with block structures has received considerable attention recently and is widely used in many practical applications such as train timetabling and vehicle routing problems. It is known to be NP-hard due to the presence of integer variables. We define a novel augmented Lagrangian function by directly penalizing the inequality constraints and establish the strong duality between the primal problem and the augmented Lagrangian dual problem. Then, a customized augmented Lagrangian method is proposed to address the block-structures. In particular, the minimization of the augmented Lagrangian function is decomposed into multiple subproblems by decoupling the linking constraints and these subproblems can be efficiently solved using the block coordinate descent method. We also establish the convergence property of the proposed method. To make the algorithm more practical, we further introduce several refinement techniques to identify high-quality feasible solutions. Numerical experiments on a few interesting scenarios show that our proposed algorithm often achieves a satisfactory solution and is quite effective.

math.OC↗

arXiv2024

ODE-based Learning to Optimize

Recent years have seen a growing interest in understanding acceleration methods through the lens of ordinary differential equations (ODEs). Despite the theoretical advancements, translating the rapid convergence observed in continuous-time models to discrete-time iterative methods poses significant challenges. In this paper, we present a comprehensive framework integrating the inertial systems with Hessian-driven damping equation (ISHD) and learning-based approaches for developing optimization methods through a deep synergy of theoretical insights. We first establish the convergence condition for ensuring the convergence of the solution trajectory of ISHD. Then, we show that provided the stability condition, another relaxed requirement on the coefficients of ISHD, the sequence generated through the explicit Euler discretization of ISHD converges, which gives a large family of practical optimization methods. In order to select the best optimization method in this family for certain problems, we introduce the stopping time, the time required for an optimization method derived from ISHD to achieve a predefined level of suboptimality. Then, we formulate a novel learning to optimize (L2O) problem aimed at minimizing the stopping time subject to the convergence and stability condition. To navigate this learning problem, we present an algorithm combining stochastic optimization and the penalty method (StoPM). The convergence of StoPM using the conservative gradient is proved. Empirical validation of our framework is conducted through extensive numerical experiments across a diverse set of optimization problems. These experiments showcase the superior performance of the learned optimization methods.

math.OC↗

arXiv2024

An Augmented Lagrangian Primal-Dual Semismooth Newton Method for Multi-Block Composite Optimization

In this paper, we develop a novel primal-dual semismooth Newton method for solving linearly constrained multi-block convex composite optimization problems. First, a differentiable augmented Lagrangian (AL) function is constructed by utilizing the Moreau envelopes of the nonsmooth functions. It enables us to derive an equivalent saddle point problem and establish the strong AL duality under the Slater's condition. Consequently, a semismooth system of nonlinear equations is formulated to characterize the optimality of the original problem instead of the inclusion-form KKT conditions. We then develop a semismooth Newton method, called ALPDSN, which uses purely second-order steps and a nonmonotone line search based globalization strategy. Through a connection to the inexact first-order steps when the regularization parameter is sufficiently large, the global convergence of ALPDSN is established. Under the regularity conditions, partial smoothness, the local error bound, and the strict complementarity, we show that both the primal and the dual iteration sequences possess a superlinear convergence rate and provide concrete examples where these regularity conditions are met. Numerical results on the image restoration with two regularization terms and the corrected tensor nuclear norm problem are presented to demonstrate the high efficiency and robustness of our ALPDSN.

math.OC↗

arXiv2024

An Improved Finite-time Analysis of Temporal Difference Learning with Deep Neural Networks

Temporal difference (TD) learning algorithms with neural network function parameterization have well-established empirical success in many practical large-scale reinforcement learning tasks. However, theoretical understanding of these algorithms remains challenging due to the nonlinearity of the action-value approximation. In this paper, we develop an improved non-asymptotic analysis of the neural TD method with a general $L$-layer neural network. New proof techniques are developed and an improved new $\tilde{\mathcal{O}}(ε^{-1})$ sample complexity is derived. To our best knowledge, this is the first finite-time analysis of neural TD that achieves an $\tilde{\mathcal{O}}(ε^{-1})$ complexity under the Markovian sampling, as opposed to the best known $\tilde{\mathcal{O}}(ε^{-2})$ complexity in the existing literature.

cs.LG↗

arXiv2024

Oracle complexities of augmented Lagrangian methods for nonsmooth manifold optimization

In this paper, we present two novel manifold inexact augmented Lagrangian methods, \textbf{ManIAL} for deterministic settings and \textbf{StoManIAL} for stochastic settings, solving nonsmooth manifold optimization problems. By using the Riemannian gradient method as a subroutine, we establish an $\mathcal{O}(ε^{-3})$ oracle complexity result of \textbf{ManIAL}, matching the best-known complexity result. Our algorithm relies on the careful selection of penalty parameters and the precise control of termination criteria for subproblems. Moreover, for cases where the smooth term follows an expectation form, our proposed \textbf{StoManIAL} utilizes a Riemannian recursive momentum method as a subroutine, and achieves an oracle complexity of $\tilde{\mathcal{O}}(ε^{-3.5})$, which surpasses the best-known $\mathcal{O}(ε^{-4})$ result. Numerical experiments conducted on sparse principal component analysis and sparse canonical correlation analysis demonstrate that our proposed methods outperform an existing method with the previously best-known complexity result. To the best of our knowledge, these are the first complexity results of the augmented Lagrangian methods for solving nonsmooth manifold optimization problems.

math.OC↗

arXiv2024

Gauss-Newton Temporal Difference Learning with Nonlinear Function Approximation

In this paper, a Gauss-Newton Temporal Difference (GNTD) learning method is proposed to solve the Q-learning problem with nonlinear function approximation. In each iteration, our method takes one Gauss-Newton (GN) step to optimize a variant of Mean-Squared Bellman Error (MSBE), where target networks are adopted to avoid double sampling. Inexact GN steps are analyzed so that one can safely and efficiently compute the GN updates by cheap matrix iterations. Under mild conditions, non-asymptotic finite-sample convergence to the globally optimal Q function is derived for various nonlinear function approximations. In particular, for neural network parameterization with relu activation, GNTD achieves an improved sample complexity of $\tilde{\mathcal{O}}(\varepsilon^{-1})$, as opposed to the $\mathcal{\mathcal{O}}(\varepsilon^{-2})$ sample complexity of the existing neural TD methods. An $\tilde{\mathcal{O}}(\varepsilon^{-1.5})$ sample complexity of GNTD is also established for general smooth function approximations. We validate our method via extensive experiments in several RL benchmarks, where GNTD exhibits both higher rewards and faster convergence than TD-type methods.

math.OC↗

arXiv2024

Convergence Analysis of Stochastic Gradient Descent with MCMC Estimators

Understanding stochastic gradient descent (SGD) and its variants is essential for machine learning. However, most of the preceding analyses are conducted under amenable conditions such as unbiased gradient estimator and bounded objective functions, which does not encompass many sophisticated applications, such as variational Monte Carlo, entropy-regularized reinforcement learning and variational inference. In this paper, we consider the SGD algorithm that employ the Markov Chain Monte Carlo (MCMC) estimator to compute the gradient, called MCMC-SGD. Since MCMC reduces the sampling complexity significantly, it is an asymptotically convergent biased estimator in practice. Moreover, by incorporating a general class of unbounded functions, it is much more difficult to analyze the MCMC sampling error. Therefore, we assume that the function is sub-exponential and use the Bernstein inequality for non-stationary Markov chains to derive error bounds of the MCMC estimator. Consequently, MCMC-SGD is proven to have a first order convergence rate $O(\log K/\sqrt{n K})$ with $K$ iterations and a sample size $n$. It partially explains how MCMC influences the behavior of SGD. Furthermore, we verify the correlated negative curvature condition under reasonable assumptions. It is shown that MCMC-SGD escapes from saddle points and reaches $(ε,ε^{1/4})$ approximate second order stationary points or $ε^{1/2}$-variance points at least $O(ε^{-11/2}\log^{2}(1/ε) )$ steps with high probability. Our analysis unveils the convergence pattern of MCMC-SGD across a broad class of stochastic optimization problems, and interprets the convergence phenomena observed in practical applications.

stat.ML↗

arXiv2023

Sharper Convergence Guarantees for Federated Learning with Partial Model Personalization

Partial model personalization, which encompasses both shared and personal variables in its formulation, is a critical optimization problem in federated learning. It balances individual client needs with collective knowledge utilization, and serves as a general formulation covering various key scenarios, ranging from fully shared to fully personalized federated learning. This paper introduces two effective algorithms, FedAvg-P and Scaffold-P, to solve this problem and provides sharp convergence analyses, quantifying the influence of gradient variance, local steps, and partial client sampling on their performance. Our established rates surpass existing results and, meanwhile, are based on more relaxed assumptions. Additionally, our analyses are also applicable to fully shared or fully personalized federated learning, matching or even outperforming their best known convergence rates. Numerical experiments corroborate our theoretical findings.

math.OC↗

arXiv2023

On the Optimal Lower and Upper Complexity Bounds for a Class of Composite Optimization Problems

We study the optimal lower and upper complexity bounds for finding approximate solutions to the composite problem $\min_x\ f(x)+h(Ax-b)$, where $f$ is smooth and $h$ is convex. Given access to the proximal operator of $h$, for strongly convex, convex, and nonconvex $f$, we design efficient first order algorithms with complexities $\tilde{O}\left(κ_A\sqrt{κ_f}\log\left(1/ε\right)\right)$, $\tilde{O}\left(κ_A\sqrt{L_f}D/\sqrtε\right)$, and $\tilde{O}\left(κ_A L_fΔ/ε^2\right)$, respectively. Here, $κ_A$ is the condition number of the matrix $A$ in the composition, $L_f$ is the smoothness constant of $f$, and $κ_f$ is the condition number of $f$ in the strongly convex case. $D$ is the initial point distance and $Δ$ is the initial function value gap. Tight lower complexity bounds for the three cases are also derived and they match the upper bounds up to logarithmic factors, thereby demonstrating the optimality of both the upper and lower bounds proposed in this paper.

math.OC↗