Spline_based_Transformers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pixelsandpointers/annotated-transformer/blob/main/Spline_based_Transformers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Implementation of Spline-based Transformers\n",
        "See paper: https://la.disneyresearch.com/publication/spline-based-transformers/  \n",
        "No official implementation, so this one is based on the additional material's pseudocode."
      ],
      "metadata": {
        "id": "uXkvg2Aopiub"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hhW21LNCEnFw"
      },
      "outputs": [],
      "source": [
        "import einops\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import torchvision.datasets as datasets\n",
        "from torch.utils.data import DataLoader\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ozRZ3S4E8LN"
      },
      "outputs": [],
      "source": [
        "class Spline(nn.Module):\n",
        "    N = torch.tensor([                                                          # characteristic matrix\n",
        "        [ 1,  0,  0,  0],\n",
        "        [-3,  3,  0,  0],\n",
        "        [ 3, -6,  3,  0],\n",
        "        [-1,  3, -3,  1]\n",
        "    ])\n",
        "    def __init__(self, trainable: bool = True):\n",
        "        super(Spline, self).__init__()\n",
        "        self.param = nn.Parameter(self.N, requires_grad=trainable)              # register parameter\n",
        "\n",
        "    def forward(self, control_points: torch.Tensor, seq_len: int):\n",
        "        # assume batch size x seq_len x emb_dim\n",
        "        device = control_points.device\n",
        "        batch_size, n_points, _ = control_points.shape\n",
        "        assert n_points == 4, \"Only implements the cubic spline\"\n",
        "        linspace = torch.linspace(0, 1, steps=seq_len, device=device)           # create domain from [0, 1] to evaluate the spline\n",
        "        linspace = einops.repeat(linspace, 't -> b t', b=batch_size)            # creates batch_size x linspace\n",
        "        powers = torch.arange(0, n_points, device=device.device).T              # power terms [1, t, t^2, t^3]\n",
        "        times = einops.rearrange(linspace, '... -> ... 1') ** powers\n",
        "        return times @ self.param @ control_points\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ngDICbfjKm4p"
      },
      "outputs": [],
      "source": [
        "class SBT(nn.Module):\n",
        "    def __init__(self,\n",
        "                 n_out: int,\n",
        "                 n_dim: int = 512,\n",
        "                 n_enc: int = 6,\n",
        "                 n_dec: int = 6,\n",
        "                 n_control_points: int = 4):\n",
        "        super(SBT, self).__init__()\n",
        "        self.n_dim = n_dim\n",
        "        self.n_enc = n_enc\n",
        "        self.n_dec = n_dec\n",
        "        self.n_control_points = n_control_points\n",
        "        self.control_points = nn.Parameter(torch.zeros(n_control_points, n_dim))\n",
        "\n",
        "        # layers\n",
        "        self.emb = nn.Sequential(\n",
        "            nn.Linear(n_dim, n_dim*2),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(n_dim*2, n_dim)\n",
        "            )\n",
        "\n",
        "        # setup Transformer Encoder (T5 was used in the paper)\n",
        "        t_enc_layer = nn.TransformerEncoderLayer(n_dim, 8, activation='gelu', batch_first=True)\n",
        "        self.t_enc = nn.TransformerEncoder(t_enc_layer, num_layers=n_enc)\n",
        "\n",
        "        # setup Transformer Decoder\n",
        "        t_dec_layer = nn.TransformerDecoderLayer(n_dim, 8, activation='gelu', batch_first=True)\n",
        "        self.t_dec = nn.TransformerDecoder(t_dec_layer, num_layers=n_dec)\n",
        "\n",
        "        # setup Spline\n",
        "        self.spline = Spline(trainable=True)\n",
        "\n",
        "        # setup classification head\n",
        "        self.head = nn.Linear(n_dim, n_out)\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:                         # x = [batch_size x seq_len x n_dim]\n",
        "        batch_size, seq_len, *_ = x.shape\n",
        "        emb = self.emb(x)\n",
        "\n",
        "        # project control points\n",
        "        control_points = einops.repeat(self.control_points, 'p d -> b p d',\n",
        "                                            b=batch_size)\n",
        "\n",
        "        enc_in = torch.concat((control_points, emb), axis=1)\n",
        "        enc_out = self.t_enc(enc_in)\n",
        "\n",
        "        control_points = enc_out[:, :self.n_control_points, :]\n",
        "        latent = self.spline(control_points, seq_len)\n",
        "        dec_out = self.t_dec(latent, memory=enc_out)\n",
        "\n",
        "        return dec_out"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "VHqc1FdSUJfh"
      },
      "outputs": [],
      "source": [
        "# Hyperparameters\n",
        "input_dim = 28 * 28  # For MNIST images\n",
        "latent_dim = 20\n",
        "batch_size = 64\n",
        "lr = 1e-3\n",
        "epochs = 10\n",
        "\n",
        "# Data\n",
        "transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Lambda(lambda x: x.view(-1))  # Flatten\n",
        "])\n",
        "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
        "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "# Model, optimizer\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = SBT(input_dim, n_dim=input_dim).to(device)\n",
        "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
        "\n",
        "# Training\n",
        "for epoch in range(epochs):\n",
        "    model.train()\n",
        "    total_loss = 0\n",
        "    for batch in train_loader:\n",
        "        x, _ = batch\n",
        "        x = x.to(device)\n",
        "        x = x.unsqueeze(0)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        recon_x = model(x)\n",
        "        loss = nn.functional.mse_loss(recon_x, x)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        total_loss += loss.item()\n",
        "\n",
        "    avg_loss = total_loss / len(train_loader.dataset)\n",
        "    print(f\"Epoch {epoch + 1}/{epochs}, Loss: {avg_loss:.4f}\")\n",
        "\n",
        "print(\"Training complete!\")\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "authorship_tag": "ABX9TyM2NFYKFndrJlrY6PVM6a6R",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}