Non-convex Training using Imaginary Paths

Most machine learning methods including deep learning are defined by non-convex optimization problems that are typically regularized for enhanced generalization yet fundamentally challenged by loss landscapes with local minima. We show that regularization by the ${\ell }_1$ and ${\ell }_2$ norms can be expressed in terms of a radial-phase decomposition for which we prove the existence of a complex path with a bounded loss leaving any nonzero real stationary point under the stated assumptions. We explore this by designing a novel complex lifting framework from ${\mathbb{R}}^P$ to ${\mathbb{C}}^P$ that introduces a bilinear skew-symmetric auxiliary loss. We impose a consensus constraint between the lifted and original parameterization and use the generic alternating direction methods of multipliers (ADMM) for the optimization defining the proposed Complex-Continuation ADMM (CC-ADMM) method for non-convex optimization. After finite time complex exploration, CC-ADMM projects the model onto the real parameter space and continues real-valued optimization, reducing its convergence analysis to that of the real optimization phase. We show that, under the stated assumptions, the CC-ADMM converges to a stationary point of the original loss with probability one. CC-ADMM uses only twice the number of parameters during training with no resulting inference overhead. Experiments on synthetic non-convex problems confirm the predicted exploration behavior and when applied on a variety of non-convex machine learning tasks including "grokking", as well as Vision Transformers and Llama-1B and 3B finetuning, we observe enhanced generalization performance. The CC-ADMM thereby provides a versatile and favorable framework for non-convex training.