【Pytorch】FM推导及其实现

因子分解机(Factorization Machine, FM, 2010年)是由Steffen Rendle提出的一种基于矩阵分解的机器学习算法。最大的特点是易于整合交叉特征、可以处理高度稀疏数据，主要应用在推荐系统及广告CTR预估等领域。

数理推导

FM的原始的模型方程为：
$$\hat{y}(x):=w_0+\sum _{i=1}^n{w}_i{x}_i+\sum _{i=1}^n\sum _{j=i+1}^n⟨v_i,v_j⟩x_i{x}_j$$
这个式子的前两项就是一个简单的线性函数，这没什么好说的。接下来主要说一下最后这一项：
$$\sum _{i=1}^n\sum _{j=i+1}^n⟨v_i,v_j⟩x_i{x}_j$$
如果直接按照上面这个公式计算的话，复杂度就是$O(n^2)$。可以对其进行矩阵分解，优化成复杂度为$O(kn)$ 的线性复杂度，推导过程如下：
$$\begin{gathered} \sum _{i=1}^n\sum _{j=i+1}^n⟨v_i,v_j⟩x_i{x}_j \\ =\frac{1}{2}\sum _{i=1}^n\sum _{j=1}^n⟨v_i,v_j⟩x_i{x}_j-\frac{1}{2}\sum _{i=1}^n⟨v_i,v_i⟩x_i{x}_i \\ =\frac{1}{2}(\sum _{i=1}^n\sum _{j=1}^n\sum _{f=1}^k{v}_{if}{v}_{jf}{x}_i{x}_j-\sum _{i=1}^n\sum _{f=1}^k{v}_{if}{v}_{if}{x}_i{x}_i) \\ =\frac{1}{2}\sum _{f=1}^k\left[(\sum _{i=1}^n{v}_{if}{x}_i)(\sum _{j=1}^n{v}_{jf}{x}_j)-\sum _{i=1}^n{v}_{if}^2{x}_i^2\right] \\ =\frac{1}{2}\sum _{f=1}^k\left[(\sum _{i=1}^n{v}_{if}{x}_i{)}^2-\sum _{i=1}^n{v}_{if}^2{x}_i^2\right] \end{gathered}$$
其中，$⟨v_i,v_j⟩={\sum }_{f=1}^k{v}_{if}{v}_{jf}$

总的FM方程就是：
$$\hat{y}(x):=w_0+\sum _{i=1}^n{w}_i{x}_i+\frac{1}{2}\sum _{f=1}^k\left[(\sum _{i=1}^n{v}_{if}{x}_i{)}^2-\sum _{i=1}^n{v}_{if}^2{x}_i^2\right]$$

实现过程

稠密型数值特征

对于原始数据的特征是数值型的任务，可以直接使用上述公式，实现过程如下：

1、首先将上述式子改写成矩阵相乘的格式，方便后续编码实现，如下。

设输入的每一个样本可以表示成：
$$X=(x_1,x_2,...,x_i,...,x_n)\in {\mathbb{R}}^{1\times n}$$
设可学习的权重矩阵：
$$W=\left(\begin{array}{c}{w}_1\\ w_2\\ ⋮\\ w_i\\ ⋮\\ w_n\end{array}\right)\in {\mathbb{R}}^{n\times 1}$$
和
$$V=\left(\begin{array}{cccccc}{v}_{11}& v_{12}& \cdots & v_{1f}& \cdots & v_{1k}\\ v_{21}& v_{22}& \cdots & v_{2f}& \cdots & v_{2k}\\ ⋮& ⋮& \ddots & ⋮& \ddots & ⋮\\ v_{i1}& v_{i2}& \cdots & v_{if}& \cdots & v_{ik}\\ ⋮& ⋮& \ddots & ⋮& \ddots & ⋮\\ v_{n1}& v_{n2}& \cdots & v_{nf}& \cdots & v_{nk}\end{array}\right)\in {\mathbb{R}}^{n\times k}$$
那么矩阵形式的FM模型方程就是：
$$\begin{gathered} \hat{y}(x):=w_0+\sum _{i=1}^n{w}_i{x}_i+\frac{1}{2}\sum _{f=1}^k\left[(\sum _{i=1}^n{v}_{if}{x}_i{)}^2-\sum _{i=1}^n{v}_{if}^2{x}_i^2\right] \\ =w_0+XW+\frac{1}{2}\sum _{f=1}^k\left\{{\left[\left(\begin{array}{c}{x}_1,x_2,...x_n\end{array}\right)\left(\begin{array}{c}{v}_{1f}\\ v_{2f}\\ ⋮\\ v_{nf}\end{array}\right)\right]}^2-\left(\begin{array}{c}{x}_1^2,x_2^2,...x_n^2\end{array}\right)\left(\begin{array}{c}{v}_{1f}^2\\ v_{2f}^2\\ ⋮\\ v_{nf}^2\end{array}\right)\right\} \\ =w_0+XW+\frac{1}{2}sum((XV)\circ (XV)-(X\circ X)(V\circ V),axis=1) \end{gathered}$$
其中，$\circ$表示哈达玛积，即两个同阶矩阵对应元素相乘。

2、Pytorch代码实现：

import torch
import torch.nn as nn


class FactorizationMachine(nn.Module):
    def __init__(self, n, k):
        super(FactorizationMachine, self).__init__()
        self.n = n
        self.k = k
        self.linear = nn.Linear(self.n, 1, bias=True)
        self.v = nn.Parameter(torch.Tensor(self.k, self.n))  # 注：权重矩阵是(k,n)的，与公式里的相反，目的是下一步能在n的维度上分布初始化
        nn.init.xavier_uniform_(self.v)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(b, n)``
        :return: Long tensor of size ``(b, 1)``
        """
        x1 = self.linear(x)
        square_of_sum = torch.mm(x, self.v.T) * torch.mm(x, self.v.T)
        sum_of_square = torch.mm(x * x, self.v.T * self.v.T)
        x2 = 0.5 * torch.sum((square_of_sum - sum_of_square), dim=-1, keepdim=True)
        x = x1 + x2
        return x

稀疏型类值特征

1、类别型特征不好直接送入fm方程，需要先将其转换成稠密型嵌入向量。如下：

仍然设输入的每一个样本为：
$$X=(x_1,x_2,...,x_i,...,x_n)\in {\mathbb{R}}^{1\times n}$$
设嵌入函数为：
$$\begin{gathered} E_{w_i}(x_i)=embedding\_lookup(w_i,x_i)\in {\mathbb{R}}^{1\times 1} \\ 则：E_W(X)\in {\mathbb{R}}^{1\times n\times 1} \end{gathered}$$
和
$$\begin{gathered} E_{v_i}(x_i)=embedding\_lookup(v_i,x_i)\in {\mathbb{R}}^{1\times emb\_dim} \\ 则：E_V(X)\in {\mathbb{R}}^{1\times n\times emb\_dim} \end{gathered}$$
其中，$x_i\in N$ 表示第$i$个特征的类值，是自然数。

则FM方程可以表示为：
$$\begin{gathered} \hat{y}(x):=w_0+\sum _{i=1}^n{E}_{w_i}(x_i)+\frac{1}{2}\sum _{f=1}^{emb\_dim}\left[(\sum _{i=1}^n{E}_{v_{if}}(x_i){)}^2-\sum _{i=1}^n{E}_{v_{if}}(x_i{)}^2\right] \\ =w_0+sum(E_W(X),axis=1)+ \\ \frac{1}{2}sum(sum(E_V(X),axis=1)\circ sum(E_V(X),axis=1)-sum(E_V(X)\circ E_V(X),axis=1),axis=1) \end{gathered}$$
2、Pytorch实现：

源代码来自https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/fm.py，这里整合了一下：

import torch


class FeaturesLinear(torch.nn.Module):
    def __init__(self, field_dims, output_dim=1):
        super().__init__()
        self.fc = torch.nn.Embedding(sum(field_dims), output_dim)
        self.bias = torch.nn.Parameter(torch.zeros((output_dim,)))
        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        return torch.sum(self.fc(x), dim=1) + self.bias
    
    
class FeaturesEmbedding(torch.nn.Module):
    def __init__(self, field_dims, embed_dim):
        super().__init__()
        self.embedding = torch.nn.Embedding(sum(field_dims), embed_dim)
        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)
        torch.nn.init.xavier_uniform_(self.embedding.weight.data)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        return self.embedding(x)
    
    
class FactorizationMachine(torch.nn.Module):
    def __init__(self, reduce_sum=True):
        super().__init__()
        self.reduce_sum = reduce_sum

    def forward(self, x):
        """
        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``
        """
        square_of_sum = torch.sum(x, dim=1) ** 2
        sum_of_square = torch.sum(x ** 2, dim=1)
        ix = square_of_sum - sum_of_square
        if self.reduce_sum:
            ix = torch.sum(ix, dim=1, keepdim=True)
        return 0.5 * ix
    
    
# 接口
class FactorizationMachineModel(torch.nn.Module):
    """
    A pytorch implementation of Factorization Machine.

    Reference:
        S Rendle, Factorization Machines, 2010.
    """

    def __init__(self, field_dims, embed_dim):
        super().__init__()
        self.embedding = FeaturesEmbedding(field_dims, embed_dim)
        self.linear = FeaturesLinear(field_dims)
        self.fm = FactorizationMachine(reduce_sum=True)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = self.linear(x) + self.fm(self.embedding(x))
        return torch.sigmoid(x.squeeze(1))