Model implementation

Factorization Machines

Factorization Machines

Factorization Machines (Rendle, 2010) — generalize MF to arbitrary feature pairs, not just (user, item). Predict from a sparse feature vector \mathbf{x} via:

\hat y(\mathbf{x}) = w_0 + \sum_i w_i x_i + \sum_{i<j} \langle \mathbf{v}_i, \mathbf{v}_j \rangle x_i x_j.

• Linear term — like logistic regression.
• Pairwise term — every pair of features contributes a bilinear interaction, with each feature represented by a k-dim latent vector \mathbf{v}_i (just like an embedding).

The crucial trick: the pairwise sum can be computed in \mathcal{O}(kn) instead of \mathcal{O}(n^2) via:

\sum_{i<j} \langle \mathbf{v}_i, \mathbf{v}_j \rangle x_i x_j = \tfrac{1}{2} \sum_f \big[ (\sum_i v_{i,f} x_i)^2 - \sum_i v_{i,f}^2 x_i^2 \big].

Powerful for CTR prediction on sparse one-hot ad features. Generalizes MF (with two features = user + item) and recovers logistic regression as a special case.

Efficient optimization

from d2l import torch as d2l
import torch
from torch import nn
import os

class FM(nn.Module):
    def __init__(self, field_dims, num_factors):
        super().__init__()
        num_inputs = int(sum(field_dims))
        self.embedding = nn.Embedding(num_inputs, num_factors)
        self.fc = nn.Embedding(num_inputs, 1)
        self.linear_layer = nn.Linear(1, 1)

    def forward(self, x):
        square_of_sum = self.embedding(x).sum(dim=1) ** 2
        sum_of_square = (self.embedding(x) ** 2).sum(dim=1)
        x = self.linear_layer(self.fc(x).sum(dim=1)) \
            + 0.5 * (square_of_sum - sum_of_square).sum(dim=1, keepdim=True)
        return x

CTR (advertising) dataset

Standard sparse-features benchmark — many one-hot categorical fields per row:

batch_size = 2048
data_dir = d2l.download_extract('ctr')
train_data = d2l.CTRDataset(os.path.join(data_dir, 'train.csv'))
test_data = d2l.CTRDataset(os.path.join(data_dir, 'test.csv'),
                           feat_mapper=train_data.feat_mapper,
                           defaults=train_data.defaults)
train_iter = torch.utils.data.DataLoader(
    train_data, shuffle=True, drop_last=True, batch_size=batch_size,
    num_workers=d2l.get_dataloader_workers())
test_iter = torch.utils.data.DataLoader(
    test_data, shuffle=False, drop_last=True, batch_size=batch_size,
    num_workers=d2l.get_dataloader_workers())

Training

Binary cross-entropy + Adam:

devices = d2l.try_all_gpus()
net = FM(train_data.field_dims, num_factors=20)

def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
    if type(m) == nn.Embedding:
        nn.init.xavier_uniform_(m.weight)

net.apply(init_weights)
lr, num_epochs = 0.02, 30
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
loss = nn.BCEWithLogitsLoss(reduction='none')
d2l.train_ch13(net, train_iter, test_iter, loss, optimizer, num_epochs, devices)

loss 0.003, train acc 1.000, test acc 0.928
341844.0 examples/sec on [device(type='cuda', index=0)]

Recap

• FMs = linear model + bilinear feature interactions, all feature pairs share latent factor structure.
• Closed-form \mathcal{O}(kn) pairwise computation makes FMs practical on sparse features with millions of fields.
• Generalizes MF; foundation for DeepFM (next deck) and many production CTR models.