import plotly.io as pio; pio.renderers.default = "notebook_connected"

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from sklearn.datasets import load_digits

torch.manual_seed(189);

# Berkeley colors and plotly layout settings
BLUE, GOLD, RED, GREEN, PURPLE, GRAY = '#002676', '#C79100', '#D9531E', '#2E8B57', '#6C3BAA', '#7F7F7F'
LAYOUT = dict(template='plotly_white', font=dict(size=13), margin=dict(l=60, r=20, t=60, b=50))
DIVERGING = 'RdBu_r'


def figure(rows=1, cols=1, titles=None, width=800, height=420, **kw):
    """A plotly figure (with subplots) in the notebook's style."""
    fig = make_subplots(rows=rows, cols=cols, subplot_titles=titles, **kw)
    fig.update_layout(width=width, height=height, **LAYOUT)
    return fig


def attention_heatmap(alpha, labels, showscale=True):
    """Heatmap trace of an N x N attention matrix: row n holds the weights alpha_n1, ..., alpha_nN."""
    alpha = np.asarray(alpha)
    text = [[f'{v:.2f}'.lstrip('0') if v >= 0.05 else '' for v in row] for row in alpha]
    return go.Heatmap(z=alpha, x=list(range(len(labels))), y=list(range(len(labels))), zmin=0, zmax=1,
                      colorscale='Blues', text=text, texttemplate='%{text}', textfont=dict(size=9),
                      showscale=showscale, colorbar=dict(title='α_ni'),
                      customdata=[[(labels[n], labels[i]) for i in range(len(labels))] for n in range(len(labels))],
                      hovertemplate='query %{customdata[0]}<br>key %{customdata[1]}<br>α = %{z:.3f}<extra></extra>')


def label_axes(fig, labels, row=1, col=1):
    """Use the words as tick labels (words repeat, so the axes are indexed by position)."""
    ticks = dict(tickmode='array', tickvals=list(range(len(labels))), ticktext=labels)
    fig.update_xaxes(row=row, col=col, tickangle=-60, **ticks)
    fig.update_yaxes(row=row, col=col, autorange='reversed', **ticks)

features = ['water', 'money', 'motion', 'place', 'grammar']
embedding = {          # toy word vectors x_n in R^5, one coordinate per feature
    'I':      [0, 0, 0, 0, 1],
    'swam':   [1, 0, 1, 0, 0],
    'walked': [0, 0, 1, 0, 0],
    'across': [0, 0, 0.5, 0, 1],
    'the':    [0, 0, 0, 0, 1],
    'river':  [2, 0, 0, 1, 0],
    'road':   [0, 0, 0, 1, 0],
    'to':     [0, 0, 0, 0, 1],
    'get':    [0, 0, 0, 0, 1],
    'cash':   [0, 2, 0, 0, 0],
    'from':   [0, 0, 0, 0, 1],
    'other':  [0, 0, 0, 0, 1],
    'bank':   [1, 1, 0, 1, 0],   # ambiguous: both water and money
}

def embed(sentence):
    """Stack the embeddings of the words into X (N x D): one row per token."""
    return torch.tensor([embedding[w] for w in sentence.split()], dtype=torch.float32)

def simple_attention(X):
    """Y = Softmax[X X^T] X (Bishop, Eq. 12.6). Returns Y and the attention weights alpha."""
    alpha = torch.softmax(X @ X.T, dim=1)    # row n: weights alpha_{n1}, ..., alpha_{nN}
    return alpha @ X, alpha

river = 'I swam across the river to get to the other bank'
cash = 'I walked across the road to get cash from the other bank'
Y_river, alpha_river = simple_attention(embed(river))
Y_cash, alpha_cash = simple_attention(embed(cash))
print('rows of alpha sum to 1:', torch.allclose(alpha_river.sum(1), torch.ones(len(river.split()))))

rows of alpha sum to 1: True

fig = figure(1, 2, ['Attention weights: river sentence', 'Attention weights: cash sentence'], width=1100, height=560,
             horizontal_spacing=0.12)
fig.add_trace(attention_heatmap(alpha_river, river.split(), showscale=False), 1, 1)
fig.add_trace(attention_heatmap(alpha_cash, cash.split()), 1, 2)
label_axes(fig, river.split(), 1, 1)
label_axes(fig, cash.split(), 1, 2)
fig.update_yaxes(title_text='token n (query)', col=1)
fig.update_xaxes(title_text='attended token i (key)')
fig.show()

x_bank = embed('bank')[0]
y_bank_river, y_bank_cash = Y_river[-1], Y_cash[-1]
fig = figure(width=800, height=400)
for vec, color, name in [(x_bank, GRAY, 'input x (bank)'), (y_bank_river, BLUE, 'output y (bank), river sentence'),
                         (y_bank_cash, GOLD, 'output y (bank), cash sentence')]:
    fig.add_trace(go.Bar(x=features, y=vec.numpy(), name=name, marker_color=color))
fig.update_layout(barmode='group', title='The output for "bank" depends on its context', yaxis_title='coordinate value')
fig.show()

def attention(Q, K, V):
    """Scaled dot-product attention (Bishop, Eq. 12.14).
    Q, K: N x D_k; V: N x D_v; row n belongs to token n."""
    D_k = Q.shape[1]
    Z = Q @ K.T / math.sqrt(D_k)       # scores z_ni = q_n^T k_i / sqrt(D_k)
    alpha = torch.softmax(Z, dim=1)    # each row sums to 1
    return alpha @ V, alpha            # outputs and weights

g = torch.Generator().manual_seed(0)
N, D, D_k, D_v = 6, 8, 4, 5
X = torch.randn(N, D, generator=g)
W_q, W_k = torch.randn(D, D_k, generator=g), torch.randn(D, D_k, generator=g)
W_v = torch.randn(D, D_v, generator=g)
Q, K, V = X @ W_q, X @ W_k, X @ W_v

Y, alpha = attention(Q, K, V)
Y_torch = F.scaled_dot_product_attention(Q, K, V)
print('output shape:', tuple(Y.shape), '  max |difference| from PyTorch:', (Y - Y_torch).abs().max().item())
print('alpha is not symmetric once W_q != W_k:', not torch.allclose(alpha, alpha.T))

output shape: (6, 5)   max |difference| from PyTorch: 1.430511474609375e-06
alpha is not symmetric once W_q != W_k: True

def softmax_jacobian_norm(alpha):
    """Frobenius norm of d alpha / d z = diag(alpha) - alpha alpha^T for each row of alpha."""
    J = torch.diag_embed(alpha) - alpha[..., :, None] * alpha[..., None, :]
    return J.flatten(-2).norm(dim=-1)

g = torch.Generator().manual_seed(1)
N, trials = 32, 200
D_ks = [2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
stats = {'unscaled': [], 'scaled': []}
for D_k in D_ks:
    Q = torch.randn(trials, N, D_k, generator=g)    # unit-variance queries and keys
    K = torch.randn(trials, N, D_k, generator=g)
    Z = Q @ K.transpose(1, 2)                        # q_n^T k_i for every pair
    for name, scores in [('unscaled', Z), ('scaled', Z / math.sqrt(D_k))]:
        alpha = torch.softmax(scores, dim=-1)
        stats[name].append((scores.var().item(), alpha.max(-1).values.mean().item(),
                            softmax_jacobian_norm(alpha).mean().item()))
stats = {k: np.array(v) for k, v in stats.items()}
for D_k, u, s in zip(D_ks, stats['unscaled'], stats['scaled']):
    if D_k in (4, 64, 1024):
        print(f'D_k = {D_k:4d}: var of scores {u[0]:7.1f} unscaled, {s[0]:.2f} scaled;  '
              f'largest weight {u[1]:.2f} unscaled, {s[1]:.2f} scaled')

D_k =    4: var of scores     4.0 unscaled, 1.01 scaled;  largest weight 0.34 unscaled, 0.16 scaled
D_k =   64: var of scores    63.8 unscaled, 1.00 scaled;  largest weight 0.82 unscaled, 0.16 scaled
D_k = 1024: var of scores  1021.6 unscaled, 1.00 scaled;  largest weight 0.96 unscaled, 0.16 scaled

fig = figure(1, 3, ['Variance of the scores', 'Average largest weight in a row', 'Size of the softmax Jacobian'],
             width=1150, height=400)
for key, color, name in [('unscaled', RED, 'scores QKᵀ'), ('scaled', BLUE, 'scores QKᵀ/√D_k')]:
    for j in range(3):
        fig.add_trace(go.Scatter(x=D_ks, y=stats[key][:, j], mode='lines+markers', line_color=color, name=name,
                                 legendgroup=key, showlegend=(j == 0)), 1, j + 1)
fig.update_xaxes(type='log', title_text='key dimension D_k')
fig.update_yaxes(type='log', title_text='var[z_ni]', col=1)
fig.update_yaxes(title_text='mean largest α_ni in a row', col=2)
fig.update_yaxes(title_text='mean ‖∂α/∂z‖ (Frobenius)', col=3)
fig.show()

def multi_head_attention(X, W_q, W_k, W_v, W_o):
    """Multi-head self-attention (Bishop, Algorithm 12.2).

    X: N x D. W_q, W_k, W_v: lists of H matrices (D x D_k, D x D_k, D x D_v). W_o: H D_v x D.
    """
    H_heads = [attention(X @ W_q[h], X @ W_k[h], X @ W_v[h])[0] for h in range(len(W_q))]
    return torch.cat(H_heads, dim=1) @ W_o      # Concat[H_1, ..., H_H] W^(o)

torch.manual_seed(0)
N, D, H = 7, 16, 4
D_k = D // H
mha = nn.MultiheadAttention(D, H, bias=False, batch_first=True)
X = torch.randn(N, D)

Wq_all, Wk_all, Wv_all = mha.in_proj_weight.detach().T.split(D, dim=1)   # each D x D
W_q = list(Wq_all.split(D_k, dim=1))       # H matrices of size D x D_k
W_k = list(Wk_all.split(D_k, dim=1))
W_v = list(Wv_all.split(D_k, dim=1))
W_o = mha.out_proj.weight.detach().T       # H D_v x D

Y_ours = multi_head_attention(X, W_q, W_k, W_v, W_o)
Y_torch = mha(X[None], X[None], X[None], need_weights=False)[0][0]
print('max |difference| from nn.MultiheadAttention:', (Y_ours - Y_torch).abs().max().item())

max |difference| from nn.MultiheadAttention: 4.470348358154297e-08

def count(module):
    return sum(p.numel() for p in module.parameters())

D, H = 768, 12                             # GPT-2 small
print(f'W_q, W_k, W_v with D_k = D_v = D:         3 D^2 = {3 * D**2:,}')
print(f'nn.MultiheadAttention({D}, {H}, bias=False): {count(nn.MultiheadAttention(D, H, bias=False)):,} = 4 D^2')
N = 1024
print(f'dense layer on all N = {N} tokens:        N^2 D^2 = {(N * D)**2:,}')

Ns = 2 ** np.arange(0, 13)
fig = figure(width=750, height=420)
fig.add_trace(go.Scatter(x=Ns, y=np.full(len(Ns), 3.0 * D**2), line=dict(color=BLUE, width=3), name='self-attention: 3D²'))
fig.add_trace(go.Scatter(x=Ns, y=(Ns * D) ** 2.0, line=dict(color=RED, width=3), name='dense layer: N²D²'))
fig.add_trace(go.Scatter(x=[N, N], y=[3 * D**2, (N * D) ** 2], mode='markers+text', marker=dict(size=10, color=[BLUE, RED]),
                         text=['1.8 million', '618 billion'], textposition='top left', name='GPT-2 small, N = 1024'))
fig.update_xaxes(type='log', title_text='number of tokens N')
fig.update_yaxes(type='log', title_text='parameters', exponentformat='power')
fig.update_layout(title=f'Parameters of one layer (D = {D})')
fig.show()

W_q, W_k, W_v with D_k = D_v = D:         3 D^2 = 1,769,472
nn.MultiheadAttention(768, 12, bias=False): 2,359,296 = 4 D^2
dense layer on all N = 1024 tokens:        N^2 D^2 = 618,475,290,624

class TransformerLayer(nn.Module):
    """One transformer layer (Bishop, Section 12.1.7): self-attention and an MLP, each with a residual
    connection and layer normalization, in the post-norm (Eq. 12.20) or pre-norm (Eq. 12.21) arrangement."""

    def __init__(self, D, H, pre_norm=True):
        super().__init__()
        self.pre_norm = pre_norm
        self.attn = nn.MultiheadAttention(D, H, batch_first=True)
        self.norm1, self.norm2 = nn.LayerNorm(D), nn.LayerNorm(D)
        self.mlp = nn.Sequential(nn.Linear(D, 4 * D), nn.ReLU(), nn.Linear(4 * D, D))

    def Y(self, X):
        """Multi-head self-attention Y(X)."""
        return self.attn(X, X, X, need_weights=False)[0]

    def forward(self, X):
        if self.pre_norm:
            Z = X + self.Y(self.norm1(X))                # Z = X + Y(LayerNorm[X])
            return Z + self.mlp(self.norm2(Z))           # X~ = Z + MLP(LayerNorm[Z])
        Z = self.norm1(self.Y(X) + X)                    # Z = LayerNorm[Y(X) + X]
        return self.norm2(self.mlp(Z) + Z)               # X~ = LayerNorm[MLP[Z] + Z]


layer = TransformerLayer(D=16, H=4)
print('input and output shapes match:', layer(torch.randn(2, 7, 16)).shape == (2, 7, 16))
print(f'parameters for D = 768: {count(TransformerLayer(768, 12)):,} '
      f'(attention {count(nn.MultiheadAttention(768, 12)):,}, MLP {count(TransformerLayer(768, 12).mlp):,})')

input and output shapes match: True
parameters for D = 768: 7,087,872 (attention 2,362,368, MLP 4,722,432)

V_count, N_count = 8, 16

def counting_batch(g, B=32):
    """Random token sequences x (B x N) and targets t_n = number of other tokens equal to x_n."""
    x = torch.randint(0, V_count, (B, N_count), generator=g)
    t = (x[:, :, None] == x[:, None, :]).sum(-1) - 1
    return x, t.clamp(max=V_count - 1)


class TransformerStack(nn.Module):
    """Token embedding, L transformer layers, and a linear output layer for each token."""

    def __init__(self, L=12, D=32, H=4, pre_norm=True):
        super().__init__()
        self.embed = nn.Embedding(V_count, D)
        self.layers = nn.ModuleList([TransformerLayer(D, H, pre_norm) for _ in range(L)])
        self.final_norm = nn.LayerNorm(D) if pre_norm else nn.Identity()
        self.out = nn.Linear(D, V_count)

    def forward(self, x):
        X = self.embed(x)
        for layer in self.layers:
            X = layer(X)
        return self.out(self.final_norm(X))


def train_counting(pre_norm, lr, steps=200, seed=0):
    """Train with Adam at a constant learning rate; return the training error after each step."""
    torch.manual_seed(seed)
    g = torch.Generator().manual_seed(seed)
    model = TransformerStack(pre_norm=pre_norm)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    errors = []
    for _ in range(steps):
        x, t = counting_batch(g)
        E = F.cross_entropy(model(x).flatten(0, 1), t.flatten())
        opt.zero_grad()
        E.backward()
        opt.step()
        errors.append(E.item())
    return np.array(errors)


learning_rates = [3e-4, 1e-3, 3e-3, 1e-2]
runs = {(pre, lr): train_counting(pre, lr) for pre in (False, True) for lr in learning_rates}

# Reference: the error of always predicting the marginal distribution of the counts.
_, t = counting_batch(torch.Generator().manual_seed(1), B=20000)
p = torch.bincount(t.flatten(), minlength=V_count).float() / t.numel()
E_marginal = -(p[p > 0] * p[p > 0].log()).sum().item()
for (pre, lr), E in runs.items():
    print(f'{"pre-norm " if pre else "post-norm"}  lr = {lr:.0e}: error over the last 40 steps {E[-40:].mean():.2f}')
print(f'error of predicting the marginal distribution of counts: {E_marginal:.2f}')

post-norm  lr = 3e-04: error over the last 40 steps 1.37
post-norm  lr = 1e-03: error over the last 40 steps 1.63
post-norm  lr = 3e-03: error over the last 40 steps 1.64
post-norm  lr = 1e-02: error over the last 40 steps 1.64
pre-norm   lr = 3e-04: error over the last 40 steps 0.86
pre-norm   lr = 1e-03: error over the last 40 steps 0.31
pre-norm   lr = 3e-03: error over the last 40 steps 0.46
pre-norm   lr = 1e-02: error over the last 40 steps 0.93
error of predicting the marginal distribution of counts: 1.62

fig = figure(1, 2, ['Final error after 200 Adam steps', 'Training curves at learning rate 1e-3'], width=1100, height=420)
for pre, color, name in [(False, RED, 'post-norm'), (True, BLUE, 'pre-norm')]:
    final = [runs[pre, lr][-40:].mean() for lr in learning_rates]
    fig.add_trace(go.Scatter(x=learning_rates, y=final, mode='lines+markers', line=dict(color=color, width=3),
                             marker_size=9, name=name, legendgroup=name), 1, 1)
    fig.add_trace(go.Scatter(y=runs[pre, 1e-3], line=dict(color=color, width=1.5), name=name, legendgroup=name,
                             showlegend=False), 1, 2)
for col in (1, 2):
    fig.add_hline(y=E_marginal, line_dash='dash', line_color=GRAY, row=1, col=col)
fig.add_annotation(x=np.log10(learning_rates[0]), y=E_marginal, text='predict the marginal of the counts', showarrow=False,
                   yshift=-12, xanchor='left', font_color=GRAY, row=1, col=1)
fig.update_xaxes(type='log', title_text='learning rate (constant, no warmup)', col=1)
fig.update_xaxes(title_text='step τ', col=2)
fig.update_yaxes(title_text='training error (cross-entropy)')
fig.show()

torch.manual_seed(0)
D = 16
words = 'the food was good not bad at all'.split()
word_vectors = {w: torch.randn(D) for w in words}
layer = TransformerLayer(D, H=4).eval()

perm = [0, 1, 2, 5, 4, 3, 6, 7]                       # swap positions 3 ("good") and 5 ("bad")
X = torch.stack([word_vectors[w] for w in words])     # N x D
X_perm = X[perm]                                      # P X
with torch.no_grad():
    Y = layer(X[None])[0]
    Y_perm = layer(X_perm[None])[0]
print('swapped sentence:', ' '.join(words[i] for i in perm))
print('max |layer(P X) - P layer(X)|:', (Y_perm - Y[perm]).abs().max().item())

swapped sentence: the food was bad not good at all
max |layer(P X) - P layer(X)|: 2.384185791015625e-07

r = 0.5 * torch.randn(len(words), D)                  # one vector r_n per position n
with torch.no_grad():
    Y_pos = layer((X + r)[None])[0]
    Y_perm_pos = layer((X_perm + r)[None])[0]
print('with positional encodings, max |layer(P X + R) - P layer(X + R)|:',
      round((Y_perm_pos - Y_pos[perm]).abs().max().item(), 3))

i_good = words.index('good')
j_good = perm.index(i_good)                           # position of "good" in the swapped sentence
fig = figure(1, 2, ['No positional encoding', 'With positional encodings'], width=1100, height=400, shared_yaxes=True)
for col, (a, b) in [(1, (Y[i_good], Y_perm[j_good])), (2, (Y_pos[i_good], Y_perm_pos[j_good]))]:
    fig.add_trace(go.Scatter(y=a.numpy(), mode='lines+markers', line=dict(color=BLUE, width=2),
                             name='"good" in "... was good not bad ..."', showlegend=(col == 1)), 1, col)
    fig.add_trace(go.Scatter(y=b.numpy(), mode='lines+markers', line=dict(color=GOLD, width=2, dash='dash'),
                             marker_symbol='square', name='"good" in "... was bad not good ..."', showlegend=(col == 1)), 1, col)
fig.update_xaxes(title_text='output coordinate')
fig.update_yaxes(title_text='output of the layer for "good"', col=1)
fig.update_layout(legend=dict(orientation='h', y=-0.25))
fig.show()

with positional encodings, max |layer(P X + R) - P layer(X + R)|: 1.382

def sinusoidal_encoding(N, D, L):
    """Position encodings r_n (rows) for n = 0, ..., N-1 from Bishop, Eq. 12.25. Returns an N x D matrix."""
    n = torch.arange(N, dtype=torch.float64)[:, None]
    i = torch.arange(D, dtype=torch.float64)[None, :]
    even = torch.sin(n / L ** (i / D))
    odd = torch.cos(n / L ** ((i - 1) / D))
    return torch.where(i % 2 == 0, even, odd)

R = sinusoidal_encoding(N=200, D=100, L=30)          # the setting of Bishop, Figure 12.10(b)
fig = figure(width=750, height=560)
fig.add_trace(go.Heatmap(z=R.numpy(), colorscale=DIVERGING, zmin=-1, zmax=1, colorbar=dict(title='r_ni'),
                         hovertemplate='position n = %{y}<br>dimension i = %{x}<br>r = %{z:.2f}<extra></extra>'))
fig.update_yaxes(autorange='reversed', title_text='position n')
fig.update_xaxes(title_text='embedding dimension i')
fig.update_layout(title='Sinusoidal encodings r_ni (D = 100, L = 30)')
fig.show()

G = R @ R.T                                           # G[n, m] = r_n^T r_m
D_r = R.shape[1]
omega = 30.0 ** (-2 * torch.arange(D_r // 2, dtype=torch.float64) / D_r)   # one frequency per (sin, cos) pair
d = torch.arange(-100, 101, dtype=torch.float64)
formula = torch.cos(d[:, None] * omega[None, :]).sum(1)                   # sum_j cos(omega_j (m - n))

fig = figure(1, 2, ['Inner products rₙᵀrₘ', 'Every pair (n, m) plotted against m − n'], width=1150, height=450, horizontal_spacing=0.17)
fig.add_trace(go.Heatmap(z=G.numpy(), colorscale='Viridis', colorbar=dict(x=0.41, title='rₙᵀrₘ'),
                         hovertemplate='n = %{y}<br>m = %{x}<br>%{z:.2f}<extra></extra>'), 1, 1)
n_idx, m_idx = torch.meshgrid(torch.arange(200), torch.arange(200), indexing='ij')
near = (m_idx - n_idx).abs() <= 100
fig.add_trace(go.Scattergl(x=(m_idx - n_idx)[near].numpy(), y=G[near].numpy(), mode='markers',
                           marker=dict(size=4, color=BLUE, opacity=0.15), name='all pairs (n, m)'), 1, 2)
fig.add_trace(go.Scatter(x=d.numpy(), y=formula.numpy(), line=dict(color=GOLD, width=2), name='Σⱼ cos(ωⱼ (m − n))'), 1, 2)
fig.update_yaxes(autorange='reversed', title_text='position n', col=1)
fig.update_xaxes(title_text='position m', col=1)
fig.update_xaxes(title_text='offset m − n', col=2)
fig.update_yaxes(title_text='rₙᵀrₘ', col=2)
fig.update_layout(legend=dict(orientation='h', y=-0.2))
fig.show()
spread = max((torch.diagonal(G, o).max() - torch.diagonal(G, o).min()).item() for o in range(-199, 200))
print(f'largest spread of r_n^T r_m among pairs with the same offset: {spread:.1e}')

largest spread of r_n^T r_m among pairs with the same offset: 1.5e-13

def rotation_matrix(k, D, L):
    """Block-diagonal R_k with r_{n+k} = R_k r_n for the encoding of Bishop, Eq. 12.25."""
    R_k = torch.zeros(D, D, dtype=torch.float64)
    for j in range(D // 2):
        omega = L ** (-2 * j / D)                       # frequency of the pair (r_{n,2j}, r_{n,2j+1})
        c, s = math.cos(omega * k), math.sin(omega * k)
        R_k[2 * j:2 * j + 2, 2 * j:2 * j + 2] = torch.tensor([[c, s], [-s, c]], dtype=torch.float64)
    return R_k

N, D, L = 200, 100, 30
R = sinusoidal_encoding(N, D, L)
for k in (1, 5, 37):
    R_k = rotation_matrix(k, D, L)
    err = (R[k:] - R[:-k] @ R_k.T).abs().max().item()   # row n: r_{n+k}^T - (R_k r_n)^T
    print(f'k = {k:2d}: max over all n of |r_(n+k) - R_k r_n| = {err:.1e}')

k =  1: max over all n of |r_(n+k) - R_k r_n| = 2.5e-14
k =  5: max over all n of |r_(n+k) - R_k r_n| = 2.2e-14
k = 37: max over all n of |r_(n+k) - R_k r_n| = 2.1e-14

k, D_fit = 5, 16
R16 = sinusoidal_encoding(N, D_fit, L)
M_fit = torch.linalg.lstsq(R16[:-k], R16[k:]).solution.T   # minimize sum_n ||M r_n - r_{n+k}||^2
print('max |M - R_k|:', f'{(M_fit - rotation_matrix(k, D_fit, L)).abs().max().item():.1e}')

fig = figure(1, 2, ['R₅ from the formula (D = 16)', 'M fitted by least squares (D = 16)'], width=900, height=440)
for col, A in [(1, rotation_matrix(k, D_fit, L)), (2, M_fit)]:
    fig.add_trace(go.Heatmap(z=A.numpy(), colorscale=DIVERGING, zmin=-1, zmax=1, showscale=(col == 2),
                             hovertemplate='row %{y}, column %{x}: %{z:.3f}<extra></extra>'), 1, col)
fig.update_yaxes(autorange='reversed', title_text='row')
fig.update_xaxes(title_text='column')
fig.show()

max |M - R_k|: 7.8e-16

def rope(X, positions=None, L=10000.0):
    """Rotate each row x_n of X (N x D_k): pair i of row n is rotated by the angle n theta_i,
    theta_i = L^(-2i/D_k). positions defaults to n = 0, ..., N-1."""
    N, D_k = X.shape
    n = torch.arange(N, dtype=X.dtype) if positions is None else positions.to(X.dtype)
    theta = L ** (-2 * torch.arange(D_k // 2, dtype=X.dtype) / D_k)
    angle = n[:, None] * theta[None, :]                  # N x D_k/2 angles n theta_i
    cos, sin = torch.cos(angle), torch.sin(angle)
    x1, x2 = X[:, 0::2], X[:, 1::2]                      # the two coordinates of each pair
    out = torch.empty_like(X)
    out[:, 0::2] = cos * x1 - sin * x2
    out[:, 1::2] = sin * x1 + cos * x2
    return out


def rope_rotation(n, D_k, L=10000.0):
    """The block-diagonal matrix R_n that rope applies at position n."""
    R_n = torch.zeros(D_k, D_k, dtype=torch.float64)
    for i in range(D_k // 2):
        c, s = math.cos(n * L ** (-2 * i / D_k)), math.sin(n * L ** (-2 * i / D_k))
        R_n[2 * i:2 * i + 2, 2 * i:2 * i + 2] = torch.tensor([[c, -s], [s, c]], dtype=torch.float64)
    return R_n

g = torch.Generator().manual_seed(0)
N, D_k = 64, 32
q = torch.randn(D_k, generator=g, dtype=torch.float64)
k = torch.randn(D_k, generator=g, dtype=torch.float64)
Q_rot = rope(q.repeat(N, 1))                              # row n: R_n q
K_rot = rope(k.repeat(N, 1))                              # row m: R_m k
S = Q_rot @ K_rot.T                                       # s_nm = (R_n q)^T (R_m k)

offsets = np.arange(-(N - 1), N)
by_offset = [torch.diagonal(S, int(d)) for d in offsets]  # entries with m - n = d
print('largest spread of s_nm among entries with the same m - n:',
      f'{max((s.max() - s.min()).item() for s in by_offset):.1e}')
print('max |s_nm - q^T R_(m-n) k| over 20 pairs:',
      f'{max(abs(S[n, m] - q @ rope_rotation(m - n, D_k) @ k).item() for n, m in np.random.default_rng(0).integers(0, N, (20, 2))):.1e}')
print('norms preserved:', torch.allclose(Q_rot.norm(dim=1), q.norm().expand(N)))

largest spread of s_nm among entries with the same m - n: 2.0e-14
max |s_nm - q^T R_(m-n) k| over 20 pairs: 5.3e-15
norms preserved: True

smax = S.abs().max().item()
fig = figure(1, 2, ['Scores (Rₙq)ᵀ(Rₘk), same q and k at every position', 'Every pair (n, m) plotted against m − n'],
             width=1150, height=450, horizontal_spacing=0.17)
fig.add_trace(go.Heatmap(z=S.numpy(), colorscale=DIVERGING, zmin=-smax, zmax=smax, colorbar=dict(x=0.41, title='s_nm'),
                         hovertemplate='n = %{y}<br>m = %{x}<br>%{z:.2f}<extra></extra>'), 1, 1)
n_idx, m_idx = torch.meshgrid(torch.arange(N), torch.arange(N), indexing='ij')
fig.add_trace(go.Scatter(x=(m_idx - n_idx).flatten().numpy(), y=S.flatten().numpy(), mode='markers',
                         marker=dict(size=5, color=PURPLE, opacity=0.25), name=f'all {N * N:,} pairs (n, m)'), 1, 2)
fig.add_trace(go.Scatter(x=offsets, y=[(q @ rope_rotation(int(o), D_k) @ k).item() for o in offsets],
                         line=dict(color=GOLD, width=2), name='qᵀ R_(m−n) k'), 1, 2)
fig.update_yaxes(autorange='reversed', title_text='query position n', col=1)
fig.update_xaxes(title_text='key position m', col=1)
fig.update_xaxes(title_text='offset m − n', col=2)
fig.update_yaxes(title_text='score', col=2)
fig.show()

theta = 10000.0 ** (-2 * np.arange(D_k // 2) / D_k)
d_all = np.arange(1, 10001)                                  # every offset from 1 to 10,000
expected_score = lambda d: 2 * np.cos(np.asarray(d)[..., None] * theta).sum(-1)
expected = expected_score(d_all)
window = np.maximum(1, d_all // 10)                          # running average over the nearest ~10% of offsets
csum = np.concatenate([[0.0], np.cumsum(expected)])
lo, hi = np.maximum(0, np.arange(len(d_all)) - window // 2), np.minimum(len(d_all), np.arange(len(d_all)) + window // 2 + 1)
trend = (csum[hi] - csum[lo]) / (hi - lo)
fig = figure(width=850, height=420)
fig.add_trace(go.Scattergl(x=d_all, y=expected, mode='lines', line=dict(color=BLUE, width=0.6), opacity=0.35,
                           name='exact: 2 Σᵢ cos(d θᵢ)'))
fig.add_trace(go.Scatter(x=d_all, y=trend, line=dict(color=GOLD, width=3), name='running average'))
fig.add_hline(y=0, line_color=GRAY, line_width=1)
fig.update_xaxes(type='log', title_text='offset d = m − n')
fig.update_yaxes(title_text='expected score E[qᵀ R_d q]')
fig.update_layout(title=f'RoPE with D_k = {D_k}, L = 10000: expected score when k = q')
fig.show()
print('expected score at offsets 0, 10, 100, 1000, 10000:', [round(float(expected_score(d)), 1) for d in (0, 10, 100, 1000, 10000)])

expected score at offsets 0, 10, 100, 1000, 10000: [32.0, 20.1, 18.7, 8.2, 3.1]

Q, K = torch.randn(N, D_k, generator=g, dtype=torch.float64), torch.randn(N, D_k, generator=g, dtype=torch.float64)
positions = torch.arange(N)
S0 = rope(Q, positions) @ rope(K, positions).T
S1000 = rope(Q, positions + 1000) @ rope(K, positions + 1000).T
print('max |score change| after shifting all positions by 1000:', f'{(S0 - S1000).abs().max().item():.1e}')

max |score change| after shifting all positions by 1000: 6.5e-13

digits = load_digits()
images = torch.tensor(digits.images, dtype=torch.float32) / 16     # pixel values in [0, 1]
labels = torch.tensor(digits.target)
order = torch.randperm(len(images), generator=torch.Generator().manual_seed(0))
train_idx, test_idx = order[:1300], order[1300:]

def patchify(images, P):
    """Split B x S x S images into non-overlapping P x P patches: returns B x (S/P)^2 x P^2, row-major order."""
    B, S, _ = images.shape
    patches = images.reshape(B, S // P, P, S // P, P).permute(0, 1, 3, 2, 4)
    return patches.reshape(B, (S // P) ** 2, P * P)

print('tokens per image:', tuple(patchify(images[:1], 2).shape[1:]))

tokens per image: (16, 4)

class TinyViT(nn.Module):
    """Patch embedding, class token, learned positional embeddings, pre-norm transformer layers, linear head."""

    def __init__(self, P=2, S=8, D=32, H=4, L=2, n_classes=10):
        super().__init__()
        self.P = P
        n_patches = (S // P) ** 2
        self.patch_embed = nn.Linear(P * P, D)
        self.class_token = nn.Parameter(torch.zeros(1, 1, D))
        self.r = nn.Parameter(0.02 * torch.randn(1, n_patches + 1, D))   # learned r_n, one per position
        self.layers = nn.ModuleList([TransformerLayer(D, H, pre_norm=True) for _ in range(L)])
        self.final_norm = nn.LayerNorm(D)
        self.head = nn.Linear(D, n_classes)

    def tokens(self, images):
        """The input X (B x (N+1) x D): class token first, then the patch tokens, plus positions."""
        X = self.patch_embed(patchify(images, self.P))
        class_tokens = self.class_token.expand(len(X), -1, -1)
        return torch.cat([class_tokens, X], dim=1) + self.r

    def forward(self, images):
        X = self.tokens(images)
        for layer in self.layers:
            X = layer(X)
        return self.head(self.final_norm(X[:, 0]))      # read the prediction from the class token


torch.manual_seed(0)
vit = TinyViT()
opt = torch.optim.Adam(vit.parameters(), lr=3e-3)
g = torch.Generator().manual_seed(0)
for epoch in range(30):
    for batch in train_idx[torch.randperm(len(train_idx), generator=g)].split(64):
        E = F.cross_entropy(vit(images[batch]), labels[batch])
        opt.zero_grad()
        E.backward()
        opt.step()
vit.eval()
with torch.no_grad():
    test_accuracy = (vit(images[test_idx]).argmax(1) == labels[test_idx]).float().mean().item()
print(f'parameters: {count(vit):,}   test accuracy on {len(test_idx)} held-out digits: {test_accuracy:.3f}')

parameters: 26,538   test accuracy on 497 held-out digits: 0.954

def class_token_attention(model, images):
    """Weights alpha_{0i} of the class token over the patch tokens in the last layer (averaged over heads)."""
    with torch.no_grad():
        X = model.tokens(images)
        for layer in model.layers[:-1]:
            X = layer(X)
        last = model.layers[-1]
        Xn = last.norm1(X)                                # pre-norm: attention sees LayerNorm[X]
        alpha = last.attn(Xn, Xn, Xn, need_weights=True, average_attn_weights=True)[1]
    return alpha[:, 0, 1:]                                # row of the class token, patch columns

examples = test_idx[:6]
alpha0 = class_token_attention(vit, images[examples])
with torch.no_grad():
    predictions = vit(images[examples]).argmax(1)
fig = figure(2, 6, [f'label {labels[e].item()}, predicted {p.item()}' for e, p in zip(examples, predictions)]
             + ['class-token attention'] * 6, width=1150, height=430, vertical_spacing=0.12, horizontal_spacing=0.02)
for j in range(6):
    fig.add_trace(go.Heatmap(z=images[examples[j]].numpy(), colorscale='Greys', showscale=False,
                             hovertemplate='pixel (%{y}, %{x}): %{z:.2f}<extra></extra>'), 1, j + 1)
    for c in (1.5, 3.5, 5.5):                             # patch boundaries
        fig.add_hline(y=c, line_color=GOLD, line_width=1, row=1, col=j + 1)
        fig.add_vline(x=c, line_color=GOLD, line_width=1, row=1, col=j + 1)
    fig.add_trace(go.Heatmap(z=alpha0[j].reshape(4, 4).numpy(), colorscale='Blues', zmin=0, showscale=False,
                             hovertemplate='patch (%{y}, %{x}): α = %{z:.2f}<extra></extra>'), 2, j + 1)
fig.update_xaxes(showticklabels=False, showgrid=False, constrain='domain')
fig.update_yaxes(showticklabels=False, showgrid=False, autorange='reversed', constrain='domain')
for r in (1, 2):
    for c in range(1, 7):
        idx = (r - 1) * 6 + c
        fig.update_yaxes(scaleanchor='x' if idx == 1 else f'x{idx}', row=r, col=c)   # square pixels in every panel
fig.update_annotations(font_size=11)
fig.show()

Notebook section	Slides
Attention as Weighted Averaging	Need to Include Context in Representations; Attention: Weighted Averaging of Context
(Lecture Demo 1) Keys, Queries, and Values	Constructing Keys and Queries through The Self Attention Layer; Demo: Self Attention in PyTorch
Why Divide by $\sqrt{D_k}$?	Scaled Softmax Inner-Product Self-Attention
Multi-Head Attention	Multi-Headed Attention; Parameters and Computation
The Transformer Layer: Post-Norm and Pre-Norm	The Transformer Layer; Post-Norm vs. Pre-Norm; Why Pre-Norm?
Permutation Equivariance	Transformers and Token Ordering; Token Position in Images and Text
Sinusoidal Positional Encoding	Sinusoidal Positional Encoding; Querying a Relative Position; The Rotation Matrix
Rotary Position Embeddings (RoPE)	Rotary Position Embeddings (RoPE); Each Pair Rotates at Its Own Speed; RoPE in Practice
A Tiny Vision Transformer	Vision Transformers; Reading the Output of the Transformer; The Encoder Transformer

Lecture 17: Transformers (CS 189, Fall 2026)

Plotting Code

Attention as Weighted Averaging

(Lecture Demo 1) Keys, Queries, and Values

Why Divide by $\sqrt{D_k}$?

Multi-Head Attention

Parameters: attention versus a dense layer

The Transformer Layer: Post-Norm and Pre-Norm

Pre-norm trains at larger learning rates

Permutation Equivariance

Sinusoidal Positional Encoding

Inner products depend only on the offset

A fixed offset is a rotation

Rotary Position Embeddings (RoPE)

A Tiny Vision Transformer