1. 预训练

根据batch_size为512，最大移动窗口为5，噪声词数量为5，获取数据。

from d2l import mxnet as d2l
from mxnet import np, npx, gluon, init, autograd
from mxnet.gluon import nn
import pandas as pd
import os 
import math
import random
from plotly import express as px
npx.set_np()
batch_size, max_window_size, num_noise_words = 512, 5, 5
data_iter, vocab = load_data_ptb(batch_size, max_window_size, num_noise_words)

1.1 使用Skip-Gram 模型

通过使用嵌入层和最小批处理乘法来实现Skip-Gram模型。这些方法也经常用于其他的自言语言处理应用。

1.1.1 嵌入层

获取嵌入单词的层称为嵌入层，可以通过nn.Embedding在Gluon中创建实例来获取该层。嵌入层的权重是一个矩阵，其行数是字典的大小（input_dim）,列数为每个单词向量的维度（output_dim）,这里我们将字典大小设置为20，单词向量维度设置为4。

embed = nn.Embedding(input_dim=20, output_dim=4)
embed.initialize()
embed.weight

#  Parameter embedding0_weight (shape=(20, 4), dtype=float32)

嵌入层的输入为单词的索引。嵌入层对输入处理之后将词向量作为行返回。输入一个（2，3）形状到嵌入层，返回（2，3，4）形状，其中4是词向量维度。

x = np.array([[1, 2, 3], [4, 5, 6]])
embed(x).shape

# (2, 3, 4)

1.1.2 小批量乘法

我们可以使用batch_dot函数将两个minibatch中的矩阵相乘。假设第一个batch中包含 $n$ 个$a\times b$形状的矩阵： ${\mathbf{X}}_1,\dots ,{\mathbf{X}}_n$ 并且 第二个batch中包含 $n$ 个 形状为$b\times c$的矩阵：${\mathbf{Y}}_1,\dots ,{\mathbf{Y}}_n$ 。这两个批次通过矩阵乘法输出 $n$个形状为 $a\times c$的矩阵 ${\mathbf{X}}_1{\mathbf{Y}}_1,\dots ,{\mathbf{X}}_n{\mathbf{Y}}_n$。 因此, 给定两个张量 ($n$, $a$, $b$) 和 ($n$, $b$, $c$), minibatch乘法的输出为 ($n$, $a$, $c$)。说白了就是，两个批次中索引对应的一对进行矩阵相乘。

X = np.ones((3, 1, 4))
Y = np.ones((3, 4, 6))
npx.batch_dot(X, Y).shape

#  (3, 1, 6)

1.1.3 正向计算

模型的输入中包含中心目标词索引center以及串联的上下文和噪声词索引contexts_and_negatives。其中，center变量具有形状（批量大小， 1），contexts_and_negatives形状为（批量大小， 最大长度）这两个变量首先通过单词嵌入层从单词索引转换为单词向量，然后通过小批量乘法获得输出形状为（批量大小， 1， 最大长度）。

def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
    v = embed_v(center)
    u = embed_u(contexts_and_negatives)
    pred = npx.batch_dot(v, u.swapaxes(1, 2))
    return pred

skip_gram(np.ones((2, 1)), np.ones((2, 4)), embed, embed).shape

#  (2, 1, 4)

2. 训练

在训练之前定义损失函数。

2.1 二元交叉熵损失函数

根据负采样中损失函数的定义使用二元交叉熵损失函数。直接Gluon中的即可。

loss = gluon.loss.SigmoidBinaryCrossEntropyLoss()

由于为了补齐到max_len，部分是填补上的0，通过mask进行了标记，在计算损失的时候这部分不要计算，即只有mask中标记为1的参与损失函数的计算。

pred = np.array([[.5]*4]*2)
label = np.array([[1, 0, 1, 0]]*2)
mask = np.array([[1, 1, 1, 1], [1, 1, 0, 0]])
loss(pred, label, mask)

#  array([0.724077 , 0.3620385])

由于每个示例的长度不同，我们可以根据长度归一化对应损失

loss(pred, label, mask) / mask.sum(axis=1) * mask.shape[1]

#  array([0.724077, 0.724077])

2.2 初始化模型参数

分别构造中心词和上下文的嵌入层，并将超参数词向量embed_size设为100。

embed_size = 100
net = nn.Sequential()
net.add(nn.Embedding(input_dim=len(vocab), output_dim=embed_size), nn.Embedding(input_dim=len(vocab), output_dim=embed_size))

2.3 模型训练

训练功能定义为如下函数。由于存在填充，因此损失函数的计算跟之间的有点不一样。

def train(net, data_iter, lr, num_epochs, device=npx.gpu()):
    net.initialize(ctx=device, force_reinit=True)
    trainer = gluon.Trainer(net.collect_params(), 'adam',{'learning_rate': lr})
    epochs_lst, loss_lst = [], []
    for epoch in range(num_epochs):
        timer = d2l.Timer()
        # 记录 loss 以及 token的个数
        metric = d2l.Accumulator(2)  
        for i, batch in enumerate(data_iter):
            center, context_negative, mask, label = [
                data.as_in_ctx(device) for data in batch]
            with autograd.record():
                pred = skip_gram(center, context_negative, net[0], net[1])
                # 这里的loss根据长度做了归一化
                l = (loss(pred.reshape(label.shape), label, mask)
                     / mask.sum(axis=1) * mask.shape[1])
            l.backward()
            trainer.step(batch_size)
            metric.add(l.sum(), l.size)
            if (i+1) % 50 == 0:
                epochs_lst.append(epoch+(i+1)/len(data_iter))
                loss_lst.append(metric[0]/metric[1])
    print(f'loss {metric[0] / metric[1]:.3f}, ' f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')
    fig = px.line(pd.DataFrame([loss_lst], columns=epochs_lst, index=['loss']).T, width=600, height=360, labels={'index':'epoch', 'value':'loss'})
    fig.show()

使用负采样训练一个Skip-Gram模型：

lr, num_epochs = 0.01, 5
train(net, data_iter, lr, num_epochs)

#  loss 0.331, 27137.8 tokens/sec on gpu(0)

3. 词嵌入模型应用

训练单词嵌入模型后，我们可以基于两个单词向量的余弦相似度来表示单词之间的意义相似度。如我们所见，在使用经过训练的单词嵌入模型时，含义最接近“chip”的单词主要与chip相关。

def get_similar_tokens(query_token, k, embed):
    W = embed.weight.data()
    x = W[vocab[query_token]]
    # 计算余弦相似度。 加1e-9可获得数值稳定性,主要是怕除以0
    cos = np.dot(W, x) / np.sqrt(np.sum(W * W, axis=1) * np.sum(x * x) + 1e-9)
    topk = npx.topk(cos, k=k+1, ret_typ='indices').asnumpy().astype('int32')
    for i in topk[1:]:  #  将输入项去掉
        print(f'cosine sim={float(cos[i]):.3f}: {vocab.idx_to_token[i]}')

get_similar_tokens('chip', 3, net[0])

4. 参考

https://d2l.ai/chapter_natural-language-processing-pretraining/word2vec-pretraining.html

5. 代码

github