Joint Bayesian Parameter and Model Order Estimation for Low-Rank Probability Mass Tensors

Technical Summary: Joint Bayesian Parameter and Model Order Estimation for Low-Rank Probability Mass Tensors

Problem Statement
In statistical signal processing and machine learning, obtaining a reliable estimate of the joint probability mass function (PMF) of a set of discrete random variables is a fundamental challenge. While the joint PMF can be represented as an $N$-dimensional tensor, direct estimation via histogram methods suffers from the "curse of dimensionality," requiring an exponential number of observations as the number of variables increases.

Recent approaches model the joint PMF as a low-rank Canonical Polyadic Decomposition (CPD), linking it to latent variable models such as the Naïve Bayes model. However, existing CPD-based estimation algorithms (e.g., maximum likelihood or coupled tensor factorization) require the tensor rank (model order) to be specified beforehand. In practice, the true rank is unknown. Current solutions involve selecting a rank from a candidate set using cross-validation (CV) or information criteria (AIC, BIC, DNML). These procedures are computationally expensive, as they require multiple training runs, and risk model misspecification if the true rank is not included in the candidate set.

Methodology
The paper proposes a novel Bayesian framework, termed VB-PMF (Variational Bayesian PMF), which simultaneously estimates the low-rank components of a joint PMF tensor and infers its rank directly from the observed data, eliminating the need for cross-validation.

1. Bayesian Model Specification:

   • The joint PMF is modeled as a CPD where the loading vector $\lambda$ represents the prior probabilities of latent states, and factor matrices $A_n$ represent conditional PMFs.
   • To enforce the probability simplex constraints (non-negativity and sum-to-one), the authors assign Dirichlet priors to the loading vector $\lambda$ and the columns of the factor matrices.
   • Crucially, a sparsity-promoting Dirichlet prior is imposed on the loading vector by setting its hyperparameter $\alpha_\lambda$ to a small value (e.g., $10^{-6}$). This encourages the posterior distribution to concentrate mass on a few components, effectively identifying and pruning irrelevant rank-one terms.

2. Variational Inference (VI) Solution:

   • Since exact Bayesian inference is computationally intractable due to high-dimensional integration, the authors employ Variational Inference.
   • A mean-field approximation is used, assuming the posterior distribution factorizes into independent distributions for the latent variables ($Z$), the loading vector ($\lambda$), and the factor matrices ($A_n$).
   • The algorithm iteratively updates the variational parameters to maximize the Evidence Lower Bound (ELBO). This yields closed-form update equations for all posterior distributions, ensuring deterministic convergence without the need for sampling (unlike MCMC methods).

3. Automatic Rank Detection:

   • The rank is inferred automatically after convergence. Components with estimated weights ($\hat{\lambda}_r$) below a derived pruning threshold ($\alpha_\lambda / T$, where $T$ is the dataset size) are considered irrelevant and removed.
   • The algorithm is initialized with a rank $R$ that is larger than or equal to the true rank (e.g., satisfying Kruskal uniqueness conditions). The sparsity-promoting prior ensures that extraneous components converge to negligible values and are pruned, leaving only the significant latent components.

Key Contributions

• Unified Bayesian Framework: The paper introduces a method to estimate both the parameters and the rank of a joint PMF tensor in a single training run, avoiding the computational cost of cross-validation or grid search.
• Sparsity-Driven Rank Inference: By leveraging the properties of the Dirichlet distribution, the method automatically identifies the correct model order through posterior pruning, removing the need for heuristic thresholding or manual model selection.
• Efficient Deterministic Algorithm: The derivation of a Variational Inference solution provides closed-form updates and deterministic convergence, offering a computationally efficient alternative to sampling-based Bayesian methods (MCMC) and iterative likelihood-based methods requiring multiple rank trials.
• Robustness to Missing Data: The framework naturally handles missing observations (outage probabilities) by marginalizing over unobserved variables within the likelihood function.

Experimental Results
The authors validate the method using synthetic data and real-world datasets (UCI classification tasks and MovieLens 10M recommendation data).

• Synthetic Data:

  • Rank Consistency: As the number of observations increases, VB-PMF consistently converges to the true tensor rank, demonstrating statistical consistency.
  • Accuracy: VB-PMF achieves Mean Kullback-Leibler Divergence (KLD) and Mean Squared Relative Error (MSRE) comparable to or better than state-of-the-art methods (SQ-AIC, SQ-BIC, SQ-DNML, CTF, and Tensor Networks), particularly in scenarios with missing data.
  • Efficiency: The runtime of VB-PMF is significantly lower than methods requiring cross-validation (e.g., CTF-ValErr, SQ-AIC) because it requires only a single training run regardless of the initial rank setting.
  • Hyperparameter Sensitivity: The method is robust to the choice of the initial rank (provided it is sufficiently large) and the hyperparameter $\alpha_\lambda$, provided $\alpha_\lambda$ is set small enough to promote sparsity.

• Real-World Applications:

  • Classification: On five UCI datasets, VB-PMF achieved classification accuracy and F1 scores comparable to Random Forests (a strong discriminative benchmark) and outperformed other PMF-based approaches. The model provided interpretable insights into latent voting blocs in the Congressional Voting dataset, which discriminative models cannot reveal.
  • Recommendation: On the MovieLens 10M dataset (predicting movie ratings), VB-PMF achieved RMSE and MAE comparable to Coupled Tensor Factorization (CTF3D) and significantly better than Biased Matrix Factorization (BMF), with a runtime roughly 10 times faster than CTF3D.

Significance and Claims
The paper claims that the proposed VB-PMF framework offers a significant advancement over existing PMF estimation techniques by unifying parameter estimation and model order selection. Its primary significance lies in:

1. Eliminating Cross-Validation: It removes the computational bottleneck of selecting a rank via CV or information criteria, making it feasible for large-scale datasets.
2. Interpretability: By estimating the full joint PMF via low-rank components, the model offers an interpretable representation of data structure (e.g., latent classes in classification) that purely discriminative models lack.
3. Computational Efficiency: The variational approach provides a deterministic, closed-form solution that is faster than both sampling-based Bayesian inference and iterative rank-selection strategies.

The authors conclude that VB-PMF provides a robust, efficient, and accurate solution for joint PMF estimation, particularly in scenarios involving high-dimensional discrete data and missing observations, without the need for manual model tuning.