通过模型分片扩展零阶预训练
Scaling Zero-Order Pretraining through Model Sharding
浏览论文内容
中文总结 AI 辅助
本文提出分片优化混合组装（SOMA）方法，通过独立训练LSTM专家并采用可分离损失，在零阶优化中降低梯度方差，实现高效的大模型预训练，并在测试损失和推理速度上优于单体基线。
中文摘要 AI 辅助
零阶优化（ZO）无需反向传播即可进行训练，这使得它适用于仅前向硬件和非可微损失函数，但其梯度方差会随着扰动维度增加而增大，从而阻碍了大模型的训练。分片优化混合组装（SOMA）方法使用同时扰动随机逼近（SPSA）在$N$个数据簇上独立训练LSTM专家，无需交换梯度、激活或优化器状态。其可分离损失消除了跨专家的扰动噪声，但代价是放弃了跨领域的联合学习表示。利用估计的80,000个RTX 5090 GPU小时，我们证明适度的分片在所有测试的单体零阶对照方法中均能提高训练计算效率。在8.44M参数和150个聚合GPU小时下，SOMA $N=2$配合64次扰动达到了1.76的测试nats/字节，而单体SPSA在64、256或1024次扰动下为2.00至2.11，EGGROLL为2.21。在WikiText-103上，这些冻结的检查点分别达到了2.07、2.25至2.36和2.49。在具有相等大小块的固定可分离目标上，我们证明了独立损失能将相对梯度方差降低到共享损失估计器的约$1/N$。在保持起始权重、数据、扰动和计算量固定的情况下，独立损失而非求和损失使SOMA $N=4$在三个种子的1,000次更新后测试损失降低了0.035 nats/字节。更大的集成提供了单独的推理优势：在相似的模型大小和top-$k$路由（$k=4$）下，SOMA $N=256$达到了2.36M tokens/s，而SOMA $N=8$为257k（包括路由在内为$9.19\ imes$），且测试损失更低（1.68对比1.71），尽管使用了$59.9\ imes$的聚合训练计算量。我们发布了所有训练和评估代码及检查点。
英文摘要
Zero-order optimization (ZO) trains without backpropagation, making it relevant to forward-only hardware and non-differentiable loss, but its gradient variance grows with perturbed dimension, inhibiting large-model training. Sharded Optimization Mixture of Assemblies (SOMA) trains LSTM experts independently on $N$ data clusters using simultaneous perturbation stochastic approximation (SPSA), without exchanging gradients, activations or optimizer state. Its separable loss removes cross-expert perturbation noise at the cost of jointly learned representations across domains. Using 80,000 estimated RTX 5090 GPU-hours, we show modest sharding improves training compute efficiency over all tested monolithic ZO controls. At 8.44M parameters and 150 aggregate GPU-hours, SOMA $N=2$ with 64 perturbations reaches 1.76 test nats/byte, versus 2.00--2.11 for monolithic SPSA at 64, 256 or 1,024 perturbations and 2.21 for EGGROLL. On WikiText-103, these frozen checkpoints reach 2.07, 2.25--2.36 and 2.49, respectively. On a fixed separable objective with equal-size blocks, we prove independent losses reduce relative gradient variance to approximately $1/N$ of a shared-loss estimator's. Holding starting weights, data, perturbations and compute fixed, independent rather than summed losses lower SOMA $N=4$ test loss by 0.035 nats/byte after 1,000 updates across three seeds. Larger ensembles offer a separate inference benefit: at similar model size with top-$k$ routing ($k=4$), SOMA $N=256$ achieves 2.36M tokens/s versus 257k for SOMA $N=8$ ($9.19\times$, including routing), at lower test loss (1.68 versus 1.71), albeit using $59.9\times$ as much aggregate training compute. We release all training and evaluation code and checkpoints.
发表机构
- Stanford University（斯坦福大学）
机构由 AI 辅助整理，请以论文原文为准。
