Wavelet-based density sketching with functional hierarchical tensor

TL;DR

This work introduces the functional hierarchical tensor under a wavelet basis (FHT-W) for high-dimensional density estimation on lattice models, addressing the limited capacity of prior FTN approaches under strong coupling by exploiting wavelet multiresolution to reduce effective rank. The method builds a tree-based FTN operating on wavelet coordinates, enabling exact normalization via tensor contractions and a sketched linear-algebra procedure to learn density components. Empirical results on 1D and 2D Ornstein–Uhlenbeck and Ginzburg–Landau models show that wavelet coordinates substantially reduce numerical rank (p(c) < p(x)) and enable accurate density estimation from samples. Extensions to Kolmogorov backward equations and stochastic control are discussed, highlighting the framework’s potential for scalable, multiscale probabilistic modeling and PDE-solving on lattice systems.

Abstract

We introduce the functional hierarchical tensor under a wavelet basis (FHT-W) ansatz for high-dimensional density estimation in lattice models. Recently, the functional tensor network has emerged as a suitable candidate for density estimation due to its ability to calculate the normalization constant exactly, a defining feature not enjoyed by neural network alternatives such as energy-based models or diffusion models. While current functional tensor network models show good performance for lattice models with weak or moderate couplings, we show that they face significant model capacity constraints when applied to lattice models with strong coupling. To address this issue, this work proposes to perform density estimation on the lattice model under a wavelet transformation. Motivated by the literature on scale separation, we perform iterative wavelet coarsening to separate the lattice model into different scales. Based on this multiscale structure, we design a new functional hierarchical tensor ansatz using a hierarchical tree topology, whereby information on the finer scale is further away from the root node of the tree. Our experiments show that the numerical rank of typical lattice models is significantly lower under appropriate wavelet transformation. Furthermore, we show that our proposed model allows one to model challenging Gaussian field models and Ginzburg-Landau models.

Figures (11)

• Figure 1: (A) A tree structure $T = (V, E)$ with $V = V_{\mathrm{ext}}= \{1, \ldots, 10\}$. (B) Tensor Diagram representation of a tree tensor network over $T$.
• Figure 1: Illustration of the FHT-W ansatz for $L = 4$. Each $v_{k, l}$ represents the tensor component $G_{v_{k, l}}$ at the external node $v_{k, l}$, and the $v_{k, l}$ corresponds to the variable $c_{k, l}$ in the wavelet transformation. Each $w_{k, l}$ represents the tensor component $G_{w_{k, l}}$ at the internal node $w_{k, l}$.
• Figure 1: Results of the $d = 2$ model $p(x_1, x_2) = \exp\left(-x_1^2/2 - x_2^2/2 - \frac{\alpha}{2}(x_1 - x_2)^2\right)$. \ref{['fig:bivariate_OU_plot_combined']}(a)-(b) shows the contour plot of $p(x)$ and $p(c)$. \ref{['fig:bivariate_OU_plot_combined']}(c) shows that the numerical rank of $p(x)$ increases with the coupling strength parameter $\alpha$.
• Figure 1: 1D Ginzburg-Landau model. Plots of the marginal distribution of $C \sim p(c)$ at $(c_{15, 5}, c_{8, 4})$ and $(c_{15, 5}, c_{9, 4})$. For illustration purposes, a scaling is performed so that the marginal distribution lies in $[-1,1]^2$.
• Figure 2: Illustration of the chosen tree structure for 2D wavelet iterative coarsening for a $d = 8^2$ lattice system. We use circles to illustrate nodes at $q=2, 0$, and we use squares to illustrate nodes at $q=1$. Internal nodes on the tree $T$ are omitted for simplicity.

Theorems & Definitions (1)

• Definition 2.1