AlphaAdam:Asynchronous Masked Optimization with Dynamic Alpha for Selective Updates

TL;DR

AlphaAdam introduces an innovative optimization framework for large language models that employs dynamic intra-layer parameter masking and adaptive update strengths, leading to faster convergence and enhanced training stability.

Contribution

It proposes a novel intra-layer asynchronous masked optimization method with dynamic alpha adjustment, improving efficiency and stability over existing optimizers for LLM training.

Findings

Outperforms AdamW in convergence speed

Enhances training stability for LLMs

Applicable to various large models like GPT-2, RoBERTa, and Llama-7B

Abstract

In the training of large language models (LLMs), updating parameters more efficiently and stably has always been an important challenge. To achieve efficient parameter updates, existing methods usually achieve performance comparable to full parameter updates through methods such as low-dimensional decomposition or layer-wise selective updates. In this work, we propose AlphaAdam, an optimization framework for LLM from the perspective of intra-layer parameter updates. By decoupling parameter updates and dynamically adjusting their strength, AlphaAdam accelerates convergence and improves training stability. We construct parameter masks based on the consistency of historical momentum and gradient direction and combine them with an adaptive mask strength strategy to ensure efficient optimization and theoretical convergence guarantees, which is also applicable to most momentum-based optimizers. Extensive experiments show that AlphaAdam outperforms state-of-the-art methods such as AdamW in terms of convergence speed and computational efficiency across tasks, including GPT-2 pre-trained and fine-tuned RoBERTa and Llama-7B. Our AlphaAdam implements an optimizer enhancement framework for LLMs through intra-layer asynchronous masked adaptive updates. Our code is available in this https://github.com/MaeChd/AlphaAdam.