[Paper Review] RaFM: Rank-Aware Factorization Machines

Xiaohong Chen, Yin Zheng+3 more|arXiv (Cornell University)|May 18, 2019

Machine Learning and ELM4 citations

TL;DR

RaFM proposes a rank-aware factorization machine that assigns feature-specific embedding ranks to model pairwise interactions, improving performance on sparse, frequency-varying datasets. It achieves better accuracy than standard FMs with reduced computational and storage costs by dynamically managing active factors, enabling efficient training and inference comparable to or better than fixed-rank FMs.

ABSTRACT

Factorization machines (FM) are a popular model class to learn pairwise interactions by a low-rank approximation. Different from existing FM-based approaches which use a fixed rank for all features, this paper proposes a Rank-Aware FM (RaFM) model which adopts pairwise interactions from embeddings with different ranks. The proposed model achieves a better performance on real-world datasets where different features have significantly varying frequencies of occurrences. Moreover, we prove that the RaFM model can be stored, evaluated, and trained as efficiently as one single FM, and under some reasonable conditions it can be even significantly more efficient than FM. RaFM improves the performance of FMs in both regression tasks and classification tasks while incurring less computational burden, therefore also has attractive potential in industrial applications.

Motivation & Objective

• To address performance degradation in standard factorization machines (FMs) when features have highly variable occurrence frequencies.
• To reduce overfitting for rare features and underfitting for frequent features by allowing feature-specific embedding ranks.
• To maintain computational and storage efficiency despite using multiple ranks per feature.
• To develop a training algorithm that efficiently handles inactive factors and ensures convergence.
• To demonstrate superior performance and efficiency in both public and industrial-scale datasets.

Proposed method

• RaFM assigns multiple embedding vectors with different ranks to each feature, where the maximum rank is determined by feature frequency.
• It models pairwise interactions using a sum over rank-specific dot products: $\left<\mathcal{V}_i, \mathcal{V}_j\right> = \sum_{k=1}^{k_i} \bm{v}_i^{(k)} \cdot \bm{v}_j^{(k)}$.
• The model leverages the fact that many high-rank components are inactive (zero when feature is zero), enabling efficient storage and evaluation.
• A unified training algorithm is proposed that avoids storing and updating inactive factors, reducing computational overhead.
• Theoretical analysis proves bounded training error and convergence rate under reasonable assumptions.
• The method is implemented using FTRL for sparsity and distributed learning for scalability in industrial settings.

Experimental results

Research questions

• RQ1 Can a factorization machine model achieve better performance on datasets with highly skewed feature frequencies by assigning different embedding ranks per feature?
• RQ2 Can such a model maintain or improve efficiency compared to standard FMs with fixed ranks?
• RQ3 Does dynamically managing active and inactive factors in multi-rank embeddings lead to reduced model size and faster training?
• RQ4 Can the proposed training algorithm converge efficiently while avoiding overfitting from excessive parameters?
• RQ5 How does RaFM compare to existing methods like DiFacto and MRMA in terms of performance and computational cost on real-world and industrial datasets?

Key findings

• RaFM improves logloss by 1% to 2% on ML 10M, ML 20M, and AMovie, and by 15% on Frappe, with statistically significant gains.
• On the Criteo dataset, RaFM reduces logloss by 0.002, exceeding the 0.001 threshold considered practically significant.
• RaFM reduces model size to 20%–66% of standard FM and training time to 24%–95% of FM on public datasets.
• On an industrial CTR dataset with 1.7B records and 120M features, RaFM increases AUC by ~1% over LR with only 1.55× the model size, compared to 7× for FM.
• RaFM achieves a 22% model size reduction relative to FM while matching or exceeding its AUC performance, demonstrating a superior efficiency-accuracy trade-off.
• The training algorithm ensures bounded error and convergence, with inactive factors excluded from computation and parameter updates, enabling efficient learning.