Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer
Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, Jianfeng Gao
NeurIPS 2021
cs.LG, cond-mat.dis-nn, cs.NE
2022-03-07
微软与OpenAI用μP把超参从一千三百万代理零样本拷到三亿五千万BERT-large,测试损失1.683对默认1.731;四千万代理调六十七亿GPT-3仅花预训练百分之七。
超参调不好,大模型直接训不起。标准参数化里,最优学习率随宽度跑,宽模型甚至可能比窄的差。Tensor Programs 系列前面几篇给出了 Maximal Update Parametrization(μP):让每层更新对激活的影响不随宽度消失或爆炸。这篇是它的工程兑现,来自 Microsoft 与 OpenAI,发表于 NeurIPS 2021。
做法叫 μTransfer:目标模型用 μP 参数化,超参在小代理模型上调,再零样本拷到大模型,大模型本身不再搜。
μP 按层宽改初始化方差和学习率。对 SGD,输入层学习率随 fan-out 放大、输出层随 fan-in 缩小、隐层按 1/fan-in;输出层初始化方差是 1/fanin²,比常见的 1/fanin 再小一截。Adam 的换算不同:隐层学习率不随宽度变,输出层仍除以 fan-in。Transformer 还要把注意力 logits 从 1/√d 改成 1/d,因为训练后 q 和 k 相关,点积按大数定律走 d 而不是 √d。
可迁移的是优化相关超参:学习率、Adam β、学习率日程、初始化方差、参数乘数。不可迁移的是正则,dropout 和 weight decay 同时依赖模型容量和数据量。可以「跨过去」的尺度包括宽度,以及经验上的深度、batch、序列长度、训练步数。深度迁移只在 pre-LN Transformer 上成立,初始化标准差跨深度也不稳,实践里往往把 init 钉死、只调别的。
代理模型要够大才像目标:文中经验是宽度至少约 256、深度约 4、batch 约 32、序列长度约 128、训练步大约 5000。再小,最优点会漂。
IWSLT14 De-En 上,4M 的 0.25× 代理往 40M 的 1× 模型拷超参,BLEU 分位 35.27 到 35.33,高于 fairseq 默认的 33.62 到 35.00;同样预算下直接在 1× 上搜,以及标准参数化的朴素迁移,都会发散或更差。WMT14 En-De 上,15M 代理拷到 211M,BLEU 25.94 到 26.34,对上 fairseq 默认 25.69。
BERT 用约 13M 的 prototype 同时服务 base 与 large,调参总 FLOPs 约等于训一次 BERT-large。相对 Megatron 默认,μTransfer 的 BERT-base 测试损失 1.970 对 1.995,MNLI 84.3/84.8 对 84.2/84.2,QQP 90.8 对 90.6;BERT-large 测试损失 1.683 对 1.731,MNLI 87.0/86.5 对 86.3/86.2,QQP 91.4 对 90.9。朴素迁移两次都发散。
GPT-3 6.7B(32 层、宽 4096)从宽 256、约 40M 的代理迁移,代理小 168 倍,调参成本约预训练的 7%。验证损失 1.98,对照按原超参重跑的 2.03。LAMBADA zero-shot 73.49,对论文 6.7B 的 70.3 和 13B 的 72.5;HellaSwag zero-shot 72.0,对 67.4 和 70.9。两个caveat 必须一起读:μP 模型因数值问题改用 FP32,原版和重跑是 FP16;重跑还误用了绝对注意力。困惑度用了自定义分词,不能直接跟文献比。
μP 下「更宽更好」贯穿训练过程;标准参数化里宽模型可以更差。作者把 GPT-3 式 Transformer 压到 4 层、宽度扫到 32,768 做压力测试,除约 1e8 token 附近一小段,更宽始终更低训练损失。
十亿参数模型没法在自身上做超参网格。μTransfer 把调参加回单卡小模型,一套超参还能复用整族宽度。pip install mup 把这件事做成了可抄的实现。它解决的是优化超参,不是正则,也不是换数据换任务之后的迁移。
初始化跨深度不好使,post-LN 的深度迁移也不成立。小代理的最优点仍会轻微漂移。正则超参明确不在范围内,小数据集微调尤其如此。GPT-3 对照混进了精度和注意力实现差异,6.7B 打过 13B 的部分数字不能当成干净消融。论文主要盯训练损失;测试变好是因为这几项实验里正则不是瓶颈。