Controlling stochastic gradient descent using stochastic approximation for robust distributed optimization
Adit Jain, Vikram Krishnamurthy · Numerical Algebra Control and Optimization · 2024
This paper deals with the problem of controlling the stochastic gradient descent, performed by multiple learners where the aim is to estimate the respective $ \arg\min f $ using noisy gradients obtained by querying a stochastic oracle. Each query has a learning cost, and the noisy gradient response has varying degrees of noise variance, the bound of which is assumed to vary in a Markovian fashion. For a single learner, the decision problem is to choose when to query the oracle such that the learning cost is minimized. A constrained Markov decision process (CMDP) is formulated to solve the decision problem of a single learner. Structural results are proven for the optimal policy for the CMDP, which is shown to be threshold decreasing in the queue state. For multiple learners, a constrained switching control game is formulated for scheduling and controlling $ N $ learners querying the same oracle, one at a time. The structural results are extended for the optimal policy achieving the Nash equilibrium. The structural results are used to propose a stochastic approximation algorithm to search for the optimal policy, which tracks the parameters of the optimal policy using a sigmoidal approximation and does not require knowledge of the underlying transition probabilities. The paper also briefly discusses applications in federated learning and numerically shows the convergence properties of the proposed algorithm.