Word2vec 的实现

本节是对前两节内容的实践。我们以“词嵌入(word2vec)”一节中的跳字模型和“近似训练”一节中的负采样为例,介绍在语料库上训练词嵌入模型的实现。我们还会介绍一些实现中的技巧,例如二次采样(subsampling)和掩码(mask)变量。

首先让我们导入实验所需的包或模块。

In [1]:
import collections
import gluonbook as gb
import math
from mxnet import autograd, gluon, nd
from mxnet.gluon import data as gdata, loss as gloss, nn
import random
import sys
import time
import zipfile

处理数据集

Penn Tree Bank(PTB)是一个常用的小型语料库 [1]。它采样自华尔街日报的文章,包括训练集、验证集和测试集。我们将在 PTB 的训练集上训练词嵌入模型。该数据集的每一行作为一个句子。句子中的每个词由空格隔开。

In [2]:
with zipfile.ZipFile('../data/ptb.zip', 'r') as zin:
    zin.extractall('../data/')

with open('../data/ptb/ptb.train.txt', 'r') as f:
    lines = f.readlines()
    # st 是 sentence 在循环中的缩写。
    raw_dataset = [st.split() for st in lines]

'# sentences: %d' % len(raw_dataset)
Out[2]:
'# sentences: 42068'

对于数据集的前三个句子,打印每个句子的词数和前五个词。这个数据集中句尾符为“<eos>”,生僻词全用“<unk>”表示,数字则被替换成了“N”。

In [3]:
for st in raw_dataset[:3]:
    print('# tokens:', len(st), st[:5])
# tokens: 24 ['aer', 'banknote', 'berlitz', 'calloway', 'centrust']
# tokens: 15 ['pierre', '<unk>', 'N', 'years', 'old']
# tokens: 11 ['mr.', '<unk>', 'is', 'chairman', 'of']

建立词语索引

为了计算简单,我们只保留在数据集中至少出现 5 次的词。

In [4]:
# tk 是 token 的在循环中的缩写。
counter = collections.Counter([tk for st in raw_dataset for tk in st])
counter = dict(filter(lambda x: x[1] >= 5, counter.items()))

然后将词映射到整数索引。

In [5]:
idx_to_token = [tk for tk, _ in counter.items()]
token_to_idx = {tk: idx for idx, tk in enumerate(idx_to_token)}
dataset = [[token_to_idx[tk] for tk in st if tk in token_to_idx]
           for st in raw_dataset]
num_tokens = sum([len(st) for st in dataset])
'# tokens: %d' % num_tokens
Out[5]:
'# tokens: 887100'

二次采样

文本数据中一般会出现一些高频词,例如英文中的“the”、“a”和“in”。通常来说,在一个背景窗口中,一个词(如“chip”)和较低频词(如“microprocessor”)同时出现比和较高频词(如“the”)同时出现对训练词嵌入模型更有益。因此,训练词嵌入模型时可以对词进行二次采样 [2]。具体来说,数据集中每个被索引词 \(w_i\) 将有一定概率被丢弃,该丢弃概率为

\[\mathbb{P}(w_i) = \max\left(1 - \sqrt{\frac{t}{f(w_i)}}, 0\right),\]

其中 \(f(w_i)\) 是数据集中词 \(w_i\) 的个数与总词数之比,常数 \(t\) 是一个超参数(实验中设为 \(10^{-4}\))。可见,只有当 \(f(w_i) > t\) 时,我们才有可能在二次采样中丢弃词 \(w_i\),并且越高频的词被丢弃的概率越大。

In [6]:
def discard(idx):
    return random.uniform(0, 1) < 1 - math.sqrt(
        1e-4 / counter[idx_to_token[idx]] * num_tokens)

subsampled_dataset = [[tk for tk in st if not discard(tk)] for st in dataset]
'# tokens: %d' % sum([len(st) for st in subsampled_dataset])
Out[6]:
'# tokens: 375509'

可以看到,二次采样后我们去掉了一半左右的词。下面比较一个词在二次采样前后出现在数据集中的次数。可见高频词“the”的采样率不足 1/20。

In [7]:
def compare_counts(token):
    return '# %s: before=%d, after=%d' % (token, sum(
        [st.count(token_to_idx[token]) for st in dataset]), sum(
        [st.count(token_to_idx[token]) for st in subsampled_dataset]))

compare_counts('the')
Out[7]:
'# the: before=50770, after=2044'

但低频词“join”则完整地保留了下来。

In [8]:
compare_counts('join')
Out[8]:
'# join: before=45, after=45'

提取中心词和背景词

我们将与中心词距离不超过背景窗口大小的词作为它的背景词。下面定义函数提取出所有中心词和它们的背景词。它每次在整数 1 和max_window_size(最大背景窗口)之间均匀随机采样一个整数作为背景窗口大小。

In [9]:
def get_centers_and_contexts(dataset, max_window_size):
    centers, contexts = [], []
    for st in dataset:
        if len(st) < 2:  # 每个句子至少要有 2 个词才可能组成一对“中心词 - 背景词”。
            continue
        centers += st
        for center_i in range(len(st)):
            window_size = random.randint(1, max_window_size)
            indices = list(range(max(0, center_i - window_size),
                                 min(len(st), center_i + 1 + window_size)))
            indices.remove(center_i)  # 将中心词排除在背景词之外。
            contexts.append([st[idx] for idx in indices])
    return centers, contexts

下面我们创建一个人工数据集,其中含有词数分别为 7 和 3 的两个句子。设最大背景窗口为 2,打印所有中心词和它们的背景词。

In [10]:
tiny_dataset = [list(range(7)), list(range(7, 10))]
print('dataset', tiny_dataset)
for center, context in zip(*get_centers_and_contexts(tiny_dataset, 2)):
    print('center', center, 'has contexts', context)
dataset [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]
center 0 has contexts [1, 2]
center 1 has contexts [0, 2, 3]
center 2 has contexts [0, 1, 3, 4]
center 3 has contexts [1, 2, 4, 5]
center 4 has contexts [2, 3, 5, 6]
center 5 has contexts [3, 4, 6]
center 6 has contexts [5]
center 7 has contexts [8]
center 8 has contexts [7, 9]
center 9 has contexts [7, 8]

实验中,我们设最大背景窗口大小为 5。下面提取数据集中所有的中心词及其背景词。

In [11]:
all_centers, all_contexts = get_centers_and_contexts(subsampled_dataset, 5)

负采样

我们使用负采样来进行近似训练。对于一对中心词和背景词,我们随机采样 \(K\) 个噪音词(实验中设 \(K=5\))。根据 word2vec 论文的建议,噪音词采样概率 \(\mathbb{P}(w)\) 设为 \(w\) 词频与总词频之比的 0.75 次方 [2]。

In [12]:
def get_negatives(all_contexts, sampling_weights, K):
    all_negatives, neg_candidates, i = [], [], 0
    population = list(range(len(sampling_weights)))
    for contexts in all_contexts:
        negatives = []
        while len(negatives) < len(contexts) * K:
            if i == len(neg_candidates):
                # 根据每个词的权重(sampling_weights)随机生成 k 个词的索引作为噪音
                # 词。为了高效计算,可以将 k 设的稍大一点。
                i, neg_candidates = 0, random.choices(
                    population, sampling_weights, k=int(1e5))
            neg, i = neg_candidates[i], i + 1
            # 噪音词不能是背景词。
            if neg not in set(contexts):
                negatives.append(neg)
        all_negatives.append(negatives)
    return all_negatives

sampling_weights = [counter[w]**0.75 for w in idx_to_token]
all_negatives = get_negatives(all_contexts, sampling_weights, 5)

读取数据

我们从数据集中提取所有中心词all_centers,以及每个中心词对应的背景词all_contexts和噪音词all_negatives。我们将通过随机小批量来读取它们。

在一个小批量数据中,第 \(i\) 个样本包括一个中心词以及它所对应的 \(n_i\) 个背景词和 \(m_i\) 个噪音词。由于每个样本的背景窗口大小可能不一样,其中背景词与噪音词个数之和 \(n_i+m_i\) 也会不同。在构造小批量时,我们将每个样本的背景词和噪音词连结在一起,并添加填充项 0 直至连结后的长度相同,即长度均为 \(\max_i n_i+m_i\)max_len)。为了避免填充项对损失函数计算的影响,我们构造了掩码变量masks,其每一个元素分别与连结后的背景词和噪音词contexts_negatives中的元素一一对应。当变量contexts_negatives中的某个元素为填充项时,相同位置的掩码变量masks中的元素取 0,否则取 1。为了区分正例和负例,我们还需要将contexts_negatives变量中的背景词和噪音词区分开来。依据掩码变量的构造思路,我们只需创建与contexts_negatives变量形状相同的标签变量labels,并将与背景词(正例)对应的元素设 1,其余清 0。

下面我们将实现这个小批量读取函数batchify。它的小批量输入data是一个长度为批量大小的列表,其中每个元素分别包含中心词center、背景词context和噪音词negative。该函数返回的小批量数据符合我们所需要的格式,例如包含了掩码变量。

In [13]:
def batchify(data):
    max_len = max(len(c) + len(n) for _, c, n in data)
    centers, contexts_negatives, masks, labels = [], [], [], []
    for center, context, negative in data:
        cur_len = len(context) + len(negative)
        centers += [center]
        contexts_negatives += [context + negative + [0] * (max_len - cur_len)]
        masks += [[1] * cur_len + [0] * (max_len - cur_len)]
        labels += [[1] * len(context) + [0] * (max_len - len(context))]
    return (nd.array(centers).reshape((-1, 1)), nd.array(contexts_negatives),
            nd.array(masks), nd.array(labels))

我们用刚刚定义的batchify函数指定DataLoader实例中小批量的读取方式。然后打印读取的第一个批量中各个变量的形状。

In [14]:
batch_size = 512
num_workers = 0 if sys.platform.startswith('win32') else 4
dataset = gdata.ArrayDataset(all_centers, all_contexts, all_negatives)
data_iter = gdata.DataLoader(dataset, batch_size, shuffle=True,
                             batchify_fn=batchify, num_workers=num_workers)
for batch in data_iter:
    for name, data in zip(['centers', 'contexts_negatives', 'masks',
                           'labels'], batch):
        print(name, 'shape:', data.shape)
    break
centers shape: (512, 1)
contexts_negatives shape: (512, 60)
masks shape: (512, 60)
labels shape: (512, 60)

跳字模型

我们将通过使用嵌入层和小批量乘法来实现跳字模型。它们也常常用于实现其他自然语言处理的应用。

嵌入层

获取词嵌入的层称为嵌入层,在 Gluon 中可以通过创建nn.Embedding实例得到。嵌入层的权重是一个矩阵,其行数为词典大小(input_dim),列数为每个词向量的维度(output_dim)。我们设词典大小为 20,词向量的维度为 4。

In [15]:
embed = nn.Embedding(input_dim=20, output_dim=4)
embed.initialize()
embed.weight
Out[15]:
Parameter embedding0_weight (shape=(20, 4), dtype=float32)

嵌入层的输入为词的索引。输入一个词的索引 \(i\),嵌入层返回权重矩阵的第 \(i\) 行作为它的词向量。下面我们将形状为(2,3)的索引输入进嵌入层,由于词向量的维度为 4,我们得到形状为(2,3,4)的词向量。

In [16]:
x = nd.array([[1, 2, 3], [4, 5, 6]])
embed(x)
Out[16]:

[[[ 0.01438687  0.05011239  0.00628365  0.04861524]
  [-0.01068833  0.01729892  0.02042518 -0.01618656]
  [-0.00873779 -0.02834515  0.05484822 -0.06206018]]

 [[ 0.06491279 -0.03182812 -0.01631819 -0.00312688]
  [ 0.0408415   0.04370362  0.00404529 -0.0028032 ]
  [ 0.00952624 -0.01501013  0.05958354  0.04705103]]]
<NDArray 2x3x4 @cpu(0)>

小批量乘法

我们可以使用小批量乘法运算batch_dot对两个小批量中的矩阵一一做乘法。假设第一个批量中包含 \(n\) 个形状为 \(a\times b\) 的矩阵 \(\boldsymbol{X}_1, \ldots, \boldsymbol{X}_n\),第二个批量中包含 \(n\) 个形状为 \(b\times c\) 的矩阵 \(\boldsymbol{Y}_1, \ldots, \boldsymbol{Y}_n\)。这两个批量的矩阵乘法输出为 \(n\) 个形状为 \(a\times c\) 的矩阵 \(\boldsymbol{X}_1\boldsymbol{Y}_1, \ldots, \boldsymbol{X}_n\boldsymbol{Y}_n\)。因此,给定两个形状分别为(\(n\)\(a\)\(b\))和(\(n\)\(b\)\(c\))的 NDArray,小批量乘法输出的形状为(\(n\)\(a\)\(c\))。

In [17]:
X = nd.ones((2, 1, 4))
Y = nd.ones((2, 4, 6))
nd.batch_dot(X, Y).shape
Out[17]:
(2, 1, 6)

跳字模型前向计算

在前向计算中,跳字模型的输入包含中心词索引center以及连结的背景词与噪音词索引contexts_and_negatives。其中center变量的形状为(批量大小,1),而contexts_and_negatives变量的形状为(批量大小,max_len)。这两个变量先通过词嵌入层分别由词索引变换为词向量,再通过小批量乘法得到形状为(批量大小,1,max_len)的输出。输出中的每个元素是中心词向量与背景词向量或噪音词向量的内积。

In [18]:
def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
    v = embed_v(center)
    u = embed_u(contexts_and_negatives)
    pred = nd.batch_dot(v, u.swapaxes(1, 2))
    return pred

训练模型

在训练词嵌入模型之前,我们需要定义模型的损失函数。

二元交叉熵损失函数

根据负采样中损失函数的定义,我们可以直接使用 Gluon 的二元交叉熵损失函数SigmoidBinaryCrossEntropyLoss

In [19]:
loss = gloss.SigmoidBinaryCrossEntropyLoss()

值得一提的是,我们可以通过掩码变量指定小批量中参与损失函数计算的部分预测值和标签:当掩码为 1 时,相应位置的预测值和标签将参与损失函数的计算;当掩码为 0 时,相应位置的预测值和标签则不参与损失函数的计算。我们之前提到,掩码变量可用于避免填充项对损失函数计算的影响。

In [20]:
pred = nd.array([[1.5, 0.3, -1, 2], [1.1, -0.6, 2.2, 0.4]])
# 标签变量 label 中的 1 和 0 分别代表背景词和噪声词。
label = nd.array([[1, 0, 0, 0], [1, 1, 0, 0]])
mask = nd.array([[1, 1, 1, 1], [1, 1, 1, 0]])  # 掩码变量。
loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)
Out[20]:

[0.8739896 1.2099689]
<NDArray 2 @cpu(0)>

作为比较,下面将从零开始实现二元交叉熵损失函数的计算,并根据掩码变量mask计算掩码为 1 的预测值和标签的损失。

In [21]:
def sigmd(x):
    return -math.log(1 / (1 + math.exp(-x)))

print('%.7f' % ((sigmd(1.5) + sigmd(-0.3) + sigmd(1) + sigmd(-2)) / 4))
print('%.7f' % ((sigmd(1.1) + sigmd(-0.6) + sigmd(-2.2)) / 3))
0.8739896
1.2099689

初始化模型参数

我们分别构造中心词和背景词的嵌入层,并将超参数词向量维度embed_size设置成 100。

In [22]:
embed_size = 100
net = nn.Sequential()
net.add(nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size),
        nn.Embedding(input_dim=len(idx_to_token), output_dim=embed_size))

训练

下面定义训练函数。由于填充项的存在,跟之前的训练函数相比,损失函数的计算稍有不同。

In [23]:
def train(net, lr, num_epochs):
    ctx = gb.try_gpu()
    net.initialize(ctx=ctx, force_reinit=True)
    trainer = gluon.Trainer(net.collect_params(), 'adam',
                            {'learning_rate': lr})
    for epoch in range(num_epochs):
        start_time, train_l_sum = time.time(), 0
        for batch in data_iter:
            center, context_negative, mask, label = [
                data.as_in_context(ctx) for data in batch]
            with autograd.record():
                pred = skip_gram(center, context_negative, net[0], net[1])
                # 使用掩码变量 mask 来避免填充项对损失函数计算的影响。
                l = (loss(pred.reshape(label.shape), label, mask) *
                     mask.shape[1] / mask.sum(axis=1))
            l.backward()
            trainer.step(batch_size)
            train_l_sum += l.mean().asscalar()
        print('epoch %d, train loss %.2f, time %.2fs'
              % (epoch + 1, train_l_sum / len(data_iter),
                 time.time() - start_time))

现在我们可以训练使用负采样的跳字模型了。

In [24]:
train(net, 0.005, 8)
epoch 1, train loss 0.46, time 5.23s
epoch 2, train loss 0.39, time 4.72s
epoch 3, train loss 0.35, time 4.67s
epoch 4, train loss 0.32, time 4.72s
epoch 5, train loss 0.30, time 4.72s
epoch 6, train loss 0.29, time 4.83s
epoch 7, train loss 0.29, time 4.68s
epoch 8, train loss 0.28, time 4.72s

应用词嵌入模型

当训练好词嵌入模型后,我们可以根据两个词向量的余弦相似度表示词与词之间在语义上的相似度。可以看到,使用训练得到的词嵌入模型时,与词“chip”语义最接近的词大多与芯片有关。

In [25]:
def get_similar_tokens(query_token, k, embed):
    W = embed.weight.data()
    x = W[token_to_idx[query_token]]
    cos = nd.dot(W, x) / nd.sum(W * W, axis=1).sqrt() / nd.sum(x * x).sqrt()
    topk = nd.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32')
    for i in topk[1:]:  # 除去输入词。
        print('cosine sim=%.3f: %s' % (cos[i].asscalar(), (idx_to_token[i])))

get_similar_tokens('chip', 3, net[0])
cosine sim=0.475: microprocessor
cosine sim=0.461: intel
cosine sim=0.458: toshiba

小结

  • 我们可以使用 Gluon 通过负采样训练跳字模型。
  • 二次采样试图尽可能减轻高频词对训练词嵌入模型的影响。
  • 我们可以将长度不同的样本填充至长度相同的小批量,并通过掩码变量区分非填充和填充,然后只令非填充参与损失函数的计算。

练习

  • 我们用batchify函数指定DataLoader实例中小批量的读取方式,并打印了读取的第一个批量中各个变量的形状。这些形状该如何计算得到?
  • 试着找出其他词的近义词。
  • 调一调超参数,观察并分析实验结果。
  • 当数据集较大时,我们通常在迭代模型参数时才对当前小批量里的中心词采样背景词和噪音词。也就是说,同一个中心词在不同的迭代周期可能会有不同的背景词或噪音词。这样训练有哪些好处?尝试实现该训练方法。

扫码直达讨论区

参考文献

[1] Penn Tree Bank. https://catalog.ldc.upenn.edu/ldc99t42

[2] Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., & Dean, J. (2013). Distributed representations of words and phrases and their compositionality. In Advances in neural information processing systems (pp. 3111-3119).