Attention and embeddings

A word embedding is a dense, relatively low-dimensional vector of real numbers that represents a word so that similar meanings end up close together in the vector space. Unlike one-hot vectors, which are huge, sparse, and treat every word as equally different from every other, embeddings capture semantic relationships and let models generalize across related words. They are usually learned from data rather than assigned by hand.

A high cosine similarity (close to 1) means the two vectors point in nearly the same direction, which the model has learned to associate with words used in similar contexts or with related meanings. Cosine similarity focuses on direction rather than magnitude, so even words appearing at different frequencies can be judged as semantically close. For example, 'king' and 'queen' typically have higher cosine similarity than 'king' and 'banana'.

It demonstrates that embeddings can capture certain relationships as consistent directions (vector offsets) in the space, so adding and subtracting vectors performs a kind of analogical reasoning. The gender difference between 'man' and 'woman' corresponds to roughly the same offset as between 'king' and 'queen'. This linear structure emerges naturally when embeddings are trained on large text corpora.

Attention lets a model decide, for each position it is processing, which other parts of the input are most relevant and weight them accordingly. Instead of treating all inputs equally, the model computes a set of weights that sum to one and forms a weighted combination of values. This allows it to focus on the most informative pieces of context.

The Query represents what the current position is looking for, the Keys describe what each input position offers, and the Values carry the actual information to be retrieved. Attention compares the Query against all Keys to produce relevance scores, then uses those scores to take a weighted sum of the Values. This is loosely analogous to looking up information in a database by matching a search query against keys.

Scaled dot-product attention is computed as softmax(QKᵀ / √dₖ) · V, where dₖ is the dimension of the key vectors. The division by √dₖ keeps the dot products from growing too large in magnitude, which would otherwise push the softmax into regions with extremely small gradients. This scaling keeps training stable and gradients useful.

import numpy as np

def attention(Q, K, V):
    dk = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(dk)
    weights = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ V

Softmax turns the raw relevance scores into a set of non-negative weights that sum to one, so they behave like a probability distribution over the input positions. This makes the output a proper weighted average of the values, where larger scores get proportionally larger influence. It also makes the mechanism differentiable so it can be trained with gradient descent.

In self-attention, the Queries, Keys, and Values all come from the same sequence, so each position attends to other positions within the same input. In cross-attention, the Queries come from one sequence while the Keys and Values come from a different sequence, such as a decoder attending to an encoder's output. Self-attention builds internal context; cross-attention links two different sequences.

Multi-head attention runs several attention operations in parallel, each with its own learned projections of the Queries, Keys, and Values, and then concatenates their outputs. This lets different heads focus on different types of relationships, such as one tracking syntax and another tracking long-range references. The result is a richer representation than any single head could produce.

Self-attention treats its inputs as an unordered set and has no built-in notion of sequence order, so without extra information it cannot tell whether a word came first or last. Positional encodings add order information to each token's embedding so the model can distinguish positions and understand sequence structure. This is essential for language, where word order changes meaning.

Skip-gram trains a model to predict the surrounding context words given a single center word, sliding this task across a large corpus. The weights that the network learns for each word become its embedding, so words appearing in similar contexts develop similar vectors. This reflects the distributional hypothesis: words with similar contexts tend to have similar meanings.

The distributional hypothesis states that words appearing in similar contexts tend to have similar meanings, often summarized as 'a word is known by the company it keeps.' Embedding methods exploit this by learning vectors from the contexts in which words occur, so context-sharing words end up nearby in the vector space. It is the core linguistic justification for learning meaning from raw text.

Static embeddings assign each word a single fixed vector regardless of the sentence it appears in, so 'bank' has the same vector whether it means a riverbank or a financial bank. Contextual embeddings are produced by a model that reads the whole sentence, so the same word gets different vectors depending on context. This lets contextual models handle polysemy (multiple meanings) far better.

An attention weight is the normalized score, produced by softmax, that tells how much a given input position contributes to the output for a particular query. Because softmax outputs are non-negative and sum to one across all positions, each individual weight lies between 0 and 1. A weight near 1 means the model is focusing almost entirely on that position.

The dot product is large and positive when two vectors point in similar directions and have substantial magnitude, so it acts as an unnormalized similarity measure. This lets attention score a Query against each Key cheaply using fast matrix multiplication. Higher dot products lead to higher attention weights after softmax, focusing the model on better-matching keys.

A causal mask prevents each position from attending to positions that come later in the sequence, typically by setting their scores to negative infinity before softmax. This is needed in autoregressive generation so that when predicting a token, the model cannot 'peek' at future tokens it is supposed to predict. It enforces that the model only uses past and present context.

The embedding dimension is the number of components in each word's vector, determining how much information the representation can hold. Larger dimensions can capture more nuanced relationships but require more memory and data and risk overfitting, while smaller ones are efficient but may lose detail. Choosing it balances expressiveness against computational cost.

Cosine similarity measures the angle between vectors, ignoring their magnitude, which often correlates with word frequency rather than meaning. This makes it focus on the direction that encodes semantic content, so frequent and rare words can still be compared fairly. Euclidean distance, by contrast, can be dominated by magnitude differences that are not semantically meaningful.

Softmax exponentiates each score and normalizes, giving e², e¹, e⁰ which are about 7.39, 2.72, and 1.0, summing to roughly 11.1. Dividing gives approximately 0.67, 0.24, and 0.09. So the first position receives the most attention, reflecting its highest score.

import numpy as np
scores = np.array([2.0, 1.0, 0.0])
w = np.exp(scores) / np.exp(scores).sum()
print(np.round(w, 2))  # [0.67 0.24 0.09]

An embedding layer is essentially a lookup table that maps each token's integer ID to a learnable vector. During training, these vectors are updated by backpropagation just like any other weights, so they adapt to the task. At inference time, the layer simply retrieves the stored vector for each input token.

import torch
import torch.nn as nn

emb = nn.Embedding(num_embeddings=10000, embedding_dim=300)
tokens = torch.tensor([1, 42, 7])
vectors = emb(tokens)  # shape: (3, 300)

RNNs process tokens sequentially and compress everything into a hidden state, so information from far-back tokens can fade or get distorted before it is needed. Attention lets any position directly access any other position in a single step, regardless of distance, giving a constant-length path between tokens. This makes capturing long-range relationships much easier and more reliable.

Permutation equivariant means that if you reorder the input tokens, the self-attention outputs are reordered in exactly the same way, with no inherent sense of position. This matters because it shows the mechanism ignores order by itself, which is why positional information must be injected separately. Without added position signals, the model would treat 'dog bites man' and 'man bites dog' identically.

Standard self-attention is O(n²) in both time and memory because every one of the n positions computes a score against every other position, forming an n×n attention matrix. This quadratic cost is why very long sequences become expensive. It motivates research into more efficient 'sparse' or 'linear' attention variants.

Attention allows all positions to be processed in parallel rather than step by step, dramatically speeding up training compared to sequential RNNs. It also creates short, direct paths between any pair of tokens, making it easier to learn long-range dependencies than with convolutions that have limited receptive fields. These advantages let Transformers scale to large datasets effectively.

Token embeddings represent individual words or subwords as vectors, while sentence or document embeddings represent an entire span of text as a single vector. Sentence embeddings are often produced by pooling or combining token embeddings, capturing overall meaning for tasks like semantic search or clustering. They operate at different granularities of language.

Although attention weights show where the model placed focus, high weight does not necessarily mean that information caused the final prediction, since other components and residual paths also shape the output. Research has shown that different attention distributions can yield the same predictions, so attention is suggestive but not a definitive explanation. They should be interpreted with caution rather than as ground-truth reasoning.

Subword tokenization breaks words into smaller frequent units, like 'playing' into 'play' and 'ing', so the vocabulary stays manageable while still covering rare or new words. Each subword unit gets its own embedding, and a word's representation is built from its pieces. This lets models handle out-of-vocabulary words gracefully instead of mapping them all to a single 'unknown' token.

In the decoder's cross-attention layer, the Queries come from the decoder's own current representations, while the Keys and Values come from the encoder's output. This lets each decoding step look back at the encoded source sequence and pull in relevant information. It is the bridge that connects the input being processed to the output being generated, as in machine translation.

Embeddings live in high-dimensional spaces (hundreds of dimensions) that humans cannot directly see, so dimensionality-reduction methods like PCA or t-SNE project them down to 2D or 3D. These plots reveal clusters of related words and help us inspect whether the embedding captured meaningful structure. The visualization is approximate, since some information is inevitably lost in the reduction.

If the training text contains social biases, the embeddings can absorb and reproduce them, for example associating certain professions more strongly with one gender. Because downstream systems use these vectors, the bias can propagate into hiring tools, search, or recommendations and cause unfair outcomes. This makes auditing and mitigating bias in embeddings an important ethical concern.

Each input embedding is multiplied by separate learned matrices (W_Q, W_K, W_V) to project it into Query, Key, and Value spaces tailored to the attention task. These projections let the model decide what aspects of each token matter for asking, matching, and carrying information. Because they are learned, the model can adapt how it forms relevance comparisons during training.

GloVe learns embeddings by factorizing a global word-word co-occurrence count matrix, directly using corpus-wide statistics about how often words appear together. Word2Vec instead learns from local prediction tasks over sliding context windows. Both aim to place similar words nearby, but GloVe emphasizes global counts while Word2Vec emphasizes local context prediction.

Setting a score to negative infinity ensures that after softmax its exponential becomes effectively zero, giving that position exactly zero attention weight. Setting the raw score to zero instead would still produce a nonzero softmax weight (e⁰ = 1), so the position would leak some influence. Negative infinity is the clean way to fully exclude positions like future tokens.

Using pre-trained embeddings means starting with vectors already learned on a large general corpus and reusing them in a new, often smaller, task. This transfers general knowledge about word meanings so the new model needs less data and trains faster than learning embeddings from scratch. It is a common and effective form of transfer learning in NLP.