试想,如果我们定义函数,其中 为有界的连续函数。与此同时,我们再定义函数 ,记 是使得该函数最小化时的最优参数。
根据 CLT 我们可以自然地推导出当 时, 就是纯粹的以 为自变量的函数,而且它是收敛的。相反,原始的函数 随着 的不同可能会发散。不仅如此,当 足够大时,对任意 都有 接近于.
类比过来也就是说,当我们想训练一个大模型,希望其损失最小化时,就相当于期望最小化函数 。此时,传统的参数化方法相当于是在对 进行寻优,而 muP 技术则是在对 寻优。并且,前者得到的最优参数不具有随宽度变化的稳定性,即不可迁移到任意更大的 的情况,而后者可以。
作者以单值线性投影 为例,简要说明了经过一层的隐藏层,然后利用 SGD 算法对参数 和 进行一次更新后,输出的分布范围。我们称之为 “activation scales”(激活规模)。
1
然后通过 SGD 算法对参数进行一次更新之后,. 这样一来,新的函数 因为其表达式内含 而产生爆炸。在利用 Transformer 模型做1
次到 5
次的更新的实验可以进一步说明该问题。与输入宽度有关是输出值 logits
, attn logits
使用 SP 方法都产生了爆炸,而 μP 则相对很稳定。其中,word embedding
是与宽度无关的层,所以二者都比较稳定。
值得一提的是,作者还指出,SP 方法因为统一了学习率,所以“如果将学习率也按照宽度缩放不就能解决SP方法的迁移问题了吗?”的说法并不可行。因为这会导致不依赖于宽度的那些层(比如 Word Embedding)的更新也一同变得缓慢了,导致学习不到有用的信息。
安装 mup
库:
1 | pip install mup |
在代码层面,做如下调整:
MuReadout(...)
torchdistx
来避免直接实例化模型set_base_shapes(...)
mup
自带的参数初始化方法 mup.init
替代 torch.init
mup
自带的优化器 mup.optim
替代 torch.optim
1 | from mup import (MuReadout, |