Beyond Linear Models: An In-Depth Guide to Factorization Machines

Machine learning on large, sparse datasets poses unique challenges. Such data arises frequently in applications like recommender systems, computational advertising, and analysis of user behavior logs. Naive linear models struggle to capture complex feature interactions, while models like polynomial regression that introduce interaction terms explode in computational and memory complexity.

Factorization machines (FMs), proposed by Steffen Rendle in 2010, have emerged as a powerful tool for these settings. By learning a low-rank representation of pairwise feature interactions, FMs balance expressiveness and scalability. In this post, we‘ll dive into the mathematical formulation of FMs, discuss optimization approaches, and walk through examples of using FMs for recommender systems and ad click-through rate (CTR) prediction.

The Trouble with Sparsity

To motivate FMs, let‘s consider the task of predicting user ratings for movies to power a recommendation engine. Our training data consists of a large, sparse matrix where each row represents a user, each column a movie, and entries are ratings on a 1-5 scale (with many missing values).

A standard linear regression approach models the rating as a weighted sum of user and movie features:

$\hat{y}(x) = w0 + \sum{i=1}^n w_i x_i$

Here $x_i$ might represent a user‘s age, gender, or previous watching history. However, this fails to capture interactions between features – e.g. that a certain genre resonates with one demographic but not another.

We could introduce interaction terms, resulting in a degree-2 polynomial model:

$\hat{y}(x) = w0 + \sum{i=1}^n w_i xi + \sum{i=1}^n \sum{j=i+1}^n w{ij} x_i x_j$

But this has several issues. Most glaringly, the number of parameters $w_{ij}$ grows quadratically with the feature dimension $n$. Furthermore, many interaction terms will not even appear in the sparse training data, leading to poorly determined weights.

Matrix Factorization to the Rescue

The key insight of FMs is to learn low-rank representations of feature interactions. Mathematically, this means factorizing the matrix of pairwise feature weights $W$ into low-dimensional vectors:

$W = VV^T$

where $V$ is a $n \times k$ matrix. Each feature $i$ is associated with a $k$-dimensional embedding $v_i$, and the interaction weight between features $i$ and $j$ is modeled as the dot product $\langle v_i, v_j \rangle$.

The FM model is then:

$\hat{y}(x) = w0 + \sum{i=1}^n w_i xi + \sum{i=1}^n \sum_{j=i+1}^n \langle v_i, v_j \rangle x_i x_j$

Despite having $O(nk)$ parameters, this can be computed in linear time in $n$ due to the identity:

$\sum{i=1}^n \sum{j=i+1}^n \langle v_i, v_j \rangle x_i xj = \frac{1}{2} \left[ \left( \sum{i=1}^n v_i xi \right)^2 – \sum{i=1}^n (v_i x_i)^2 \right]$

To build intuition, let‘s return to our movie ratings example. Suppose $k=5$ and the learned embeddings for the Sci-Fi genre and the Age 18-24 demographic are:

$v_{\text{Sci-Fi}} = [0.8, -0.2, 0.1, 0.5, 0.3]$

$v_{\text{Age 18-24}} = [0.6, 0.7, -0.4, 0.1, 0.2]$

Then their interaction weight is:

$\langle v{\text{Sci-Fi}}, v{\text{Age 18-24}} \rangle = 0.8 \cdot 0.6 – 0.2 \cdot 0.7 + 0.1 \cdot (-0.4) + 0.5 \cdot 0.1 + 0.3 \cdot 0.2 = 0.47$

This positive value captures that young adults rate sci-fi movies more highly than the population average. The learned embeddings map similar features to nearby vectors, revealing latent structure in the data.

Training Factorization Machines

To learn the parameters ${w_0, w_1, \dots, w_n, V}$ of an FM, we minimize an empirical loss function over the training data $\mathcal{D} = {(x^{(i)}, y^{(i)})}$:

$\min_{w0, w, V} \sum{(x, y) \in \mathcal{D}} \mathcal{L}(y, \hat{y}(x)) + \lambda (\lVert w \rVert^2 + \lVert V \rVert^2)$

The loss $\mathcal{L}$ could be squared error for regression or log loss for classification. The $\ell_2$ regularization terms prevent overfitting.

This optimization problem is typically solved with stochastic gradient descent (SGD). In each iteration, we sample a batch of data points, calculate the gradients of the loss with respect to the parameters, and update them in the negative gradient direction. Pseudocode is shown below:

def train_fm(data, k, nu, lambda_, num_iters):
    w0 = 0
    w = np.zeros(n)
    V = np.random.normal(0, 0.01, (n, k)) 

    for iter in range(num_iters):
        x, y = sample_batch(data) 

        y_pred = predict_fm(w0, w, V, x)
        loss = calc_loss(y, y_pred)

        g_w0 = grad_w0(y, y_pred)
        g_w = grad_w(y, y_pred, x) + 2*lambda_*w
        g_V = grad_V(y, y_pred, x, V) + 2*lambda_*V

        w0 -= nu * g_w0
        w -= nu * g_w
        V -= nu * g_V

    return w0, w, V  

There are several libraries that efficiently implement FMs in various languages, including libfm, xLearn, and lightfm for Python. Many even support the more general class of field-aware factorization machines (FFMs), described next.

Field-Aware Factorization Machines

In the standard FM, each feature has a single $k$-dimensional embedding that is used in all pairwise interactions. The FFM introduces more flexibility by learning a separate embedding for each feature depending on what field the other feature in the interaction belongs to.

For example, consider a CTR prediction problem with features like website, ad position, device type, etc. An FFM would learn a "website" embedding of the Yahoo.com feature to use when modeling interactions with ad position features, and a separate one for interactions with device features.

Mathematically, the FFM model is:

$\hat{y}(x) = w0 + \sum{i=1}^n w_i xi + \sum{i=1}^n \sum{j=i+1}^n \langle v{i, fj}, v{j, f_i} \rangle x_i x_j$

where $f_i$ denotes the field that feature $i$ belongs to. If feature $i$ appears in $f$ distinct fields, it will have $f \cdot k$ parameters in the FFM as opposed to just $k$ in the FM.

This expressiveness has led FFMs to achieve state-of-the-art results on many CTR benchmarks. A popular library for training FFMs is xLearn, which provides a concise Python API:

import xlearn as xl

ffm_model = xl.create_ffm()
ffm_model.setTrain("./small_train.txt")  
ffm_model.setValidate("./small_test.txt")  

param = {
    ‘task‘:‘binary‘, 
    ‘lr‘:0.2, 
    ‘lambda‘:0.002,
    ‘metric‘:‘auc‘
}

ffm_model.fit(param, "./model.out") 

Conclusion and Future Directions

Factorization machines have become an essential tool for machine learning on sparse data, powering recommender systems, dynamic pricing engines, and more. By embedding features into a latent space and modeling their pairwise interactions, FMs achieve strong performance while keeping computational and memory costs under control.

There are several exciting areas of ongoing research that build on the ideas behind FMs:

• Higher-order FMs (HOFMs) that model not just pairwise, but also 3-way, 4-way, etc. interactions between features. Efficient training algorithms that exploit tensor algebra have been proposed.

• Neural factorization machines that combine the embedding approach of FMs with the non-linear modeling capabilities of deep learning. These often stack a neural network on top of the FM interaction terms.

• Techniques for making FMs more interpretable by adding sparsity or non-negativity constraints on the embeddings, or through post-hoc analysis of the learned latent factors.

To dive deeper into the world of FMs, check out Steffen Rendle‘s original paper, this excellent tutorial on Coursera, or browse the documentation of libraries like xLearn. Happy factorizing!