用低秩扰动加速进化策略,百亿参数模型训练快100倍
Evolution Strategies at the Hyperscale
- 将随机扰动设计为低秩矩阵,提升计算密集度
- 百亿参数模型训练速度提升百倍,接近纯批量推理吞吐量
- 适合大规模强化学习和语言模型预训练场景
进化策略(ES)是一类强大的黑箱优化方法,具备高度并行性,可处理非可微和噪声目标。然而,朴素的ES在大规模GPU上因批量矩阵乘法的算术强度过低而变得极其昂贵。本文提出基于低秩学习的进化引导通用优化方法(EGGROLL),通过将个体扰动结构化为秩-r矩阵,使算术强度显著提升,在大种群规模下对百亿参数模型的训练速度提升百倍,达到纯批量推理吞吐量的91%。我们对高维参数空间下的高斯进化策略进行了严格理论分析,揭示了其在高维下具有线性化效应,并证明了EGGROLL与标准ES在参数维度增大时的一致性。实验表明:(1) EGGROLL可稳定预训练仅使用整数数据类型的非线性循环语言模型;(2) 在推理任务中,其性能与GRPO相当;(3) 在无初始化强化学习设置下,性能不逊于原始ES,但速度更快。
原文摘要 · Abstract (English)
Evolution Strategies (ES) is a class of powerful black-box optimisation methods that are highly parallelisable and can handle non-differentiable and noisy objectives. However, naïve ES becomes prohibitively expensive at scale on GPUs due to the low arithmetic intensity of batched matrix multiplications with unstructured random perturbations. We introduce Evolution Guided GeneRal Optimisation via Low-rank Learning (EGGROLL), which improves arithmetic intensity by structuring individual perturbations as rank-$r$ matrices, resulting in a hundredfold increase in training speed for billion-parameter models at large population sizes, achieving up to 91% of the throughput of pure batch inference. We provide a rigorous theoretical analysis of Gaussian ES for high-dimensional parameter objectives, investigating conditions needed for ES updates to converge in high dimensions. Our results reveal a linearising effect, and proving consistency between EGGROLL and ES as parameter dimension increases. Our experiments show that EGGROLL: (1) enables the stable pretraining of nonlinear recurrent language models that operate purely in integer datatypes, (2) is competitive with GRPO for post-training LLMs on reasoning tasks, and (3) does not compromise performance compared to ES in tabula rasa RL settings, despite being faster.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。