提出可迁移超参数的MoE模型新参数化方法，让大模型训练更省力。
Hyperparameter Transfer with Mixture-of-Expert Layers
- 基于动态平均场理论设计新参数化，支持跨模型维度迁移超参数
- 在51M到20亿参数模型间实现稳定超参数转移，仅需少量调参
- 小模型调好参数后，可直接用于训练大模型长序列任务
混合专家（MoE）层通过将总参数量与前向传播激活参数解耦，成为扩展现代神经网络的重要工具。然而，稀疏MoE引入新训练参数（路由器权重）和架构尺度（专家数量与大小），增加了训练复杂性。为使超参数选择更高效可靠，本文提出一种适用于宽度、深度、专家数量及专家隐层大小扩展的Transformer模型新参数化方式。该参数化基于新颖的动力学平均场理论分析。在固定令牌预算下，实证表明该方法可在5100万至超过20亿总参数模型间实现可靠的超参数迁移。进一步地，利用小模型在短序列上搜索出的超参数，成功训练了大模型在长序列上的表现，展现出优异性能。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) layers have emerged as an important tool in scaling up modern neural networks by decoupling total trainable parameters from activated parameters in the forward pass for each token. However, sparse MoEs add complexity to training due to (i) new trainable parameters (router weights) that, like all other parameter groups, require hyperparameter (HP) tuning; (ii) new architecture scale dimensions (number of and size of experts) that must be chosen and potentially taken large. To make HP selection cheap and reliable, we propose a new parameterization for transformer models with MoE layers when scaling model width, depth, number of experts, and expert (hidden) size. Our parameterization is justified by a novel dynamical mean-field theory (DMFT) analysis. When varying different model dimensions trained at a fixed token budget, we find empirically that our parameterization enables reliable HP transfer across models from 51M to over 2B total parameters. We further take HPs identified from sweeping small models on a short token horizon to train larger models on longer horizons and report performant model behaviors.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。