Xmixers(1) 初始化对于模型的影响

在本博客中，我们将正式开始xmixers的旅程，我们将讨论一个问题，初始化对于模型的收敛性的影响。

模型架构

这里我会选择四个模型架构进行研究（这里的LLaMA使用了LayerNorm，事实上，我发现LayerNorm和RMSNorm效果基本一致）：

LLaMA
- 事实上，这里和标准的Llama有一些不同，我会在最后讨论该问题；
GPT；
- Transformer with learnable pe，和nanoGPT中相同；
GPT-sincos;
- Transformer with sincos pe；

另一方面，为了验证我的实现是否正确，我还将nanoGPT中的GPT模型作为我的baseline。

我们会测试三个尺寸的模型，分别是：

small: 124M
medium: 350M
large: 774M

具体参数请参考nanoGPT.

实验

Method1

Method1来自于huggingface的初始化策略：

使用$\mathcal N(0, 0.02^2)$初始化所有层；

下面是学习曲线的对比：

从上图中我们可以看到几个现象：

GPT和GPT-baseline的loss相差非常多，但是这两者应该是一样的模型，是什么造成了这个不同？
LLaMA相比于GPT-baseline更好；
GPT-sincos的loss不太稳定，并且效果远逊于GPT-baseline，这体现了Learnable PE的优势；

因为发现了现象1，我不得不仔细对比nanoGPT和xmixers关于GPT的实现不同，经过仔细对比，我发现了只有初始化部分略有不同，nanoGPT添加了如下的初始化

# apply special scaled init to the residual projections, per GPT-2 paper
for pn, p in self.named_parameters():
	if pn.endswith('c_proj.weight'):
		torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))

即将Attention, FFN层的最后一个Linear层用$\mathcal N(0, (0.02/\sqrt{2L})^2)$初始化，其中$L$是Transformer Block的数量。

Method2

结合之前的分析，我们得到了Method2：

$\mathrm{std} = 0.02$；
使用$\mathcal N(0, {\mathrm{std}}^2)$初始化Attention, FFN层的最后一个Linear层以外的层以及Embedding层；
使用$\mathcal N(0, ({\mathrm{std}}/\sqrt{2L})^2)$初始化Attention, FFN层的最后一个Linear层；

我们重新进行实验，下面是学习曲线的对比：

可以观察到如下现象：

GPT和GPT-baseline的loss非常接近，说明关于层的scale起了作用；
LLaMA的loss变化不大，说明关于层的scale对于LLaMA影响不大；
GPT-sincos直接spike了，看起来关于层的scale对于GPT-sincos起了反作用；

因为我希望对齐GPT和GPT-basline，所以在后续的实验中，我将默认添加关于层的scale，然后通过别的办法解决GPT-sincos的spike的问题。

Method3

最近一篇论文指出，使用truncated normal来初始化weight会让训练效果更好，并且减少spike，代码为：

elif (
    self.config.init_type == 1
):  # credit to https://arxiv.org/pdf/2409.02060#page=14.84
    std = self.config.init_std
    trunc_std = 3 * std
    if isinstance(module, nn.Linear):
        nn.init.trunc_normal_(
            module.weight, mean=0.0, std=std, a=-trunc_std, b=trunc_std
        )
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.trunc_normal_(
            module.weight, mean=0.0, std=std, a=-trunc_std, b=trunc_std
        )
        if module.padding_idx is not None:
            nn.init.zeros_(module.weight[module.padding_idx])

由此我们得到了Method3：

$\mathrm{std} = 0.02$；
使用$\mathrm{clip}(\mathcal N(0, {\mathrm{std}}^2),-3\mathrm{std}, 3\mathrm{std})$初始化Attention, FFN层的最后一个Linear层以外的层以及Embedding层；
使用$\mathrm{clip}(\mathcal N(0, ({\mathrm{std}}/\sqrt{2L})^2),-3{\mathrm{std}}/\sqrt{2L}, 3{\mathrm{std}}/\sqrt{2L})$初始化Attention, FFN层的最后一个Linear层；

我们进行实验，下面是学习曲线的对比：

很不幸的是，GPT-sincos的loss依然spike了。然后我们进一步对比一下method2和method3的效果：

从上图中我们可以看出，method3效果相比于method2更好。

Method4

我们先将注意力转移一下，研究Attention部分的初始化，这里我测试了fla和fairseq关于Attention的初始化Method，具体分别为：

def _initialize_weights(self, module):
    if getattr(module, "_is_hf_initialized", False):
        return

    if self.token_mixer_init_type == 0:
        pass
    elif self.token_mixer_init_type == 1:  # fla init
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight, gain=2**-2.5)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
    elif self.token_mixer_init_type == 2:  # fairseq init
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight, gain=2**-0.5)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
    module._is_hf_initialized = True

method 4和method2结合，由此我们得到了Method4.1：

$\mathrm{std} = 0.02$；
使用$\mathcal U(-a, a)$初始化$W_q,W_k, W_v$；
- $a$根据方法的不同进行选择；
使用$\mathcal N(0, {\mathrm{std}}^2)$初始化FFN层的最后一个Linear层以外的层以及Embedding层；
使用$\mathcal N(0, ({\mathrm{std}}/\sqrt{2L})^2)$初始化Attention, FFN层的最后一个Linear层；

这里我们只测试了LLaMA模型，下面是学习曲线的对比：

可以看到，效果上：fairseq > 无特殊初始化 > fla。

作为了一个补充实验，我们将fairseq attention init和method 3进行结合，得到Method4.2：

$\mathrm{std} = 0.02$；
使用$\mathcal U(-a, a)$初始化$W_q,W_k, W_v$；
使用$\mathrm{clip}(\mathcal N(0, {\mathrm{std}}^2),-3\mathrm{std}, 3\mathrm{std})$初始化Attention, FFN层的最后一个Linear层以外的层以及Embedding层；
使用$\mathrm{clip}(\mathcal N(0, ({\mathrm{std}}/\sqrt{2L})^2),-3{\mathrm{std}}/\sqrt{2L}, 3{\mathrm{std}}/\sqrt{2L})$初始化Attention, FFN层的最后一个Linear层；

学习曲线的对比如下：