From Scratch · C++ · Transformers · Attention Variants · Part 1 of 13

SelfAttention: The Core Pattern

April 2026 · C++ / LibTorch · Transformers

Introduction

Attention is the mechanism that lets each token in a sequence "look at" every other token and decide how much to integrate from each one. Self-attention is the case where all the look directions (queries), lookup keys, and values come from the same input sequence. This is the building block for every attention variant that follows.

In this post, we implement a concrete C++ LibTorch class for self-attention: define three learnable projections (Q, K, V), compute scaled dot-product similarity, apply softmax to get weights, and combine values according to those weights.

Operations Pipeline

Here is the complete computation, from input embeddings to attended output:

x (batch, T, d_model) input tokens |-- Linear Q --> (batch, T, d_k) queries |-- Linear K --> (batch, T, d_k) keys |-- Linear V --> (batch, T, d_v) values | |-- Q @ K^T / sqrt(d_k) --> (batch, T, T) scaled scores |-- softmax (each row) --> (batch, T, T) attention weights |-- @ V --> (batch, T, d_v) output

In many designs, $d_k = d_v = d_{model}$ (or a fraction of it per head). We compute similarity between every query and every key, scale to stabilize gradient flow, convert to probabilities, and use those to mix value vectors.

The C++ Class

struct SelfAttentionImpl : public torch::nn::Module { torch::nn::Linear wq{nullptr}, wk{nullptr}, wv{nullptr}; int64_t d_model, d_k; SelfAttentionImpl(int64_t d_model_) : d_model(d_model_), d_k(d_model_) { wq = register_module("wq", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); wk = register_module("wk", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); wv = register_module("wv", torch::nn::Linear( torch::nn::LinearOptions(d_model, d_k).bias(false) ) ); } torch::Tensor forward(torch::Tensor x) { auto Q = wq(x); // (batch, T, d_k) auto K = wk(x); // (batch, T, d_k) auto V = wv(x); // (batch, T, d_k) auto scores = torch::matmul(Q, K.transpose(-2, -1)); // (batch, T, T) scores = scores / sqrt((float)d_k); // scaled by sqrt(d_k) auto attn_weights = torch::softmax(scores, -1); // (batch, T, T), rows sum to 1 auto output = torch::matmul(attn_weights, V); // (batch, T, d_k) return output; } }; TORCH_MODULE(SelfAttention);

Q, K, V Projections

All three projections start from the same input $x$. No interaction between samples yet; this is purely per-token dimension reduction.

Operation	Input Shape	Output Shape	Interpretation
$Q = W_Q \cdot x$	(batch, T, d_model)	(batch, T, d_k)	What each token is looking for
$K = W_K \cdot x$	(batch, T, d_model)	(batch, T, d_k)	What each token represents
$V = W_V \cdot x$	(batch, T, d_model)	(batch, T, d_k)	What each token contributes if attended

The three weight matrices are independent. This means each token can be "viewed" differently through each lens: what it looks for, how it matches, and what it provides are learned separately.

Scaled Dot-Product Scores

Once we have Q and K, the core operation is:

scores = (Q @ K^T) / sqrt(d_k) Shape: (batch, T, d_k) @ (batch, d_k, T) = (batch, T, T)

Each entry $\text{scores}[i,j]$ is the dot product of token $i$'s query with token $j$'s key, normalized by $\sqrt{d_k}$. Without this scaling, when $d_k$ is large (say, 512), the dot products grow large, and softmax becomes nearly one-hot, killing gradient flow during backprop. Dividing by $\sqrt{d_k}$ brings the average magnitude back down to a reasonable range.

Softmax and Weighted Sum

After scaling, apply softmax to each row independently:

attn_weights[i, :] = softmax(scores[i, :]) Each row becomes a probability distribution over all tokens.

Then multiply by value vectors:

output[i] = sum_j ( attn_weights[i, j] * V[j] ) For each token i, take a weighted sum of all value vectors.

Tokens with high attention weight contribute more to the output; those with low weight contribute less. The softmax ensures weights are positive and sum to 1, making the combination interpretable as a distribution.

Complete Shape Trace

Step	Variable	Shape	Notes
Input	x	(batch, T, d_model)	Batch of token embeddings
Project Q	Q	(batch, T, d_k)	Linear wq applied
Project K	K	(batch, T, d_k)	Linear wk applied
Project V	V	(batch, T, d_k)	Linear wv applied
Transpose K	K^T	(batch, d_k, T)	Last two dims swapped
Scores	scores	(batch, T, T)	Q @ K^T; all (query, key) pairs
Scale	scores	(batch, T, T)	Divide by sqrt(d_k)
Softmax	attn_weights	(batch, T, T)	Each row sums to 1
Output	output	(batch, T, d_k)	Weighted sum of values

Interactive Pipeline Animation

▶ Shape Transformation through Self-Attention

Step through the tensor shapes and operations as input flows through the attention module.

Key Insights

Self-attention links all tokens in one pass. Unlike RNNs that process sequentially, or CNNs with fixed receptive fields, attention allows any token to directly influence any other token with a learned weight. The computational complexity is $O(T^2 \cdot d)$ (where $T$ is sequence length), which is quadratic in $T$ but linear in embedding dimension.

Queries, keys, and values are learned separately. The three weight matrices mean the model can learn what each token should look for (Q), what it should match against (K), and what it should contribute (V). This triple independence is a core reason attention is so expressive.

Scaling by $\sqrt{d_k}$ is not optional. Without it, large embedding dimensions cause the softmax to collapse to one-hot distributions, blocking gradient flow and making the model hard to train. This is a simple but critical empirical finding from the original Transformer paper.

What Comes Next

Self-attention assumes every token can attend to every other token. But in many settings, this is too much: the attention matrix becomes huge, or certain attention patterns shouldn't be allowed. The next posts explore variants that modify or restrict this behavior:

Softmax Attention — break down softmax into its three stable components.
Causal Attention — add a triangular mask to prevent tokens from seeing the future.
Cross-Attention — separate keys/values from queries (encoder-decoder pattern).
Sliding-Window Attention — limit the attention span to a band around each token.
Global Attention — combine sliding-window with a few "global" tokens that see all.
Linear Attention — replace softmax with a positive feature map to reduce complexity from $O(T^2)$ to $O(T \cdot d^2)$.