{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.ffm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FFM\n",
    "> A pytorch implementation of Field-aware Factorization Machine.\n",
    "\n",
    "Despite effectiveness, FM can be hindered by its modelling of all feature interactions with the same weight, as not all feature interactions are equally useful and predictive. For example, the interactions with useless features may even introduce noises and adversely degrade the performance.\n",
    "\n",
    "|  | For each | Learn |\n",
    "| --- | --- | --- |\n",
    "| Linear | feature | a weight |\n",
    "| Poly | feature pair | a weight |\n",
    "| FM | feature | a latent vector |\n",
    "| FFM | feature | multiple latent vectors |\n",
    "\n",
    "Field-aware factorization machine (FFM) is an extension to FM. It was originally introduced in [2]. The advantage of FFM over FM is that it uses different factorized latent factors for different groups of features. The \"group\" is called \"field\" in the context of FFM. Putting features into fields resolves the issue that the latent factors shared by features that intuitively represent different categories of information may not well generalize the correlation.\n",
    "\n",
    "FFM addresses this issue by splitting the original latent space into smaller latent spaces specific to the fields of the features.\n",
    "\n",
    "$$\\phi(\\pmb{w}, \\pmb{x}) = w_0 + \\sum\\limits_{i=1}^n w_i x_i + \\sum\\limits_{i=1}^n \\sum\\limits_{j=i + 1}^n \\langle \\mathbf{v}_{i, f_{2}} \\cdot \\mathbf{v}_{j, f_{1}} \\rangle x_i x_j$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from recohut.models.layers.common import FeaturesLinear"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FieldAwareFactorizationMachine(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim):\n",
    "        super().__init__()\n",
    "        self.num_fields = len(field_dims)\n",
    "        self.embeddings = torch.nn.ModuleList([\n",
    "            torch.nn.Embedding(sum(field_dims), embed_dim) for _ in range(self.num_fields)\n",
    "        ])\n",
    "        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)\n",
    "        for embedding in self.embeddings:\n",
    "            torch.nn.init.xavier_uniform_(embedding.weight.data)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        x = x + x.new_tensor(self.offsets).unsqueeze(0)\n",
    "        xs = [self.embeddings[i](x) for i in range(self.num_fields)]\n",
    "        ix = list()\n",
    "        for i in range(self.num_fields - 1):\n",
    "            for j in range(i + 1, self.num_fields):\n",
    "                ix.append(xs[j][:, i] * xs[i][:, j])\n",
    "        ix = torch.stack(ix, dim=1)\n",
    "        return ix\n",
    "\n",
    "class FFM(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Field-aware Factorization Machine.\n",
    "    Reference:\n",
    "        Y Juan, et al. Field-aware Factorization Machines for CTR Prediction, 2015.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim):\n",
    "        super().__init__()\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        self.ffm = FieldAwareFactorizationMachine(field_dims, embed_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        ffm_term = torch.sum(torch.sum(self.ffm(x), dim=1), dim=1, keepdim=True)\n",
    "        x = self.linear(x) + ffm_term\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- Y Juan, et al. Field-aware Factorization Machines for CTR Prediction, 2015.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/ffm.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p recohut"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
