examples/parametrisations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic imports and some hyperparameters\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Matrices of size 4 x 4 or 4 x 5\n",
    "N = 4\n",
    "M = 5\n",
    "# Batch size of 3\n",
    "B = 3\n",
    "x = torch.rand(B, N, N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parametrisations ([PR #33344](https://github.com/pytorch/pytorch/pull/33344))\n",
    "\n",
    "This notebook provides an introduction to the design of parametrisations in PyTorch. Parametrisations are the way `geotorch` works behind the scenes, so having some grip on how they work should greatly help in using `geotorch` effectively.\n",
    "\n",
    "## Motivating Example\n",
    "\n",
    "Given a function `f` and a `Parameter` `X` which is registered on a module, we would like to be able to use `f(X)` in place of `X`.\n",
    "\n",
    "This is easier understood with an example. Suppose that we want to have a linear layer whose matrix is symmetric. We could write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Symmetric(nn.Module):\n",
    "    def __init__(self, n_features):\n",
    "        super().__init__()\n",
    "        self.weight = nn.Parameter(torch.rand(n_features, n_features))\n",
    "\n",
    "    def forward(self, x):\n",
    "        A = self.weight.triu()\n",
    "        A = A + A.T\n",
    "        #print(A)           # A is symmetric\n",
    "        return x @ A\n",
    "layer = Symmetric(N);\n",
    "layer(x);  # It works as expected"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This implementation has clearly two components. A reimplmenentation of `nn.Linear` and a parametrisation of the symmetric matrices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SymmetricParametrization(nn.Module):\n",
    "    def forward(X):\n",
    "        A = X.triu()\n",
    "        return A + A.T "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective\n",
    "\n",
    "We would like to separate these two, and have a mechanism to be able to inject a parametrisation onto a parameter or a buffer in a neural network. In particular, we would like to be able to do the following:\n",
    "\n",
    "```python\n",
    "layer = nn.Linear(N, N)\n",
    "torch.register_parametrization(layer, \"weight\", SymmetricParametrization())\n",
    "# layer now behaves as an object from the `Symmetric` class\n",
    "print(layer.weight)  # Prints the symmetric matrix\n",
    "layer(x)             # Multiplies the vectors `x` by the symmetric matrix layer.weight \n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "### Symmetric layers\n",
    "\n",
    "(see above)\n",
    "\n",
    "### Pruning\n",
    "When doing pruning, one samples a boolean mask of the size of the parameter and does an element-wise multiplication. It seems that one may train a neural network and then make it somewhat sparse, and everything magically works. This is called the \"lottery ticket hypothesis\". (see `torch.nn.utils.prune`)\n",
    "\n",
    "A simple pruning method that prunes an entry of the tensor with some given probability could go as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PruningParametrization(nn.Module):\n",
    "    def __init__(self, X, p_drop=0.2):\n",
    "        # sample zeros with probability p_drop\n",
    "        mask = torch.full_like(X, 1.0 - p_drop)\n",
    "        self.mask = torch.bernoulli(mask)\n",
    "\n",
    "    def forward(self, X):\n",
    "        return X * self.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would like to use it as:\n",
    "```python\n",
    "cnn = nn.Conv2D(8, 16, (3, 3))\n",
    "torch.register_parametrization(cnn, \"weight\", PruningParametrization(cnn.weight, p_drop=0.1))\n",
    "# 10% of the entires of the tensor cnn.weight have now been zeroed out\n",
    "```\n",
    "### Other examples:\n",
    "- `torch.weight_norm`\n",
    "- `torch.spectral_norm` (to regularise the Lipschitz constant of a layer)\n",
    "- Optimisation with orthogonal constraints / invertible layers / Symmetric Positive Defininite layers... More on this later\n",
    "\n",
    "## Implementing `torch.register_parametrization`\n",
    "### A first approximation\n",
    "A moment's reflection shows that it is possible to implement `Symmetric` without having to reimplement `nn.Linear` by using inheritance and properties.\n",
    "\n",
    "```python\n",
    "class SymmetricRevisited(nn.Linear):\n",
    "    def __init__(self, n_features):\n",
    "        super().__init__(n_features, n_features, bias=False)\n",
    "        # Rename weight attribute to _weight\n",
    "        self._weight = self.weight\n",
    "        delattr(self, \"weight\")\n",
    "    \n",
    "    @property\n",
    "    def weight(self):\n",
    "        A = self._weight.triu()\n",
    "        return A + A.T\n",
    "```\n",
    "\n",
    "Note: This code does not work! It is possible to make it work using metaclasses (for example), but we will skip that.\n",
    "\n",
    "### A caching system\n",
    "\n",
    "Sometimes we use the same layer many times in the forward pass of a neural network (e.g., in the recurrent kernel of an RNN). In those cases, we would not want to recompute `layer.weight` every time we execute it. We would like to compute it a the beginning of the forward pass and cache the result throughout the whole forward pass.\n",
    "\n",
    "We can achieve that by implementing a caching system as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "_cache_enabled = 0\n",
    "_cache = {}\n",
    "\n",
    "@contextmanager\n",
    "def cached():\n",
    "    global _cache\n",
    "    global _cache_enabled\n",
    "    _cache_enabled += 1\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        _cache_enabled -= 1\n",
    "        if not _cache_enabled:\n",
    "            _cache = {}\n",
    "\n",
    "class SymmetricCached(nn.Module):\n",
    "    def __init__(self, n_features):\n",
    "        super().__init__()\n",
    "        # Rename weight attribute to _weight\n",
    "        self._weight = nn.Parameter(torch.rand(n_features, n_features))\n",
    "        \n",
    "    def parametrization(self, X):\n",
    "        print(\"Computing\")\n",
    "        A = X.triu()\n",
    "        return A + A.T\n",
    "\n",
    "    @property\n",
    "    def weight(self):\n",
    "        global _cache\n",
    "\n",
    "        key = (id(self), \"weight\")\n",
    "        if key not in _cache:\n",
    "            _cache[key] = self.parametrization(self._weight)\n",
    "        return _cache[key]\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x @ self.weight.T\n",
    "\n",
    "# Usage:\n",
    "layer = SymmetricCached(N)\n",
    "with cached():\n",
    "    # Just computes the parametrization once\n",
    "    print(layer.weight - layer.weight.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### A generic implementation\n",
    "\n",
    "Now, all we need to do is to implement a function that, given a module, a name, and a parametrisation (i.e., another module), injects a property similar to how we did it manually in `SymmetricCached`. In particular, we have to write a function with signature\n",
    "```python\n",
    "def register_parametrization(module: Module, tensor_name: str, parametrization: Module) -> None:\n",
    "```\n",
    "that does:\n",
    "\n",
    "- Rename the tensor from `tensor_name` to `f\"_{tensor_name}\"`\n",
    "- Saves `parametrization` within `module` to use it in the forward pass\n",
    "- Injects a property with the name `tensor_name` that computes `parametrization(module[tensor_name])` when called\n",
    "\n",
    "The first two things are direct. To implement the third one, we use the `type` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inject_property(module, tensor_name):\n",
    "    # We create a new class so that we can inject properties in it\n",
    "    cls_name = \"Parametrized\" + module.__class__.__name__\n",
    "\n",
    "    # Define the getter\n",
    "    def getter(module):\n",
    "        global _cache\n",
    "\n",
    "        key = _key(module, tensor_name)\n",
    "        # If the _cache is not enabled or the caching was not enabled for this\n",
    "        # tensor, this function just evaluates the parametrization\n",
    "        if _cache_enabled and key in _cache:\n",
    "            if _cache[key] is None:\n",
    "                _cache[key] = module.parametrizations[tensor_name]()\n",
    "            return _cache[key]\n",
    "        else:\n",
    "            return module.parametrizations[tensor_name]()\n",
    "\n",
    "    # Define the setter\n",
    "    def setter(module, value):\n",
    "        module.parametrizations[tensor_name].initialize(value)\n",
    "        \n",
    "    # Create a new class that inherits from `module.__class__` and has a property called `tensor_name`\n",
    "    param_cls = type(cls_name, (module.__class__,), {\n",
    "        tensor_name: property(getter, setter)\n",
    "    })\n",
    "    module.__class__ = param_cls\n",
    "\n",
    "layer = nn.Linear(3, 4)\n",
    "inject_property(layer, \"weight\")\n",
    "print(type(layer))\n",
    "print(type(layer).weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other things that `torch.register_parametrization` implements:\n",
    "\n",
    "- If the module implements a `right_inverse` method (similar to a right-inverse of forward, more on this below), it allows initialising the parametrised buffer/parameter\n",
    "- It allows putting several parametrisations on the same buffer/parameter\n",
    "- It allows removing the parametrisations and leave the original parameter or the parametrised parameter\n",
    "- Any combination of the above\n",
    "\n",
    "## More applications of parametrizations\n",
    "\n",
    "- Constrained optimisation on manifold using `geotorch`!\n",
    "- Normalising flows. The `right_inverse` method can be implemented as a right-inverse of forward. \n",
    "    - In the simplest case, if `forward` is a diffeomorphism, then this reduces to the usual normalising flows framework.\n",
    "    - The general case comes when the forward is a [submersion](https://en.wikipedia.org/wiki/Submersion_(mathematics)) (a function with differentiable local right-inverses). An example of this is a linear layer from `R^n` to `R^k` with `n > k` that is full rank (e.g. a `k x n` matrix with orthogonal rows). Using a submersion, one may construct a generalisation of normalising flows that allows for dimensionality reduction. The simplest case of this setting comes from projecting a vector in `R^n` onto its first `k` compontents. This is called in the normalising flows literature \"multi-scale architecture\", and it was introduced in the model [real NVP](https://arxiv.org/abs/1605.08803).\n",
    "    \n",
    "## Examples of some simple parametrisations, composing them, and initialising them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This part assumes that you have `geotorch` installed. You can install it doing\n",
    "# pip install git+https://github.com/Lezcano/geotorch/\n",
    "import geotorch.parametrize as P\n",
    "\n",
    "class Skew(nn.Module):\n",
    "    def forward(self, X):\n",
    "        X = X.triu(1)\n",
    "        return X - X.T\n",
    "\n",
    "    def is_skew(self, X):\n",
    "        return torch.allclose(X, -X.T)\n",
    "\n",
    "    def right_inverse(self, X):\n",
    "        if not self.is_skew(X):\n",
    "            raise ValueError(\"This  matirx is not skew-symmetric!\")\n",
    "        return X.triu(1)\n",
    "    \n",
    "# Skew.forward(Skew.right_inverse(X)) == X\n",
    "# In functional notation: Skew.forward o Skew.right_inverse = Id\n",
    "# In other words, right_inverse is a right inverse of forward.\n",
    "\n",
    "model = nn.Linear(5, 5)\n",
    "P.register_parametrization(model, \"weight\", Skew())\n",
    "# Just computes `model.weight` once\n",
    "with P.cached():\n",
    "    assert(torch.allclose(model.weight, -model.weight.T))\n",
    "# Sample a skew matrix X and initialise the parametrised model.weight\n",
    "X = torch.rand(5,5)\n",
    "X = X - X.T\n",
    "model.weight = X\n",
    "assert(torch.allclose(model.weight, X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Orthogonal(nn.Module):\n",
    "    def __init__(self, n):\n",
    "        super().__init__()\n",
    "        self.register_buffer(\"B\", torch.eye(n))\n",
    "\n",
    "    def forward(self, A):\n",
    "        # Cayley map: (I + A)(I - A)^{-1}\n",
    "        # This is orthogonal whenever A is skew-symmetric\n",
    "        Id = torch.eye(A.size(0))\n",
    "        return self.B @ torch.solve(Id - A, Id + A).solution\n",
    "\n",
    "    def is_orthogonal(self, X):\n",
    "        Id = torch.eye(X.size(0))\n",
    "        return torch.allclose(X.T @ X, Id, atol=1e-6)\n",
    "\n",
    "    def right_inverse(self, X):\n",
    "        if not self.is_orthogonal(X):\n",
    "            raise ValueError(\"This matrix is not orthogonal!\")\n",
    "        # cayley(0) == Id, so B @ cayley(0) == B\n",
    "        self.B = X\n",
    "        return torch.zeros_like(X)\n",
    "\n",
    "\n",
    "model = nn.Linear(5,5)\n",
    "P.register_parametrization(model, \"weight\", Skew())\n",
    "P.register_parametrization(model, \"weight\", Orthogonal(5))\n",
    "\n",
    "# Sample an orthogonal matrix and initialise the layer\n",
    "X = torch.empty_like(model.weight)\n",
    "nn.init.orthogonal_(X)\n",
    "model.weight = X\n",
    "\n",
    "# model.weight == X\n",
    "assert(torch.allclose(model.weight, X))\n",
    "\n",
    "# A more programmatic way of initialising the weight\n",
    "model.weight = nn.init.orthogonal_(model.weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}