PyTorch
版本的魔改 《Transformer Quality in Linear Time》在中文 NLU 任务上的复现、评测试验。
参考:JunnYu 的 实现 与苏神的 实现 与 笔记
更多细节见整理的文章:
原论文提出了一种 将 Attention 与 FFN 相结合 的方案:GAU。
标准的 FFN 是两层 MLP :
$$
\begin{equation}\boldsymbol{O}=\phi(\boldsymbol{X}\boldsymbol{W}_u)\boldsymbol{W}_o\end{equation}
$$
其中 $\boldsymbol{X}\in\mathbb{R}^{n\times d},\boldsymbol{W}u\in\mathbb{R}^{d\times d{ff}},\boldsymbol{W}o\in\mathbb{R}^{d{ff}\times d}$。而后,有研究表明将 FFN 换成如下的 GLU 效果更好:
$$
\begin{equation}\boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{V})\boldsymbol{W}_o,\quad \boldsymbol{U}=\phi_u(\boldsymbol{X}\boldsymbol{W}_u),\quad\boldsymbol{V}=\phi_v(\boldsymbol{X}\boldsymbol{W}_v)\end{equation}
$$
其中
我们知道,在 Transformer 结构中,Attention 的作用是捕捉 Token-to-Token 间的关系;而 FFN 则是增强模型的非线性表达能力。而 GLU 只是用 自身 去 gate 自身,忽略了 Token 间关系,那么一个自然的想法是用 Token-to-Token 的关系也即 注意力矩阵 A 去 gate:
$$
\begin{equation}\boldsymbol{O}=(\boldsymbol{U}\odot\boldsymbol{A}\boldsymbol{V})\boldsymbol{W}_o\label{eq:mix}\end{equation}
$$
在原论文中,注意力矩阵被简化至了如下的形式:
$$
\begin{equation}\boldsymbol{A}=\text{relu}^2\left(\frac{\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}}{n}\right)=\frac{1}{n^2}\text{relu}^2\left(\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}\right),\quad \boldsymbol{Z}=\phi_z(\boldsymbol{X}\boldsymbol{W}_z)\label{eq:relu-att}\end{equation}
$$
其中
而后每一层只用一个 GAU 单元就行了,原来一层的计算量大约等于两层 GAU。而且 GAU 只需要一个头就行了。
原论文的主要思路便是这些了,很简洁但是也有很多不合理的地方:比如说缩放因子取的是
鉴于此,本文使用了苏神提出的改进版的 softmax:
$$
\begin{equation}\boldsymbol{A}=softmax\left(\frac{\log_{512} n}{\sqrt{d}}\mathcal{Q}(\boldsymbol{Z})\mathcal{K}(\boldsymbol{Z})^{\top}\right)\end{equation}
$$
- 预训练
- 2022/05/07 相关工作整理成文
- 2022/04/27 添加模型笔记、显存、推理速度比较结果。
- 2022/04/23 增加 CLUE 评测
- 2022/04/18 两种归一化策略比较
- 2022/04/14 重构预训练代码
WWM,结巴分词
基于 CLUECorpusSmall,数据处理教程 来源
数据集简介:可用于语言建模、预训练或生成型任务等,数据量超过 14G,近 4000 个定义良好的 txt 文件、50 亿字。主要部分来自于 nlp_chinese_corpus 项目
包含如下子语料库(总共 14G 语料):新闻语料news2016zh_corpus.zip, 社区互动语料webText2019zh_corpus.zip,维基百科语料wiki2019zh_corpus.zip,评论数据语料comment2019zh_corpus.zip。
(1) 数据集解压
unzip comment2019zh_corpus.zip -d /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/comment2019zh_corpus
unzip news2016zh_corpus.zip -d /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/news2016zh_corpus
unzip webText2019zh_corpus.zip -d /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/webText2019zh_corpus
unzip wiki2019zh_corpus.zip -d /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/wiki2019zh_corpus
(2) 将txt文件转换为jsonl格式
cd data
python trans_to_json.py --input_path /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data --output_path /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/small_data.jsonl
(3) 使用 rjieba 进行中文分词
会得到refids.txt
和reftext.txt
两个文件,并组合refids.txt
和reftext.txt
两个文件保存成huggingface
的dataset
,存放在 clue_small_wwm_data 文件夹下。
python run_chinese_ref.py --model_name /root/autodl-tmp/models/GAU-Base-Full --input_path /root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data/small_data.jsonl
TRAIN_DIR=/root/autodl-tmp/GAU-PyTorch/clue_small_wwm_data
BATCH_SIZE=32
ACCUMULATION=8
LR=2e-4
ACTIVATION_FUNCTION=roformerv2
python run_mlm_wwm.py \
--do_train \
--model_type roformerv2 \
--tokenizer_name /root/autodl-tmp/models/roformerv2_chinese_base \
--train_dir $TRAIN_DIR \
--output_dir /root/autodl-tmp/GAU-PyTorch/outputs/$ACTIVATION_FUNCTION/ \
--logging_dir /root/tf-logs/$ACTIVATION_FUNCTION/ \
--per_device_train_batch_size $BATCH_SIZE \
--gradient_accumulation_steps $ACCUMULATION \
--learning_rate $LR \
--weight_decay 0.01 \
--adam_epsilon 1e-6 \
--max_steps 30000 \
--warmup_steps 3000 \
--logging_steps 50 \
--save_steps 2000 \
--seed 1234 \
--max_grad_norm 1.0 \
--dataloader_num_workers 6 \
--fp16 \
--activation_function $ACTIVATION_FUNCTION \
--scaling_factor n \
--pad_to_max_length False \
根据 苏神的分析,采用了两种注意力矩阵的归一化策略:
$$
softmax_plus=softmax\left(\frac{\log_{512} n}{\sqrt{d}}QK^{\top}\right)V
$$
与 FLASH 原论文提出的:
$$
squared_relu=\frac{1}{n^2} relu^2\left(\frac{QK^{\top}}{\sqrt{d}}\right)V
$$
在预训练阶段:squared relu 出现了较为严重的波动,几次 checkpoint 也未能恢复正常。排查原因后发现是波动附近 batch 的 sql_len 与其他 batch 接近 512 的长度有较大差异。这进一步说明了 squared relu 在样本长度方面的迁移能力不好。而苏神的 softmax plus 则训练稳定,效果较好,原因与推导详见 blog
以下均是在 A40*1 取得的结果。
在 CLUE 分类数据集上对 RoFormerV1、RoFormerV2(多任务)、RoFormerV2(MLM,3W步)、GAU(3W步)、GAU(完整训练)(以上均为 Base 模型)进行对比:
模型 | AFQMC | CMNLI | CSL | IFLYTEK | TNews | WSC | Score |
---|---|---|---|---|---|---|---|
RoFormerV1 | 74.21 | 81.50 | 83.13 | 60.17 | 58.07 | 83.22 | 73.38 |
RoFormerV2 | 75.96 | 81.41 | 84.81 | 63.24 | 59.39 | 83.93 | 74.79 |
GAU(3W) | 73.14 | 76.73 | 80.44 | 59.81 | 54.66 | 75.9 | 70.19 |
RoFormerV2(3W) | 73.97 | 77.82 | 79.65 | 58.73 | 53.3 | 76.41 | 69.96 |
GAU(Full) | 74.51 | 81.97 | 83.7 | 62.72 | 57.93 | 82.89 | 73.95 |
可见 GAU 虽比拟不了 RoFormerV2 的结果,但超出 V1 约 0.6 个点,而且比对两个训练 3W 步的模型可知,GAU 的拟合能力是不逊于 RoFormerV2 的。如果再仔细设计一下预训练,增多语料很有可能会取得更好的结果。
以下是在 CLUE CMNLI 任务上取 bs=8 sql_len=512 的推理速度:
模型 | Time Cost | Percent |
---|---|---|
RoFormerV1 | 03:37:48 | 1.64× |
RoFormerV2 | 03:03:09 | 1.379× |
GAU | 02:12:48 | 1× |
以下是在 CLUE CMNLI 任务上取 bs=64 sql_len=512 的最大显存占用:
模型 | VRam Cost | Percent |
---|---|---|
RoFormerV1 | 18487M | 100% |
RoFormerV2 | 15807M | 85.5% |
GAU | 17453M | 94.4% |
以下是各模型的参数量对比:(XForSequenceClassification)
模型 | Parametrers | Percent |
---|---|---|
RoFormerV1 | 124,147,970 | 128.63% |
RoFormerV2 | 94,777,090 | 98.20% |
GAU | 96,519,170 | 100% |