High-Probability PL-SGD with Markovian Noise: Optimal Mixing and Tail Dependence
Dhruv Sarkar, Aprameyo Chakrabartty, Vaneet Aggarwal
Unverified — Be the first to reproduce this paper.
ReproduceAbstract
We study first-order methods for smooth objectives satisfying the Polyak-Łojasiewicz (PL) condition when gradient samples are generated by an exogenous Markov chain. In the light-tailed setting, prior uniform-in-time high-probability bounds for ordinary Stochastic Gradient Descent (SGD) under a standard growth envelope scale as O(t_mix^2/k), leaving a gap with the O(t_mix/k) expectation bounds. We close this gap using a lag-blocking argument to establish a uniform high-probability guarantee with a leading stochastic term of O(t_mix/(k+K_0)) under geometric mixing. We prove this linear dependence on the mixing time is optimal via a matching Ω(σ^2 t_mix/k) lower bound on a quadratic objective driven by a persistent two-state chain. We then extend this framework to heavy-tailed Markovian gradients satisfying a stationary finite-p-moment condition, p (1,2]. We design an all-samples clipped block method that uses every Markov transition while mitigating Markovian bias. Under a transition budget T, this algorithm achieves a high-probability stochastic error of O(σ_p^2(t_mix/T)^2(p-1)/p). We establish a matching lower bound by reducing PL optimization to heavy-tailed mean estimation for a sticky Markov chain. Ultimately, this work tightly characterizes the optimal polynomial dependence on mixing time for light-tailed PL-SGD, and the optimal heavy-tail exponent and effective-sample-size dependence in the robust regime.