annotated_hyena.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df886051-bc38-43d1-8cb7-4f9b605325cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "print(sys.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc28dd4b-c24a-4e4b-9539-6950d05c7d0b",
   "metadata": {},
   "source": [
    "# The Annotated Hyena\n",
    "\n",
    "A didactic annotation of the architecture introduced in the paper _Hyena Hierarchy: Towards Larger Convolutional Language Models_ by Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon and Christopher Ré.\n",
    "\n",
    "The Hyena architecture is an exciting development that acts as a drop-in replacement for attention layers in transformer models enabling long-range sequence modeling. It matches and surpasses the language-modeling performance of attention-based transformers with similar parameter counts and can support context lengths over a 100K tokens with a speedup of 100x over the FlashAttention transformer already at a context length of 64K. Here we walk through its construction with annotated code in a style inspired by _[The Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/)_. The paper's [reference implementation](https://github.com/HazyResearch/safari/blob/main/src/models/sequence/hyena.py) was used to debug the below code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "401b95c9-f8d7-46e2-b850-4c8274345894",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "* Preliminaries\n",
    "* Part 1: The Hyena Operator\n",
    "  * Operator definition\n",
    "  * FFT\n",
    "  * Interpretation as attention\n",
    "* Part 2: Defining the Hyena filter\n",
    "* Part 3: A Working Model\n",
    "  * Dataset\n",
    "  * Training\n",
    "  * Evaluation\n",
    "* Appendix: Comparison with Attention\n",
    "* Appendix: Hyena as described in the paper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b9c01f-156f-4073-b648-aaafcf58ef13",
   "metadata": {},
   "source": [
    "## Preliminaries\n",
    "\n",
    "The below code imports the necessary packages and defines some utility functions. Understanding it is not important for understanding hyena-based models. If you are using Jupyter, you will want to begin by installing the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e710b89e-e689-476e-b72f-3fef36563dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets lightning numpy regex torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3506f01-65d3-480e-ac14-447ec4ff09c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d00001a0-8108-48dd-a492-b1586fa333a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict, List, Optional, Tuple, Type, Union\n",
    "\n",
    "import dataclasses\n",
    "import json\n",
    "import math\n",
    "import random\n",
    "\n",
    "import datasets\n",
    "import lightning\n",
    "import numpy as np\n",
    "import regex\n",
    "import torch\n",
    "import wandb\n",
    "\n",
    "\n",
    "def prettify(z: Union[complex, List[complex], np.array]) -> str:\n",
    "  def fix_float(f: float) -> Union[float, int]:\n",
    "    if abs(round(f) - f) < 1e-6:\n",
    "      return round(f)\n",
    "    return f\n",
    "\n",
    "  if isinstance(z, complex) or isinstance(z, float) or isinstance(z, int):\n",
    "    re = fix_float(z.real)\n",
    "    im = fix_float(z.imag)\n",
    "    if im == 1:\n",
    "      im_str = \"i\"\n",
    "    elif im == -1:\n",
    "      im_str = \"-i\"\n",
    "    else:\n",
    "      im_str = f\"{im}i\" if isinstance(im, int) else f\"{im:.3f}i\"\n",
    "  \n",
    "    re_str = f\"{re}\" if isinstance(re, int) else f\"{re:.3f}\"\n",
    "    \n",
    "    if not re:\n",
    "      if not im:\n",
    "        return \"0\"\n",
    "      return im_str\n",
    "    elif not im:\n",
    "      return re_str\n",
    "    else:\n",
    "      return f\"{re_str} + {im_str}\"\n",
    "  elif isinstance(z, list):\n",
    "    return prettify(np.array(z))\n",
    "  else:\n",
    "    if len(z.shape) == 1:\n",
    "      return \"  \".join([prettify(a) for a in z.tolist()])\n",
    "    elif len(z.shape) == 2:\n",
    "      return \"\\n\".join([\n",
    "        \"\\t\".join([prettify(a) for a in row])\n",
    "        for row in z.tolist()\n",
    "      ])\n",
    "    else:\n",
    "      raise NotImplementedError(\"3rd and higher order tensors unsupported\")\n",
    "\n",
    "def average_probability(\n",
    "    model: lightning.LightningModule,\n",
    "    ds: datasets.Dataset,\n",
    "    tokens: int = 256\n",
    "  ) -> float:\n",
    "  \"\"\"Average probability of predicting the next token in the dataset\"\"\"\n",
    "  model.eval()\n",
    "  model = model.to(\"cpu\")\n",
    "  total_p = 0\n",
    "  row = -1\n",
    "  context_length = len(ds[0][\"curr_id\"])\n",
    "  for t in range(tokens):\n",
    "    L = t % context_length\n",
    "    if L == 0:\n",
    "      row += 1\n",
    "      row_ids = ds[row][\"curr_id\"]\n",
    "    input_ids = torch.LongTensor([row_ids[:L + 1]])\n",
    "    logits, _ = model.forward((input_ids, None))\n",
    "    expected_id = ds[row][\"next_id\"][L]\n",
    "    p = torch.softmax(logits[0, L, :], dim=-1)[expected_id].item()\n",
    "    total_p += p\n",
    "  return total_p / tokens\n",
    "  \n",
    "def generate(\n",
    "    model: lightning.LightningModule,\n",
    "    context_length: int,\n",
    "    prompt: str,\n",
    "    max_tokens: int = 32,\n",
    "    method: str = \"topk\",\n",
    "    k: int = 3\n",
    "  ) -> str:\n",
    "  model.eval()\n",
    "  model = model.to(\"cpu\")\n",
    "  result = prompt\n",
    "  prompt = prompt[:context_length]\n",
    "  L = len(prompt)\n",
    "  for _ in range(max_tokens):\n",
    "    input_ids = torch.LongTensor([[tok2id[ch] for ch in prompt[:L]]])\n",
    "    logits, _ = model.forward((input_ids, None))\n",
    "    match method:\n",
    "      case \"sample\":\n",
    "        output_id = torch.multinomial(\n",
    "          torch.exp(logits[0, L - 1, :]),\n",
    "          num_samples=1\n",
    "        ).item()\n",
    "      case \"topk\":\n",
    "        values, indices = torch.topk(torch.exp(logits[0, L - 1, :]), k=k)\n",
    "        p = torch.sparse_coo_tensor(\n",
    "          indices.unsqueeze(0),\n",
    "          values,\n",
    "          (logits.shape[-1],)\n",
    "        )\n",
    "        output_id = torch.multinomial(p.to_dense(), num_samples=1).item()\n",
    "      case \"greedy\":\n",
    "        output_id = torch.argmax(logits[0, L - 1, :]).item()\n",
    "    prompt += vocabulary[output_id]\n",
    "    result += vocabulary[output_id]\n",
    "    L = min(L + 1, context_length)\n",
    "  return result\n",
    "\n",
    "@dataclasses.dataclass(kw_only=True)\n",
    "class Config:\n",
    "  learning_rate: float\n",
    "  epochs: int\n",
    "  betas: Tuple[float, float]\n",
    "  weight_decay: float\n",
    "  device_type: str\n",
    "  precision: str\n",
    "  batch_size: int\n",
    "  num_workers: int\n",
    "\n",
    "\n",
    "@dataclasses.dataclass(kw_only=True)\n",
    "class HyenaConfig(Config):\n",
    "  d_model: int\n",
    "  n_layers: int\n",
    "  vocab_size: int\n",
    "  d_embed: int\n",
    "  d_filter_mlp: int\n",
    "  n_filter_layers: int\n",
    "  context_length: int\n",
    "  short_conv_size: int\n",
    "  order: int\n",
    "  pdrop_hyena: float\n",
    "  pdrop_embed: float\n",
    "  omega: Optional[int]\n",
    "\n",
    "\n",
    "@dataclasses.dataclass(kw_only=True)\n",
    "class AttentionConfig(Config):\n",
    "  d_model: int\n",
    "  n_layers: int\n",
    "  vocab_size: int\n",
    "  d_embed: int\n",
    "  n_head: int\n",
    "  context_length: int\n",
    "  pdrop_attn: float\n",
    "  pdrop_embed: float \n",
    "\n",
    "\n",
    "torch.set_float32_matmul_precision(\"medium\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5318ed43-3340-4b79-8127-dd8c265cd40b",
   "metadata": {},
   "source": [
    "## Part 1: The Hyena Operator\n",
    "\n",
    "The Hyena Operator alternates between convolutions and Hadamard (elementwise) products. The convolutions are with filters (vectors) which will be defined later. The Hadamard products are with functions of the input. In other words, the Hyena Operator of order N is\n",
    "$$\n",
    "H_N(u) = x_N(u) \\cdot (h_N \\ast (x_{N - 1}(u) \\cdot (h_{N - 1} \\ast (\\cdots (h_1 \\ast x_0(u)) \\cdots )))) \\qquad x_i(u) = w_i \\star (A_i u)\n",
    "$$\n",
    "where $u$ is the input vector, $h_i$ are vectors called filters, and $x_i$ is a matrix multiplication followed by a padded cross-correlation known in machine learning as a \"depthwise convolution\". Here $\\star$ represents cross-correlation and $\\ast$ represents convolution.\n",
    "\n",
    "If you are already lost, don't worry! The parts of this definition will be explained below.\n",
    "\n",
    "### Hadamard product\n",
    "\n",
    "To begin, the Hadamard product is just the elementwise product of two vectors. In other words, if $a$ and $b$ are vectors, then\n",
    "$$\n",
    "a \\cdot b = (a_1 b_1, \\ldots, a_n b_n) \\qquad a = (a_1, \\ldots, a_n) \\quad b = (b_1, \\ldots, b_n)\n",
    "$$\n",
    "For those who prefer code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd308c33-2daf-4cc1-96bb-ffd7019697aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hadamard(a: List[complex], b: List[complex]) -> List[complex]:\n",
    "  assert len(a) == len(b)\n",
    "  return [a[i] * b[i] for i in range(len(a))]\n",
    "\n",
    "hadamard([0, 1, 2, 3], [0, 2, 4, 6])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e6defb0-c0de-436f-90c0-e03c6826e399",
   "metadata": {},
   "source": [
    "### Convolution\n",
    "\n",
    "A convolution is a way of multiplying two sequences. Let $a$ and $b$ be infinite sequences. Then their convolution is a doubly infinite sequence whose $n$th element is\n",
    "$$\n",
    "(a \\ast b)_n = \\sum_{k=-\\infty}^{\\infty} a_k b_{n - k} \\qquad a = (\\ldots, a_{-1}, a_0, a_1, \\ldots) \\quad b = (\\ldots, b_{-1}, b_0, b_1, \\ldots)\n",
    "$$\n",
    "So, for example, the $0$th element of the convolution is\n",
    "$$\n",
    "\\sum_{k=-\\infty}^{\\infty} a_k b_{-k}\n",
    "$$\n",
    "We can define the convolution of two finite-dimensional vectors by viewing them as sequences padded by infinitely many zeros to the right and to the left. For example, if $a = (-1, -2, -3)$ and $b = (1, 2, 3)$, their convolution is\n",
    "$$\n",
    "\\begin{array}{rcl}\n",
    "a \\ast b & = & (-1 * 1, (-1 * 2) + (-2 * 1), (-1 * 3) + (-2 * 2) + (-3 * 1), (-2 * 3) + (-3 * 2), -3 * 3) \\\\\n",
    "& = & (-1, -4, -10, -12, -9)\n",
    "\\end{array}\n",
    "$$\n",
    "In code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56de8a74-78c4-46d5-9694-fdfbc2d06229",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convolution(a: List[complex], b: List[complex]) -> List[complex]:\n",
    "  M = len(a)\n",
    "  N = len(b)\n",
    "  return [\n",
    "    sum([a[k] * b[n - k] for k in range(M) if 0 <= n - k < N])\n",
    "    for n in range(N + M - 1)\n",
    "  ]\n",
    "\n",
    "a = [-1, -2, -3]\n",
    "b = [1, 2, 3]\n",
    "convolution(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8208b2bd-1aac-4c1e-a4a3-a1e8a447d9b7",
   "metadata": {},
   "source": [
    "#### Discrete Fourier Transform\n",
    "\n",
    "The convolution theorem says that we can calculate convolutions using the Discrete Fourier Transform (DFT) and the Hadamard product. This is important because using the definition, calculating a convolution is slow for large vectors and has complexity $O(n^2)$. However there is a fast way to calculate the DFT in $O(n \\log n)$ time aptly called the Fast Fourier Transform. Together with the Hadamard product which requires $O(n)$ time, this means that we can calculate convolutions in $O(n \\log n)$ time. Let's look at the details.\n",
    "\n",
    "To start, the formula is\n",
    "$$\n",
    "a \\ast b = IDFT(DFT(\\overline{a}) \\cdot DFT(\\overline{b}))\n",
    "$$\n",
    "where IDFT stands for the Inverse Discrete Fourier Transform and $\\overline{a}$ and $\\overline{b}$ are $a$ and $b$ padded with $N - 1$ zeros on the right, $N$ being the length of the vectors.\n",
    "\n",
    "The DFT of a column vector $a$ of length $N$ is defined as the multiplication of $a$ on the left by the matrix with entries\n",
    "$$\n",
    "(DFT_N)_{jk} = e^{-2\\pi i j k / N}\n",
    "$$\n",
    "or in other words multiplication by\n",
    "$$\n",
    "DFT_N = \\left(\n",
    "\\begin{array}{cccccc}\n",
    "1 & 1 & 1 & 1 & \\cdots & 1 \\\\\n",
    "1 & \\omega_N & \\omega_N^2 & \\omega_N^3 & \\cdots & \\omega_N^{N-1} \\\\\n",
    "1 & \\omega_N^2 & \\omega_N^4 & \\omega_N^6 & \\cdots & \\omega_N^{N-2} \\\\\n",
    "1 & \\omega_N^3 & \\omega_N^6 & \\omega_N^9 & \\cdots & \\omega_N^{N-3} \\\\\n",
    "\\vdots & & & & & \\\\\n",
    "1 & \\omega_N^{N-1} & \\omega_N^{N-2} & \\omega_N^{N-3} & \\cdots & \\omega_N\n",
    "\\end{array}\n",
    "\\right) \\qquad \\omega_N = e^{-2\\pi i / N}\n",
    "$$\n",
    "The inverse of this matrix, i.e., the matrix of the IDFT, turns out to be its elementwise complex conjugate divided by $N$. In code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3533ff8-6cd9-4b42-8a6d-8c577bc7f0d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def omega(N: int) -> complex:\n",
    "  return pow(math.e, -2 * math.pi * 1j / N)\n",
    "  \n",
    "def DFT_matrix(N: int) -> np.array:\n",
    "  return np.array([[pow(omega(N), j * k) for k in range(N)] for j in range(N)])\n",
    "\n",
    "def DFT(a: List[complex]) -> np.array:\n",
    "  return DFT_matrix(len(a)) @ np.array(a)\n",
    "  \n",
    "def IDFT_matrix(N: int) -> List[List[complex]]:\n",
    "  return np.conjugate(DFT_matrix(N)) / N\n",
    "\n",
    "def IDFT(a: List[complex]) -> np.array:\n",
    "  return IDFT_matrix(len(a)) @ np.array(a)\n",
    "  \n",
    "print(\"DFT_matrix(4):\")\n",
    "print(prettify(DFT_matrix(4)))\n",
    "print(\"IDFT_matrix(4):\")\n",
    "print(prettify(IDFT_matrix(4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3650b000-37b0-45f0-a264-ca16381780a2",
   "metadata": {},
   "source": [
    "We can now check the convolution theorem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4236b1-dc9e-40d7-b98f-b17883dde20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [-1, -2, -3, 0, 0]\n",
    "b = [1, 2, 3, 0, 0]\n",
    "print(prettify(IDFT(hadamard(DFT(a), DFT(b)))))\n",
    "print(prettify(convolution(a[:3], b[:3])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea18bd3-35c7-428f-9789-7b334f9776f5",
   "metadata": {},
   "source": [
    "Or here's a check with random vectors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692f8bf0-f55a-4c5a-ab7d-8cc04c9f909a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#random.seed(0)\n",
    "\n",
    "a = [random.random() for _ in range(8)]\n",
    "b = [random.random() for _ in range(8)]\n",
    "a_bar = a + [0] * (len(a) - 1)\n",
    "b_bar = b + [0] * (len(b) - 1)\n",
    "print(prettify(IDFT(hadamard(DFT(a_bar), DFT(b_bar)))))\n",
    "print(prettify(convolution(a, b)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9908d16a-a3c0-4c6b-8798-1c9a7a65514a",
   "metadata": {},
   "source": [
    "#### Fast Fourier Transform\n",
    "\n",
    "The idea for how to calculate the Discrete Fourier Transform in $O(n \\log n)$ time is to use the divide and conquer technique. This means noticing that\n",
    "$$\n",
    "DFT_N(a_s)_k = DFT_{N/2}(a_{2s})_k + e^{-2\\pi i k/N} DFT_{N/2}(a_{2s+1})_k, \\qquad k = 0, \\ldots, N / 2 - 1\n",
    "$$\n",
    "and\n",
    "$$\n",
    "DFT_N(a_s)_k = DFT_{N/2}(a_{2s})_{k-N/2} + e^{-2\\pi i k/N} DFT_{N/2}(a_{2s+1})_{k-N/2}, \\qquad k = N/2, \\ldots, N - 1\n",
    "$$\n",
    "where, for example, $DFT_N(a_{2s})_k$ means the $k$th coordinate of the Discrete Fourier Transform of the even entries (using 0-based indexing) of the vector $a$. The code might be easier to read than the formula: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5394737a-7481-40f2-a79f-66315e179c6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def roots_of_unity(n: int) -> List[complex]:\n",
    "  return [pow(math.e, -2 * math.pi * 1j * k / n) for k in range(n)]\n",
    "\n",
    "def _FFT(x: List[complex], W: List[complex]) -> Tuple[List[complex], int]:\n",
    "  k = len(x)\n",
    "  if k == 1:\n",
    "    return x, 0\n",
    "  else:\n",
    "    n = len(W)\n",
    "    \n",
    "    X_even, ops_even = _FFT([x[i] for i in range(0, k, 2)], W)\n",
    "    X_odd, ops_odd = _FFT([x[i] for i in range(1, k, 2)], W)\n",
    "    W_k = [W[i] for i in range(0, n, n // k)]\n",
    "\n",
    "    X_left = [X_even[i] + W_k[i] * X_odd[i] for i in range(k // 2)]\n",
    "    X_right = [X_even[i] + W_k[k // 2 + i] * X_odd[i] for i in range(k // 2)]\n",
    "    ops = 2 * k + ops_even + ops_odd\n",
    "    return X_left + X_right, ops\n",
    "    \n",
    "def FFT(x: List[complex], verbose: bool = False) -> np.array:\n",
    "  n = len(x)\n",
    "  assert n == n & ~(n - 1), \"only vectors of length 2**n are supported\"\n",
    "\n",
    "  W = roots_of_unity(n)\n",
    "\n",
    "  X, ops = _FFT(x, W)\n",
    "  if verbose:\n",
    "    print(\"n:\", n, \"ops:\", ops)\n",
    "  return np.array(X)\n",
    "\n",
    "for k in range(2, 6):\n",
    "  a = list(range(2**k))\n",
    "  FFT(a, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2100feb0-1b24-48be-80f8-a1cee5d95309",
   "metadata": {},
   "source": [
    "Notice that the algorithm uses $2n \\log_2 n$ operations not including the calculation of the roots of unity so we've achieved our goal of subquadratic complexity. Let's check that it gives the same answer as the DFT:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71997015-02e5-4743-bb0e-0d27ddee58f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#random.seed(0)\n",
    "\n",
    "a = [random.random() for _ in range(8)]\n",
    "print(prettify(DFT(a)))\n",
    "print(prettify(FFT(a)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b63c59eb-0e4d-47dc-9e1f-97acc1e0df36",
   "metadata": {},
   "source": [
    "### Padded cross-correlation\n",
    "The last concept we need to understand the Hyena operator is the padded cross-correlation, a.k.a. the depthwise convolution. Even though it is commonly called a convolution in machine learning, it is not a convolution in the usual mathematical sense that we defined above. If $x$ and $f$ are vectors, then their cross-correlation is the vector with $n$th coordinate\n",
    "$$\n",
    "(f \\star x)_n = \\sum_{k=1}^d f_k x_{n + k} \\qquad x = (x_0, \\ldots, x_{N - 1}) \\quad f = (f_0, \\ldots, f_{d - 1})\n",
    "$$\n",
    "The cross-correlation has length $N - d + 1$. For those who prefer code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc3de60b-2dc6-4a8c-97c5-b4dd686a99cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def crosscorrelation1d(x: List[float], filter: List[float]) -> List[float]:\n",
    "  N = len(x)\n",
    "  d = len(filter)\n",
    "  return [\n",
    "    sum(x[n + k] * filter[k] for k in range(d))\n",
    "    for n in range(N - d + 1)\n",
    "  ]\n",
    "\n",
    "x = [1, 2, 3, 4, 5, 6, 7]\n",
    "f = [-1, 2, 1]\n",
    "print(prettify(crosscorrelation1d(x, f)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "089a520a-753e-429f-babe-2f11b0355b93",
   "metadata": {},
   "source": [
    "The Hyena operator will be used to create a model that takes in text and attempts to predict the next token, so we need it to be _causal_. In other words, we do not want a value at position $n$ in the sequence receiving information from values in future positions, e.g., at $n + 1$. Otherwise the model will use the information about the future tokens to predict the next token.\n",
    "\n",
    "The cross-correlation will not be causal because, for example, given a length three vector $f$, the $0$th coordinate of the convolution is $(f \\star x)_0 = f_0 x_0 + f_1 x_1 + f_2 x_2$. So the $0$th value will have information about $x_1$ and $x_2$ which are values in future positions relative to the $0$th value. To make the cross-correlation causal, we just need to pad the vector on the left with zeros. If $f$ has length $d$ then we need $d - 1$ zeros. In our example, that would mean that $(f \\star x)_0 = f_0 0 + f_1 0 + f_2 x_0$.\n",
    "\n",
    "### Adding dimensions\n",
    "\n",
    "So far we have considered the Hyena operator as an operator on vectors. However in practice we will want it to operate on tensors that have a batch axis and an embedding axis, in other words on tensors with shape $(b, d, L)$ where $b$ is the number of samples in a batch and $d$ is the embedding dimension. We just apply the operations we have described in parallel across the batch and embedding dimensions.\n",
    "\n",
    "### Putting it all together\n",
    "\n",
    "Using PyTorch it is straightforward to define the matrix multiplication followed by padded cross-correlation. We will call it the `Projection` module because it projects the input embeddings to $x_1, \\ldots, x_N$. As a reminder, it is implementing\n",
    "$$\n",
    "x_i(u) = w_i \\star A_i u\n",
    "$$\n",
    "Instead of implementing it in a loop, we do the operation for all $i$ at once and then split the result into separate $x_i$.\n",
    "As explained in the previous section, the cross-correlation (`Conv1d`) is executed in parallel across the embedding dimension which is accomplished by setting `groups=d_model * (N + 1)`, i.e., one group per dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d76dc219-f586-445c-ab59-9212fa15faf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Projection(torch.nn.Module):\n",
    "  def __init__(self, d_model: int, N: int, conv_len: int):\n",
    "    super().__init__()\n",
    "    self.d_model = d_model\n",
    "    self.N = N\n",
    "    self.linear = torch.nn.Linear(d_model, d_model * (N + 1))\n",
    "    self.conv = torch.nn.Conv1d(\n",
    "      in_channels=d_model * (N + 1),\n",
    "      out_channels=d_model * (N + 1),\n",
    "      kernel_size=conv_len,\n",
    "      groups=d_model * (N + 1),  # Depthwise convolution\n",
    "      padding=conv_len - 1,\n",
    "    )\n",
    "    \n",
    "  def forward(self, u: torch.Tensor) -> List[torch.Tensor]:\n",
    "    z = self.linear(u)\n",
    "    z = z.transpose(1, 2)  # Channels (embedding dim) needs to come first\n",
    "    \n",
    "    L = z.shape[2]\n",
    "    z = self.conv(z)[..., :L]\n",
    "    \n",
    "    x = torch.split(z, self.d_model, dim=1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "390c16b9-e965-4edf-b34e-83d0c9562b74",
   "metadata": {},
   "source": [
    "Next we define the convolution using FFT and the Hadamard product as discussed above. Because our vectors are real-valued, we can use a special version of FFT optimized for real numbers called `torch.fft.rfft`. It turns out that the FFT of a real-valued vector has the property that every conjugate of a coordinate is also a coordinate, and `rfft` drops the conjugates because they are superfluous. See below for an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9bd7f4-6dd9-41c8-b49f-cc9de779d700",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.Tensor([random.random() for _ in range(8)])\n",
    "b = torch.Tensor([random.random() for _ in range(8)])\n",
    "a_rf = torch.fft.rfft(a, norm=\"forward\")\n",
    "b_rf = torch.fft.rfft(b, norm=\"forward\")\n",
    "c_rf = a_rf * b_rf\n",
    "cr = torch.fft.irfft(c_rf, norm=\"forward\")\n",
    "\n",
    "a_f = torch.fft.fft(a, norm=\"forward\")\n",
    "b_f = torch.fft.fft(b, norm=\"forward\")\n",
    "c_f = a_f * b_f\n",
    "c = torch.fft.irfft(c_f, norm=\"forward\")\n",
    "\n",
    "print(\"a_rf:\", prettify(a_rf.numpy()))\n",
    "print(\"a_f:\", prettify(a_f.numpy()))\n",
    "print(\"convolution using rfft:\", prettify(cr.numpy()))\n",
    "print(\"convolution using fft:\", prettify(cr.numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b042c485-2475-44c5-8c41-548d60f4efc6",
   "metadata": {},
   "source": [
    "We also add in a skip connection, i.e., we actually compute\n",
    "$$\n",
    "h_i \\ast x_i + B_i x_i\n",
    "$$\n",
    "for some matrix $B_i$. This improves gradient flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e224f2-4d7a-4eb6-9d49-de7debca4420",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FFTConv(torch.nn.Module):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "\n",
    "  def forward(\n",
    "      self,\n",
    "      h: torch.Tensor,\n",
    "      x: torch.Tensor,\n",
    "      B: torch.Tensor\n",
    "    ) -> torch.Tensor:\n",
    "    L = h.shape[-1]\n",
    "    h_f = torch.fft.rfft(h, n=2 * L, norm=\"forward\")\n",
    "    x_f = torch.fft.rfft(x.to(dtype=h.dtype), n=2 * L)\n",
    "    y = torch.fft.irfft(h_f * x_f, n=2 * L, norm=\"forward\")[..., :L]\n",
    "    y = y + x * B\n",
    "    y = y.to(dtype=h.dtype)  # y is ComplexFloat but we need it to be float\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c4aadcc-1849-45e7-a43c-9eddaf26a7f1",
   "metadata": {},
   "source": [
    "Now we are ready to define the Hyena block which is an $N$th order Hyena Operator followed by a linear output layer. It will be a drop-in replacement for an attention block in a transformer model.\n",
    "\n",
    "We make two changes which are different from the paper to address vanishing and exploding gradients:\n",
    "\n",
    "1. We add three skip connections which are marked by comments below. By a skip connection we just mean that we add the input tensor to the output of an operator. See the [Resnet paper](https://arxiv.org/pdf/1512.03385.pdf) for more details.\n",
    "2. We normalize $x_i$ across the embedding dimension.\n",
    "\n",
    "Note that it is important that the softmax is taken across the embedding dimension and not across the sequence dimension because taking it across the sequence dimension would mean the operator is no longer causal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7c1e37d-64a1-4ddd-9d9c-fa515e77e21b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HyenaBlock(torch.nn.Module):\n",
    "  def __init__(self, config: HyenaConfig):\n",
    "    super().__init__()\n",
    "    self.proj_input = Projection(config.d_model, config.order, config.short_conv_size)\n",
    "    self.proj_output = torch.nn.Linear(config.d_model, config.d_model)\n",
    "    self.filter = HyenaFilter(\n",
    "      config.d_model,\n",
    "      config.d_filter_mlp,\n",
    "      config.d_embed,\n",
    "      config.order,\n",
    "      config.n_filter_layers,\n",
    "      config.context_length,\n",
    "      config.omega,\n",
    "    )\n",
    "    self.dropout = torch.nn.Dropout(config.pdrop_hyena)\n",
    "    self.fft_conv = FFTConv()\n",
    "    self.B = torch.nn.Parameter(torch.randn((config.order, 1, config.d_model, 1)))\n",
    "\n",
    "  def forward(self, u: torch.Tensor) -> torch.Tensor:\n",
    "    L = u.shape[1]\n",
    "    \n",
    "    *x, v = self.proj_input(u)\n",
    "    v = v + u.transpose(1, 2)  # skip connection\n",
    "    \n",
    "    h = self.filter(L)\n",
    "    \n",
    "    for i, x_i in enumerate(x):\n",
    "      h_i = h[i].unsqueeze(0)\n",
    "      v = v + torch.nn.functional.normalize(x_i, dim=1) * self.fft_conv(h_i, v, self.B[i])  # skip connection\n",
    "    \n",
    "    v = v.transpose(1, 2)\n",
    "    y = v + self.proj_output(v)  # skip connection\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2606fa1d-b1ed-406a-8757-527f6309a7ee",
   "metadata": {},
   "source": [
    "## Part 2: Hyena Filter\n",
    "\n",
    "It remains to define the vectors $h_i$. We will call these vectors \"filters\" because we are convolving the input with them.\n",
    "\n",
    "The filters are parameters of the model, but we force them to have a special form that decays exponentially along the sequence so that the model will pay more attention to close context than far away context.\n",
    "$$\n",
    "h_i = \\mathrm{norm}(g_i \\cdot (e^{\\alpha \\cdot t} + b))\n",
    "$$\n",
    "where $\\mathrm{norm}(x) = x / |x|$ along the sequence axis. The vector $t$ is fixed to be equal to $(0, 1/(L-1), 2/(L-1), \\ldots, 1)$ for sequence length $L$. The vectors $\\alpha$ and $b$ are parameters of the model that vary along the embedding axis (they are constant along the sequence axis). The addition of the $b$ vector allows the model to prevent the exponential decay from approaching zero. The vector $g_i$ is also a parameter of the model. The dot again denotes the Hadamard product.\n",
    "\n",
    "The paper has a more complex definition of $h_i$ that is described below in an appendix, but in our limited testing, the additional complexity did not win any performance improvement.\n",
    "\n",
    "Putting this into code gives:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ebac238-276a-4468-a0ee-1bd5dcac5b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Window(torch.nn.Module):\n",
    "  def __init__(\n",
    "      self,\n",
    "      d_model: int,\n",
    "      max_seq_len: int,\n",
    "      fast_decay_pct: float = 0.3,\n",
    "      slow_decay_pct: float = 1.5,\n",
    "      target: float = 1e-2,\n",
    "    ):\n",
    "    super().__init__()\n",
    "    self.b = torch.nn.Parameter(torch.zeros((1, d_model, 1)))\n",
    "    min_decay = math.log(target) / slow_decay_pct\n",
    "    max_decay = math.log(target) / fast_decay_pct\n",
    "    self.alphas = torch.nn.Parameter(\n",
    "      torch.linspace(\n",
    "        start=min_decay,\n",
    "        end=max_decay,\n",
    "        steps=d_model\n",
    "      )[None, :, None]\n",
    "    )\n",
    "    self.t = torch.nn.Parameter(\n",
    "      torch.linspace(\n",
    "        start=0,\n",
    "        end=1,\n",
    "        steps=max_seq_len\n",
    "      )[None, None, :], requires_grad=False\n",
    "    )\n",
    "      \n",
    "  def forward(self, x):\n",
    "    L = x.shape[2]\n",
    "    c = torch.exp(self.alphas * self.t)[:, :, :L]\n",
    "    x = x * (c + self.b)\n",
    "    return x\n",
    "      \n",
    "class HyenaFilter(torch.nn.Module):\n",
    "  def __init__(\n",
    "      self,\n",
    "      d_model: int,\n",
    "      d_mlp: int,\n",
    "      d_embed: int,\n",
    "      N: int,\n",
    "      n_layers: int = 4,\n",
    "      max_seq_len: int = 128,\n",
    "      omega: int = 8,\n",
    "    ):\n",
    "    assert n_layers >= 2, \"n_layers must be at least 2\"\n",
    "    super().__init__()\n",
    "\n",
    "    self.N = N\n",
    "    self.d_model = d_model\n",
    "\n",
    "    # Making this a parameter, even though it is not trained, ensures\n",
    "    # it will be moved to the gpu with the rest of the model\n",
    "    self.h = torch.nn.Parameter(torch.randn((N, d_model, max_seq_len)))\n",
    "      \n",
    "    self.window = Window(d_model, max_seq_len)\n",
    "  \n",
    "  def forward(self, L: int) -> torch.Tensor:\n",
    "    h = self.h[:, :, :L]\n",
    "    h = self.window(h)\n",
    "    \n",
    "    h = h / torch.norm(h, dim=-2, p=1, keepdim=True)\n",
    "        \n",
    "    return h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e48f80b6-3d42-4c9a-ac01-f1f0dbf94ca8",
   "metadata": {},
   "source": [
    "## Part 3: A Working Model\n",
    "\n",
    "The `HyenaBlock` defined above is a drop-in replacement for a self-attention block, so we can just use the standard GPT-2 architecture to make a full model. Recall that GPT-2 has randomly initialized token and position embeddings with the input token embedding weights tied (the same variables as) to the logit output weights. We apply dropout to the embeddings.\n",
    "\n",
    "We train using the cross-entropy loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bc15afb-470b-4e63-91c0-9dd4bad15a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPModel(lightning.LightningModule):\n",
    "  def __init__(self, config: Config, block_cls: torch.nn.Module):\n",
    "    super().__init__()\n",
    "    self.config = config\n",
    "    self.tok_emb = torch.nn.Embedding(config.vocab_size, config.d_model)\n",
    "    self.pos_emb = torch.nn.Parameter(\n",
    "      torch.randn(1, config.context_length, config.d_model)\n",
    "    )\n",
    "    self.drop = torch.nn.Dropout(config.pdrop_embed)\n",
    "    self.layers = torch.nn.Sequential(*[\n",
    "      block_cls(config) for _ in range(config.n_layers)\n",
    "    ])\n",
    "    self.ln = torch.nn.LayerNorm(config.d_model)\n",
    "    self.head = torch.nn.Linear(\n",
    "      config.d_model,\n",
    "      config.vocab_size,\n",
    "      bias=False\n",
    "    )\n",
    "    # input embedding and logit output weights are tied\n",
    "    self.head.weight = self.tok_emb.weight\n",
    "\n",
    "  def forward(\n",
    "      self,\n",
    "      batch: Tuple[torch.Tensor, torch.Tensor]\n",
    "    ) -> torch.Tensor:\n",
    "    x, y = batch\n",
    "\n",
    "    token_embeddings = self.tok_emb(x)\n",
    "    position_embeddings = self.pos_emb[:, :token_embeddings.shape[1], :]\n",
    "\n",
    "    x = self.drop(token_embeddings + position_embeddings)\n",
    "    x = self.layers(x)\n",
    "    logits = self.head(self.ln(x))\n",
    "\n",
    "    return logits, y\n",
    "    \n",
    "  def calculate_loss(\n",
    "      self,\n",
    "      logits: torch.Tensor,\n",
    "      targets: torch.Tensor\n",
    "    ) -> float:\n",
    "    loss = torch.nn.functional.cross_entropy(\n",
    "      logits.transpose(1, 2), targets\n",
    "    )\n",
    "     \n",
    "    return loss\n",
    "\n",
    "  def training_step(\n",
    "      self,\n",
    "      batch: Tuple[torch.Tensor, torch.Tensor],\n",
    "      batch_idx: int\n",
    "    ) -> float:\n",
    "    #with torch.autograd.detect_anomaly():\n",
    "    logits, targets = self.forward(batch)\n",
    "    loss = self.calculate_loss(logits, targets)\n",
    "    self.log(\n",
    "      \"train_loss\", loss, prog_bar=True, on_step=True, on_epoch=False\n",
    "    )\n",
    "    return loss\n",
    "\n",
    "  def validation_step(\n",
    "      self,\n",
    "      batch: Tuple[torch.Tensor, torch.Tensor],\n",
    "      batch_idx: int\n",
    "    ) -> float:\n",
    "    logits, targets = self.forward(batch)\n",
    "    loss = self.calculate_loss(logits, targets)\n",
    "    self.log(\"val_loss\", loss, prog_bar=True, on_step=False, on_epoch=True)\n",
    "    return loss\n",
    "    \n",
    "  def configure_optimizers(self):\n",
    "    return torch.optim.AdamW(\n",
    "      self.parameters(),\n",
    "      lr=self.config.learning_rate,\n",
    "      betas=self.config.betas,\n",
    "      weight_decay=self.config.weight_decay,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df82d20c-bb86-428c-b9e1-de0f635e6115",
   "metadata": {},
   "source": [
    "### Dataset\n",
    "\n",
    "Now we just need a dataset. We choose a dataset of Karpathy composed of text from Shakespeare. The below code downloads it from HuggingFace. It is a single very long row of data.\n",
    "\n",
    "Our tokens will be the individual characters of the text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc232808-ce30-4445-bed4-58bd1ffe55cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = datasets.load_dataset(\"tiny_shakespeare\", split=\"train\")\n",
    "ds = ds.map(\n",
    "  lambda x: {\"char\": regex.findall(r\"\\X\", x[\"text\"])},\n",
    "  remove_columns=[\"text\"]\n",
    ")\n",
    "\n",
    "vocabulary = sorted(set(ds[0][\"char\"]))  # Entire dataset is a single row\n",
    "print(vocabulary)\n",
    "\n",
    "tok2id = {ch: i for i, ch in enumerate(vocabulary)}\n",
    "print(tok2id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a866e1-81eb-4b1f-83d4-481e643f3c39",
   "metadata": {},
   "source": [
    "We process the dataset to tokenize it and then split it into lines of a size specified by the config. Then we split the lines into training and validation sets. Finally we save it to disk so next time we will not need to download it.\n",
    "\n",
    "Crucially we need to choose a context length that will be the maximum number of tokens our models can process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad503eb1-6913-4f47-a75b-3d0da986fa97",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONTEXT_LENGTH = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aa7e047-2a6f-4a48-8ef4-2bf17baff3b0",
   "metadata": {},
   "source": [
    "With that decided, we can process the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86093d88-3af2-4e69-a12f-95fc6394f196",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of rows to use for the validation set\n",
    "val_size = 3200\n",
    "  \n",
    "def process_dataset(\n",
    "    ds: datasets.Dataset\n",
    "  ) -> Tuple[datasets.Dataset, datasets.Dataset]:\n",
    "  \n",
    "  def chunk(lst: List[Any], n: int) -> List[List[Any]]:\n",
    "    \"\"\"Break `lst` into length n chunks, dropping final chunk\"\"\"\n",
    "    return [\n",
    "      lst[i:i + n]\n",
    "      for i in range(0, len(lst) - n, n)\n",
    "    ]\n",
    "  \n",
    "  def create_batches(x: Dict[str, List[Any]]) -> Dict[str, List[Any]]:\n",
    "    return {\n",
    "      \"curr_id\": chunk(x[\"curr_id\"][0], CONTEXT_LENGTH),\n",
    "      \"next_id\": chunk(x[\"next_id\"][0], CONTEXT_LENGTH),\n",
    "    }\n",
    "  \n",
    "  def tokenize(x: Dict[str, List[Any]]) -> Dict[str, List[Any]]:\n",
    "    return {\n",
    "      \"id\": [tok2id[ch] for ch in x[\"char\"]]\n",
    "    }\n",
    "  \n",
    "  ds = ds.map(tokenize, remove_columns=[\"char\"])\n",
    "  ds = ds.map(\n",
    "    lambda x: {\"curr_id\": x[\"id\"][:-1], \"next_id\": x[\"id\"][1:]},\n",
    "    remove_columns=[\"id\"]\n",
    "  )\n",
    "  ds = ds.map(create_batches, batched=True, batch_size=1)\n",
    "  ds = ds.shuffle(seed=0)\n",
    "  \n",
    "  val_ds = ds.select(range(0, val_size))\n",
    "  train_ds = ds.select(range(val_size, len(ds)))\n",
    "  \n",
    "  return train_ds, val_ds\n",
    "\n",
    "train_ds, val_ds = process_dataset(ds)\n",
    "train_ds.save_to_disk(f\"data/train-{CONTEXT_LENGTH}\")\n",
    "val_ds.save_to_disk(f\"data/val-{CONTEXT_LENGTH}\")\n",
    "\n",
    "with open(\"data/vocabulary.json\", \"w\") as f:\n",
    "  json.dump(vocabulary, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd355c59-5678-4471-82bc-858c6266bd33",
   "metadata": {},
   "source": [
    "If you have previously downloaded the data and just want to load it from disk, then use this code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09257f89-e9ac-440c-a3f0-6626484cb0da",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = datasets.Dataset.load_from_disk(f\"data/train-{CONTEXT_LENGTH}\")\n",
    "val_ds = datasets.Dataset.load_from_disk(f\"data/val-{CONTEXT_LENGTH}\")\n",
    "\n",
    "with open(\"data/vocabulary.json\", \"r\") as f:\n",
    "  vocabulary = json.load(f)\n",
    "  print(vocabulary)\n",
    "\n",
    "tok2id = {ch: i for i, ch in enumerate(vocabulary)}\n",
    "print(tok2id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e94032a-33bb-481c-863f-21dbb5505586",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "The training loop can be defined using PyTorch Lightning. If you plan to experiment with the architecture, the tool from Weights and Biases might be useful. It requires a free account and uncommenting the Weights and Biases (`wandb`) code. It allows you to view charts of the gradients to diagnose problems like vanishing gradients as well as charts of the loss that will allow you to compare different runs.\n",
    "\n",
    "If you do use Weights and Biases, the first step is to log in (assuming you have installed the `wandb` package above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53252a3a-ed39-49f3-8189-a0b7c69346ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b7b56ff-a5a8-4279-97bb-985dacd811b1",
   "metadata": {},
   "source": [
    "In any case, Pytorch Lightning makes the training loop easy to write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3e73fe8-972e-4469-8254-2f5078e8dc5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning.pytorch.loggers\n",
    "\n",
    "def collate(\n",
    "    data: List[Dict[str, List[int]]]\n",
    ") -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "    curr_ids = torch.LongTensor([d[\"curr_id\"] for d in data])\n",
    "    next_ids = torch.LongTensor([d[\"next_id\"] for d in data])\n",
    "    return curr_ids, next_ids\n",
    "\n",
    "\n",
    "def train(model: lightning.LightningModule, config: Config) -> None:\n",
    "    wandb_logger = lightning.pytorch.loggers.WandbLogger(\n",
    "      project=\"hyena-gpt-shakespeare\"\n",
    "    )\n",
    "    wandb_logger.experiment.config.update(dataclasses.asdict(config))\n",
    "    wandb_logger.watch(model, log=\"all\", log_freq=1)\n",
    "  \n",
    "    trainer = lightning.Trainer(\n",
    "        accelerator=config.device_type,\n",
    "        precision=config.precision,\n",
    "        max_epochs=config.epochs,\n",
    "        gradient_clip_val=0.2,\n",
    "        logger=wandb_logger,\n",
    "    )\n",
    "\n",
    "    train_dl = torch.utils.data.DataLoader(\n",
    "        train_ds,\n",
    "        collate_fn=collate,\n",
    "        shuffle=True,\n",
    "        pin_memory=True,\n",
    "        batch_size=config.batch_size,\n",
    "        num_workers=config.num_workers,\n",
    "    )\n",
    "\n",
    "    val_dl = torch.utils.data.DataLoader(\n",
    "        val_ds,\n",
    "        collate_fn=collate,\n",
    "        shuffle=False,\n",
    "        pin_memory=True,\n",
    "        batch_size=config.batch_size,\n",
    "        num_workers=config.num_workers,\n",
    "    )\n",
    "\n",
    "    trainer.fit(model, train_dl, val_dl)\n",
    "\n",
    "    wandb.finish()\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ec061d3-314d-4b3a-a5dd-241e6cd29eb7",
   "metadata": {},
   "source": [
    "All that's left to do is set the hyperparameters and train. If you are using a cpu, then change the device type to \"cpu\" and the precision to \"32\". You might also want to decrease the size of the model to speed up training and decrease the number of epochs. If you are using a gpu, it may not support bf16 precision, in which case you can change it to \"16\" or \"32\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f13192f-e162-4439-b236-f54c61bd4d9d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "hyena_config = HyenaConfig(\n",
    "  d_model=386,\n",
    "  n_layers=6,\n",
    "  vocab_size=len(vocabulary),\n",
    "  d_embed=33,\n",
    "  d_filter_mlp=64,\n",
    "  n_filter_layers=4,\n",
    "  context_length=CONTEXT_LENGTH,\n",
    "  short_conv_size=3,\n",
    "  order=2,\n",
    "  pdrop_hyena=0.0,\n",
    "  pdrop_embed=0.2,\n",
    "  omega=12,\n",
    "  epochs=40,\n",
    "  learning_rate=6e-4,\n",
    "  betas=(0.9, 0.98),\n",
    "  weight_decay=0.4,\n",
    "  device_type=\"gpu\",  # cpu, gpu\n",
    "  precision=\"bf16\",  # 32, 16, 16-mixed, bf16\n",
    "  batch_size=64,\n",
    "  num_workers=4,\n",
    ")\n",
    "\n",
    "hyena_model = GPModel(hyena_config, HyenaBlock)\n",
    "hyena_model = train(hyena_model, hyena_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6691b7c9-e5b0-4504-a54a-b9eaf73e1f97",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "\n",
    "Once the model is trained we should check that it can predict the next token reasonably well and that its generated text shows some resemblance to its training data. If there is a bug that, for example, violates causility, the loss could have decreased but the model will not do prediction or generation well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3c7dea3-4ae5-4d87-8b3e-b5080864b5c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "  \"hyena average probability of next token:\", \n",
    "  average_probability(hyena_model, val_ds, tokens=256)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5de62c9-17fb-445f-9597-2c61666ae1cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "  generate(\n",
    "    hyena_model,\n",
    "    hyena_config.context_length,\n",
    "    \"Wherefore art thou \",\n",
    "    method=\"topk\",\n",
    "    max_tokens=100,\n",
    "    k=2\n",
    "  )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb096aa-9978-4d2c-bd98-2d14dc73e7dc",
   "metadata": {},
   "source": [
    "## Appendix: Comparison with Attention\n",
    "\n",
    "We can define a traditional causal self-attention layer as follows (Cf. [The Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d2fe136-787c-4367-9297-f8c4ba53a010",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GELU(torch.nn.Module):\n",
    "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        return torch.nn.functional.gelu(input)\n",
    "\n",
    "class CausalSelfAttention(torch.nn.Module):\n",
    "  def __init__(self, d_embed: int, n_head: int, pdrop_attn: float):\n",
    "    super().__init__()\n",
    "    assert d_embed % n_head == 0\n",
    "    self.d_embed = d_embed\n",
    "    self.n_head = n_head\n",
    "\n",
    "    self.mask = torch.zeros((1, 1), dtype=torch.bool)\n",
    "    self.attn = torch.nn.MultiheadAttention(\n",
    "      embed_dim=d_embed, num_heads=n_head, dropout=pdrop_attn, batch_first=True\n",
    "    )\n",
    "\n",
    "  def forward(self, x: torch.Tensor, padding: Optional[torch.Tensor] = None) -> torch.Tensor:\n",
    "    seq_len = x.shape[1]\n",
    "    if self.mask.shape[0] != seq_len:\n",
    "      self.mask = torch.tril(\n",
    "        torch.ones((seq_len, seq_len), dtype=torch.bool, device=x.device), diagonal=-1\n",
    "      ).T\n",
    "    return self.attn(\n",
    "      x, x, x, key_padding_mask=padding, need_weights=False, attn_mask=self.mask\n",
    "    )[0]\n",
    "\n",
    "class SelfAttentionBlock(torch.nn.Module):\n",
    "  def __init__(self, config: AttentionConfig):\n",
    "    super().__init__()\n",
    "    self.ln1 = torch.nn.LayerNorm(config.d_embed)\n",
    "    self.ln3 = torch.nn.LayerNorm(config.d_embed)\n",
    "    self.self_attn = CausalSelfAttention(\n",
    "      config.d_embed,\n",
    "      config.n_head,\n",
    "      config.pdrop_attn,\n",
    "    )\n",
    "    self.mlp = torch.nn.Sequential(\n",
    "      torch.nn.Linear(config.d_embed, 4 * config.d_embed),\n",
    "      GELU(),\n",
    "      torch.nn.Linear(4 * config.d_embed, config.d_embed),\n",
    "      torch.nn.Dropout(config.pdrop_attn),\n",
    "    )\n",
    "\n",
    "  def forward(self, x: torch.Tensor, padding: Optional[torch.Tensor] = None) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Do layer normalization before attention/MLP according to\n",
    "    https://arxiv.org/pdf/2002.04745.pdf\n",
    "    \"\"\"\n",
    "    # padding is for x (the key)\n",
    "    z = x + self.self_attn(self.ln1(x), padding)\n",
    "    z = z + self.mlp(self.ln3(z))\n",
    "    return z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e450d8-0e17-4126-9fa7-a63f0c416659",
   "metadata": {},
   "source": [
    "Then we can simply pass the layer to the model constructor which will use it to build the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea7ea129-96f8-4dd8-afdc-0157993be00e",
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_config = AttentionConfig(\n",
    "  d_model=384,\n",
    "  n_layers=6,\n",
    "  vocab_size=len(vocabulary),\n",
    "  d_embed=384,\n",
    "  n_head=6,\n",
    "  context_length=CONTEXT_LENGTH,\n",
    "  pdrop_attn=0.2,\n",
    "  pdrop_embed=0.2,\n",
    "  learning_rate=3e-4,\n",
    "  epochs=40,\n",
    "  betas=(0.9, 0.98),\n",
    "  weight_decay=0.1,\n",
    "  device_type=\"gpu\",\n",
    "  precision=\"bf16\",\n",
    "  batch_size=64,\n",
    "  num_workers=4,\n",
    ")\n",
    "\n",
    "attn_model = GPModel(attn_config, SelfAttentionBlock)\n",
    "attn_model = train(attn_model, attn_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3cd428-55d5-429d-ab8b-d1f5ff615d0f",
   "metadata": {},
   "source": [
    "In our experiments, we found that the attention-based model attains a loss of about 1.6 compared to less than 1.5 for the hyena-based model despite the attention-based model having more parameters (10M vs 8M). However in this short (128) context regime it does execute more quickly than the hyena-based model.\n",
    "\n",
    "Finally we evaluate the prediction and generation of the newly trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89b06c8d-5ed6-4bf9-ba00-45bda406f26f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\n",
    "  \"attention average probability of next token:\",\n",
    "  average_probability(attn_model, val_ds, tokens=256)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a4a9e61-5e40-4271-a404-903ae1d4f497",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "  generate(\n",
    "    attn_model,\n",
    "    attn_config.context_length,\n",
    "    \"Wherefore art thou \",\n",
    "    method=\"topk\",\n",
    "    max_tokens=80,\n",
    "    k=2\n",
    "  )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbac3275-ed9a-4d11-a664-5d124c07cf2e",
   "metadata": {},
   "source": [
    "## Appendix: Hyena as described in the paper\n",
    "\n",
    "__The below modules are labeled as \"authentic\" because they correspond to the description in the paper. The authors may have omitted details that greatly improve the performance, so don't take the below performance as an indicator of the authors' work.__\n",
    "\n",
    "The paper's Hyena filter is more complex than the simplified version presented above. It works in three steps:\n",
    "\n",
    "1. There is a trainable positional embedding initialized with `sin` and `cos` values\n",
    "2. The positional embeding is passed through a few linear layers with sinusoidal activation functions\n",
    "3. The output is multiplied elementwise by trainable exponentially decaying vectors (\"windows\")\n",
    "\n",
    "The positional embedding is initialized to a matrix whose $t$th row is\n",
    "$$\n",
    "[t / L, \\cos(2\\pi*0*t/L), \\ldots, \\cos(2\\pi*(K-1)*t/L), \\sin(2\\pi*0*t/L), \\ldots \\sin(2\\pi*(K-1)*t/L]\n",
    "$$\n",
    "for a hyperparameter $K$ which determines the size of the embedding.\n",
    "\n",
    "The frequency of the `sin` activation function of the linear layers is another hyperparameter of the model.\n",
    "\n",
    "The window function is the same as described above except, at least in the [reference implementation](https://github.com/HazyResearch/safari/blob/main/src/models/sequence/hyena.py), the shift $b$ is an untrainable scalar.\n",
    "$$\n",
    "h = h \\cdot (e^{\\alpha \\cdot t} + b)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb27d49-4002-4e20-babe-7963876060f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEmbedding(torch.nn.Module):\n",
    "  def __init__(self, d_embed: int, max_seq_len: int):\n",
    "    assert d_embed % 2 == 1, \"only odd dimensional positional embeddings are supported\"\n",
    "    assert d_embed > 1, \"positional embedding must be at least 3\"\n",
    "    super().__init__()\n",
    "    \n",
    "    t = torch.linspace(start=0, end=1, steps=max_seq_len)[:, None]\n",
    "\n",
    "    tp = torch.linspace(start=0, end=max_seq_len - 1, steps=max_seq_len)[:, None]\n",
    "    K = (d_embed - 1) // 2   \n",
    "    k = torch.linspace(start=0, end=K - 1, steps=K)[None, :]\n",
    "    z = torch.exp(2 * math.pi * 1j * k * tp / max_seq_len)\n",
    "    self.time_emb = torch.nn.Parameter(t.transpose(0, 1).unsqueeze(0), requires_grad=False)\n",
    "    self.pos_emb = torch.nn.Parameter(torch.cat([t, z.real, z.imag], dim=-1), requires_grad=True)\n",
    "\n",
    "  def forward(self, L: int) -> Tuple[torch.Tensor, torch.Tensor]: \n",
    "    return self.time_emb[:, :, :L], self.pos_emb[:L]    \n",
    "\n",
    "\n",
    "class Sin(torch.nn.Module):\n",
    "  def __init__(self, d_model: int, omega: int = 8, trainable: bool = False):\n",
    "    super().__init__()\n",
    "    self.freq = torch.nn.Parameter(omega * torch.ones(1, d_model), requires_grad=trainable)\n",
    "\n",
    "  def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "    return torch.sin(self.freq * x)\n",
    "\n",
    "\n",
    "class AuthenticWindow(torch.nn.Module):\n",
    "  def __init__(\n",
    "      self,\n",
    "      d_model: int,\n",
    "      fast_decay_pct: float = 0.3,  # Defaults from the official implementation\n",
    "      slow_decay_pct: float = 1.5,\n",
    "      target: float = 1e-2,\n",
    "      shift: float = 0.0,\n",
    "    ):\n",
    "    super().__init__()\n",
    "    self.shift = shift\n",
    "    min_decay = math.log(target) / slow_decay_pct\n",
    "    max_decay = math.log(target) / fast_decay_pct\n",
    "    self.alphas = torch.nn.Parameter(\n",
    "      torch.linspace(\n",
    "        start=min_decay,\n",
    "        end=max_decay,\n",
    "        steps=d_model\n",
    "      )[None, :, None], requires_grad=True)\n",
    "        \n",
    "  def forward(self, t, x):\n",
    "    L = x.shape[2]\n",
    "    c = torch.exp(self.alphas * t)[:, :, :L]\n",
    "    x = x * (c + self.shift)\n",
    "    return x\n",
    "      \n",
    "class AuthenticHyenaFilter(torch.nn.Module):\n",
    "  def __init__(\n",
    "      self,\n",
    "      d_model: int,\n",
    "      d_mlp: int,\n",
    "      d_embed: int,\n",
    "      N: int,\n",
    "      n_layers: int = 4,\n",
    "      max_seq_len: int = 128,\n",
    "      omega: int = 8,\n",
    "    ):\n",
    "    assert n_layers >= 2, \"n_layers must be at least 2\"\n",
    "    super().__init__()\n",
    "\n",
    "    self.N = N\n",
    "    self.d_model = d_model\n",
    "      \n",
    "    self.pos_emb = PositionalEmbedding(d_embed, max_seq_len)\n",
    "    \n",
    "    self.mlp = torch.nn.Sequential(\n",
    "      torch.nn.Linear(d_embed, d_mlp),\n",
    "      Sin(d_mlp, omega),\n",
    "    )\n",
    "    for _ in range(n_layers - 2):\n",
    "      self.mlp.append(torch.nn.Linear(d_mlp, d_mlp))\n",
    "      self.mlp.append(Sin(d_mlp, omega))\n",
    "    self.mlp.append(torch.nn.Linear(d_mlp, N * d_model, bias=False))\n",
    "\n",
    "    self.t = torch.nn.Parameter(\n",
    "      torch.linspace(\n",
    "        start=0,\n",
    "        end=1,\n",
    "        steps=max_seq_len\n",
    "      )[None, None, :], requires_grad=False)\n",
    "    self.h = torch.nn.Parameter(torch.randn((N, d_model, max_seq_len)))\n",
    "      \n",
    "    self.window = AuthenticWindow(d_model)\n",
    "\n",
    "  def forward(self, L: int) -> torch.Tensor:\n",
    "    t, z = self.pos_emb(L)\n",
    "    h = self.mlp(z)\n",
    "\n",
    "    h = h.transpose(0, 1)\n",
    "    h = h.reshape(self.N, self.d_model, L)\n",
    "    \n",
    "    h = self.h[:, :, :L]\n",
    "    h = self.window(self.t, h)\n",
    "    \n",
    "    return h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e0c8078-bff1-4ed6-af77-46f4edce8468",
   "metadata": {},
   "source": [
    "The operator definition from the paper does not have the three skip connections present in the version above. Otherwise it is the same as described above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0070388f-36f3-422e-a141-9abfbbdfd11e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AuthenticHyenaBlock(torch.nn.Module):\n",
    "  def __init__(self, config: HyenaConfig):\n",
    "    super().__init__()\n",
    "    self.proj_input = Projection(config.d_model, config.order, config.short_conv_size)\n",
    "    self.proj_output = torch.nn.Linear(config.d_model, config.d_model)\n",
    "    self.filter = AuthenticHyenaFilter(\n",
    "      config.d_model,\n",
    "      config.d_filter_mlp,\n",
    "      config.d_embed,\n",
    "      config.order,\n",
    "      config.n_filter_layers,\n",
    "      config.context_length,\n",
    "      config.omega,\n",
    "    )\n",
    "    self.dropout = torch.nn.Dropout(config.pdrop_hyena)\n",
    "    self.fft_conv = FFTConv()\n",
    "    self.B = torch.nn.Parameter(torch.randn((config.order, 1, config.d_model, 1)))\n",
    "\n",
    "  def forward(self, u: torch.Tensor) -> torch.Tensor:\n",
    "    L = u.shape[1]\n",
    "    \n",
    "    *x, v = self.proj_input(u)\n",
    "    \n",
    "    h = self.filter(L)\n",
    "\n",
    "    # The reference code for the paper does the product with x_i first\n",
    "    # but we follow the paper eq (6) here in putting it after the convolution\n",
    "    for i, x_i in enumerate(x):\n",
    "      h_i = h[i].unsqueeze(0)\n",
    "      v = x_i * self.fft_conv(h_i, v, self.B[i])\n",
    "    \n",
    "    v = v.transpose(1, 2)\n",
    "    y = self.proj_output(v)\n",
    "\n",
    "    return y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8663cfd7-70d4-48de-a184-f095299ef866",
   "metadata": {},
   "source": [
    "Now we can train the hyena-based model. We reduce the number of layers due to vanishing/exploding gradients and the epochs due to overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7322e239-5e16-4988-b540-90154b5e1472",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyena_config = HyenaConfig(\n",
    "  d_model=386,\n",
    "  n_layers=2,\n",
    "  vocab_size=len(vocabulary),\n",
    "  d_embed=33,\n",
    "  d_filter_mlp=64,\n",
    "  n_filter_layers=4,\n",
    "  context_length=CONTEXT_LENGTH,\n",
    "  short_conv_size=3,\n",
    "  order=2,\n",
    "  pdrop_hyena=0.0,\n",
    "  pdrop_embed=0.2,\n",
    "  omega=12,\n",
    "  epochs=10,\n",
    "  learning_rate=6e-4,\n",
    "  betas=(0.9, 0.98),\n",
    "  weight_decay=1,\n",
    "  device_type=\"gpu\",  # cpu, gpu\n",
    "  precision=\"bf16\",  # 32, 16, 16-mixed, bf16\n",
    "  batch_size=64,\n",
    "  num_workers=4,\n",
    ")\n",
    "\n",
    "authentic_hyena_model = GPModel(hyena_config, AuthenticHyenaBlock)\n",
    "authentic_hyena_model = train(authentic_hyena_model, hyena_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4418ecfa-8d9a-47f6-abc4-5be094cbecbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "  \"hyena average probability of next token:\",\n",
    "  average_probability(authentic_hyena_model, val_ds, tokens=256)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635d7d18-fe5d-4a45-8d46-0f80c652cae6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "  generate(\n",
    "    authentic_hyena_model,\n",
    "    hyena_config.context_length,\n",
    "    \"Wherefore art thou \",\n",
    "    method=\"topk\",\n",
    "    max_tokens=100,\n",
    "    k=2\n",
    "  )\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}