Attention Mechanisms & The Transformer

Before 2017, sequence-to-sequence modeling relied on recurrent architectures that compressed variable-length source sentences into a single fixed-size bottleneck vector $h_T$.

Vaswani et al. (2017) eliminated recurrence entirely in "Attention Is All You Need", introducing the Transformer. By replacing sequential unrolling with pairwise Scaled Dot-Product Self-Attention, Transformers enable massive GPU parallelization and direct $O(1)$ path length between any two tokens across time.

1. Scaled Dot-Product Attention Formulation#

Let sequence representations be packed into matrix $X \in \mathbb{R}^{T \times d_{\text{model}}}$ of $T$ tokens.

We project $X$ using three learnable parameter matrices:

• Queries: $Q = X W_Q \in \mathbb{R}^{T \times d_k}$
• Keys: $K = X W_K \in \mathbb{R}^{T \times d_k}$
• Values: $V = X W_V \in \mathbb{R}^{T \times d_v}$

The canonical Scaled Dot-Product Attention equation is:

Where $\text{softmax}(A)_{i,j} = \frac{\exp(A_{i,j})}{\sum_{m=1}^T \exp(A_{i,m})}$.

2. Mathematical Proof: The $\frac{1}{\sqrt{d_k}}$ Scaling Factor#

Why is the dot-product matrix $Q K^T$ scaled specifically by the inverse square root of key dimension $d_k$?

Consider a query vector $q \in \mathbb{R}^{d_k}$ and key vector $k \in \mathbb{R}^{d_k}$. Assume their scalar components $q_i$ and $k_i$ are independent random variables with zero mean and unit variance:

The dot product is the sum of $d_k$ scalar products:

Scaling by $\frac{1}{\sqrt{d_k}}$ preserves unit variance into the softmax regardless of feature dimensionality $d_k$.

3. Sinusoidal Positional Encoding & Linear Relative Shifts#

Because self-attention computes dot products symmetrically across all pairs, the permutation of input tokens $X$ yields identical attention scores (permutation equivariance).

To supply positional order, Vaswani et al. add fixed sinusoidal encodings $PE \in \mathbb{R}^{T \times d_{\text{model}}}$:

Where $pos \in \{0, \dots, T-1\}$ is token position and $i \in \{0, \dots, d_{\text{model}}/2 - 1\}$ is feature channel.

Mathematical Proof of Relative Linear Shift#

For any fixed offset $k$, can the model attend to relative position $pos + k$ using a linear transformation?

Recall the trigonometric addition formulas:

Let $\omega_i = \frac{1}{10000^{2i / d_{\text{model}}}}$. We can represent $PE_{(pos + k, 2i)}$ and $PE_{(pos + k, 2i+1)}$ as a rotation matrix acting on $PE_{pos}$:

Because the rotation matrix depends solely on offset $k$ and not on absolute position $pos$, the self-attention mechanism can easily learn to attend to relative distances.

4. From-Scratch Implementation: Zero Dependencies#

The following implementations construct the complete Scaled Dot-Product Attention mechanism (including numerically stable softmax and value aggregation) from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import math
from typing import List

class AttentionFromScratch:
    @staticmethod
    def _dot(a: List[float], b: List[float]) -> float:
        return sum(x * y for x, y in zip(a, b))

    @staticmethod
    def _softmax(row: List[float]) -> List[float]:
        # Numerical stability: subtract max(row)
        max_val = max(row)
        exp_vals = [math.exp(x - max_val) for x in row]
        sum_exp = sum(exp_vals)
        return [v / sum_exp for v in exp_vals]

    @classmethod
    def scaled_dot_product_attention(
        cls,
        Q: List[List[float]],
        K: List[List[float]],
        V: List[List[float]]
    ) -> List[List[float]]:
        # Q: [T, d_k], K: [T, d_k], V: [T, d_v]
        T = len(Q)
        d_k = len(Q[0])
        d_v = len(V[0])
        scale = 1.0 / math.sqrt(d_k)

        # 1. Compute Q * K^T / sqrt(d_k)
        scores = [[0.0] * T for _ in range(T)]
        for i in range(T):
            for j in range(T):
                scores[i][j] = cls._dot(Q[i], K[j]) * scale

        # 2. Row-wise Softmax to compute attention weights
        attention_weights = [cls._softmax(row) for row in scores]

        # 3. Multiply weights by V: [T, T] * [T, d_v] -> [T, d_v]
        output = [[0.0] * d_v for _ in range(T)]
        for i in range(T):
            for col in range(d_v):
                output[i][col] = sum(attention_weights[i][j] * V[j][col] for j in range(T))

        return output

if __name__ == "__main__":
    # Sequence of 3 tokens with d_k = 4, d_v = 2
    Q = [
        [1.0, 0.0, 1.0, 0.0],
        [0.0, 2.0, 0.0, 1.0],
        [1.0, 1.0, 1.0, 1.0]
    ]
    K = [
        [1.0, 0.0, 1.0, 0.0],
        [0.0, 2.0, 0.0, 1.0],
        [1.0, 1.0, 1.0, 1.0]
    ]
    V = [
        [10.0, 1.0],
        [20.0, 2.0],
        [30.0, 3.0]
    ]

    out = AttentionFromScratch.scaled_dot_product_attention(Q, K, V)
    print("Scaled Dot-Product Attention Output (3x2):")
    for row in out:
        print(" ", [round(v, 4) for v in row])

5. Active Knowledge Checks#