Wavelet-based density sketching with functional hierarchical tensor

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.