chapter_attention-mechanisms/self-attention-and-positional-encoding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 0
   },
   "source": [
    "# Self-Attention and Positional Encoding\n",
    ":label:`sec_self-attention-and-positional-encoding`\n",
    "\n",
    "In deep learning,\n",
    "we often use CNNs or RNNs to encode a sequence.\n",
    "Now with attention mechanisms.\n",
    "imagine that we feed a sequence of tokens\n",
    "into attention pooling\n",
    "so that\n",
    "the same set of tokens\n",
    "act as queries, keys, and values.\n",
    "Specifically,\n",
    "each query attends to all the key-value pairs\n",
    "and generates one attention output.\n",
    "Since the queries, keys, and values\n",
    "come from the same place,\n",
    "this performs\n",
    "*self-attention* :cite:`Lin.Feng.Santos.ea.2017,Vaswani.Shazeer.Parmar.ea.2017`, which is also called *intra-attention* :cite:`Cheng.Dong.Lapata.2016,Parikh.Tackstrom.Das.ea.2016,Paulus.Xiong.Socher.2017`.\n",
    "In this section,\n",
    "we will discuss sequence encoding using self-attention,\n",
    "including using additional information for the sequence order.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load ../utils/djl-imports\n",
    "%load ../utils/plot-utils\n",
    "%load ../utils/Functions.java\n",
    "%load ../utils/PlotUtils.java\n",
    "\n",
    "%load ../utils/attention/Chap10Utils.java\n",
    "%load ../utils/attention/DotProductAttention.java\n",
    "%load ../utils/attention/MultiHeadAttention.java\n",
    "%load ../utils/attention/PositionalEncoding.java"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NDManager manager = NDManager.newBaseManager();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 3
   },
   "source": [
    "## Self-Attention\n",
    "\n",
    "Given a sequence of input tokens\n",
    "$\\mathbf{x}_1, \\ldots, \\mathbf{x}_n$ where any $\\mathbf{x}_i \\in \\mathbb{R}^d$ ($1 \\leq i \\leq n$),\n",
    "its self-attention outputs\n",
    "a sequence of the same length\n",
    "$\\mathbf{y}_1, \\ldots, \\mathbf{y}_n$,\n",
    "where\n",
    "\n",
    "$$\\mathbf{y}_i = f(\\mathbf{x}_i, (\\mathbf{x}_1, \\mathbf{x}_1), \\ldots, (\\mathbf{x}_n, \\mathbf{x}_n)) \\in \\mathbb{R}^d$$\n",
    "\n",
    "according to the definition of attention pooling $f$ in\n",
    ":eqref:`eq_attn-pooling`.\n",
    "Using multi-head attention,\n",
    "the following code snippet\n",
    "computes the self-attention of a tensor\n",
    "with shape (batch size, number of time steps or sequence length in tokens, $d$).\n",
    "The output tensor has the same shape.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int numHiddens = 100;\n",
    "int numHeads = 5;\n",
    "MultiHeadAttention attention = new MultiHeadAttention(numHiddens, numHeads, 0.5f, false);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int batchSize = 2;\n",
    "int numQueries = 4;\n",
    "NDArray validLens = manager.create(new float[] {3, 2});\n",
    "NDArray X = manager.ones(new Shape(batchSize, numQueries, numHiddens));\n",
    "ParameterStore ps = new ParameterStore(manager, false);\n",
    "NDList input = new NDList(X, X, X, validLens);\n",
    "attention.initialize(manager, DataType.FLOAT32, input.getShapes());\n",
    "NDList result = attention.forward(ps, input, false);\n",
    "result.get(0).getShape()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 7
   },
   "source": [
    "## Comparing CNNs, RNNs, and Self-Attention\n",
    ":label:`subsec_cnn-rnn-self-attention`\n",
    "\n",
    "Let us\n",
    "compare architectures for mapping\n",
    "a sequence of $n$ tokens\n",
    "to another sequence of equal length,\n",
    "where each input or output token is represented by\n",
    "a $d$-dimensional vector.\n",
    "Specifically,\n",
    "we will consider CNNs, RNNs, and self-attention.\n",
    "We will compare their\n",
    "computational complexity, \n",
    "sequential operations,\n",
    "and maximum path lengths.\n",
    "Note that sequential operations prevent parallel computation,\n",
    "while a shorter path between\n",
    "any combination of sequence positions\n",
    "makes it easier to learn long-range dependencies within the sequence :cite:`Hochreiter.Bengio.Frasconi.ea.2001`.\n",
    "\n",
    "\n",
    "![Comparing CNN (padding tokens are omitted), RNN, and self-attention architectures.](https://raw.githubusercontent.com/d2l-ai/d2l-en/master/img/cnn-rnn-self-attention.svg)\n",
    ":label:`fig_cnn-rnn-self-attention`\n",
    "\n",
    "Consider a convolutional layer whose kernel size is $k$.\n",
    "We will provide more details about sequence processing\n",
    "using CNNs in later chapters.\n",
    "For now,\n",
    "we only need to know that\n",
    "since the sequence length is $n$,\n",
    "the numbers of input and output channels are both $d$,\n",
    "the computational complexity of the convolutional layer is $\\mathcal{O}(knd^2)$.\n",
    "As :numref:`fig_cnn-rnn-self-attention` shows,\n",
    "CNNs are hierarchical  so \n",
    "there are $\\mathcal{O}(1)$ sequential operations\n",
    "and the maximum path length is $\\mathcal{O}(n/k)$.\n",
    "For example, $\\mathbf{x}_1$ and $\\mathbf{x}_5$\n",
    "are within the receptive field of a two-layer CNN\n",
    "with kernel size 3 in :numref:`fig_cnn-rnn-self-attention`.\n",
    "\n",
    "When updating the hidden state of RNNs,\n",
    "multiplication of the $d \\times d$ weight matrix\n",
    "and the $d$-dimensional hidden state has \n",
    "a computational complexity of $\\mathcal{O}(d^2)$.\n",
    "Since the sequence length is $n$,\n",
    "the computational complexity of the recurrent layer\n",
    "is $\\mathcal{O}(nd^2)$.\n",
    "According to :numref:`fig_cnn-rnn-self-attention`,\n",
    "there are $\\mathcal{O}(n)$ sequential operations\n",
    "that cannot be parallelized\n",
    "and the maximum path length is also $\\mathcal{O}(n)$.\n",
    "\n",
    "In self-attention,\n",
    "the queries, keys, and values \n",
    "are all $n \\times d$ matrices.\n",
    "Consider the scaled dot-product attention in\n",
    ":eqref:`eq_softmax_QK_V`,\n",
    "where a $n \\times d$ matrix is multiplied by\n",
    "a $d \\times n$ matrix,\n",
    "then the output $n \\times n$ matrix is multiplied\n",
    "by a $n \\times d$ matrix.\n",
    "As a result,\n",
    "the self-attention\n",
    "has a $\\mathcal{O}(n^2d)$ computational complexity.\n",
    "As we can see in :numref:`fig_cnn-rnn-self-attention`,\n",
    "each token is directly connected\n",
    "to any other token via self-attention.\n",
    "Therefore,\n",
    "computation can be parallel with $\\mathcal{O}(1)$ sequential operations\n",
    "and the maximum path length is also $\\mathcal{O}(1)$.\n",
    "\n",
    "All in all,\n",
    "both CNNs and self-attention enjoy parallel computation\n",
    "and self-attention has the shortest maximum path length.\n",
    "However, the quadratic computational complexity with respect to the sequence length\n",
    "makes self-attention prohibitively slow for very long sequences.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "## Positional Encoding\n",
    ":label:`subsec_positional-encoding`\n",
    "\n",
    "\n",
    "Unlike RNNs that recurrently process\n",
    "tokens of a sequence one by one,\n",
    "self-attention ditches\n",
    "sequential operations in favor of \n",
    "parallel computation.\n",
    "To use the sequence order information,\n",
    "we can inject\n",
    "absolute or relative\n",
    "positional information\n",
    "by adding *positional encoding*\n",
    "to the input representations.\n",
    "Positional encodings can be \n",
    "either learned or fixed.\n",
    "In the following, \n",
    "we describe a fixed positional encoding\n",
    "based on sine and cosine functions :cite:`Vaswani.Shazeer.Parmar.ea.2017`.\n",
    "\n",
    "Suppose that\n",
    "the input representation $\\mathbf{X} \\in \\mathbb{R}^{n \\times d}$ contains the $d$-dimensional embeddings for $n$ tokens of a sequence.\n",
    "The positional encoding outputs\n",
    "$\\mathbf{X} + \\mathbf{P}$\n",
    "using a positional embedding matrix $\\mathbf{P} \\in \\mathbb{R}^{n \\times d}$ of the same shape,\n",
    "whose element on the $i^\\mathrm{th}$ row \n",
    "and the $(2j)^\\mathrm{th}$\n",
    "or the $(2j + 1)^\\mathrm{th}$ column is\n",
    "\n",
    "$$\\begin{aligned} p_{i, 2j} &= \\sin\\left(\\frac{i}{10000^{2j/d}}\\right),\\\\p_{i, 2j+1} &= \\cos\\left(\\frac{i}{10000^{2j/d}}\\right).\\end{aligned}$$\n",
    ":eqlabel:`eq_positional-encoding-def`\n",
    "\n",
    "At first glance,\n",
    "this trigonometric-function\n",
    "design looks weird.\n",
    "Before explanations of this design,\n",
    "let us first implement it in the following `PositionalEncoding` class.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "public class PositionalEncoding extends AbstractBlock {\n",
    "\n",
    "    private Dropout dropout;\n",
    "    public NDArray P;\n",
    "\n",
    "    public PositionalEncoding(int numHiddens, float dropout, int maxLen, NDManager manager) {\n",
    "        this.dropout = Dropout.builder().optRate(dropout).build();\n",
    "        addChildBlock(\"dropout\", this.dropout);\n",
    "\n",
    "        // Create a long enough `P`\n",
    "        P = manager.zeros(new Shape(1, maxLen, numHiddens));\n",
    "        NDArray X =\n",
    "                manager.arange(maxLen)\n",
    "                        .reshape(-1, 1)\n",
    "                        .div(\n",
    "                                manager.create(10000)\n",
    "                                        .pow(manager.arange(0, numHiddens, 2).div(numHiddens)));\n",
    "        P.set(new NDIndex(\":, :, {}::{}\", 0, 2), X.sin());\n",
    "        P.set(new NDIndex(\":, :, {}::{}\", 1, 2), X.cos());\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    protected NDList forwardInternal(\n",
    "            ParameterStore parameterStore,\n",
    "            NDList inputs,\n",
    "            boolean training,\n",
    "            PairList<String, Object> params) {\n",
    "        NDArray X = inputs.get(0);\n",
    "        X = X.add(P.get(\":, :{}, :\", X.getShape().get(1)));\n",
    "        return new NDList(\n",
    "                dropout.forward(parameterStore, new NDList(X), training, params).get(0));\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    public Shape[] getOutputShapes(Shape[] inputShapes) {\n",
    "        throw new UnsupportedOperationException(\"Not implemented\");\n",
    "    }\n",
    "\n",
    "    @Override\n",
    "    public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {\n",
    "        try (NDManager sub = manager.newSubManager()) {\n",
    "            NDArray X = sub.zeros(inputShapes[0], dataType);\n",
    "            X = X.add(P.get(\":, :{}, :\", X.getShape().get(1)));\n",
    "            dropout.initialize(manager, dataType, X.getShape());\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 10
   },
   "source": [
    "In the positional embedding matrix $\\mathbf{P}$,\n",
    "rows correspond to positions within a sequence\n",
    "and columns represent different positional encoding dimensions.\n",
    "In the example below,\n",
    "we can see that\n",
    "the $6^{\\mathrm{th}}$ and the $7^{\\mathrm{th}}$\n",
    "columns of the positional embedding matrix \n",
    "have a higher frequency than \n",
    "the $8^{\\mathrm{th}}$ and the $9^{\\mathrm{th}}$\n",
    "columns.\n",
    "The offset between \n",
    "the $6^{\\mathrm{th}}$ and the $7^{\\mathrm{th}}$ (same for the $8^{\\mathrm{th}}$ and the $9^{\\mathrm{th}}$) columns\n",
    "is due to the alternation of sine and cosine functions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int encodingDim = 32;\n",
    "int numSteps = 60;\n",
    "PositionalEncoding posEncoding = new PositionalEncoding(encodingDim, 0, 1000, manager);\n",
    "input = new NDList(manager.zeros(new Shape(1, numSteps, encodingDim)));\n",
    "X = posEncoding.forward(ps, input, false).get(0);\n",
    "NDArray P = posEncoding.P.get(new NDIndex(\":, :{}, :\", X.getShape().get(1)));\n",
    "\n",
    "double[][] plotX = new double[4][];\n",
    "double[][] plotY = new double[4][];\n",
    "for (int i = 0; i < 4; i++) {\n",
    "    if (i == 0) {\n",
    "        plotX[i] = manager.arange(numSteps).toType(DataType.FLOAT64, false).toDoubleArray();\n",
    "    } else {\n",
    "        plotX[i] = plotX[i - 1];\n",
    "    }\n",
    "    plotY[i] =\n",
    "            Functions.floatToDoubleArray(\n",
    "                    P.get(new NDIndex(\"0, :, {},\", i + 6)).toFloatArray());\n",
    "}\n",
    "\n",
    "\n",
    "PlotUtils.plot(\n",
    "        plotX,\n",
    "        plotY,\n",
    "        new String[] {\"Col6\", \"Col7\", \"Col8\", \"Col9\"},\n",
    "        \"Row (position)\",\n",
    "        \"\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 13
   },
   "source": [
    "### Absolute Positional Information\n",
    "\n",
    "To see how the monotonically decreased frequency\n",
    "along the encoding dimension relates to absolute positional information,\n",
    "let us print out the binary representations of $0, 1, \\ldots, 7$.\n",
    "As we can see,\n",
    "the lowest bit, the second-lowest bit, and the third-lowest bit alternate on every number, every two numbers, and every four numbers, respectively.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for (int i = 0; i < 8; i++) {\n",
    "    System.out.println(i + \" in binary is \" + Integer.toBinaryString(i));\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 15
   },
   "source": [
    "In binary representations,\n",
    "a higher bit has a lower frequency than a lower bit.\n",
    "Similarly,\n",
    "as demonstrated in the heat map below,\n",
    "the positional encoding decreases\n",
    "frequencies along the encoding dimension\n",
    "by using trigonometric functions.\n",
    "Since the outputs are float numbers,\n",
    "such continuous representations\n",
    "are more space-efficient\n",
    "than binary representations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = P.get(new NDIndex(\"0, :, :\")).expandDims(0).expandDims(0);\n",
    "PlotUtils.showHeatmaps(\n",
    "        P, \"Column (encoding dimension)\", \"Row (position)\", new String[] {\"\"}, 500, 700);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "origin_pos": 18
   },
   "source": [
    "### Relative Positional Information\n",
    "\n",
    "Besides capturing absolute positional information,\n",
    "the above positional encoding\n",
    "also allows\n",
    "a model to easily learn to attend by relative positions.\n",
    "This is because\n",
    "for any fixed position offset $\\delta$,\n",
    "the positional encoding at position $i + \\delta$\n",
    "can be represented by a linear projection\n",
    "of that at position $i$.\n",
    "\n",
    "\n",
    "This projection can be explained\n",
    "mathematically.\n",
    "Denoting\n",
    "$\\omega_j = 1/10000^{2j/d}$,\n",
    "any pair of $(p_{i, 2j}, p_{i, 2j+1})$ \n",
    "in :eqref:`eq_positional-encoding-def`\n",
    "can \n",
    "be linearly projected to $(p_{i+\\delta, 2j}, p_{i+\\delta, 2j+1})$\n",
    "for any fixed offset $\\delta$:\n",
    "\n",
    "$$\\begin{aligned}\n",
    "&\\begin{bmatrix} \\cos(\\delta \\omega_j) & \\sin(\\delta \\omega_j) \\\\  -\\sin(\\delta \\omega_j) & \\cos(\\delta \\omega_j) \\\\ \\end{bmatrix}\n",
    "\\begin{bmatrix} p_{i, 2j} \\\\  p_{i, 2j+1} \\\\ \\end{bmatrix}\\\\\n",
    "=&\\begin{bmatrix} \\cos(\\delta \\omega_j) \\sin(i \\omega_j) + \\sin(\\delta \\omega_j) \\cos(i \\omega_j) \\\\  -\\sin(\\delta \\omega_j) \\sin(i \\omega_j) + \\cos(\\delta \\omega_j) \\cos(i \\omega_j) \\\\ \\end{bmatrix}\\\\\n",
    "=&\\begin{bmatrix} \\sin\\left((i+\\delta) \\omega_j\\right) \\\\  \\cos\\left((i+\\delta) \\omega_j\\right) \\\\ \\end{bmatrix}\\\\\n",
    "=& \n",
    "\\begin{bmatrix} p_{i+\\delta, 2j} \\\\  p_{i+\\delta, 2j+1} \\\\ \\end{bmatrix},\n",
    "\\end{aligned}$$\n",
    "\n",
    "where the $2\\times 2$ projection matrix does not depend on any position index $i$.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* In self-attention, the queries, keys, and values all come from the same place.\n",
    "* Both CNNs and self-attention enjoy parallel computation and self-attention has the shortest maximum path length. However, the quadratic computational complexity with respect to the sequence length makes self-attention prohibitively slow for very long sequences.\n",
    "* To use the sequence order information, we can inject absolute or relative positional information by adding positional encoding to the input representations.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Suppose that we design a deep architecture to represent a sequence by stacking self-attention layers with positional encoding. What could be issues?\n",
    "1. Can you design a learnable positional encoding method?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java",
   "language": "java",
   "name": "java"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "Java",
   "pygments_lexer": "java",
   "version": "14.0.2+12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}