The high level idea of word2vec is that words used in similar contexts are similar. We’ll find the words “fridge” and “freezer” in similar places in sentences, so those two words are likely related in some way.

One of the popular approaches is to train a neural network to predict neighbouring words. So if we give it e.g. the word “maintenance” it may predict that the surrounding words are likely to contain “package” or “burden”. If we give it the word “understanding”, it may predict that the surrounding words will contain “clearer” or “reached”.

Here’s the one weird trick: we pick a very specific shape for this neural network, namely one of three parts:

1. Input vector: one-hot encoded word given as input.
2. Hidden layer with k neurons.
3. Output vector: softmax scaled prediction of neighbouring words.

Now when the network is trained to predict neighbouring words, each possible input word will be associated with weights that control the vaules of the hidden layer for that word. Those weights become the vector for that word. That’s how word embeddings are extracted from the neural network.

Intuitively, we can imagine the hidden layer as encoding the underlying concept the input word corresponds to, and the neural network uses that underlying concept to guess what the neighbouring words are. This means the word vector will be a representation of what general concept a word belongs to. So for the example of the word “fridge”, the underlying concept may include elements of noun, appliance, common item, large, etc.

This also explains why we can use the weights from the hidden layer as word vectors: when comparing those word vectors, we are really comparing the underlying general concepts!

By reducing the number of neurons in the hidden layer, we can force a dimensionality reduction in the concept space.

The idea of gradient descent makes complete sense: if someone tells us in which direction to nudge the inputs of a function to reduce loss, then obviously we just nudge inputs in that direction. The gradient is the mathematical operation that produces the information on which direction to nudge inputs in.

Computing the gradient is usually done through backpropagation.

The idea of backpropagation also makes complete sense. A neural network is really just chained matrix multiplications and activation functions (see above). To compute the gradient of chained function applications, we can do it one step at a time. We could do it forwards, but doing it backwards makes sense from a dynamic programming perspective, because it allows us to reuse results from previous computations.

Then comes actually implementing it in practise. It’s definitely not rocket surgery, but – at least for me – it really required concentration and a good chunk of pencil-and-paper thinking.

But, then, suddenly2² These things always seem to happen suddenly to me. I have no sense of how far off success still is – I’m just working through bugs and understandings until, well, suddenly it works., there it was! A neural network that could predict neighbouring words from a small set of training pairs.

I can’t find any output from this version, but it was exhilerating to have something working, simple as it was. With a small enough vocabulary, it could perfectly predict the other half of the word pairs I fed it.3³ In other words, its reliability was excellent; it’s (external) validity was abysmal. That’s part of the purpose of the hidden layer. By going through concepts, it makes the network a little bit more robust. That’s a low bar to clear, but it was my first neural network. And it worked.

The network is – if we ignore its structure5⁵ There’s an alternative architectural choice here: make the structure part of the network data. That is the far more general implementation, but I was aiming for quickly implementing a specific network, not generality. The 2020-era word2vec type network will never change, so there’s no point in making code that can adapt to such change! – defined by just two things: the weights used when going from input to hidden layer, and the weights used when going from hidden to output layer. That results in a short function to create the network:

sub init_network {
    my $w_hidden = ones $vocab_width, $concept_width;
    my $w_out = ones $concept_width, $vocab_width;
    for (0..1000) {
        $w_hidden(int(rand($vocab_width)), int(rand($concept_width))) .= rand();
        $w_out(int(rand($concept_width)), int(rand($vocab_width))) .= rand();
    }
    return {
        w_hidden => $w_hidden,
        w_out => $w_out,
    };
}

Note that in this implementation, we have been lazy and not bothered looking up how to create a matrix initialised to random values. Instead, we look up 1000 random locations in the matrices and set those to random values.

This is one of the simpler steps: multiply the matrices and apply softmax to the output layer. The intermediate values are all returned since at least some of them will be needed to compute gradients.

sub forward_network {
    my ($network, $input) = @_;
    my $hidden = $network->{w_hidden} x $input;
    my $output = $network->{w_out} x $hidden;
    my $prediction = $output->exp / $output->exp->sum;
    return {
        input => $input,
        hidden => $hidden,
        output => $output,
        prediction => $prediction,
    };
}

In the code, we will see some confusing names used. The layer that sits between the hidden layer and the prediction is named output. The same name is used for the second half of the training data pairs! This did lead to some mistakes, but not enough that I bothered going through the code and renaming one of the things.

It seemed sensible to me to implement the scoring of the prediction as a separate function, because then we can run th enetwork forward with only an input word. This computes the cross-entropy of the prediction and the output part of the training pair.

sub evaluate_forward {
    my ($forward, $output) = @_;
    my $loss = - sum ($output * $forward->{prediction}->log);
    return {
        forward => $forward,
        output => $output,
        loss => $loss,
    };
}

This part of the code was by far the trickiest to get right6⁶ And to be honest, not entirely sure I have gotten it completely right in the published code here, either.. To avoid accidentally writing a how to backpropagation tutorial here (see links at the bottom of this article), let’s just see the code and move on.

The variables named d_x_y refer to the gradient of x with respect to y. The terms whidden and wout are the weights associated with the hidden and output layers respectively.

sub propagate_backward {
    my ($network, $evaluation) = @_;

    my $d_loss_pred = - $evaluation->{output} * 1/$evaluation->{forward}{prediction};

    my $d_pred_out = do {
        my $i_v = identity $vocab_width;
        my $s_i = $evaluation->{forward}{prediction}->dup(0, $vocab_width);
        my $s_j = $evaluation->{forward}{prediction}->transpose->dup(1, $vocab_width);
        $s_i * ($i_v - $s_j);
    };
    my $d_loss_out = $d_pred_out x $d_loss_pred;

    my $d_out_wout = $evaluation->{forward}{hidden}->transpose->dup(1, $vocab_width);
    my $d_loss_wout = $d_out_wout * $d_loss_out;

    my $d_out_hidden = $network->{w_out}->transpose;
    my $d_loss_hidden = $d_out_hidden x $d_loss_out;

    my $d_hidden_whidden = $evaluation->{forward}{input}->transpose->dup(1, $concept_width);
    my $d_loss_whidden = $d_hidden_whidden * $d_loss_hidden;

    return {
        d_loss_whidden => $d_loss_whidden,
        d_loss_wout => $d_loss_wout,
    };
}

This part of the code is by far the slowest part. When profiling, the computation of $d_pred_out really stood out as slowest of them all.7⁷ Turns out it’s very expensive to create huge matrices only to multiply most of them away again. Who would have thought? I created a slightly optimised version of that computation that computes $d_pred_out directly without any intermediate states, and the computation for that variable looks like

my $d_loss_out = do {
    my ($i) = @{unpdl which $evaluation->{output}};
    my $p = $evaluation->{forward}{prediction}->copy;
    my ($d_loss_pred_ish) = @{unpdl (-1/$p(:,($i)))};
    my $d = zeroes $vocab_width;
    $d($i) .= 1;
    $d = $d->transpose - $p(:, ($i))->dup(1, $vocab_width);
    $p = $p * $d;
    $p * $d_loss_pred_ish;
};

This is a bit tricky to follow, but it’s based on the fact that the output word is one-hot encoded, so the cross-entropy loss is zero for almost all components of the prediction except in one location.

The important part is that it improved training time by almost a factor of 20! With the slower code, training 1000 pairs takes 30 seconds, and with the newer code it takes just under two seconds.

Once we have the hard-earned gradients for the weights, the actual update to walk down the gradient is a matter of arithmetic. Since the gradient points upslope (toward higher loss), we apply a small fraction of negative gradient to shimmy downslope.

sub gradient_descent {
    my ($network, $gradients) = @_;
    my $step_size = 0.01;
    my $w_hidden_next = $network->{w_hidden} + -$step_size * $gradients->{d_loss_whidden};
    my $w_out_next = $network->{w_out} + -$step_size * $gradients->{d_loss_wout};
    return {
        w_hidden => $w_hidden_next,
        w_out => $w_out_next,
    };
}

Calling these functions in sequence yields an iteration of training. We can do these iterations many times to hopefully get something sensible.

my $n = init_network;
for my $iters (0..$max_iters-1) {
    # Pick a random word pair to train on.
    my $word = int(rand($max_samples));

    # Input the word and run the network forward.
    my $f = forward_network $n, $inputs(:, ($word))->transpose;

    # Compare to the actual output and compute cross entropy loss.
    my $e = evaluate_forward $f, $outputs(:, ($word))->transpose;

    # Run backpropagation to find relevant gradients.
    my $d = propagate_backward $n, $e;

    # Apply gradients to nudge weights in the right direction.
    $n = gradient_descent $n, $d;
}

For a while I had what I thought was a working neural network, except it kept predicting nonsense. I tried adjusting parameters, trying to train it for longer, reduced the vocabulary to just a few words, but nothing helped. After spending way too long trying to find a problem in the backpropagation, I realised I had given the network a nonsense vector for the input, rather than the one-hot vector corresponding to the word I thought I had asked it to predict.

This is a really good lesson. When code misbehaves, it’s almost always due to a bug in the code, rather than faulty input, so it’s easy to develop a heuristic to start debugging immediately. However, checking for faulty input is so much easier than debugging code that works, so it’s still a better first response to verify the input, on the off chance that this was one of those one-in-a-hundred times when the input really is the problem.

We saw one optimisation in the section on backpropagation. I often profile code manually, by injecting counting logic in the code I suspect is slow, and then analyse the rates at which things happen.

For this code, I tried Perl’s NYTProf. It was a magical experience. It showed not only which functions were slower than the rest, but also which statements were slow. It even displayed which parts of each statement was slow!

Without that, it would have taken me much longer to figure out that the innocent-looking statement $s_i * ($i_v - $s_j) was the main cause of slow execution8⁸ Blessing and curse of working with matrices. We get low-effort vectorisation, but it’s also easy to accidentally work on huge data structures without it being obvious in the code.. Optimising that statement away made training go 20× faster.