\section{Introduction}
\begin{wrapfigure}[22]{r}{0.5\textwidth}
    \centering
    \vspace{-18pt}
    \includegraphics[width=0.5\textwidth]{figures/conceptual}
    \caption{
        Conceptual, interpolated map from investigated hyperparameters of proposed algorithm \textsc{nre-c}\xspace to a measurement of posterior exactness using the Classifier Two-Sample Test. Best 0.5, worst 1.0. Red dot indicates \textsc{nre-a}\xspace's hyperparameters, $\gamma = 1$ and $K = 1$ \cite{Hermans2019}. Purple line implies \textsc{nre-b}\xspace \cite{Durkan2020} with $\gamma = \infty$ and $K \geq 1$. \textsc{nre-c}\xspace covers the entire plane, generalizing other methods. Best performance occurs with $K > 1$ and $\gamma \approx 1$, in contrast with the settings of existing algorithms.
    }
    \label{fig:conceptual-hyperparameters-to-c2st}
\end{wrapfigure}

We begin with a motivating example: Consider the task of inferring the mass of an exoplanet $\boldsymbol{\theta}_o$ from the light curve observations $\boldsymbol{x}_o$ of a distant star. We design a computer program that maps hypothetical mass $\boldsymbol{\theta}$ to a simulated light curve $\boldsymbol{x}$ using relevant physical theory. Our simulator computes $\boldsymbol{x}$ from $\boldsymbol{\theta}$, but the inverse mapping is unspecified and likely intractable. \emph{Simulation-based inference} (\textsc{sbi}\xspace) puts this problem in a probabilistic context \cite{sisson2018handbook, Cranmer2020}. Although we cannot analytically evaluate it, we assume that the simulator is sampling from the conditional probability distribution $p(\boldsymbol{x} \mid \boldsymbol{\theta})$. After specifying a prior $p(\boldsymbol{\theta})$, the inverse amounts to estimating the posterior $p(\boldsymbol{\theta} \mid \boldsymbol{x}_o)$. This problem setting occurs across scientific domains \cite{cole2021fast, alsing2018massive, brehmer2018constraining, hermans2020towards, lensing} where $\boldsymbol{\theta}$ generally represents input parameters of the simulator and $\boldsymbol{x}$ the simulated output observation. Our design goal is to produce a surrogate model $\hat{p}(\boldsymbol{\theta} \mid \boldsymbol{x})$ approximating the posterior for any data $\boldsymbol{x}$ while limiting excessive simulation.

\begin{wrapfigure}[32]{R}{0.5\textwidth}
    \centering
    \vspace{-15pt}
    \includegraphics[width=0.5\textwidth]{figures/loss-diagram.png}
    \caption{
        Schematic depicts how the loss is computed in \textsc{nre}\xspace algorithms. ($\boldsymbol{\theta}$, $\boldsymbol{x}$) pairs are sampled from distributions at the top of the figure, entering the loss functions as depicted. \textsc{nre-c}\xspace controls the number of contrastive classes with $K$ and the weight of independent and dependent terms with $p_0$ and $p_K$. \textsc{nre-c}\xspace generalizes other algorithms. Hyperparameters recovering \textsc{nre-a}\xspace and \textsc{nre-b}\xspace are listed next to the name within the dashed areas. Notation details are defined in Section~\ref{sec:nrec}.
    }
    \label{fig:loss-diagram}
\end{wrapfigure}

Density estimation \cite{bishop, papamakarios2017masked, papamakarios2019normalizing} can fit the likelihood \cite{drovandi2018approximating, papamakarios2019sequential, alsing2019fast, lueckmann2019likelihood} or posterior \cite{blum2010non, papamakarios2016fast, lueckmann2017flexible, greenberg2019automatic} directly; however, an appealing alternative for practitioners is estimating a \emph{ratio} between distributions \cite{izbicki2014high, Cranmer2015, thomas2016likelihood, Hermans2019, Durkan2020}. Specifically, the likelihood-to-evidence ratio $\frac{p(\boldsymbol{\theta} \mid \boldsymbol{x})}{p(\boldsymbol{\theta})} = \frac{p(\boldsymbol{x} \mid \boldsymbol{\theta})}{p(\boldsymbol{x})} = \frac{p(\boldsymbol{\theta}, \boldsymbol{x})}{p(\boldsymbol{\theta}) p(\boldsymbol{x})}$.
Unlike the other methods, ratio estimation enables easy aggregation of independent and identically drawn data $\boldsymbol{x}$. Ratio and posterior estimation can compute bounds on the mutual information and an importance sampling diagnostic.

Estimating $\frac{p(\boldsymbol{x} \mid \boldsymbol{\theta})}{p(\boldsymbol{x})}$ can be formulated as a binary classification task \cite{Hermans2019}, where the classifier $\sigma \circ f_{\boldsymbol{w}}(\boldsymbol{\theta},\boldsymbol{x})$ distinguishes between pairs $(\boldsymbol{\theta}, \boldsymbol{x})$ sampled either from the joint distribution $p(\boldsymbol{\theta}, \boldsymbol{x})$ or the product of its marginals $p(\boldsymbol{\theta}) p(\boldsymbol{x})$.
We call it \textsc{nre-a}\xspace. The optimal classifier has
\begin{align}
     f_{\boldsymbol{w}}(\boldsymbol{\theta},\boldsymbol{x}) \approx \log \frac{p(\boldsymbol{\theta} \mid \boldsymbol{x})}{p(\boldsymbol{\theta})}.
\end{align}
Here, $\sigma$ represents the sigmoid function, $\circ$ implies function composition, and $f_{\boldsymbol{w}}$ is a neural network with weights $\boldsymbol{w}$. As a part of an effort to unify different \textsc{sbi}\xspace methods and to improve simulation-efficiency, \citet{Durkan2020} reformulated the classification task to identify which of $K$ possible $\boldsymbol{\theta}_k$ was responsible for simulating $\boldsymbol{x}$. We refer to it as \textsc{nre-b}\xspace.
At optimum
\begin{align}
    g_{\boldsymbol{w}}(\boldsymbol{\theta}, \boldsymbol{x}) \approx \log \frac{p(\boldsymbol{\theta} \mid \boldsymbol{x})}{p(\boldsymbol{\theta})} + c_{\boldsymbol{w}}(\boldsymbol{x}),
\end{align}
where an additional bias, $c_{\boldsymbol{w}}(\boldsymbol{x})$, appears. $g_{\boldsymbol{w}}$ represents another neural network. The $c_{\boldsymbol{w}}(\boldsymbol{x})$ term nullifies many of the advantages ratio estimation offers. $c_{\boldsymbol{w}}(\boldsymbol{x})$ can be arbitrarily pathological in $\boldsymbol{x}$, meaning that the normalizing constant can take on extreme values. This limits the applicability of verification tools like the importance sampling-based diagnostic in Section~\ref{sec:diagnostic}.

The $c_{\boldsymbol{w}}(\boldsymbol{x})$ term also arises in contrastive learning \cite{gutmann2012noise, van2018representation} with \citet{ma2018noise} attempting to estimate it in order to reduce its impact. We will propose a method that discourages this bias instead. Further discussion in Appendix~\ref{apndx:mutual-information}.

There is a distinction in deep learning-based \textsc{sbi}\xspace between \emph{amortized} and \emph{sequential} algorithms which produce surrogate models that estimate any posterior $p(\boldsymbol{\theta} \mid \boldsymbol{x})$ or a specific posterior $p(\boldsymbol{\theta} \mid \boldsymbol{x}_o)$ respectively. Amortized algorithms sample parameters from the prior, while sequential algorithms use an alternative proposal distribution--increasing efficiency at the expense of flexibility. Amortization is usually necessary to compute diagnostics that do not require samples from $p(\boldsymbol{\theta} \mid \boldsymbol{x}_o)$ and amortized estimators are empirically more reliable \cite{hermans2021averting}. Our study therefore focuses on amortized algorithms.

\paragraph{Contribution}
We design a more general formulation of likelihood-to-evidence ratio estimation as a multiclass problem in which the bias inherent to \textsc{nre-b}\xspace is discouraged by the loss function and it does not appear at optimum. Figure~\ref{fig:conceptual-hyperparameters-to-c2st} diagrams the interpolated performance as a function of hyperparameters. It shows which settings recover \textsc{nre-a}\xspace and \textsc{nre-b}\xspace, also indicating that highest performance occurs with settings distant from these. Figure~\ref{fig:loss-diagram} shows the relationship of the loss functions. We call our framework \textsc{nre-c}\xspace \footnote{The code for our project can be found at \href{https://github.com/bkmi/cnre}{https://github.com/bkmi/cnre} under the Apache License 2.0.} and expound the details in Section~\ref{sec:methods}.

An existing importance sampling diagnostic \cite{Hermans2019} tests whether a classifier can distinguish $p(\boldsymbol{x} \mid \boldsymbol{\theta})$ samples from from samples from $p(\boldsymbol{x})$ weighted by the estimated ratio. We demonstrate that, when estimating accurate posteriors, our proposed \textsc{nre-c}\xspace passes this diagnostic while \textsc{nre-b}\xspace does not.

Taking inspiration from mutual information estimation \cite{poole2019variational}, we propose applying a variational bound on the mutual information between $\boldsymbol{\theta}$ and $\boldsymbol{x}$ in a novel way--as an informative metric measuring a lower bound on the Kullback-Leibler divergence between surrogate posterior estimate $p_{\boldsymbol{w}}(\boldsymbol{\theta} \mid \boldsymbol{x})$ and $p(\boldsymbol{\theta} \mid \boldsymbol{x})$, averaged over $p(\boldsymbol{x})$. Unlike with two-sample testing methods commonly used in machine learning literature \cite{sbibm}, our metric samples only from $p(\boldsymbol{\theta}, \boldsymbol{x})$, which is always available in \textsc{sbi}\xspace, and does not require samples from the intractable $p(\boldsymbol{\theta} \mid \boldsymbol{x})$. Our metric is meaningful to scientists working on problems with intractable posteriors. The technique requires estimating the partition function, which can be expensive. We find the metric to be well correlated with results from two-sample tests.

We evaluate \textsc{nre-b}\xspace and \textsc{nre-c}\xspace in a fair comparison in several training regimes in Section~\ref{sec:experiments}. We perform a hyperparameter search on three simulators with tractable likelihood by benchmarking the behavior when (a) jointly drawn pairs $(\boldsymbol{\theta}, \boldsymbol{x})$ are unlimited or when jointly drawn pairs $(\boldsymbol{\theta}, \boldsymbol{x})$ are fixed but we (b) can draw from the prior $p(\boldsymbol{\theta})$ without limit or (c) are restricted to the initial pairs. We also perform the \textsc{sbi}\xspace benchmark of \citet{sbibm} with our recommended hyperparameters.

\section{Methods}
\label{sec:methods}
The ratio between probability distributions can be estimated using the ``likelihood ratio trick'' by training a classifier to distinguish samples \cite{hastie2009elements, sugiyama2012density, goodfellow2014generative, Cranmer2015, thomas2016likelihood, Mohamed2016, Hermans2019}. We first summarize the loss functions of \textsc{nre-a}\xspace and \textsc{nre-b}\xspace which approximate the \emph{intractable} likelihood-to-evidence ratio $r(\boldsymbol{x} \mid \boldsymbol{\theta}) \coloneqq \frac{p(\boldsymbol{x} \mid \boldsymbol{\theta})}{p(\boldsymbol{x})}$. We then elaborate on our proposed generalization, \textsc{nre-c}\xspace. Finally, we explain how to recover \textsc{nre-a}\xspace and \textsc{nre-b}\xspace within our framework and comment on the normalization properties.

\paragraph{\textsc{nre-a}\xspace} \citet{Hermans2019} train a binary classifier to distinguish $(\boldsymbol{\theta}, \boldsymbol{x})$ pairs drawn dependently $p(\boldsymbol{\theta}, \boldsymbol{x})$ from those drawn independently $p(\boldsymbol{\theta}) p(\boldsymbol{x})$. This classifier is parameterized by a neural network $f_{\boldsymbol{w}}$ which approximates $\log r(\boldsymbol{x} \mid \boldsymbol{\theta})$. We seek optimal network weights

\begin{equation}
    \boldsymbol{w}^{*} \in \argmin_{\boldsymbol{w}} - \frac{1}{2 B}
    \left[ \sum_{b=1}^{B} \log \left(
        1 - \sigma \circ f_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b)}, \boldsymbol{x}^{(b)})
    \right)
    + \sum_{b'=1}^{B}\log \left(
        \sigma \circ f_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b')}, \boldsymbol{x}^{(b')})
    \right) \right]
\end{equation}

$\boldsymbol{\theta}^{(b)}, \boldsymbol{x}^{(b)} \sim p(\boldsymbol{\theta}) p(\boldsymbol{x})$ and $\boldsymbol{\theta}^{(b')}, \boldsymbol{x}^{(b')} \sim p(\boldsymbol{\theta}, \boldsymbol{x})$ over $B$ samples.
\textsc{nre-a}\xspace's ratio estimate converges to $f_{\boldsymbol{w}^*} = \log \frac{p(\boldsymbol{x} \mid \boldsymbol{\theta})}{p(\boldsymbol{x})}$ given unlimited model flexibility and data. Details can be found in Appendix~\ref{apndx:other-sbi}.
\looseness=-1

\paragraph{\textsc{nre-b}\xspace}
\citet{Durkan2020} train a classifier that selects from among $K$ parameters $(\boldsymbol{\theta}_1, \ldots, \boldsymbol{\theta}_K)$ which could have generated $\boldsymbol{x}$, in contrast with \textsc{nre-a}\xspace's binary possibilities. One of these parameters $\boldsymbol{\theta}_k$ is \emph{always} drawn jointly with $\boldsymbol{x}$. The classifier is parameterized by a neural network $g_{\boldsymbol{w}}$ which approximates $\log r(\boldsymbol{x} \mid \boldsymbol{\theta})$. Training is done over $B$ samples by finding

\begin{equation}
\label{eqn:nreb-loss}
    \boldsymbol{w}^* \in \argmin_{\boldsymbol{w}}
    	\left[ - \frac{1}{B} \sum_{b'=1}^{B} \log \frac{ \exp \circ g_{\boldsymbol{w}}(\boldsymbol{\theta}_{k}^{(b')}, \boldsymbol{x}^{(b')}) }{ \sum_{i=1}^{K} \exp \circ g_{\boldsymbol{w}}(\boldsymbol{\theta}_i^{(b')}, \boldsymbol{x}^{(b')}) } \right]
\end{equation}

where $\boldsymbol{\theta}_{1}^{(b')}, \ldots, \boldsymbol{\theta}_{K}^{(b')} \sim p(\boldsymbol{\theta})$ and $\boldsymbol{x}^{(b')} \sim p(\boldsymbol{x} \mid \boldsymbol{\theta}_k^{(b')})$. Given unlimited model flexibility and data \textsc{nre-b}\xspace's ratio estimate converges to ${g_{\boldsymbol{w}^*}(\boldsymbol{\theta}, \boldsymbol{x}) = \log \frac{ p(\boldsymbol{\theta} \mid \boldsymbol{x}) }{p(\boldsymbol{\theta})} + c_{\boldsymbol{w}^*}(\boldsymbol{x})}$. Details are in Appendix~\ref{apndx:other-sbi}.
\looseness=-1

\subsection{Contrastive Neural Ratio Estimation}
\label{sec:nrec}
Our proposed algorithm \textsc{nre-c}\xspace trains a classifier to identify which $\boldsymbol{\theta}$ among $K$ candidates is responsible for generating a given $\boldsymbol{x}$, inspired by \textsc{nre-b}\xspace. We added another option that indicates $\boldsymbol{x}$ was drawn independently, inspired by \textsc{nre-a}\xspace. The introduction of the additional class yields a ratio without the specific $c_{\boldsymbol{w}}(\boldsymbol{x})$ bias at optimum. Define $\boldsymbol{\Theta} \coloneqq (\boldsymbol{\theta}_1, ..., \boldsymbol{\theta}_K)$ and conditional probability

\begin{equation}
    p_{\textsc{nre-c}\xspace}(\boldsymbol{\Theta}, \boldsymbol{x} \mid y = k) \coloneqq
    \begin{cases}
        p(\boldsymbol{\theta}_1) \cdots p(\boldsymbol{\theta}_K) p(\boldsymbol{x}) & k=0 \\
    	p(\boldsymbol{\theta}_1) \cdots p(\boldsymbol{\theta}_K) p(\boldsymbol{x} \mid \boldsymbol{\theta}_k) & k = 1, \ldots, K
    \end{cases}.
\end{equation}

We set marginal probabilities $p(y = k) \coloneqq p_{K}$ for all $k \geq 1$ and ${p(y = 0) \coloneqq p_{0}}$, yielding the relationship $p_{0} = 1 - K p_{K}$. Let the odds of any pair being drawn dependently to completely independently be $\gamma \coloneqq \frac{K p_{K}}{p_{0}}$. We now use Bayes' formula to compute the conditional probability

\begin{equation}
\begin{aligned}
\label{eqn:cnre-posterior}
    p(y=k \mid \boldsymbol{\Theta}, \boldsymbol{x})
    &= \frac{p(y=k) \, p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=k)/p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)}{\sum_{i=0}^{K} p(y=i) \,  p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=i)/p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)} \\
    &= \frac{p(y=k) \, p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=k)/p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)}{p(y=0) + \sum_{i=1}^{K} p(y=i) \, p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=i)/p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)} \\
    &= \begin{cases}
    	\frac{K}{K + \gamma \sum_{i=1}^{K} r(\boldsymbol{x} \mid \boldsymbol{\theta}_i)} & k=0 \\
    	\frac{\gamma \, r(\boldsymbol{x} \mid \boldsymbol{\theta}_k)}{K + \gamma \sum_{i=1}^{K} r(\boldsymbol{x} \mid \boldsymbol{\theta}_i)} & k=1, \ldots, K
    \end{cases}.
\end{aligned}
\end{equation}

We dropped the \textsc{nre-c}\xspace subscript and substituted in $\gamma$ to replace the $p(y)$ class probabilities. We train a classifier, parameterized by neural network $h_{\boldsymbol{w}}(\boldsymbol{\theta}, \boldsymbol{x})$ with weights $\boldsymbol{w}$, to approximate \eqref{eqn:cnre-posterior} by

\begin{equation}
\label{eqn:classifier}
    q_{\boldsymbol{w}}(y = k \mid \boldsymbol{\Theta}, \boldsymbol{x}) =
    \begin{cases}
    	\frac{K}{K + \gamma \sum_{i=1}^{K} \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}_i, \boldsymbol{x})} & k = 0 \\
    	\frac{\gamma \, \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}_k,\boldsymbol{x}))}{K + \gamma \sum_{i=1}^{K} \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}_i,\boldsymbol{x})} & k = 1, \ldots, K.
    \end{cases}.
\end{equation}

We note that \eqref{eqn:classifier} still satisfies $\sum_{k=0}^K q_{\boldsymbol{w}}(y = k \mid \boldsymbol{\Theta}, \boldsymbol{x}) =1$, no matter the parameterization.

\paragraph{Optimization}
We design a loss function that encourages $h_{\boldsymbol{w}}(\boldsymbol{\theta}, \boldsymbol{x}) = \log \frac{p(\boldsymbol{x} \mid \boldsymbol{\theta})}{p(\boldsymbol{x})}$ at convergence, and holds at optimum with unlimited flexibility and data. We introduce the cross entropy loss
\begin{equation}
\begin{aligned}
    \ell(\boldsymbol{w})
    &\coloneqq \mathbb{E}_{p(y, \boldsymbol{\Theta}, \boldsymbol{x})} \left[ -\log q_{\boldsymbol{w}}(y \mid \boldsymbol{\Theta}, \boldsymbol{x})\right] \\
    &= - p_{0} \mathbb{E}_{p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)} \left[
        \log q_{\boldsymbol{w}}(y =0 \mid \boldsymbol{\Theta}, \boldsymbol{x})
    \right]
    - p_{K} \sum_{k=1}^{K} \mathbb{E}_{p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=k)} \left[
        \log q_{\boldsymbol{w}}(y = k \mid \boldsymbol{\Theta}, \boldsymbol{x})
    \right] \\
    &= - p_{0} \mathbb{E}_{p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=0)} \left[
    \log q_{\boldsymbol{w}}(y =0 \mid \boldsymbol{\Theta}, \boldsymbol{x})
    \right]
    - K p_{K} \mathbb{E}_{p(\boldsymbol{\Theta}, \boldsymbol{x} \mid y=K)} \left[
    \log q_{\boldsymbol{w}}(y = K \mid \boldsymbol{\Theta}, \boldsymbol{x})
    \right]
\end{aligned}
\end{equation}
and minimize it towards $\boldsymbol{w}^* \in \argmin_{\boldsymbol{w}} \ell(\boldsymbol{w})$.
We point out that the final term is symmetric up to permutation of $\boldsymbol{\Theta}$, enabling the replacement of the sum by multiplication with $K$. When $\gamma$ and $K$ are known, $p_{0} = \frac{1}{1 + \gamma}$ and $p_{K} = \frac{1}{K} \frac{\gamma}{1 + \gamma}$ under our constraints. Without loss of generality, we let $\boldsymbol{\theta}_1, \ldots, \boldsymbol{\theta}_{K} \sim p(\boldsymbol{\theta})$ and  $\boldsymbol{x} \sim p(\boldsymbol{x} \mid \boldsymbol{\theta}_{K})$. An empirical estimate of the loss on $B$ samples is therefore
\begin{equation}
\label{eqn:nrec-empirical-loss}
\begin{aligned}
    \hat{\ell}_{\gamma, K}(\boldsymbol{w})
    &\coloneqq - \frac{1}{B}
    \Bigg[ \frac{1}{1 + \gamma} \sum_{b=1}^{B} \log q_{\boldsymbol{w}} \left(y = 0 \mid \boldsymbol{\Theta}^{(b)}, \boldsymbol{x}^{(b)} \right) \\
    &\phantom{\approx - \frac{1}{B} \Bigg[ } + \frac{\gamma}{1 + \gamma} \sum_{b'=1}^{B} \log q_{\boldsymbol{w}} \left(y = K \mid \boldsymbol{\Theta}^{(b')}, \boldsymbol{x}^{(b')} \right) \Bigg].
\end{aligned}
\end{equation}

In the first term, the classifier sees a completely independently drawn sample of $\boldsymbol{x}$ and $\boldsymbol{\Theta}$ while $\boldsymbol{\theta}_K$ is drawn jointly with $\boldsymbol{x}$ in the second term. In both terms, the classifier considers $K$ choices. In practice, we bootstrap both $\boldsymbol{\theta}_1^{(b)}, \ldots, \boldsymbol{\theta}_{K}^{(b)}$ and $\boldsymbol{\theta}_1^{(b')}, \ldots, \boldsymbol{\theta}_{K-1}^{(b')}$ from the same mini-batch and compare them to the same $\boldsymbol{x}$, similarly to \textsc{nre-a}\xspace and \textsc{nre-b}\xspace. Proof of the above is in Appendix~\ref{apndx:proof}.

\paragraph{Recovering \textsc{nre-a}\xspace and \textsc{nre-b}\xspace}
\textsc{nre-c}\xspace is general because specific hyperparameter settings recover \textsc{nre-a}\xspace and \textsc{nre-b}\xspace. To recover \textsc{nre-a}\xspace one should set $\gamma = 1$ and $K=1$ in \eqref{eqn:nrec-empirical-loss} yielding

\begin{equation}
\begin{aligned}
    \hat{\ell}_{1, 1}(\boldsymbol{w})
    &= - \frac{1}{2 B}
    \Bigg[
    \sum_{b=1}^{B} \log \frac{1}{1 + \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b)}, \boldsymbol{x}^{(b)})}
    + \sum_{b'=1}^{B}  \log \frac{\exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b')}, \boldsymbol{x}^{(b')})}{1 + \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b')}, \boldsymbol{x}^{(b')})}
    \Bigg] \\
    &= - \frac{1}{2B} \left[- \sum_{b=1}^{B} \log \left(
        1 - \sigma \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b)}, \boldsymbol{x}^{(b)})
    \right)
    + \sum_{b'=1}^{B}\log \left(
        \sigma \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}^{(b')}, \boldsymbol{x}^{(b')})
    \right) \right]
\end{aligned}
\end{equation}
where we dropped the lower index. Recovering \textsc{nre-b}\xspace requires taking the limit $\gamma \to \infty$ in the loss function. In that case, the first term goes to zero, and second term converges to the softmax function.
\looseness=-1

\begin{equation}
\begin{aligned}
    \hat\ell_{\infty, K}(\boldsymbol{w})
    &= \lim_{\gamma \to \infty} \hat\ell_{\gamma, K}(\boldsymbol{w})
    &= - \frac{1}{B}
    \Bigg[ \sum_{b'=1}^{B} \log
    \frac{\exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}_k,\boldsymbol{x}))}{\sum_{i=1}^{K} \exp \circ h_{\boldsymbol{w}}(\boldsymbol{\theta}_i,\boldsymbol{x})}
    \Bigg]
\end{aligned}
\end{equation}

is determined by substitution into \eqref{eqn:nrec-empirical-loss}. Both equations are obviously the same as their counterparts.

\paragraph{Estimating a normalized posterior}
In the limit of infinite data and infinite neural network capacity (width, depth) the optimal classifier trained using \textsc{nre-c}\xspace (with $\gamma \in \mathbb{R}^{+}$) satisfies the equality:
\begin{align}
    h_{\boldsymbol{w}^*}(\boldsymbol{\theta},\boldsymbol{x}) & = \log \frac{p(\boldsymbol{\theta}\mid\boldsymbol{x})}{p(\boldsymbol{\theta})}.
\end{align}
In particular, we have that the following normalizing constant is trivial:
\begin{align}
  Z(\boldsymbol{x}) & := \int \exp\left(h_{\boldsymbol{w}^*}(\boldsymbol{\theta},\boldsymbol{x})\right) p(\boldsymbol{\theta})\,d\boldsymbol{\theta} = \int p(\boldsymbol{\theta}\mid\boldsymbol{x})\,d\boldsymbol{\theta} =  1.
\end{align}
This is a result of Lemma~\ref{lem:optimal=log-ratio} in Appendix~\ref{apndx:proof}. However, practitioners never operate in this setting, rather they use finite sample sizes and neural networks with limited capacity that are optimized locally. The non-optimal function $\exp(h_{\boldsymbol{w}}(\boldsymbol{\theta},\boldsymbol{x}))$ does not have a direct interpretation as a ratio of probability distributions, rather as the function to weigh the prior $p(\boldsymbol{\theta})$ to approximate the unnormalized posterior. In other words, we find the following approximation for the posterior $p(\boldsymbol{\theta}\mid\boldsymbol{x})$:
\begin{align}
   \label{eqn:normalizing-constant}
   p_{\boldsymbol{w}}(\boldsymbol{\theta}\mid\boldsymbol{x}) &:= \frac{\exp(h_{\boldsymbol{w}}(\boldsymbol{\theta},\boldsymbol{x}))}{Z_{\boldsymbol{w}}(\boldsymbol{x})} p(\boldsymbol{\theta}), &
  Z_{\boldsymbol{w}}(\boldsymbol{x}) & := \int \exp\left(h_{\boldsymbol{w}}(\boldsymbol{\theta},\boldsymbol{x})\right) p(\boldsymbol{\theta})\, d\boldsymbol{\theta},
\end{align}
where in general the normalizing constant is not trivial, i.e.\ $Z_{\boldsymbol{w}}(\boldsymbol{x}) \neq 1$. As stated above, the \textsc{nre-c}\xspace (and \textsc{nre-a}\xspace) objective encourages $Z_{\boldsymbol{w}}(\boldsymbol{x})$ to converge to $1$. This is in sharp contrast to  \textsc{nre-b}\xspace, where even at optimum with an unrestricted function class a non-trivial $\boldsymbol{x}$-dependent bias term can appear.

There is no restriction on how pathological the \textsc{nre-b}\xspace bias $c_{\boldsymbol{w}}(\boldsymbol{x})$ can be. Consider a minimizer of \eqref{eqn:nreb-loss}, the \textsc{nre-b}\xspace loss function, $h_{\boldsymbol{w}^{\ast}} + c_{\boldsymbol{w}^{\ast}}(\boldsymbol{x})$. Adding any function $d(\boldsymbol{x})$ cancels out in the fraction and is also a minimizer of \eqref{eqn:nreb-loss}. This freedom complicates any numerical computation of the normalizing constant and renders the importance sampling diagnostic from Section~\ref{sec:diagnostic} generally inapplicable. We report Monte Carlo estimates of $Z_{\boldsymbol{w}}(\boldsymbol{x})$ on a test problem across hyperparameters in Figure~\ref{fig:partition-function}.

\subsection{Measuring performance \& ratio estimator diagnostics}
\label{sec:diagnostic}
\textsc{sbi}\xspace is difficult to verify because, for many use cases, the practitioner cannot compare surrogate $p_{\boldsymbol{w}}(\boldsymbol{\theta} \mid \boldsymbol{x})$ to the intractable ground truth $p(\boldsymbol{\theta} \mid \boldsymbol{x})$. Incongruous with the practical use case for \textsc{sbi}\xspace, much of the literature has focused on measuring the similarity between surrogate and posterior using two-samples tests on tractable problems. For comparison with literature, we first reference a two-sample exactness metric which requires a tractable posterior. We then discuss diagnostics which do not require samples from $p(\boldsymbol{\theta} \mid \boldsymbol{x})$, commenting on the relevance for each \textsc{nre}\xspace algorithm with empirical results. Further, we find that a known variational bound to the mutual information is tractable to estimate within \textsc{sbi}\xspace, that it bounds the average Kullback-Leibler divergence between surrogate and posterior, and propose to use it for model comparison on intractable inference tasks.

\paragraph{Comparing to a tractable posterior with estimates of exactness}
Assessments of approximate posterior quality are available when samples can be drawn from both the posterior $\boldsymbol{\theta} \sim p(\boldsymbol{\theta} \mid \boldsymbol{x})$ and the approximation $\boldsymbol{\theta} \sim q(\boldsymbol{\theta} \mid \boldsymbol{x})$. In the deep learning-based \textsc{sbi}\xspace literature, exactness is measured as a function of computational cost, usually simulator calls. We investigate this with \textsc{nre-c}\xspace in Section~\ref{sec:benchmark}.

Based on the recommendations of \citet{sbibm} our experimental results are measured using the Classifier Two-Sample Test (C2ST) \cite{friedman2003multivariate, lehmann2005testing, lopez2017revisiting}. A classifier is trained to distinguish samples from either the surrogate or the ground truth posterior. An average classification probability on holdout data of 1.0 implies that samples from each distribution are easily identified; 0.5 implies either the distributions are the same or the classifier does not have the capacity to distinguish them.

\paragraph{Importance sampling diagnostic}
\begin{figure}[hbt]
\centering
    \begin{subfigure}[b]{0.3\textwidth}
        \includegraphics[width=\textwidth]{figures/importance-diagnostic-posterior}
        \caption{Posteriors}
        \end{subfigure}
    \hfill
        \begin{subfigure}[b]{0.3\textwidth}
        \includegraphics[width=\textwidth]{figures/importance-diagnostic-likelihoods}
        \caption{Likelihood and reweighted marginal generative models.}
        \end{subfigure}
    \hfill
        \begin{subfigure}[b]{0.3\textwidth}
        \includegraphics[width=\textwidth]{figures/importance-diagnostic-roc}
        \caption{Receiver operating characteristic.}
    \end{subfigure}
    \caption{
        The figures visualize the importance sampling diagnostic on ratio estimators trained using \textsc{nre-b}\xspace and \textsc{nre-c}\xspace. (a) Both methods produce satisfactory posterior estimates that agree with $p(\boldsymbol{\theta} \mid \boldsymbol{x})$. (b) $p(\boldsymbol{x} \mid \boldsymbol{\theta})$ is shown along with $p(\boldsymbol{x})$ samples weighted by \textsc{nre-a}\xspace $\exp \circ f_{\boldsymbol{w}}(\boldsymbol{\theta}, \boldsymbol{x})$ and \textsc{nre-b}\xspace $\exp \circ g_{\boldsymbol{w}}(\boldsymbol{\theta}, \boldsymbol{x})$. Each plot corresponds to a different $\boldsymbol{\theta}$. Despite high posterior accuracy, the \textsc{nre-b}\xspace estimates are distinct from $p(\boldsymbol{x} \mid \boldsymbol{\theta})$. (c) Two classifier's \textsc{roc}\xspace curves, each trained to distinguish $p(\boldsymbol{x} \mid \boldsymbol{\theta})$ samples from $p(\boldsymbol{x})$ samples weighted by the corresponding \textsc{nre}\xspace's $\hat{r}$ estimate. The classifier failed to distinguish likelihood samples from the \textsc{nre-c}\xspace weighted data samples, but successfully identified \textsc{nre-b}\xspace weighted samples. \textsc{nre-b}\xspace accurately approximates the posterior, but fails the diagnostic. \textsc{nre-c}\xspace produces an accurate posterior surrogate and passes the diagnostic.
    }
    \label{fig:importance-sampling-diagnostic}
\end{figure}
An accurate likelihood-to-evidence weight transforms the data distribution into the likelihood by ${p(\boldsymbol{x} \mid \boldsymbol{\theta}) = p(\boldsymbol{x}) r(\boldsymbol{x} \mid \boldsymbol{\theta})}$. Since \textsc{nre}\xspace necessitates simulator access, we can test the ratio estimator by training a classifier to distinguish unweighted $p(\boldsymbol{x} \mid \boldsymbol{\theta})$ samples from weighted $p(\boldsymbol{x}) \hat{r}(\boldsymbol{x} \mid \boldsymbol{\theta})$ samples, where $\hat{r}$ implies an estimate. Indistinguishability between samples implies either that the approximate ratio is accurate for parameter $\boldsymbol{\theta}$ or that the classifier does not have sufficient power to find predictive features. Issues with classification power can be detected by assessing the classifier's ability to distinguish $p(\boldsymbol{x})$ from $p(\boldsymbol{x} \mid \boldsymbol{\theta})$. The performance can be visualized in a receiver operating curve (\textsc{roc}\xspace) or measured by the area under the curve (\textsc{roc}\xspace \textsc{auc}\xspace). This diagnostic has been used for ratio estimators before \cite{Cranmer2015, Hermans2019} but it comes from training models under covariate shift \cite{shimodaira2000improving}. It is particularly appealing because it does not require samples from $p(\boldsymbol{\theta} \mid \boldsymbol{x})$.
\looseness=-1

\citet{Durkan2020} do not mention this diagnostic in their paper, but due to its intrinsic bias \textsc{nre-b}\xspace does not fulfill the identity necessary for this diagnostic to hold at optimum. The unknown factor that depends on $\boldsymbol{x}$ implies ${p(\boldsymbol{x} \mid \boldsymbol{\theta}) \neq p(\boldsymbol{x}) \exp\circ g_{\boldsymbol{w}}(\boldsymbol{x} \mid \boldsymbol{\theta})}$. We provide empirical evidence of this issue in Figure~\ref{fig:importance-sampling-diagnostic}. Although \textsc{nre-b}\xspace accurately approximates the true posterior, it demonstrably fails the diagnostic. Given the limited options for verification of \textsc{sbi}\xspace results, this presents a major problem by significantly limiting the trustworthiness of \textsc{nre-b}\xspace on any problem with an intractable posterior. In Appendix~\ref{apndx:proof}, we show that the unrestricted \textsc{nre-b}\xspace-specific $c_{\boldsymbol{w}}(\boldsymbol{x})$ bias means approximating $p(\boldsymbol{x} \mid \boldsymbol{\theta})$ with normalized importance weights will not solve the issue.

\paragraph{Mutual information bound}
Selecting the surrogate model most-similar to the target posterior remains intractable without access to $p(\boldsymbol{\theta} \mid \boldsymbol{x}_o)$. Nevertheless, practitioners must decide which surrogate should approximate the posterior across training and hyperparameter search. Unfortunately, the validation losses between different versions of \textsc{nre}\xspace and different $K$ and $\gamma$ settings are not comparable. A good heuristic is to choose the model which minimizes the Kullback-Leibler divergence \emph{on average} over possible data $p(\boldsymbol{x})$.
In Appendix~\ref{apndx:mutual-information}, we prove the relationship between $I(\boldsymbol{\theta}; \boldsymbol{x})$, the \emph{mutual information} with respect to $p(\boldsymbol{\theta}, \boldsymbol{x})$, our models' variational bound $I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$, and the average \textsc{kld}\xspace
\begin{equation}
\label{eqn:average-kld-is-mi}
    \mathbb{E}_{p(\boldsymbol{x})} \left[ \textsc{kld}\xspace(p(\boldsymbol{\theta} \mid \boldsymbol{x}) \, \Vert \, p_{\boldsymbol{w}}(\boldsymbol{\theta} \mid \boldsymbol{x})) \right]
    = I(\boldsymbol{\theta}; \boldsymbol{x}) - I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x}),
\end{equation}
\begin{equation}
\label{eqn:define-mi-0-in-paper}
    I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})
    \coloneqq \mathbb{E}_{p(\boldsymbol{\theta}, \boldsymbol{x})} \left[ \log \hat{r}(\boldsymbol{x} \mid \boldsymbol{\theta}) \right] - \mathbb{E}_{p(\boldsymbol{x})} \left[ \log \mathbb{E}_{p(\boldsymbol{\theta})} [\hat{r}(\boldsymbol{x} \mid \boldsymbol{\theta})] \right].
\end{equation}
The non-negativity of all terms in \eqref{eqn:average-kld-is-mi} implies $I(\boldsymbol{\theta}; \boldsymbol{x}) \geq I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$; that means the model which minimizes $-I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$ best satisfies our heuristic. We propose to approximate $-I_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$ with Monte Carlo using held-out data as a metric for model selection during training and across hyperparameters. The expectation values average over $p(\boldsymbol{\theta}, \boldsymbol{x})$, $p(\boldsymbol{\theta})$, and $p(\boldsymbol{x})$. We can sample from all of these distributions in the \textsc{sbi}\xspace context. Since the second term in \eqref{eqn:define-mi-0-in-paper} computes the average log partition function, our metric can compare \textsc{nre-b}\xspace-based surrogates to \textsc{nre-c}\xspace-based ones.
However, the metric comes with the normal challenges of estimating the log partition function which can be very expensive or high variance in some cases.
We go into more depth in Appendix~\ref{apndx:mutual-information}, including mentioning the relevance to Neural Posterior Estimation \cite{papamakarios2016fast, lueckmann2017flexible, greenberg2019automatic, Durkan2020}. While the application to \textsc{sbi}\xspace is novel, bounds on the mutual information have been broadly investigated for contrastive representation learning and mutual information estimation \cite{gutmann2010noise, gutmann2012noise, gutmann2022statistical, belghazi2018mine, van2018representation, poole2019variational}.

\paragraph{Empirical expected coverage probability}
For a candidate distribution to qualify as the posterior, integrating over data must return the prior. A measurement that follows from calibration to the prior is called expected coverage probability. Expected coverage probability can be estimated with samples from $p(\boldsymbol{\theta}, \boldsymbol{x})$ and any amortized \textsc{sbi}\xspace method. Although important, ability to compute this metric does not distinguish \textsc{nre-c}\xspace. We refer the interested reader to \citet{hermans2021averting}. We note that popular sequential techniques generally render this diagnostic inapplicable, with exceptions \cite{miller2020simulation, miller2021truncated, cole2021fast}.

\section{Experiments}
\label{sec:experiments}
We perform experiments in three settings to
measure the exactness of surrogate models under various hyperparameters and training regimes. Section~\ref{sec:unlimited-joint} aims to identify whether data or architecture is the bottleneck in accuracy by drawing from the joint with every mini-batch. The next experiments aim to determine how to optimally extract inference information given a limited amount of simulation data. Section~\ref{sec:unlimited-prior} leverages a cheap prior by drawing new contrastive parameters with every mini-batch. Finally, Section~\ref{sec:benchmark} applies the commonplace training regime for deep learning-based \textsc{sbi}\xspace literature of fixed data and bootstrapped contrastive parameters. In this setting, we also benchmark \textsc{nre-c}\xspace on ten inference tasks from \citet{sbibm} using our recommended hyperparameters.

On all hyperparameter searches we consider three simulators from the simulation-based inference benchmark, namely SLCP, Two Moons, and Gaussian Mixture \cite{sbibm}. SLCP is a difficult inference task which has a simple likelihood and complex posterior \cite{greenberg2019automatic, papamakarios2019sequential}. Parameters are five dimensional and the data are four samples from a two dimensional random variable. Two Moons is two dimensional with a crescent-shaped bimodal posterior. Finally, the Gaussian Mixture data draws from a mixture of two, two-dimensional normal distributions with extremely different covariances \cite{sisson2018handbook, Beaumont2009-gl, Toni2009-fd, simola2020adaptive}.

Our surrogate models are parameterized by one of these architectures: \emph{Small NN} is like the benchmark with 50 hidden units and two residual blocks. \emph{Large NN} has 128 hidden units and three residual blocks. We use batch normalization, unlike the benchmark. We compare their performance on a grid of $\gamma$ and $K$ values. We report post-training results using the C2ST, and mid-training validation loss for insight into convergence rate. We generally use residual networks \cite{resnet} with batch normalization \cite{ioffe2015batch} and train them using adam \cite{kingma2015adam}. We also run $\textsc{nre-b}\xspace$ with the same architecture for comparison. To compare with \textsc{nre-b}\xspace we set the number of total contrastive parameters equal.

What does fair comparison mean in our experimentation? We compare models across fixed number of gradient steps. This implies that models with more classes, i.e., greater $K$, evaluate the classifier-in-training on more pairs at a given training step than a model with fewer classes. An alternative which we do \emph{not} apply in this paper: Vary the number of gradient steps, holding the number of pair evaluations fixed, i.e., a model with higher $K$ sees the same number of pairs as a model with lower $K$ but the model with lower $K$ has been updated more times. We leave this analysis for future work.

\subsection{Behavior with unlimited data}
\label{sec:unlimited-joint}
What is responsible for inaccuracies of the surrogate model when training has saturated? (a) the amount of training data (b) the flexibility of the model? In this section we provide new simulation and parameter data with every mini-batch and train until saturation, thereby eliminating the possibility of (a). The newly drawn mini-batch parameters $\boldsymbol{\Theta}$ are bootstrapped for the contrastive pairs.

The setting is similar to \textsc{rej-abc}\xspace where simulations are drawn until the posterior has converged. The results of this study will provide a baseline to compare with limited-data results and help us understand how the deep learning architecture's limitations are affected by our introduced hyperparameters. The results are reported in the top row of Figure~\ref{fig:c2st-average}.

The trend is that increasing the number of contrastive examples helps \textsc{nre-b}\xspace and \textsc{nre-c}\xspace.
There is not a clear trend with respect to $\gamma$ in the C2ST. Study of the detailed validation losses in Appendix~\ref{apndx:experimental-details} reveals that high $\gamma$ is associated with higher variance validation loss during training, but here it does not seem to strongly affect convergence rate.
Saturation is reached before the maximum number of epochs have elapsed and Large NN converges to a better result. Although obfuscated in the averaged results in Figure~\ref{fig:c2st-average}, the result is obvious on SLCP, the most difficult inference task that provides an opportunity for a more flexible architecture to improve.

We argue that generally the performance bottleneck has to do with the number of contrastive parameters, network flexibility, or training details rather than the amount of data. Despite saturating validation losses, the C2ST improves based on these factors. Intuitively a more flexible model could continue to extract information from this unlimited set of jointly drawn $(\boldsymbol{\theta}, \boldsymbol{x})$ pairs. This has some consequence because network flexibility may limit performance in the benchmark case. Appendix~\ref{apndx:experimental-details} contains more detailed results.

\subsection{Leveraging fast priors (drawing theta)}
\label{sec:unlimited-prior}
In practice, the simulator is often slow but one can often draw from the prior at the same pace as training a neural network for one mini-batch. Our goal is to understand how our hyperparameter are affected by this setting and whether it is valuable to use this technique in practice. To our knowledge, this training regime has not been explored in the deep learning-based \textsc{sbi}\xspace literature.

Initially we draw a fixed set of around 20,000 samples $(\boldsymbol{\theta}, \boldsymbol{x}) \sim p(\boldsymbol{\theta}, \boldsymbol{x})$. For every mini-batch during training, we sample all necessary contrastive parameters from the prior. For each term in $\hat{\ell}_{\gamma, K}$ we take the same batch of contrastive parameters and reshuffle them, thereby equalizing the number of samples seen by \textsc{nre-b}\xspace and \textsc{nre-c}\xspace, up to bootstrap. The averaged C2ST results from the three inference tasks are reported in the middle row of Figure~\ref{fig:c2st-average}.

The resulting estimators are markedly less sensitive to the number of contrastive parameters than in Section~\ref{sec:unlimited-joint}. The number of network parameters has a positive effect, although again this is best seen on the hardest task, SLCP, in the Appendix~\ref{apndx:experimental-details}. Empirically, lower values of $\gamma$ improve performance with Large NN; however, there is still high variance and this result may be from noise. In this setting, $\gamma \neq 1$ has a very small negative effect on convergence rate compared to unity, as seen in the validation loss plots in Appendix~\ref{apndx:experimental-details}. Since training has saturated, we claim this result implies drawing contrastive parameters from the prior helps extract the maximum amount of information from fixed simulation data $\boldsymbol{x}$--without being strongly affected by the number of contrastive parameters $K$.

\begin{figure}[htb]
\centering
    \begin{subfigure}[b]{0.218\textwidth}
        \includegraphics[width=\textwidth]{figures/mi0-slcp-nreb-summary.pdf}
        \caption{\textsc{nre-b}\xspace}
        \end{subfigure}
    \hfill
        \begin{subfigure}[b]{0.775\textwidth}
        \includegraphics[width=\textwidth]{figures/mi0-slcp-nrec-summary.pdf}
        \caption{\textsc{nre-c}\xspace}
        \end{subfigure}
    \caption{
        Our proposed metric, a negative bound on the mutual information $-\hat{I}_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$, for the SLCP task estimated over the validation set versus training epochs using (a) \textsc{nre-b}\xspace and (b) \textsc{nre-c}\xspace with various values of $\gamma$ and $K$, a Large NN architecture, and fixed training data. The bound permits visualization of the convergence rates and pairwise comparison across models. When $K$ is fixed, small $\gamma$ slows convergences but leads to better optima than large $\gamma$. When $\gamma$ is fixed, small $K$ slows convergence but leads to better optima than high $K$. Unlike when computing the validation loss directly, this metric does not exhibit a $\gamma$ and $K$ dependent bias as seen in Figures~\ref{fig:grid-validation-joint}, \ref{fig:grid-validation-prior}, and \ref{fig:grid-validation-bench}.
    }
    \label{fig:mutual-information-0-slcp-summary}
\end{figure}
\subsection{Simulation-based inference benchmark}
\label{sec:benchmark}

In this section, we assume the traditional literature setting of limited simulation and prior budget. Once we've selected hyperparameters based on a grid search of a subset of the \textsc{sbi}\xspace benchmark, we perform the entire benchmark with those hyperparameters.

\begin{wrapfigure}[38]{R}{0.5\textwidth}
    \centering
    \vspace{-13pt}
    \includegraphics[width=0.5\textwidth]{figures/rel-avg}
    \caption{%
    Exactness of \textsc{nre}\xspace posterior surrogates is computed for various contrastive parameter counts, $\gamma$ values, and architectures on an \emph{average} of three tasks from the \textsc{sbi}\xspace benchmark \cite{sbibm}. C2ST assigns 1.0 to inaccurate and 0.5 to accurate posterior approximation.
    (top) $p(\boldsymbol{\theta}, \boldsymbol{x})$ was sampled at every mini-batch during training. The accuracy strongly depends on $K$.
    (mid) A fixed number of dependent $(\boldsymbol{\theta}, \boldsymbol{x})$ pairs were drawn, but $p(\boldsymbol{\theta})$ was sampled at every mini-batch during training. In this regime, $K$ has a smaller effect.
    (bot) The training data is completely fixed. Contrastive parameters are drawn in a bootstrap from the mini-batch. On the problems with fixed simulation data $\boldsymbol{x}$, higher $K$ improves accuracy and small $\gamma$ with larger architectures slightly improves performance. The effects of the architecture are more clearly seen on difficult problems like SLCP in Appendix~\ref{apndx:experimental-details}.
    }
    \label{fig:c2st-average}
\end{wrapfigure}

\paragraph{Hyperparameter search}
To compare to the previous two sections, the amount of training data was fixed such that each epoch was comparable to Section~\ref{sec:unlimited-joint}, this amounts to about 20,000 samples. We first inspect the C2ST results in the last row of Figure~\ref{fig:c2st-average}. The sensitivity to $K$ appears less than in Section~\ref{sec:unlimited-joint} but more than in Section~\ref{sec:unlimited-prior}. Smaller $\gamma$ slightly improves Large NN's performance but the noise makes this result uncertain. The larger network performs better, see Appendix~\ref{apndx:experimental-details}.

We do a Monte Carlo estimate of our proposed metric, denoted $-\hat{I}_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$, to show performance on the SLCP task as a function of epochs in Figure~\ref{fig:mutual-information-0-slcp-summary}. For both \textsc{nre-b}\xspace and \textsc{nre-c}\xspace, increasing $K$ tends to positively affect the convergence rate.
Generally, with fixed $\gamma$ and high $K$, peak performance is negatively impacted, in contrast to results from Figure~\ref{fig:c2st-average} using the C2ST. Generally, with fixed $K$, small $\gamma$ slows convergences but leads to better optima than large $\gamma$. When $\gamma$ is fixed, small $K$ slows convergence but leads to better optima than high $K$. See Figure~\ref{fig:mutual-information-0-slcp}.

Our take-away is that $\gamma \in [0.1, 10]$ has a small effect, otherwise learning can become unstable or very slow. $\gamma=1$ has a good compromise on high convergence rate without sacrificing too much performance. Generally, we saw improved performance on the C2ST by increasing contrastive parameters $K$. However, larger $K$ also increases the magnitude of the normalizing constant, and extreme values reduced the performance on $-\hat{I}_{\boldsymbol{w}}^{(0)}(\boldsymbol{\theta}; \boldsymbol{x})$, see Appendix~\ref{apndx:mutual-information}. Since the benchmark compares C2ST, we optimized our architecture based on that metric. Due to bootstrapping parameters from the batched $\boldsymbol{\theta}$, the maximum $K$ is $B/2$ without reusing any $\boldsymbol{\theta}$ to compare with $\boldsymbol{x}$. This is considering both terms in our loss function.

\paragraph{Benchmark}
We performed the benchmark tasks from \citet{sbibm} with \textsc{nre-c}\xspace. Our architecture and hyperparameters deviate based on the results of the last paragraph. We applied the largest number of computationally practical contrastive parameters, namely $K=99$, and set $\gamma = 1$. Since we identified potential architecture bottlenecks in Sections~\ref{sec:unlimited-joint} and \ref{sec:benchmark}, we used Large NN for our results. The averaged results are visible in Table~\ref{table:sbibm}. Sequential estimates are prepended with an \textsc{s}, all other estimates are amortized. We trained a five amortized estimators per task with five seeds, then we predicted the posterior for each of the ten observations $\boldsymbol{x}_1, \ldots, \boldsymbol{x}_{10}$. \citet{sbibm} trained an estimator for every $\boldsymbol{x}$ with a new seed. The reported C2ST for our method is averaged across ten posteriors per seed with five seeds, i.e., fifty C2ST computations per task. The other methods are averaged across one posterior per seed with ten seeds, i.e., ten C2ST computations per task.

\textsc{nre-c}\xspace performed better than \textsc{nre-b}\xspace at all budgets and generally well among the amortized algorithms. The high performance at $10^{5}$ simulation budget implies that our design has high capacity and may scale well. Sequential methods naturally lead to high exactness since they are tailored to the specific observation $\boldsymbol{x}_o$, but they come with drawbacks since they cannot practically perform empirical expected coverage tests \cite{hermans2021averting, miller2021truncated} nor bound the mutual information, as discussed in Section~\ref{sec:diagnostic}. Detailed results per task can be found in Appendix~\ref{apndx:experimental-details}.

\begin{wraptable}{R}{.5\textwidth}
    \centering
    \vspace{-12pt}
    \caption{Averaged C2ST for various budgets on the simulation-based inference benchmark \cite{sbibm}. An \textsc{s} implies a sequential i.e. non-amortized method.}
    \label{table:sbibm}
    \begin{tabular}{lrrr}
    \toprule
     & \multicolumn{3}{r}{C2ST} \\
    Simulation budget & $10^3$ & $10^4$ & $10^5$ \\
    Algorithm &  &  &  \\
    \midrule
    \textsc{nre-c}\xspace (ours) & 0.853 & 0.762 & 0.680 \\
    \midrule
    \textsc{rej-abc}\xspace & 0.965 & 0.920 & 0.871 \\
    \textsc{nle}\xspace & 0.826 & 0.753 & 0.723 \\
    \textsc{npe}\xspace & 0.838 & 0.736 & 0.654 \\
    \textsc{nre}\xspace (\textsc{nre-b}\xspace) & 0.867 & 0.811 & 0.762 \\
    \textsc{smc-abc}\xspace & 0.948 & 0.873 & 0.816 \\
    \textsc{snle}\xspace & 0.783 & 0.704 & 0.655 \\
    \textsc{snpe}\xspace & 0.796 & 0.677 & 0.615 \\
    \textsc{snre}\xspace (\textsc{snre-b}\xspace) & 0.788 & 0.703 & 0.610 \\
    \bottomrule
    \end{tabular}
\end{wraptable}

\section{Conclusion}
\label{sec:conclusion}
We introduced a generalization of \textsc{nre-a}\xspace and \textsc{nre-b}\xspace which discourages the \textsc{nre-b}\xspace specific $c_{\boldsymbol{w}}(\boldsymbol{x})$ bias term and produces an unbiased estimate of the likelihood-to-evidence ratio at optimum. This property has implications for the importance sampling diagnostic, which is critical to practitioners.
We suggested using a variational bound on the mutual information as a model-selection metric that could replace the C2ST averaged over several observations. It is significantly more practical since it does not require access to the ground truth posterior; however, it remains a lower bound so we don't know the overall quality of our estimator. Further discussions and derivations related to the mutual information are available in Appendix~\ref{apndx:mutual-information}.

In the context of \textsc{nre-c}\xspace, we found that increasing the number of contrastive parameters $K$ improved the C2ST in most cases, at the price of increasing the normalizing constant. Setting $\gamma \approx 1$ was generally the fastest to saturate and lower values seemed to slightly improve performance with the Large NN and fixed simulation data.
These results indicate that better performance can be achieved by using \textsc{nre-c}\xspace with these hyperparameters than the other version of \textsc{nre}\xspace. When both training parameters $\boldsymbol{\theta}$ and simulation data $\boldsymbol{x}$ are unlimited, the architecture size plays an important role in the quality of surrogate model.
We tried drawing contrastive parameters from the prior, a commonly available practical use-case, and found that it damped the effect of $K$ on the C2ST. According to the C2ST, for highest accuracy with a fixed budget set $K$ large, up to the mini-batch size divided by 2, and $\gamma \approx 1.0$ for the best convergence rate. When the prior is sampled for every mini-batch, applying a higher $K$ may be less valuable. In situations where the normalizing constant should be very close to unity, i.e., for practitioners who want to run diagnostics, drawing $\boldsymbol{\theta}$ from the prior with every mini-batch and using a smaller $K$ and $\gamma$ is best.

\paragraph{Broader Impact} The societal implications of \textsc{sbi}\xspace are similar to other inference methods. The methods are primarily scientific and generally lead to interesting discoveries; however, one must be careful to use an accurate generative model and to carefully test empirical results. Model mismatch and untested inference can lead to incorrect conclusions. That is why we emphasize the importance of the diagnostics in our paper. This nuance can be missed by practitioners doing inference in any field; however, special care should be taken when producing inference results that may be used for making decisions in areas like predicting hidden variables that describe human behavior or determining what factors are responsible for climate change. This list is non-comprehensive and not specific to \textsc{sbi}\xspace.

\begin{ack}
	We thank Noemi Anau Montel, James Alvey, Kosio Karchev for reading; Teodora Pandeva for mathematics consultation; Maud Canisius for design consultation; Marco Federici and Gilles Louppe for discussions. This work uses \texttt{numpy} \cite{harris2020array}, \texttt{scipy} \cite{2020SciPy-NMeth}, \texttt{seaborn} \cite{Waskom2021}, \texttt{matplotlib} \cite{Hunter:2007}, \texttt{pandas} \cite{reback2020pandas, mckinney-proc-scipy-2010}, \texttt{pytorch} \cite{pytorch}, and \texttt{jupyter} \cite{jupyter}.

    We want to thank the DAS-5 computing cluster for access to their TitanX GPUs. DAS-5 is funded by the NWO/NCF (the Netherlands Organization for Scientific Research). We thank SURF (\href{www.surf.nl}{www.surf.nl}) for the support in using the Lisa Compute Cluster.

	Benjamin Kurt Miller is funded by the University of Amsterdam Faculty of Science (FNWI), Informatics Institute (IvI), and the Institute of Physics (IoP).
	We received funding from the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (Grant agreement No. 864035 -- UnDark).
\end{ack}

\clearpage

\bibliographystyle{abbrvnat}


\section{Introduction}
The interplay between topology and correlations represents
a highly fruitful area in condensed matter physics.
However,
exploring unconventional states of matter requires
identifying systems where electronic correlations, topology
and electronic dispersions can be realistically controlled.
In this line, twisted van der Waals materials\cite{Cao2018super,Yankowitz2019,Liu2014,Lu2019,Rickhaus2018,PhysRevB.88.121408,Liao2020,Shimazaki2020}
provide
a powerful solid state platform to realize exotic quantum phenomena.
The tunability
of twisted van der Waals materials
stems from the emergence of a band structure that
can be controlled
by the twist between different two-dimensional materials.\cite{PhysRevLett.99.256802,PhysRevB.82.121407}
In particular,
the different quantum states in twisted graphene systems stem
from the possibility of controlling the ratio between kinetic and
interaction terms. In twisted graphene bilayers,
such tunability allowed to realize
superconducting,\cite{Cao2018super,Yankowitz2019,Lu2019} correlated
insulators, topological networks,\cite{Rickhaus2018,PhysRevB.88.121408}
Chern insulators\cite{Serlin2019}
and quasicrystals.\cite{Ahn2018,PhysRevB.99.165430,Yu2019,2020arXiv200110427P}
As a result, current experimental efforts are focusing on
exploring new twisted van der Waals materials,
with the aim of finding platforms that allow for an
even higher degree of control.\cite{2020arXiv200411340C,2020arXiv200411353P}

From the quantum engineering point of view,
applying a perpendicular bias between layers\cite{PhysRevLett.99.216802} provides a versatile
way of tuning correlated states in twisted graphene multilayers.
This has been demonstrated
in paradigmatic examples of correlated states in
twisted tetralayers (double bilayer)\cite{2019arXiv190308130L,Shen2020}
and twisted trilayers
(monolayer/bilayer).\cite{2020arXiv200411340C,2020arXiv200411353P}
Moreover, interlayer bias is known to generate internal
valley currents in twisted graphene bilayers,\cite{Rickhaus2018,PhysRevB.88.121408,PhysRevLett.121.146801,PhysRevLett.123.096802}
creating
topological networks at low angles\cite{Rickhaus2018,PhysRevB.88.121408,PhysRevLett.121.146801} and
generating valley fluxes in flat bands regimes.\cite{PhysRevLett.123.096802}
This interplay of correlations and topology in twisted graphene
multilayers makes these materials a powerful platform to explore
exotic states of matter\cite{PhysRevLett.124.106803,2020arXiv200409522L,2019arXiv191211469R,2019arXiv191209634L}
in a realistically feasible manner.

From the theoretical point of view,
electronic structure calculation of
twisted graphene bilayer conducted with
real-space tight-binding models\cite{PhysRevB.82.121407} or continuum Dirac descriptions\cite{PhysRevLett.99.256802}
capture
the fundamental features of the electronic dispersion.
Nevertheless, internal coordinates optimization can quantitatively modify the electronic
dispersion.\cite{PhysRevB.98.224102,PhysRevB.98.195432,PhysRevB.96.075311,2019arXiv191012805L,Brihuega2017,PhysRevB.100.075416,PhysRevX.9.041010,Rickhaus2019} Well known examples of this are the growth of AB/BA regions
in twisted bilayers.\cite{PhysRevB.96.075311,PhysRevB.100.155426}
It is important to note that studying twisted graphene multilayers
from first-principles represents a remarkable
challenge, due to the large amount of atoms
present in a unit cell.

\begin{figure}[t!]
\centering
    \includegraphics[width=\columnwidth]{fig1.pdf}

\caption{(a,b) Sketch of the structure a twisted graphene trilayer,
in which the top and bottom layers are aligned, and twisted with respect the middle
one.
Panel (c) shows the color map of a DFT-based relaxed 2$\times$2 supercell of twisted trilayer. Colored pattern corresponds to the difference $\Delta = z-z_{av}$ between the z coordinate of an atom in surface layer and the average coordinate $z_{av}$ of all atoms in the same layer. AA stacking regions are colored with darker tones whereas AB/BA stacking regions are colored with lighter tones. Superimposed is the unit cell of the trilayer.
Panel (d) shows the histogram shows the number of occurrences of each distance to the average deviation value. Tiny deviations from flatness occur in the central layer.
Calculations performed with DFT.
}
\label{figureContour}
\label{fig:sketch}
\end{figure}

Here, by combining first-principles calculations and low-energy models, we
show that twisted graphene trilayers\cite{PhysRevB.87.125414}
host flat bands whose bandwidth
can be controlled electrically. We address the
impact of an interlayer bias both from first-principles
and effective models, showing that van Hove singularities can be
merged electrically.
We show that associated with the interlayer bias, bulk valley currents emerge,
that are impacted by the existence of chemical impurities in the system.
We finally address the superconducting states in these doped trilayers,
showing that the non-uniform superfluid density has an impact
in high-energy bands. Our manuscript is organized as follows,
in Sec. \ref{sec:bands} we show the electronic structure
of twisted graphene trilayer both from first-principles
and low-energy models, in Sec. \ref{sec:field} we explore
in detail the impact of an electric field,
in Sec. \ref{sec:imp} we explore the effect of chemical impurities,
and
in Sec. \ref{sec:sc} we address
the impact of an emergent non-uniform superfluid density. Finally,
in Sec. \ref{sec:sum} we summarize our conclusions.

\section{Electronic structure of twisted trilayer graphene}
\label{sec:bands}

The electronic structure of twisted graphene trilayers\cite{PhysRevB.87.125414,2019arXiv190712338L,PhysRevB.100.085109,2019arXiv190700952C,PhysRevLett.123.026402,2020arXiv200501258P} shows
different features in comparison
with twisted graphene bilayers.\cite{PhysRevB.82.121407,Bistritzer2011}
The electronic structure of small "magic" angle tBLG
features four flat bands lying around the Fermi level, with
band splitting at the $\Gamma$ point of the Brillouin zone of
the emergent moir{\'e} superlattice.\cite{PhysRevB.82.121407,Bistritzer2011}
Density functional theory (DFT) calculations have shown
consistent results with those effective models
in twisted graphene multilayers, yet
quantitative modifications are observed when including
relaxation of the atomic coordinates\cite{PhysRevB.96.075311,PhysRevB.98.235137,2019arXiv191012805L,PhysRevB.99.195419,2020arXiv200414323C}
and crystal-field effects\cite{Rickhaus2019,Haddadi2020}.
Therefore, to benchmark the electronic properties
of twisted graphene multilayers, it is essential to start from a
correct description that takes into account the
geometric corrugation and ab-initio
electrostatics of the moir{\'e} system.

We consider a twisted trilayer structure in which
the upper and lower layers are aligned,
and the middle one is twisted with an angle $\theta$ with respect to those (Fig. \ref{fig:sketch}ab).
In particular, in the following we will consider a twisted trilayer whose middle layer has
a twisting angle of 1.9$^\circ$ with respect to the external layers.
A twisted multilayer like this can be created with standard tear, rotate and stack
techniques.\cite{Kim2016rotate}
The system of 5514 atoms
is fully relaxed allowing for lateral and vertical displacement of the C atoms in the structure.
Figure \ref{figureContour}c shows the color map of one of the two equivalent external layer relaxation. The color scheme
represents the vertical variation of each C atom at the surface with respect to the average deviation within each
layer. The darker areas indicate a displacement of atoms out of the surface
which occurs predominantly in the AA stacking
region. The histogram of Fig. \ref{figureContour}d shows
that the number of atoms in the upper layer whose vertical coordinate is above the
average is twice as large as the atoms displaced in the opposite direction
below the average.

The first-principles
electronic band diagram of the fully relaxed structure of
$\theta=1.9^\circ$ TTG
is shown in Fig. \ref{figEfields}.
In contrast with twisted bilayer graphene, two
highly dispersive bands coexist with
four low-dispersive bands grouped at the Fermi energy.
The Dirac-like crossing above the charge neutrality point
leads to a small charge transfer between the flat and dispersive bands
even at half filling.
Additional flattening of the localized states can be induced by application of an external
electric field
perpendicular to the TTG surface. A field of 0.03 eV/\AA\ reduces
the dispersion of all electronic states in the vicinity of
the Fermi energy without inducing any inter-band charge transfer.
Increasing the strength up to 0.25 eV/\AA\ a disruption
of the linear bands is observed and a hybridization of the flat bands
with neighboring state increase their dispersion. It is observed that
the application of an interlayer bias generates Dirac crossings above and
below charge neutrality, besides a variety of anticrossings
(Fig. \ref{figEfields}).

The first-principles calculations above show that
the electronic structure of twisted trilayer graphene show strong differences
with the one of twisted graphene bilayers.
In particular,
a highly dispersive set of
bands coexists
with the nearly
flat bands at charge neutrality. In order to explore more in detail
the physics and twisted trilayers, in the following we will
exploit a low-energy model.
We find that the tight binding model qualitatively reproduces
the important features of the band structure without including
relaxations and additional charge transfer effects.
Therefore, for the sake of simplicity
we now take an unrelaxed structure for our tight binding calculations.
We take a single orbital per carbon atom, yielding
a tight-binding Hamiltonian of the form

\def{\bold{r}}{{\bold{r}}}

\begin{eqnarray}
        \mathcal{H}_0 = -t \sum_{\langle i, j \rangle,s} c_{i,s}^\dagger c_{j,s} -
        \sum_{i,j,s} \bar t_{\perp} ({\bold{r}}_i,{\bold{r}}_j) c_{i,s}^\dagger c_{j,s},
\end{eqnarray}
with
$
\bar t_{\perp}(\bold r_i,\bold r_j) =
t_{\perp}
\frac{(z_i - z_j)^2 }{|\bold r_i - \bold r_j|^2}
e^{-\beta (|\bold r_i - \bold r_j|-d)}$, where
$d$ is the interlayer distance and $\beta$
controls the decay of the interlayer hopping.
As a reference, for twisted graphene multialyers
$t \approx 3$ eV and $t_\perp \approx 0.15 t$.
\footnote{At low energies the spectra is invariant upon rescaling
of the interlayer coupling, which allows to explore effective smaller angles
with smaller unit cells.\cite{PhysRevB.98.195101,PhysRevLett.119.107201}
Our calculations are performed with a resscaled $t_\perp = 0.4 t$.}
Similar real-space models were used to study
a variety of twisted graphene multilayers,\cite{PhysRevB.82.121407,PhysRevB.87.125414,PhysRevB.82.121407,2019arXiv191101347C,PhysRevB.91.035441}
providing a simple formalism to study the effect of dopands
and impurities.\cite{PhysRevB.99.245118,PhysRevMaterials.3.084003}
However, in contrast
to continuum models,\cite{PhysRevB.86.155449,PhysRevLett.99.256802,Bistritzer2011}
measuring of valley related quantities
with a real-space based formalism is non-trivial.

\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{fig2.pdf}

\caption{Electronic band diagrams of twisted trilayer graphene. Localized states at the Fermi energy coexist with highly dispersive states bands forming a Dirac cone in the conduction band. Applying a small electric field perpendicular to the TTG a reduction of all bands dispersion is induced. A stronger electric field removes the Dirac cones and induces hybridization between electronic states.
Calculations performed with DFT}
\label{figEfields}
\label{fig:dftf}
\end{figure}

Twisted graphene multilayers own an approximate
symmetry associated to the valley quantum number.\cite{PhysRevLett.99.256802}
Valley physics in both
DFT calculations and tight-binding are emergent
symmetries, in the sense that valley are not
easily defined in terms of real-space chemical orbitals.
This limitation can be overcome by defining the so-called
valley operator\cite{PhysRevLett.120.086603,PhysRevLett.121.146801,PhysRevB.99.245118}
in the tight-binding description.
With the valley operator
the expectation value of the valley can be computed in a real
space representation.\cite{Wolf2019,2020arXiv200305163M}
The details of the
valley operator\cite{PhysRevLett.120.086603,PhysRevLett.121.146801,PhysRevB.99.245118}
are given in Sec. \ref{sec:valley}, in the following
we will take as starting point the valley operator $\mathcal{V}_z$
With the previous operator, we can compute the valley flavor of each
eigenstate of the twisted trilayer supercell
within the real-space formalism as
$\langle \mathcal{V}_z \rangle = \langle \Psi | \mathcal{V}_z | \Psi \rangle | $.
It is worth noting that
this operator can be easily defined in the tight-binding
basis but not in the DFT basis. We finally note that the
valley operator $\mathcal{V}_z$ in the twisted moire system
will show the valley flavor in the original Brillouin zone
of graphene, not the mini-Brillouin zone of the twisted system.

With the previous formalism, we now compute the electronic structure
of the low-energy model at the same angle $\theta = 1.9^\circ$ as
in the first-principles calculations (Fig. \ref{fig:tbnof}).
We note that, besides some additional splittings observed
in the first-principles calculations
(Fig. \ref{fig:dftf}), the band-structure
obtained with the low-energy model (Fig. \ref{fig:tbnof}ab)
gives comparable results.
As shown in Fig. \ref{fig:tbnof}a, in the absence of an
interlayer bias the system shows nearly flat bands
coexisting with highly dispersive states. Whereas the
nearly flat bands are degenerate in valleys in the
path shown, the dispersive states belong to different
valleys in different parts of the Brillouin zone.
It is also observed that the dispersive
Dirac cones are slightly displaced
from charge neutrality, as observed in the first
principles results.

\begin{figure}[t!]
\centering
    \includegraphics[width=\columnwidth]{fig3.pdf}

\caption{
(a) Band structure of a twisted trilayer at
angle of 1.9$^\circ$, without (a) and with (b) interlayer bias,
showing analogous dispersion to those found by first
principles in Fig. \ref{fig:dftf}.
Panel (c) shows the band structure
at a lower angle of
$\theta = 1.6^\circ$, showing a flattening of the bands.
Panel (d) shows the local density of states at the Fermi energy,
showing a triangular supercell pattern.
Note that part of the k-path is different
in comparison with Fig. \ref{fig:dftf}.
Calculations performed with tight-binding.
}
\label{fig:tbnof}
\end{figure}

We now move on to considered the effect of the itnerlayer bias.
An interlayer bias can be easily
included in the low-energy
tight-binding model by means of

\begin{equation}
    \mathcal{H}_V = V/(2d)\sum_{i,s} z_i c^\dagger_{i,s} c_{i,s}
\end{equation}
where $z_i$ is the $z$-position of site $i$, $d$ is the interlayer distance
and $V$ is the strength of the interlayer bias. The full Hamiltonian
in the presence of interlayer bias is thus $\mathcal{H} = \mathcal{H}_0 + \mathcal{H}_V$.
When an interlayer bias is turned on (Fig. \ref{fig:tbnof}b),
the valley degeneracy in the $G-K$ path is lifted, and the
dispersive bands show two crossings above and below the
nearly flat bands, analogously to the
first-principles results. It is also observed that the nearly
flat bands are slightly modified by the interlayer bias.
As we will show below, such effect becomes stronger for even smaller
angles.

A last interesting point is related with the
the localization of the different states in the supercell.
This can be characterized by means
of the inverse participation ratio (IPR), defined as
$IPR = \sum_i |\Psi_k (i)|^4$. Large values of the IPR correspond to states localized in the
moir{\'e} supercell, whereas low values correspond to delocalized in the moir{\'e}
supercell.
Focusing now in a structure with
$\theta = 1.6^\circ$ twisting angle,
it is clearly observed that the nearly flat bands show a substantially
higher degree of localization than
the dispersive bands (Fig. \ref{fig:tbnof}c).
In particular, the flat band states are associated to an emergent
triangular lattice in the supercell, as shown in the local
density of states (LDOS) of Fig. \ref{fig:tbnof}d. In the
following we will see how these flat band states can
be electrically controlled, and how the bias
creates bulk valley currents associated with the valley
splittings.

\section{Band flattening and valley currents by an interlayer bias}
\label{sec:field}

We now move on to systematically
analyze the effect of an interlayer bias in the
twisted graphene trilayer, and in particular its effect on the low
energy density of states.
Electric biases in twisted graphene multilayers are known to give rise
to valley Hall currents,\cite{PhysRevB.74.161403,PhysRevLett.100.036804,Zhang2013,PhysRevB.75.155424,PhysRevB.88.121408,PhysRevLett.121.146801,Rickhaus2018}
and represent an effective knob to control
the low-energy electronic structure.\cite{PhysRevLett.99.216802}
In the following we focus on the structure at $\theta = 1.6^\circ$.
As shown in Fig. \ref{fig:tbf}a, the application of an interlayer
bias merges the original two van Hove singularities of zero bias in
a single one, dramatically enhancing the low-energy density of states.
Interestingly,
this merging of van Hove singularities is similar to the evolution with the twist angle in
twisted
bilayers,\cite{PhysRevB.82.121407,Bistritzer2011,PhysRevLett.108.216802}
with the key difference
that in the present case the merging is electrical,
turning the process highly controllable in-situ.

\begin{figure}[t!]
\centering
    \includegraphics[width=\columnwidth]{fig4.pdf}

\caption{
(a) Evolution of the density of states as a function of the interlayer bias
	for $\theta = 1.6^\circ$.
Panel (b) shows the band structure at a finite electric field depicting the
localization of each state of each state.
Panel (c) shows the valley polarization of the low-energy states, whereas
panel (d) their layer polarization.
Panel (e) shows the spatial localization of the nearly flat bands.
Panel (f) shows the local valley flux at one electron per unit cell.
The spatial maps in (e,f) are summed over the three layers.
Calculations performed with tight-binding.
}
\label{fig:tbf}
\end{figure}

We now focus on the flattest regime with finite bias,
when the two van Hove singularities get merged.
Fig. \ref{fig:tbf}bcd shows the band structure of that regime .
It is observed that almost perfectly flat bands are present at the
Fermi energy (Fig. \ref{fig:tbf}bcd), accounting for the
van Hove singularity at charge neutrality (Fig. \ref{fig:tbf}a).
It is also observed that the states remain highly localized
in the unit cell as highlighted by the IPR, whereas at higher energies the
dispersive bands are delocalized in the moir{\'e} unit cell (Fig.
\ref{fig:tbf}b).
It is also observed that the bands show perfect valley polarization
(Fig. \ref{fig:tbf}c), so that the interlayer bias
does create any intervalley scattering in agreement
with continuum models. Given that the
interlayer bias breaks the symmetry between the top
and bottom layers, it is interesting to look at the
layer polarization of the states $\langle \mathcal{L} \rangle$,
defined by the layer polarization operator
$\mathcal{L} = \sum_{i,s} z_i c^\dagger_{i,s} c_{i,s}$,
with $z_i$ the $z$-component of the site $i$.
As shown in Fig. \ref{fig:tbf}d, the Dirac cones
at the $K,K'$ points
above and below the flat bands have a slightly
opposite layer polarization. However, such
layer polarization is reversed at the $\Gamma$ point (Fig. \ref{fig:tbf}d),
highlighting that the states remain highly entangled between
all the layers.

We now
explore how the band flattening is accompanied with the emergence of bulk valley currents.
The emergence of valley currents associated with interlayer biases
is a well known effect in aligned graphene
bilayers\cite{PhysRevB.74.161403,PhysRevLett.100.036804,Zhang2013,PhysRevB.75.155424}
and tiny-angle twisted bilayers.\cite{PhysRevB.88.121408,PhysRevLett.121.146801,Rickhaus2018}
Such valley currents arise due to the emergence of valley
a non-zero local valley Chern numbers,\cite{Ren2016} whose quantization
is associated to the emergent valley conservation.\cite{RevModPhys.81.109}
In twisted systems, the local Chern number is expected to change
from region to region,\cite{PhysRevB.88.121408} due to the locally
modulated Hamiltonian. To address this local
topological property, in the following
we compute the local
valley flux by means of the Berry flux density $\chi_V$.
The real-space valley flux $\chi_V$ defines the valley Chern number as
$C_V = C_K - C_{K'} = \int \chi_V(\mathbf r,\omega) d^2 \mathbf r d\omega$.
The real-space valley flux $\chi_V$ can be computed as\cite{PhysRevB.84.205137,Wolf2019,2020arXiv200305163M}

\begin{align}
    \chi_V(\mathbf r,\omega) = \int \frac{d^2 \mathbf k}{(2\pi)^2} \frac{\epsilon_{\alpha \beta}}{2}
 \langle \mathbf r  | G_V (\partial_{k_\alpha}G_V^{-1}) (\partial_{k_\beta}G_V)| \mathbf r \rangle
\end{align}
where $\epsilon_{\alpha\beta}$ denotes the Levi-Civita tensor, $G_V$ the
valley Green's function $G_V = [\omega - \mathcal{H}_{\mathbf k}+i0^+]^{-1}
\mathcal{V}_z$,
$\mathcal{H}_k$ the Bloch Hamiltonian,
and $\mathcal{V}_z$ the valley polarization operator
of Eq. \ref{eq:valley}.

Fig. \ref{fig:tbf}f shows the spatial profile of the valley flux density
at the Fermi energy $\chi_V(\mathbf r) \equiv \chi_V(\mathbf r,\omega=\epsilon_F)$,
for a chemical potential with one hole per unit cell. It can
be observed clearly that sizable valley currents
appear in regions in the complementary regions to the low-energy
states, similarly to other moir{\'e} systems.\cite{Wolf2019,2020arXiv200305163M}.
Since the emergence of such currents relies on valley conservation,
terms in the system creating intervalley mixing are expected
to substantially impact them,\cite{PhysRevB.88.121408,PhysRevB.86.205414,PhysRevB.93.041407,PhysRevB.94.134201}
as we address in the next section.

\section{Impact of chemical doping}
\label{sec:imp}

In this section we address the modification of the electronic structure under both electrical and
chemical doping. First, we start with the study of the electric doping from first-principles, and
afterwards we move on to consider its effect in the low-energy effective model.

\begin{figure}[t!]
\centering
\includegraphics[width=\columnwidth]{fig5.pdf}

\caption{Electronic band diagrams of three doped TTG configurations. Left panel shows
electrostatic doping with one extra electron. Central panel shows
that doping with a   N atom causes a similar effect but
removes the Dirac point and converts TTG in a tiny-gap semiconductor.
Right panel shows that doping with a
B atom empties some bands at the Fermi level including the highly dispersive bands.
Calculations performed with DFT.
}
\label{figDoping}
\label{fig:dopdft}
\end{figure}

We first address the impact of doping from first-principles.
First-principles DFT calculations are employed to study the modification of the
band diagram of TTG upon doping with both an extra charge and one of the layers with foreign species.
Uniform doping of the TTG is realized by introducing an extra electron which is compensated with an
equally uniform background charge of opposite sign.
This can be realized experimentally
via field effect.
Figure \ref{fig:dopdft} shows that the extra electron fills the flat states and align the Dirac point
with the Fermi energy. Small modifications in the electronic band
structure as a result of new  electrostatic contributions is observed, similarly to other twisted multilayer
systems.\cite{PhysRevB.100.205113,Guinea2018,2020arXiv200305050P,2020arXiv200414784G}

A similar effect can be induced with chemical doping, namely adding one N atom in substitution of a C atom.
The effects of chemical substitution have been extensively studied in graphene,\cite{PhysRevB.93.245420,PhysRevB.83.245403,Liu2011,LopezBezanilla2009,Agnoli2016,Georgakilas2012,PhysRevLett.109.226802,PhysRevLett.104.187201,LopezBezanilla2009_2,PhysRevLett.97.236802}
including twisted bilayers.\cite{PhysRevMaterials.3.084003,2019arXiv190809555Y}
Figure \ref{fig:dopdft} shows the effect on the bands structure of one N atom in one of the surface layers.
The extra charge supplied the N atom shifts the chemical potential, increasing
the filling of the flat bands. Major
difference with respect to the electrostatic doping is the opening
of a meV large band gap as a result of the
symmetry breaking imposed by the impurity which
also induces a mixing of the linear states with the less
dispersive states. These new anticrossings are
due to the intervalley scattering created by the chemical
impurity, and they do not appear
in the case of electrostatic doping.
A similar hybridization of electronic states is also
observed when one electron is removed by means of
doping with a B atom, as shown in Fig. \ref{fig:dopdft}. The chemical potential shift occurs in the opposite
direction and the filling of flat band states decreases.

\begin{figure}[t!]
\centering
    \includegraphics[width=\columnwidth]{fig6.pdf}

\caption{
(a) Band structure in the presence of a N impurity, and (b) valley current at charge neutrality.
It is observed that, although the presence of a dopand creates a non-zero intervalley scattering,
the low-energy bands remain relatively flat. The emergence of intervalley scattering can be
seen in the depletion of the bulk valley currents in (b).
The valley flux in (b) is summed over the three layers.
Calculations performed with tight-binding.
}
\label{fig:doptb}
\end{figure}

We now consider the effect of a chemical
impurity in the tight-binding model. We model the addition
of a chemical impurity by adding to the Hamiltonian $\mathcal{H}_D = w \sum_s c^\dagger_{i,s} c_{i,s}$,
where $i$ in the site that has been chemically replaced.
We take $w=-2t$, which is the typical energy scale expected for a $N$
dopand, and yields results comparable with the first principle calculations.
For the sake of concreteness we will focus on the case with interlayer bias,
so our full Hamiltonian will be $\mathcal{H} = \mathcal{H}_0 + \mathcal{H}_V + \mathcal{H}_D$.
With the previous Hamiltonian,
we now compute the electronic band structure and project each eigenstate onto the valley operator.
The result is shown in Fig. \ref{fig:doptb}a, where we see that small anticrossings
appear as in the first-principles calculations. The valley projection
clearly shows that such anticrossings are
associated with intervalley mixing. The effect of the impurity in terms of intervalley mixing can
also be readily seen in the bulk valley currents. In particular, the existence of the impurity
is expected to strongly perturb the original valley fluxes in the unit cell,
depleting the local value of the valley Chern number.
This is verified
in Fig. \ref{fig:doptb}b, where we observe that the original valley fluxes are impacted
by the presence of the impurity.

Interestingly, despite the effect in terms of intervalley mixing, the low-energy bands
remain relatively flat, retaining their associated large density of states. This suggests
that correlated states can still appear in chemical doped twisted trilayers. It is worth
to emphasize that despite this large density of states, correlated phases relying on
valley coherent states will be strongly suppressed due to the impurity-induced
intervalley mixing. In particular, valley ferromagnet states, and valley triplet
superconducting states will be depleted due to chemical dopands. Nevertheless,
conventional spin-singlet valley-singlet states are not affected by the presence
of intervalley scattering. Motivated by this, in the next section we address the
emergence of spin/valley singlet superconductivity, and show
how the superfluid density impacts the high energy dispersive states.

\section{Superconducting state}
\label{sec:sc}

The large density of states close to charge neutrality suggests that the
twisted graphene trilayer can have superconducting instabilities,
similarly to twisted bilayers and tetra-layers.
As shown above, both with electrostatic and chemical doping the
system shows divergent density of state close to charge neutrality.
For the sake of concreteness we will now focus on the
electrostatically doped system, yet we have verified that our results
remain qualitatively unchanged with chemical doping.

An emergent superconducting state
is associated with a Fermi surface instability, yet its effect can give rise
to second order perturbations above the Fermi energy.
In order to understand the potential impact at high-energies,
we first briefly analyze the structure of the low-energy states.
This can be done by comparing the spatial distribution of the low-energy states
with respect to the Fermi surface states. In particular, we define the projection
over the Fermi surface states as
$
    \mathcal{R} = \int d^2 \mathbf r |\Psi(\mathbf r)|^2 \rho_F (\mathbf r)
$
with $\rho_F (\mathbf r)$ the local density at the Fermi energy
$\rho_F (\mathbf r) \sim \int \langle \mathbf r | \delta (E_F - H) | \mathbf r \rangle$.
The quantity $\mathcal{R}$ allows to qualitatively distinguish which states are localized
in the same region at the Fermi surface states.

By computing the Fermi surface state projector $\mathcal{R}$, it observed
that the states of the flat band are localized in similar regions.
In contrast, as one departs from the
Fermi energy, the states start to delocalize to other regions of the unit
cell (Fig. \ref{fig:figsc}a). This highlights the different orbital
nature of the nearly flat bands (in purple) and low-energy
dispersive bands (in green). Importantly, and in strike
contrast with twisted graphene
bilayers, the flat bands of this system are not decoupled from the dispersive
states, suggesting that the superconducting states of this system will have
a genuine multiorbital nature.

Given the unavoidable entanglement between the flat and dispersive bands,
an effective model description of this twisted trilayer
cannot be easily performed. Therefore, in the following
we will study the emergent superconducting state by exploiting the
full atomistic model, including all the bands in our calculation.
A variety of mechanisms
have been suggested to give rise to attractive interactions in these
systems, including phonon,\cite{PhysRevB.98.220504,PhysRevLett.121.257001,PhysRevLett.122.257002}
Coulomb\cite{PhysRevLett.122.026801}
and magnon fluctuations.\cite{PhysRevLett.121.087001,PhysRevLett.121.217001}
For the sake of concreteness, we will focus on
effective local attractive interactions,
as employed in other twisted graphene systems
\cite{PhysRevLett.97.230404,PhysRevLett.98.146801,PhysRevLett.100.246808,PhysRevB.101.060505,Heikkil2019,PhysRevB.98.220504,PhysRevLett.121.257001,PhysRevLett.123.237002}

\begin{figure}[t!]
\centering
    \includegraphics[width=\columnwidth]{fig7.pdf}

\caption{
	Normal state band structure (a) in the absence of interactions,
	with the color denoting
	the projection of the states on the density
	of states at the Fermi energy.
	Upon introducing interactions,
	a non-zero superfluid density
	appears as shown in
	panel (b).
	Panel (c) shows the associated
	electronic dispersion (c) to the selfconsistent
	solution of (b), where
	we have projected the eigenstates over the
	the electron sector. Panel
	(d) shows the electronic dispersion
	projected in the electron sector
	for a uniform
	superfluid density, highlighting
	that the high energy anticrossings stem
	from the non-uniform pairing.
	Calculations performed with tight-binding.
}
\label{fig:figsc}
\end{figure}

\begin{equation}
\mathcal{H}_I = - \sum_i
	g c_{i,\uparrow}^\dagger c_{i,\uparrow}
	c_{i,\downarrow}^\dagger c_{i,\downarrow}
\end{equation}

that we solve at the mean field level
$
	\mathcal{H}_I^{MF} = -g\sum_i
	\langle c_{i,\uparrow}^\dagger
	c_{i,\downarrow}^\dagger \rangle
	c_{i,\uparrow}
	c_{i,\downarrow}
	+ h.c.
$
where $\langle c_{i,\uparrow}^\dagger
	c_{i,\downarrow}^\dagger \rangle$ is computed self-consistently
	for the full Hamiltonian
	$\mathcal{H} = \mathcal{H}_0 + \mathcal{H}_V + \mathcal{H}^{MF}_I$,
	that we solve
	with the Bogoliubov-de Gennes (BdG) formalism.
The local attractive interaction $g$ will give rise to
a net superfluid density, with a moir{\'e} momentum structure of
$s$-wave symmetry.
In the following we take $g=t$, and we verified
that our results remain qualitatively similar with
smaller interaction strengths.
We note that although interactions are local,
the would lead to a non-trivial multiorbital structural in the moir{\'e}
orbital space.

By solving the previous selfconsistent problem, we find that the
superfluid density is non-uniform in the moir{\'e} unit cell
(Fig. \ref{fig:figsc}b), stemming from the
non-uniformity of the low-energy states.
By projecting the BdG eigenstates in the electron sector
via the electron projector $P_e$,
we observe
that the net non-uniform superfluid density
$\Delta_{\uparrow\downarrow} (\mathbf{r}_i) \sim \langle c^\dagger_{i,\uparrow} c^\dagger_{i,\downarrow} \rangle$
gives rise to a full gap in the Brillouin zone
in the superconducting state. This phenomenology is similar to the
one found in twisted graphene bilayers.
More interestingly, besides the gap opening at the
chemical potential, anticrossings appear in high energy bands when the
selfconsistent pairing is included (Fig. \ref{fig:figsc}c). It is worth
to emphasize that,
in the presence of a uniform pairing artificially imposed, the high
energy anticrossings disappear, leading only to the gap opening at charge
neutrality (Fig. \ref{fig:figsc}d).
The emergence of gap openings away from charge neutrality is associated to
the intrinsically multiorbital nature of the superconducting state, and
stems from the non-unitarity of the superconducting matrix
in orbital subspace.\cite{PhysRevResearch.1.033107,2019arXiv190711602Z}
Interestingly, this shows that signatures of the superconducting state
can be obtained by analyzing the system away from the chemical potential,
and could provide powerful spectroscopic
signatures\cite{PhysRevB.97.165414,2020arXiv200202289L} of the superconducting state.

\section{Conclusions}
\label{sec:sum}

By combining first-principles calculation and low-energy effective models,
we have shown that twisted graphene trilayers realize tunable electronic systems.
In particular, is was shown that nearly perfect flat bands can be
electrically controlled, that coexist with highly dispersive
states. Interestingly, such
electric flattening of the bands is accompanied by the emergence of
bulk valley currents. We have found both from first
principles and low-energy calculations that chemical doping
does not destroy the flat bands, yet it substantially impacts
the bulk valley currents. This suggests that chemical doping
of twisted graphene trilayers could provide an intrinsic way of providing
the necessary electronic doping required for the emergence of
a superconducting state. We finally demonstrated that
an emergent superconducting state
would give rise to spectroscopic changes int he high energy bands,
associated to the non-uniform superfluid density. Our results
highlight the rich physics of twisted graphene trilayers,
and provide a starting point to explore the interplay between
flat bands, correlation and dispersive states in twisted
graphene multilayers.

\section*{Acknowledgments}
Los Alamos National Laboratory is managed by Triad National Security, LLC, for the National Nuclear Security Administration of the U.S. Department of Energy under Contract No. 89233218CNA000001.
This work was supported by the U.S. DOE Office of Basic Energy Sciences Program (E3B5).
A.L.-B. acknowledges the computing resources provided on Bebop, the high-performance
computing clusters operated by the Laboratory Computing Resource Center at Argonne
National Laboratory. J.L.L. thanks T. Wolf, G. Blatter, M. Sigrist,
O. Zilberberg, W. Chen, T. Neupert, A. Ramires and T. Heikkil\"a
for fruitful discussions. J.L.L. acknowledges the computational
resources provided by the Aalto Science-IT project.

\section*{Appendix}