Ever since Watson left me speechless years ago, I've been trying to closely follow Natural Language Processing (NLP). Lately, the standard was to use Recurrent Neural Networks (RNN) or LSTMs. The problem is they process text like humans do: word by word, from left to right. If you have a long sentence, the network "forgets" the beginning when it reaches the end, and the worst part: you can't parallelize the computation on GPUs. You have to wait to process word 1 to process word 2.
Earlier this summer, a group of researchers from Google Brain published a paper with a cocky title: "Attention Is All You Need". I printed it out, I've been reading it in my spare time for a month, and I think I've finally understood the mathematical genius it proposes: the Transformer architecture.
The magic of Self-Attention
The authors propose throwing recurrent networks in the trash. Instead, the Transformer injects the entire sentence into the neural network at once. To maintain context, they use a mechanism called "Self-Attention".
The idea is that each word in the sentence simultaneously looks at the rest of the words and mathematically calculates how much "attention" it should pay them to understand its own meaning in that specific context. In the phrase "The park bank is broken", the word "bank" will pay a lot of attention to "park" and "broken", discarding the financial meaning.
Technically, this is solved with pure algebraic matrix multiplication. Three vectors are assigned to each word: Query (Q), Key (K), and Value (V).
The formula from the paper, which I already know by heart, is: $$ Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V $$
In Python pseudocode (using Numpy), the conceptual basis of an attention layer would look like this:
import numpy as np
def self_attention(Q, K, V, d_k):
# 1. Multiplicamos Queries por Keys transpuestas para ver la "afinidad"
scores = np.matmul(Q, K.transpose())
# 2. Escalamos para evitar que los gradientes exploten
scores = scores / np.sqrt(d_k)
# 3. Aplicamos Softmax para convertir los scores en probabilidades (pesos de 0 a 1)
pesos = softmax(scores)
# 4. Multiplicamos por los Values para obtener el vector contextualizado final
return np.matmul(pesos, V)
Reflection: The end of sequences
This simple matrix multiplication changes everything. Since we don't have to process word by word, we can shove gigantic matrices into those GPUs we used for computer vision and train massive models in a fraction of the time.
The paper uses it for machine translation, but the implications are dizzying. If we can parallelize language training, well-funded labs are going to start training models with billions of parameters sucking up half the internet. We've just broken the NLP bottleneck.