# ALX: Large Scale Matrix Factorization on TPUs

arXiv: https://arxiv.org/abs/2112.02194 · HTML: https://mlsanity.com/paper/2112.02194/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 21 references that are in the map.
- First posted on arXiv 3 Dec 2021.

## Builds on (most cited first, 8 of 8)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Asymmetric LSH (ALSH) for Sublinear Time Maximum Inner Product Search (MIPS)](https://mlsanity.com/paper/1405.5869/) (2014)
- [On the Difficulty of Evaluating Baselines: A Study on Recommender Systems](https://mlsanity.com/paper/1905.01395/) (2019)
- [Revisiting the Performance of iALS on Item Recommendation Benchmarks](https://mlsanity.com/paper/2110.14037/) (2021)
- [Efficient Training on Very Large Corpora via Gramian Estimation](https://mlsanity.com/paper/1807.07187/) (2018)
- [Faster and Cheaper: Parallelizing Large-Scale Matrix Factorization on GPUs](https://mlsanity.com/paper/1603.03820/) (2016)
- [Factorbird - a Parameter Server Approach to Distributed Matrix Factorization](https://mlsanity.com/paper/1411.0602/) (2014)
- [Matrix Factorization on GPUs with Memory Optimization and Approximate Computing](https://mlsanity.com/paper/1808.03843/) (2018)

## Built on by (most cited first, 2 of 2)

- [Efficient Mixed Dimension Embeddings for Matrix Factorization](https://mlsanity.com/paper/2205.11248/) (2022)
- [A latent linear model for nonlinear coupled oscillators on graphs](https://mlsanity.com/paper/2311.14910/) (2023)

## Abstract (arXiv)

We present ALX, an open-source library for distributed matrix factorization using Alternating Least Squares, written in JAX. Our design allows for efficient use of the TPU architecture and scales well to matrix factorization problems of O(B) rows/columns by scaling the number of available TPU cores. In order to spur future research on large scale matrix factorization methods and to illustrate the scalability properties of our own implementation, we also built a real world web link prediction dataset called WebGraph. This dataset can be easily modeled as a matrix factorization problem. We created several variants of this dataset based on locality and sparsity properties of sub-graphs. The largest variant of WebGraph has around 365M nodes and training a single epoch finishes in about 20 minutes with 256 TPU cores. We include speed and performance numbers of ALX on all variants of WebGraph. Both the framework code and the dataset is open-sourced.