Technical Summary: Low-rank Distributional Matrix Completion

Problem Definition

This paper addresses the problem of distributional matrix completion, a generalization of the classical low-rank matrix completion task. In this setting, the target matrix $\rho^* \in \mathcal{P}(\mathcal{X})^{m_1 \times m_2}$ consists of entries that are probability distributions over a domain $\mathcal{X}$, rather than scalar values.

The observation model is characterized by two layers of incompleteness:

1. Partial Observation: Only a subset of matrix entries is observed, indicated by a binary mask $T(i,j)$.
2. Indirect Observation: Even for observed entries, the underlying distribution $\rho^*(i,j)$ is not directly accessible. Instead, one observes a finite set of i.i.d. samples $\{X^{(ij)}_1, \dots, X^{(ij)}_{N_{ij}}\}$ drawn from $\rho^*(i,j)$, where the sample size $N_{ij}$ may vary across entries.

The goal is to recover the full distribution-valued matrix $\rho^*$ (or its representation) from these partial and indirect observations. This problem arises in applications such as modeling traffic flow variability between locations or analyzing the distribution of earnings forecasts across different fiscal quarters and institutions.

Methodology

Kernel Mean Embeddings (KME)

To handle the infinite-dimensional nature of probability distributions and avoid manifold constraints associated with density functions, the authors represent each distribution $\rho^*(i,j)$ via its Kernel Mean Embedding (KME) in a Reproducing Kernel Hilbert Space (RKHS) $\mathcal{H}$:
$$\mu^*(i, j, \cdot) = \mathbb{E}_{X \sim \rho^*(i,j)} [K(X, \cdot)] \in \mathcal{H}$$
Recovering $\rho^*$ is equivalent to recovering the KME matrix $\mu^* \in \mathcal{H}^{m_1 \times m_2}$. The distance between distributions is measured using the Maximum Mean Discrepancy (MMD).

Low-Rank Modeling via Functional Unfolding

The authors propose a Tucker rank structure for the distributional matrix. They view $\mu^*$ as a three-way tensor in the space $\mathbb{R}^{m_1} \otimes \mathbb{R}^{m_2} \otimes \mathcal{H}$. Since $\mathcal{H}$ is often infinite-dimensional, standard tensor rank definitions do not directly apply.

To address this, the paper introduces functional unfolding operators ($U_1, U_2, U_3$) that map the tensor into products of two spaces (e.g., $U_1: \mathbb{R}^{m_1} \otimes \mathbb{R}^{m_2} \otimes \mathcal{H} \to \mathbb{R}^{m_1} \otimes (\mathbb{R}^{m_2} \otimes \mathcal{H})$). These operators induce linear operators whose ranks define the multilinear rank $(r_1, r_2, r_3)$.

• Structural Implication: A low multilinear rank implies that the functional components of the KMEs lie in a low-dimensional subspace, and the associated expectation functional matrices (e.g., mean matrices) are low-rank.

Estimation Framework

The authors propose a regularized M-estimator $\hat{\mu}$ defined by minimizing a loss function plus a nuclear-norm penalty:
$$\hat{\mu} = \arg\min_{\mu} \left[ \frac{1}{m_1 m_2} \sum_{(i,j): T(i,j)=1} \left\| \mu(i, j, \cdot) - \hat{\mu}_{emp}(i, j, \cdot) \right\|_{\mathcal{H}}^2 + \lambda \Psi(\mu) \right]$$

• Loss Term: The squared MMD between the estimated KME and the empirical KME derived from observed samples.
• Regularization: $\Psi(\mu) = \sum_{j=1}^3 \beta_j \|U_j \mu\|_*$, where $\|\cdot\|_*$ denotes the nuclear norm of the induced linear operator. This promotes low-rankness across all three modes.

Computational Solution

Despite the infinite-dimensional nature of $\mathcal{H}$, the Representer Theorem ensures that the solution $\hat{\mu}$ lies in the finite-dimensional span of the observed kernel functions.

1. The problem is re-parameterized into a finite-dimensional tensor optimization problem over a coefficient tensor $B \in \mathbb{R}^{m_1 \times m_2 \times \tilde{N}}$.
2. To handle large sample sizes, the authors suggest approximating the Gram matrix using kernel approximation techniques (e.g., Nyström method or Random Fourier Features).
3. The optimization is solved using an accelerated Alternating Direction Method of Multipliers (ADMM).

Key Contributions

1. Novel Problem Formulation: The paper formalizes matrix completion for distribution-valued entries, addressing the challenge of recovering distributions from partial and noisy sample data.
2. Theoretical Framework: It introduces a Tucker rank definition for distribution-valued tensors using functional unfolding operators, bridging the gap between infinite-dimensional functional data and finite-dimensional tensor analysis.
3. Algorithmic Innovation: The proposed estimator utilizes nuclear-norm regularization on induced linear operators. Crucially, the authors prove that the solution admits a finite-dimensional representation, making the problem computationally tractable via convex optimization.
4. Statistical Guarantees: The paper establishes non-asymptotic error bounds for the aggregate squared MMD.
   • The convergence rate matches that of classical nuclear-norm regularized matrix completion for scalar matrices.
   • Unlike existing nearest-neighbor approaches (e.g., Feitelberg et al., 2024), the proposed method's consistency does not require the number of samples per entry ($N_{ij}$) to diverge. Instead, consistency is achieved as the matrix dimensions ($m_1, m_2$) grow, leveraging the global low-rank structure for information sharing.

Experimental Results

Synthetic Data

• Setup: Simulated matrices with Gaussian mixture distributions. Comparisons were made against a nearest-neighbor (NN) method (Feitelberg et al., 2024) and classical matrix completion (SOFTIMPUTE) applied to sample means/variances.
• Findings:
  • The proposed method (LRKME) consistently outperformed the NN method in both observed and unobserved entries, particularly in estimating variance and higher-order statistics.
  • LRKME also outperformed SOFTIMPUTE, demonstrating the advantage of utilizing full distributional information rather than just scalar summaries.
  • In 2D settings, the NN method was excluded due to computational intractability of multivariate Wasserstein distances, while LRKME remained effective.

Real-World Application

• Dataset: NYC taxi trip data, modeling the distribution of daily trip counts between 265 source and destination locations.
• Challenge: High missing rate (>39%) and highly variable sample sizes per entry (1 to 3672).
• Findings: LRKME achieved a significantly lower empirical relative error (0.0015) compared to the NN method (0.0212) on the test set. Visual inspection of specific cases showed that LRKME better captured heavy-tailed distributions and concentration patterns compared to the NN estimator.

Significance and Claims

The paper claims to provide the first method for distributional matrix completion that:

1. Accommodates general observation patterns with varying sample sizes.
2. Extends naturally to multidimensional distributions without the computational bottlenecks of Wasserstein-based methods.
3. Offers global consistency guarantees that do not rely on the asymptotic growth of samples per entry, relying instead on the low-rank structure of the matrix.
4. Provides a computationally efficient solution via finite-dimensional representation and convex optimization.

The authors position their work as a significant step forward from existing nearest-neighbor approaches, offering a more robust, theoretically grounded, and scalable framework for imputing missing distributional data.