一个轻量级的 Albert 模型

一个轻量级的 Albert 模型

BERT-base 一上来就是 1 亿多参数，部署到线上要占不少显存，想在手机这类资源紧张的环境里跑更是奢望。其实参数多不一定是好事，很多参数学完之后利用率很低。ALBERT 是 Google 2019 年拿出的瘦身方案，网络层数和隐藏维度跟 BERT 保持一致，只动参数的分配方式，就把 base 版本压到约 1200 万参数，下游成绩却几乎不掉。这种少花钱多办事的思路，是我们理解大模型怎么压缩的一个好入口。

接下来，我们先讲 BERT 的参数到底浪费在哪，再拆开 ALBERT 的两刀，分解词嵌入矩阵和跨层参数共享，然后看它顺手换掉的两个预训练任务，最后给一段加载现成模型的代码。读这篇之前，你需要知道 BERT 的整体结构，明白词嵌入和 Transformer 编码器层各是干嘛的。

1. 瘦身动机

动手拆之前，我们先看 BERT 的参数都花在哪了。BERT-base 有 12 层 Transformer，隐藏维度 768，中文词表约 21128。粗看参数分布，会发现两个疑点。第一处是词嵌入矩阵，它要为词表里每个词学一个 768 维向量，参数量是 21128 乘 768，约 1623 万。可训练时每个批次只碰到词表中的一小撮词，这张巨表绝大多数参数一轮里根本不更新，摊薄下来利用率很低。第二处是 12 层编码器，BERT 每一层都有自己独立的注意力和前馈网络参数，层数越深参数越多，而论文里观察到，各层其实在重复学类似的东西。ALBERT 的两刀，就分别砍向这两处。

2. 分解嵌入

第一刀砍向词嵌入矩阵，思路是别让每个词直接背一个 768 维向量。ALBERT 把一步到位拆成两步。先让每个词只学一个 128 维的低维向量，词嵌入矩阵就从 21128 乘 768 缩成 21128 乘 128。低维向量进编码器之前，再过一个线性层，把 128 维映回 768 维，后面的网络结构一个字都不用改。为什么 128 就够了呢？论文做过对比实验，嵌入维度从小往大试，128 在跨层共享的配置下收益最好，再往上加提升微乎其微。直觉上也说得通，词向量只需要记录这个词本身是什么，不带上下文，低维就足够，而 768 维那种容量，是留给自注意力算完之后的上下文表示用的。

2.1 参数对比

我们把两边的嵌入部分算在一起对比，差别一目了然。

项目BERT-baseALBERT-base
词表大小2112821128
词向量维度768128
词嵌入矩阵21128×768，约 1623 万21128×128，约 270 万
升维投影无128×768，约 10 万
嵌入部分合计约 1623 万约 280 万

算一笔账，嵌入部分从约 1623 万降到约 280 万，少了约 83%，而论文里的实验显示，下游性能只有很小损失。

2.2 代码拆解

我们看 transformers 库里的实现差别。BERT 直接建一张大嵌入表，词向量维度就是隐藏维度：

class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, config):
        super().__init__()
        # config.vocab_size 21128
        # config.hidden_size 768
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        ......

ALBERT 的嵌入表小一号，用单独的 embedding_size 来控制：

class AlbertEmbeddings(nn.Module):
    """
    Construct the embeddings from word, position and token_type embeddings.
    """

    def __init__(self, config):
        super().__init__()
        # config.vocab_size 21128
        # config.embedding_size 128
        self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)
        ......

升维发生在编码器入口，一个 Linear 层把 128 维抬回 768 维：

self.embedding_hidden_mapping_in = nn.Linear(config.embedding_size, config.hidden_size)
hidden_states = self.embedding_hidden_mapping_in(hidden_states)

3. 参数共享

嵌入矩阵省下来的只是零头，第二刀才是大头，砍向 12 层编码器的重复参数。BERT 的 12 层各学各的，参数随层数线性涨。ALBERT 让 12 层共用同一套参数，注意力权重、前馈网络权重全部共享，第 1 层和第 12 层用的是同一组数字。训练时不管堆多少层，可训练参数只算一份。

论文试过几种共享粒度，只共享注意力、只共享前馈、全部共享，最后默认选了全共享，效果最稳。这里要补一句，层数还是 12 层，前向计算照样跑 12 次，省的是参数存储和更新量，不是推理时间。两刀合起来，ALBERT-base 的可训练参数压到约 1200 万，约为 BERT-base 的九分之一。

这里本该放一张跨层参数共享的示意图，对比 BERT 每层独立权重和 ALBERT 多层共用一套权重。

4. 任务改造

参数瘦下来之后，论文团队发现 BERT 的预训练任务本身也有水分，于是又动了两处手术，一处换句子关系任务，一处改掩码方式。

4.1 替换 NSP

先看 NSP。BERT 用它判断第二句是不是第一句的下一句。正样本取语料里连续的两句，标为 IsNext，负样本随机挑一句不相关的，标为 NotNext。ALBERT 认为这个任务太缺乏难度，两句话话题不一样模型就猜得出来，根本不需要真正理解段落衔接，对性能提升有限。

于是它换成 SOP，也就是句子顺序预测。正样本还是连续两句，让模型判断 B 是不是 A 的下一句。负样本直接把 A 和 B 的顺序对调，变成判断 A 是不是 B 的下一句。这样负样本和正样本话题完全一致，区别只在逻辑顺序，模型必须真正读懂句间关系才能答对。

4.2 掩码与丢弃

第二处手术在掩码上。BERT 随机遮单个字，被遮的字之间彼此独立。可语言是连贯的，要猜的往往是一个短语片段。ALBERT 按 n-gram 整块遮，最多遮连续 3 个词，每次从 1-gram、2-gram、3-gram 里挑一种，选中概率分别是 6/11、3/11、2/11。单字遮得最多，三字短语遮得最少，模型被迫学会根据上下文猜一整块连续的词，而不是孤立地填一个空。

还有个反直觉的细节。ALBERT 把 dropout 基本关了，默认配置里注意力 dropout 和隐藏层 dropout 都是 0。原因是作者发现模型训得更久也不过拟合，留着 dropout 反而白丢信息。transformers 库里大尺寸版本的默认配置打印出来是这样：

AlbertConfig {
  "attention_probs_dropout_prob": 0,
  "bos_token_id": 2,
  "classifier_dropout_prob": 0.1,
  "embedding_size": 128,
  "eos_token_id": 3,
  "hidden_act": "gelu_new",
  "hidden_dropout_prob": 0,
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "inner_group_num": 1,
  "intermediate_size": 16384,
  "layer_norm_eps": 1e-12,
  "max_position_embeddings": 512,
  "model_type": "albert",
  "num_attention_heads": 64,
  "num_hidden_groups": 1,
  "num_hidden_layers": 12,
  "pad_token_id": 0,
  "position_embedding_type": "absolute",
  "transformers_version": "4.22.2",
  "type_vocab_size": 2,
  "vocab_size": 30000
}

我们看三处，attention_probs_dropout_prob 和 hidden_dropout_prob 都是 0，embedding_size 是 128，这正是 ALBERT 的典型特征。这份配置对应的是大尺寸版本，hidden_size 拉到了 4096，base 版本则是 768。

5. 加载使用

机制讲完，我们说点落地的。ALBERT 预训练好的权重在 transformers 库里直接下载，微调接口和 BERT 一模一样。我们以文本分类为例，接一个二分类头：

from transformers import AlbertTokenizer, AlbertForSequenceClassification

tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
model = AlbertForSequenceClassification.from_pretrained('albert-base-v2', num_labels=2)

text = '这家餐厅的菜真好吃'
inputs = tokenizer(text, return_tensors='pt')
out = model(**inputs)                 # out.logits 是两个类别的得分
label = out.logits.argmax().item()    # 0 负面, 1 正面

如果是从本地加载，就指定本地模型目录，换成你自己的路径即可：

AlbertForSequenceClassification.from_pretrained('model/albert_chinese_tiny')

微调时学习率要调小，预训练权重已经学得很好，步子太大容易把知识冲掉，2e-5 到 5e-5 这个量级比较常用。

6. 误区澄清

第一个坑，以为共享参数会让 12 层退化成 1 层。前向时数据确实过了 12 遍同一套权重，但每层的输入来自上一层的输出，算出来的中间表示仍然逐层变化，并不是简单重复。

第二个坑，以为 128 维词向量表达能力不够。它后面紧跟一个升到 768 维的线性投影，参与上下文计算的依然是 768 维，被压掉的只是存储词表本身的参数量。

第三个坑，以为 ALBERT 推理更快。它省的是参数量，12 层的计算量一次没少，跑得快不快主要看层数和隐藏维度。

我们回到开头的问题。ALBERT 给 BERT 瘦了三刀。分解嵌入把词表矩阵拆成低维查表加线性投影，嵌入参数省了约 83%。跨层全共享让 12 层共用一套权重，把 base 版本压到约 1200 万参数，约为 BERT-base 的九分之一。SOP 和 n-gram 掩码换掉了过于简单的 NSP 和单字掩码，再关掉 dropout 补回丢掉的信息。它证明了 BERT 的参数里有相当一部分是冗余的。原论文见 https://arxiv.org/pdf/1909.11942.pdf

©2026 老孟讲编程 京ICP备2021009384号anbei冀公网安备13050302001966号