15.4. Предобучение word2vec

Мы переходим к реализации модели skip-gram, определённой в Section 15.1. Затем мы обучим word2vec с использованием негативной выборки на наборе данных PTB. Прежде всего, получим итератор данных и словарь для этого набора данных, вызвав функцию d2l.load_data_ptb, описанную в Section 15.3

import math
import torch
from torch import nn
from d2l import torch as d2l

batch_size, max_window_size, num_noise_words = 512, 5, 5
data_iter, vocab = d2l.load_data_ptb(batch_size, max_window_size,
                                     num_noise_words)

15.4.1. Модель Skip-Gram

Мы реализуем модель skip-gram, используя слои эмбеддингов и пакетные матричные умножения. Сначала рассмотрим, как работают слои эмбеддингов.

15.4.1.1. Слой эмбеддингов

Как описано в Section 10.7, слой эмбеддингов отображает индекс токена в его вектор признаков. Весами этого слоя является матрица, число строк которой равно размеру словаря (input_dim), а число столбцов — размерности вектора для каждого токена (output_dim). После обучения модели эмбеддингов именно эти веса нам и нужны.

embed = nn.Embedding(num_embeddings=20, embedding_dim=4)
print(f'Parameter embedding_weight ({embed.weight.shape}, '
      f'dtype={embed.weight.dtype})')

Входом слоя эмбеддинга является индекс токена (слова). Для любого индекса токена \(i\) его векторное представление можно получить из строки \(i^\textrm{th}\) матрицы весов слоя эмбеддинга. Поскольку размерность вектора (output_dim) равна 4, слой эмбеддинга возвращает векторы формы (2, 3, 4) для мини-батча индексов токенов формы (2, 3).

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
embed(x)

15.4.1.2. Определение прямого распространения

В прямом распространении вход модели skip-gram включает индексы центрального слова center формы (размер батча, 1) и конкатенированные индексы контекстных и шумовых слов contexts_and_negatives формы (размер батча, max_len), где max_len определяется в Section 15.3.5. Сначала эти две переменные преобразуются из индексов токенов в векторы с помощью слоя эмбеддинга, затем их матричное умножение по батчам (описанное в Section 11.3.2.2) возвращает выход формы (размер батча, 1, max_len). Каждый элемент выхода представляет собой скалярное произведение вектора центрального слова и вектора контекстного или шумового слова.

def skip_gram(center, contexts_and_negatives, embed_v, embed_u):
    v = embed_v(center)
    u = embed_u(contexts_and_negatives)
    pred = torch.bmm(v, u.permute(0, 2, 1))
    return pred

Давайте выведем форму результата этой skip_gram функции для некоторых примеров входных данных.

skip_gram(torch.ones((2, 1), dtype=torch.long),
          torch.ones((2, 4), dtype=torch.long), embed, embed).shape

15.4.2. Обучение

Прежде чем приступать к обучению модели skip-gram с отрицательной выборкой, определим её функцию потерь.

15.4.2.1. Функция потерь бинарной перекрёстной энтропии

Согласно определению функции потерь для отрицательной выборки в Section 15.2.1, мы будем использовать функцию потерь бинарной перекрёстной энтропии.

class SigmoidBCELoss(nn.Module):
    # Binary cross-entropy loss with masking
    def __init__(self):
        super().__init__()

    def forward(self, inputs, target, mask=None):
        out = nn.functional.binary_cross_entropy_with_logits(
            inputs, target, weight=mask, reduction="none")
        return out.mean(dim=1)

loss = SigmoidBCELoss()

Вспомним описание переменной маски и переменной метки в Section 15.3.5. Ниже вычисляется функция потерь бинарной перекрёстной энтропии для указанных переменных.

pred = torch.tensor([[1.1, -2.2, 3.3, -4.4]] * 2)
label = torch.tensor([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]])
mask = torch.tensor([[1, 1, 1, 1], [1, 1, 0, 0]])
loss(pred, label, mask) * mask.shape[1] / mask.sum(axis=1)

Ниже показано, как вычисляются приведённые выше результаты (менее эффективным способом) с использованием функции активации сигмоиды в функции потерь бинарной перекрёстной энтропии. Можно рассматривать эти два выхода как два нормализованных значения потерь, усреднённых по немаскированным предсказаниям.

def sigmd(x):
    return -math.log(1 / (1 + math.exp(-x)))

print(f'{(sigmd(1.1) + sigmd(2.2) + sigmd(-3.3) + sigmd(4.4)) / 4:.4f}')
print(f'{(sigmd(-1.1) + sigmd(-2.2)) / 2:.4f}')

15.4.2.2. Инициализация параметров модели

Мы определяем два слоя эмбеддинга для всех слов в словаре: один для центральных слов, а другой — для слов контекста. Размерность вектора слова embed_size установлена равной 100.

embed_size = 100
net = nn.Sequential(nn.Embedding(num_embeddings=len(vocab),
                                 embedding_dim=embed_size),
                    nn.Embedding(num_embeddings=len(vocab),
                                 embedding_dim=embed_size))

15.4.2.3. Определение цикла обучения

Цикл обучения определен ниже. Из-за наличия дополнения вычисление функции потерь несколько отличается от предыдущих функций обучения.

def train(net, data_iter, lr, num_epochs, device=d2l.try_gpu()):
    def init_weights(module):
        if type(module) == nn.Embedding:
            nn.init.xavier_uniform_(module.weight)
    net.apply(init_weights)
    net = net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[1, num_epochs])
    # Sum of normalized losses, no. of normalized losses
    metric = d2l.Accumulator(2)
    for epoch in range(num_epochs):
        timer, num_batches = d2l.Timer(), len(data_iter)
        for i, batch in enumerate(data_iter):
            optimizer.zero_grad()
            center, context_negative, mask, label = [
                data.to(device) for data in batch]

            pred = skip_gram(center, context_negative, net[0], net[1])
            l = (loss(pred.reshape(label.shape).float(), label.float(), mask)
                     / mask.sum(axis=1) * mask.shape[1])
            l.sum().backward()
            optimizer.step()
            metric.add(l.sum(), l.numel())
            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
                animator.add(epoch + (i + 1) / num_batches,
                             (metric[0] / metric[1],))
    print(f'loss {metric[0] / metric[1]:.3f}, '
          f'{metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')

Теперь мы можем обучить модель skip-gram с использованием негативной выборки.

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

15.4.3. Применение векторных представлений слов

После обучения модели word2vec мы можем использовать косинусное сходство векторов слов из обученной модели для поиска слов из словаря, наиболее семантически близких к входному слову.

def get_similar_tokens(query_token, k, embed):
    W = embed.weight.data
    x = W[vocab[query_token]]
    # Compute the cosine similarity. Add 1e-9 for numerical stability
    cos = torch.mv(W, x) / torch.sqrt(torch.sum(W * W, dim=1) *
                                      torch.sum(x * x) + 1e-9)
    topk = torch.topk(cos, k=k+1)[1].cpu().numpy().astype('int32')
    for i in topk[1:]:  # Remove the input words
        print(f'cosine sim={float(cos[i]):.3f}: {vocab.to_tokens(i)}')

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

15.4.4. Итоги

  • Мы можем обучить модель skip-gram с отрицательной выборкой, используя слои эмбеддингов и функцию потерь бинарной перекрёстной энтропии.

  • Применение векторных представлений слов включает поиск слов, семантически близких к заданному, на основе косинусного сходства векторов слов.

15.4.5. Упражнения

  1. Используя обученную модель, найдите семантически похожие слова для других входных слов. Можно ли улучшить результаты, настроив гиперпараметры?

  2. Когда обучающий корпус огромен, мы часто выбираем контекстные и шумовые слова для центральных слов в текущем мини-батче при обновлении параметров модели. Иными словами, одно и то же центральное слово может иметь разные контекстные или шумовые слова в разных эпохах обучения. Каковы преимущества этого метода? Попробуйте реализовать этот подход к обучению.

Обсуждения