Stochastic tensor contraction for quantum chemistry

Abstract

Many computational methods in ab initio quantum chemistry are formulated in terms of high-order tensor contractions, whose cost determines the size of system that can be studied. We introduce stochastic tensor contraction to perform such operations with greatly reduced cost, and present its application to the gold-standard quantum chemistry method, coupled cluster theory with up to perturbative triples. For total energy errors more stringent than chemical accuracy, we reduce the computational scaling to that of mean-field theory, while starting to approach the mean-field absolute cost, thereby challenging the existing cost-to-accuracy landscape. Benchmarks against state-of-the-art local correlation approximations further show that we achieve an order-of-magnitude improvement in both total computation time and error, with significantly reduced sensitivity to system dimensionality and electron delocalization. We conclude that stochastic tensor contraction is a powerful computational primitive to accelerate a wide range of quantum chemistry.

Figures (7)

• Figure 1: Workflow of stochastic tensor contraction in quantum chemistry. Starting from the chemical structure and an initial mean-field calculation, input tensors are created for electronic correlation computations. The physical quantities (energy) are contractions of these tensors, typically involving loopy contractions. The loopy contractions are computed stochastically by importance sampling from a distribution constructed by breaking the loops in the contraction.
• Figure 2: (A) Illustration of optimal sampling for tree tensor contractions. An example $S_{ilm}=\sum_{jk}A_{ijk}B_{jl}C_{km}$ is considered. One can exactly sample the indices from the optimal probability distribution $\tilde{p}_{ijklm}^{\text{opt}}=|A_{ijk}B_{jl}C_{km}|$ by the procedure (1) sample $i$, (2) sample $j,k$ conditioned on $i$, (3) sample $l$ conditioned on $j$, (4) sample $m$ conditioned on $k$. All exact marginal and conditional probability tables can be constructed with a cost proportional to the tensor sizes. With the tables constructed, one can sample a set of indices with $O(1)$ cost. (B) Illustration of the loop breaking strategy for general loopy tensor contractions with the example $S=\sum_{ijk} A_{ik}B_{ij}C_{jk}$. We apply the approximate decomposition $|\boldsymbol{A}|\leq\boldsymbol{P}\otimes\boldsymbol{Q}$ to the optimal sampling distribution $\tilde{p}_{ijk}^{\text{opt}}=|A_{ik}B_{ij}C_{jk}|$, which results in $\tilde{p}_{ijk}^{\prime}=|P_{i}B_{ij}C_{jk}Q_{k}|\geq\tilde{p}_{ijk}^{\text{opt}}$. $\tilde{p}_{ijk}^{\prime}$ has a tree structure and thus can be exactly sampled by the previous strategy. (C) Two representative CCSD tensor contraction terms in $\mathcal{C}_{\mathrm{CC}}^{(2)}$ (Eq. \ref{['eq:CCSD_contraction']}), and a representative (T) tensor contraction term in $\mathcal{C}_{\text{(T)}}$ (Eq. \ref{['eq:(T)_contraction']}).
• Figure 3: Numerical scalings and computation time of STC-CCSD(T) and comparison with exact CCSD(T) demonstrated on water clusters with $2\sim30$ water molecules. (A) Scaling of $N_{\text{sample}}^{\epsilon}$ and $N_{\text{critical}}$ for STC-CCSD with local and canonical (labeled as "Can.") bases. (B) Scaling of the number of floating-point operations (FLOP) for STC and exact CCSD and (T). (C) Total computation time of the complete STC and exact CCSD(T) calculation. The reported number of samples or floating point operations of CCSD is for each iteration. In STC calculations, the target total error is fixed to be 0.2m$E_{h}$ for both CCSD and (T). All structures are lowest-energy structures, and all calculations use the 6-31G basis and are performed on the same CPU with 8 cores.
• Figure 4: Statistical histograms of (A) STC-CCSD and (B) STC-(T) energy errors from $n=2500$ independent STC-CCSD(T) calculations for a single benzene molecule in the cc-pVTZ basis, relative to exact CCSD(T) energies. The target per-sample statistical error is set to $\epsilon=0.25\text{m}E_{h}$ for STC-CCSD and $\epsilon=0.2\text{m}E_{h}$ for STC-(T) with input tensors from STC-CCSD. Black vertical dashed lines indicate the sample means of the 2500 calculations, which are $2.53\mu E_{h}$ for STC-CCSD and $3.40\mu E_{h}$ for STC-(T). Black solid lines are the predicted error distributions $N(0, \epsilon^2)$ from the target error $\epsilon$. The sample standard deviations, $\sigma=$0.257 m$E_{h}$ for STC-CCSD and $\sigma=$0.206 m$E_{h}$ for STC-(T), are consistent with the target errors. Gray shaded regions denote the standard error of the mean (SEM), given by $\sigma/\sqrt{n}$, corresponding to $5\mu E_{h}$ and $4\mu E_{h}$ for STC-CCSD and STC-(T), respectively. In both cases, zero lies within the SEM, indicating the absence of detectable bias in the $\mu E_{h}$ range.
• Figure 5: Dependence of STC-CCSD(T) on system dimensionality and orbital delocalization compared with DLPNO-CCSD(T), demonstrated on finite hydrogen-terminated h-BN (H-hBN) and polycyclic aromatic hydrocarbon (PAH) clusters. Four geometries, 1× 13, 2× 8, 3× 5, and 4× 4, are used to capture the transition from quasi-one-dimensional to two-dimensional geometries, thereby varying the effective dimensionality. The Hartree-Fock gaps of H-hBN and PAH are $\sim 0.5$ and $\sim 0.2$$E_{h}$ for the 4 geometries, respectively, yielding different levels of delocalization. The STC-CCSD(T) target error is set to be $0.2$ kcal/mol, while DLPNO-CCSD(T) calculations use the TightPNO setting. (A) Total computation time of STC and DLPNO CCSD(T) for H-hBN and PAH. (B) Total energy error of STC and DLPNO CCSD(T) for H-hBN and PAH. (C,D) Normalized number of samples in (C) iterative STC-CCSD and (D) STC-(T) for H-hBN and PAH calculations. An example of the 3x5 H-hBN geometry is shown in Fig. \ref{['fig:locality']}(B). All computations are performed with the frozen-core approximation, cc-pVDZ basis and 8 CPU cores. In all STC-CCSD(T) calculations, the true energy errors compared to the exact reference are consistent with the target error.