svd.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Singular Value Decomposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "_ = torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate a rank-deficient matrix W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-1.0797,  0.5545,  0.8058, -0.7140, -0.1518,  1.0773,  2.3690,  0.8486,\n",
      "         -1.1825, -3.2632],\n",
      "        [-0.3303,  0.2283,  0.4145, -0.1924, -0.0215,  0.3276,  0.7926,  0.2233,\n",
      "         -0.3422, -0.9614],\n",
      "        [-0.5256,  0.9864,  2.4447, -0.0290,  0.2305,  0.5000,  1.9831, -0.0311,\n",
      "         -0.3369, -1.1376],\n",
      "        [ 0.7900, -1.1336, -2.6746,  0.1988, -0.1982, -0.7634, -2.5763, -0.1696,\n",
      "          0.6227,  1.9294],\n",
      "        [ 0.1258,  0.1458,  0.5090,  0.1768,  0.1071, -0.1327, -0.0323, -0.2294,\n",
      "          0.2079,  0.5128],\n",
      "        [ 0.7697,  0.0050,  0.5725,  0.6870,  0.2783, -0.7818, -1.2253, -0.8533,\n",
      "          0.9765,  2.5786],\n",
      "        [ 1.4157, -0.7814, -1.2121,  0.9120,  0.1760, -1.4108, -3.1692, -1.0791,\n",
      "          1.5325,  4.2447],\n",
      "        [-0.0119,  0.6050,  1.7245,  0.2584,  0.2528, -0.0086,  0.7198, -0.3620,\n",
      "          0.1865,  0.3410],\n",
      "        [ 1.0485, -0.6394, -1.0715,  0.6485,  0.1046, -1.0427, -2.4174, -0.7615,\n",
      "          1.1147,  3.1054],\n",
      "        [ 0.9088,  0.1936,  1.2136,  0.8946,  0.4084, -0.9295, -1.2294, -1.1239,\n",
      "          1.2155,  3.1628]])\n"
     ]
    }
   ],
   "source": [
    "d, k = 10, 10\n",
    "\n",
    "# This way we can generate a rank-deficient matrix\n",
    "W_rank = 2\n",
    "W = torch.randn(d,W_rank) @ torch.randn(W_rank,k)\n",
    "print(W)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluate the rank of the matrix W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rank of W: 2\n"
     ]
    }
   ],
   "source": [
    "W_rank = np.linalg.matrix_rank(W)\n",
    "print(f'Rank of W: {W_rank}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the SVD decomposition of the W matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of B: torch.Size([10, 2])\n",
      "Shape of A: torch.Size([2, 10])\n"
     ]
    }
   ],
   "source": [
    "# Perform SVD on W (W = UxSxV^T)\n",
    "U, S, V = torch.svd(W)\n",
    "\n",
    "# For rank-r factorization, keep only the first r singular values (and corresponding columns of U and V)\n",
    "U_r = U[:, :W_rank]\n",
    "S_r = torch.diag(S[:W_rank])\n",
    "V_r = V[:, :W_rank].t()  # Transpose V_r to get the right dimensions\n",
    "\n",
    "# Compute B = U_r * S_r and A = V_r\n",
    "B = U_r @ S_r\n",
    "A = V_r\n",
    "print(f'Shape of B: {B.shape}')\n",
    "print(f'Shape of A: {A.shape}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the same input, check the output using the original W matrix and the matrices resulting from the decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original y using W:\n",
      " tensor([ 7.2684e+00,  2.3162e+00,  7.7151e+00, -1.0446e+01, -8.1639e-03,\n",
      "        -3.7270e+00, -1.1146e+01,  2.0207e+00, -9.6258e+00, -4.1163e+00])\n",
      "\n",
      "y' computed using BA:\n",
      " tensor([ 7.2684e+00,  2.3162e+00,  7.7151e+00, -1.0446e+01, -8.1638e-03,\n",
      "        -3.7270e+00, -1.1146e+01,  2.0207e+00, -9.6258e+00, -4.1163e+00])\n"
     ]
    }
   ],
   "source": [
    "# Generate random bias and input\n",
    "bias = torch.randn(d)\n",
    "x = torch.randn(d)\n",
    "\n",
    "# Compute y = Wx + bias\n",
    "y = W @ x + bias\n",
    "# Compute y' = (B*A)x + bias\n",
    "y_prime = (B @ A) @ x + bias\n",
    "\n",
    "print(\"Original y using W:\\n\", y)\n",
    "print(\"\")\n",
    "print(\"y' computed using BA:\\n\", y_prime)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total parameters of W:  100\n",
      "Total parameters of B and A:  40\n"
     ]
    }
   ],
   "source": [
    "print(\"Total parameters of W: \", W.nelement())\n",
    "print(\"Total parameters of B and A: \", B.nelement() + A.nelement())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-lora",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}