轻量级 BERT 模型 ALBERT

BERT 有很多改进版本,ALBERT 是一个轻量化版本。ALBERT 源自 2020 年的发表论文《ALBERT: A LITE BERT FOR SELF-SUPERVISED LEARNING OF LANGUAGE REPRESENTATIONS》,论文地址:https://arxiv.org/pdf/1909.11942.pdf。从题目可以看出,论文重点是轻量化 BERT 模型,以及优化了半监督学习,本文主要关注轻量化模型部分。

目前 NLP 的高级模型 BERT、GPT-2 都基于 Pretrain/fine-tune 模式,先使用无监督学习的海量文本预训练出一个带有“语言常识”的大模型,然后再根据具体任务调优,这就完美解决了具体任务训练集不足的问题。模型往往拥有千万或亿级的参数,目前的趋势是随着模型功能能越来越强大,模型的规模也越来越大,这使普通开发者越发无法企及。

随着 BERT 模型越来越频繁地被使用到真实场景中,模型的速度、规模、硬件逐渐成为瓶颈。预训练模型一般都规模庞大,在后期调优以及调用模型时也都需要花费大量的内存、算力和时间。ALBERT 在基本保证模型质量的前提下把模型缩小到 BERT-LARGE 的 1/18(18M/334M),训练提速 1.7 倍。

原理

ALBERT 全称是 A Lite BERT,轻量级的 BERT,它使用了两种技术来简化参数:

分解了词嵌入(Embedding)参数

从原理来看,词嵌入表示的是词含义与上下文无关,而隐藏层表示的是词之间的组合含义,句子越长它的信息量越丰富,后者更加复杂,也需要更多参数。如果词表 V 非常大(词多),需要 V×E 空间存储参数,如果保持词嵌入层大小 E 与隐藏词 H 的大小一致,当 H 变大时,VxE 也跟着迅速变大。因此,论文作者将原来词嵌入层大小恒等于隐藏层大小:E≡H,改为:H>>E。这样就减少了参数,也可以支持未来使用更大的词表,比如对中文来说,词为单位比以字为单位效果更佳,但词量又比较大。当前 BERT、GPT-2 模型的隐藏层大小动辄 512、768 维,而腾讯词向量 200 维就可以很好地描述 800 多万词条的词性。因此,论文作者将词过嵌入过程从原来的 V->H 分解成:先把词 V 映射到 E(V->E),再由 E 映射到 H(E->H),使其复杂度从 O(V×H) 降到了 O(V×E+E×H),当 E 比 H 小很多时其优势尤为明显。

共享了交叉层(cross-layer)参数

ALBERT 与 BERT 原理一样,都沿用了 Transformer 模型的方案,堆叠了多个 Encoder 层,BERT 标准版 12 层,BERT-LARGE 版 24 层来说。ALBERT 共享了所有层的参数,使参数数据大幅减少。论文作者对比了 LARGE 模型 24 层,每一层的输入与输出的 L2 距离以及 cosine 相似性,如下图所示。

可以看到,共享了参数之后 ALBERT 比 BERT 表现更加平滑稳定,笔者个人认为:多层共享参数方法有点像 RNN,BERT 中每一层(layer)提取出不同层次的关系(有的在词法层,有的在词义层),由于参数足够多,存储在参数矩阵中的有效信息可能是稀疏的,因此将多层参数存储在一起,并没有产生严重的互相影响。

除了简化参数,ALBERT 还在 Next-sentence prediction (NSP) 方面对 BERT 做了一些训练方法上的优化,不是本篇重点,有兴趣的读者请直接查看论文。效果

下面列出了 ALBERT 与 BERT 的对比效果,最后两列分别是二者在多者任务中的平均分和速度对比。可以看到,其模型参数显著减少,而计算提度相对没有那么明显,这是由于模型由于共享了参数,而层数和对应层的计算量没变,只减小了 Embedding 的计算量。也可以从表中看到,简化后模型效果相对于 BERT 变化不大。文中还针对不同的共享方法,以及不同 Embedding 大小提供了大量实验数据,详见论文。

代码

TensorFlow 实现(官方源码):

https://github.com/google-research/ALBERT

Pytorch 实现(2.3.0 以上版本支持 ALBERT):

https://github.com/huggingface/transformers

中文版 TensorFlow:

https://github.com/brightmart/albert_zh

中文模型:

https://storage.googleapis.com/albert_models/albert_base_zh.tar.gz

ALBERT 用法与 BERT 基本一致。