{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.ultragcn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# UltraGCN\n",
    "> An efficient graph-convolutional recommendation model.\n",
    "\n",
    "Industrial recommender systems usually involve massive graphs due to the large numbers of users and items. However, current GCN-based models are hard to train with large graphs, which hinders their wide adoption in industry. This brings efficiency and scalability challenges for model designs. Some efforts have been made to simplify the design of GCN-based CF models, mainly by removing feature transformations and non-linear activations that are not necessary for CF. These proved beneficial. But it has been seen that message passing (i.e., neighborhood aggregation) on a large graph is usually time-consuming for CF. In particular, stacking multiple layers of message passing could lead to the slow convergence of GCN-based models on CF tasks. For example, in our experiments, three-layer LightGCN takes more than 700 epochs to converge to its best result on the Amazon Books dataset, which would be unacceptable in an industrial setting.\n",
    "\n",
    "UltraGCN has the following optimization objective:\n",
    "\n",
    "$$\\mathcal{L} = \\mathcal{L}_O + \\lambda \\mathcal{L}_C$$\n",
    "\n",
    "where 𝜆 is the hyper-parameter to control the importance weights of two loss terms.\n",
    "\n",
    "Moreover, except for user-item relationships, some other relationships (e.g., item-item and user-user relationships) also greatly contribute to the effectiveness of GCN-based models on CF. However, in conventional GCN-based models, these relationships are implicitly learned through the same message passing layers with user-item relationships. This not only leads to the unreasonable edge weight assignments, but also fails to capture the relative importances of different types of relationships.\n",
    "\n",
    "<img src='https://github.com/recohut/reco-static/raw/master/media/images/models/ultragcn.png'>\n",
    "\n",
    "### Constraint Loss\n",
    "\n",
    "Instead of performing explicit message passing, LightGCN aims to directly approximate the convergence state by normalizing the embeddings to unit vectors and then maximize the dot product of both terms, which is equivalent to maximize the cosine similarity between $𝑒_u$ and $e_i$.\n",
    "\n",
    "$$\\max \\sum_{i \\in \\mathcal{N}(u)} \\beta_{u,i}e_u^Te_i,\\ \\forall u \\in U$$\n",
    "\n",
    "For ease of optimization, we further incorporate sigmoid activation and negative log likelihood, and derive the following loss:\n",
    "\n",
    "$$\\mathcal{L}_C = - \\sum_{u \\in U} \\sum_{i \\in \\mathcal{N}(u)} \\beta_{u,i} \\log(\\sigma(e_u^Te_i)),$$\n",
    "\n",
    "To avoid over-smoothing (users and items could easily converge to the same embeddings), UltraGCN perform negative sampling (inspired by Word2vec). The constraint loss would then become:\n",
    "\n",
    "$$\\mathcal{L}_C = - \\sum_{(u,i) \\in N^+} \\beta_{u,i} \\log(\\sigma(e_u^Te_i)) - \\sum_{(u,j) \\in N^-} \\beta_{u,j} \\log(\\sigma(e_u^Te_j))$$\n",
    "\n",
    "### Optimization Loss\n",
    "\n",
    "Typically, CF models perform item recommendation by applying either pairwise BPR (Bayesian personalized ranking) loss or pointwise BCE (binary cross-entropy) loss for optimization. UltraGCN formulate CF as a link prediction problem in graph learning, and therefore chooses the following BCE loss as the main optimization objective. It is also consistent with the loss format of $\\mathcal{L}_C$:\n",
    "\n",
    "$$\\mathcal{L}_O = - \\sum_{(u,i) \\in N^+} \\log(\\sigma(e_u^Te_i)) - \\sum_{(u,j) \\in N^-} \\log(\\sigma(e_u^Te_j))$$\n",
    "\n",
    "### Learning on Item-Item\n",
    "\n",
    "Moreover, except for user-item relationships, some other relationships (e.g., item-item and user-user relationships) also greatly contribute to the effectiveness of GCN-based models on CF. However, in conventional GCN-based models, these relationships are implicitly learned through the same message passing layers with user-item relationships. This not only leads to the unreasonable edge weight assignments, but also fails to capture the relative importances of different types of relationships.\n",
    "\n",
    "UltraGCN is flexible to extend to model many different relation graphs, such as user-user graphs, item-item graphs, and even knowledge graphs. For now, we will focus on the item-item co-occurrence graph, which has been shown to be useful for recommendation.\n",
    "\n",
    "For each positive (𝑢, 𝑖) pair, we first construct 𝐾 weighted positive (𝑢, 𝑗) pairs, for 𝑗 ∈ 𝑆 (𝑖). Then, we penalize the learning of these pairs with the more reasonable similarity score $\\omega_{𝑖,𝑗}$ and derive the constraint loss $\\mathcal{L}_𝐼$ on the item-item graph as follow:\n",
    "\n",
    "$$\\mathcal{L}_I = \\sum_{(u,i) \\in N^+}\\sum_{j \\in S(i)} \\omega_{i,j} \\log (\\sigma(e_u^Te_j))$$\n",
    "\n",
    "We can omit the negative sampling here as the negative sampling in $\\mathcal{L}_𝐶$ and $\\mathcal{L}_O$ has already enabled UltraGCN to counteract over-smoothing. With this constraint loss, we can extend UltraGCN to better learn item-item relationships, and finally derive the following training objective of UltraGCN,\n",
    "\n",
    "$$\\mathcal{L} = \\mathcal{L}_O + \\lambda \\mathcal{L}_C + \\gamma\\mathcal{L}_I$$\n",
    "\n",
    "where 𝜆 and 𝛾 are hyper-parameters to adjust the relative importances of user-item and item-item relationships, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class UltraGCN(nn.Module):\n",
    "    def __init__(self, args, constraint_mat, ii_constraint_mat, ii_neighbor_mat):\n",
    "        super(UltraGCN, self).__init__()\n",
    "        self.user_num = args.user_num\n",
    "        self.item_num = args.item_num\n",
    "        self.embedding_dim = args.embedding_dim\n",
    "        self.w1 = args.w1\n",
    "        self.w2 = args.w2\n",
    "        self.w3 = args.w3\n",
    "        self.w4 = args.w4\n",
    "\n",
    "        self.negative_weight = args.negative_weight\n",
    "        self.gamma = args.gamma\n",
    "        self.lambda_ = args.lambda_\n",
    "\n",
    "        self.user_embeds = nn.Embedding(self.user_num, self.embedding_dim)\n",
    "        self.item_embeds = nn.Embedding(self.item_num, self.embedding_dim)\n",
    "\n",
    "        self.constraint_mat = constraint_mat\n",
    "        self.ii_constraint_mat = ii_constraint_mat\n",
    "        self.ii_neighbor_mat = ii_neighbor_mat\n",
    "\n",
    "        self.initial_weight = args.initial_weight\n",
    "\n",
    "        self.initial_weights()\n",
    "\n",
    "    def initial_weights(self):\n",
    "        nn.init.normal_(self.user_embeds.weight, std=self.initial_weight)\n",
    "        nn.init.normal_(self.item_embeds.weight, std=self.initial_weight)\n",
    "\n",
    "    def get_omegas(self, users, pos_items, neg_items):\n",
    "        device = self.get_device()\n",
    "        if self.w2 > 0:\n",
    "            pos_weight = self.constraint_mat[users * self.item_num + pos_items].to(device)\n",
    "            pow_weight = self.w1 + self.w2 * pos_weight\n",
    "        else:\n",
    "            pos_weight = self.w1 * torch.ones(len(pos_items)).to(device)\n",
    "        \n",
    "        users = (users * self.item_num).unsqueeze(0)\n",
    "\n",
    "        if self.w4 > 0:\n",
    "            neg_weight = self.constraint_mat[torch.cat([users] * neg_items.size(1)).transpose(1, 0) + neg_items].flatten().to(device)\n",
    "            neg_weight = self.w3 + self.w4 * neg_weight\n",
    "        else:\n",
    "            neg_weight = self.w3 * torch.ones(neg_items.size(0) * neg_items.size(1)).to(device)\n",
    "\n",
    "        weight = torch.cat((pow_weight, neg_weight))\n",
    "        return weight\n",
    "\n",
    "    def cal_loss_L(self, users, pos_items, neg_items, omega_weight):\n",
    "        device = self.get_device()\n",
    "        user_embeds = self.user_embeds(users)\n",
    "        pos_embeds = self.item_embeds(pos_items)\n",
    "        neg_embeds = self.item_embeds(neg_items)\n",
    "      \n",
    "        pos_scores = (user_embeds * pos_embeds).sum(dim=-1) # batch_size\n",
    "        user_embeds = user_embeds.unsqueeze(1)\n",
    "        neg_scores = (user_embeds * neg_embeds).sum(dim=-1) # batch_size * negative_num\n",
    "\n",
    "        neg_labels = torch.zeros(neg_scores.size()).to(device)\n",
    "        neg_loss = F.binary_cross_entropy_with_logits(neg_scores, neg_labels, weight = omega_weight[len(pos_scores):].view(neg_scores.size()), reduction='none').mean(dim = -1)\n",
    "        \n",
    "        pos_labels = torch.ones(pos_scores.size()).to(device)\n",
    "        pos_loss = F.binary_cross_entropy_with_logits(pos_scores, pos_labels, weight = omega_weight[:len(pos_scores)], reduction='none')\n",
    "\n",
    "        loss = pos_loss + neg_loss * self.negative_weight\n",
    "      \n",
    "        return loss.sum()\n",
    "\n",
    "    def cal_loss_I(self, users, pos_items):\n",
    "        device = self.get_device()\n",
    "        neighbor_embeds = self.item_embeds(self.ii_neighbor_mat[pos_items].to(device))    # len(pos_items) * num_neighbors * dim\n",
    "        sim_scores = self.ii_constraint_mat[pos_items].to(device)     # len(pos_items) * num_neighbors\n",
    "        user_embeds = self.user_embeds(users).unsqueeze(1)\n",
    "        \n",
    "        loss = -sim_scores * (user_embeds * neighbor_embeds).sum(dim=-1).sigmoid().log()\n",
    "      \n",
    "        # loss = loss.sum(-1)\n",
    "        return loss.sum()\n",
    "\n",
    "    def norm_loss(self):\n",
    "        loss = 0.0\n",
    "        for parameter in self.parameters():\n",
    "            loss += torch.sum(parameter ** 2)\n",
    "        return loss / 2\n",
    "\n",
    "    def forward(self, users, pos_items, neg_items):\n",
    "        omega_weight = self.get_omegas(users, pos_items, neg_items)\n",
    "        \n",
    "        loss = self.cal_loss_L(users, pos_items, neg_items, omega_weight)\n",
    "        loss += self.gamma * self.norm_loss()\n",
    "        loss += self.lambda_ * self.cal_loss_I(users, pos_items)\n",
    "        return loss\n",
    "\n",
    "    def test_foward(self, users):\n",
    "        items = torch.arange(self.item_num).to(users.device)\n",
    "        user_embeds = self.user_embeds(users)\n",
    "        item_embeds = self.item_embeds(items)\n",
    "         \n",
    "        return user_embeds.mm(item_embeds.t())\n",
    "\n",
    "    def get_device(self):\n",
    "        return self.user_embeds.weight.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-29 04:08:41\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "IPython: 5.5.0\n",
      "torch  : 1.10.0+cu111\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
