24 Mar

高阶muP:更简明但更高明的谱条件缩放

在文章《初探muP:超参数的跨模型尺度迁移规律》中,我们基于前向传播、反向传播、损失增量和特征变化的尺度不变性推导了muP(Maximal Update Parametrization)。可能对于部分读者来说,这一过程还是显得有些繁琐,但实际上它比原始论文已经明显简化。要知道,我们是在单篇文章内相对完整地介绍的muP,而muP的论文实际上是作者Tensor Programs系列论文的第5篇!

不过好消息是,作者在后续的研究《A Spectral Condition for Feature Learning》中,发现了一种新的理解方式(下称“谱条件”),它比muP的原始推导和笔者的推导都更加直观和简洁,但却能得到比muP更丰富的结果,可谓muP的高阶版本,简明且不失高明的代表作。

准备工作

顾名思义,谱条件(Spectral Condition)跟谱范数(Spectral Norm)相关,它的出发点是谱范数的一个基本不等式:
\begin{equation}\Vert\boldsymbol{x}\boldsymbol{W}\Vert_2\leq \Vert\boldsymbol{x}\Vert_2 \Vert\boldsymbol{W}\Vert_2\label{neq:spec-2}\end{equation}

点击阅读全文...

28 Mar

MoE环游记:4、难处应当多投入

前两篇文章我们都在讨论负载均衡,其中在《MoE环游记:3、换个思路来分配》介绍Loss-Free方案时,笔者留了一个悬念:它引入的Bias项有一个冗余的自由度,这个自由度可以用来做另外有趣的事情。这篇文章我们就来讨论这件事。

我们知道,MoE是为每个Token只选择最匹配的$k$个Expert来进行计算,从而在增大参数量的同时还节省了计算量。然而,当我们仔细思考就会发现,这个策略实际上有明显的可改进之处:直观来看,每个Token的难度并不一样,所以更合理的方案应该是难的Token分配更多的计算资源,简单的token分配更少的资源,这样或许能在同样有限的资源下将效果最大化。

而刚才提到的Bias的额外自由度,恰好可以用来简单地实现这个目标。

点击阅读全文...