positional_encoding_torch.ipynb

Kernel: Python 3

Please find jax implementation of this notebook here: https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/book1/15/positional_encoding_jax.ipynb

Posititional encoding for transformers.

We show how to implement positional encoding. Based on sec 10.6 of http://d2l.ai/chapter_attention-mechanisms/self-attention-and-positional-encoding.html.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(seed=1)
import math
import collections

try:
    import torch
except ModuleNotFoundError:
    %pip install -qq torch
    import torch
from torch import nn
from torch.nn import functional as F

!mkdir figures # for saving plots

!wget https://raw.githubusercontent.com/d2l-ai/d2l-en/master/d2l/torch.py -q -O d2l.py
import d2l

class PositionalEncoding(nn.Module):
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # Create a long enough `P`
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1) / torch.pow(
            10000, torch.arange(0, num_hiddens, 2, dtype=torch.float32) / num_hiddens
        )
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, : X.shape[1], :].to(X.device)
        return self.dropout(X)

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, : X.shape[1], :]
d2l.plot(
    torch.arange(num_steps),
    P[0, :, 6:10].T,
    xlabel="Row (position)",
    figsize=(6, 2.5),
    legend=["Col %d" % d for d in torch.arange(6, 10)],
)
plt.savefig("positionalEncodingSinusoids.pdf", dpi=300)

P = P[0, :, :].unsqueeze(0).unsqueeze(0)
d2l.show_heatmaps(P, xlabel="Column (encoding dimension)", ylabel="Row (position)", figsize=(3.5, 4), cmap="Blues")
plt.savefig("positionalEncodingHeatmap.pdf", dpi=300)