harry's blog

1. 为什么要添加新 token

在使用预训练模型时，我们有时需要使用一些自定义 token 来增强输入，例如使用 <IMG> 和 </IMG> 在文本中标记出实体。由于自定义 token 并不在预训练模型原来的词表中，因此直接运用分词器 (Tokenizer) 处理输入就会出现问题

例如直接使用 BERT 分词器处理下面的句子：

>>> from transformers import AutoTokenizer, AutoModel
>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
>>> sentence = "<IMG>image</IMG> Describe this image."
>>> print(tokenizer(sentence).tokens())
['[CLS]', '<', 'I', '##MG', '>', 'image', '<', '/', 'I', '##MG', '>', 'Des', '##cribe', 'this', 'image', '.', '[SEP]']

由于分词器无法识别 <IMG> 和 </IMG>，将它们都当作未知字符处理，例如 <IMG> 被切分成 '<' 'I' '##MG' '>' 一共 4 个 token。此外，有时我们还会遇到一些领域相关词汇，例如医学领域的文本通常会包含大量的医学术语，它们可能并不在模型的词表中（例如一些术语是使用多个词语的缩写拼接而成），这时也会出现上面的问题。此时我们就需要将这些新 token 添加到模型的词表中，让分词器与模型可以识别并处理这些 token。

2. 如何添加新 token

Huggingface 的 Transformers 库提供了两种方式来添加新 token

add_tokens()
添加普通 token：添加新 token 列表，如果 token 不在词表中，就会被添加到词表的最后

>>> num_added_toks = tokenizer.add_tokens(["new_tok1", "new_tok2"])
>>> print("We have added", num_added_toks, "tokens")
We have added 2 tokens

此时我们在往里再次添加这两个新词，就会返回 0，说明 token 已经被添加过了

>>> new_tokens = ["new_tok1", "new_tok2"]
>>> new_tokens = set(new_tokens) - set(tokenizer.vocab.keys())
>>> tokenizer.add_tokens(list(new_tokens))
0

add_special_tokens()
添加特殊 token：添加包含特殊 token 的字典，键值从 bos_token eos_token unk_token sep_token pad_token cls_token mask_token additional_special_tokens 中选择，如果 token 不在词表中，就会被添加到词表的最后。添加后，可以通过特殊属性来访问这些 token，例如 tokenizer.cls_token 就指向 cls_token

>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
>>> special_tokens_dict = {"cls_token": "[NEW_CLS]"}
>>> num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
>>> print("We have added", num_added_toks, "tokens")
We have added 1 tokens
>>> assert tokenizer.cls_token == "[NEW_CLS]"

也可以使用 add_tokens() 添加特殊 token，只需要额外设置参数 special_tokens=True

>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
>>> num_added_toks = tokenizer.add_tokens(["[NEW_tok1]", "[NEW_tok2]"])
>>> num_added_toks = tokenizer.add_tokens(["[NEW_tok3]", "[NEW_tok4]"], special_tokens=True)
>>> print("We have added", num_added_toks, "tokens")
We have added 2 tokens
>>> print(tokenizer('[NEW_tok1] Hello [NEW_tok2] [NEW_tok3] World [NEW_tok4]!').tokens())
['[CLS]', '[new_tok1]', 'hello', '[new_tok2]', '[NEW_tok3]', 'world', '[NEW_tok4]', '!', '[SEP]']

特殊 token 的标准化 (normalization) 过程与普通 token 有一些不同，比如不会被小写。这里我们使用的是不区分大小写的 BERT 模型，因此分词后添加的普通 token [NEW_tok1] 和 [NEW_tok2] 都被处理为了小写，而特殊 token [NEW_tok3] 和 [NEW_tok4] 则维持大写，与 [CLS] 等自带特殊 token 保持一致

3. 调整 embedding 矩阵

无论使用哪种方式向词表中添加新 token 后，都需要重置模型 token embedding 矩阵的大小，也就是向矩阵中添加新 token 对应的 embedding

该操作通过调用预训练模型的 resize_token_embeddings() 函数来实现，例如对于上面的例子：

>>> from transformers import AutoTokenizer, AutoModel
>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
>>> model = AutoModel.from_pretrained("bert-base-uncased")
>>> len(tokenizer)
30522
>>> tokenizer.add_tokens(['<IMG>', '</IMG>'], special_tokens=True)
2
>>> len(tokenizer)
30524
>>> model.resize_token_embeddings(len(tokenizer))
Embedding(30524, 768)
>>> model.embeddings.word_embeddings.weight.size()
torch.Size([30524, 768])
>>> model.embeddings.word_embeddings.weight[-2:, :]
tensor([[ 0.0151, -0.0102,  0.0080,  ..., -0.0353,  0.0043, -0.0220],
        [-0.0016, -0.0073, -0.0046,  ...,  0.0048,  0.0065, -0.0096]],
       grad_fn=<SliceBackward0>)

可以看到，在添加了特殊 token <IMG> 和 </IMG> 之后，分词器的词表大小从 30522 增加到了 30524，并且模型的 token embedding 矩阵大小也成功调整为了$30524\times 768$

我们打印出新添加 token 对应的 embedding。因为新 token 会添加在词表的末尾，因此只需打印出矩阵最后两行。这里的 <IMG> 和 </IMG> 的 embedding 是不同的。这是因为在默认情况下，这些新 token 的 embedding 是随机初始化的。

4. token embedding 初始化

如果有充分的训练语料对模型进行微调或者继续预训练，那么将新添加 token 初始化为随机向量没什么问题。但是如果训练语料较少，这种做法就可能存在问题。研究表明，在训练数据不够多的情况下，这些新添加 token 的 embedding 只会在初始值附近小幅波动。换句话说，即使经过训练，它们的值事实上还是随机的。因此，在很多情况下，我们需要手工初始化这些新 token 的 embedding。对于 Transformers 库来说，可以通过直接对 embedding 矩阵赋值来实现。

初始化为全零向量

>>> with torch.no_grad():
...     model.embeddings.word_embeddings.weight[-2:, :] = torch.zeros([2, model.config.hidden_size], requires_grad=True)
...
>>> model.embeddings.word_embeddings.weight[-2:, :]
tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]], grad_fn=<SliceBackward0>)

注意，初始化 embedding 的过程并不可导，因此这里通过 torch.no_grad() 暂停梯度的计算

初始化为已有 token 的值
比较常见的操作是根据新添 token 的语义，将其值初始化为训练好 token 的 embedding。例如对于上面的例子，我们可以将 <IMG> 和 </IMG> 的值都初始化为 image 对应的 embedding。因为 token id 就是 token 在矩阵中的索引，因此我们可以直接通过 weight[token_id] 取出 image 对应的 embedding。

>>> from transformers import AutoTokenizer, AutoModel
>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
>>> model = AutoModel.from_pretrained("bert-base-cased")
>>> num_added_tokens = tokenizer.add_tokens(["<IMG>", "</IMG>"])
>>> token_id = tokenizer.convert_tokens_to_ids('image')
>>> token_id
3077
>>> token_embedding = model.embeddings.word_embeddings.weight[token_id]
>>> with torch.no_grad():
...     for i in range(1, num_added_tokens+1):
...         model.embeddings.word_embeddings.weight[-i:, :] = token_embedding.clone().detach().requires_grad_(True)
...
>>> print(model.embeddings.word_embeddings.weight[-2:, :])
tensor([[ 0.0811,  0.0543, -0.0313,  ...,  0.0036,  0.0189, -0.0743],
        [ 0.0811,  0.0543, -0.0313,  ...,  0.0036,  0.0189, -0.0743]],
       grad_fn=<SliceBackward0>)

可以看到最终结果符合我们的预期，<IMG> 和 </IMG> 被初始化为相同的 embedding。

根据新 token 的语义初始化
另一种常见的做法是根据新 token 的语义，使用对应的描述文本来完成初始化。例如将值初始化为描述文本中所有 token 的平均值，假设新 token $t_{i}$的描述文本为$w_{i,1},w_{i,2},…,w_{i,n}$，那么$t_{i}$的初始化 embedding 为：
$$
E_{t_{i}} = \frac{1}{n} \sum_{j=1}^{n} E(w_{i,j})
$$
$E$表示预训练模型的 token embedding 矩阵。对于上面的例子，我们可以分别为 <IMG> 和 </IMG> 编写对应的描述，然后再对它们的值进行初始化：

>>> descriptions = ['start of image', 'end of image']
>>> with torch.no_grad():
...     for i, token in enumerate(reversed(descriptions), start=1):
...         tokenized = tokenizer.tokenize(token)
...         print(tokenized)
...         tokenized_ids = tokenizer.convert_tokens_to_ids(tokenized)
...         new_embedding = model.embeddings.word_embeddings.weight[tokenized_ids].mean(axis=0)
...         model.embeddings.word_embeddings.weight[-i, :] = new_embedding.clone().detach().requires_grad_(True)
...
['end', 'of', 'image']
['start', 'of', 'image']
>>> model.embeddings.word_embeddings.weight[-2:, :]
tensor([[ 0.0303,  0.0285, -0.0171,  ..., -0.0325,  0.0048, -0.0147],
        [ 0.0039,  0.0290, -0.0326,  ..., -0.0260,  0.0074, -0.0047]],
       grad_fn=<SliceBackward0>)

可以看到，这里成功地将 <IMG> 初始化为 start of image embedding 的平均值，将 </IMG> 初始化为 end of image embedding 的平均值。

本文由 Yonghui Wang 创作，采用知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Dec 19, 2024 12:18 pm

Huggingface 向词表中添加 token

1. 为什么要添加新 token

2. 如何添加新 token

3. 调整 embedding 矩阵

4. token embedding 初始化