[论文速览] Small-scale proxies for large-scale Transformer training instabilities

Pre

title: Small-scale proxies for large-scale Transformer training instabilities
source: ICLR 2024
paper: https://arxiv.org/abs/2309.14322
code:
ref: 小尺度Transformer如何Scale

Idea

这篇文章想通过小模型来研究大模型（Transformer）训练不稳定的原因

Method（Model）

图1 橙色的是不使用Qk-layernorm，蓝色则是启用，N是参数量，LR sensitivity是作者提出用于衡量学习率变化时与最优值的预期偏差，越小越好，qk-layernorm可以减低但无法阻止该值随着参数量增加而上涨。

注：qk-layernorm就是在qk相乘之前对二者分别做一次layernorm

图2 研究不同参数量和学习率设置下 attention logit 增长的不稳定性，这里的attention logits 就是z，是qk相乘后还没过softmax的结果。

图3 输出logit发散不稳定的例子，实线使用了z-loss，效果明显比weight decay（权重衰减来得好）

输出logit一般先过softmax得到类别概率，而作者说不稳定性是由于logit发散且变成比较小的负数，通过将softmax的分母记为Z，z-loss就是 log Z的平方。

Critique

qk-layernorm, z-loss 有用
论文挺长，后面还有Warm-up什么的，有机会再看