Laplace Approximation for Bayesian Tensor Network Kernel Machines

1. Problem Statement

Machine learning systems require robust uncertainty estimation (UE) to handle ambiguous inputs, out-of-distribution data, and noisy features. While Gaussian Processes (GPs) provide principled probabilistic frameworks for UE, their computational complexity scales cubically ($O(N^3)$) with the number of data points, limiting them to small-to-medium datasets.

Conversely, Tensor Network Kernel Machines (TNKMs) offer a scalable alternative by parameterizing model weights as tensor networks (specifically Canonical Polyadic Decomposition, or CPD). This reduces the number of parameters from exponential to linear relative to the input dimension. However, standard TNKMs typically rely on Maximum A Posteriori (MAP) estimation. The non-linear, multilinear structure of tensor networks breaks the Gaussianity assumption required for closed-form Bayesian inference, making it difficult to derive principled uncertainty estimates without resorting to expensive sampling methods (like MCMC) or complex variational inference.

The Core Question: How can Tensor Network Kernel Machines provide principled, scalable uncertainty estimates comparable to GPs or Bayesian Neural Networks (BNNs)?

2. Methodology: LA-TNKM

The authors propose LA-TNKM (Laplace Approximation Tensor Network Kernel Machine), a framework that unifies variational inference with a (linearized) Laplace approximation to perform tractable Bayesian inference on tensor networks.

Key Components:

1. Model Formulation:

   • The model is defined as $f(x) = \phi(x)^T w$, where features $\phi(x)$ have a tensor-product structure and weights $w$ are represented as a CPD tensor network.
   • A Gaussian likelihood and a Gaussian prior are assumed over the parameters.

2. Laplace Approximation:

   • Instead of integrating over the high-dimensional posterior (intractable), the authors approximate the posterior distribution $p(v|D)$ with a Gaussian distribution centered at the MAP estimate $v^*$:
     $$q(v) = \mathcal{N}(v | v^*, H_{v^*}^{-1})$$
   • Here, $H_{v^*}$ is the Hessian matrix of the loss function evaluated at the optimum.

3. Hessian Approximations:
   Computing the full Hessian for tensor networks is memory-intensive. The paper systematically evaluates five approximation strategies to balance memory and accuracy:

   • Full Hessian: Exact but computationally expensive ($O((DIR)^2)$).
   • Generalized Gauss-Newton (GGN): Positive semi-definite approximation, often used in BNNs.
   • Block-Diagonal: Assumes independence between CPD cores (analogous to layers in BNNs).
   • Diagonal: Assumes all weights are independent (mean-field).
   • Last Core: Applies Bayesian inference only to the last CPD core, treating earlier cores as fixed. This is shown to be highly efficient ($O(I^2R^2)$) and effective.

4. Linearized Laplace Approximation (LLA):

   • The authors address the mismatch between the non-linear model and the Gaussian posterior by linearizing the model function $f(x, v)$ around $v^*$.
   • This allows the predictive distribution to be computed analytically using the linearized model, significantly improving predictive performance (Negative Log-Likelihood) compared to the non-linearized Laplace approximation.

5. Eigenvalue Truncation:

   • To handle ill-conditioned Hessians (where eigenvalues are near zero or negative), the authors apply a truncated eigenvalue decomposition. A hyperparameter $\hat{t}$ is tuned via cross-validation to filter out small eigenvalues, preventing the predictive variance from exploding.

3. Key Contributions

• Novel Framework (LA-TNKM): Introduction of the first Bayesian Tensor Network Kernel Machine that enables uncertainty-aware predictions while maintaining computational costs comparable to standard MAP-based TNKMs.
• Systematic Hessian Analysis: A comprehensive comparison of Hessian approximation techniques (Full, GGN, Block, Diagonal, Last Core) within the tensor network context, providing guidelines on memory usage and predictive performance.
• Linearization Strategy: Demonstration that linearizing the prediction function (LLA) is crucial for achieving high-quality uncertainty estimates in tensor networks, aligning with recent findings in BNNs.
• Empirical Validation: Extensive experiments showing that LA-TNKM matches or surpasses GPs and BNNs on diverse regression tasks.

4. Experimental Results

The method was evaluated on nine UCI regression datasets and a synthetic cubic regression task ($y = x^3 + \epsilon$).

• Performance vs. Baselines:

  • LA-TNKM (specifically using the Last Core approximation) consistently achieved the lowest Negative Log-Likelihood (NLL) on 4 out of 9 datasets and performed competitively on the rest.
  • It outperformed standard BNNs (Bayes-by-Backprop, Variational Dropout) and other tensor-based methods (MF-BTN, SP-BTN).
  • In terms of Uncertainty Calibration (measured by ECP-95 and Relative Calibration Error), LA-TNKM produced intervals that were sharper (lower WCPI-95) and better calibrated than BNNs, closely matching the behavior of Gaussian Processes.

• Synthetic Data Findings:

  • On the synthetic $x^3$ task, LA-TNKM successfully captured the uncertainty in regions with sparse data, whereas SP-BTN tended to overestimate variance (overly wide intervals) and LA-BNN showed signs of model mismatch during extrapolation.

• Ablation Studies:

  • Linearization: The Linearized Laplace Approximation (LLA) consistently outperformed the standard Laplace approximation in terms of NLL.
  • Hessian Choice: While the "Last Core" approximation was the most memory-efficient and generally performed best, the optimal choice was found to be data-dependent. However, "Last Core" and "Block-Diagonal" offered the best trade-off between efficiency and accuracy.

5. Significance and Conclusion

This work bridges the gap between the scalability of Tensor Networks and the robustness of Bayesian Inference.

• Scalability: By utilizing tensor decompositions and efficient Hessian approximations, LA-TNKM avoids the $O(N^3)$ bottleneck of GPs, making principled uncertainty estimation feasible for larger datasets and higher-dimensional inputs.
• Practicality: The method provides a "plug-and-play" probabilistic layer for existing tensor network architectures without requiring complex sampling or variational families.
• Future Impact: The results suggest that for many regression tasks, a localized Bayesian treatment (e.g., Last Core) combined with linearization is sufficient to achieve state-of-the-art uncertainty quantification, offering a computationally efficient alternative to both GPs and full BNNs.

In summary, LA-TNKM demonstrates that tensor network kernel machines can be effectively Bayesian, offering a scalable, robust, and principled approach to uncertainty estimation in modern machine learning.