{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YIlPrhp9bxFY"
      },
      "source": [
        "# Understanding Su-scaled Rotary Position Embeddings (RoPE)\n",
        "\n",
        "The Phi-3-Vision model leverages Su-scaled RoPE, a variant of Rotary Position Embeddings, to effectively extend its context length to an impressive 128K tokens, enabling it to process and reason over longer sequences of text and visual information.\n",
        "\n",
        "```python\n",
        "class Phi3SuScaledRotaryEmbedding(nn.Module):\n",
        "    def __init__(self, dim, config):\n",
        "        self.dim = dim\n",
        "        self.base = config.rope_theta\n",
        "        self.short_factor = config.rope_scaling[\"short_factor\"]\n",
        "        self.long_factor = config.rope_scaling[\"long_factor\"]\n",
        "        self.original_max_position_embeddings = config.original_max_position_embeddings\n",
        "        self.scaling_factor = math.sqrt(1 + math.log(config.max_position_embeddings / config.original_max_position_embeddings) / math.log(config.original_max_position_embeddings))\n",
        "        self.inv_freq = None\n",
        "\n",
        "    def __call__(self, position_ids):\n",
        "        seq_len = position_ids.max() + 1\n",
        "        ext_factors = mx.array(self.long_factor, dtype=mx.float32) if seq_len > self.original_max_position_embeddings else mx.array(self.short_factor, dtype=mx.float32)\n",
        "        inv_freq_shape = mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim\n",
        "        self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)\n",
        "        inv_freq_expanded = mx.repeat(self.inv_freq[None, :, None], position_ids.shape[0], axis=0)\n",
        "        position_ids_expanded = mx.array(position_ids, dtype=mx.float32)[:, None, :]\n",
        "        freqs = mx.matmul(inv_freq_expanded, position_ids_expanded).transpose(0, 2, 1)  \n",
        "        emb = mx.concatenate([freqs, freqs], axis=-1)  \n",
        "        cos = mx.cos(emb) * self.scaling_factor\n",
        "        sin = mx.sin(emb) * self.scaling_factor\n",
        "        return cos, sin\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sNYgH65tQXca"
      },
      "source": [
        "## 1. Inverse Frequencies\n",
        "\n",
        "Each frequency is associated with a specific dimension in the embedding space.\n",
        "\n",
        "## 2. Position IDs:\n",
        "\n",
        "Each token in the sequence is assigned a unique position ID (0, 1, 2, etc.)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 469
        },
        "id": "hs6zKMNMQ2KW",
        "outputId": "aa1e7046-8efc-40dd-8ac9-3b0255520538"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x7d9ef4319180>]"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 20616 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "import numpy as np\n",
        "import math\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "position_ids= np.arange(15)[None]\n",
        "dim = 3072/32 # = hidden_size / num_heads\n",
        "inv_freq_shape = np.arange(0, dim, 2) / dim\n",
        "inv_freq_vanilla = 1.0 / (10000.0**inv_freq_shape)\n",
        "\n",
        "inv_freq = inv_freq_vanilla\n",
        "\n",
        "plt.figure()\n",
        "plt.title('inv_freq = 10000**-d')\n",
        "plt.plot(inv_freq, c='k', alpha = 0.5, label = 'None')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuCkILWAVl4g"
      },
      "source": [
        "## 3. Frequency-Position Interaction:\n",
        "\n",
        "The inverse frequencies are multiplied by the position IDs, creating unique sinusoidal patterns for each position.\n",
        "\n",
        "These patterns are interleaved to produce two separate vectors: one for the sine component and one for the cosine component.\n",
        "\n",
        "These vectors are then added to the original input embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 469
        },
        "id": "JNJtNE-vXbTZ",
        "outputId": "e945ea48-b321-4512-ac82-c35b05bfe579"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<matplotlib.legend.Legend at 0x7d9ef4129e10>"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 74120 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "inv_freq_expanded = np.repeat(inv_freq[None, :, None], position_ids.shape[0], axis=0)\n",
        "position_ids_expanded = position_ids[:, None, :]\n",
        "freqs = np.matmul(inv_freq_expanded, position_ids_expanded).transpose(0, 2, 1)  # (1,15,48)\n",
        "\n",
        "plt.figure()\n",
        "plt.title('freqs = inv_freq @ position')\n",
        "for i, freq_i in enumerate(freqs[0]):\n",
        "  plt.plot(freq_i.flatten(), label = i)\n",
        "plt.legend(title=\"Position IDs\", loc=\"upper right\", frameon=False, fontsize=10)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 469
        },
        "id": "YLHlcGdAVgtJ",
        "outputId": "c69dd301-739b-4263-8567-bbb606eb43c2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<matplotlib.legend.Legend at 0x7d9ee4fab3a0>"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 87856 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "emb = np.concatenate([freqs, freqs], axis=-1)\n",
        "plt.figure()\n",
        "plt.title('emb = concatenate([freqs, freqs])')\n",
        "for i in emb[0]:\n",
        "  plt.plot(i.flatten())\n",
        "for i, emb_i in enumerate(emb[0]):\n",
        "  plt.plot(emb_i.flatten(), label = i)\n",
        "plt.legend(title=\"Position IDs\", loc=\"upper right\", frameon=False, fontsize=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epHp7J7vZi75"
      },
      "source": [
        "Imagine each frequency as a wave with a different wavelength.  Lower frequencies have longer wavelengths and capture broader patterns, while higher frequencies have shorter wavelengths and capture finer details. By combining these waves with the position IDs, you get unique patterns that the model can learn to interpret as positional information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 447
        },
        "id": "vW2qqJRbYiu0",
        "outputId": "b8cf3eb7-e207-4c0f-c6cd-f9734fe238f4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x7d9ee420a920>]"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 53600 characters]",
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "scaling_factor = 1.1902380714238083 # math.sqrt(1 + math.log(max_position_embeddings / original_max_position_embeddings) / math.log(original_max_position_embeddings)) where max_position_embeddings, original_max_position_embeddings = 131072, 4096\n",
        "cos = np.cos(emb) * scaling_factor # (1, 15, 96)\n",
        "sin = np.sin(emb) * scaling_factor\n",
        "plt.figure()\n",
        "plt.subplot(211)\n",
        "plt.plot(cos[0][2], c='g')\n",
        "plt.plot(sin[0][2], c='m')\n",
        "plt.subplot(212)\n",
        "plt.plot(cos[0][-1], c='g')\n",
        "plt.plot(sin[0][-1], c='m')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toH4hF8UZVo8"
      },
      "source": [
        "## 4. Su-scaled RoPE:\n",
        "\n",
        "In Su-scaled RoPE, the base inverse frequencies are multiplied by scaling factors.\n",
        "These scaling factors are designed to adjust the frequencies based on the sequence length, allowing the model to better handle both short and long sequences.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 493
        },
        "id": "zMUst0kJRR8w",
        "outputId": "f40ec9e7-2035-443a-9685-65b40787667d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<matplotlib.legend.Legend at 0x7d9ee4103a90>"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 42272 characters]",
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "short_factors = np.array([1.05, 1.05, 1.05, 1.1, 1.1, 1.1, 1.2500000000000002, 1.2500000000000002, 1.4000000000000004, 1.4500000000000004, 1.5500000000000005, 1.8500000000000008, 1.9000000000000008, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.000000000000001, 2.1000000000000005, 2.1000000000000005, 2.2, 2.3499999999999996, 2.3499999999999996, 2.3499999999999996, 2.3499999999999996, 2.3999999999999995, 2.3999999999999995, 2.6499999999999986, 2.6999999999999984, 2.8999999999999977, 2.9499999999999975, 3.049999999999997, 3.049999999999997, 3.049999999999997 ])\n",
        "long_factors = np.array([ 1.0299999713897705, 1.0499999523162842, 1.0499999523162842, 1.0799999237060547, 1.2299998998641968, 1.2299998998641968, 1.2999999523162842, 1.4499999284744263, 1.5999999046325684, 1.6499998569488525, 1.8999998569488525, 2.859999895095825, 3.68999981880188, 5.419999599456787, 5.489999771118164, 5.489999771118164, 9.09000015258789, 11.579999923706055, 15.65999984741211, 15.769999504089355, 15.789999961853027, 18.360000610351562, 21.989999771118164, 23.079999923706055, 30.009998321533203, 32.35000228881836, 32.590003967285156, 35.56000518798828, 39.95000457763672, 53.840003967285156, 56.20000457763672, 57.95000457763672, 59.29000473022461, 59.77000427246094, 59.920005798339844, 61.190006256103516, 61.96000671386719, 62.50000762939453, 63.3700065612793, 63.48000717163086, 63.48000717163086, 63.66000747680664, 63.850006103515625, 64.08000946044922, 64.760009765625, 64.80001068115234, 64.81001281738281, 64.81001281738281 ])\n",
        "\n",
        "plt.figure()\n",
        "plt.subplot(211)\n",
        "plt.title('Su factors')\n",
        "plt.plot(short_factors, c = 'r', alpha = 0.5, label = 'Short')\n",
        "plt.plot(long_factors, c = 'b', alpha = 0.5, label = 'Long')\n",
        "plt.legend(title=\"Factors\", loc=\"upper right\", frameon=False, fontsize=10)\n",
        "\n",
        "inv_freq_short = 1.0 / (short_factors * 10000.0**inv_freq_shape)\n",
        "inv_freq_long = 1.0 / (long_factors * 10000.0**inv_freq_shape)\n",
        "\n",
        "plt.subplot(212)\n",
        "plt.title('inv_freq')\n",
        "plt.plot(inv_freq_vanilla, c='k', alpha = 0.5, label = 'None')\n",
        "plt.plot(inv_freq_short, c='b', alpha = 0.5, label = 'Short')\n",
        "plt.plot(inv_freq_long, c='r', alpha = 0.5, label = 'Long')\n",
        "plt.xlabel('Dimension', fontsize=12)\n",
        "plt.legend(title=\"Factors\", loc=\"upper right\", frameon=False, fontsize=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fIEEBj7RXyl"
      },
      "source": [
        "### 4.1 Short factor scaled RoPE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 485
        },
        "id": "LOo06nxsRR6M",
        "outputId": "6548b32e-f012-418b-8bd5-4106736565b8"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "<ipython-input-7-4f7465b2260c>:11: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
            "  plt.subplot(211)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 117208 characters]",
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "inv_freq = inv_freq_short\n",
        "\n",
        "inv_freq_expanded = np.repeat(inv_freq[None, :, None], position_ids.shape[0], axis=0)\n",
        "position_ids_expanded = position_ids[:, None, :]\n",
        "freqs = np.matmul(inv_freq_expanded, position_ids_expanded).transpose(0, 2, 1)  # (1,15,48)\n",
        "emb = np.concatenate([freqs, freqs], axis=-1)\n",
        "scaling_factor = 1.1902380714238083 # math.sqrt(1 + math.log(max_position_embeddings / original_max_position_embeddings) / math.log(original_max_position_embeddings)) where max_position_embeddings, original_max_position_embeddings = 131072, 4096\n",
        "cos = np.cos(emb) * scaling_factor # (1, 15, 96)\n",
        "plt.figure()\n",
        "plt.title('sin, cos')\n",
        "plt.subplot(211)\n",
        "for i in cos[0]:\n",
        "  plt.plot(i.flatten())\n",
        "sin = np.sin(emb) * scaling_factor\n",
        "plt.subplot(212)\n",
        "for i in sin[0]:\n",
        "  plt.plot(i.flatten())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0T8BVW23attj"
      },
      "source": [
        "### 4.2 Long factor scaled RoPE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 485
        },
        "id": "_td0ePLlasRd",
        "outputId": "1aa2f384-5e8d-47d4-ae82-c83253a4ec9c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "<ipython-input-8-4391cbc3eff3>:11: MatplotlibDeprecationWarning: Auto-removal of overlapping axes is deprecated since 3.6 and will be removed two minor releases later; explicitly call ax.remove() as needed.\n",
            "  plt.subplot(211)\n"
          ]
        },
        {
          "data": {
            "image/png": "[encoded data, 101752 characters]",
            "text/plain": [
              "<Figure size 640x480 with 2 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "inv_freq = inv_freq_long\n",
        "\n",
        "inv_freq_expanded = np.repeat(inv_freq[None, :, None], position_ids.shape[0], axis=0)\n",
        "position_ids_expanded = position_ids[:, None, :]\n",
        "freqs = np.matmul(inv_freq_expanded, position_ids_expanded).transpose(0, 2, 1)  # (1,15,48)\n",
        "emb = np.concatenate([freqs, freqs], axis=-1)\n",
        "scaling_factor = 1.1902380714238083 # math.sqrt(1 + math.log(max_position_embeddings / original_max_position_embeddings) / math.log(original_max_position_embeddings)) where max_position_embeddings, original_max_position_embeddings = 131072, 4096\n",
        "cos = np.cos(emb) * scaling_factor # (1, 15, 96)\n",
        "plt.figure()\n",
        "plt.title('sin, cos')\n",
        "plt.subplot(211)\n",
        "for i in cos[0]:\n",
        "  plt.plot(i.flatten())\n",
        "sin = np.sin(emb) * scaling_factor\n",
        "plt.subplot(212)\n",
        "for i in sin[0]:\n",
        "  plt.plot(i.flatten())"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
