Distributed Sign Momentum with Local Steps for Training Transformers

• URL: http://arxiv.org/abs/2411.17866v2
• Date: Fri, 07 Mar 2025 19:35:00 GMT
• Title: Distributed Sign Momentum with Local Steps for Training Transformers
• Authors: Shuhua Yu, Ding Zhou, Cong Xie, An Xu, Zhi Zhang, Xin Liu, Soummya Kar,
• Abstract summary: Pre-training Transformer models are resource-intensive.<n>Recent studies have shown that sign momentum is an efficient technique for training large-scale deep learning models.<n>This paper investigates a novel communication momentum with multiple broad steps.
• Score: 21.046099659465508
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Pre-training Transformer models is resource-intensive, and recent studies have shown that sign momentum is an efficient technique for training large-scale deep learning models, particularly Transformers. However, its application in distributed training remains underexplored. This paper investigates a novel communication-efficient distributed sign momentum method with multiple local steps, to cope with the scenarios where communicating at every step is prohibitive. Our proposed method allows for a broad class of base optimizers for local steps, and uses sign momentum in the global step, where momentum is generated from differences accumulated during local steps. For generic base optimizers, by approximating the sign operator with a randomized version that acts as a continuous analog in expectation, we present a general convergence analysis, which specializes to an $O(1/\sqrt{T})$ rate for a particular instance. When local step is stochastic gradient descent, we show an optimal $O(1/T^{1/4})$ rate in terms of $\ell_1$ gradient norm for nonconvex smooth cost functions. We extensively evaluate our method on the pre-training of various sized GPT-2 models from scratch, and the empirical results show significant improvement compared to other distributed methods with multiple local steps.

Related papers

• Benefits of Learning Rate Annealing for Tuning-Robustness in Stochastic Optimization [29.174036532175855]
  Learning rate in gradient methods is a critical hyperspecification that is notoriously costly to tune via standard grid search. We identify a theoretical advantage of learning rate annealing schemes that decay the learning rate to zero at a rate, such as the widely-used cosine schedule.
  arXiv  Detail & Related papers  (2025-03-12T14:06:34Z)
• Distributed Gradient Descent with Many Local Steps in Overparameterized Models [20.560882414631784]
  In distributed training of machine learning models, gradient descent with local iterative steps is a popular method. We try to explain this good performance from a viewpoint of implicit bias in Local Gradient Descent (Local-GD) with a large number of local steps.
  arXiv  Detail & Related papers  (2024-12-10T23:19:40Z)
• In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics [34.458004744956334]
  We prove that there exists a transformer capable of achieving a prediction error of order $mathcalO(sqrtd/n)$ with high probability. We also analyze the training dynamics of transformers with single linear self-attention layers, demonstrating that, with appropriately parameters, gradient flow optimization over the population mean square loss converges to a global optimum.
  arXiv  Detail & Related papers  (2024-10-18T05:28:47Z)
• Convergence of Distributed Adaptive Optimization with Local Updates [3.895864050325129]
  We study distributed adaptive algorithms with local updates (intermittent communication) We prove that em Local SGD em with momentum (em Local em SGDM) and em Local em Adam can outperform their minibatch counterparts in convex and weakly convex settings.
  arXiv  Detail & Related papers  (2024-09-20T01:45:10Z)
• Adaptive Federated Learning Over the Air [108.62635460744109]
  We propose a federated version of adaptive gradient methods, particularly AdaGrad and Adam, within the framework of over-the-air model training. Our analysis shows that the AdaGrad-based training algorithm converges to a stationary point at the rate of $mathcalO( ln(T) / T 1 - frac1alpha ).
  arXiv  Detail & Related papers  (2024-03-11T09:10:37Z)
• Towards Seamless Adaptation of Pre-trained Models for Visual Place Recognition [72.35438297011176]
  We propose a novel method to realize seamless adaptation of pre-trained models for visual place recognition (VPR) Specifically, to obtain both global and local features that focus on salient landmarks for discriminating places, we design a hybrid adaptation method. Experimental results show that our method outperforms the state-of-the-art methods with less training data and training time.
  arXiv  Detail & Related papers  (2024-02-22T12:55:01Z)
• Multiply Robust Estimation for Local Distribution Shifts with Multiple Domains [9.429772474335122]
  We focus on scenarios where data distributions vary across multiple segments of the entire population. We propose a two-stage multiply robust estimation method to improve model performance on each individual segment. Our method is designed to be implemented with commonly used off-the-shelf machine learning models.
  arXiv  Detail & Related papers  (2024-02-21T22:01:10Z)
• Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization [87.21285093582446]
  Diffusion Generative Flow Samplers (DGFS) is a sampling-based framework where the learning process can be tractably broken down into short partial trajectory segments. Our method takes inspiration from the theory developed for generative flow networks (GFlowNets)
  arXiv  Detail & Related papers  (2023-10-04T09:39:05Z)
• Provable and Practical: Efficient Exploration in Reinforcement Learning via Langevin Monte Carlo [104.9535542833054]
  We present a scalable and effective exploration strategy based on Thompson sampling for reinforcement learning (RL) We instead directly sample the Q function from its posterior distribution, by using Langevin Monte Carlo. Our approach achieves better or similar results compared with state-of-the-art deep RL algorithms on several challenging exploration tasks from the Atari57 suite.
  arXiv  Detail & Related papers  (2023-05-29T17:11:28Z)
• Model-Based Reinforcement Learning with Multinomial Logistic Function Approximation [10.159501412046508]
  We study model-based reinforcement learning (RL) for episodic Markov decision processes (MDP) We establish a provably efficient RL algorithm for the MDP whose state transition is given by a multinomial logistic model. To the best of our knowledge, this is the first model-based RL algorithm with multinomial logistic function approximation with provable guarantees.
  arXiv  Detail & Related papers  (2022-12-27T16:25:09Z)
• Generalized Differentiable RANSAC [95.95627475224231]
  $nabla$-RANSAC is a differentiable RANSAC that allows learning the entire randomized robust estimation pipeline. $nabla$-RANSAC is superior to the state-of-the-art in terms of accuracy while running at a similar speed to its less accurate alternatives.
  arXiv  Detail & Related papers  (2022-12-26T15:13:13Z)
• Distributionally Robust Models with Parametric Likelihood Ratios [123.05074253513935]
  Three simple ideas allow us to train models with DRO using a broader class of parametric likelihood ratios. We find that models trained with the resulting parametric adversaries are consistently more robust to subpopulation shifts when compared to other DRO approaches.
  arXiv  Detail & Related papers  (2022-04-13T12:43:12Z)
• Parallel Successive Learning for Dynamic Distributed Model Training over Heterogeneous Wireless Networks [50.68446003616802]
  Federated learning (FedL) has emerged as a popular technique for distributing model training over a set of wireless devices. We develop parallel successive learning (PSL), which expands the FedL architecture along three dimensions. Our analysis sheds light on the notion of cold vs. warmed up models, and model inertia in distributed machine learning.
  arXiv  Detail & Related papers  (2022-02-07T05:11:01Z)
• MDPGT: Momentum-based Decentralized Policy Gradient Tracking [29.22173174168708]
  We propose a momentum-based decentralized policy gradient tracking (MDPGT) for multi-agent reinforcement learning. MDPGT achieves the best available sample complexity of $mathcalO(N-1epsilon-3)$ for converging to an $epsilon-stationary point of the global average of $N$ local performance functions. This outperforms the state-of-the-art sample complexity in decentralized model-free reinforcement learning.
  arXiv  Detail & Related papers  (2021-12-06T06:55:51Z)
• On Second-order Optimization Methods for Federated Learning [59.787198516188425]
  We evaluate the performance of several second-order distributed methods with local steps in the federated learning setting. We propose a novel variant that uses second-order local information for updates and a global line search to counteract the resulting local specificity.
  arXiv  Detail & Related papers  (2021-09-06T12:04:08Z)
• The Gradient Convergence Bound of Federated Multi-Agent Reinforcement Learning with Efficient Communication [20.891460617583302]
  The paper considers independent reinforcement learning (IRL) for collaborative decision-making in the paradigm of federated learning (FL) FL generates excessive communication overheads between agents and a remote central server. This paper proposes two advanced optimization schemes to improve the system's utility value.
  arXiv  Detail & Related papers  (2021-03-24T07:21:43Z)
• Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks [78.76880041670904]
  In neural networks with binary activations and or binary weights the training by gradient descent is complicated. We propose a new method for this estimation problem combining sampling and analytic approximation steps. We experimentally show higher accuracy in gradient estimation and demonstrate a more stable and better performing training in deep convolutional models.
  arXiv  Detail & Related papers  (2020-06-04T21:51:21Z)
• Non-asymptotic Convergence of Adam-type Reinforcement Learning Algorithms under Markovian Sampling [56.394284787780364]
  This paper provides the first theoretical convergence analysis for two fundamental RL algorithms of policy gradient (PG) and temporal difference (TD) learning. Under general nonlinear function approximation, PG-AMSGrad with a constant stepsize converges to a neighborhood of a stationary point at the rate of $mathcalO(log T/sqrtT)$. Under linear function approximation, TD-AMSGrad with a constant stepsize converges to a neighborhood of the global optimum at the rate of $mathcalO(log T/sqrtT
  arXiv  Detail & Related papers  (2020-02-15T00:26:49Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.