# xTC: An efficient treatment of three-body interactions in transcorrelated
methods

Evelin Martine Christlmaier<EMAIL_ADDRESS>Max Planck Institute
for Solid State Research, Heisenbergstraße 1, 70569 Stuttgart, Germany  Thomas
Schraivogel  Max Planck Institute for Solid State Research, Heisenbergstraße
1, 70569 Stuttgart, Germany  Pablo López Ríos  Max Planck Institute for Solid
State Research, Heisenbergstraße 1, 70569 Stuttgart, Germany  Ali Alavi  Max
Planck Institute for Solid State Research, Heisenbergstraße 1, 70569
Stuttgart, Germany  Yusuf Hamied Department of Chemistry, University of
Cambridge, Lensfield Road, Cambridge CB2 1EW, United Kingdom  Daniel Kats
<EMAIL_ADDRESS>Max Planck Institute for Solid State Research,
Heisenbergstraße 1, 70569 Stuttgart, Germany

###### Abstract

An efficient implementation for approximate inclusion of the three-body
operator arising in transcorrelated methods via exclusion of explicit three-
body components (xTC) is presented and tested against results in the “HEAT”
benchmark set [A. Tajti et al., J. Chem. Phys. 121, 11599 (2004)]. Using
relatively modest basis sets and computationally simple methods, total,
atomization, and formation energies within near-chemical accuracy from HEAT
results were obtained. The xTC ansatz reduces the nominal scaling of the
three-body part of transcorrelation by two orders of magnitude to
$\mathcal{O}(N^{5})$ and can readily be used with almost any quantum chemical
correlation method.

##  I Introduction

Achieving chemical accuracy in computational quantum chemistry calculations is
still a significant challenge for all but the most simple systems. Convergence
of the correlation energy with basis set size is often slow and high-quality
correlation methods are required to obtain accurate results, leading to
rapidly escalating computational cost.

One of the reasons for the slow convergence of the energy with the size of the
basis set are the discontinuities in the first derivative of the wave function
at the electron coalescence points ${\mathbf{r}}_{ij}=0$, i.e., the Kato cusp.
Kato (1957) Explicitly correlated methods aim to address this issue by
introducing explicit dependence on the inter-electronic distances into the
wave function, leading to significantly improved convergence to the basis set
limit. Kutzelnigg (1985); Kutzelnigg and Klopper (1991); Klopper and Samson
(2002); Manby (2003); Ten-no (2004a, b); Valeev (2004); Tew and Klopper
(2005); Kedžuch, Milko, and Noga (2005a); Fliegl, Klopper, and Hättig (2005);
Fliegl, Hättig, and Klopper (2006); Werner, Adler, and Manby (2007a); Noga,
Kedžuch, and Šimunek (2007); Adler, Knizia, and Werner (2007); Tew  _et al._
(2007); Knizia and Werner (2008); Shiozaki  _et al._ (2008a, b); Noga  _et
al._ (2008); Tew, Klopper, and Hättig (2008); Valeev (2008); Valeev and
Crawford (2008); Torheyden and Valeev (2008); Bokhan, Ten-no, and Noga (2008);
Knizia, Adler, and Werner (2009a); Werner, Knizia, and Manby (2011); Ten-no
(2012); Hättig  _et al._ (2012); Kong, Bischoff, and Valeev (2012) However,
extension to higher-order methods is non-trivial.

Alternatively, the electronic wavefunction $\Psi$ can be factorized into a
Jastrow factor $e^{\tau}$ and a smoother function $\Phi$

$\displaystyle\Psi=e^{\tau}\Phi$ (1)

with the correlation factor

$\displaystyle\tau=\sum_{i<j}u({\mathbf{r}}_{i},{\mathbf{r}}_{j})$ (2)

describing the correlation of the electron pairs via the symmetric functions
$u({\mathbf{r}}_{i},{\mathbf{r}}_{j})$. Jastrow (1955) Transferring the
Jastrow factor to the Hamiltonian via a similarity transformation

$\displaystyle e^{-\tau}He^{\tau}\Phi={\tilde{H}}\Phi=E\Phi$ (3)

gives rise to the non-hermitian transcorrelated Hamiltonian, $\tilde{H}$. The
transcorrelated ansatz was pioneered by Boys and HandyBoys and Handy (1969)
and has seen renewed interest in recent years. Ten-no (2000); Hino, Tanimura,
and Ten-no (2002); Umezawa and Tsuneyuki (2003); Yanai and Chan (2006, 2007);
Yanai and Shiozaki (2012); Tsuneyuki (2008); Ochi  _et al._ (2012); Ochi and
Tsuneyuki (2014, 2015); Ochi  _et al._ (2016); Wahlen-Strothman  _et al._
(2015); Luo and Alavi (2018); Dobrautz, Luo, and Alavi (2019); Cohen  _et al._
(2019); Baiardi and Reiher (2020); Khamoshi  _et al._ (2021); Giner (2021);
Guther  _et al._ (2021); Liao  _et al._ (2021); Schraivogel  _et al._ (2021);
Liao  _et al._ (2023); Haupt  _et al._ (2023); Schraivogel  _et al._ (2023)
Compared to the explicitly correlated methods, the transcorrelated ansatz has
the advantage that it not only improves the basis set convergence, but also
the quality of the underlying correlation method.

A major bottleneck of transcorrelated methods is the calculation of the three-
body integrals arising from the similarity transformation of the Hamiltonian.
Further, they require the implementation of additional modifications in the
subsequent correlation methods that cost both additional CPU and developer
time. Recently, approximate transcorrelated coupled cluster Čížek (1966);
Purvis and Bartlett (1982); Raghavachari  _et al._ (1989) and distinguishable
cluster Kats and Manby (2013); Kats (2014); Kats  _et al._ (2015a); Kats
(2016, 2018) methods that drop the pure three-body part of the normal ordered
similarity transformed Hamiltonian have been proposed and shown to introduce
only negligible errors in the final energies for atoms and
moleculesSchraivogel  _et al._ (2021, 2023) as well as the three-dimensional
uniform electron gasLiao  _et al._ (2021). However, explicit calculation of
the costly three-body integrals via numerical integration was still required.

In this work, we introduce an efficient algorithm for the inclusion of the
transcorrelated three-body operator excluding explicit three-body components
(xTC) via modifications of the zero, one- and two-body integrals. Instead of
calculating the three-body integrals explicitly, the corrections to the
integrals are obtained directly from the intermediates of the numerical
integration, leading to a reduction of the scaling of the matrix-element
evaluation stage by two orders of magnitude with respect to a naïve
implementation. Extensive benchmarking results using different correlation
methods, basis sets and additional approximations are presented.

The paper is organized as follows. In section II we will derive the working
equations for the xTC ansatz in spin-orbital basis. Section III shows the
results of the benchmarking calculations. Section IV will offer a brief
summary of the results. Finally, spin-integrated working equations are shown
in the appendix.

##  II Theory

###  II.1 Transcorrelation

The transcorrelated Hamiltonian in second quantization is

$\displaystyle{\tilde{H}}=$ $\displaystyle
h_{P}^{Q}a^{P}_{Q}+\frac{1}{2}(V_{PR}^{QS}-K_{PR}^{QS})a^{PR}_{QS}$
$\displaystyle-\frac{1}{6}L_{PRT}^{QSU}a^{PRT}_{QSU}+E_{\text{nuc}}$ (4)

where $P$, $Q$, $R$, … are general spin-orbital indices and

$\displaystyle a^{P}_{Q}=a^{P}a_{Q}=a_{P}^{\dagger}a_{Q}$ (5a) $\displaystyle
a^{PR}_{QS}=a^{P}a^{R}a_{S}a_{Q}$ (5b) $\displaystyle
a^{PRT}_{QSU}=a^{P}a^{R}a^{T}a_{U}a_{S}a_{Q}.$ (5c)

are products of creation and annihilation operators. Note that here and
throughout the article we use the Einstein convention, i.e., we implicitly sum
over repeated indices on the right hand side of the equation as long as they
do not appear on the left hand side.

In addition to the standard integrals

$\displaystyle h_{P}^{Q}=\langle{P}|h|{Q}\rangle$ (6a) $\displaystyle
V_{PR}^{QS}=\langle{PR}{|}{QS}\rangle=({PQ}{|}{RS}),$ (6b)

the Baker-Campbell-Hausdorff expansion of the similarity transformed
Hamiltonian gives rise to the non-hermitian two-body integrals

$\displaystyle K_{PR}^{QS}=\frac{1}{2}\langle{PR}|\mathcal{P}^{1}_{2}($
$\displaystyle\nabla_{1}^{2}u({\mathbf{r}}_{1},{\mathbf{r}}_{2})+(\nabla_{1}u({\mathbf{r}}_{1},{\mathbf{r}}_{2}))^{2}$
$\displaystyle+2\nabla_{1}u({\mathbf{r}}_{1},{\mathbf{r}}_{2})\cdot\nabla_{1})|{QS}\rangle$
(7)

and the three-body integrals

$\displaystyle L_{PRT}^{QSU}=\langle{PRT}|\mathcal{P}_{2,3}^{1}($
$\displaystyle\nabla_{1}u({\mathbf{r}}_{1},{\mathbf{r}}_{2})$
$\displaystyle\cdot\nabla_{1}u({\mathbf{r}}_{1},{\mathbf{r}}_{3}))|{QSU}\rangle$
(8)

where we have used the symmetric permutation operators

$\displaystyle\mathcal{P}^{1}_{2}f(1,2)=f(1,2)+f(2,1)$ (9a)
$\displaystyle\mathcal{P}^{1}_{2,3}f(1,2,3)=f$ $\displaystyle(1,2,3)+f(2,1,3)$
$\displaystyle+f(3,2,1).$ (9b)

Aside from evaluation of the costly three-body integrals, inclusion of the
three-body operator $L_{N}$ requires modification of any theory using
transcorrelation to include terms interacting with a three-body operator in
the Hamiltonian. In order to find approximations that will solve these issues,
we will first introduce normal ordering in the next section.

###  II.2 Generalized normal ordering

Instead of defining normal-order with respect to a single determinant,
Kutzelnigg and Mukherjee proposed a generalized normal ordering with respect
to an arbitrary reference function $|{\Phi_{0}}\rangle$. Kutzelnigg and
Mukherjee (1997, 1999)

In the general normal ordering, the antisymmetrized products of Kronecker
deltas are replaced by the density matrices of the reference function

$\displaystyle\gamma^{P...}_{Q...}=\langle{\Phi_{0}}|a^{P...}_{Q...}|{\Phi_{0}}\rangle.$
(10)

For a three-body operator, this yields

$\displaystyle a^{PRT}_{QSU}=$
$\displaystyle\tilde{a}^{PRT}_{QSU}+\sum(-1)^{\mathcal{P}}\gamma^{P}_{Q}\tilde{a}^{RT}_{SU}$
$\displaystyle+\sum(-1)^{\mathcal{P}}\gamma^{PR}_{QS}\tilde{a}^{T}_{U}+\gamma^{PRT}_{QSU}$
(11)

with $\tilde{a}$ being normal ordered operators and the sums going over all
nine non-redundant permutations of (PRT) and (QSU) with the appropriate sign.

Introducing the combined two-electron integrals

$\displaystyle U_{PR}^{QS}=V_{PR}^{QS}-K_{PR}^{QS}$ (12)

the normal ordered Hamiltonian can be separated into effective one-, two- and
three-body operators

$\displaystyle{\tilde{H}_{N}=\tilde{H}-\langle{\Phi_{0}}|\tilde{H}|{\Phi_{0}}\rangle=}F_{N}+V_{N}+L_{N}$
(13)

with

$\displaystyle F_{N}=\bigg{[}$ $\displaystyle
h_{P}^{Q}+\big{(}U_{PR}^{QS}-U_{PR}^{SQ}\big{)}\gamma^{R}_{S}$
$\displaystyle-\frac{1}{2}\big{(}L_{PRT}^{QSU}-L_{PRT}^{SQU}-L_{PRT}^{USQ}\big{)}\gamma^{RT}_{SU}\bigg{]}\tilde{a}^{P}_{Q}$
(14a) $\displaystyle V_{N}=\frac{1}{2}\bigg{[}$ $\displaystyle
U_{PR}^{QS}-\big{(}L_{PRT}^{QSU}$ $\displaystyle-
L_{PRT}^{QUS}-L_{PRT}^{USQ}\big{)}\gamma^{T}_{U}\bigg{]}\tilde{a}^{PR}_{QS}$
(14b) $\displaystyle L_{N}=-\frac{1}{6}L_{PRT}^{QSU}\tilde{a}^{PRT}_{QSU}$
(14c)

###  II.3 The xTC approximation: excluding explicit three-body components

We can now exclude the explicit three-body components $L_{N}$ (xTC) and
incorporate the remaining 3-body contributions via a change of the two-, one-
and zero-body integrals.

$\displaystyle\bar{U}_{PR}^{QS}\leftarrow U_{PR}^{QS}+\Delta U_{PR}^{QS}$
(15a) $\displaystyle\bar{h}_{P}^{Q}\leftarrow h_{P}^{Q}+\Delta h_{P}^{Q}$
(15b) $\displaystyle\bar{E}_{\text{nuc}}\leftarrow
E_{\text{nuc}}+\langle{\Phi_{0}}|L|{\Phi_{0}}\rangle$ (15c)

Neglect of the pure normal ordered three-body operator eq. (14c) has been
shown to introduce only minor errors, even for normal ordering with respect to
the Hartree-Fock determinant. Schraivogel  _et al._ (2021, 2023)

The correction to the two-electron integrals is trivially obtained as

$\displaystyle\Delta
U_{PR}^{QS}=-\big{(}L_{PRT}^{QSU}-L_{PRT}^{QUS}-L_{PRT}^{USQ}\big{)}\gamma^{T}_{U}.$
(16)

Inserting $\bar{U}$ instead of $U$ into $F_{N}$ yields a new generalized Fock-
operator

$\displaystyle F_{N}^{\prime}=\bigg{[}h_{P}^{Q}+$
$\displaystyle\big{(}U_{PR}^{QS}-U_{PR}^{SQ}\big{)}\gamma^{R}_{S}$
$\displaystyle+$
$\displaystyle\big{(}L_{PRT}^{QSU}-L_{PRT}^{SQU}-L_{PRT}^{USQ}\big{)}$
$\displaystyle\times\big{(}-\gamma^{R}_{S}\gamma^{T}_{U}+\gamma^{R}_{U}\gamma^{T}_{S}-\frac{1}{2}\gamma^{RT}_{SU}\big{)}\bigg{]}\tilde{a}^{P}_{Q}$
(17)

that differs from the original generalized Fock-operator $F_{N}$ according to

$\displaystyle F_{N}=F_{N}^{\prime}-$
$\displaystyle\big{(}L_{PRT}^{QSU}-L_{PRT}^{SQU}-L_{PRT}^{USQ}\big{)}$
$\displaystyle\times\big{(}-\gamma^{R}_{S}\gamma^{T}_{U}+\gamma^{R}_{U}\gamma^{T}_{S}\big{)})\tilde{a}^{P}_{Q}.$
(18)

In order to ensure that the generalized Fock-operator remains invariant with
respect to the change in integrals, we combine the correction of the change
due to $\Delta U_{PR}^{QS}$ in $F_{N}$ with the one-body correction arising
due to the normal ordering of the three-body operator to obtain

$\displaystyle\Delta h_{P}^{Q}=$
$\displaystyle\big{(}L_{PRT}^{QSU}-L_{PRT}^{SQU}-L_{PRT}^{USQ}\big{)}$
$\displaystyle\times\big{(}\gamma^{R}_{S}\gamma^{T}_{U}-\gamma^{R}_{U}\gamma^{T}_{S}-\frac{1}{2}\gamma^{RT}_{SU}\big{)}.$
(19)

The zero-body correction is simply the expectation value of the three-body
operator

$\displaystyle\langle{\Phi_{0}}|L|{\Phi_{0}}\rangle=-\frac{1}{6}L_{PRT}^{QSU}\gamma^{PRT}_{QSU}.$
(20)

So far these equations are valid for an arbitrary reference function
$|{\Phi_{0}}\rangle$. In the next section we will show how a single-
determinant reference function greatly simplifies the equations.

###  II.4 Single determinant case

If the reference function $|{\Phi_{0}}\rangle$ is a single determinant, the
higher order density matrices are simply antisymmetrized products of the one-
body density matrix

$\displaystyle\gamma^{PR...}_{QS...}=\mathcal{A}(\gamma^{P}_{Q}\gamma^{R}_{S}...).$
(21)

In this case, the one- and zero-body corrections are trivially obtained from
contraction of the higher-body correction with the one-body density matrix

$\displaystyle\Delta h_{P}^{Q}=$
$\displaystyle\frac{1}{2}\big{(}L_{PRT}^{QSU}-L_{PRT}^{SQU}-L_{PRT}^{USQ}\big{)}$
$\displaystyle\times\big{(}\gamma^{R}_{S}\gamma^{T}_{U}-\gamma^{R}_{U}\gamma^{T}_{S}\big{)}$
$\displaystyle=$ $\displaystyle-\frac{1}{2}\big{(}\Delta U_{PR}^{QS}-\Delta
U_{PR}^{SQ}\big{)}\gamma^{R}_{S}$ (22)

$\displaystyle\langle{\Phi_{0}}|L|{\Phi_{0}}\rangle=-\frac{1}{3}\Delta
h_{P}^{Q}\gamma^{P}_{Q}$ (23)

While these two equations are only exact for a single-determinant reference,
they can still be used approximately with correlated densities (i.e.,
negelectng the higher order cumulants in the cumulant expansion of the
respective densityKutzelnigg and Mukherjee (1997, 1999)), which offers the
potential for additional flexibility in the xTC ansatz.

The equations given above still require explicit calculation of the L-matrix.
In the next section we will describe how we can obtain the corrections in (15)
without the need of costly six-index intermediates.

###  II.5 Decomposition of two-body correction

In our implementationCohen  _et al._ (2019), the three-body integrals are
obtained via numerical integration over the grid points $x_{i}$ with the
weights $w_{i}$

$\displaystyle L_{PRT}^{QSU}=P^{(PQ)}_{(RS),(TU)}\sum_{i}w_{i}$
$\displaystyle\phi_{P}^{*}(i){\mathbf{V}}_{R}^{S}(i)\cdot{\mathbf{V}}_{T}^{U}(i)\phi^{Q}(i)$
(24)

and the intermediate

$\displaystyle{\mathbf{V}}_{R}^{S}(i)=\sum_{j}w_{j}\phi_{R}^{*}(j)\big{(}\nabla
u({\mathbf{r}}_{i},{\mathbf{r}}_{j})\big{)}\phi^{S}(j).$ (25)

We use boldface here to indicate that ${\mathbf{V}}_{R}^{S}(i)$ and some of
the following intermediates are three-dimensional spatial vectors and use the
scalar product to indicate contraction over the spatial coordinates.

Making use of the decomposition eq. (24) allows us to avoid calculation of the
three-body integrals entirely and instead calculate $\Delta U$ directly. By
inserting eq. (21) into eq. (16) and changing the order of summation, we
obtain

$\displaystyle\Delta U_{PR}^{QS}=-\mathcal{P}^{(PQ)}_{(RS)}\big{(}$
$\displaystyle\rho_{P}^{Q}(i)A_{R}^{S}(i)+{\mathbf{V}}_{P}^{Q}(i)\cdot{\mathbf{B}}_{R}^{S}(i)\big{)}$
(26)

with the intermediates

$\displaystyle\rho_{P}^{Q}(i)=w_{i}\phi_{P}^{*}(i)\phi^{Q}(i)$ (27a)
$\displaystyle A_{R}^{S}(i)=\tilde{V}_{R}^{S}(i)-\tilde{Z}_{R}^{S}(i)$ (27b)
$\displaystyle\tilde{V}_{R}^{S}(i)={\mathbf{W}}(i)\cdot{\mathbf{V}}_{R}^{S}(i)$
(27c) $\displaystyle{\mathbf{W}}(i)={\mathbf{V}}_{T}^{U}(i)\gamma^{T}_{U}$
(27d)
$\displaystyle\tilde{Z}_{R}^{S}(i)={\mathbf{V}}_{R}^{U}(i)\cdot{\mathbf{X}}^{S}_{U}(i)$
(27e)
$\displaystyle{\mathbf{X}}^{S}_{U}(i)={\mathbf{V}}_{T}^{S}(i)\gamma^{T}_{U}$
(27f)
$\displaystyle{\mathbf{B}}_{R}^{S}(i)=\frac{1}{2}\tilde{W}(i){\mathbf{V}}_{R}^{S}(i)-{\mathbf{Z}}_{R}^{S}(i)$
(27g) $\displaystyle\tilde{W}(i)=\rho_{T}^{U}(i)\gamma^{T}_{U}$ (27h)
$\displaystyle{\mathbf{Z}}_{R}^{S}(i)=\rho_{R}^{U}(i){\mathbf{X}}^{S}_{U}(i)+{\mathbf{Y}}_{R}^{T}(i)\rho_{T}^{S}(i)$
(27i)
$\displaystyle{\mathbf{Y}}_{R}^{T}(i)={\mathbf{V}}_{R}^{U}(i)\gamma^{T}_{U}$
(27j)

Whereas the formal scaling of evaluation of the three-body integrals is of
order $N_{\text{orb}}^{6}N_{\text{grid}}$, the most expensive step of the xTC
ansatz presented here, i.e., the assembly of the two-body correction eq. (26),
scales only as $N_{\text{orb}}^{4}N_{\text{grid}}$.

##  III Results

###  III.1 Computational details

The “HEAT” series of papers Tajti  _et al._ (2004); Bomble  _et al._ (2006);
Harding  _et al._ (2008); Thorpe  _et al._ (2019) provide bechmark-quality
energetics for 31 atoms and molecules obtained with computationally expensive,
accurate correlation methods and large basis sets, employing basis-set
extrapolation to produce the final results. We compute the total, atomization,
and formation energies of these systems using the transcorrelation ansatz
excluding explicit three-body components (xTC) with moderate basis-set sizes
without extrapolation and compare them against the corresponding
nonrelativistic electronic energies from HEAT, reported as Hartree-Fock and
correlation energy contributions in Table I of Ref. Tajti  _et al._ , 2004,
which we treat as “exact” values.

Note that carbon is a solid under standard conditions, so “proper” formation
energies of compounds containing carbon atoms cannot be calculated from our
results. For the purpose of comparing with Table III of Ref. Tajti  _et al._ ,
2004 we compute formation energies with respect to carbon monoxide instead,
e.g., the formation energy for C2H2 is obtained from the reaction

$\displaystyle
2\text{CO}+\text{H${}_{2}$}\rightarrow\text{C${}_{2}$H${}_{2}$}+2\text{O}\;.$
(28)

Table 1 lists all reactions considered in the evaluation of formation
energies.

Table 1:  Chemical reactions used for formation energies. CO - O | $\rightarrow$ | C
---|---|---
$\frac{1}{2}$F2 | $\rightarrow$ | F
$\frac{1}{2}$H2 | $\rightarrow$ | H
$\frac{1}{2}$N2 | $\rightarrow$ | N
$\frac{1}{2}$O2 | $\rightarrow$ | O
$2$CO - $2$O + H2 | $\rightarrow$ | C2H2
$2$CO - $2$O + $\frac{1}{2}$H2 | $\rightarrow$ | CCH
CO - O + H2 | $\rightarrow$ | CH2
CO - O + $\frac{1}{2}$H2 | $\rightarrow$ | CH
CO - O + $\frac{3}{2}$H2 | $\rightarrow$ | CH3
CO - O + O2 | $\rightarrow$ | CO2
H2 \+ O2 | $\rightarrow$ | H2O2
H2 \+ $\frac{1}{2}$O2 | $\rightarrow$ | H2O
CO - O + $\frac{1}{2}$H2 \+ $\frac{1}{2}$O2 | $\rightarrow$ | HCO
$\frac{1}{2}$H2 \+ $\frac{1}{2}$F2 | $\rightarrow$ | HF
$\frac{1}{2}$H2 \+ O2 | $\rightarrow$ | HO2
$\frac{1}{2}$N2 \+ $\frac{1}{2}$O2 | $\rightarrow$ | NO
$\frac{1}{2}$H2 \+ $\frac{1}{2}$O2 | $\rightarrow$ | OH
$\frac{1}{2}$H2 \+ $\frac{1}{2}$N2 \+ $\frac{1}{2}$O2 | $\rightarrow$ | HNO
CO - O + $\frac{1}{2}$N2 | $\rightarrow$ | CN
$\frac{1}{2}$H2 \+ CO - O + $\frac{1}{2}$N2 | $\rightarrow$ | HCN
CO - O + $\frac{1}{2}$F2 | $\rightarrow$ | CF
$\frac{1}{2}$N2 \+ H2 | $\rightarrow$ | NH2
$\frac{1}{2}$N2 \+ $\frac{3}{2}$H2 | $\rightarrow$ | NH3
$\frac{1}{2}$N2 \+ $\frac{1}{2}$H2 | $\rightarrow$ | NH
$\frac{1}{2}$O2 \+ $\frac{1}{2}$F2 | $\rightarrow$ | OF

For the Jastrow factors, we employ the Drummond-Towler-Needs form

$\displaystyle u({\mathbf{r}}_{i},{\mathbf{r}}_{j})=$ $\displaystyle
v(r_{ij})+\frac{1}{N_{el}-1}\sum_{I}\big{[}\chi(r_{iI})+\chi(r_{jI})\big{]}$
$\displaystyle+f(r_{ij},r_{iI},r_{jI})$ (29)

with electron-electron, electron-nucleus, and electron-electron-nucleus
contributions $v$, $\chi$, and $f$, respectively, expanded in natural powers.
Drummond, Towler, and Needs (2004); López Ríos  _et al._ (2012) The Jastrow
functions have been optimized by minimizing the variance of the reference
energy within variational Monte Carlo, with the details being described
elsewhere. Haupt  _et al._ (2023)

Numerical quadrature over the direct product of atom centered grids built from
Treutler-Ahlrichs radial grids and Lebedev angular grids obtained from PySCF
has been employed. Sun  _et al._ (2018) We observe that results are
essentially converged for grid level 2 (see supplementary material) and this
is the choice we have used throughout the paper. Final results were obtained
using unrestriced coupled cluster (CC) and distinguishable cluster (DC)
methods implemented in the Molpro quantum chemistry package. Werner  _et al._

In the following, we investigate the quality of the energies obtained with
different coupled cluster based correlation methods as well as different types
and sizes of basis sets. A comparison to the well established explicitly
correlated methods represented by the F12a approximation with the diagonal
ansatz 3C(D) as implemented in MolproKnizia, Adler, and Werner (2009b); Kats
_et al._ (2015b); Kedžuch, Milko, and Noga (2005b); Werner, Adler, and Manby
(2007b) has been included as well. Finally, we investigate the effect of
further approximations to the xTC ansatz itself. Unless otherwise specified,
the density matrix used in the xTC contractions is the HF density as described
in appendix A for open-shell molecules and appendix B for closed-shell
molecules, which corresponds to approximation B in previous work.Schraivogel
_et al._ (2021, 2023) Furthermote, if not otherwise indicated, all-electron
calculations have been carried out.

In order to simplify the comparison, we will only present the signed mean
error (ME), standard deviation (STD) and maximum error (MaxE) for total
energies and the mean absolute error (MAE), root mean square error (RMSE) and
MaxE for relative energies in the body of this work. Individual results and
additional figures can be found in the supplementary material.

###  III.2 Method dependence

Figure 1:  Method dependence of total energies (aug-cc-pVTZ) compared to HEAT. Dotted lines indicate chemical accuracy. The shaded area corresponds to the sum of gaussians centered at the data points with the width of the gaussians chosen such that equidistantly distributed gaussians would be contained to 95% in the corresponding segment.  Figure 2:  Method dependence of atomization energies (aug-cc-pVTZ) compared to HEAT. Dotted lines indicate chemical accuracy. The shaded area corresponds to the sum of gaussians centered at the data points with the width of the gaussians chosen such that equidistantly distributed gaussians would be contained to 95% in the corresponding segment.  Figure 3:  Method dependence of formation energies (aug-cc-pVTZ) compared to HEAT. Dotted lines indicate chemical accuracy. The shaded area corresponds to the sum of gaussians centered at the data points with the width of the gaussians chosen such that equidistantly distributed gaussians would be contained to 95% in the corresponding segment.  Table 2: Method dependence of energies (aug-cc-pVTZ) compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
CCSD-F12a | 28.5 | 16.3 | 62.5 | 23.91 | 28.34 | -53.47 | 15.21 | 18.58 | -33.18
DCSD-F12a | 20.0 | 10.9 | 44.0 | 7.82 | 9.50 | -23.13 | 5.42 | 7.32 | -15.47
CCSD(T)-F12a | 16.4 | 8.7 | 32.9 | 1.47 | 1.90 | -4.01 | 1.60 | 1.89 | -4.37
VMC | 105.3 | 61.9 | 206.1 | 102.50 | 126.81 | -221.50 | 88.67 | 105.30 | -204.19
xTC-HF | 181.2 | 105.3 | 373.5 | 193.37 | 230.16 | -374.34 | 143.38 | 166.77 | -285.83
xTC-CCSD | 13.3 | 8.9 | 30.9 | 11.36 | 14.17 | -27.62 | 10.99 | 12.87 | -24.21
xTC-DCSD | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
xTC-CCSDT | 8.2 | 6.0 | 24.7 | 2.27 | 2.58 | 5.88 | 1.48 | 1.82 | -3.94

The results for transcorrelated and explicitly correlated F12a methods are
shown in Table 2 and Figures 1, 2 and 3. Additionally, the variational Monte-
Carlo (VMC) energies obtained in the Jastrow optimization as well as the mean-
field energies obtained with the xTC integrals (xTC-HF) have been included in
Table 2. Note that we have used CCSDT instead of CCSD(T) for the
transcorrelated calculations since use of the xTC ansatz would require an
iterative approach to the perturbative inclusion of the triples. Furthermore,
there are several possibilities for non-iterative (T) corrections in
transcorrelated methods, which will be investigated in a separate publication.

For total energies, transcorelated methods are on average by a factor of
roughly two closer to the HEAT reference than the respective explicitly
correlated methods. For relative energies however, the F12a methods benefit to
a much larger degree from error compensation than the transcorrelated methods.
Nevertheless, xTC-CCSD and xTC-DCSD on average still perform better than their
explicitly correlated counterparts. While CCSD is unable to provide
satisfactory results both with the transcorralated and the explicitly
correlated ansatz, DCSD performs significantly better. xTC-DCSD in particular
yields average errors close to or even within chemical accuracy (1 kcal/mol =
4.2 kJ/mol = 1.6 mhartree) that are however still about twice as large as
those of CCSD(T)-F12a. Nevertheless, it provides a reasonable compromise
between accuracy and computational cost. xTC-CCSDT further improves both
absolute and relative energies compared to xTC-DCSD, so the success of xTC-
DCSD is unlikely to be due to fortuitous error compensation alone.
CCSD(T)-F12a emerges as the overall best method for the treatment of relative
energies, with xTC-CCSDT yielding slightly worse results for atomization
energies and almost identical results for formation energies. However, unlike
CCSD(T)-F12a where explicit correlation in the triples part enters only
implicitly via the doubles, inclusion of explicit correlation via xTC for
higher order excitations is trivial in xTC-CCSDT and subsequent non-
perturbative methods in the coupled cluster hierarchy and comes with no
additional computational cost beyond generation of the xTC integrals.

###  III.3 Cardinal number and diffuse functions

Table 3: Basis set dependence of xTC-DCSD energies compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
cc-pVDZ | 64.7 | 39.5 | 148.5 | 41.26 | 45.67 | -74.59 | 20.20 | 23.63 | -44.96
aug-cc-pVDZ | 49.0 | 29.7 | 108.1 | 31.18 | 36.51 | -69.50 | 18.69 | 24.15 | -51.58
cc-pVTZ | 13.7 | 9.1 | 35.4 | 5.59 | 6.39 | -11.37 | 3.97 | 4.73 | -9.36
aug-cc-pVTZ | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
cc-pVQZ | 3.6 | 2.8 | 8.6 | 6.00 | 6.90 | -13.14 | 4.13 | 5.31 | -11.76
aug-cc-pVQZ | 2.8 | 2.3 | 7.5 | 5.01 | 6.13 | -13.40 | 3.91 | 5.31 | -12.30

A comparison of the results for xTC-DCSD with the basis sets cc-pVXZ and aug-
cc-pVXZ (X = D, T, Q)T.H. Dunning (1989) can be found in Table 3. The total
energies become progressively better with increasing cardinal number, but on
average do not reach mhartree accuracy. Inclusion of diffuse functions further
improves the results for the total energies, but the average errors remain
larger than for the unaugmented basis set with the next higher cardinal
number.

This clear trend does not translate to the atomization and formation energies.
The double zeta basis sets yield results with average errors of up to roughly
ten times chemical accuracy and are therefore clearly unsuitable. Triple zeta
basis sets on the other hand are already able to achieve chemical accuracy for
many molecules, especially if diffuse functions are included. However, despite
the further improvement in the total energies for quadruple zeta basis sets,
no further improvement in the relative energies is observed. Instead, the
results become slightly worse while still remaining close to chemical
accuracy, indicating that we benefit from favorable error compensation for the
triple zeta basis sets.

###  III.4 Core correlated basis sets

Table 4: Effect of core-correlated basis sets on xTC-DCSD energies compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
cc-pwCVDZ | 55.1 | 36.2 | 135.9 | 40.41 | 45.42 | -81.43 | 20.27 | 23.50 | -42.29
aug-cc-pwCVDZ | 39.7 | 26.1 | 95.9 | 30.64 | 35.71 | -67.56 | 17.83 | 22.71 | -47.84
cc-pVTZ | 13.7 | 9.1 | 35.4 | 5.59 | 6.39 | -11.37 | 3.97 | 4.73 | -9.36
aug-cc-pVTZ | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
cc-pwCVTZ | 11.5 | 7.7 | 29.2 | 8.36 | 9.39 | -15.35 | 4.49 | 5.15 | -10.37
aug-cc-pwCVTZ | 7.7 | 5.4 | 20.3 | 5.28 | 7.24 | -17.02 | 4.40 | 5.98 | -13.06

The results for the use of the core-correlated basis sets cc-pwCVXZ and aug-
cc-pwCVXZ (X = D, T)Peterson and T. H. Dunning (2002) are shown in Table 4.
For double zeta basis sets, we see improvement in both total and relative
energies, but not nearly enough to justify their use. For triple zeta basis
sets, the total energies improve compared to the respective basis sets without
additional core functions. For atomization energies however, use of core-
correlated basis sets leads to larger errors outside of chemical accuracy and
aug-cc-pVTZ results remain closest to the reference for the systems under
investigation. For formation energies, the deterioration of results is less
pronounced and the results remain close to chemical accuracy.

###  III.5 Effect of F12 basis sets

Table 5: Effect of F12 basis sets on xTC-DCSD energies compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
cc-pVDZ-F12 | 23.8 | 15.8 | 57.0 | 18.57 | 22.94 | -50.61 | 11.02 | 14.85 | -31.19
aug-cc-pVDZ-F12 | 23.2 | 15.5 | 55.8 | 17.99 | 22.41 | -50.94 | 11.54 | 15.39 | -31.55
cc-pVTZ | 13.7 | 9.1 | 35.4 | 5.59 | 6.39 | -11.37 | 3.97 | 4.73 | -9.36
aug-cc-pVTZ | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
cc-pVTZ-F12 | 2.6 | 2.4 | 7.5 | 5.13 | 6.42 | -14.36 | 4.00 | 5.56 | -12.75
aug-cc-pVTZ-F12 | 2.5 | 2.4 | 7.6 | 4.71 | 6.00 | -14.03 | 3.83 | 5.37 | -12.48

Results for the basis sets cc-pVXZ-F12 and aug-cc-pVXZ-F12 (X = D, T)Peterson,
Adler, and Werner (2008); Sylvetsky, Kesharwaniand, and Martin (2017)
optimized for use with explicitly correlated methods are shown in Table 5.
While total energies are greatly improved compared to the respective basis
sets discussed previously, the relative energies paint an ambiguous picture.
Atomization and formation energies for double zeta basis sets are greatly
enhanced but still far from chemical accuracy. For triple zeta basis sets
however we observe slight deterioration of the results despite the more
accurate total energies. Therefore, we come to the conclusion that the triple
zeta F12 basis sets offer no advantage here compared to the smaller aug-cc-
pVTZ basis set.

###  III.6 Frozen core approximation

Table 6: Effect of frozen core approximation on energies (aug-cc-pVTZ) compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
FC-xTC-DCSD | 102.0 | 48.0 | 210.3 | 9.65 | 13.02 | -26.18 | 6.27 | 8.19 | -19.72
xTC-FC-DCSD | 11.3 | 6.7 | 25.8 | 3.79 | 5.06 | -12.13 | 3.68 | 4.80 | -12.02
xTC-DCSD | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
FC-DCSD-F12a | 98.7 | 46.9 | 204.9 | 12.79 | 15.16 | -35.50 | 6.09 | 8.13 | -19.27
DCSD-F12a | 20.0 | 10.9 | 44.0 | 7.82 | 9.50 | -23.13 | 5.42 | 7.32 | -15.47
FC-CCSD(T)-F12a | 95.3 | 44.4 | 194.0 | 5.47 | 6.53 | -12.78 | 1.91 | 2.57 | -7.05
CCSD(T)-F12a | 16.4 | 8.7 | 32.9 | 1.47 | 1.90 | -4.01 | 1.60 | 1.89 | -4.37

An obvious approach to reduce the cost of the calculation is the frozen core
approximation. The orbitals can be frozen either before generating the
additional integrals (FC-xTC) or after calculating the xTC corrections (xTC-
FC).

In the first approach, we employ the standard frozen core approximation after
evaluation of the one- and two-electron integrals eqs. (6a) and (6b),
respectively, and evaluate the transcorrelated integrals only for the
remaining orbital space. However, this means that we simply dismiss the
contribution of the core orbitals to the transcorrelation ansatz.

For the xTC-FC approach on the other hand, we first evaluate the xTC
corrections to the integrals and only afterwards apply the frozen core
approximation according to the standard equations, leading to inclusion of
additional correlation arising from the core orbitals. Note that for xTC-FC,
the one-electron integrals are no longer hermitian due to the contraction and
addition of non-hermitian two-electron integrals onto the one-electron
integrals.

While FC-xTC obviously reduces the cost of the generation of the
transcorrelated integrals, the effect on the efficiency of the subsequent
correlation treatment is identical for both approaches. Being able to employ
the FC-xTC approach would enable treatment of larger systems, but that would
require the additional correlation included in xTC-FC via the transcorrelated
integrals to be negligible.

The results for both approaches can be found in Table 6 together with results
for DCSD-F12a and CCSD(T)-F12a with and without frozen core approximation. For
all molecules, the 1s-shell was frozen for each C, N, O and F atom in the
respective molecule. The FC-xTC approximation increases the average error in
the total energies by an order of magnitude and the average errors in the
atomization and formation energies by a factor of about three and two,
respectively. For the F12a methods use of the frozen core approximation also
significantly deteriorates the quality of the total energies. The effect on
the atomization and formation energies however is far less severe, showing
that the F12 approach benefits again from excellent error compensation.

The total energies for FC-xTC and the FC-F12a methods are remarkably close to
each other (see also the additional data provided in the supplementary
material) despite the all electron calculations providing vastly different
results. The xTC-FC approach however introduces comparatively small errors in
both total and relative energies. This indicates that most of the correlation
due to core-core and core-valence interactions is already included in the xTC
ansatz and the subsequent correlation treatment primarily accounts for the
valence-valence interactions.

###  III.7 Density and orbital choice

Table 7: Effect of using spin-averaged (SA) approach for open-shell systems and choice of orbitals and density on xTC-DCSD energies (aug-cc-pVTZ) compared to HEAT. |  $E_{\rm tot}$ (mhartee) |  $\Delta E_{\rm at}$ (kJ/mol) |  $\Delta E_{\rm form}$ (kJ/mol)
---|---|---|---
| ME | STD | MaxE | MAE | RMSE | MaxE | MAE | RMSE | MaxE
HF/HF | 9.4 | 6.6 | 25.7 | 2.69 | 3.35 | -8.49 | 3.47 | 4.39 | -10.16
HF/HF(SA) | 9.6 | 6.5 | 25.7 | 2.67 | 3.09 | 6.14 | 2.52 | 3.02 | 6.14
HF/DCSD | 9.2 | 6.4 | 25.1 | 2.23 | 2.70 | -6.44 | 3.35 | 4.22 | -8.42
HF/DCSD(SA) | 9.4 | 6.3 | 25.1 | 3.28 | 3.99 | 8.69 | 2.38 | 2.90 | 7.36
DCSD/DCSD | 9.4 | 6.5 | 25.4 | 2.73 | 3.53 | -10.49 | 4.37 | 5.65 | -12.13
DCSD/DCSD(SA) | 9.7 | 6.4 | 25.4 | 2.93 | 3.53 | 8.40 | 3.03 | 3.82 | 8.59

The xTC ansatz offers some additional flexibility through the choice of the
density matrix used in the contractions. We are going to investigate three
cases: HF orbitals with HF density (HF/HF), HF orbitals with DCSD density
(HF/DCSD) and DCSD natural orbitals with DCSD density (DCSD/DCSD).
Furthermore, it would be appealing to avoid using the restricted open-shell
algorithm (appendix A) for open-shell molecules, since it leads not only to
increased computational cost compared to the restricted closed-shell algorithm
(appendix B) but also to spin-specific one- and two-body integrals, requiring
additional storage space and the subsequent correlation methods to be able to
handle spin-specific integrals. This can be accomplished by using the density
of an open-shell system as the density in the closed-shell algorithm, which we
will refer to as the spin-averaged (SA) approximation.

The results for the three orbital and density choices with and without SA
approximation are listed in Table 7. The orbital and density choice has only a
marginal impact on the total energies. For atomization and formation energies,
HF/DCSD and DCSD/DCSD yield slightly better or worse results, respectively, on
average. These changes are unlikely to be meaningful and are likely a result
of changes in error cancellation.

Curiously, using the SA approximation does not significantly change the
quality of either total or relative energies. It can even lead to slightly
better results, as can be seen for the HF/HF orbital choice using the SA
approximation, which shows the overall best results for the choices
investigated here. Therefore, considering that other error sources in the
current approach lead to significantly larger deviations, there is no
compelling practical reason not to use the SA approximation.

##  IV Conclusion and Outlook

We have implemented an approximation to the full transcorrelation treatment
which excludes the full three-body term of the normal ordered transcorrelated
Hamiltonian (xTC) and includes the effective zero-, one- and two-body
contributions of the three-body term via modification of the remaining
integrals and the nuclear repulsion term. Explicit calculation of three-body
integrals is avoided by obtaining the correction to the integrals via
contractions of a density matrix with the three-index intermediates for the
numerical integration of the three-body integrals, reducing the scaling by two
orders of magnitude while introducing only minor errors.

The scheme presented here has the added benefit of no longer requiring
explicit contractions with the three-body integrals in the subsequent
correlation treatment, thus removing the need for costly and complicated
modifications in conventional correlation methods in order to allow the use of
the transcorrelated method. In contrast to F12 methods, inclusion of explicit
correlation via xTC in multireference methods or methods using triples and
higher excitations is therefore trivial as long as one accomodates non-
hermitian integrals.

Furthermore, unlike the frozen core approximation for F12 methods, freezing
the core orbitals after generation of the xTC integrals (xTC-FC) allows
partial inclusion of explicit correlation for the core orbitals. This has been
shown to lead to only minor deviations from the all electron calculations
allowing for computational savings in the subsequent correlation treatment
without significant loss in accuracy.

Benchmark calculations have been carried out on the molecules of the HEAT set
and it has been demonstrated that chemical accuracy for atomization and
formation energies is within reach for the comparatively cheap xTC-DCSD method
and the aug-cc-pVTZ basis set. For methods that do not include triple
excitations, results are generally better than those obtained with the
respective explicitly correlated F12 method while methods that include triples
excitations yield similar results for xTC and F12. While improving total
energies, the more expensive basis sets (aug-)cc-pVQZ, (aug-)cc-pwCVTZ and
(aug-)cc-pVTZ-F12 fail to show improvement in the relative energies compared
to results obtained with aug-cc-pVTZ.

The xTC ansatz in combination with the DCSD method enables the application of
transcorrelation to systems of several hundred orbitals at a reasonable
computational cost. Further development should focus on additional efficiency
improvements and the development of more balanced Jastrow factors in order to
further improve the quality of relative energies. Since the DCSD method has
proven to be less unstable for systems with strong static correlation than
conventional coupled cluster methods, the study of strongly correlated systems
may be an interesting application as well.

## Supplementary Material

Raw data and additional figures can be found in the supplementary material.

###### Acknowledgements.

Funded by the Deutsche Forschungsgemeinschaft (DFG, German Research
Foundation) – 455145945. Financial support from the Max-Planck Society is
gratefully acknowledged. P.L.R. and A.A. acknowledge support from the European
Centre of Excellence in Exascale Computing TREX, funded by the Horizon 2020
program of the European Union under grant no. 952165. Any views and opinions
expressed are those of the authors only and do not necessarily reflect those
of the European Union or the European Research Executive Agency. Neither the
European Union nor the granting authority can be held responsible for them.

## Data Availability Statement

The data that support the findings of this study are available within the
article and its supplementary material.

## Author Declarations

The authors have no conflicts to disclose.

## References

## References

  * Kato (1957) T. Kato, Comm. Pure Appl. Math. 10, 151 (1957).
  * Kutzelnigg (1985) W. Kutzelnigg, Theor. Chim. Acta 68, 445 (1985).
  * Kutzelnigg and Klopper (1991) W. Kutzelnigg and W. Klopper, J. Chem. Phys. 94, 1985 (1991).
  * Klopper and Samson (2002) W. Klopper and C. C. M. Samson, J. Chem. Phys. 116, 6397 (2002).
  * Manby (2003) F. R. Manby, J. Chem. Phys. 119, 4607 (2003).
  * Ten-no (2004a) S. Ten-no, Chem. Phys. Lett. 398, 56 (2004a).
  * Ten-no (2004b) S. Ten-no, J. Chem. Phys. 121, 117 (2004b).
  * Valeev (2004) E. F. Valeev, Chem. Phys. Lett. 395, 190 (2004).
  * Tew and Klopper (2005) D. P. Tew and W. Klopper, J. Chem. Phys. 123, 074101 (2005).
  * Kedžuch, Milko, and Noga (2005a) S. Kedžuch, M. Milko, and J. Noga, Int. J. Quantum Chem. 105, 929 (2005a).
  * Fliegl, Klopper, and Hättig (2005) H. Fliegl, W. Klopper, and C. Hättig, J. Chem. Phys. 122, 084107 (2005).
  * Fliegl, Hättig, and Klopper (2006) H. Fliegl, C. Hättig, and W. Klopper, Int. J. Quantum Chem. 106, 2306 (2006).
  * Werner, Adler, and Manby (2007a) H.-J. Werner, T. B. Adler, and F. R. Manby, J. Chem. Phys. 126, 164102 (2007a).
  * Noga, Kedžuch, and Šimunek (2007) J. Noga, S. Kedžuch, and J. Šimunek, J. Chem. Phys. 127, 034106 (2007).
  * Adler, Knizia, and Werner (2007) T. B. Adler, G. Knizia, and H.-J. Werner, J. Chem. Phys. 127, 221106 (2007).
  * Tew  _et al._ (2007) D. P. Tew, W. Klopper, C. Neiss, and C. Hättig, Phys. Chem. Chem. Phys. 9, 1921 (2007).
  * Knizia and Werner (2008) G. Knizia and H.-J. Werner, J. Chem. Phys. 128, 154103 (2008).
  * Shiozaki  _et al._ (2008a) T. Shiozaki, M. Kamiya, S. Hirata, and E. F. Valeev, J. Chem. Phys. 129, 071101 (2008a).
  * Shiozaki  _et al._ (2008b) T. Shiozaki, M. Kamiya, S. Hirata, and E. F. Valeev, Phys. Chem. Chem. Phys. 10, 3358 (2008b).
  * Noga  _et al._ (2008) J. Noga, S. Kedžuch, J. Šimunek, and S. Ten-no, J. Chem. Phys. 128, 174103 (2008).
  * Tew, Klopper, and Hättig (2008) D. P. Tew, W. Klopper, and C. Hättig, Chem. Phys. Lett. 452, 326 (2008).
  * Valeev (2008) E. F. Valeev, Phys. Chem. Chem. Phys. 10, 106 (2008).
  * Valeev and Crawford (2008) E. F. Valeev and T. D. Crawford, J. Chem. Phys. 128, 244113 (2008).
  * Torheyden and Valeev (2008) M. Torheyden and E. F. Valeev, Phys. Chem. Chem. Phys. 10, 3410 (2008).
  * Bokhan, Ten-no, and Noga (2008) D. Bokhan, S. Ten-no, and J. Noga, Phys. Chem. Chem. Phys. 10, 3320 (2008).
  * Knizia, Adler, and Werner (2009a) G. Knizia, T. B. Adler, and H.-J. Werner, J. Chem. Phys. 130, 054104 (2009a).
  * Werner, Knizia, and Manby (2011) H.-J. Werner, G. Knizia, and F. R. Manby, Mol. Phys. 109, 407 (2011).
  * Ten-no (2012) S. Ten-no, Theor. Chim. Acta 131, 1070 (2012).
  * Hättig  _et al._ (2012) C. Hättig, W. Klopper, A. Köhn, and D. P. Tew, Chem. Rev. 112, 4 (2012).
  * Kong, Bischoff, and Valeev (2012) L. Kong, F. A. Bischoff, and E. F. Valeev, Chem. Rev. 112, 75 (2012).
  * Jastrow (1955) R. Jastrow, Phys. Rev. 98, 1479 (1955).
  * Boys and Handy (1969) S. F. Boys and N. C. Handy, Proc. R. Soc. A 310, 63 (1969).
  * Ten-no (2000) S. Ten-no, Chem. Phys. Lett. 330, 169 (2000).
  * Hino, Tanimura, and Ten-no (2002) O. Hino, Y. Tanimura, and S. Ten-no, Chem. Phys. Lett. 353, 317 (2002).
  * Umezawa and Tsuneyuki (2003) N. Umezawa and S. Tsuneyuki, J. Chem. Phys. 119, 10015 (2003).
  * Yanai and Chan (2006) T. Yanai and G. K.-L. Chan, J. Chem. Phys. 124, 194106 (2006).
  * Yanai and Chan (2007) T. Yanai and G. K.-L. Chan, J. Chem. Phys. 127, 104107 (2007).
  * Yanai and Shiozaki (2012) T. Yanai and T. Shiozaki, J. Chem. Phys. 136, 084107 (2012).
  * Tsuneyuki (2008) S. Tsuneyuki, Prog. Theor. Phys. Suppl. 176, 134 (2008).
  * Ochi  _et al._ (2012) M. Ochi, K. Sodeyama, R. Sakuma, and S. Tsuneyuki, J. Chem. Phys. 136, 094108 (2012).
  * Ochi and Tsuneyuki (2014) M. Ochi and S. Tsuneyuki, J. Chem. Theory Comput. 10, 4098 (2014).
  * Ochi and Tsuneyuki (2015) M. Ochi and S. Tsuneyuki, Chem. Phys. Lett. 621, 177 (2015).
  * Ochi  _et al._ (2016) M. Ochi, Y. Yamamoto, R. Arita, and S. Tsuneyuki, J. Chem. Phys. 144, 104109 (2016).
  * Wahlen-Strothman  _et al._ (2015) J. M. Wahlen-Strothman, C. A. Jiménez-Hoyos, T. M. Henderson, and G. E. Scuseria, Phys. Rev. B 91, 041114 (2015).
  * Luo and Alavi (2018) H. Luo and A. Alavi, J. Chem. Theory Comput. 14, 1403 (2018).
  * Dobrautz, Luo, and Alavi (2019) W. Dobrautz, H. Luo, and A. Alavi, Phys. Rev. B 99, 075119 (2019).
  * Cohen  _et al._ (2019) A. J. Cohen, H. Luo, K. Guther, W. Dobrautz, D. P. Tew, and A. Alavi, J. Chem. Phys. 151, 061101 (2019).
  * Baiardi and Reiher (2020) A. Baiardi and M. Reiher, J. Chem. Phys. 153, 164115 (2020).
  * Khamoshi  _et al._ (2021) A. Khamoshi, G. P. Chen, T. M. Henderson, and G. E. Scuseria, J. Chem. Phys. 154, 074113 (2021).
  * Giner (2021) E. Giner, J. Chem. Phys. 154, 084119 (2021).
  * Guther  _et al._ (2021) K. Guther, A. J. Cohen, H. Luo, and A. Alavi, The Journal of Chemical Physics 155, 011102 (2021), https://doi.org/10.1063/5.0055575 .
  * Liao  _et al._ (2021) K. Liao, T. Schraivogel, H. Luo, D. Kats, and A. Alavi, Phys. Rev. Research 3, 033072 (2021).
  * Schraivogel  _et al._ (2021) T. Schraivogel, A. J. Cohen, A. Alavi, and D. Kats, J. Chem. Phys. 155, 191101 (2021).
  * Liao  _et al._ (2023) K. Liao, H. Zhai, E. M. Christlmaier, T. Schraivogel, P. L. Ríos, D. Kats, and A. Alavi, J. Chem. Theory Comput. 19, 1734 (2023).
  * Haupt  _et al._ (2023) J. P. Haupt, S. M. Hosseini, P. L. Ríos, W. Dobrautz, A. Cohen, and A. Alavi,  (2023), arXiv:2302.13683 [physics].
  * Schraivogel  _et al._ (2023) T. Schraivogel, E. M. Christlmaier, P. L. Rıos, A. Alavi, and D. Kats, J. Chem. Phys. 158, 214106 (2023).
  * Čížek (1966) J. Čížek, J. Chem. Phys. 45, 4256 (1966).
  * Purvis and Bartlett (1982) G. D. Purvis and R. J. Bartlett, J. Chem. Phys. 76, 1910 (1982).
  * Raghavachari  _et al._ (1989) K. Raghavachari, G. W. Trucks, J. A. Pople, and M. Head-Gordon, Chem. Phys. Lett. 157, 479 (1989).
  * Kats and Manby (2013) D. Kats and F. R. Manby, J. Chem. Phys. 139, 021102 (2013).
  * Kats (2014) D. Kats, J. Chem. Phys. 141, 061101 (2014).
  * Kats  _et al._ (2015a) D. Kats, D. Kreplin, H.-J. Werner, and F. R. Manby, J. Chem. Phys. 142, 064111 (2015a).
  * Kats (2016) D. Kats, J. Chem. Phys. 144, 044102 (2016).
  * Kats (2018) D. Kats, Mol. Phys. 116, 1435 (2018).
  * Kutzelnigg and Mukherjee (1997) W. Kutzelnigg and D. Mukherjee, J. Chem. Phys. 107, 432 (1997).
  * Kutzelnigg and Mukherjee (1999) W. Kutzelnigg and D. Mukherjee, J. Chem. Phys. 110, 2800 (1999).
  * Tajti  _et al._ (2004) A. Tajti, P. G. Szalay, A. G. Császár, M. Kállay, J. Gauss, E. F. Valeev, B. A. Flowers, J. Vázquez, and J. F. Stanton, J. Chem. Phys. 121, 11599 (2004).
  * Bomble  _et al._ (2006) Y. J. Bomble, J. Vázquez, M. Kállay, C. Michauk, P. G. Szalay, A. G. Császár, J. Gauss, and J. F. Stanton, J. Chem. Phys. 125, 064108 (2006).
  * Harding  _et al._ (2008) M. E. Harding, J. Vázquez, B. Ruscic, A. K. Wilson, J. Gauss, and J. F. Stanton, J. Chem. Phys. 128, 114111 (2008).
  * Thorpe  _et al._ (2019) J. H. Thorpe, C. A. Lopez, T. L. Nguyen, J. H. Baraban, D. H. Bross, B. Ruscic, and J. F. Stanton, J. Chem. Phys. 150, 224102 (2019).
  * Drummond, Towler, and Needs (2004) N. D. Drummond, M. D. Towler, and R. J. Needs, Phys. Rev. B 70, 235119 (2004).
  * López Ríos  _et al._ (2012) P. López Ríos, P. Seth, N. D. Drummond, and R. J. Needs, Phys. Rev. E 86, 036703 (2012).
  * Sun  _et al._ (2018) Q. Sun, T. C. Berkelbach, N. S. Blunt, G. H. Booth, S. Guo, Z. Li, J. Liu, J. D. McClain, E. R. Sayfutyarova, S. Sharma, S. Wouters, and G. K.-L. Chan, WIREs Comput. Mol. Sci. 8, e1340 (2018).
  * (74) H.-J. Werner, P. J. Knowles, G. Knizia, F. R. Manby, M. Schütz, P. Celani, W. Györffy, D. Kats, T. Korona, R. Lindh, A. Mitrushenkov, G. Rauhut, K. R. Shamasundar, T. B. Adler, R. D. Amos, S. J. Bennie, A. Bernhardsson, A. Berning, D. L. Cooper, M. J. O. Deegan, A. J. Dobbyn, F. Eckert, E. Goll, C. Hampel, A. Hesselmann, G. Hetzer, T. Hrenar, G. Jansen, C. Köppl, S. J. R. Lee, Y. Liu, A. W. Lloyd, Q. Ma, R. A. Mata, A. J. May, S. J. McNicholas, W. Meyer, T. F. Miller III, M. E. Mura, A. Nicklass, D. P. O’Neill, P. Palmieri, D. Peng, K. Pflüger, R. Pitzer, M. Reiher, T. Shiozaki, H. Stoll, A. J. Stone, R. Tarroni, T. Thorsteinsson, M. Wang, and M. Welborn, “Molpro, version 2022.3.0, a package of ab initio programs,” See https://www.molpro.net.
  * Knizia, Adler, and Werner (2009b) G. Knizia, T. B. Adler, and H.-J. Werner, J. Chem. Phys. 130, 054104 (2009b).
  * Kats  _et al._ (2015b) D. Kats, D. Kreplin, H.-J. Werner, and F. R. Manby, J. Chem. Phys. 142, 064111 (2015b).
  * Kedžuch, Milko, and Noga (2005b) S. Kedžuch, M. Milko, and J. Noga, Int. J. Quantum Chem. 105, 929 (2005b).
  * Werner, Adler, and Manby (2007b) H.-J. Werner, T. B. Adler, and F. R. Manby, J. Chem. Phys. 126, 164102 (2007b).
  * T.H. Dunning (1989) J. T.H. Dunning, J. Chem. Phys. 90, 1007 (1989).
  * Peterson and T. H. Dunning (2002) K. A. Peterson and J. T. H. Dunning, J. Chem. Phys. 117, 10548 (2002).
  * Peterson, Adler, and Werner (2008) K. A. Peterson, T. B. Adler, and H.-J. Werner, J. Chem. Phys. 128, 084102 (2008).
  * Sylvetsky, Kesharwaniand, and Martin (2017) N. Sylvetsky, M. K. Kesharwaniand, and J. M. L. Martin, J. Chem. Phys. 147, 134106 (2017).

##  Appendix A Restricted open-shell working equations

Integrating equations (II.4), (23), (26) and (27) over spin and assuming use
of the same spatial orbitals for $\alpha$\- and $\beta$-spin, we can simplify
the working equations. In our current implementation we also assume

$\displaystyle[\rho_{p}^{q}]^{\sigma}=\rho_{p}^{q}$
$\displaystyle\forall\sigma\in\alpha,\beta$ (30a)
$\displaystyle[{\mathbf{V}}_{p}^{q}]^{\sigma\tau}={\mathbf{V}}_{p}^{q}$
$\displaystyle\forall\sigma,\tau\in\alpha,\beta,$ (30b)

i.e. in addition to using the same orbitals for $\alpha$\- and $\beta$-spin we
also use the same Jastrow factors. While this removes much of the spin-
dependency, we still obtain spin-specific sets of one- and two-electron
integrals due to the use of spin-specific density matrices.

$\displaystyle[\Delta h_{p}^{q}]^{\sigma}=-\frac{1}{2}\bigg{(}$
$\displaystyle\sum_{\tau}[\Delta
U_{pr}^{qs}]^{\sigma\tau}[\gamma^{r}_{s}]^{\tau}$ $\displaystyle-[\Delta
U_{pr}^{sq}]^{\sigma\sigma}[\gamma^{r}_{s}]^{\sigma}\bigg{)}$ (31)

$\displaystyle\langle{\Psi_{0}}|L|{\Psi_{0}}\rangle=-\frac{1}{3}\sum_{\sigma}[\Delta
h_{p}^{q}]^{\sigma}[\gamma^{p}_{q}]^{\sigma}$ (32)

$\displaystyle[\Delta
U_{pr}^{qs}]^{\sigma\tau}=-\mathcal{P}^{(pq\sigma)}_{(rs\tau)}\big{(}$
$\displaystyle\rho_{p}^{q}(i)[A_{r}^{s}(i)]^{\tau}$
$\displaystyle+{\mathbf{V}}_{p}^{q}(i)\cdot{\mathbf{[}}B_{r}^{s}(i)]^{\tau}\big{)}$
(33) $\displaystyle\rho_{p}^{q}(i)=w_{i}\phi_{p}^{*}(i)\phi^{q}(i)$ (34a)
$\displaystyle[A_{r}^{s}(i)]^{\tau}=\tilde{V}_{r}^{s}(i)-[\tilde{Z}_{r}^{s}(i)]^{\tau}$
(34b)
$\displaystyle\tilde{V}_{r}^{s}(i)={\mathbf{W}}(i)\cdot{\mathbf{V}}_{r}^{s}(i)$
(34c)
$\displaystyle{\mathbf{W}}(i)={\mathbf{V}}_{t}^{u}(i)\sum_{\rho}[\gamma_{u}^{t}]^{\rho}$
(34d)
$\displaystyle[\tilde{Z}_{r}^{s}(i)]^{\tau}={\mathbf{V}}_{r}^{u}(i)\cdot[{\mathbf{X}}_{u}^{s}(i)]^{\tau}$
(34e)
$\displaystyle[{\mathbf{X}}_{u}^{s}(i)]^{\tau}={\mathbf{V}}_{t}^{s}(i)[\gamma^{t}_{u}]^{\tau}$
(34f)
$\displaystyle[{\mathbf{B}}_{r}^{s}(i)]^{\tau}=\frac{1}{2}\tilde{W}(i){\mathbf{V}}_{r}^{s}(i)-[{\mathbf{Z}}_{r}^{s}(i)]^{\tau}$
(34g)
$\displaystyle\tilde{W}(i)=\rho_{t}^{u}(i)\sum_{\rho}[\gamma^{t}_{u}]^{\rho}$
(34h)
$\displaystyle[{\mathbf{Z}}_{r}^{s}(i)]^{\tau}=\rho_{r}^{u}(i)[{\mathbf{X}}_{u}^{s}(i)]^{\tau}+[{\mathbf{Y}}_{r}^{t}(i)]^{\tau}\rho_{t}^{s}(i)$
(34i)
$\displaystyle[{\mathbf{Y}}_{r}^{t}]^{\tau}(i)={\mathbf{V}}_{r}^{u}(i)[\gamma^{t}_{u}]^{\tau}$
(34j)

##  Appendix B Restricted closed-shell working equations

Further assuming a closed-shell system, we can drop the remaining spin-
dependencies. The following equations have been used for the closed-shell
systems investigated in this study. They can double as a spin-averaged
approximation to the restricted open-shell case if the density matrices are
taken as the average of $\alpha$\- and $\beta$-density matrices. This can be
desirable, since unlike the former equations, there are no spin-specific sets
of integrals.

Further note that the equations given here deviate from the standard
convention of including the factor of two from the spin-integration into the
density matrix. Instead this factor is explicitly included in the equations
presented here.

$\displaystyle\Delta h_{p}^{q}=-\frac{1}{2}\bigg{(}2\Delta
U_{pr}^{qs}\gamma^{r}_{s}-\Delta U_{pr}^{sq}\gamma^{r}_{s}\bigg{)}$ (35)

$\displaystyle\langle{\Psi_{0}}|L|{\Psi_{0}}\rangle=-\frac{2}{3}\Delta
h_{p}^{q}\gamma^{p}_{q}$ (36)

$\displaystyle\Delta U_{pr}^{qs}=-\mathcal{P}^{(pq)}_{(rs)}\big{(}$
$\displaystyle\rho_{p}^{q}(i)A_{r}^{s}(i)+{\mathbf{V}}_{p}^{q}(i)\cdot{\mathbf{B}}_{r}^{s}(i)\big{)}$
(37)

$\displaystyle\rho_{p}^{q}(i)=w_{i}\phi_{p}^{*}(i)\phi^{q}(i)$ (38a)
$\displaystyle A_{r}^{s}(i)=\tilde{V}_{r}^{s}(i)-\tilde{Z}_{r}^{s}(i)$ (38b)
$\displaystyle\tilde{V}_{r}^{s}(i)={\mathbf{W}}(i)\cdot{\mathbf{V}}_{r}^{s}(i)$
(38c) $\displaystyle{\mathbf{W}}(i)=2{\mathbf{V}}_{t}^{u}(i)\gamma^{t}_{u}$
(38d)
$\displaystyle\tilde{Z}_{r}^{s}(i)={\mathbf{V}}_{r}^{u}(i)\cdot{\mathbf{X}}_{u}^{s}(i)$
(38e)
$\displaystyle{\mathbf{X}}_{u}^{s}(i)={\mathbf{V}}_{t}^{s}(i)\gamma^{t}_{u}$
(38f)
$\displaystyle{\mathbf{B}}_{r}^{s}(i)=\frac{1}{2}\tilde{W}(i){\mathbf{V}}_{r}^{s}(i)-{\mathbf{Z}}_{r}^{s}(i)$
(38g) $\displaystyle\tilde{W}(i)=2\rho_{t}^{u}(i)\gamma^{t}_{u}$ (38h)
$\displaystyle{\mathbf{Z}}_{r}^{s}(i)=\rho_{r}^{u}(i){\mathbf{X}}_{u}^{s}(i)+{\mathbf{Y}}_{r}^{t}(i)\rho_{t}^{s}(i)$
(38i)
$\displaystyle{\mathbf{Y}}_{r}^{t}(i)={\mathbf{V}}_{r}^{u}(i)\gamma^{t}_{u}$
(38j)


# ESGCN: Edge Squeeze Attention Graph Convolutional Network
for Traffic Flow Forecasting

Sangrok Lee  Ha Young Kim Yonsei University
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Traffic forecasting is a highly challenging task owing to the dynamical
spatio-temporal dependencies of traffic flows. To handle this, we focus on
modeling the spatio-temporal dynamics and propose a network termed Edge
Squeeze Graph Convolutional Network (ESGCN) to forecast traffic flow in
multiple regions. ESGCN consists of two modules: W-module and ES module.
W-module is a fully node-wise convolutional network. It encodes the time-
series of each traffic region separately and decomposes the time-series at
various scales to capture fine and coarse features. The ES module models the
spatio-temporal dynamics using Graph Convolutional Network (GCN) and generates
an Adaptive Adjacency Matrix (AAM) with temporal features. To improve the
accuracy of AAM, we introduce three key concepts. 1) Using edge features to
directly capture the spatio-temporal flow representation among regions. 2)
Applying an edge attention mechanism to GCN to extract the AAM from the edge
features. Here, the attention mechanism can effectively determine important
spatio-temporal adjacency relations. 3) Proposing a novel node contrastive
loss to suppress obstructed connections and emphasize related connections.
Experimental results show that ESGCN achieves state-of-the-art performance by
a large margin on four real-world datasets (PEMS03, 04, 07, and 08) with a low
computational cost.

##  1 Introduction

Traffic flow forecasting is a core component of intelligent transportation
systems. It is essential for analyzing traffic situations and aims at
predicting the future traffic flow of regions using historical traffic data.
However, this task is challenging because of the heterogeneity and dynamic
spatio-temporal dependence of traffic data.

Traffic data can be modeled using Graph Neural Networks (GNNs). In such
networks, the regions are represented as nodes and flows between regions as
edges. Graph Convolutional Network (GCN) is a type of GNN commonly used to
handle traffic flow forecasting tasks. It can adequately leverage the graph
structure and aggregate node information Wu et al. (2019, 2020); Kong et al.
(2020); Song et al. (2020); Li and Zhu (2021). Because the edges define the
intensity of the adjacency matrix in a graph operation, accurate edge graph is
an important factor in determining the performance of a GCN. Recent studies
focus on capturing the connection patterns via an adaptive adjacency matrix
(AAM) that is found in the training process BAI et al. (2020a); Wu et al.
(2019). In our study, we focus on enhancing the AAM with the following three
distinctive features.

First, we propose an Edge Squeeze (ES) module that directly uses spatio-
temporal flows with edge features to construct an AAM. Recent studies in GCN
revealed that edge features are equally important as node features Gong and
Cheng (2019); Chen and Chen (2021). While the edge features can be used to
simulate the traffic flows between regions, to the best of our knowledge this
is the first study to build an adjacency matrix from the edge features in this
task. Existing methods use node embeddings which cannot accurately reflect the
relationship among nodes because the embedding vectors are fixed for inference
and only represent spatial nodes. Therefore, they cannot handle dynamic
patterns occurred in the inference and are unable to accurately capture
spatio-temporal features. However, the ES module leverages the temporal
features directly to construct the AAM and reflects the changes of temporal
features in the inference. ES module creates three-dimensional (3D) spatio-
temporal correlations beyond the spatial-specific embeddings.

Second, we develop a novel edge attention mechanism. We further explore the
edge features with the attention mechanism to refine the AAM. Because the edge
features represent the adjacency relations, we apply the attention mechanism
to activate meaningful edges and suppress the others of the adjacency matrix.
The edge attention exploits feature’s channel information such as SENet Hu et
al. (2020), referred to as the squeeze attention. Existing methods apply
transformer attention mechanism Vaswani et al. (2017) which presents high
computational burden. The channel attention mechanism provides relatively
lower computation cost and can generate more refined AAM.

Third, we introduce a novel node contrastive loss. Previous studies computed
the similarities between the node embeddings that were trained with a
forecasting objective function for the AAM. This method generated the
adjacency matrix without an explicit objective for the shape of the graph,
consequently inducing suboptimal performance. To overcome this, we maximize
the difference between related and unrelated nodes to facilitate separation of
the forecasting relevant nodes from the residuals. This aids in preventing the
propagation of information on unrelated nodes through AAM.

Additionally, we propose a backbone network, W-module, to extract multi-scale
temporal features for traffic flow forecasting. W-module is a fully
convolutional network that consists of node-wise convolution to handle time-
series features of each node separately. Owing to its non-autoregressive
attributes and receptive field of convolution layers, the W-module can extract
multiple levels of temporal features from shallow to deep layers and provide a
hierarchical decomposition. We combine the ES module with W-module and propose
ESGCN to forecast traffic flow. The main contributions of this study are
summarized as follows:

  * •

We propose an end-to-end framework (ESGCN) using two novel modules: ES module
and W-module. ESGCN effectively learns hidden and dynamic spatio-temporal
relationships using edge features.

  * •

We introduce an edge attention mechanism and node contrastive loss to
construct an AAM which captures accurately the relationships among nodes.

  * •

We perform extensive experiments on four real-world datasets (PEMS03, 04, 07,
and 08); the results shows that ESGCN achieves state-of-the-art performance by
a large margin with a low computational cost.

##  2 Related Work

GCN for traffic forecasting. GCN is a special type of convolutional neural
network that is widely used in traffic forecasting tasks Hechtlinger et al.
(2017); Kipf and Welling (2017). Recurrent-based GCN adopts a recurrent
architecture, such as LSTM and replaces inner layers with GCN Li et al.
(2018); Bai et al. (2020b). This type of GCN handles spatial and temporal
features recurrently, however it has long-range memory loss problem. STGCN Yu
et al. (2017), GSTNet Fang et al. (2019), and Graph WaveNet Wu et al. (2019)
use fully convolutional architecture and graph operation. They exploit spatial
and temporal convolution separately to model spatio-temporal data. These
methods show relatively fast inference speed and improved performance on long-
range temporal data.

Attention mechanism for traffic forecasting. Attention mechanisms are used to
effectively capture spatio-temporal dynamics for traffic forecasting. ASTGCN
Guo et al. (2019), GMAN Zheng et al. (2020), and STGRAT Park et al. (2020)
exploited the attention mechanism which considers changes in road speed and
diverse influence of spatial and temporal network. Existing approaches
leveraged transformer attention Vaswani et al. (2017) that computes key, query
and value relations. Contrastingly, our proposed network adopts a channel
attention such as SENet Hu et al. (2020) and CBAM Woo et al. (2018) which have
relatively low computation cost and high speed.

Adaptive Adjacency Matrix Construction. Previous studies used a predefined
adjacency matrix for GCN Li et al. (2018); Yu et al. (2017); Yao et al.
(2019); Guo et al. (2019). In a previous study, a spatial AAM was proposed as
a supplementary for the predefined adjacency matrix for graph WaveNet Wu et
al. (2019). STGAT Kong et al. (2020) also utilized the AAM, however, they were
limited by spatial dependencies. AGCRN Bai et al. (2020b) exploited an AAM
solely based on spatial node embedding. In this study, we introduce a spatio-
temporal based AAM that captures accurate relations.

Figure 1: Overall architecture of ESGCN. ESGCN comprises the ES module and
W-module consisting of four stages. Each stage of the W-module consists of
multiple W-blocks, namely gated node-wise convolution layers with a layer
normalization. The ES module computes the spatio-temporal correlation,
generates an adaptive adjacency matrix with attention mechanism, and finally
conducts GCN.

##  3 Problem Definition

To predict future traffic, univariate time-series data from each region
${\\{X_{1},X_{2},X_{3},\dots\\}}$ are provided where $X_{i}\in\mathbb{R}^{n}$
is the traffic record at time step $i$ and $n$ is the number of regions. The
purpose is to find a function $\mathcal{E}$ that is capable of predicting the
future of length $o$ by analyzing existing $T$-length past data. Traffic flows
are generated for multiple regions and the problem is formulated as follows:

$\\{X_{t+1},\dots,X_{t+o}\\}=\mathcal{E}(X_{t-T+1},\dots,X_{t}),$ (1)

where $t$ is an arbitrary time step.

##  4 Proposed Method

This section describes the edge squeeze graph convolution network (ESGCN). The
network comprises two modules: W-module and edge squeeze (ES) module, as shown
in Fig. 1. W-module extracts region-specific temporal features and ES module
produces an AAM using temporal features with an edge attention and 3D spatio-
temporal relations.

###  4.1 W-Module

W-module extracts the time-series features of each node and is split into four
W-block groups as illustrated in Fig. 1. We call each group as stage. To
organize the W-module, we use a combination of gated node-wise convolutions
(GNC) and a layer normalization Ba et al. (2016), namely W-block as the
smallest unit. GNC has two node-wise convolution layers: one for embedding
features with a tanh function and the other for gating with a sigmoid
function, as shown in Fig. 1. The two outputs are multiplied element-wise and
summarized as follows:

$GNC(X)=sigmoid(\mathcal{C}(X))\circ tanh(\mathcal{C}(X)),$ (2)

where $\mathcal{C}$ is a 2D convolution layer with a 1 $\times$ 3 kernel and
0$\times$1 padding, and $\circ$ is element-wise multiplication. $X$ is an
input feature of $X\in\mathbb{R}^{c\times n\times t}$, where $c$ is the number
of channels, $n$ is the number of nodes, and $t$ is an arbitrary temporal
dimension. This 1$\times$3 convolution layer is referred to as node-wise
convolution. Given that the convolution kernel has a size of one in the node
dimension and expands the receptive field in the time dimension, the W-block
is only responsible for extracting temporal features. The output features of
the GNC are fed to the layer normalization layer. Each stage has 1, 2, 2, and
2 W-blocks with a stride of 1 for first stage and a stride of 2 for the
others. The output of $i^{th}$ stage is defined as
$F_{i}\in\mathbb{R}^{~{}c_{i}\times n\times\frac{l}{2^{i-1}}}$, where $c_{i}$
is specified as a hyperparameter and $l$ is input time-series length.

The receptive field of convolution layer expands as the feature passes through
the various stages. Therefore, the features of the early stages capture local
signals and subtle changes and the features of late stages are learned for
global signals and overall movements. To take this advantages in handling
multiple levels of signal size, the outputs of the intermediate stages are
also used with the last output features for the final prediction. In W-module,
each stage has 3, 9, 12, and 12 receptive field sizes in a sequential order.
This can decompose the time-series hierarchically and each stage is
efficiently trained to be responsible for possible signals.

###  4.2 Edge Squeeze Module

The ES module is a spatio-temporal feature extractor, which enhances temporal
features with relations from the W-module. As described in Fig. 1, it reflects
the node relations through three steps: spatio-temporal correlation
computation, relational feature extraction, and GCN operation.

Spatio-temporal correlation computation. In this step, we construct spatio-
temporal correlations to model flows among nodes. The modeling involves three
steps. First, we feed the last features of W-module, $F_{4}$, to a single
convolution layer to reduce the channel size and computation cost. This is
defined as:

$F_{c}=\mathcal{C}(F_{4}),F_{c}\in\mathbb{R}^{~{}\frac{c_{4}}{4}\times n\times
l_{4}},$ (3)

where $\mathcal{C}$ is $1\times 1$ convolution of output channels and $l_{4}$
is $\frac{l}{8}$.

Second, we extract a temporal representative from $l_{4}$ temporal nodes in
each region to compute spatio-temporal correlations between the
representatives ($n$) and all nodes ($n\times l_{4}$). The node is set in the
last time step as a representative node. This is inspired by the Markov
decision process (MDP) Bellman and Howard (1961), which considers only a state
of current time step to decide an action at the subsequent time, and the
representative of node is the closest value to the prediction value of the
last element on the time dimension. The representative function $\mathcal{T}$
is as follows:

$F_{l}=\mathcal{T}(F_{c}),~{}F_{l}\in\mathbb{R}^{~{}c_{4}\times n},$ (4)

where $F_{l}$ denotes the last node on the time dimension.

Finally, the relations between the features and the representatives are
computed. The original feature, $F_{c}$ has $n\times l_{4}$ nodes, and
representative, $F_{l}$ has $n$ nodes. The spatio-temporal correlations
between $F_{c}$ and $F_{l}$ are in $n\times n\times l_{4}$ space. Inspired by
previous studies Wu et al. (2019); Li and Zhu (2021); Kong et al. (2020), a
similarity function is adopted to measure the correlations between spatio-
temporal nodes. The cosine similarity is adopted for this study. The result of
cosine similarity is bounded from -1 to 1, and it is suitable for comparing
high-dimensional features Luo et al. (2018). The cosine similarity $cos$ is
defined as follows:

$cos(x_{1},x_{2})=\frac{x_{1}\cdot x_{2}}{\|x_{1}\|_{2}\cdot\|x_{2}\|_{2}},$
(5)

where $\cdot$ is dot product and $\|\cdot\|_{2}$ is Euclidean norm. The final
spatio-temporal correlations $S$ is defined as:

$S=cos(F_{l},F_{c}),~{}S\in\mathbb{R}^{n\times n\times l_{4}}$ (6)

Relational features extraction. Based on the previous step, each node has
$n\times l_{4}$ correlations. We employ these correlations to weight $F_{4}$
and generate correlation-aware features, $\in\mathbb{R}^{c_{4}\times n\times
l_{4}}$, for each node. Subsequently, we aggregate the temporal channel
features, $l_{4}$, to reflect temporal contexts. To facilitate computation, we
expand the dimension of $S$ to $S^{e}\in\mathbb{R}^{n\times c_{4}\times
n\times l_{4}}$ (1 to $c_{4}$) and define $S_{i}^{e}\in\mathbb{R}^{c_{4}\times
n\times l_{4}}$ as the correlation of the $i^{th}$ node. In following
equations, $z(:,:,*)$ denotes $*^{th}$ tensor of $z$ on the last dimension.
The computation process is defined as:

$R(:,:,k)=\sum_{j\in l_{4}}S_{k}^{e}(:,:,j)\circ F_{4}(:,:,j),$ (7)

where $R(:,:,k)\in\mathbb{R}^{c_{4}\times n}$ is the $k^{th}$ node’s
relational features between all nodes, $k=1,2,\dots,n$.

We use the relational features as a feature matrix for graph convolution. This
has an advantage over other features as it reflects temporal relations. GCN
conducts operation only considering of spatial node relations. However, in
spatio-temporal data, temporal relations need to be included. In the
relational features, each node has $n$ different nodes created by considering
temporal contexts. Thus, we use $R(:,:,k)$ tensor as $n$ neighbor nodes with
$c_{4}$ dimension to predict $k^{th}$ node future flow. It can reflect the
temporal relations in the relational features and spatial relations in GCN
operation.

GCN operation. GCN requires a feature matrix and an adjacency matrix. As
previously mentioned, the relational features are used as the feature matrix.
To construct the adjacency matrix, we apply an attention mechanism to the
relational features. Inspired by SENet Hu et al. (2020) and CBAM Woo et al.
(2018), we adopt a channel attention mechanism known as squeeze attention to
refine the AAM from the relational features. The squeeze attention activates
important spatial and channel positions. Therefore, we use this mechanism to
extract the edge positions’ importance. To squeeze the edge features, we feed
the relational features to max operation, tanh, and ReLU activation as
following:

$A=ReLU(tanh(max(R)))$ (8)

The outcome, $A\in n\times n$, is the generated AAM. The graph operation with
$A$ and $R$ is defined as:

$F_{g}(:,k)=W\otimes R(:,:,k)\otimes A(k,:)+B,$ (9)

where $F_{g}(:,k)\in\mathbb{R}^{c_{4}}$ represents features for $k^{th}$ node,
$W$ is a learnable weight denoted as $W\in\mathbb{R}^{c_{4}\times c_{4}}$, $B$
is a bias as $B\in R^{c_{4}}$, and $\otimes$ is matrix multiplication.

###  4.3 ESGCN

The proposed framework consists of the W-module and ES module. The outputs of
the first three stages of the W-module and ES module are employed to access
all levels of features. The computation process is summarized as follows:

$P=\sum_{i=1}^{t-1}\mathcal{C}^{~{}i}(F_{i})+\mathcal{C}^{e}(F_{g}),$ (10)

where $t$ is the number of stages, $\mathcal{C}^{~{}i}$ and $\mathcal{C}^{e}$
are a $1\times 1$ convolution layer for the $i$-th stage and ES module
respectively.

To predict future flow, two fully connected layers with ReLU activation
function are used. $P$ is fed to the two fully connected layers.

$\hat{Y}=W_{b}\otimes ReLU(W_{a}\otimes
P+B_{a})+B_{b},~{}\hat{Y}\in\mathbb{R}^{~{}h\times n},$ (11)

where $W_{a}$ and $W_{b}$ are weights and $B_{a}$ and $B_{b}$ are bias of the
fully connected layers and $h$ is the number of time steps in the prediction.

Figure 2: Proposed the node contrastive loss implementation. Note the branch
with red dashed line is solely used for loss computation and not for
inference.

###  4.4 Loss Function

We adopt the Huber loss and the proposed node contrastive loss. Huber loss is
defined as follows:

$h\big{(}\hat{Y},Y\big{)}=\left\\{\begin{aligned}
&\frac{1}{2}(\hat{Y}-Y)^{2},&|\hat{Y}-Y|<\delta\\\
&\delta|\hat{Y}-Y|-\frac{1}{2}\delta^{2},&|\hat{Y}-Y|\geq\delta\end{aligned}\right.,$
(12)

where $\delta$ is set to 1, $\hat{Y}$ denotes predicted values and $Y$ is
ground truth values. The objective function, $L_{h}$ is defined as follows:

$L_{h}=\frac{1}{b\times
o}\sum\limits_{k=1}^{b}\sum\limits_{i=1}^{o}h(\hat{Y}_{i}^{k},Y_{i}^{k}),$
(13)

where $\hat{Y}_{i}^{k}$ and $Y_{i}^{k}$ are the predicted value and ground
truth of $i$-th time step of $k$-th sample in a mini batch respectively, and
$b$ is the number of samples in a mini batch.

Node contrastive loss. This is used to effectively separate related and
unrelated nodes. As shown in Fig. 2, the reversed adjacency matrix, $A_{r}$,
is generated. We modify the squeeze function as follows:

$A_{r}=ReLU(-tanh(max(R)))$ (14)

Subsequently, the unrelated features are extracted using the graph operation
in the reversed adjacency matrix. Finally, we maximize the distance between
the related features, $F_{g}$, with the original adjacency matrix and the
unrelated features, $F_{gr}$, with the reversed adjacency matrix. The node
contrastive loss is defined as:

$L_{n}=\frac{1}{n}tr(F_{g}^{T}\otimes F_{gr}),$ (15)

where $tr$ is the trace of a matrix and $F_{g}^{T}\in\mathbb{R}^{~{}n\times
c_{4}}$ denotes a transposed matrix. Minimizing an orthogonal of the
multiplied matrix can be viewed as reducing the similarity between
corresponding nodes based on a dot product.

The final loss function is calculated as follows:

$L=L_{h}+\lambda L_{n},$ (16)

where $\lambda$ is set to 0.1 in our experiments.

Dataset |  Metric | STGCN$\dagger$ | ASTGCN$\dagger$ |  Graph WaveNet$\dagger$ | GMAN* | STSGCN$\dagger$ |  AGCRN* | STFGNN$\dagger$ | ESGCN
---|---|---|---|---|---|---|---|---|---
PEMS03 | RMSE | 30.12 $\pm$ 0.70 | 29.66 $\pm$ 1.68 | 32.94 $\pm$ 0.18 | 27.92 $\pm$ 1.15 | 29.21 $\pm$ 0.56 | 28.19 $\pm$ 0.20 | 28.34$\pm$ 0.46 | 25.01 $\pm$ 0.34
MAE | 17.49 $\pm$ 0.46 | 17.69 $\pm$ 1.43 | 19.85 $\pm$ 0.03 | 16.87 $\pm$ 0.31 | 17.48 $\pm$ 0.15 | 15.93 $\pm$ 0.08 | 16.77 $\pm$ 0.09 | 15.05 $\pm$ 0.07
MAPE | 17.15 $\pm$ 0.45 | 19.40 $\pm$ 2.24 | 19.31 $\pm$ 0.49 | 18.23 $\pm$ 0.01 | 16.78 $\pm$ 0.20 | 14.97 $\pm$ 0.17 | 16.30 $\pm$ 0.09 | 14.32 $\pm$ 0.01
PEMS04 | RMSE | 35.55 $\pm$ 0.75 | 35.22 $\pm$ 1.90 | 39.70 $\pm$ 0.04 | 30.71 $\pm$ 0.09 | 33.65 $\pm$ 0.20 | 32.35 $\pm$ 0.22 | 31.88$\pm$ 0.14 | 30.38 $\pm$ 0.10
MAE | 22.70 $\pm$ 0.64 | 22.93 $\pm$ 1.29 | 25.45 $\pm$ 0.03 | 19.16 $\pm$ 0.04 | 21.19 $\pm$ 0.10 | 19.76 $\pm$ 0.09 | 19.83 $\pm$ 0.06 | 18.88 $\pm$ 0.05
MAPE | 14.59 $\pm$ 0.21 | 16.56 $\pm$ 1.36 | 17.29 $\pm$ 0.24 | 13.40 $\pm$ 0.01 | 13.90 $\pm$ 0.05 | 13.13 $\pm$ 0.35 | 13.02 $\pm$ 0.05 | 12.70 $\pm$ 0.01
PEMS07 | RMSE | 38.78 $\pm$ 0.58 | 42.57 $\pm$ 3.31 | 42.78 $\pm$ 0.07 | - | 39.03 $\pm$ 0.27 | 35.09 $\pm$ 0.10 | 35.80 $\pm$ 0.18 | 33.93 $\pm$ 0.28
MAE | 25.38 $\pm$ 0.49 | 28.05 $\pm$ 2.34 | 26.85 $\pm$ 0.05 | - | 24.26 $\pm$ 0.14 | 21.17 $\pm$ 0.13 | 22.07 $\pm$ 0.11 | 20.55 $\pm$ 0.10
MAPE | 11.08 $\pm$ 0.18 | 13.92 $\pm$ 1.65 | 12.12 $\pm$ 0.41 | - | 10.21 $\pm$ 1.65 | 8.97 $\pm$ 0.07 | 9.21 $\pm$ 0.07 | 8.61 $\pm$ 0.01
PEMS08 | RMSE | 27.83 $\pm$ 0.20 | 28.16 $\pm$ 0.48 | 31.05 $\pm$ 0.07 | 24.71 $\pm$ 0.13 | 26.80 $\pm$ 0.18 | 25.69 $\pm$ 0.21 | 26.22 $\pm$ 0.15 | 24.59 $\pm$ 0.15
MAE | 18.02 $\pm$ 0.14 | 18.61 $\pm$ 0.40 | 19.13 $\pm$ 0.08 | 15.69 $\pm$ 0.02 | 17.13 $\pm$ 0.09 | 16.22 $\pm$ 0.13 | 16.64 $\pm$ 0.09 | 15.56 $\pm$ 0.10
MAPE | 11.40 $\pm$ 0.10 | 13.08 $\pm$ 1.00 | 12.68 $\pm$ 0.57 | 10.04 $\pm$ 0.01 | 10.96 $\pm$ 0.07 | 10.50 $\pm$ 0.20 | 10.60 $\pm$ 0.06 | 9.88 $\pm$ 0.08
Comp. cost | # Param. | 384,243 | 450,031 | 311,400 | 229,569 | 2,872,686 | 748,810 | 3,873,580 | 199,062
Train | 9.14 | 31.79 | 29.30 | 86.9 | 43.16 | 21.57 | 43.02 | 26.65
Test | 6.92 | 3.69 | 2.19 | 9.0 | 48.45 | 2.43 | 48.45 | 3.02

Table 1: Performance and computation cost comparison of ESGCN and baseline
models on PEMS03, PEMS04, PEMS07 and PEMS08 datasets. The values in bold
indicates the best performances. The experiments are repeated ten times. MAPE
metric is measured as a percentage. Comp. cost denotes computational cost. The
Train is average time in seconds per epoch and the Test is total time of
running test dataset in seconds. The computation costs are computed on the
same machine with the same mini batch size on PEMS04 dataset. $\dagger$:
Experimental results are excerpted from [Li and Zhu, 2021] *: Experiment is
conducted by us. For all metrics, a smaller value denotes superior
performance.

##  5 Experiments

###  5.1 Implementation Details

The proposed model is trained with an Adam optimizer Kingma and Ba (2015) for
50 epochs. The initial learning rate is 0.0003 and reduced by 0.7 every 5
epochs. The weight decay factor for L2 regulation is set to 0.0001, and the
batch size is set to 30 for PEMS07 and 64 for PEMS03, 04, and 08. Training
sessions are conducted on an NVIDIA Tesla V100 and Intel Xeon Gold 5120 CPU.

Figure 3: Heatmap of AAM with and without node contrastive loss on the first
45 nodes.

###  5.2 Datasets

Datasets | Nodes | MissingRatio | Range
---|---|---|---
PEMS03 | 358 | 0.672% | 9/1/2018 - 11/30/2018
PEMS04 | 307 | 3.182% | 1/1/2018 - 2/28/2018
PEMS07 | 883 | 0.452% | 5/1/2017 - 8/31/2017
PEMS08 | 170 | 0.696% | 7/1/2016 - 8/31/2016

Table 2: Dataset description and statistics.

The framework is validated on four real-world traffic datasets, namely:
PEMS03, PEMS04, PEMS07, and PEMS08 Song et al. (2020). Table 2 shows the
description of each dataset. These four datasets contain generated traffic
flows in four different regions of California using the Caltrans Performance
Measurement System. Time-series data are collected at 5-minute intervals.
Standard normalization and linear interpolation are used for stable training.
For a fair comparison, all datasets are split into training, validation, and
test data at a ratio of 6 : 2 : 2. Twelve time steps (1 h) are used to predict
the next 12 time steps (1 h) and all experiments are repeated 10 times with
random seeds. The test data performance is verified by selecting the model of
the epoch that showed the best performance in the validation data.

###  5.3 Baseline Methods

ESGCN is compared with the following models on the same hyperparameters and
official implementations:

  * •

STGCN: Spatio-temporal graph convolutional networks, which comprise spatial
and temporal dilated convolutions Yu et al. (2017).

  * •

ASTGCN: Attention-based spatial temporal graph convolutional networks, which
adopt spatial and temporal attention into the model Guo et al. (2019).

  * •

GraphWaveNet: Graph WaveNet exploits an adaptive adjacency graph and dilated
1D convolution Wu et al. (2019).

  * •

GMAN: Graph multi-attention network uses spatial and temporal attention in
graph neural network Zheng et al. (2020).

  * •

STSGCN: Spatial–temporal synchronous graph convolutional networks, which
utilize a spatio-temporal graph that extends the spatial graph to the temporal
dimension Song et al. (2020).

  * •

AGCRN: Adaptive graph convolutional recurrent network for traffic forecasting.
This model exploits node adaptive parameter learning and an adaptive graph Bai
et al. (2020b).

  * •

STFGNN: Spatial–temporal fusion graph neural networks, which leverage fast-DTW
to construct a spatiotemporal graph Li and Zhu (2021).

Case | ES module | Node contrastive loss | Attention Op. | $\lambda$ | Rep. function | RMSE | MAE | MAPE(%)
---|---|---|---|---|---|---|---|---
1 |  |  | - | - | - | 30.69 | 19.12 | 12.91
2 | ✓ |  | Max | - | Last | 30.45 | 18.94 | 12.80
3(ours) | ✓ | ✓ | Max | 0.1 | Last | 30.38 | 18.88 | 12.72
4 | ✓ | ✓ | Max | 0.3 | Last | 30.41 | 18.89 | 12.78
5 | ✓ | ✓ | Max | 0.5 | Last | 30.62 | 18.98 | 12.74
6 | ✓ | ✓ | Max | 0.7 | Last | 30.40 | 18.92 | 12.79
7 | ✓ | ✓ | Max | 0.9 | Last | 30.64 | 18.99 | 12.80
8 | ✓ | ✓ | Avg | 0.1 | Last | 30.46 | 18.98 | 12.91
9 | ✓ | ✓ | Max + L | 0.1 | Last | 30.48 | 18.89 | 12.87
10 | ✓ | ✓ | Max | 0.1 | Middle | 30.55 | 19.04 | 12.80
11 | ✓ | ✓ | Max | 0.1 | First | 30.64 | 19.12 | 12.87

Table 3: Ablation experiments on the PEMS04 dataset. Attention Op. ,
$\lambda$, and Rep. function denote Attention Operation, $\lambda$ weight for
node contrastive loss, and representative function. Max, Avg, Max + L denote
Max operation, Average operation, and Max operation with learnable layers.
Last, Middle, and Last are representative node positions.

###  5.4 Comparison with the Baseline Methods

The proposed model is compared with state-of-the-art models. ESGCN
outperformed all other baselines in terms of RMSE, MAE, and MAPE, as shown in
Table 1. Compared to the current best performing models in each dataset
(PEMS07: AGCRN, PEMS03,04, and 08: GMAN), ESGCN yields 4%, 2.7%, and 3.8%
relative improvements on average for all datasets in RMSE, MAE, and MAPE.
Graph WaveNet and AGCRN employ the AAM and STFGNN uses a non-adaptive spatio-
temporal adjacency matrix. Compared to these methods, the proposed model
showed superior performance and guaranteed the effectiveness of our AAM which
is refined with attention and node contrastive loss. Based on the experimental
results, ESGCN, has improved representation ability and exhibits promising
forecasting performance.

###  5.5 Computational Cost

To evaluate the computational cost, we compare the number of parameters, the
training time, and the inference time of our model with those of baselines in
Table 1. ESGCN has the smallest parameters compared to other baselines. In
training time, ESGCN is faster than STFGNN and slightly slower than AGCRN.
ESGCN has the third-fastest inference speed. Although AGCRN shows faster
training and inference time, specifically, 5 and 0.6 s in training and
inference, respectively, the differences are insignificant. Additionally,
AGCRN requires three times more parameters than ours. Especially, our model
which has the channel attention shows faster running speed and includes fewer
parameters than GMAN which leverages the transformer attention. Considering
its superior performance, ESGCN has an acceptable computational cost.

###  5.6 Ablation Study

####  5.6.1 Components

To validate the proposed ES module and the node contrastive loss, we conduct
experiments (cases 1-3) as shown in Table 3. The model with only W-module
(case 1) shows the lowest performance. However, the W-module with the ES
module which reflects traffic flows using GCN (case 2) achieves significant
improvement. ESGCN, consisting of the W-module, ES module, and the node
contrastive loss (case 3) outperforms the others. This highlights that the
feature enhancement ability of ES module and the proposed loss function.

####  5.6.2 Attention operation

To extract important adjacency relations, our edge attention used max
operation to squeeze channel features. However, CBAM Woo et al. (2018) also
uses average operation to squeeze channel features. In this section, we
conduct an ablation study (cases 3, 8, and 9) on the attention operation in
Table 3. We empirically discovered that the max operation is the most optimal
setting. Notably, the attention function with learnable layer (case 9) shows a
lower performance because dimension reduction of the layer disturbs to extract
the AAM Wang et al. (2020).

####  5.6.3 Node contrastive loss

We use a hyperparameter $\lambda$ to balance the Huber loss and node
contrastive loss. Although the node contrastive loss assists in the
construction of the refined AAM, an excessive effect of this loss can
negatively affect forecasting performance. We empirically determine the
magnitude of $\lambda$. The experimental results (cases 3-7) are shown in
Table 3. Based on the results, the optimal value of $\lambda$ is 0.1. We set
$\lambda$ as 0.1 across in all the experiments. In Fig 3, it shows
additionally sparse AAM in which unrelated connections are removed.

####  5.6.4 Representative function

The representative function of ES module extracts temporal representatives on
time dimension. Inspired by MDP, the representative function returns the last
step node of $F_{c}$. To test whether the last node can be a temporal
representative, we conduct an experiment by replacing the last node with the
first node, $F_{c}(:,:,0)$, and the middle node, $F_{c}(:,:,l_{4}/2)$. The
experiment results (cases 3,10, and 11) are shown in Table 3. The closer the
node to the initial time step is, the lower the performance. Our
representative function is based on the assumption the closest value to the
prediction value can be the representatives. This ablation study empirically
shows the assumption is reasonable.

##  6 Conclusion and Future Work

This study proposes a novel method, ESGCN, to address traffic flow
forecasting. Experiments show that the proposed model achieves state-of-the-
art performance on four real-world datasets and has the smallest parameters
with relatively faster inference and training speed. In the future, given that
ESGCN is designed as a general framework to handle spatio-temporal data, it
can be applied to other applications that have spatio-temporal data structures
such as regional housing market prediction and electricity demand forecasting.

## References

  * Ba et al. [2016] Jimmy Ba, J. Kiros, and Geoffrey E. Hinton.  Layer normalization.  ArXiv, abs/1607.06450, 2016.
  * BAI et al. [2020a] LEI BAI, Lina Yao, Can Li, Xianzhi Wang, and Can Wang.  Adaptive graph convolutional recurrent network for traffic forecasting.  In Advances in Neural Information Processing Systems, volume 33, pages 17804–17815, 2020.
  * Bai et al. [2020b] Lei Bai, Lina Yao, Can Li, Xianzhi Wang, and Can Wang.  Adaptive graph convolutional recurrent network for traffic forecasting.  arXiv preprint arXiv:2007.02842, 2020.
  * Bellman and Howard [1961] R. Bellman and R. Howard.  Dynamic programming and markov processes.  Mathematics of Computation, 15:217, 1961.
  * Chen and Chen [2021] Jun Chen and Hao-Peng Chen.  Edge-featured graph attention network.  ArXiv, abs/2101.07671, 2021.
  * Fang et al. [2019] Shen Fang, Qi Zhang, Gaofeng Meng, Shiming Xiang, and Chunhong Pan.  Gstnet: Global spatial-temporal network for traffic flow prediction.  In IJCAI, pages 2286–2293, 2019.
  * Gong and Cheng [2019] Liyu Gong and Qiang Cheng.  Exploiting edge features for graph neural networks.  2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 9203–9211, 2019.
  * Guo et al. [2019] Shengnan Guo, Youfang Lin, Ning Feng, Chao Song, and Huaiyu Wan.  Attention based spatial-temporal graph convolutional networks for traffic flow forecasting.  In Proceedings of the AAAI Conference on Artificial Intelligence, pages 922–929, 2019.
  * Hechtlinger et al. [2017] Yotam Hechtlinger, Purvasha Chakravarti, and Jining Qin.  A generalization of convolutional neural networks to graph-structured data.  arXiv preprint arXiv:1704.08165, 2017.
  * Hu et al. [2020] Jie Hu, Li Shen, Samuel Albanie, Gang Sun, and Enhua Wu.  Squeeze-and-excitation networks.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 42:2011–2023, 2020.
  * Kingma and Ba [2015] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  CoRR, abs/1412.6980, 2015.
  * Kipf and Welling [2017] Thomas Kipf and M. Welling.  Semi-supervised classification with graph convolutional networks.  ArXiv, abs/1609.02907, 2017.
  * Kong et al. [2020] Xiangyuan Kong, Weiwei Xing, Xiang Wei, Peng Bao, Jia kui Zhang, and Wei Lu.  Stgat: Spatial-temporal graph attention networks for traffic flow forecasting.  IEEE Access, 8:134363–134372, 2020.
  * Li and Zhu [2021] Mengzhang Li and Zhanxing Zhu.  Spatial-temporal fusion graph neural networks for traffic flow forecasting.  In AAAI, 2021.
  * Li et al. [2018] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu.  Diffusion convolutional recurrent neural network: Data-driven traffic forecasting.  In International Conference on Learning Representations (ICLR ’18), 2018.
  * Luo et al. [2018] Chunjie Luo, Jianfeng Zhan, Lei Wang, and Qiang Yang.  Cosine normalization: Using cosine similarity instead of dot product in neural networks.  ArXiv, abs/1702.05870, 2018.
  * Park et al. [2020] Cheonbok Park, Chunggi Lee, Hyojin Bahng, Yunwon Tae, Seungmin Jin, Kihwan Kim, Sungahn Ko, and Jaegul Choo.  St-grat: A novel spatio-temporal graph attention networks for accurately forecasting dynamically changing road speed.  Proceedings of the 29th ACM International Conference on Information & Knowledge Management, 2020.
  * Song et al. [2020] Chao Song, Youfang Lin, S. Guo, and Huaiyu Wan.  Spatial-temporal synchronous graph convolutional networks: A new framework for spatial-temporal network data forecasting.  In AAAI, 2020.
  * Vaswani et al. [2017] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  ArXiv, abs/1706.03762, 2017.
  * Wang et al. [2020] Qilong Wang, Banggu Wu, Pengfei Zhu, P. Li, Wangmeng Zuo, and Qinghua Hu.  Eca-net: Efficient channel attention for deep convolutional neural networks.  2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11531–11539, 2020.
  * Woo et al. [2018] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In-So Kweon.  Cbam: Convolutional block attention module.  In ECCV, 2018.
  * Wu et al. [2019] Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, and Chengqi Zhang.  Graph wavenet for deep spatial-temporal graph modeling.  arXiv preprint arXiv:1906.00121, 2019.
  * Wu et al. [2020] Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, Xiaojun Chang, and C. Zhang.  Connecting the dots: Multivariate time series forecasting with graph neural networks.  Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2020.
  * Yao et al. [2019] Huaxiu Yao, Xianfeng Tang, Hua Wei, Guanjie Zheng, and Z. Li.  Revisiting spatial-temporal similarity: A deep learning framework for traffic prediction.  In AAAI, 2019.
  * Yu et al. [2017] Bing Yu, Haoteng Yin, and Zhanxing Zhu.  Spatio-temporal graph convolutional networks: A deep learning framework for traffic forecasting.  arXiv preprint arXiv:1709.04875, 2017.
  * Zheng et al. [2020] Chuanpan Zheng, Xiaoliang Fan, C. Wang, and Jianzhong Qi.  Gman: A graph multi-attention network for traffic prediction.  In AAAI, 2020.