Combining monte carlo and tensor-network methods for partial differential equations by sketching

Our primary objective in this paper is to obtain an MPS/TT representation of the solution of the initial value problem (3.1) at any time $t\geq 0$. Since the technique works for $u(x,t)$ at any given $t\geq 0$, we omit $t$ from the expression and use $u(x)$ to denote an arbitrary $d$-dimensional function:

And the state spaces $X_{1},\cdots,X_{d}\subseteq\mathbb{R}$.

We present the tensor diagram depicting the MPS/TT representation in Fig. 1(a). In this diagram, each tensor core is represented by a node, and it possesses one exposed leg that represents the degree of freedom associated with the corresponding dimension.

The tensor network diagram corresponding to the MPO representation is depicted in Fig. 1(b). For more comprehensive discussions on tensor networks and tensor diagrams, we refer interested readers to [tt_committor].

Finally, we introduce some indexing conventions. For two integers $m,n\in\mathbb{N}$ where $n>m$, we use the MATLAB notation $m:n$ to denote the set $\{m,m+1,\cdots,n\}$. When working with high-dimensional functions, it is often convenient to group the variables into two subsets and think of the resulting object as a matrix. We call these matrices unfolding matrices. In particular, for $k=1,\cdots,d-1$, we define the $k$-th unfolding matrix by $u(x_{1},\cdots,x_{k};x_{k+1},\cdots,x_{d})$ or $u(x_{1:k};x_{k+1:d})$, which is obtained by grouping the first $k$ and last $d-k$ variables to form rows and columns, respectively.

In this subsection, we introduce several tensor network operations of high importance in our applications using the example of MPS/TT. Similar operators can be extended to more general tensor networks.

In this subsection, we present a parallel method for obtaining the tensor cores for representing a $d$-dimensional function $u(x)$ that is discretely valued, i.e. each $x_{j}$ takes on finite values in a set $X_{j}$, as an MPS/TT. This is done via an MPS/TT sketching technique proposed in [hur2022generative] where the key idea is to solve a sequence of core determining equations. Let $u_{k}:X_{1}\times\cdots\times X_{k-1}\times\Gamma_{k-1}\rightarrow\mathbb{R},k=2,\ldots,d$ be a set of functions such that $\text{Range}(u_{k}(x_{1:k};\gamma_{k}))=\text{Range}(p(x_{1:k};x_{k+1:d}))$. In this case, a representation of $u$ as Definition 1 can be obtained via solving $\mathcal{G}_{k}$ from the following set of equations

based on knowledge of $u_{k}$’s.

However (2.6) is still inefficient to solve since each $u_{k}$ is exponentially sized, moreover, such a size prohibits it to be obtained/estimated in practice. Notice that (2.6) is over-determined, we further reduce the row dimensions by applying a left sketching function to (2.6),

where $S_{k-1}:X_{1}\times\cdots\times X_{k-1}\times\Xi_{k-1}\rightarrow\mathbb{R}$ is the left sketching function which compresses over variables $x_{1},\cdots,x_{k-1}$.

Now to obtain $u_{k}$ where $\text{Range}(u_{k}(x_{1:k};\gamma_{k}))=\text{Range}(u(x_{1:k};x_{k+1:d}))$, we use a right-sketching by sketching the dimensions $x_{k+1:d}$, i.e.

where $T_{k+1}:X_{k+1}\times\cdots\times X_{d}\times\Gamma_{k}\rightarrow\mathbb{R}$ is the right sketching function which compresses $u$ by contracting out variables $x_{k+1},\cdots,x_{d}$. Plugging such a $u_{k}$ into (2.6), we get

where

and we can readily solve for $\mathcal{G}_{k}$.

Many different types of sketch functions can be used, e.g. random tensor sketches or cluster basis sketches [ahle2019almost, wang2015fast]. Take a single $S_{k}(x_{1:k},\xi_{k})$ as an example, we choose a separable form $S_{k}(x_{1:k},\xi_{k})=h_{1}(x_{1})\cdots h_{k}(x_{k})$ for some $h_{1},\ldots,h_{k}$ (and similarly for $T_{k}$’s) in order to perform fast tensor operations. When the state space is discrete, random tensor sketch amounts to taking $h_{1},\cdots,h_{k}$ to be random vectors of size $|X_{1}|,\ldots,|X_{k}|$.

We give special focus to cluster basis sketch, defined as the following:

A similar construct is used in [peng2023generative]. In this case, $A_{k}$ is of size estimate ${k-1\choose c}n^{c}\times{d-k+1\choose c}n^{c}$, and $B_{k}$ is of size ${k-1\choose c}n^{c}\times|X_{k}|\times{d-k+2\choose c}n^{c}$. In principle, we only need ${d-k+2\choose c}n^{c}>r$ for determining a rank-$r$ MPS/TT, therefore in practice $c\leq 2$ meets the need. The most important property we look for is that the variance $A_{k}[\hat{u}],B_{k}[\hat{u}]$ is small, if $\hat{u}$ is an unbiased estimator of $u$. As we shall see in Section 5, when $u$ is a density and $\hat{u}$ is an empirical distribution with $N$ samples approximating $u$, the choice of such a cluster basis results in the standard deviation of $A_{k}[\hat{u}],B_{k}[\hat{u}]$ being $O\left(\frac{n^{2c+1}}{\sqrt{N}}\right)$. In contrast, if one chooses a randomized function that involves all $d$ variables as a sketch, the variance would be $O(n^{d}/\sqrt{N})$.

The concept of MPS/TT sketching or general tensor network sketching for density estimations has been extensively explored in [hur2022generative, ren2022high, tang2022generative]. The choice of tensor network representation in practice depends on the problem’s structure. In this work, we demonstrate the workflow using MPS/TT, but the framework can be readily extended to other tensor networks.

A crucial assumption underlying our approach is that each particle $f(x;x^{i})$ exhibits a simple structure and can be efficiently represented or approximated in MPS/TT format. For instance, in the case of a vanilla Markov Chain Monte Carlo (MCMC), we have $f(x;x^{i})=\delta(x-x^{i})=\prod_{j=1}^{d}\delta(x_{j}-x_{j}^{i})$, which can be expressed as a rank-1 MPS/TT. Consequently, the right-hand side of (3.2) becomes a summation of $N$ MPS/TTs. The tensor diagram illustrating this general particle approximation is shown in Fig. 1(a). A naive approach for representing this sum as an MPS/TT consists of directly adding these MPS/TTs, which yields an MPS/TT with a rank of $N$ [oseledets2011tensor]. This growth in rank with the number of samples can lead to a high computational complexity of $O(\text{poly}(N))$ when performing tensor operations.

On the other hand, depending on the problem, if we know that the intrinsic MPS/TT rank of the solution is small, it should be possible to represent the MPS/TT with a smaller size. To achieve this, we apply the method described in Section 2.3 to estimate a low-rank tensor directly from the sum of $N$ particles. Specifically, we construct $A_{k}[\hat{\phi}_{t+1}]$ and $B_{k}[\hat{\phi}_{t+1}]$ from the empirical samples $\hat{\phi}_{t+1}$ and use them to solve for the tensor cores through (2.10), resulting in $\phi_{\theta_{t+1}}$. This process is illustrated in Fig. 1(b), Fig. 1(c), and Fig. 1(d). It is important to note that we approximate the ground truth $\phi_{t+1}$ using stochastic samples $\hat{\phi}_{t+1}$. The estimated tensor cores form an MPS/TT representation of $\phi_{t+1}$, denoted as $\phi_{\theta_{t+1}}$.

In this section, we assume we are given $\hat{\phi}_{t+1}$ in (3.1) as $N$ constant rank MPS/TT. In terms of samples, forming $A_{k}[\hat{\phi}_{t+1}])$ and $B_{k}[\hat{\phi}_{t+1}]$ in (2.11) and (2.12) can be done as

$B_{k}$ is a $3$-tensor of shape $|\Xi_{k-1}|\times|X_{k}|\times|\Gamma_{k}|$. $A_{k}$ is a matrix of shape $|\Xi_{k-1}|\times|\Gamma_{k}|$. The size of the linear system is independent of the number of dimensions $d$ and the total number of particles $N$. The number of sketch functions $|\Xi_{k}|$ and $|\Gamma_{k}|$ are hyperparameters we can control. Let $n=\max_{k}|X_{k}|$, $\tilde{r}=\max_{k}\{|\Xi_{k}|,|\Gamma_{k}|\}$. First we consider the complexity of forming the core determining equations, i.e. evaluating $A_{k}$ and $B_{k}$. We note that the complexity is dominated by evaluating the tensor contractions between sketch functions $S_{k}$’s, right sketch functions $T_{k}$’s and samples $f(\cdot;x^{i})$’s. If the sample $f(\cdot;x^{i})$ and sketch functions $S_{k}(\cdot,\xi_{k})$, $T_{k}(\cdot,\gamma_{k-1})$ are in low rank MPS/TT representations, then the tensor contractions in the square bracket in (3.3) can be evaluated in $O(d)$ time. Each term in the summation $\sum_{i=1}^{N}$ can be done independently, potentially even in parallel, giving our final $A_{k}$’s and $B_{k}$’s. Taking all $d$ dimensions into account, the total complexity of evaluating $A_{k}$’s and $B_{k}$’s is $O(n\tilde{r}^{2}Nd)$.

Next, the complexity of solving the linear system (2.10) is $O(n\tilde{r}^{3})$. For all $d$ dimensions, the total complexity for solving the core determining equations is $O(n\tilde{r}^{3}d)$. Combing everything together, the total computational complexity for MPS/TT sketching for a discrete particle system is

In this subsection, we apply the proposed framework to ground state energy estimation problems in quantum mechanics for the spin system. Statistical sampling-based approaches have been widely applied to these types of problems, see e.g. [lee2022unbiasing]. In this problem, we want to solve

where $\phi(\cdot,t):\{\pm 1\}^{d}\rightarrow\mathbb{C}$, and $H$ is the Hamiltonian operator. Let $O:=[O_{i}]^{d}_{i=1}$, where

for some $\tilde{O}\in\mathbb{C}^{2\times 2}$, $O_{i}^{2}=I_{2^{d}}$, usually $H$ takes the form of $H_{1}[O]=\sum_{i=1}^{d}O_{i}$ or $H_{2}[O]=\sum_{i,j=1}^{d}J_{ij}O_{i}O_{j}$ or both. When $t\rightarrow\infty$, $\phi(\cdot,t)=\frac{\exp(-Ht)\phi_{0}}{\|\exp(-Ht)\phi_{0}\|_{2}}$ gives the lowest eigenvector of $H$. This can be done with the framework detailed in Section 3 with $A=H$. In what follows, we detail how $\exp(-\delta tH)\phi_{\theta_{t}}$ can be approximated as a sum of $N$ functions $f(x;x^{i})$ via a specific version of quantum Monte Carlo, the auxiliary-field quantum Monte Carlo (AFQMC). The AFQMC method [afqmc1, afmc2] is a powerful numerical technique that has been developed to overcome some of the limitations of traditional Monte Carlo simulations. AFQMC is based on the idea of introducing auxiliary fields to decouple the correlations between particles by means of the application of the Hubbard–Stratonovich transformation [hubbard1959calculation]. This reduces the many-body problem to the calculation of a sum or integral over all possible auxiliary-field configurations. The method has been successfully applied to a wide range of problems in statistical mechanics, including lattice field theory, quantum chromodynamics, and condensed matter physics [zhang2013auxiliary, carlson2011auxiliary, shi2021some, qin2016benchmark, lee2022twenty]. However, an issue of AFQMC is that the random walkers are biased towards a mean-field solution [qin2016coupling] in order to reduce the variance caused by the “sign problem”, giving rise to bias when determining the ground state energy. To solve this issue, [qin2016coupling] alternates between running AFQMC and determining a new mean-field solution in order to remove such a bias. Our approach can be regarded as a generalization to the philosophy in [qin2016coupling], where the mean-field is now replaced by a more general MPS/TT representation.

We first focus on how to apply $\exp(-\delta tH_{2}[O])$ where $H_{2}[O]=\sum^{d}_{i,j=1}J_{ij}O_{i}O_{j}$. We assume $J$ is positive semi-definite, which can be easily achieved by adding a diagonal to $J$ since $O_{i}^{2}=I_{2^{d}}$. To decouple the interactions through orthogonal transformations, we first perform an eigenvalue decomposition $J=U\text{diag}(\lambda)U^{T}$ and compute

where $U=[u_{1},\ldots,u_{j}],u_{j}\cdot O:=\sum_{i=1}^{d}u_{ji}O_{i}$. The approximation follows from Trotter-decomposition [hatano2005finding]. Using the identity

we can write (4.3) as

Instead of sampling directly from the spin space, AFQMC approximates the $d$-dimensional integration in (4.5) by a Monte Carlo integration in the $k$-space. Replacing the integration with Monte Carlo samples, we get the following approximation,

where $N$ is the total number of Monte Carlo samples, $k^{i}_{j}$ is the $i$-th sample for the frequency variable of the $j$-th dimension, and $\tilde{k}^{i}=\sum^{d}_{j=1}\sqrt{2\delta t\lambda_{j}}k^{i}_{j}u_{j}$. All samples $\{k^{i}_{j}\}_{j=1,\cdots,d,\ i=1,\cdots,N}$ are drawn from i.i.d. standard normal distribution. The structure of such an approximation to $\exp(-\delta tH_{2}[O])$ is illustrated in Fig. 1(a), which is a sum of rank-1 MPO. The application of $\exp(-\delta tH_{2}[O])$ to a function $u$ that is already represented as an MPS/TT in the form of Def. 1 can written as

Notably, since each $\exp(\sqrt{-1}(\tilde{k}^{i}\cdot O))$ is a rank-1 MPO, $\exp(\sqrt{-1}(\tilde{k}^{i}\cdot O))\phi_{t}$ has the same rank as $\phi_{t}$. Lastly, the application of $\exp(-\delta tH_{1}[O])$ to some $\phi_{t}$ that is represented as an MPS can be done easily, since $\exp(-\delta tH_{1}[O])=\prod_{i=1}^{d}\exp(-\delta tO_{i})$ is a rank-1 MPO.

To summarize, we can approximate the imaginary time propagator $\exp(-\delta t(H_{1}[O]+H_{2}[O]))$ with a summation of $N$ rank-$1$ MPOs. Applying the propagator to many-body wavefunction $\phi_{t}$ reduces to MPO-MPS contractions. The corresponding tensor diagram is shown in Fig. 1(b).

In this subsection, we demonstrate the application of the proposed framework to numerical simulations of parabolic PDEs, specifically focusing on the overdamped Langevin process and its corresponding Fokker-Planck equations. We consider a particle system governed by the following overdamped Langevin process,

where $x_{t}\in\Omega\subseteq\mathbb{R}^{d}$ is the state of the system, $V:\Omega\subset\mathbb{R}^{d}\rightarrow\mathbb{R}$ is a smooth potential energy function, $\beta=1/T$ is the inverse of the temperature $T$, and $W_{t}$ is a $d$-dimensional Wiener processs. If the potential energy function $V$ is confining for $\Omega$ (see, e.g., [bhattacharya2009stochastic, Definition 4.2]), it can be shown that the equilibrium probability distribution of the Langevin dynamics (4.8) is the Boltzmann-Gibbs distribution,

where $Z_{\beta}=\int_{\Omega}\exp(-\beta V(x))\,dx$ is the partition function. Moreover, the evolution of the distribution of the particle system can be described by the corresponding time-dependent Fokker-Planck equation,

where $\phi_{0}$ is the initial distribution. $\|\phi(\cdot,t)\|_{1}=1$ ensures that the $\int|\phi(x,t)|dx=1$. Therefore, (4.10) is the counterpart of (3.1) in our framework.

Now we need to be able to approximate $\exp(-A\delta t)\phi_{t}$ as particle systems. Assuming the current density $\phi_{t}$ is a MPS/TT, we use the following procedures to generate a particle approximation $\hat{\phi}_{t+1}$:

1. (1)

   We apply conditional sampling on the current density estimate $\phi_{\theta_{t}}$ (Section 2.2.3) to generate $N$ i.i.d. samples $x^{1},\ldots,x^{N}\sim\phi_{\theta_{t}}$.

2. (2)

   Then, we simulate the overdamped Langevin dynamics (4.8) using Euler-Maruyama method over time interval $\delta t$ for each of the $N$ initial stochastic samples $x^{1},\ldots x^{N}\sim\phi_{\theta_{t}}$. By the end of $\delta t$ we have final particle positions $x^{1},\ldots,x^{N}\sim\phi_{t+1}$, and

   (4.11)

   $\displaystyle\hat{\phi}_{t+1}(x)$

   $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\delta(x-x^{i}),$