Field-matrixed Factorization Machines implementation for pytorch

fmfm.py

# https://arxiv.org/pdf/2102.12994.pdf

from typing import Dict
from itertools import permutations
import torch
from torch import nn
from math import sqrt


class FMFM(nn.Module):
    def __init__(self,
                embedding_dims: Dict[str, int],
                output_dim: int,
                norm_dim: bool = True):
        super().__init__()

        self.field_linears = nn.ModuleDict({
                (i[0] + '_' + j[0]): nn.Linear(i[1], j[1], bias=False) \
            for i, j in permutations(embedding_dims.items(), 2)})

        self.projection = nn.Linear(sum(embedding_dims.values()) * (len(embedding_dims) - 1),
                                    output_dim)
        self.norm_dim = norm_dim

    def forward(self, x):
        if(self.norm_dim):
            for k in x.keys():
                x[k] /= sqrt(x[k].shape[-1])

        transformed, embeddings = tuple(zip(*[
                (self.field_linears[i[0] + '_' + j[0]](i[1]), j[1]) \
            for i, j in permutations(sorted(x.items()), 2)]))

        return self.projection(torch.cat(transformed, dim=-1) * \
                                torch.cat(embeddings, dim=-1))