从零实现深度学习框架——深入浅出Word2vec(下)

引言

本着“凡我不能创造的，我就不能理解”的思想，本系列文章会基于纯Python以及NumPy从零创建自己的深度学习框架，该框架类似PyTorch能实现自动求导。

要深入理解深度学习，从零开始创建的经验非常重要，从自己可以理解的角度出发，尽量不使用外部完备的框架前提下，实现我们想要的模型。本系列文章的宗旨就是通过这样的过程，让大家切实掌握深度学习底层实现，而不是仅做一个调包侠。

前面介绍的CBOW和Skip-gram模型有一个重大的缺点，就是计算量太大了。主要是在最终的多分类问题上，我们经过了一个Softmax操作，想象一下百万级的词汇量，那么Softmax需要计算百万次。

针对这个问题有两种优化方法，分别是层次Softmax和负采样。本文我们主要介绍带负采样的Skip-gram。

负采样

SGNS(Skip-Gram with Negative-Sampling)，即带负采样的Skip-gram。

它将多分类任务简化为二分类任务，即不预测每个单词附近会出现某个单词，而是判断某个单词是否会在$w$附近出现。同样，训练完之后，我们需要的是学习到的权重。

二分类任务其实就是一个逻辑回归分类器，它的训练过程如下：

1. 将目标词和一个上下文单词组成正例
2. 随机采样词典中的其他单词与目标词组成负例
3. 训练逻辑回归分类器去区分正例和负例
4. 使用学到的权重作为嵌入

逻辑回归分类器

假设窗口大小为$2$，目标词为$w_t$，对于句子 $\cdots w_{t-2}{w}_{t-1}\underline{w_t}{w}_{t+1}{w}_{t+2}\cdots$。

它的上下文单词为$c\in \{w_{t-2},w_{t-1},w_{t+1},w_{t+2}\}$，假设$c$代表其中任意上下文单词，和目标词组成元组$(w_t,c)$。那么分类器输出$c$是$w_t$上下文单词的概率：
$$P(+|w_t,c)\text{\hspace{1em}(1)}$$
为了让它是一个概率，那么如果$c$不是上下文的概率就可以用$1$减去上式得到，以保证这两个事件概率之和为$1$：
$$P(-|w_t,c)=1-P(+|w_t,c)\text{\hspace{1em}(2)}$$
现在问题是我们如何计算这个概率呢？可能你已经看出来了，对，就是通过Sigmoid函数。具体做法为，还是计算$w_t$和$c$这两个词嵌入向量的点积得到一个(相似度)得分，然后传入Sigmoid函数，得到一个概率：
$$P(+|w_t,c)=\sigma (w_t\cdot c)=\frac{1}{1+\exp (-w_t\cdot c)}\text{\hspace{1em}(3)}$$
同时我们要满足$(11)$，即$c$不是上下文单词的概率为：
$$\begin{array}{rl}P(-|w_t,c)& =1-P(+|w_t,c)\\ & =\sigma (-w_t\cdot c)=\frac{1}{1+\exp (w\cdot c)}\end{array}\text{\hspace{1em}(4)}$$
其中$1-\frac{1}{1+\exp (-x)}=\frac{1}{1+\exp (x)}$很好证明，这里就不展开了。

这样我们得到了其中一个上下文单词的概率，但是窗口内包含很多个($L=2k$)上下文单词。Skip-gram简化为所有上下文单词都是独立的假设，我们只需要让它们的概率相乘：
$$P(+|w_t,c_{1:L})=\prod _{i=1}^L\sigma (w_t\cdot c_i)\text{\hspace{1em}(5)}$$
我们使用取对数的基操，变成连加，防止数值溢出：
$$\log P(+|w_t,c_{1:L})=\sum _{i=1}^L\log \sigma (w_t\cdot c_i)\text{\hspace{1em}(6)}$$
其中$w_t$和$c_i$都表示词嵌入向量，计算方法在之前的Skip-gram模型中有介绍。

我们的模型定义好了，接下来看如何训练。

训练

如果我们想普通的Skip-gram模型一样，光有正例是不够的，那你的模型直接输出$1$就好了。因此，我们需要负例，这也是负采样的由来。

我们需要让模型为正例尽可能输出$1$，为负例尽可能输出$0$。

我们考虑一个简单的例子：

... I love natural language processing ...

这里假设窗口大小$k=2$，有一个目标词natural和$4$个上下文单词，我们可以得到$4$个正例：

$w$	$c_{pos}$
natural	I
natural	love
natural	language
natural	processing

这里$w$表示目标词，$c_{pos}$表示真正的上下文单词，$(w,c_{pos})$组成正例。上面说到，我们也需要负例。实际上SGNS使用了比正例数量更多的负例(有参数$K$控制)。对于上面的每个正例，我们创建$K$个负例，每个包含目标词和一个随机噪声单词。

噪声单词从词典中随机采样，但不能是上下文单词。这里的采样有一定的技巧。

使用加权unigram频率$P_{\alpha }(w)$采样，其中$\alpha$是一个权重。

那为什么需要加权呢？

我们对比下未加权的方法和加权的unigram。假设我们根据未加权频率$P(w)$进行采样，假设一个很罕见的单词aardvark，其概率$P(aardvark)=0.01$。为了看到效果，夸张一点，假设另一个单词the出现的概率$P(the)=0.99$。

未加权说的是，我们只有$1\%$的概率抽取到单词aardvark。

再看下加权的情况，一般令$\alpha =0.75$。那么有：
$$P_{\alpha }(w)=\frac{\text{f}(w{)}^{\alpha }}{{\sum }_{w^{\prime }}\text{f}(w^{\prime }{)}^{\alpha }}\text{\hspace{1em}(7)}$$
我们计算：
$$\begin{array}{r}{P}_{\alpha }(the)=\frac{0.99^{0.75}}{0.99^{0.75}+0.1^{0.75}}=0.97\\ P_{\alpha }(aardvark)=\frac{0.01^{0.75}}{0.99^{0.75}+0.1^{0.75}}=0.03\end{array}$$
这样，有更高的概率采样罕见单词。

假设我们令$K=2$，即对于每个正例，我们采样$2$个负例，假设采样的负例为：

$w$	$c_{neg}$
natural	where
natural	if
natural	jam
natural	ping
natural	coaxial
natural	oh
natural	pang
natural	bang

由于我们有$4$个正例，我们就采样了$8$个负例。

那么目标就是训练这个分类器，使得

• 最大化正例中目标词和上下文词对$(w,c_{pos})$出现的概率
• 最小化负例中$(w,c_{neg})$词对出现的概率

假设我们考虑一个目标词/上下文词对$(w,c_{pos})$和$K$个噪音单词$c_{ne{g}_1},\cdots ,c_{ne{g}_K}$，那么基于公式$(15)$，我们需要最小化损失(所以加了个负号)，并同时考虑这两个目标：
$$\begin{array}{rl}{L}_{CE}& =-\log \left[P(+|w,c_{pos})\prod _{i=1}^{K}P(-|w,c_{ne{g}_i})\right]\\ & =-\left[\log P(+|w,c_{pos})+\sum _{i=1}^K\log P(-|w,c_{ne{g}_i})\right]\\ & =-\left[\log P(+|w,c_{pos})+\sum _{i=1}^K\log \left(1-P(+|w,c_{ne{g}_i})\right)\right]\\ & =-\left[\log \sigma (c_{pos}\cdot w)+\sum _{i=1}^K\log \sigma (-c_{ne{g}_i}\cdot w)\right]\end{array}\text{\hspace{1em}(8)}$$
这里分为两项，第一项我们希望分类器给正例很高的概率判断为$+$；第二项希望给负例很高的概率判断为$-$。

Sigmoid函数中计算了点积，我们想要最大化目标词与真正上下文单词的点积，同时最小化目标词与$K$个负样本的点积。

到此为止就已经可以实现模型了，但是我们深入一步，推导一下对每种嵌入的梯度。

首先看对$c_{pos}$的梯度：
$$\begin{array}{rl}\frac{\partial L_{CE}}{\partial c_{pos}}& =-\frac{\partial \left[\log \sigma (c_{pos}\cdot w)+{\sum }_{i=1}^K\log \sigma (-c_{ne{g}_i}\cdot w)\right]}{\partial c_{pos}}\\ & =-\frac{\sigma (c_{pos}\cdot w{)}^{\prime }\cdot w}{\sigma (c_{pos}\cdot w)}\\ & =-\frac{\sigma (c_{pos}\cdot w)[1-\sigma (c_{pos}\cdot w)]\cdot w}{\sigma (c_{pos}\cdot w)}\\ & =-[1-\sigma (c_{pos}\cdot w)]\cdot w\\ & =[\sigma (c_{pos}\cdot w)-1]\cdot w\end{array}\text{\hspace{1em}(9)}$$
然后是对第$i$个负样本$c_{ne{g}_i}$的梯度：
$$\begin{array}{rl}\frac{\partial L_{CE}}{\partial c_{ne{g}_i}}& =-\frac{\partial \left[\log \sigma (c_{pos}\cdot w)+{\sum }_{i=1}^K\log \sigma (-c_{ne{g}_i}\cdot w)\right]}{\partial c_{ne{g}_i}}\\ & =-\frac{\sigma (-c_{ne{g}_i}\cdot w{)}^{\prime }\cdot (-w)}{\sigma (-c_{ne{g}_i}\cdot w)}\\ & =-\frac{\sigma (-c_{ne{g}_i}\cdot w)[1-\sigma (-c_{ne{g}_i}\cdot w)]\cdot (-w)}{\sigma (-c_{ne{g}_i}\cdot w)}\\ & =\sigma (c_{ne{g}_i}\cdot w)\cdot w\end{array}\text{\hspace{1em}(10)}$$
最后看一下对中心词$w$的梯度：
$$\begin{array}{rl}\frac{\partial L_{CE}}{\partial w}& =-\frac{\partial \left[\log \sigma (c_{pos}\cdot w)+{\sum }_{i=1}^K\log \sigma (-c_{ne{g}_i}\cdot w)\right]}{\partial w}\\ & =-\frac{\sigma (c_{pos}\cdot w{)}^{\prime }\cdot c_{pos}}{\sigma (c_{pos}\cdot w)}-\sum _{i=1}^K\left[\frac{\sigma (-c_{ne{g}_i}\cdot w{)}^{\prime }\cdot (-c_{ne{g}_i})}{\sigma (-c_{ne{g}_i}\cdot w)}\right]\\ & =-\frac{\sigma (c_{pos}\cdot w)[1-\sigma (c_{pos}\cdot w)]\cdot c_{pos}}{\sigma (c_{pos}\cdot w)}-\sum _{i=1}^K\left[\frac{\sigma (-c_{ne{g}_i}\cdot w)[1-\sigma (-c_{ne{g}_i}\cdot w)]\cdot (-c_{ne{g}_i})}{\sigma (-c_{ne{g}_i}\cdot w)}\right]\\ & =-[1-\sigma (c_{pos}\cdot w)]\cdot c_{pos}+\sum _{i=1}^K\sigma (c_{ne{g}_i}\cdot w)\cdot c_{ne{g}_i}\\ & =[\sigma (c_{pos}\cdot w)-1]\cdot c_{pos}+\sum _{i=1}^K\sigma (c_{ne{g}_i}\cdot w)\cdot c_{ne{g}_i}\end{array}\text{\hspace{1em}(11)}$$
因为$w$参与了这两项，所以它的式子也由两项组成。

代码实现

首先构建SGNS数据集，对于每个训练(正)样本，需要根据某个负采样概率分布生成相应的负样本，同时需要保证负样本不包含当前上下文中的词。
一种实现方式是，在构建训练数据的过程中就完成负样本的生成，这样在训练时直接读取负样本即可。这么做的优点是训练过程无需再进行采样，因此效率较高；缺点是每次迭代使用的是同样的负样本，缺乏多样性。
这里采用在训练过程中实时进行负采样的实现方式，通过以下类的collate_fn函数完成负采样。

class SGNSDataset(Dataset):
    def __init__(self, corpus, vocab, window_size=2, n_negatives=5, ns_dist=None):
        self.data = []
        self.bos = vocab[BOS_TOKEN]
        self.eos = vocab[EOS_TOKEN]
        self.pad = vocab[PAD_TOKEN]

        for sentence in tqdm(corpus, desc='Dataset Construction'):
            sentence = [self.bos] + sentence + [self.eos]
            for i in range(1, len(sentence) - 1):
                # 模型输入：(w, context)
                # 输出：0/1，表示context是否为负样本
                w = sentence[i]
                left_context_index = max(0, i - window_size)
                right_context_index = min(len(sentence), i + window_size)
                context = sentence[left_context_index:i] + sentence[i + 1:right_context_index + 1]
                context += [self.pad] * (2 * window_size - len(context))
                self.data.append((w, context))

        # 负样本数量
        self.n_negatives = n_negatives
        # 负采样分布：若参数ns_dist为None，则使用uniform分布
        self.ns_dist = ns_dist if ns_dist is not None else Tensor.ones(len(vocab))

        self.data = np.asarray(self.data)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        return self.data[i]

    def collate_fn(self, examples):
        words = Tensor([ex[0] for ex in examples])
        contexts = Tensor([ex[1] for ex in examples])

        batch_size, window_size = contexts.shape
        neg_contexts = []
        # 对batch内的样本分别进行负采样
        for i in range(batch_size):
            # 保证负样本不包含当前样本中的context
            ns_dist = self.ns_dist.index_fill_(0, contexts[i], .0)
            neg_contexts.append(Tensor.multinomial(ns_dist, self.n_negatives * window_size, replace=True))
        neg_contexts = F.stack(neg_contexts, axis=0)
        return words, contexts, neg_contexts

在模型类中需要维护两个词向量w_embeddings和c_embeddings，分别用于词与上下文的向量表示。 同时因为word2vec模型比较特殊，我们不关心模型的输出，而是它学到的权重。为了简单，我们这里在forward中直接输出损失，使用公式$(8)$来计算总损失：

class SGNSModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        # 目标词嵌入
        self.w_embeddings = nn.Embedding(vocab_size, embedding_dim)
        # 上下文嵌入
        self.c_embeddings = nn.Embedding(vocab_size, embedding_dim)

    def forward(self, target_words, pos_contexts, neg_contexts) -> Tensor:
        '''
        word2vec模型比较特殊，我们不关心模型的输出，而是它学到的权重
        为了简单，我们这里直接输出损失
        '''
        batch_size = target_words.shape[0]
        n_negatives = neg_contexts.shape[-1]

        word_embeds = self.w_embeddings(target_words)  # (batch_size, embedding_dim)
        context_embeds = self.c_embeddings(pos_contexts)  # (batch_size, window_size * 2, embedding_dim)
        neg_context_embeds = self.c_embeddings(neg_contexts)  # (batch_size, window_size * n_negatives, embedding_dim)

        word_embeds = word_embeds.unsqueeze(2)

        # 正样本的对数似然
        context_loss = F.logsigmoid((context_embeds @ word_embeds).squeeze(2))
        context_loss = context_loss.mean(axis=1)
        # 负样本的对数似然
        neg_context_loss = F.logsigmoid((neg_context_embeds @ word_embeds).squeeze(axis=2).neg())
        neg_context_loss = neg_context_loss.reshape((batch_size, -1, n_negatives)).sum(axis=2)
        neg_context_loss = neg_context_loss.mean(axis=1)

        # 总损失： 负对数似然
        loss = -(context_loss + neg_context_loss).mean()

        return loss

但我们还需要编写从训练语料库中统计Unigram出现次数，并计算概率分布。以此概率为基础进行负采样：

def get_unigram_distribution(corpus, vocab_size):
    # 从给定语料中统计unigram概率分布
    token_counts = Tensor([.0] * vocab_size)
    total_count = .0
    for sentence in corpus:
        total_count += len(sentence)
        for token in sentence:
            token_counts[token] += 1
    unigram_dist = token_counts / total_count
    return unigram_dist

下面是具体的训练过程：

	  embedding_dim = 64
    window_size = 2
    batch_size = 10240
    num_epoch = 10
    min_freq = 3  # 保留单词最少出现的次数
    n_negatives = 10  # 负采样数

    # 读取数据
    corpus, vocab = load_corpus('../data/xiyouji.txt', min_freq)
    # 计算unigram概率分布
    unigram_dist = get_unigram_distribution(corpus, len(vocab))
    # 根据unigram分布计算负采样分数： p(w)**0.75
    negative_sampling_dist = unigram_dist ** 0.75
    # 构建数据集
    dataset = SGNSDataset(corpus, vocab, window_size=window_size, ns_dist=negative_sampling_dist)
    # 构建数据加载器
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        collate_fn=dataset.collate_fn,
        shuffle=True
    )

    device = cuda.get_device("cuda:0" if cuda.is_available() else "cpu")

    print(f'current device:{device}')

    # 构建模型
    model = SGNSModel(len(vocab), embedding_dim)
    model.to(device)

    optimizer = SGD(model.parameters())
    with debug_mode():
        for epoch in range(num_epoch):
            total_loss = 0
            for batch in tqdm(data_loader, desc=f'Training Epoch {epoch}'):
                words, contexts, neg_contexts = [x.to(device) for x in batch]
                optimizer.zero_grad()
                loss = model(words, contexts, neg_contexts)
                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            print(f'Loss: {total_loss:.2f}')

    save_pretrained(vocab, model.embeddings.weight, 'sgns.vec')

完整代码

https://github.com/nlp-greyfoss/metagrad

References

1. 从零实现Word2Vec
2. 自然语言处理：基于预训练模型的方法
3. Speech and Language Processing