ALX: Large Scale Matrix Factorization on TPUs
ALX: Large Scale Matrix Factorization on TPUs
复制标题
ALX：TPU 上的大规模矩阵分解
DOI：
--
复制
发表时间：
2021
期刊：
影响因子：
--
通讯作者：
Li Zhang
中科院分区：
文献类型：
--
作者：
Harsh Mehta;Steffen Rendle;Walid Krichene;Li Zhang
我们提出了ALX，一个开源的库分布式矩阵分解使用交替最小二乘法，写在JAX。我们的设计允许有效地使用TPU架构，并通过缩放可用TPU核的数量来很好地扩展到O（B）行/列的矩阵分解问题。为了促进未来对大规模矩阵分解方法的研究，并说明我们自己的实现的可扩展性，我们还建立了一个真实的世界网络链接预测数据集，称为WebGraph。这个数据集可以很容易地建模为矩阵分解问题。我们根据子图的局部性和稀疏性属性创建了这个数据集的几个变体。WebGraph的最大变体拥有大约3.65亿个节点，使用256个TPU核心，训练一个epoch大约在20分钟内完成。我们在WebGraph的所有变体上都包含了ALX的速度和性能数据。框架代码和数据集都是开源的。
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.