Больше нет Адама: масштабирование скорости обучения при инициализации - все, что вам нужно.

No More Adam: Learning Rate Scaling at Initialization is All You Need

December 16, 2024

Авторы: Minghao Xu, Lichuan Xiang, Xu Cai, Hongkai Wen

cs.AI

Аннотация

В данной работе мы рассматриваем необходимость адаптивных методов градиентов для обучения глубоких нейронных сетей. SGD-SaI представляет собой простое, но эффективное улучшение стохастического градиентного спуска с импульсом (SGDM). SGD-SaI выполняет масштабирование скорости обучения при инициализации (SaI) для различных групп параметров, руководствуясь их соответствующими отношениями сигнал-шум градиента (g-SNR). Путем корректировки скоростей обучения без использования адаптивного второго порядка импульса, SGD-SaI помогает предотвратить дисбаланс обучения уже с первой итерации и сокращает использование памяти оптимизатором наполовину по сравнению с AdamW. Несмотря на свою простоту и эффективность, SGD-SaI последовательно соответствует или превосходит AdamW при обучении различных задач на основе трансформеров, успешно преодолевая давнюю проблему использования SGD для обучения трансформеров. SGD-SaI проявляет себя в классификации ImageNet-1K с использованием Vision Transformers (ViT) и предварительном обучении GPT-2 для больших языковых моделей (LLM, только декодер трансформера), демонстрируя устойчивость к вариациям гиперпараметров и пригодность для различных приложений. Мы также проверили его устойчивость на задачах, таких как донастройка LoRA для LLM и моделей диффузии, где он последовательно превосходит оптимизаторы нового поколения. С точки зрения эффективности использования памяти, SGD-SaI достигает значительных экономий памяти для состояний оптимизатора, сокращая использование памяти на 5,93 ГБ для GPT-2 (1,5 млрд параметров) и на 25,15 ГБ для Llama2-7B по сравнению с AdamW в настройках обучения с полной точностью.

English

In this work, we question the necessity of adaptive gradient methods for training deep neural networks. SGD-SaI is a simple yet effective enhancement to stochastic gradient descent with momentum (SGDM). SGD-SaI performs learning rate Scaling at Initialization (SaI) to distinct parameter groups, guided by their respective gradient signal-to-noise ratios (g-SNR). By adjusting learning rates without relying on adaptive second-order momentum, SGD-SaI helps prevent training imbalances from the very first iteration and cuts the optimizer's memory usage by half compared to AdamW. Despite its simplicity and efficiency, SGD-SaI consistently matches or outperforms AdamW in training a variety of Transformer-based tasks, effectively overcoming a long-standing challenge of using SGD for training Transformers. SGD-SaI excels in ImageNet-1K classification with Vision Transformers(ViT) and GPT-2 pretraining for large language models (LLMs, transformer decoder-only), demonstrating robustness to hyperparameter variations and practicality for diverse applications. We further tested its robustness on tasks like LoRA fine-tuning for LLMs and diffusion models, where it consistently outperforms state-of-the-art optimizers. From a memory efficiency perspective, SGD-SaI achieves substantial memory savings for optimizer states, reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW in full-precision training settings.