🔢 mathematics

Variational inference and density estimation with non-negative tensor of hierarchical tucker format

This paper proposes a two-stage, linear-complexity methodology that compresses high-dimensional discrete probability tensors into a non-negative hierarchical Tucker format using interpolation followed by a tailored second-order optimization, enabling efficient variational inference and density estimation in high-dimensional settings.

Original authors: Xun Tang, Haoxuan Chen, Lexing Ying

Published 2026-06-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xun Tang, Haoxuan Chen, Lexing Ying

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, multi-dimensional library of information. In the world of probability, this library is a "tensor"—a giant grid of numbers representing the likelihood of every possible combination of events happening. If you have just 10 variables with 100 possibilities each, your library has $100^{10}$ pages. That's too big to store, let alone read.

This paper proposes a clever way to shrink that giant library down into a tiny, manageable backpack without losing the essential story. They call this method Variational Inference and Density Estimation with Non-Negative Hierarchical Tucker Format.

Here is the simple breakdown of how they do it, using everyday analogies.

The Problem: The "Sign" Trouble

In math, when you try to compress these giant libraries, you often use a technique that breaks the data into smaller pieces (factors). However, standard math allows these pieces to have "negative" numbers.

Think of probability like a pile of sand. You can't have "-5 grains of sand." If your compression method creates negative numbers, you end up with a "signed" pile of sand—some parts are sand, and some parts are "anti-sand." This breaks the rules of probability. You can't calculate the total weight of the pile, and you can't use it to make predictions.

The authors' goal is to compress the data while ensuring every single number remains positive, just like real sand.

The Solution: A Two-Stage Construction Project

The authors built a two-stage machine to solve this. Think of it like renovating a house.

Stage 1: The Rough Draft (Interpolation)

First, they take the giant, uncompressed library and create a "rough draft" version of it.

How they do it: They use a technique similar to taking a few key photos of a landscape to guess what the whole view looks like. They pick specific "pivot" points (key pages in the library) and use a method called Hierarchical Tucker (HT) to stitch them together.
The Catch: This rough draft is fast to make, but it's "signed." It might have those problematic negative numbers. It's a good sketch, but it's not a finished, usable house yet.

Stage 2: The Renovation (Fitting)

Now, they take that rough draft and force it to become a "Non-Negative" version. This is the paper's main innovation.

The Goal: They want to reshape the rough draft into a new structure (called NHT) where every single number is positive, but it still looks exactly like the original rough draft.
The Trick: They use a "second-order" method. Imagine you are trying to fit a puzzle piece into a hole. A simple method might just push the piece in blindly. This paper uses a "smart push" (a Newton step) that calculates exactly how much to push and in which direction to get the perfect fit without breaking the "no negative numbers" rule.
The Secret Sauce (Warm Start): Usually, when you try to fix a puzzle, you might get stuck in a local trap (a piece that fits okay but isn't the best fit). The authors invented a "Warm Initialization" strategy. Before they start the hard work, they do a quick, smart pre-game to arrange the pieces in a good position. This prevents them from getting stuck and helps them find the perfect solution much faster.

Why Use a "Tree" Structure?

The paper uses a Hierarchical Tucker format, which is based on a binary tree (like a family tree or a decision tree).

The Old Way (Train): Previous methods used a "Train" structure (Tensor Train), where variables are linked in a single long line. This works great for data where things only affect their immediate neighbors (like a line of people passing a message).
The New Way (Tree): The authors' "Tree" structure is better for data where things affect each other in complex, 2D patterns (like a grid of people in a room where everyone talks to their neighbors in all directions). The tree structure naturally captures these complex "2D lattice" relationships, which the "Train" structure struggles with.

The Results

The authors tested this on two types of problems:

Variational Inference: Where they have a formula and can ask questions about it directly.
Density Estimation: Where they only have a bag of random samples and have to guess the shape of the distribution.

In both cases, their method:

Compressed the data efficiently (keeping the file size small).
Kept all numbers positive (ensuring it was a valid probability model).
Converged (finished the job) much faster and more accurately than older methods, especially for complex 2D grid problems.

Summary

Think of this paper as inventing a new, smarter way to fold a giant, complex map into your pocket.

They first make a quick, rough sketch of the map (Stage 1).
Then, they use a special, smart folding technique (Stage 2) that ensures the map is folded perfectly without any "negative" creases, using a tree-like folding pattern that handles complex shapes better than old straight-line folding methods.
They also figured out how to start the folding process in the right spot so you don't waste time trying to fix a bad fold later.

The result is a highly efficient, mathematically sound way to store and understand massive amounts of probability data.

Technical Summary: Variational Inference and Density Estimation with Non-Negative Tensor of Hierarchical Tucker Format

Problem Statement
The paper addresses the challenge of compressing high-dimensional discrete probability distributions, represented as tensors $P: [n]^d \to \mathbb{R}_{\geq 0}$ , into tractable low-rank formats. Storing such a tensor directly requires $O(n^d)$ complexity, which becomes infeasible even for moderate dimensions $d$ (e.g., tens of variables). The goal is to approximate $P$ within a parametric family that allows for efficient storage, sampling, and likelihood evaluation. The work considers two access models:

Variational Inference (VI): Access to $P$ via an oracle that provides unnormalized values $P(i_1, \dots, i_d)$ up to a constant.
Density Estimation (DE): Access to $P$ via a collection of samples drawn from the distribution.

Existing tensor network models, such as the Tensor-Train (TT) format, are effective for distributions satisfying an "area law" but struggle with 2D lattice models or complex interaction structures where the area law does not hold. Furthermore, standard tensor networks do not inherently guarantee non-negativity, which is critical for probability distributions to avoid "sign problems" in downstream tasks like Kullback-Leibler divergence calculation or sampling.

Methodology
The authors propose a two-stage end-to-end algorithm to compress $P$ into a Non-Negative Hierarchical Tucker (NHT) format.

Stage 1: Hierarchical Tucker Interpolation (Reference Construction)
- Goal: Construct a signed Hierarchical Tucker (HT) approximation $\tilde{P} \approx P$ .
- VI Case: Uses existing interpolation methods based on interpolative decomposition (CUR-like decomposition) to query entries of $P$ and reconstruct $\tilde{P}$ .
- DE Case: Uses existing density estimation algorithms to construct $\tilde{P}$ from samples.
- Output: A reference tensor $\tilde{P}$ with signed components. This stage does not enforce non-negativity.
Stage 2: Non-Negative Hierarchical Tucker Fitting
- Goal: Fit a non-negative ansatz $P_G$ (parameterized by non-negative tensor factors $G = \{G_q\}$ ) to the signed reference $\tilde{P}$ .
- Optimization: The problem is formulated as minimizing a loss function $\ell(G)$ consisting of the squared Frobenius norm between $P_G$ and $\tilde{P}$ , plus a log-barrier penalty to enforce non-negativity:
  $\ell(G) = \|P_G - \tilde{P}\|_F^2 + \sum_q \mu_q \phi_q(G_q)$
  where $\phi_q$ is a log-barrier on the entries of $G_q$ .
- Algorithm: An alternating minimization approach is employed. For each tensor component $G_q$ $G_{q}$ , the algorithm updates the variable using a second-order Newton step.
  - Gradient/Hessian: Computed via message passing over the binary tree structure, utilizing Kronecker products of Gram matrices.
  - Efficiency: To avoid the $O(\rho^9)$ cost of inverting dense Hessians, the method uses Conjugate Gradient (CG) for middle nodes and decoupled solves for leaf nodes.
- Warm Initialization: A critical component is a "gauge warm initialization" strategy (Algorithm 2.3). Instead of naive thresholding of the signed reference, this method leverages the gauge freedom in tensor networks. It performs a decoupled sweep (leaves to root) to find orthogonal gauge matrices $Q$ and positive components $G_q$ that best fit the reference components $F_q$ in a "leave-one-out" fashion. This significantly improves convergence and avoids local minima.

Key Contributions

NHT Format for High-Dimensional Distributions: The paper introduces the use of the Non-Negative Hierarchical Tucker format, which is better suited for 2D lattice models and complex interaction structures than the Tensor-Train (TT) format, while maintaining the non-negativity required for probability distributions.
Second-Order Fitting Procedure: The authors design a tailored second-order optimization method for the NHT fitting stage. This includes a novel warm initialization strategy that addresses the difficulty of reaching local minima in high-dimensional settings, often improving final accuracy by an order of magnitude compared to standard initialization.
Computational Complexity: Both stages of the proposed methodology admit $O(d)$ computational complexity, making the approach scalable to high-dimensional settings.
End-to-End Algorithm: The work provides a complete pipeline for both Variational Inference and Density Estimation scenarios, bridging the gap between constructing a signed reference and fitting a physically valid non-negative model.

Experimental Results
Numerical experiments were conducted on both access models:

Variational Inference: Tested on a 2D Ginzburg-Landau model ( $d=32$ ) and a heavy-tailed model ( $d=30$ ). The adaptive-barrier preconditioned Newton iteration significantly outperformed a multiplicative-update benchmark, achieving lower relative squared Frobenius loss within the same wall-clock time.
Density Estimation: Tested on a 2D Ising model ( $d=32$ ) and a 2D Transverse-Field Ising Model ( $d=64$ ). The method successfully approximated the ground-truth distribution, with Negative Log-Likelihood (NLL) values close to the exact values.
Ablation Studies:
- NHT vs. NTT: In the 2D Ginzburg-Landau model, the NHT ansatz achieved better fitting accuracy than the Non-negative Tensor-Train (NTT) ansatz, even when the NTT model was given double the number of parameters.
- Initialization: The proposed gauge warm initialization (Algorithm 2.3) consistently yielded better fitting results than initializing with a multiplicative update algorithm (Algorithm 3.1).

Significance and Claims
The paper claims that the proposed NHT methodology offers a robust solution for compressing high-dimensional discrete distributions where traditional tensor-train models fail due to structural constraints (e.g., 2D lattices). By ensuring non-negativity through the ansatz and utilizing a sophisticated second-order fitting procedure with a specialized warm start, the method enables accurate density estimation and variational inference without the sign problems associated with unconstrained tensor networks. The authors note that while the method is illustrated on complete binary trees, the code supports incomplete trees, and future work could extend this to functional hierarchical tensor ansatzes or incorporate moment conservation constraints.