Multi-Head Attention and Transformer Block

hard · transformers, multi-head, layer-norm

Multi-Head Attention and Transformer Block

Implement multi-head attention and a pre-norm transformer block.

Components to implement

1) `MultiHeadAttention(Module)`

class MultiHeadAttention(Module):
    def __init__(self, embed_dim: int, num_heads: int):
        # embed_dim must be divisible by num_heads
        pass

    def forward(self, x: List[List[Value]], causal: bool = True) -> List[List[Value]]:
        # x: (T, D)
        # returns: (T, D)
        pass

Rules:

Split the embedding dimension into num_heads heads of size head_dim = embed_dim // num_heads.
Run self-attention independently per head.
Concatenate head outputs and apply a final output projection.

2) `LayerNorm(Module)`

Normalize a single vector of length D and apply learnable scale/shift.

3) `FeedForward(Module)`

Two-layer MLP applied per position:

D -> 4D -> D
Use ReLU activation

4) `TransformerBlock(Module)`

Pre-norm block with residuals:

y = x + MHA(LN(x))
z = y + FFN(LN(y))

Notes

Use list-of-Value math; no NumPy.
num_heads is validated with assert embed_dim % num_heads == 0.
Input/output shapes must match (T, D).

from __future__ import annotations
import math
import random
from typing import List, Iterator, Tuple, Callable, Set

class Value:
    """Scalar Value with autograd support."""

def __init__(self, data: float, _children: Tuple["Value", ...] = (), _op: str = "") -> None:
        self.data = float(data)
        self.grad = 0.0
        self._backward: Callable[[], None] = lambda: None
        self._prev: Set["Value"] = set(_children)
        self._op = _op

def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), "+")
        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward
        return out

def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), "*")
        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward
        return out

def __pow__(self, n):
        out = Value(self.data ** n, (self,), f"**{n}")
        def _backward():
            self.grad += n * (self.data ** (n - 1)) * out.grad
        out._backward = _backward
        return out

def __neg__(self):
        return self * -1

def __sub__(self, other):
        return self + (-other)

def __truediv__(self, other):
        return self * (other ** -1)

def __radd__(self, other):
        return self + other

def __rmul__(self, other):
        return self * other

def exp(self):
        out = Value(math.exp(self.data), (self,), "exp")
        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward
        return out

def log(self):
        out = Value(math.log(max(self.data, 1e-10)), (self,), "log")
        def _backward():
            self.grad += (1.0 / max(self.data, 1e-10)) * out.grad
        out._backward = _backward
        return out

def relu(self):
        out = Value(max(0.0, self.data), (self,), "relu")
        def _backward():
            self.grad += (1.0 if self.data > 0 else 0.0) * out.grad
        out._backward = _backward
        return out

def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1.0
        for node in reversed(topo):
            node._backward()

class Module:
    """Module base class."""

def __init__(self) -> None:
        object.__setattr__(self, "_parameters", {})
        object.__setattr__(self, "_modules", {})
        object.__setattr__(self, "training", True)

def __setattr__(self, name: str, value) -> None:
        if isinstance(value, Value):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        elif isinstance(value, (list, tuple)):
            self._register_list(name, value)
        object.__setattr__(self, name, value)

def _register_list(self, name: str, values) -> None:
        for i, v in enumerate(values):
            if isinstance(v, Value):
                self._parameters[f"{name}.{i}"] = v
            elif isinstance(v, Module):
                self._modules[f"{name}.{i}"] = v
            elif isinstance(v, (list, tuple)):
                self._register_list(f"{name}.{i}", v)

def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

def forward(self, *args, **kwargs):
        raise NotImplementedError()

def parameters(self) -> List[Value]:
        params = list(self._parameters.values())
        for module in self._modules.values():
            params.extend(module.parameters())
        return params

class Linear(Module):
    """Linear layer."""

def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super().__init__()
        k = 1.0 / (in_features ** 0.5)
        self.weight = [
            [Value(random.uniform(-k, k)) for _ in range(in_features)]
            for _ in range(out_features)
        ]
        self.bias = [Value(0.0) for _ in range(out_features)] if bias else None
        self.out_features = out_features

def forward(self, x: List[Value]) -> List[Value]:
        out = []
        for i in range(self.out_features):
            val = sum(w * xi for w, xi in zip(self.weight[i], x))
            if self.bias is not None:
                val = val + self.bias[i]
            out.append(val)
        return out

def softmax(scores: List[Value]) -> List[Value]:
    max_score = max(s.data for s in scores)
    exp_scores = [(s - max_score).exp() for s in scores]
    sum_exp = sum(exp_scores)
    return [e / sum_exp for e in exp_scores]

def attention_scores(Q: List[List[Value]], K: List[List[Value]], d_k: int) -> List[List[Value]]:
    seq_len = len(Q)
    scale = d_k ** 0.5
    scores = []
    for i in range(seq_len):
        row = []
        for j in range(seq_len):
            dot = sum(q * k for q, k in zip(Q[i], K[j]))
            row.append(dot / scale)
        scores.append(row)
    return scores

def causal_mask(seq_len: int) -> List[List[float]]:
    mask = []
    for i in range(seq_len):
        row = []
        for j in range(seq_len):
            row.append(0.0 if j <= i else -1e9)
        mask.append(row)
    return mask

def apply_attention(weights: List[List[Value]], V: List[List[Value]]) -> List[List[Value]]:
    seq_len = len(weights)
    d_v = len(V[0])
    output = []
    for i in range(seq_len):
        attended = []
        for k in range(d_v):
            val = Value(0.0)
            for j in range(seq_len):
                val = val + weights[i][j] * V[j][k]
            attended.append(val)
        output.append(attended)
    return output

class SelfAttention(Module):
    """Single-head self-attention."""

def __init__(self, embed_dim: int) -> None:
        super().__init__()
        self.embed_dim = embed_dim
        self.W_q = Linear(embed_dim, embed_dim)
        self.W_k = Linear(embed_dim, embed_dim)
        self.W_v = Linear(embed_dim, embed_dim)
        self.W_o = Linear(embed_dim, embed_dim)

def forward(self, x: List[List[Value]], causal: bool = True) -> List[List[Value]]:
        seq_len = len(x)
        Q = [self.W_q(pos) for pos in x]
        K = [self.W_k(pos) for pos in x]
        V = [self.W_v(pos) for pos in x]
        scores = attention_scores(Q, K, self.embed_dim)
        if causal:
            mask = causal_mask(seq_len)
            for i in range(seq_len):
                for j in range(seq_len):
                    scores[i][j] = scores[i][j] + mask[i][j]
        weights = [softmax(row) for row in scores]
        attended = apply_attention(weights, V)
        return [self.W_o(pos) for pos in attended]

class MultiHeadAttention(Module):
    """TODO: Implement multi-head attention."""

def __init__(self, embed_dim: int, num_heads: int) -> None:
        raise NotImplementedError()

def forward(self, x: List[List[Value]], causal: bool = True) -> List[List[Value]]:
        raise NotImplementedError()

class LayerNorm(Module):
    """TODO: Implement layer normalization."""

def __init__(self, dim: int, eps: float = 1e-5) -> None:
        raise NotImplementedError()

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class FeedForward(Module):
    """TODO: Implement feed-forward network."""

def __init__(self, embed_dim: int, hidden_dim: int = None) -> None:
        raise NotImplementedError()

def forward(self, x: List[Value]) -> List[Value]:
        raise NotImplementedError()

class TransformerBlock(Module):
    """TODO: Implement pre-norm transformer block."""

def __init__(self, embed_dim: int, num_heads: int) -> None:
        raise NotImplementedError()

def forward(self, x: List[List[Value]]) -> List[List[Value]]:
        raise NotImplementedError()