Rank regularization and Bayesian inference for tensor completion and extrapolation

TL;DR

This paper introduces a Bayesian tensor completion method using a new rank regularizer for PARAFAC factors, capable of accurately recovering missing data in multi-way arrays with different data distributions.

Contribution

It proposes a novel rank regularizer within a Bayesian framework for tensor completion, with algorithms tailored for Gaussian and Poisson data, improving accuracy and flexibility.

Findings

Successfully recovers ground-truth tensor rank on synthetic data.

Completes brain imaging data with 50% missing entries, achieving -10dB error.

Completes yeast gene expression data with 15% missing entries, achieving -15dB error.

Abstract

A novel regularizer of the PARAFAC decomposition factors capturing the tensor's rank is proposed in this paper, as the key enabler for completion of three-way data arrays with missing entries. Set in a Bayesian framework, the tensor completion method incorporates prior information to enhance its smoothing and prediction capabilities. This probabilistic approach can naturally accommodate general models for the data distribution, lending itself to various fitting criteria that yield optimum estimates in the maximum-a-posteriori sense. In particular, two algorithms are devised for Gaussian- and Poisson-distributed data, that minimize the rank-regularized least-squares error and Kullback-Leibler divergence, respectively. The proposed technique is able to recover the "ground-truth'' tensor rank when tested on synthetic data, and to complete brain imaging and yeast gene expression datasets with 50% and 15% of missing entries respectively, resulting in recovery errors at -10dB and -15dB.