grad_matrix_products.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Different ways of computing gradient of trace of infinite matrix products\n",
    "$\\nabla_T \\ln Z = \\nabla _T \\ln \\mathrm{tr}\\left[\\underbrace{T\\times T \\times \\cdots \\times T}_{\\mathrm{\\infty}}\\right]$\n",
    "We present several methods for computing the gradients with:\n",
    "- Autograd of $\\ln Z$ that is computed using the environment matrix and RG\n",
    "- Autograd of the leading eigenvalue of the transfer matrix\n",
    "- Autograd of $\\ln Z$ that is computed using the environment vector and power method\n",
    "- Using cavity of the matrix product\n",
    "- Using the leading eigenvector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.manual_seed(1)\n",
    "N = 10 # size of the matrix\n",
    "A = torch.rand(N, N, dtype=torch.float64) \n",
    "T = torch.nn.Parameter(A@A.t()) # a symmetric positive definite transfer matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autograd of $\\ln Z$ using environment matrix\n",
    "Consider a infinite-length chain as product of matrices $T$. Analogous to real-space RG for the one-dimensional Ising model, we merge all pairs of consequtive matrices at each RG step, resulting to a chain with half length (although still infinite), then truncate the spectrum of merged matrices to original dimension for a low-rank approximation. By keep doing this RG step, merged matrices will finally converge to environments (for each tensor $T$) $M$ which alsoindicates that $Z=M\\times T \\times M$, and $$ \\ln Z=\\ln(M\\times T \\times M).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "lnZ = 0.0\n",
    "M = T.clone() # M will be the converged envioment tensor\n",
    "Niter = 20 # number of iterations in RG. That is, 2^Niter matrices will be contracted finally.\n",
    "for i in range(Niter): # after i steps, there are totally 2^i matrices contracted.\n",
    "    s = M.norm() # This is the normalization of a matrix contracted of 2^i T.\n",
    "    lnZ = lnZ + torch.log(s)/2**i # Notice that we can only record a density of logarithm of the results, for contraction of infinite matrices.\n",
    "    M = M/s\n",
    "    M = M@M\n",
    "lnZ = lnZ + torch.trace(M)/(2**Niter) # trace(M) is the trace of contraction of all tensors.\n",
    "\n",
    "lnZ.backward()\n",
    "RG_grad = T.grad.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autograd of the leading eigenvalue of the transfer matrix\n",
    "In the last section, we have evaluated $\\ln Z$ using RG. Actually, this can be evaluated analytically as\n",
    "$$\\ln Z =\n",
    "\\frac{1}{t}\\lim_{t\\to\\infty}\\ln \\mathrm{tr}(T^t)=\\frac{1}{t}\\lim_{t\\to\\infty}\\ln\\sum_{i=1}^N \\lambda_i^t=\n",
    "\\ln\\lambda_\\mathrm{max},$$\n",
    "where $\\lambda_{\\mathrm{max}}$ is the leading eigenvalue. Thus we can do back propagation directly on $\\lambda_{\\mathrm{max}}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "w, v = torch.symeig(T, eigenvectors=True) # W is an arry of eigenvalues, v is a matrix with each row storing an eigenvector\n",
    "T.grad.zero_()\n",
    "lambda_max = torch.log(w[-1])\n",
    "lambda_max.backward()\n",
    "eigenvalue_grad = ((T.grad + T.grad.t())/2) # need to symmetrize since it is an upper triangular matrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using cavity of the matrix product\n",
    "Using the property of the environment matrix $M$, we have $Z=\\mathrm{tr}\\left[\\underbrace{T\\times T \\times \\cdots \\times T}_{\\mathrm{\\infty}}\\right]=\\mathrm{Tr}(M\\times T \\times M)$. Then the gradient with respect to a specific $T$ can be written as\n",
    "$$\\nabla_T \\ln Z = \\frac{M\\times M}{\\mathrm{tr}(M\\times T\\times M)}.$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "impurity_grad = (M@M).t()/torch.trace(M@T@M)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autograd of $\\ln Z$ using vector environment\n",
    "Here comes the trick: Infinite tensor product with a Periodic Boundary Condition (PBC) is equivalent to that with a Open Boundary Condition (OBC), which is written as\n",
    "$\\ln Z=\\ln \\left[u\\times\\underbrace{T\\times T \\times \\cdots \\times T}_{\\mathrm{\\infty}}\\times u\\right]=\\ln(m\\times T\\times m)$, where $m$ is the environment for the OBC tensor product, a vector. Obviously, $m$ is the leading eigenvector of the transfer matrix $T$ hence is normalized, and the gradient can be written as $$\\nabla_T\\ln Z=\\exp(-\\lambda_\\mathrm{max}) m\\otimes m.$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "v=v[:,-1] # the leading eigenvector of the transfer matrix $T$\n",
    "eigenvector_grad=v[:,None]@v[None,:]/torch.exp(lambda_max) # outer product of the leading eigenvector and its transpose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparing these gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.168404344971009e-18\n",
      "3.469446951953614e-18\n",
      "4.336808689942018e-18\n"
     ]
    }
   ],
   "source": [
    "# comparing RG_grad, eigenvector_grad, eigenvalue_grad and impurity_grad\n",
    "print ((impurity_grad-RG_grad).abs().max().item())\n",
    "print ((impurity_grad-eigenvalue_grad).abs().max().item())\n",
    "print ((impurity_grad-eigenvector_grad).abs().max().item())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}