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

引言

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

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

本文我们来探讨word2vec。本文介绍word2vec词嵌入，它一种稠密向量模型，向量的元素值是实数，甚至可以是负数。

值得注意的是，word2vec是一种静态嵌入(static embeddings)模型，即为词典中的单词生成的是固定嵌入，而不是像BERT那样根据上下文生成动态嵌入。

看完本文，你应该可以一次性掌握word2vec的原理以及实现。

word2vec

word2vec是一种高效的训练词向量的模型。它的想法直接，如果两个单词的上下文相似，那么这两个单词(词向量)也应该是相似的。比如，“A dog is running in the room"和"A cat is running in the room”。这两个句子，只是"cat"和"dog"不同，word2vec认为它们是相似的，而n-gram模型做不到这一点。

这里的词向量是什么？为了便于计算机处理，我们需要把文档、单词向量化。而且除了向量化之后，还希望单词的表达能计算相似词信息。

word2vec有两种计算嵌入的方法：skip-gram和CBOW。

我们先来看看CBOW。

CBOW

CBOW(Continuous Bag-of-Words，连续词袋)模型的基本思想是根据上下文对中心词(目标词,target word)进行预测。例如，对于文本$\cdots w_{t-2}{w}_{t-1}\underline{w_t}{w}_{t+1}{w}_{t+2}\cdots$，CBOW模型的任务是根据一定窗口大小内的上下文$C_t$(这里窗口大小为$2$，则$C_t=\{w_{t-2},w_{t-1},w_{t+1},w_{t+2}\}$）对$t$时刻的单词$w_t$​进行预测。

但是要注意的是，CBOW模型不考虑单词的顺序，实际上是一个词袋模型，这就是它名字的由来。

CBOW模型可以表示成下图所示的前馈神经网络结构。但不同于一般的前馈神经网络，CBOW模型的隐藏层只是执行对词向量取平均的操作，而没有线性变换和非线性激活过程。这也是CBOW模型训练效率高的原因。

CBOW示意图，来自参考1

这里给定词典大小为$V$，每个输入向量$x$和输出向量$y$都是维度为$V$的独热编码。隐藏层的大小为$N$，表示得到的词嵌入维度。

实际上我们的输入可以直接是索引，这一点在代码实现可以看到。

输入层(input)： 以大小为$2$的窗口为例，在目标词$w_t$左右各取$2$个单词作为模型的输入。那么输入就由$4$个维度为词典大小$V$的独热编码向量构成。

隐藏层(Hidden)：隐藏层所做的事情就是对上下文$C_t$中的所有词向量取平均，得到一个上下文表示。具体来说，首先输入层中每个单词的独热编码向量经过矩阵$W\in {\mathbb{R}}^{V\times N}$映射到词向量空间：
$$v_{w_i}=W^T{x}_i\text{\hspace{1em}(1)}$$
这里$x_i$表示第$i$个单词的独热编码向量，维度是$V\times 1$；而$W$为$V\times N$的权重矩阵；

上式的结果是得到一个$N\times 1$​的向量，其实就是取矩阵$W$​的第$i$​行，也就是单词$w_i$的词向量。

其实这里用不着独热编码，直接取$W$​的第$i$​行这种索引操作就行了，Pytorch提供了nn.Embedding来实现这一点。

而$C_t=\{w_{t-2},w_{t-1},w_{t+1},w_{t+2}\}$表示所有$w_t$的上下文单词的集合，这里对$C_t$中所有单词的词向量取均值，作为$w_t$的上下文表示：
$$v_{C_t}=\frac{1}{|C_t|}\sum _{w\in C_t}{v}_w\text{\hspace{1em}(2)}$$
这里$|C_t|$表示该集合中单词的总数，这样$v_{C_t}$的维度还是$N\times 1$。

输出层(Output)：输出层做的就是一个多分类问题，与前馈神经网络类似，但是也丢弃了线性变换和偏置。输出层有一个不同的权重矩阵$W^{\prime }$，它的维度是$N\times V$的。如果说$W$是表示中心词的权重矩阵，那么$W^{\prime }$​就是表示上下文词的权重矩阵。它的每一列代表一个上下文单词的词向量。

令$v_{w_i}^{\prime }$为$w_i$在$W^{\prime }$中对应的列向量，维度为$N\times 1$。那么现在做的就是用$v_{C_t}$与每一个列向量做点积，得到一个分数，这个分数可以理解为衡量中心词$w_t$与输出词的相似度。

我们可以一次计算所有单词的得分：
$$u=W^{\prime T}\cdot v_{C_t}\text{\hspace{1em}(3)}$$
得到一个$V\times 1$的列向量，其中每个元素代表对应单词与中心词的相似得分。最终经过Softmax得到一个概率分布，再和实际的中心词独热编码做一个交叉熵计算损失，我们希望损失越小越好。

若是展开来看，并加上Softmax，那么输出中心词$w_t$的概率可计算为：
$$P(w_t|v_{C_t})=\frac{\exp (v_{C_t}\cdot v_{w_t}^{\prime })}{{\sum }_{w^{\prime }\in \mathbb{V}}\exp (v_{C_t}\cdot v_{w^{\prime }}^{\prime })}\text{\hspace{1em}(4)}$$
损失函数

从公式$(4)$可以看出，这其实是一个多分类问题。除了可以用交叉熵来作为损失函数，也可以用负对数似然损失：
$$L(\theta )=\sum _{t=1}^T\log P(w_t|C_t)\text{\hspace{1em}(5)}$$
其中$C_t=\{w_{t-k},\cdots ,w_{t-1},w_{t+1},\cdots ,w_{t+k}\}$表示窗口大小为$k$的上下文单词的集合。

有了正向传播过程和损失函数，我们就可以利用metagrad进行代码实现了，而不用关心反向传播过程。

如果想了解反向传播的推导，可以参考从零实现Word2Vec[^2]。

在CBOW模型的参数中，矩阵$W$和$W^{\prime }$都可以作为词向量矩阵，它们分别描述了词典中的词在作为目标词或上下文词的不同性质。在实际中通常只用$W$就可以满足应用需求。

在介绍Skip-gram模型之前，我们先来实现已经学习的CBOW模型。

代码实现

在实现上面的权重矩阵$W$时，可以基于没有偏置的线性层Linear来实现，这样输入就如上面所说的one-hot向量。但是还有一种更常用的实现，那么就是通过嵌入层Embedding。

如果通过one-hot向量加线性层实现，就是用one-hot向量与$W$进行矩阵运算，实际上就是取的$W$中的第$k$行，假设one-hot向量中第$k$个元素为$1$。那么与其进行这么复杂的运算，不如直接传入索引$k$，拿到$W$的第$k$行。嵌入层就是实现这个功能的。

嵌入层的实现

class Embedding(Module):
    def __init__(self, num_embeddings: int, embedding_dim: int, _weight: Optional[Tensor] = None,
                 dtype=None, device=None) -> None:
        '''
        一个存储固定大小词汇表嵌入的查找表，可以通过索引(列表)直接访问，而不是one-hot向量。
        :param num_embeddings: 词汇表大小
        :param embedding_dim:  嵌入维度
        '''

        super(Embedding, self).__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim

        # 也可以传预训练好的权重进来
        if _weight is None:
            self.weight = Parameter(Tensor.empty((num_embeddings, embedding_dim), dtype=dtype, device=device))
            self.reset_parameters()
        else:
            assert list(_weight.shape) == [num_embeddings, embedding_dim], \
                'Shape of weight does not match num_embeddings and embedding_dim'
            self.weight = Parameter(_weight, device=device)

    def reset_parameters(self) -> None:
        init.uniform_(self.weight)

    def forward(self, input: Tensor) -> Tensor:
        return F.embedding(self.weight, input)

    @classmethod
    def from_pretrained(cls, embeddings: Tensor, freeze=True):
        assert embeddings.ndim == 2, \
            'Embeddings parameter is expected to be 2-dimensional'
        rows, cols = embeddings.shape
        embedding = cls(num_embeddings=rows, embedding_dim=cols, _weight=embeddings)
        embedding.weight.requires_grad = not freeze
        return embedding

代码也不复杂，这里还提供了从已经训练好的权重中加载的功能。在forward中直接调用embedding函数。显然核心逻辑在该函数里面，我们来实现看。

class Embedding(Function):
    def forward(ctx, weight: NdArray, indices: NdArray) -> NdArray:
        ctx.save_for_backward(weight.shape, indices)
        return weight[indices]

    def backward(ctx, grad: NdArray) -> Tuple[NdArray, None]:
        w_shape, indices = ctx.saved_tensors

        xp = get_array_module(grad)

        bigger_grad = xp.zeros(w_shape, dtype=grad.dtype)

        if xp is np:
            np.add.at(bigger_grad, indices, grad)
        else:
            bigger_grad.scatter_add(indices, grad)

        # 因为它有两个输入，防止错误地拆开bigger_grad
        # indices 不需要梯度
        return bigger_grad, None

def embedding(weight: Tensor, indices: Tensor) -> Tensor:
    return Embedding.apply(Embedding, weight, indices)

实现起来类似我们之前的slice函数，毕竟操作上本质是一样的嘛。

当然还有必不可少的单元测试，相关代码请参考完整代码。

那么接下来我们就可以实现CBOW模型了。

模型实现

首先我们要构建词典：

BOS_TOKEN = ""  # 句子开始标记
EOS_TOKEN = ""  # 句子结束标记
PAD_TOKEN = ""  # 填充标记
UNK_TOKEN = ""  # 未知词标记


class Vocabulary:
    def __init__(self, tokens=None):
        self._idx_to_token = list()
        self._token_to_idx = dict()

        # 如果传入了去重单词列表
        if tokens is not None:
            if UNK_TOKEN not in tokens:
                tokens = tokens + [UNK_TOKEN]
            # 构建id2word和word2id
            for token in tokens:
                self._idx_to_token.append(token)
                self._token_to_idx[token] = len(self._idx_to_token) - 1

            self.unk = self._token_to_idx[UNK_TOKEN]

    @classmethod
    def build(cls, text, min_freq=2, reserved_tokens=None):
        '''
        构建词表
        :param text: 处理好的(分词、去掉特殊符号等)text
        :param min_freq: 最小单词频率
        :param reserved_tokens: 预先保留的标记
        :return:
        '''
        token_freqs = defaultdict(int)
        for sentence in text:
            for token in sentence:
                token_freqs[token] += 1

        unique_tokens = [UNK_TOKEN] + (reserved_tokens if reserved_tokens else [])
        unique_tokens += [token for token, freq in token_freqs.items() \
                          if freq >= min_freq and token != UNK_TOKEN]
        return cls(unique_tokens)

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

    def __getitem__(self, token):
        '''得到token对应的id'''
        return self._token_to_idx.get(token, self.unk)

    def token(self, idx):
        assert 0 <= idx < len(self._idx_to_token)
        '''根据索引获取token'''
        return self._idx_to_token[idx]

    def to_ids(self, tokens):
        return [self[token] for token in tokens]

    def to_tokens(self, indices):
        return [self._idx_to_token[index] for index in indices]

然后我们需要自定义数据集：

class CBOWDataset(Dataset):
    def __init__(self, corpus, vocab, window_size=2):
        self.data = []
        self.bos = vocab[BOS_TOKEN]
        self.eos = vocab[EOS_TOKEN]

        for sentence in tqdm(corpus, desc='Dataset Construction'):
            sentence = [self.bos] + sentence + [self.eos]
            # 如果句子长度不足以构建(上下文,目标词)训练样本，则跳过
            if len(sentence) < window_size * 2 + 1:
                continue
            for i in range(window_size, len(sentence) - window_size):
                # 分别取i左右window_size个单词
                context = sentence[i - window_size:i] + sentence[i + 1:i + window_size + 1]
                # 目标词：当前词
                target = sentence[i]
                self.data.append((context, target))

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

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

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

    @staticmethod
    def collate_fn(examples):
        '''
        自定义整理函数
        :param examples:
        :return:
        '''
        inputs = Tensor([ex[0] for ex in examples])
        targets = Tensor([ex[1] for ex in examples])
        return inputs, targets

构建(上下文,目标词)训练样本，并且实现自定义的整理函数。

下面我们就可以构建模型了，

class CBOWModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        # 词向量层，即权重矩阵W
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        # 输出层，包含权重矩阵W'
        self.output = nn.Linear(embedding_dim, vocab_size, bias=False)

    def forward(self, inputs: Tensor) -> Tensor:
        # 得到所有上下文嵌入向量
        embeds = self.embeddings(inputs)
        # 计算均值，得到隐藏层向量，作为目标词的上下文表示
        hidden = embeds.mean(axis=1)
        output = self.output(hidden)
        return output

参考上面的描述图，其实就是两个权重矩阵。我们一个用嵌入层实现，另一个用不带偏置项的线性层实现。

在训练之前，我们需要构建词典对象，和处理好的语料。

def load_corpus(corpus_path):
    '''
    从corpus_path中读取预料
    :param corpus_path: 处理好的文本路径
    :return:
    '''
    with open(corpus_path, 'r', encoding='utf8') as f:
        lines = f.readlines()
    # 去掉空行，将文本转换为单词列表
    text = [[word for word in sentence.split(' ')] for sentence in lines if len(sentence) != 0]
    # 构建词典
    vocab = Vocabulary.build(text, reserved_tokens=[PAD_TOKEN, BOS_TOKEN, EOS_TOKEN])
    print(f'vocab size:{len(vocab)}')
    # 构建语料:将单词转换为ID
    corpus = [vocab.to_ids(sentence) for sentence in text]

    return corpus, vocab

最后就可以开始训练了：

 	embedding_dim = 64
    window_size = 3
    batch_size = 2048
    num_epoch = 2000
    min_freq = 3  # 保留单词最少出现的次数

    corpus, vocab = load_corpus('../../data/xiyouji.txt', min_freq)
    # 构建数据集
    dataset = CBOWDataset(corpus, vocab, window_size=window_size)
    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}')

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

    optimizer = SGD(model.parameters(), 1)
    for epoch in range(num_epoch):
        total_loss = 0
        for batch in tqdm(data_loader, desc=f'Training Epoch {epoch}'):
            inputs, targets = [x.to(device) for x in batch]
            optimizer.zero_grad()
            output = model(inputs)
            loss = loss_func(output, targets)
            loss.backward()
            optimizer.step()
            total_loss += loss

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

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

数据集采用的是《西游记》，经过分词、去掉标点符号预处理。
能看到这里的都是粉丝，这里直接放出处理好的数据集。
数据集下载 → 提取码：nap4

为了加速，我们使用GPU进行训练。所配的参数如下：

 	embedding_dim = 64
    window_size = 3
    batch_size = 2048
    num_epoch = 2000
    min_freq = 3 

最终的Loss为：800+

实验效果：

> search('观音', embeddings, vocab)
故此: 0.5987884141294884
观音菩萨: 0.5976461631931431
菩萨: 0.5212316212655066
> search('孙悟空', embeddings, vocab)
齐天大圣: 0.5778116509661732
名字: 0.5639390829512272
那方: 0.5528188565550192
> search('呆子', embeddings, vocab)
八戒: 0.6547101716347253
行者: 0.6176272985497067
沙僧: 0.5527797715535391
> search('如来', embeddings, vocab)
佛祖: 0.7029674765576888
佛: 0.5874278308846846
菩萨: 0.5716678406707916
> search('唐僧', embeddings, vocab)
长老: 0.7743582601251642
圣僧: 0.7191300108695816
那怪: 0.6567922349528186

训练好的模型下载 → 提取码: p7ye

测试方法，运行examples/embeddings/load_and_test.py即可。

Skip-gram模型

CBOW模型使用上下文窗口词中的集合作为输入来预测目标词，即$P(w_t|C_t)$。而Skip-gram模型是根据当前词$w_t$来预测上下文词$C_t$。

Skig-gram示意图，来自参考1

这里给定词典大小为$V$，输入向量$x$和输出向量$y$都是维度为$V$的独热编码。隐藏层的大小为$N$，表示得到的词嵌入维度。

输入层 这里也以窗口大小$k=2$为例，输入$w_i\in R^{V\times 1}$是$V$维的独热编码，也记为$w_t$。

隐藏层 $w_t$通过矩阵$W\in R^{N\times V}$投影到隐藏层，这里隐藏层向量即为$w_t$的词向量$v_{w_t}\in R^{N\times 1}$：

$$v_{w_t}=W^{T}x\text{\hspace{1em}(6)}$$

输出层 输出层利用线性变换矩阵$W^{\prime }$​对上下文窗口内的单词进行预测：

具体做法是，假设$v_c^{\prime }$某个上下文单词在$W^{\prime }\in R^{N\times V}$中对应的列向量，维度为$N\times 1$。那么也是用当前词的词向量$v_{w_t}$与上下文单词的词向量$v_c^{\prime }$做一个点积，得到一个数值作为得分，也可以看成相似度。

那么我们也可以一次计算词典中所有单词的得分(其实就是公式$(8)$中的：
$$u=W^{\prime T}\cdot v_{w_t}\text{\hspace{1em}(7)}$$
得到的$u\in R^{V\times 1}$，每个元素代表对应单词与当前词的相似得分，最终经过Softmax得到一个概率分布。若展开来看，那么由中心词计算上下文词$c$的概率为：
$$P(c|w_t)=\frac{\exp (v_{w_t}\cdot v_c^{\prime })}{{\sum }_{w^{\prime }}\exp (v_{w_t}\cdot v_{w^{\prime }}^{\prime })}\text{\hspace{1em}(8)}$$
其中$c\in \{w_{t-2},w_{t-1},w_{t+1},w_{t+2}\}$​。

损失函数

Skip-gram模型的负对数似然损失函数为：
$$L(\theta )=-\sum _{t=1}^T\sum _{-k\le j\le k,j\ne 0}\log P(w_{t+j}|w_t)\text{\hspace{1em}(9)}$$
即希望基于$w_t$预测得到的上下文单词$w_{t+j}$出现的概率越高越好。

代码实现

有了上面的基础，我们直接进行模型实现。

Skip-gram模型的输入输出与CBOW模型接近，主要区别在于Skip-gram的输入输出都是单个单词，即在一定上下文窗口大小内共现的词对，而CBOW模型的输入是多个上下文单词与一个中心词组成的词对。

我们首先构建这种数据集：

class SkipGramDataset(Dataset):
    def __init__(self, corpus, vocab, window_size=2):
        self.data = []
        self.bos = vocab[BOS_TOKEN]
        self.eos = vocab[EOS_TOKEN]

        for sentence in tqdm(corpus, desc='Dataset Construction'):
            sentence = [self.bos] + sentence + [self.eos]

            for i in range(1, len(sentence) - 1):
                # 模型输入：当前词
                w = sentence[i]
                # 模型输出： 窗口大小内的上下文
                # max 和 min 防止越界取到非预期的单词
                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]
                self.data.extend([(w, c) for c in context])

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

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

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

    @staticmethod
    def collate_fn(examples):
        '''
        自定义整理函数
        :param examples:
        :return:
        '''
        inputs = Tensor([ex[0] for ex in examples])
        targets = Tensor([ex[1] for ex in examples])
        return inputs, targets

从代码可以看出，假设窗口大小为$k$，那么我们一次就得到了$2k$个训练样本(中心词，上下文词)。所以还是单类别多分类问题。

模型实现就更简单了，不需要求均值：

class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
        self.output = nn.Linear(embedding_dim, vocab_size)

    def forward(self, inputs: Tensor) -> Tensor:
        # 得到输入词向量
        embeds = self.embeddings(inputs)
        # 根据输入词向量，对上下文进行预测，得到每个单词的得分，但是我们只关注样本中与中心词对应的上下文词的得分,期望越高越好。
        output = self.output(embeds)
        return output

最后的训练代码为：

embedding_dim = 64
    window_size = 3
    batch_size = 1024
    num_epoch = 10
    min_freq = 3  # 保留单词最少出现的次数

    # 读取文本数据，构建Skip-gram模型训练数据集
    corpus, vocab = load_corpus('data/xiyouji.txt', min_freq)
    dataset = SkipGramDataset(corpus, vocab, window_size=window_size)
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        collate_fn=dataset.collate_fn,
        shuffle=True
    )

    loss_func = CrossEntropyLoss()
    # 构建Skip-gram模型，并加载至device
    device = cuda.get_device("cuda:0" if cuda.is_available() else "cpu")
    model = SkipGramModel(len(vocab), embedding_dim)
    model.to(device)
    optimizer = SGD(model.parameters(), lr=1)

    for epoch in range(num_epoch):
        total_loss = 0
        for batch in tqdm(data_loader, desc=f"Training Epoch {epoch}"):
            inputs, targets = [x.to(device) for x in batch]
            optimizer.zero_grad()
            output = model(inputs)
            loss = loss_func(output, targets)
            loss.backward()
            optimizer.step()
            total_loss += loss
        print(f"Loss: {total_loss.item():.2f}")

完整代码

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

References

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