BaTFLED: Bayesian Tensor Factorization Linked to External Data

TL;DR

BaTFLED is a Bayesian tensor factorization method that predicts multi-dimensional responses from external features, effectively capturing interactions and enabling feature selection, outperforming traditional models on simulated and real biological data.

Contribution

Introduces BaTFLED, a novel Bayesian tensor factorization algorithm that links external data to multi-dimensional responses using a Tucker decomposition with sparsity priors.

Findings

Outperforms elastic net and neural networks on cold start tasks

Successfully predicts dose-response in breast cancer cell lines

Captures interactions between latent factors effectively

Abstract

The vast majority of current machine learning algorithms are designed to predict single responses or a vector of responses, yet many types of response are more naturally organized as matrices or higher-order tensor objects where characteristics are shared across modes. We present a new machine learning algorithm BaTFLED (Bayesian Tensor Factorization Linked to External Data) that predicts values in a three-dimensional response tensor using input features for each of the dimensions. BaTFLED uses a probabilistic Bayesian framework to learn projection matrices mapping input features for each mode into latent representations that multiply to form the response tensor. By utilizing a Tucker decomposition, the model can capture weights for interactions between latent factors for each mode in a small core tensor. Priors that encourage sparsity in the projection matrices and core tensor allow for feature selection and model regularization. This method is shown to far outperform elastic net and neural net models on 'cold start' tasks from data simulated in a three-mode structure. Additionally, we apply the model to predict dose-response curves in a panel of breast cancer cell lines treated with drug compounds that was used as a Dialogue for Reverse Engineering Assessments and Methods (DREAM) challenge.