FedPAGE Probabilistic Local Stochastic Gradient Method for Federated Learning

A communication-efficient local stochastic gradient method for federated convex and nonconvex optimization (Algorithm 1) that generalizes the optimal Probabilistic Gradient Estimator (PAGE) to multiple local update steps K >= 1: in each round r, with probability p_r = S/N the server queries all N clients for minibatch gradients of size b_1, and with probability 1 - p_r the server broadcasts (x^r, x^{r-1}, g^{r-1}) to a sampled subset S^r of size S where each client i initializes g_{i,0}^r = nabla_{I_2} f_i(x^r) - nabla_{I_2} f_i(x^{r-1}) + g^{r-1} and runs K - 1 local SARAH/PAGE recursive steps g_{i,k}^r = nabla_{I_3} f_i(y_{i,k}^r) - nabla_{I_3} f_i(y_{i,k-1}^r) + g_{i,k-1}^r, achieving O(L(sqrt(N) + S) / (S eps^2)) communication rounds in nonconvex settings and O(N^{3/4} L / (S eps)) in convex settings.