constant integer, the average sample complexity of gradient, would be in the
order of $\mathcal{O}(T^{4})$. As
$T\geq\frac{\left[16(C_{g}+C_{H})C^{3}+C_{g}+C_{H}\right]^{1/2}}{\epsilon^{1/2}}$,
the average sample complexity of gradient is at least

$\mathcal{O}\left(\frac{2\sigma_{1}^{2}\left[16(C_{g}+C_{H})C^{3}+C_{g}+C_{H}\right]^{2}}{\epsilon^{2}}\right)$

in order to get $\delta_{T}\leq\epsilon$.

Moreover, for the average sample complexity of Hessian
$\mathbb{E}\left[\sum_{t=1}^{T}n_{H}^{t}\right]$, we have:

$\begin{split}&\sum_{\mod(t,S)=0}\frac{(8e\cdot\log
d)^{2}\cdot\sigma_{2,1}^{2}}{(t)^{-2}}+\sum_{\mod(t,S)\neq
0}\frac{2(8e\cdot\log d)^{2}\cdot L^{\prime
2}_{2}S\mathbb{E}[\|\mathbf{\Delta}_{t-1}\|^{2}]}{(\lceil t/S\rceil
S)^{-2}}\\\ &=\sum_{k=1}^{\lceil T/S\rceil}{(8e\cdot\log
d)^{2}\cdot\sigma_{2,1}^{2}}{(kS)^{2}}+\sum_{k=1}^{\lceil
T/S\rceil}\sum_{t=(k-1)S+1}^{kS}{2(8e\cdot\log d)^{2}\cdot L^{\prime
2}_{2}S}{(kS)^{2}}\mathbb{E}[\|\mathbf{\Delta}_{t-1}\|^{2}]\\\
&\leq{(8e\cdot\log d)^{2}\cdot\sigma_{2,1}^{2}S^{2}}\sum_{k=1}^{\lceil
T/S\rceil}{k^{2}}+\\\ &\sum_{k=1}^{\lceil T/S\rceil}{2(8e\cdot\log d)^{2}\cdot
L^{\prime
2}_{2}S}{(kS)^{2}}C^{\prime\prime}\left[S^{1/3}\left(\frac{16(C_{g}+C_{H})C^{3}}{(kS)^{2}}+\frac{C_{g}+C_{H}}{(kS)^{2}}\right)^{2/3}+\frac{(c_{g}+c_{H})}{(kS)^{2}}\right]\\\
&\leq{(8e\cdot\log d)^{2}\cdot\sigma_{2,1}^{2}}{T^{3}}+\\\ &2(8e\cdot\log
d)^{2}\cdot L^{\prime
2}_{2}S^{1/3}(T)^{3-4/3}C^{\prime\prime}\left([16(C_{g}+C_{H})C^{3}+(C_{g}+C_{H})]^{2/3}+\frac{c_{g}+c_{H}}{(kS)^{2(1-2/3)}}\right)\end{split}$
(135)

By setting $S=\lfloor\frac{T}{q}\rfloor$ where $q$ is a positive integer
constant, the average sample complexity of Hessian, would be in the order of
$\mathcal{O}(T^{2})$. As $T=\mathcal{O}(\epsilon^{-\frac{1}{2}})$, the average
sample complexity of Hessian is in the order of
$\mathcal{O}\left(\frac{1}{\epsilon}\right)$.

∎

###  A.4 Proof of RL results

Two important classes of policies: 1) Scalar-action, fixed-variance Gaussian
policy:

$\displaystyle\pi_{\theta}(a|s)=\frac{1}{\sigma\sqrt{2\pi}}\exp\left\\{-\frac{1}{2}\left(\frac{a-\theta^{T}\phi(s)}{\sigma}\right)^{2}\right\\}.$
(136)

2) Softmax tabular policy:

$\displaystyle\pi_{\theta}(a|s)=\frac{\exp(\theta_{s,a})}{\sum_{a^{\prime}\in\mathcal{A}}\exp(\theta_{s,a^{\prime}})}$
(137)

where $\theta=\\{\theta_{s,a}\\}_{(s,a)\in\mathcal{S}\times\mathcal{A}}$ are
parameters of the policy.

####  A.4.1 Discussion on gradient dominance property with $\alpha=1$ for
Policies

######  Assumption 13 (Fisher-non-degenerate.).

For all $\theta\in\mathbb{R}^{d}$, there exists $\mu_{F}>0$ such that the
Fisher information matrix $F_{\rho}(\theta)$ induced by policy $\pi_{\theta}$
and initial distribution $\rho$ satisfies

$\displaystyle F_{\rho}(\theta):=\mathbb{E}_{(s,a)\sim
v_{\rho}^{\pi_{\theta}}}[\nabla_{\theta}\log\pi_{\theta}(a|s)\nabla_{\theta}\log\pi_{\theta}(a|s)^{T}]\succeq\mu_{F}I_{d\times
d},$ (138)

where
$v_{\rho}^{\pi_{\theta}}(s,a):=(1-\gamma)\mathbb{E}_{s_{0}\sim\rho}\sum_{t=0}^{\infty}\gamma^{t}\mathbb{P}(s_{t}=s,a_{t}=a|s_{0},\pi_{\theta})$
is the state-action visitation measure.

This assumption is commonly used in the literature [20, 9]. The Fisher-non-
degenerate setting implicitly guarantees that the agent is able to explore the
state-action space under the considered policy class. Similar conditions of
the Fisher-non-degeneracy is also required in other global optimum convergence
framework (Assumption 6.5 in [1] on the relative condition number and
Assumption 3 in [6] on the regularity of the parametric model). Assumption 13
is satisfied by a wide families of policies, including the Gaussian policy
(136) and certain neural policy. Without the non-degenerate Fisher information
matrix condition, the global optimum convergence of more general
parameterizations would be hard to analyze without introducing the additional
exploration procedures in the non-tabular setting [9, Sec. 8].

######  Assumption 14.

For all $\theta\in\mathbb{R}^{d}$, there exists $\epsilon_{bias}>0$ such that
the transferred compatible function approximation error with $(s,a)\sim
v_{\rho}^{\pi_{\theta^{*}}}$ satisfies

$\displaystyle\mathbb{E}[(A^{\pi_{\theta}}(s,a)-(1-\gamma){u^{*}}^{T}\nabla_{\theta}\log\pi_{\theta}(a|s))^{2}]\leq\epsilon_{bias},$
(139)

where $v_{\rho}^{\pi_{\theta^{*}}}$ is the state-action distribution induced
by an optimal policy, and $u^{*}=(F_{\rho}(\theta))^{\dagger}\nabla
J(\theta)$.

This is also a common assumption [35, 1, 41, 20, 9]. In particular, when
$\pi_{\theta}$ is a soft-max tabular policy (137), $\epsilon_{bias}$ is 0 [9];
Moreover, it can be shown that when $\pi_{\theta}$ is a rich neural policy,
$\epsilon_{bias}$ is small [35].

By Lemma 4.7 in [9], Assumptions 13 and 14 yield relaxed gradient dominance
property with $\alpha=1$ (See Assumption 6). We provide the proof in the
sequel:

From performance difference lemma [14], we have

$\displaystyle\mathbb{E}_{(s,a)\sim
v_{\rho}^{\pi_{\theta^{*}}}}[A^{\pi_{\theta_{t}}}(s,a)]=(1-\gamma)(J_{\rho}(\theta^{*})-J_{\rho}(\theta_{t})),$
(140)

and from Assumption 14 and Jensen’s inequality

$\displaystyle\mathbb{E}[A^{\pi_{\theta}}(s,a)-(1-\gamma){u^{*}}^{T}\nabla_{\theta}\log\pi_{\theta}(a|s)]\leq\sqrt{\epsilon_{bias}}.$
(141)

Substituting (140) into (141), we get

$\displaystyle
J_{\rho}(\theta^{*})-J_{\rho}(\theta_{t})\leq\frac{1}{1-\gamma}\sqrt{\epsilon_{bias}}+\mathbb{E}[{u^{*}}^{T}\nabla_{\theta}\log\pi_{\theta}(a|s)]$
$\displaystyle\leq\frac{1}{1-\gamma}\sqrt{\epsilon_{bias}}+\frac{G_{1}}{\mu_{F}}\|\nabla
J_{\rho}(\theta_{t})\|,$ (142)

where the last inequality comes from following Cauchy-Schwartz inequality

$\displaystyle\mathbb{E}[{u^{*}}^{T}\nabla_{\theta}\log\pi_{\theta}(a|s)]=\mathbb{E}[\nabla
J(\theta)^{T}(F_{\rho}(\theta))^{\dagger}\nabla_{\theta}\log\pi_{\theta}(a|s)]$
$\displaystyle\leq\mathbb{E}\|\nabla
J(\theta)\|_{2}\mathbb{E}\|(F_{\rho}(\theta))^{\dagger}\nabla_{\theta}\log\pi_{\theta}(a|s)\|_{2}\leq\frac{1}{\mu_{F}}\mathbb{E}\|\nabla_{\theta}\log\pi_{\theta}(a|s)\|_{2}\mathbb{E}\|\nabla
J(\theta)\|_{2}\leq\frac{G_{1}}{\mu_{F}}\mathbb{E}\|\nabla J(\theta)\|_{2}$

where two last inequalities are from Assumptions 13 and 7. Thus, by setting
$\epsilon^{\prime}:=\frac{1}{1-\gamma}\sqrt{\epsilon_{bias}}$ and
$\tau_{J}:=\frac{G}{\mu_{F}}$,

$\displaystyle J^{*}-J(\theta_{t})\leq\tau_{J}\|\nabla
J(\theta_{t})\|_{2}+\epsilon^{\prime}.$ (143)

In the past few years, several work have attempted to establish that PG and
SPG converge to a global optimal point for various classes of policies [43,
44]. For instance, for Fisher non-degenerate policies, Liu et al. [20]
proposed a variance-reduced SPG method that converges to a global optimal
point with sample complexity of $\tilde{\mathcal{O}}(\epsilon^{-3})$. By
considering a momentum term, Ding et al. [9] showed that the sample
complexities of SPG for the soft-max policy and Fisher-non-degenerate policies
are in the order of $\tilde{\mathcal{O}}(\epsilon^{-4.5})$ and
$\tilde{\mathcal{O}}(\epsilon^{-3})$, respectively. In the case of weak
gradient dominant functions with $\alpha=1$ (See Assumption 6), Yuan et al.
[41] showed that the sample complexity of SPG is bounded by
$\tilde{\mathcal{O}}(\epsilon^{-3})$.

Theorem 4. For a policy $\pi_{\theta}$ satisfying Assumptions 7, 8, and the
corresponding objective function $J(\theta)$ satisfying Assumption 6, SCRN
outputs the solution $\theta_{T}$ such that
$J^{*}-\mathbb{E}[J(\theta_{T})]\leq\epsilon+\epsilon^{\prime}$ and the sample
complexity (the number of observed state action pairs) is: $T\times
m\times\mathsf{H}=\tilde{\mathcal{O}}(\epsilon^{-2.5})$ for
$\epsilon^{\prime}=0$ and $T\times
m\times\mathsf{H}=\tilde{\mathcal{O}}(\epsilon^{-0.5}\epsilon^{\prime-2})$ for
$\epsilon^{\prime}>0$.

####  A.4.2 Proof of Theorem 4

In order to obtain Lipschitz Hessian for $J(\theta)$, we use the following
lemma in [46].

######  Lemma 13.

Under Assumptions 7, and 8, we have

$\displaystyle\|\nabla^{2}J(\theta)-\nabla^{2}J(\theta^{\prime})\|\leq\tilde{L}\|\theta-\theta^{\prime}\|_{2},$
(144)

where

$\tilde{L}=\frac{R_{max}G_{1}G_{2}}{(1-\gamma)^{2}}+\frac{R_{max}G_{1}^{3}(1+\gamma)}{(1-\gamma)^{3}}+\frac{R_{max}G_{1}}{1-\gamma}\max\left\\{G_{2},\frac{\gamma
G_{1}^{2}}{1-\gamma},\frac{\bar{L}_{2}}{G_{1}},\frac{G_{2}\gamma}{1-\gamma},\frac{G_{1}(1+\gamma)+G_{2}\gamma(1-\gamma)}{1-\gamma^{2}}\right\\}.$

It is good to mention that $\tilde{L}=\mathcal{O}(1/(1-\gamma)^{3})$. Using
(21) for $\alpha=1$ in Lemma 5 for $F(\theta):=-J(\theta)$, we get

$\displaystyle J^{*}-J(\theta_{t+1})-\epsilon^{\prime}\leq\tau_{J}\|\nabla
J(\theta_{t+1})\|$ $\displaystyle\leq
C(J(\theta_{t+1})-J(\theta_{t}))^{2/3}+C_{g}\|\nabla
J(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|+C_{H}\|\nabla^{2}J(\theta_{t})-\hat{\nabla}_{m}^{2}J(\theta_{t})\|^{2},$
(145)

where

$\displaystyle\begin{split}C&=3^{1/3}\tau_{F}\Big{(}\frac{M+\tilde{L}+1}{2}\Big{)}\Big{(}\frac{12}{3M-2\tilde{L}-8}\Big{)}^{2/3}\\\
C_{g}&=2^{2/3}\times
3^{\frac{-2}{3}}\tau_{J}\Big{(}\frac{M+\tilde{L}+1}{2}\Big{)}\Big{(}\frac{12}{3M-2\tilde{L}-8}\Big{)}^{2/3}+\tau_{J}\\\
C_{H}&=2^{-2/3}\times
3^{\frac{-2}{3}}\tau_{J}\Big{(}\frac{M+\tilde{L}+1}{2}\Big{)}\Big{(}\frac{12}{3M-2\tilde{L}-8}\Big{)}^{2/3}+2^{-1}\tau_{J}.\end{split}$
(146)

We know that $\hat{\nabla}_{m}J(\theta_{t})$ and
$\hat{\nabla}^{2}_{m}J(\theta_{t})$ are in fact unbiased estimators of $\nabla
J_{\mathsf{H}}(\theta_{t})$ and $\nabla^{2}J_{\mathsf{H}}(\theta_{t})$. Thus,

$\displaystyle J^{*}-J(\theta_{t+1})\leq
C(J(\theta_{t+1})-J(\theta_{t}))^{2/3}+C_{g}\|\nabla
J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|+2C_{H}\|\nabla^{2}J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}^{2}_{m}J(\theta_{t})\|^{2}+\epsilon^{\prime}$
$\displaystyle+C_{g}\|\nabla J(\theta_{t})-\nabla
J_{\mathsf{H}}(\theta_{t})\|+2C_{H}\|\nabla^{2}J(\theta_{t})-\nabla^{2}J_{\mathsf{H}}(\theta_{t})\|^{2}$
$\displaystyle\leq C(J(\theta_{t+1})-J(\theta_{t}))^{2/3}+C_{g}\|\nabla
J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|+2C_{H}\|\nabla^{2}J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}^{2}_{m}J(\theta_{t})\|^{2}+\epsilon^{\prime}$
$\displaystyle+C_{g}D_{g}\gamma^{H}+2C_{H}D_{H}\gamma^{2\mathsf{H}}.$

where $C_{g}$ and $C_{H}$ is defined in Lemma 3. We take a expectation from
both sides given $\theta_{t}$ and use Jensen’s inequality as follows:

$\displaystyle J^{*}-\mathbb{E}J(\theta_{t+1})\leq
C(\mathbb{E}J(\theta_{t+1})-\mathbb{E}J(\theta_{t}))^{2/3}+C_{g}\mathbb{E}[\|\nabla
J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|]$
$\displaystyle+2C_{H}\mathbb{E}\left[\|\nabla^{2}J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}^{2}_{m}J(\theta_{t})\|^{2}\right]+C_{g}D_{g}\gamma^{\mathsf{H}}+2C_{H}D_{H}\gamma^{2\mathsf{H}}+\epsilon^{\prime}.$
(147)

We use the following lemma to bound the error terms of the gradient estimator
and the Hessian estimator.

######  Lemma 14.

Under Assumption 7, we have

$\begin{split}&\mathbb{E}\|\hat{\nabla}_{m}J(\theta)-\nabla
J_{\mathsf{H}}(\theta)\|^{2}\leq\frac{\mathsf{H}G_{1}^{2}R_{\max}^{2}}{m(1-\gamma)^{2}},\\\
&\mathbb{E}\|\nabla^{2}J_{\mathsf{H}}(\theta)-\hat{\nabla}^{2}_{m}J(\theta)\|^{2}\leq\frac{4e\cdot\max\\{2,\log
d\\}R_{\max}^{2}(\mathsf{H}^{2}G_{1}^{4}+G_{2}^{2})}{m(1-\gamma)^{4}}.\end{split}$
(148)

The proof of Lemma 14 is given in Appendix A.4.3.

Now we have from Jensen’s inequality

$\mathbb{E}\|\nabla
J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|\leq\left(\mathbb{E}\|\nabla
J_{\mathsf{H}}(\theta_{t})-\hat{\nabla}_{m}J(\theta_{t})\|^{2}\right)^{1/2}\leq\frac{\mathsf{H}GR_{\max}}{(1-\gamma)\sqrt{m}}.$

Then we get

$\displaystyle J^{*}-\mathbb{E}J(\theta_{t+1})\leq
C(\mathbb{E}J(\theta_{t+1})-\mathbb{E}J(\theta_{t}))^{2/3}+C_{g}\frac{\mathsf{H}GR_{\max}}{(1-\gamma)\sqrt{m}}+2C_{H}\frac{4e\cdot\max\\{2,\log
d\\}R_{\max}^{2}(F^{2}+G^{4})}{(1-\gamma)^{4}m}$
$\displaystyle+\epsilon^{\prime}+C_{g}D_{g}\gamma^{\mathsf{H}}+2C_{H}D_{H}\gamma^{2\mathsf{H}}.$
(149)

Let define a stationary value for $J^{*}-\mathbb{E}J(\theta_{t})$ as
$P(m,\epsilon^{\prime})$ which is as follows:

$\displaystyle
P(m,\epsilon^{\prime}):=C_{g}\frac{GR_{\max}}{(1-\gamma)^{3/2}\sqrt{m}}+2C_{H}\frac{4e\cdot\max\\{2,\log
d\\}R_{\max}^{2}(F^{2}+G^{4})}{(1-\gamma)^{4}m}+\epsilon^{\prime}+C_{g}D_{g}\gamma^{\mathsf{H}}+2C_{H}D_{H}\gamma^{2\mathsf{H}}.$
(150)

We define:

$\delta_{t}:=\frac{J^{*}-\mathbb{E}J(\theta_{t})-P(m,\epsilon^{\prime})}{C^{3}}.$

Then we get following recursion from (A.4.2):

$\delta_{t+1}\leq(\delta_{t}-\delta_{t+1})^{2/3}.$

We know from (42) in the proof of Theorem 1 that

$\displaystyle\delta_{T}\leq\frac{4}{(DT)^{2}}$ (151)

where $D=\min\\{1/2,2(2^{1/3}-1)\delta^{-1/2}_{0}\\}$. Hence,

$J^{*}-\mathbb{E}J(\theta_{T})-P(m,\epsilon)\leq\frac{4C^{3}}{(DT)^{2}}.$

Therefore, $J^{*}-\mathbb{E}J(\theta_{t})$ converges to $P(m,\epsilon)$ with
the convergence rate of $\mathcal{O}(1/T^{2})$. Let

$\displaystyle
m\geq\max\left\\{\frac{C^{2}_{g}\mathsf{H}^{2}G_{1}^{2}R_{\max}^{2}}{(1-\gamma)^{2}\epsilon^{\prime
2}},\frac{C_{H}4e\cdot\max\\{2,\log
d\\}R^{2}_{\max}(G_{2}^{2}+G_{1}^{4})}{(1-\gamma)^{4}\epsilon^{\prime}}\right\\},$
$\displaystyle\mathsf{H}\geq\max\left\\{\frac{\log\left(\frac{C_{g}D_{g}}{\epsilon^{\prime}}\right)}{\log(1/\gamma)},\frac{\log\left(\frac{2C_{H}D_{H}}{\epsilon^{\prime}}\right)}{2\log(1/\gamma)}\right\\}.$

Then, $\theta_{T}$ satisfies:

$\displaystyle J^{*}-\mathbb{E}J(\theta_{T})\leq 5\epsilon^{\prime}+\epsilon$
(152)

where $T\geq\frac{2C^{3/2}D}{\sqrt{\epsilon}}$. Finally, the number of
observed state-action pairs are required to get (152) is as follows:

$\displaystyle T\times\mathsf{H}\times
m\geq\mathcal{O}\left(\frac{C^{3/2}DC^{2}_{g}G_{1}^{2}R_{\max}^{2}\log^{2}\left(\frac{C_{g}D_{g}}{\epsilon^{\prime}}\right)}{(1-\gamma)^{2}\log^{2}(1/\gamma)\epsilon^{\prime
2}\epsilon^{0.5}}\right)\overset{(a)}{=}\mathcal{O}\left(\frac{G^{2+\frac{7}{6}}_{1}G_{2}^{\frac{7}{6}}R_{\max}^{2+\frac{7}{6}}}{(1-\gamma)^{4+\frac{7}{3}}\epsilon^{\prime
2}\epsilon^{0.5}}\right)$ (153)

where (a) comes from $C=\mathcal{O}(M^{1/3})$ and $C_{g}=\mathcal{O}(M^{1/3})$
and the fact that $M=\mathcal{O}(\tilde{L})$ and from Lemma 13,
$\tilde{L}=\frac{R_{\max}G_{1}G_{2}}{(1-\gamma)^{2}}$.

####  A.4.3 Proof of Lemma 14

In [41, Lemma 4.2], it has been shown that

$\begin{split}\mathbb{E}\|\hat{\nabla}_{m}J(\theta)-\nabla
J_{\mathsf{H}}(\theta)\|^{2}\leq\frac{\mathsf{H}G_{1}^{2}R_{\max}^{2}}{m(1-\gamma)^{2}}.\end{split}$
(154)

We know from [29, Lemma 4.1] that

$\displaystyle\|\hat{\nabla}^{2}(\theta,\tau)\|^{2}\leq\frac{R_{\max}^{2}(\mathsf{H}^{2}G_{1}^{4}+G_{2}^{2})}{(1-\gamma)^{4}},$
(155)

where $\hat{\nabla}^{2}(\theta,\tau):=\nabla\Phi(\theta;\tau)\nabla\log
p(\tau|\pi_{\theta})^{T}+\nabla^{2}\Phi(\theta;\tau)$. We have an upper bound
on the variance of Hessian of value function for each trajectory as follows:

$\mathbb{E}[\|\nabla^{2}J_{\mathsf{H}}(\theta)-\hat{\nabla}^{2}(\theta,\tau)\|^{2}]\leq\mathbb{E}[\|\hat{\nabla}^{2}(\theta,\tau)\|^{2}]\leq\frac{R_{\max}^{2}(\mathsf{H}^{2}G_{1}^{4}+G_{2}^{2})}{(1-\gamma)^{4}}.$

Denote
$\sigma^{2}_{2,1}:=\frac{R_{\max}^{2}(\mathsf{H}^{2}G_{1}^{4}+G_{2}^{2})}{(1-\gamma)^{4}}$.
Then from Equation (18), we can get

$\displaystyle\mathbb{E}[\|\nabla^{2}J_{\mathsf{H}}(\theta)-\hat{\nabla}^{2}_{m}J(\theta)\|^{2}]\leq\frac{4e\cdot\max\\{2,\log
d\\}R_{\max}^{2}(\mathsf{H}^{2}G_{1}^{4}+G_{2}^{2})}{m(1-\gamma)^{4}}$ (156)

####  A.4.4 Soft-max policies satisfy Lipschitz Hessian

Recall that soft-max tabular policy as follows:

$\pi_{\theta}(a|s)=\frac{\exp(\theta_{s,a})}{\sum_{a\in\mathcal{A}}\exp(\theta_{s,a})}.$

For any $(s,a,a^{\prime})\in\mathcal{S}\times\mathcal{A}\times\mathcal{A}$
with $a^{\prime}\neq a$, we get the following partial derivatives for the
soft-max tabular policy

$\displaystyle\frac{\partial\pi_{\theta}(a|s)}{\partial\theta_{s,a}}=\pi_{\theta}(a|s)(1-\pi_{\theta}(a|s)),\quad\frac{\partial\pi_{\theta}(a|s)}{\partial\theta_{s,a^{\prime}}}=-\pi_{\theta}(a|s)\pi_{\theta}(a^{\prime}|s).$
(157)

Note that for $s^{\prime}\in\mathcal{S}$ with $s^{\prime}\neq s$, we have
$\frac{\partial\pi_{\theta}(a|s)}{\partial\theta_{s^{\prime},a}}=0$. We denote
$\theta_{s}=[\theta_{s,a}]_{a\in\mathcal{A}}\in\mathbb{R}^{|\mathcal{A}|}$. We
obtain the gradient and Hessian of $\log\pi_{\theta}(a|s)$ w.r.t. $\theta_{s}$
as follows:

$\displaystyle\frac{\partial\log\pi_{\theta}(a|s)}{\partial\theta_{s}}=\boldsymbol{1}_{a}-\pi_{\theta}(\cdot|s),\quad\frac{\partial^{2}\log\pi_{\theta}(a|s)}{\partial\theta_{s}^{2}}=\pi_{\theta}(\cdot|s)\pi_{\theta}(\cdot|s)^{T}-\text{\rm{Diag}}(\pi_{\theta}(\cdot|s))$
(158)

where $\boldsymbol{1}_{a}\in\mathbb{R}^{|\mathcal{A}|}$ is a vector with zero
entries except one non-zero entry $1$ corresponding to the action $a$ and
$\pi_{\theta}(\cdot|s)=[\pi_{\theta}(a|s)]_{a\in\mathcal{A}}\in\mathbb{R}^{|\mathcal{A}|}$
is a vector consists of all policies with the same state $s$. Now we have

$\displaystyle\|\nabla^{2}\log\pi_{\theta^{\prime}}(a|s)-\nabla^{2}\log\pi_{\theta}(a|s)\|=\left\|\frac{\partial^{2}}{\partial\theta^{2}_{s}}\log\pi_{\theta^{\prime}}(a|s)-\frac{\partial^{2}}{\partial\theta^{2}_{s}}\log\pi_{\theta}(a|s)\right\|$
$\displaystyle\leq\underbrace{\left\|\pi_{\theta^{\prime}}(\cdot|s)\pi_{\theta^{\prime}}(\cdot|s)^{T}-\pi_{\theta}(\cdot|s)\pi_{\theta}(\cdot|s)^{T}\right\|}_{(1)}+\underbrace{\left\|\text{\rm{Diag}}(\pi_{\theta^{\prime}}(\cdot|s))-\text{\rm{Diag}}(\pi_{\theta}(\cdot|s))\right\|}_{(2)}.$
(159)

The first term (1) can be bounded as follows:

$\displaystyle(1)=\left\|\pi_{\theta}(\cdot|s)\pi_{\theta}(\cdot|s)^{T}-\pi_{\theta^{\prime}}(\cdot|s)\pi_{\theta}(\cdot|s)^{T}+\pi_{\theta^{\prime}}(\cdot|s)\pi_{\theta}(\cdot|s)^{T}-\pi_{\theta^{\prime}}(\cdot|s)\pi_{\theta^{\prime}}(\cdot|s)^{T}\right\|$
$\displaystyle\leq(\|\pi_{\theta}(\cdot|s)\|_{2}+\|\pi_{\theta^{\prime}}(\cdot|s)\|_{2})\|\pi_{\theta}(\cdot|s)-\pi_{\theta^{\prime}}(\cdot|s)\|_{2}$
$\displaystyle\overset{(a)}{\leq}2\|\pi_{\theta}(\cdot|s)-\pi_{\theta^{\prime}}(\cdot|s)\|_{2}\overset{(b)}{\leq}4\|\theta-\theta^{\prime}\|_{2}.$
(160)

where (a) comes from the fact that as $\sum_{a}\pi(a|s)=1$ then the
$\|\pi_{\theta}(\cdot|s)\|_{2}\leq 1$ is less than one. Inequality (b) is
derived by $\max_{\theta}\|\nabla_{\theta}\pi_{\theta}(a|s)\|\leq 2$ which
yields $\pi_{\theta}(a,s)$ is 2-Lipschitz.

The second term (2) is bounded by $2\|\theta-\theta^{\prime}\|_{2}$. Hence,
$\log\pi_{\theta}(a|s)$ for soft-max policy is 6-Lipschitz Hessian.

Lemma 2. Under Assumption 7, we have $\|\nabla J(\theta)-\nabla
J_{\mathsf{H}}(\theta)\|\leq D_{g}\gamma^{\mathsf{H}}$ and
$\|\nabla^{2}J(\theta)-\nabla^{2}J_{\mathsf{H}}(\theta)\|\leq
D_{H}\gamma^{\mathsf{H}},$ where
$D_{g}=\frac{G_{1}R_{\max}}{1-\gamma}\sqrt{\frac{1}{1-\gamma}+\mathsf{H}}$ and
$D_{H}=\frac{R_{\max}(G_{2}+G_{1}^{2})}{1-\gamma}\left(\mathsf{H}+\frac{1}{1-\gamma}\right)$.

####  A.4.5 Proof of Lemma 3

Under Assumption 7, it has been shown that
$D_{g}=\frac{G_{1}R_{\max}}{1-\gamma}\sqrt{\frac{1}{1-\gamma}+\mathsf{H}}$
[41, Lemma 4.5]. We will show that under Assumption 7,
$D_{H}=\frac{R_{\max}(G_{2}+G_{1}^{2})}{1-\gamma}\left(\mathsf{H}+\frac{1}{1-\gamma}\right)$.

It can be shown that [41, Proof of Lemma 4.4]

$\displaystyle\nabla^{2}J(\theta)=\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=0}^{\infty}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]$
$\displaystyle+\mathbb{E}_{\tau}\left[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla^{2}_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right]$

and then

$\displaystyle\nabla^{2}J_{\mathsf{H}}(\theta)=\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=0}^{\mathsf{H}-1}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]$
$\displaystyle+\mathbb{E}_{\tau}\left[\sum_{t=0}^{\mathsf{H}-1}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla^{2}_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right].$

Therefore,

$\displaystyle\nabla^{2}J(\theta)-\nabla^{2}J_{\mathsf{H}}(\theta)=\underbrace{\mathbb{E}_{\tau}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla^{2}_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right]}_{(1)}$
$\displaystyle+\underbrace{\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=\mathsf{H}}^{\infty}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]}_{(2)}$
$\displaystyle+\underbrace{\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]}_{(3)}$

First we bound $\|(1)\|_{2}$ as follows:

$\displaystyle\|(1)\|_{2}\leq\mathbb{E}_{\tau}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}|R(s_{t},a_{t})|\left(\sum_{k=0}^{t}\|\nabla_{\theta}\log\pi_{\theta}(a_{k|s_{k}})\|\right)\right]$
$\displaystyle\leq
R_{\max}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\left(\sum_{k=0}^{t}\mathbb{E}_{\tau}\|\nabla_{\theta}\log\pi_{\theta}(a_{k|s_{k}})\|\right)\right]$
$\displaystyle\leq
G_{2}R_{\max}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}(1+t)=\gamma^{\mathsf{H}}G_{2}R_{\max}\sum_{t=0}^{\infty}\gamma^{t}(1+t+\mathsf{H})=\gamma^{H}G_{2}R_{\max}\left(\frac{1}{(1-\gamma)^{2}}+\frac{\mathsf{H}}{1-\gamma}\right).$

To bound the second term (2), it is good to note that for $0\leq t<t^{\prime}$

$\mathbb{E}_{\tau}\left[\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]=0.$

Because

$\displaystyle\mathbb{E}_{\tau}\left[\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]$
$\displaystyle=\mathbb{E}_{s_{0:t^{\prime}},a_{0:(t^{\prime}-1)}}\left[\underbrace{\mathbb{E}_{a_{t^{\prime}}}\left[\nabla_{\theta^{\prime}}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})|s_{t^{\prime}}\right]}_{=0}\cdot\left(\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]=0$
(161)

where

$\mathbb{E}_{a_{t^{\prime}}}\left[\nabla_{\theta^{\prime}}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})|s_{t^{\prime}}\right]=\int\nabla_{\theta}\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})da_{t^{\prime}}=\nabla_{\theta}\underbrace{\int\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})da_{t^{\prime}}}_{=1}=0.$

Hence, we get

$\displaystyle\|(2)\|$
$\displaystyle=\left\|\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=\mathsf{H}}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]\right\|$
$\displaystyle\overset{(a)}{=}\left\|\mathbb{E}_{\tau}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{t^{\prime}=\mathsf{H}}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\right)\left(\sum_{k=\mathsf{H}}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)^{T}\right]\right\|$
$\displaystyle\leq
R_{\max}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\mathbb{E}_{\tau}\left[\left\|\sum_{k=\mathsf{H}}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right\|^{2}\right]$
$\displaystyle\overset{(b)}{=}R_{\max}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\sum_{k=\mathsf{H}}^{t}\mathbb{E}_{\tau}\left[\left\|\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right\|^{2}\right]$
$\displaystyle\leq
R_{\max}G_{1}^{2}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}(t-\mathsf{H}+1)=R_{\max}G_{1}^{2}\gamma^{\mathsf{H}}\sum_{t=0}^{\infty}\gamma^{t}(t+1)\leq\frac{G_{1}^{2}R_{\max}\gamma^{\mathsf{H}}}{(1-\gamma)^{2}}$

where (a) comes from the similar argument to (A.4.5) and (b) comes from the
fact that for any $k\neq k^{\prime}$

$\displaystyle\mathbb{E}_{\tau}\left[\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})^{T}\nabla_{\theta}\log\pi_{\theta}(a_{k^{\prime}}|s_{k^{\prime}})\right]=0.$
(162)

To bound the third term (3), first, with the similar argument to (A.4.5), we
can rewrite (3) as follows:

$\displaystyle\|(3)\|$
$\displaystyle=\left\|\mathbb{E}_{\tau}\left[\sum_{t^{\prime}=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\left(\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)\right)^{T}\right]\right\|$
$\displaystyle=\left\|\mathbb{E}_{\tau}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\cdot\left(\sum_{t^{\prime}=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\right)\cdot\left(\sum_{k=0}^{t}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)^{T}\right]\right\|$
$\displaystyle\overset{(a)}{=}\left\|\mathbb{E}_{\tau}\left[\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}R(s_{t},a_{t})\cdot\left(\sum_{t^{\prime}=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{t^{\prime}}|s_{t^{\prime}})\right)\cdot\left(\sum_{k=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right)^{T}\right]\right\|$
$\displaystyle\leq
R_{\max}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\cdot\mathbb{E}_{\tau}\left[\left\|\sum_{k=0}^{\mathsf{H}-1}\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right\|^{2}\right]$
$\displaystyle\leq
R_{\max}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\cdot\sum_{k=0}^{\mathsf{H}-1}\mathbb{E}_{\tau}\left[\left\|\nabla_{\theta}\log\pi_{\theta}(a_{k}|s_{k})\right\|^{2}\right]\leq
R_{\max}G_{1}^{2}\sum_{t=\mathsf{H}}^{\infty}\gamma^{t}\mathsf{H}=\frac{\gamma^{\mathsf{H}}R_{\max}G_{1}^{2}\mathsf{H}}{1-\gamma}$

where (a) comes from (162).

Using triangle inequality and combining all bounds for (1), (2), and (3), we
have

$\displaystyle\|\nabla^{2}J(\theta)-\nabla^{2}J_{\mathsf{H}}(\theta)\|\leq\|(1)\|+\|(2)\|+\|(3)\|$
$\displaystyle\leq\gamma^{\mathsf{H}}\left(G_{2}R_{\max}\left(\frac{1}{(1-\gamma)^{2}}+\frac{\mathsf{H}}{1-\gamma}\right)+\frac{G_{1}^{2}R_{\max}}{(1-\gamma)^{2}}+\frac{R_{\max}G_{1}^{2}\mathsf{H}}{1-\gamma}\right).$
(163)

####  A.4.6 Analysis of the variance-reduced SCRN under weak gradient
dominance property with $\alpha=1$

######  Definition 3 (Importance Sampling).

For any given trajectory $\tau$ and two parameters
$\theta,\theta^{\prime}\in\mathbb{R}^{d}$, we define the importance sampling
weight as

$\displaystyle
w(\tau|\theta^{\prime},\theta)=\frac{p(\tau|\pi_{\theta^{\prime}})}{p(\tau|\pi_{\theta})}=\prod_{i=1}^{\mathsf{H}}\frac{\pi_{\theta^{\prime}}(a_{h}|s_{h})}{\pi_{\theta}(a_{h}|s_{h})},$
(164)

where $\mathsf{H}$ is the horizon length.

######  Assumption 15.

There exists a constant $W\in(0,\infty)$, such that, for any pair of policies
generated by Algorithm 3, the following holds

$\displaystyle\mathop{\rm Var}\nolimits(w(\tau|\theta^{\prime},\theta))\leq
W,\quad\theta,\theta^{\prime}\in\mathbb{R}^{d},\,\tau\sim
p(\cdot|\pi_{\theta^{\prime}}).$ (165)

In following, we present a version of Algorithm 2 adapted for RL setting by
utilizing importance sampling. In what follows,
$\hat{\nabla}J(\theta^{\prime};\tau)$ and
$\hat{\nabla}^{2}J(\theta^{\prime};\tau)$ are gradient and Hessian estimators
of the value function, respectively which are sampled from distribution
$p(\tau|\pi_{\theta})$ for some $\theta$.

Algorithm 3 IS-VR stochastic cubic regularized Newton method

Input: Maximum number of iterations $T$, batch sizes
$\\{n^{t}_{g}\\}_{t=1}^{T}$, $\\{n^{t}_{H}\\}_{t=1}^{T}$, the period length
$S$, initial point $\theta_{0}$, and cubic penalty parameter $M$.

1:$t\leftarrow 0$

2:while $t\leq T$ do

3: Sample trajectory set
$\mathcal{J}_{t}=\\{\tau^{t}_{j}\\}^{n^{t}_{g}}_{j=1}$ and trajectory set
$\mathcal{I}_{t}=\\{\tau^{t}_{i}\\}^{n^{t}_{H}}_{i=1}$ according to policy
$p(\tau|\pi_{\theta_{t}})$

4: ${\bf
v}_{t}\leftarrow\begin{cases}\frac{1}{n^{t}_{g}}\sum_{j=1}^{n^{t}_{g}}\hat{\nabla}J(\theta_{t};\tau^{t}_{j}),&\text{mod}(t,S)=0\\\
\frac{1}{n^{t}_{g}}\sum_{j=1}^{n^{t}_{g}}[\hat{\nabla}J(\theta_{t};\tau^{t}_{j})-w(\tau^{t}_{j}|\theta_{t-1},\theta_{t})\hat{\nabla}J(\theta_{t-1};\tau^{t}_{j})]+{\bf
v}_{t-1},&\text{else}\end{cases}$

5: ${\bf
U}_{t}\leftarrow\begin{cases}\frac{1}{n^{t}_{H}}\sum_{i=1}^{n^{t}_{H}}\hat{\nabla}^{2}J(\theta_{t};\tau^{t}_{i}),\quad&\text{mod}(t,S)=0\\\
\frac{1}{n^{t}_{H}}\sum_{i=1}^{n^{t}_{H}}[\hat{\nabla}^{2}J(\theta_{t};\tau^{t}_{i})-w(\tau^{t}_{i}|\theta_{t-1},\theta_{t})\hat{\nabla}^{2}J(\theta_{t-1};\tau^{t}_{i})]+{\bf
U}_{t-1},\quad&\text{else}\end{cases}$

6:
${\bf\Delta}_{t}\leftarrow\operatorname*{arg\,min}_{{\bf\Delta}\in\mathbb{R}^{d}}\langle{\bf
v}_{t},{\bf\Delta}\rangle+\frac{1}{2}\langle{\bf\Delta},{\bf
U}_{t}{\bf\Delta}\rangle+\frac{M}{6}\|{\bf\Delta}\|^{3}$

7: $\theta_{t+1}\leftarrow\theta_{t}+{\bf\Delta}_{t}$

8: $t\leftarrow t+1$

9:end while

10:return $\theta_{t}$

Assumptions 3 with $\alpha=1$ and Assumption 12 are satisfied in the
considered RL setting. To see this,
$\|\hat{\nabla}^{2}J(\theta,\tau)\|^{2}\leq\frac{\mathsf{H}^{2}G_{1}^{4}R^{2}_{\max}+G_{2}^{2}R^{2}_{\max}}{(1-\gamma)^{4}}$
from [29, Lemma 4.1]. Then, $\hat{\nabla}J(\theta,\tau)$ is Lipschitz.
Moreover, similar to the arguments in Lemma 13, one can show that
$\|\hat{\nabla}^{2}J(\theta,\tau)-\hat{\nabla}^{2}J(\theta^{\prime},\tau)\|\leq\tilde{L}\|\theta-\theta^{\prime}\|_{2}$.
Hence, we can imply that Lipschitzness of $\hat{\nabla}J(\theta,\tau)$ and
$\hat{\nabla}^{2}J(\theta,\tau)$ yield Assumption 12. Furthermore, from proof
of Lemma 3 (Section A.4.3), the variances of $\hat{\nabla}J(\theta,\tau)$ and
$\hat{\nabla}^{2}J(\theta,\tau)$ are bounded.

Now, with

$\mathsf{H}\geq\max\left\\{\frac{\log\left(\frac{C_{g}D_{g}}{\epsilon}\right)}{\log(1/\gamma)},\frac{\log\left(\frac{2C_{H}D_{H}}{\epsilon}\right)}{2\log(1/\gamma)}\right\\},$

Equation (A.4.2) can be rewritten as follows:

$\displaystyle J^{*}-\mathbb{E}J(\theta_{t+1})\leq
C(\mathbb{E}J(\theta_{t+1})-\mathbb{E}J(\theta_{t}))^{2/3}+C_{g}\mathbb{E}[\|\nabla
J_{\mathsf{H}}(\theta_{t})-{\bf v}_{t}\|]$
$\displaystyle+2C_{H}\mathbb{E}\left[\|\nabla^{2}J_{\mathsf{H}}(\theta_{t})-{\bf
U}_{t}\|^{2}\right]+2\epsilon.$ (166)

######  Lemma 15 (Lemma 6.1 in [38]).

For a policy $\pi_{\theta}$ satisfying Assumptions 7 and 15, and for any
$\theta,\theta^{\prime}\in\mathbb{R}^{d}$, we have

$\displaystyle\mathop{\rm Var}\nolimits(w(\tau|\theta,\theta^{\prime}))\leq
C_{w}\|\theta-\theta^{\prime}\|^{2},$ (167)

where $C_{w}=\mathsf{H}(2\mathsf{H}G_{1}^{2}+G_{2})(W+1)$.

In order to prove Theorem 3, we must show that with the following number of
samples at checkpoints:

$n_{g}^{t}=\Theta\left(\frac{1}{\epsilon^{2}}\right),n_{H}^{t}=\Theta\left(\frac{(\log
d)^{2}}{\epsilon}\right),$ (168)

and the following number of samples at the other iterations:

$n_{g}^{t}(\theta_{t},\theta_{t-1})=\Theta\left(\frac{S\|\mathbf{\Delta}_{t-1}\|^{2}}{\epsilon^{2}}\right),n_{H}^{t}(\theta_{t},\theta_{t-1})=\Theta\left(\frac{(\log
d)^{2}\cdot S\|\mathbf{\Delta}_{t-1}\|^{2}}{\epsilon}\right),$ (169)

we have:

$\mathbb{E}[\|\nabla J_{\mathsf{H}}(\theta_{t})-{\bf
v}_{t}\|]\leq\epsilon,\qquad\mathbb{E}[\|\nabla^{2}J_{\mathsf{H}}(\theta_{t})-{\bf
U}_{t}\|^{2}]\leq\epsilon.$ (170)

We use the same argument in the proof of Lemma 12 and we borrow the notations
in Lemma 12. Let us define ${\bf u}_{k}$ and ${\bf V}_{k}$. For the gradient
estimator ${\bf v}_{t}$, we have ${\bf v}_{t}-\nabla
J_{\mathsf{H}}(\theta_{t})=\sum_{k=\lfloor t/S\rfloor S}^{t}{\bf u}_{k}$ such
that:

${\bf
u}_{k}=\begin{cases}\frac{1}{n^{k}_{g}}\sum_{j=1}^{n^{k}_{g}}\hat{\nabla}J(\theta_{k};\tau^{k}_{j})-\nabla
J_{\mathsf{H}}(\theta_{k}),\quad&k=\lfloor t/S\rfloor S,\\\
\frac{1}{n^{k}_{g}}\sum_{j=1}^{n^{k}_{g}}[\hat{\nabla}J(\theta_{k};\tau^{k}_{j})-w(\tau^{k}_{j}|\theta_{k-1},\theta_{k})\hat{\nabla}J(\theta_{k-1};\tau^{k}_{j})]-\nabla
J_{\mathsf{H}}(\theta_{k})+\nabla J_{\mathsf{H}}(\theta_{k-1}),\quad&k>\lfloor
t/S\rfloor S,\end{cases}$

and for the Hessian estimator ${\bf U}_{t}$, we have ${\bf
U}_{t}-\nabla^{2}J_{\mathsf{H}}(\theta_{t})=\sum_{k=\lfloor t/S\rfloor
S}^{t}{\bf V}_{k}$ such that

${\bf
V}_{k}=\begin{cases}\frac{1}{n^{k}_{H}}\sum_{i=1}^{n^{k}_{H}}\hat{\nabla}^{2}J(\theta_{k};\tau^{k}_{i})-\nabla^{2}J_{\mathsf{H}}(\theta_{k}),\quad&k=\lfloor
t/S\rfloor S,\\\
\frac{1}{n^{k}_{H}}\sum_{i=1}^{n^{k}_{H}}[\hat{\nabla}^{2}J(\theta_{k};\tau^{k}_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1};\tau^{k}_{i})]-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1}),\quad&k>\lfloor
t/S\rfloor S.\end{cases}$

To prove that (170), as the importance sampling terms appears in ${\bf u}_{k}$
and ${\bf V}_{k}$ for $k>\lfloor t/S\rfloor S$, we must only modify the proof
of $\mathbb{E}[\|{\bf
u}_{k}\|^{2}|\mathcal{G}_{k}]\leq\frac{C_{1}}{n^{k}_{g}}\|\theta_{k}-\theta_{k-1}\|^{2}$
and $\mathbb{E}[\|{\bf
V}_{k}\|^{2}|\mathcal{H}_{k}]\leq\frac{C_{2}}{n^{k}_{g}}\|\theta_{k}-\theta_{k-1}\|^{2}$
for $k>\lfloor t/S\rfloor S$. We prove them as follows:

$\displaystyle\mathbb{E}[\|{\bf
u}_{k}\|^{2}|\mathcal{G}_{k}]=\mathbb{E}\left\|\frac{1}{n^{k}_{g}}\sum_{j=1}^{n^{k}_{g}}[\hat{\nabla}J(\theta_{k};\tau^{k}_{j})-w(\tau^{k}_{j}|\theta_{k-1},\theta_{k})\hat{\nabla}J(\theta_{k-1};\tau^{k}_{j})]-\nabla
J_{\mathsf{H}}(\theta_{k})+\nabla J_{\mathsf{H}}(\theta_{k-1})\right\|^{2}$
$\displaystyle\overset{(a)}{=}\frac{1}{n^{k}_{g}}\mathbb{E}\left\|\hat{\nabla}J(\theta_{k};\tau^{k}_{1})-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})\hat{\nabla}J(\theta_{k-1};\tau^{k}_{1})-\nabla
J_{\mathsf{H}}(\theta_{k})+\nabla J_{\mathsf{H}}(\theta_{k-1})\right\|^{2}$
$\displaystyle\overset{(b)}{\leq}\frac{2}{n^{k}_{g}}\mathbb{E}\|\hat{\nabla}J(\theta_{k};\tau^{k}_{1})-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})\hat{\nabla}J(\theta_{k-1};\tau^{k}_{1})\|^{2}+\frac{2}{n^{k}_{g}}\mathbb{E}\|\nabla
J_{\mathsf{H}}(\theta_{k})-\nabla J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(c)}{\leq}\frac{4}{n^{k}_{g}}\mathbb{E}\|\hat{\nabla}J(\theta_{k};\tau^{k}_{1})-\hat{\nabla}J(\theta_{k-1};\tau^{k}_{1})\|^{2}+\frac{4}{n^{k}_{g}}\mathbb{E}[(w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})-1)^{2}\|\hat{\nabla}J(\theta_{k-1};\tau^{k}_{1})\|^{2}]$
$\displaystyle+\frac{2}{n^{k}_{g}}\mathbb{E}\|\nabla
J_{\mathsf{H}}(\theta_{k})-\nabla J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(d)}{\leq}\frac{4}{n^{k}_{g}}\mathbb{E}\|\hat{\nabla}J(\theta_{k};\tau^{k}_{1})-\hat{\nabla}J(\theta_{k-1};\tau^{k}_{1})\|^{2}+\frac{4G_{1}^{2}}{n^{k}_{g}}\mathbb{E}[(w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})-1)^{2}]+\frac{2}{n^{k}_{g}}\mathbb{E}\|\nabla
J_{\mathsf{H}}(\theta_{k})-\nabla J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(e)}{\leq}\frac{6L^{\prime
2}_{1}+4G^{2}_{1}C_{w}}{n^{k}_{g}}\|\theta_{k}-\theta_{k-1}\|^{2}$ (171)

where (a) comes from
$\hat{\nabla}J(\theta_{k};\tau^{k}_{j})-w(\tau^{k}_{j}|\theta_{k-1},\theta_{k})\hat{\nabla}J(\theta_{k-1};\tau^{k}_{j})-\nabla
J_{\mathsf{H}}(\theta_{k})+\nabla J_{\mathsf{H}}(\theta_{k-1})$’s are i.i.d
conditioned on $\mathcal{G}_{k}$ for $1\leq i\leq n^{k}_{g}$. Inequalities (b)
and (c) are from $(a+b)^{2}\leq 2a^{2}+2b^{2}$. (d) is due to Assumption 7.
(e) is derived by Assumption 12 and Lemma 15.

Similarly, for Hessian estimator we have

$\displaystyle\mathbb{E}[\|{\bf
V}_{k}\|^{2}|\mathcal{H}_{k}]=\mathbb{E}\left\|\frac{1}{n^{k}_{H}}\sum_{i=1}^{n^{k}_{H}}[\hat{\nabla}^{2}J(\theta_{k},\tau_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{i})]-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\right\|^{2}$
$\displaystyle\overset{(a)}{\leq}8e\log
d\cdot\mathbb{E}\left\|\frac{1}{(n^{k}_{H})^{2}}\sum_{i=1}^{n^{k}_{H}}\left[[\hat{\nabla}^{2}J(\theta_{k},\tau_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{i})]-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\right]^{2}\right\|$
$\displaystyle\overset{(b)}{\leq}\frac{8e\log
d}{n^{k}_{H}}\cdot\frac{1}{n^{k}_{H}}\sum_{i=1}^{n^{k}_{H}}\mathbb{E}\left\|[\hat{\nabla}^{2}J(\theta_{k},\tau_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{i})-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})]^{2}\right\|$
$\displaystyle\overset{(c)}{\leq}\frac{8e\log
d}{n^{k}_{H}}\mathbb{E}\left\|\hat{\nabla}^{2}J(\theta_{k},\tau_{1})-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{1})-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\right\|^{2}$
$\displaystyle\overset{(d)}{\leq}\frac{16e\log
d}{n^{k}_{H}}\mathbb{E}\|\hat{\nabla}^{2}J(\theta_{k},\tau_{1})-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{1})\|^{2}+\frac{16e\log
d}{n^{k}_{H}}\|\nabla^{2}J_{\mathsf{H}}(\theta_{k})-\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(e)}{\leq}\frac{32e\log
d}{n^{k}_{H}}\mathbb{E}\|\hat{\nabla}^{2}J(\theta_{k},\tau_{1})-\hat{\nabla}^{2}J(\theta_{k-1},\tau_{1})\|^{2}+\frac{32e\log
d}{n^{k}_{H}}\mathbb{E}\left[(1-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k}))^{2}\|\hat{\nabla}^{2}J(\theta_{k-1},\tau_{1})\|^{2}\right]$
$\displaystyle+\frac{16e\log
d}{n^{k}_{H}}\|\nabla^{2}J_{\mathsf{H}}(\theta_{k})-\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(f)}{\leq}\frac{32e\log
d}{n^{k}_{H}}\mathbb{E}\|\hat{\nabla}^{2}J(\theta_{k},\tau_{1})-\hat{\nabla}^{2}J(\theta_{k-1},\tau_{1})\|^{2}+\frac{32e\cdot
G^{2}_{2}\log
d}{n^{k}_{H}}\mathbb{E}\left[(1-w(\tau^{k}_{1}|\theta_{k-1},\theta_{k}))^{2}\right]$
$\displaystyle+\frac{16e\log
d}{n^{k}_{H}}\|\nabla^{2}J_{\mathsf{H}}(\theta_{k})-\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})\|^{2}$
$\displaystyle\overset{(g)}{\leq}\frac{32e\log d(2L^{\prime
2}_{2}+C_{w}G^{2}_{2})}{n^{k}_{H}}\|\theta_{k}-\theta_{k-1}\|^{2},$ (172)

where (a) comes from Lemma 4 and the fact that
$[\hat{\nabla}^{2}J(\theta_{k},\tau_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{i})-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})]$’s
are i.i.d conditioned on $\mathcal{H}_{k}$ for $1\leq i\leq n^{k}_{H}$. (b)
comes from Jensen’s inequality for operator norm $\|\cdot\|$. (c) is derived
by $\|AB\|\leq\|A\|\|B\|$ and the fact that
$[\hat{\nabla}^{2}J(\theta_{k},\tau_{i})-w(\tau^{k}_{i}|\theta_{k-1},\theta_{k})\hat{\nabla}^{2}J(\theta_{k-1},\tau_{i})-\nabla^{2}J_{\mathsf{H}}(\theta_{k})+\nabla^{2}J_{\mathsf{H}}(\theta_{k-1})]$’s
are i.i.d conditioned on $\mathcal{H}_{k}$ for $1\leq i\leq n^{k}_{H}$.
Inequalities (d) and (e) are from inequality $(a+b)^{2}\leq 2a^{2}+2b^{2}$ and
(f) is derived by Assumption 7. (g) comes from Assumption 12 and Lemma 15.
Similar to the proof of Lemma 12, we can argue that for Algorithm 3, (170)
holds.

Now, based on what we proved above, with the same arguments in Theorem 3,
Algorithm 3 achieves $\epsilon$-global stationary point by querying
${\mathcal{O}}(\epsilon^{-2})$ stochastic gradients and
${\mathcal{O}}(\epsilon^{-1})$ stochastic Hessian on average. More precisely,
the following average number of queries of stochastic gradient and Hessian are
required in order to achieve $\epsilon$-global stationary point:

$2\mathsf{H}\times\sum_{i=1}^{T}n_{g}^{t}=\mathcal{O}\left(\frac{\sigma_{1}^{2}\mathsf{H}\left[4^{2}(C_{g}+C_{H})C^{3}+C_{g}+C_{H}\right]^{2}}{\epsilon^{2}}\right)=\mathcal{O}\left(\frac{(G_{1}^{2+\frac{8}{3}}G_{2}^{\frac{8}{3}}R_{\max}^{2+\frac{8}{3}}}{(1-\gamma)^{4+\frac{16}{3}}\epsilon^{2}}\right).$

where (a) comes from
$\sigma_{1}^{2}=\frac{\mathsf{H}G_{1}^{2}R_{\max}^{2}}{(1-\gamma)^{2}}$ Lemma
14 and $C=\mathcal{O}(M^{1/3})$, $C_{g}=\mathcal{O}(M^{1/3})$, and
$C_{H}=\mathcal{O}(M^{1/3})$ and the fact that $M=\mathcal{O}(\tilde{L})$ and
from Lemma 13,
$\tilde{L}=\mathcal{O}\left(\frac{R_{\max}G_{1}G_{2}}{(1-\gamma)^{2}}\right)$.

###  A.5 Further experimental results

Figure 6: Comparison of SCRN with first-order methods in maze environments. In
random shape maze, the percentages of successful instances for SPG, SCRN,
REINFORCE, SPG with entropy regularization, and REINFORCE with entropy
regularization are $86\%$, $100\%$, $95.3\%$, $92.2\%$, and $95.3\%$,
respectively. In random maze, the percentages of successful instances for SPG,
SCRN, REINFORCE, SPG with entropy regularization, and REINFORCE with entropy
regularization are $45.3\%$, $97\%$, $40.6\%$, $54.7\%$, and $54.7\%$,
respectively.

Discussion on results: In Fig. 6 (a), in the random shape maze environment,
SCRN finds paths with average length of $25$ after about $600$ episodes while
SPG and REINFORCE can only achieve the average path lengths of greater than
$25$ after $3000$ episodes. In the random maze environment, SCRN again
outperforms the other two algorithms and achieves the average path length of
$25$ after about $2000$ episodes. However, for SPG and REINFROCE, the average
path lengths are more than $100$ which also indicates the low percentages of
successful instances for these two algorithms in this environment.
Furthermore, as shown in Fig. 6 (c,d), in both random shape maze and random
maze environments, SCRN has better performance than SPG and REINFORCE in terms
of average episode return. We also provide the results for the random shape
maze and random maze environments where we add entropy regularization term to
the objective functions of SPG and REINFORCE. As can be seen, the results with
entropy regularization term are slightly better both in the average path
length and average episode return.

Figure 7: Comparison of SCRN and its variance-reduced version with REINFORCE
and variance-reduced SPG methods in MuJoCo environments.

In Fig. 7, for control tasks in MuJoCo simulator, we compared SCRN and its
variance-reduced version with first-order methods such as REINFORCE, and two
state-of-the-art representatives of variance-reduced PG methods, HAPG [29] and
MBPG [13], both with guaranteed convergence to $\epsilon$-FOSP in general non-
convex settings. In HalfCheetah environment, SCRN and variance-reduced SCRN
outperform the other methods by achieving a score of around $600$. In the
Humanoid environment, variance-reduced SCRN achieves the best performance
among other algorithms. In Walker environment, the performances of SCRN, its
variance-reduced version, and HAPG perform better than MBPG and REINFORCE. In
Reacher environment, the variance-reduced SCRN has the best performance among
the considered algorithms.

Figure 8: Comparison of SCRN and its variance-reduced version (without using
the trick in updates) with REINFORCE and variance-reduced SPG methods in
MuJoCo environments.

###  A.6 Implementation details

Environments:

  * $\bullet$

Environments with finite state and action spaces: We considered two grid world
environments in our experiments: [31, Example 6.6], and random mazes [51]. In
cliff walking, the agent’s aim is to reach a goal state from a start state,
avoiding a region of cells called “cliff”. The episode is terminated if the
agent enters the cliff region or the number of steps exceeds $100$ without
reaching the goal. The reward is $-0.1$ in all transitions except those into
the cliff where the reward is $-100$. The reward of reaching the goal is
$100$. Moreover, we considered a softmax tabular policy for all the
experiments of this part.

In random mazes, the size of each maze is $10\times 10$. In the random shape
maze, random shape blocks are placed on a grid and the agent tries to reach
the goal state finding the shortest path, avoiding blocks. The reward is
$-0.1$ in all transitions except if the agent tries to go to a cell which
belongs to a block where the reward is $-1$. Moreover, the reward of reaching
the goal is $1$. An episode is terminated if the agent could not reach the
goal after $200$ steps.

  * $\bullet$

Environments with continuous state and action spaces: We considered the
following control tasks in MuJoCo simulator [32]: Walker, Humanoid, Reacher,
and HalfCheeta. We compared SCRN with first-order methods such as REINFORCE,
and two state-of-the-art representatives of variance-reduced PG methods, HAPG
[29] and MBPG [13]. For each task, we utilized a Gaussian multi-layer
perceptron (MLP) policy whose mean and variance are parameterized by an MLP
with two hidden layers of 64 neurons. For a fair comparison, we considered the
same network architecture for all the methods.

Algorithms:

For the environments with finite state and action spaces, we provided an
implementation of SCRN, SPG, and REINFORCE with NumPy where the gradient and
Hessian (just for SCRN) of a given trajectory are computed based on their
closed forms. We also implemented SCRN with PyTorch in Garage library in order
to execute it on environments with continuous state and action spaces. In our
experiments, we used a Linux server with Intel Xeon Gold 6240 CPU (36 cores)
operating at 2.60GHz with 377 GB DDR4 of memory, and Nvidia Titan X GPU.

Regarding SPG and REINFORCE in grid world environments, we adapted a time-
varying learning rate by checking various forms and the form of $a/(\lfloor
t/P\rfloor+b)$ provided the best performance for the first-order methods where
$a,b,$ and $P$ are some constants that are needed to be tuned. The parameter
$P$ shows the number of episodes that the learning rate remains unchanged.

A trick in updates of SCRN and its variance-reduced version: Experiments in
Figures 3 and 7 are conducted using the following trick to prevent updating
based on a large $\|{\bf\Delta}_{t}\|$. If $\|{\bf\Delta}_{t}\|$ precedes a
threshold, we neglect that whole iteration and do not update ${\bf x}_{t}$.
Please note that we take the number of observed state-actions of that
neglected iteration into account in computing system probes.

For sake of completeness, in Figure 8, we also provided the performances of
SCRN and its variance-reduced version without using the trick described above,
and compared with other variance-reduced methods.

Demonstrations:

We studied how the parameters of the softmax tabular policy are evolving over
time in cliff walking environment. For each cell, we considered an arrow with
four directions: up, down, left, and right. At any time, the direction of each
arrow shows which one of the four directions has the highest probability of
being taken by the agent in the corresponding cell. The color of the arrow
becomes darker as this probability increases. For each algorithm, we run $100$
episodes and each episode, we demonstrated how the parameters are being
updated. For SPG and REINFORCE, for a long period of time, at the start state,
the agent takes the “up” action in order to avoid falling off the cliff. In
fact, it takes many episodes until the parameters of the start state are
updated such that the agent tries to find a path to the goal. In contrast,
SCRN finds a path to the goal after few episodes and then tries to improve the
path length.


# Rethinking Bayesian Deep Learning Methods for
Semi-Supervised Volumetric Medical Image Segmentation

Jianfeng Wang
Department of Computer Science
University of Oxford
<EMAIL_ADDRESS>Thomas Lukasiewicz
Department of Computer Science
University of Oxford
<EMAIL_ADDRESS>
###### Abstract

Recently, several Bayesian deep learning methods have been proposed for semi-
supervised medical image segmentation. Although they have achieved promising
results on medical benchmarks, some problems are still existing. Firstly,
their overall architectures belong to the discriminative models, and hence, in
the early stage of training, they only use labeled data for training, which
might make them overfit to the labeled data. Secondly, in fact, they are only
partially based on Bayesian deep learning, as their overall architectures are
not designed under the Bayesian framework. However, unifying the overall
architecture under the Bayesian perspective can make the architecture have a
rigorous theoretical basis, so that each part of the architecture can have a
clear probabilistic interpretation. Therefore, to solve the problems, we
propose a new generative Bayesian deep learning (GBDL) architecture. GBDL
belongs to the generative models, whose target is to estimate the joint
distribution of input medical volumes and their corresponding labels.
Estimating the joint distribution implicitly involves the distribution of
data, so both labeled and unlabeled data can be utilized in the early stage of
training, which alleviates the potential overfitting problem. Besides, GBDL is
completely designed under the Bayesian framework, and thus we give its full
Bayesian formulation, which lays a theoretical probabilistic foundation for
our architecture. Extensive experiments show that our GBDL outperforms
previous state-of-the-art methods in terms of four commonly used evaluation
indicators on three public medical datasets.

##  1 Introduction

Deep neural networks are a powerful tool for visual learning, and they have
outperformed previous methods that are based on hand-crafted features in many
classical computer vision tasks. Due to the power of deep neural networks,
researchers also try to use them to segment pathology regions from medical
images. Training segmentation networks usually relies on large labeled medical
image datasets, but they are scarce and difficult to obtain, since not
everyone is qualified to annotate medical data, and there is only a limited
number of medical experts. Therefore, semi-supervised segmentation has become
an important research direction in the area of medical computer vision.

Several methods [1, 22, 5, 9, 6, 2, 19, 10, 11, 21, 18, 14, 15, 17] have been
proposed in recent years.111Note that we only consider previous methods that
build their architectures with 3D CNNs, since our architecture is also based
on 3D CNNs. In most cases, using 3D CNNs to process medical volumetric data
performs better than using 2D CNNs, as the relationships among neighboring
slices are preserved in 3D CNNs. Thus, comparing our method with 2D-CNN-based
methods is unfair to them. Among these works, Bayesian deep-learning-based
methods [21, 18, 14, 15, 17] are closely related to this work, as they are
based on Monte Carlo (MC) dropout [4], which is an approximation of Bayesian
neural networks (BNNs). These methods are mainly built with the teacher-
student architecture, which can be regarded as the discriminative model.
Specifically, they initially build the model $P(Y|X)$ only from labeled data,
and then apply the model to generate pseudo labels that are refined or
rectified based on the epistemic uncertainty [7] provided by MC dropout for
unlabeled ones. Then, the pseudo labels are further combined with unlabeled
and labeled data to further train the overall architecture.

The works [21, 18, 14, 15, 17] mainly have two issues. Firstly, their models
are built only with labeled data in the early stage of training and may
overfit to them, since the quantity of labeled data is limited. Consequently,
the new pseudo labels generated by the models might be inaccurate, which
adversely impacts the subsequent training process. Secondly, their models are
only partially based on Bayesian deep learning and are not designed under the
Bayesian framework. Hence, they are unable to give a Bayesian formulation with
regard to their models, lacking a solid theoretical basis. As a result, some
modules in their models are empirically designed, and the functions of those
modules remain unclear.

In this work, we aim to fix these two problems by proposing a new generative
model, named generative Bayesian deep learning (GBDL) architecture, as well as
its corresponding full Bayesian formulation. Unlike teacher-student-based
architectures, GBDL estimates the joint probability distribution $P(X,Y)$ via
both labeled and unlabeled data, based on which more reliable pseudo-labels
can be generated, since the generative model is more faithful about the
overall data distribution $P(X)$ and the intrinsic relationship between inputs
and their corresponding labels (modeled by $P(X,Y)$), alleviating the
overfitting problem. Moreover, GBDL is entirely constructed under the Bayesian
framework, and its related full Bayesian formulation lays a solid theoretical
foundation, so that every part of GBDL has its corresponding probabilistic
formulation and interpretation.

The main contributions of this paper are as follows:

  * •

We propose a new generative Bayesian deep learning (GBDL) architecture. GBDL
aims to capture the joint distribution of inputs and labels, which moderates
the potential overfitting problem of the teacher-student architecture in
previous works [21, 18, 14, 15, 17].

  * •

Compared with previous methods [21, 18, 14, 15, 17], GBDL is completely
designed under the Bayesian framework, and its full Bayesian formulation is
also given here. Thus, it has a theoretical probabilistic foundation, and the
choice of each module in GBDL and the construction of the loss functions are
more mathematically grounded.

  * •

In extensive experiments, GBDL outperforms previous methods with a healthy
margin in terms of four evaluation indicators on three public medical
datasets, illustrating that GBDL is a superior architecture for semi-
supervised volumetric medical image segmentation.

The rest of this paper is organized as follows. In Section 2, we review
related methods. Section 3 presents our GBDL architecture and its Bayesian
formulation, followed by the experimental settings and results in Section 4.
In Section 5, we give a summary and an outlook on future research.

The source code is available at https://github.com/Jianf-Wang/GBDL.

##  2 Related Work

In this section, we briefly review related Bayesian and other deep learning
methods for semi-supervised volumetric medical image segmentation in the past
few years.

###  2.1 Bayesian Deep Learning Methods

Previous Bayesian deep learning methods simply use MC dropout as a tool to
detect when and where deep models might make false predictions [21, 18, 14,
15, 17]. Specifically, an uncertainty-aware self-ensembling model [21] was
proposed based on the teacher-student model and MC dropout. In particular, a
teacher model and a student model were built, where the latter learnt from the
former by minimizing the segmentation loss on the labeled data and the
consistency loss on all the input data. MC dropout was leveraged to filter out
the unreliable predictions and to preserve the reliable ones given by the
teacher model, and the refined predictions can help the student model learn
from unlabeled data. Based on this uncertainty-aware teacher-student model, a
double-uncertainty weighted method [18] was proposed, which takes both the
prediction and the feature uncertainty into consideration for refining the
predictions of the teacher model during training. Sedai et al. [14] used MC
dropout for training the teacher model as well, but they designed a novel loss
function to guide the student model by adaptively weighting regions with
unreliable soft labels to improve the final segmentation performance, rather
than simply removing the unreliable predictions. Shi et al. [15] improved the
way of uncertainty estimation by designing two additional models, which are
called object conservation model and object-radical model, respectively. Based
on these two models, certain region masks and uncertain region masks can be
obtained to improve pseudo labels and to prevent the possible error
propagation during training. Since the multi-task learning can boost the
performance of segmentation, and it has not been considered in previous works,
Wang et al. [17] combined it into their architecture, and imposed the
uncertainty estimation on all tasks to get a tripled-uncertainty that is
further used to guide the training of their student model.

###  2.2 Other Deep Learning Methods

Concerning other deep learning methods proposed recently, they are mainly
classified as three categories: new training and learning strategies [1, 22],
shape-aware or structure-aware based methods [5, 9, 6], and consistency
regularization based methods [2, 19, 10, 11].

The new training and learning strategies aim to gradually improve the quality
of pseudo labels during training. For example, Bai et al. [1] proposed an
iterative training strategy, in which the network is first trained on labeled
data, and then it predicts pseudo labels for unlabeled data, which are further
combined with the labeled data to train the network again. During the
iterative training process, a conditional random field (CRF) [8] was used to
refine the pseudo labels. Zeng et al. [22] proposed a reciprocal learning
strategy for their teacher-student architecture. The strategy contains a
feedback mechanism for the teacher network via observing how pseudo labels
would affect the student, which is omitted in previous teacher-student based
models.

The shape-aware or structure-aware based methods encourage the network to
explore complex geometric information in medical images. For instance, Hang et
al. [5] proposed a local and global structure-aware entropy-regularized mean
teacher (LG-ER-MT) architecture for semi-supervised segmentation. It extracts
local spatial structural information by calculating inter-voxel similarities
within small volumes and global geometric structure information by utilizing
weighted self-information. Li et al. [9] introduced a shape-aware semi-
supervised segmentation strategy that integrates a more flexible geometric
representation into the network for boosting the performance. Huang et al. [6]
designed a 3D Graph Shape-aware Self-ensembling Network (3D Graph-$S^{2}$Net),
which is composed by a multi-task learning network and a graph-based module.
The former performs the semantic segmentation task and predicts the signed
distance map that encodes rich features of object shape and surface, while the
latter explores co-occurrence relations and diffuse information between these
two tasks.

The consistency-regularization-based methods try to add different consistency
constraints for unlabeled data. In particular, Bortsova et al. [2] proposed a
model that is implemented by a Siamese architecture with two identical
branches to receive differently transformed versions of the same images. Its
target is to learn the consistency under different transformations of the same
input. Luo et al. [10] designed a novel dual-task-consistency model, whose
basic idea is to encourage consistent predictions of the same input under
different tasks. Xie et al. [19] proposed a pairwise relation-based semi-
supervised (PRS2) model for gland segmentation on histology tissue images. In
this model, a supervised segmentation network (S-Net) and an unsupervised
pairwise relation network (PR-Net) are built. The PR-Net learns both semantic
consistency and image representations from each pair of images in the feature
space for improving the segmentation performance of S-Net. Luo et al. [11]
proposed a pyramid (i.e., multi-scale) architecture that encourages the
predictions of an unlabeled input at multiple scales to be consistent, which
serves as a regularization for unlabeled data.

##  3 Methodology

In this section, we start from giving the Bayesian formulation of our GBDL
architecture. Thereafter, the architecture and related loss functions are
introduced in detail.

###  3.1 Bayesian Formulation

#### Learning Procedure.

The objective of semi-supervised volumetric medical image segmentation is to
learn a segmentation network, whose weights are denoted by $W$, with partially
labeled data, and the Bayesian treatment of this target is to learn the
posterior distribution, namely, $P(W|X,Y_{L})$. Now, we use $X$ to denote the
input volumes, which contains labeled input volumes ($X_{L}$) and unlabeled
input volumes ($X_{U}$), i.e., $X=\\{X_{L},X_{U}\\}$, and we use
$Y=\\{Y_{L},Y_{U}\\}$ to denote the ground-truth labels with respect to $X$.
Each label in $Y$ is a voxel-wise segmentation label map that has the same
shape as its corresponding input volume. Note that $Y_{U}$ represents the
ground-truth labels of $X_{U}$, which are not observed in the training data.

The whole learning procedure of $P(W|X,Y_{L})$ can be written as:

$\displaystyle P(W|X,Y_{L})$ (1) $\displaystyle=\iiint
P(W|X,Y)P(X,Y|Z)P(Z|X,Y_{L})dZdXdY,$

where $Z$ is the latent representation that governs the joint distribution of
$X$ and $Y$, denoted by $P(X,Y|Z)$. To estimate the joint distribution, one
should learn $Z$ from $X$ and $Y_{L}$. Concerning that Eq. (LABEL:eqn:overall)
is intractable, we take the MC approximation of it as follows:

$\displaystyle
P(W|X,Y_{L})=\frac{1}{MN}\sum\nolimits_{i=0}^{N-1}\sum\nolimits_{j=0}^{M-1}P(W|X_{(i,j)},Y_{(i,j)}).$
(2)

In this approximation, $M$ latent representations $Z$ are drawn from
$P(Z|X,Y_{L})$. Then, $N$ pairs of input volumes and labels are obtained from
$P(X,Y|Z)$, which are further used to get the posterior distribution. Thus, to
generate $X$ and $Y$, one should obtain the joint distribution $P(X,Y)$, i.e.,
estimating the distribution of $Z$. Overall, the learning procedure contains
two steps:

  * •

Learning the distribution of $Z$ that governs the joint distribution $P(X,Y)$.

  * •

Learning the posterior distribution $P(W|X,Y)$ based on $X$ and $Y$ sampled
from $P(X,Y)$.

#### Inference Procedure.

The inference procedure can be formulated as follows:

$\displaystyle P(Y_{pred}|X_{test},X,Y_{L})$ (3) $\displaystyle=\int
P(Y_{pred}|X_{test},W)P(W|X,Y_{L})dW,$

where $X_{test}$ and $Y_{pred}$ denote the test input and the predicted
result, respectively. The posterior distribution $P(W|X,Y_{L})$ is learnt
based on $X$ and $Y$ sampled from $P(X,Y)$, and thus the posterior can also be
replaced by $P(W|X,Y)$. Since Eq. (LABEL:eqn:testing) is intractable as well,
its MC approximation can be written as:

$\displaystyle
P(Y_{pred}|X_{test},X,Y_{L})=\frac{1}{T}\sum\nolimits_{i=0}^{T-1}P(Y_{pred}|X_{test},W_{i}),$
(4)

where $T$ models are drawn from $P(W|X,Y)$, which is usually implemented via
MC dropout [4] with $T$ times feedforward passes. For each model ($W_{i}$)
sampled from $P(W|X,Y)$, a prediction result can be obtained, and the final
prediction of $X_{test}$ can be calculated by averaging the $T$ results. In
addition, we get the epistemic uncertainty by calculating the entropy of these
predictions.

Figure 1: GBDL for semi-supervised volumetric medical image segmentation,
including a latent representation learning (LRL) architecture (in the green
dotted box) and a regular 3D-UNet with MC dropout (in the red dotted box).
Only the regular 3D-UNet with MC dropout is used during testing. For
simplicity, the shortcut connections between the paired 3D-UNet encoder and
decoder are omitted.

###  3.2 GBDL and Loss Functions

Under the guidance of the given Bayesian formulation, we start to introduce
the details of GBDL and loss functions. Concerning the two steps mentioned in
the learning procedure above, GBDL has a latent representation learning (LRL)
architecture and a regular 3D-UNet with MC dropout, which are shown in the
green dotted box and the red dotted box in Figure 1, respectively. LRL is
designed to learn the distribution of $Z$ and to capture the joint
distribution $P(X,Y)$, while the regular 3D-UNet with MC dropout is the
Bayesian deep model used for parameterizing the posterior distribution
$P(W|X,Y)$.

As for LRL, we first assume a variational probability distribution $Q(Z)$ to
represent the distribution of the latent representation $Z$. In this work, we
follow the conditional variational auto-encoder (cVAE) [16] to use the input
to modulate $Z$, and thus, $Q(Z)$ can be rewritten as $Q(Z|X)$. Note that this
work focuses on using 3D CNNs to process volumetric medical data, and
therefore, each input volume contains several slices, $Q(Z|X)$ is actually the
joint distribution of all these slices, i.e.,
$Q(Z|X)=q(z_{0},z_{1},z_{2},...,z_{n-1}|x_{0},x_{1},x_{2},...,x_{n-1})$, where
$n$ is the number of slices in an input volume, $x_{i}$ and $z_{i}$ are a
slice and its latent representation, respectively. To obtain the latent
representation of each slice $z_{i}$, a 3D-UNet encoder222In this paper, all
3D upsampling and downsampling layers used in the 3D-UNet encoder and the
3D-UNet decoder only perform on the spatial size, and therefore the depth of
the input volumes keeps unchanged, and we can obtain the latent representation
of each slice. More details about the network configuration are given in the
supplementary material. is used in LRL. As the downsampling layers of the
3D-UNet encoder do not perform on the depth, and the 3D convolutional layers
used in the encoder are of the kernel size $3$, the padding size $1$, and the
stride $1$ (shown in the supplementary material), each $z_{i}$ is actually
conditioned on $n_{rf}$ slices, where $n_{rf}$ is the total receptive field
along the depth dimension, which is determined by the number of 3D
convolutional layers in the encoder. For ease of calculation, we assume that
the distribution of the latent representation of each slice follows a
multivariate Gaussian distribution and that the latent representations of
slices are independent from each other. It is worth noting that such
independence assumption on the latent representations of slices is reasonable,
as in this way, the latent representation of each slice just rests on $n_{rf}$
slices, which is consistent with the fact that each slice is closely relevant
to its neighboring slices and distant slices may have no contribution to the
slice. Therefore, based on the independence assumption, $Q(Z|X)$ can be
decomposed into $\prod_{i=0}^{n-1}q(z_{i}|x_{(n_{rf})})$, where $x_{(n_{rf})}$
denotes the input slices that contribute to the $z_{i}$. After obtaining the
latent representation $Z$, we can write the evidence lower bound (ELBO) as
follows (with proof in the supplementary material):

$\displaystyle logP(X,Y)\geq$ (5)
$\displaystyle\mathbb{E}_{Q}[logP(Y|X,Z)+logP(X|Z)]-\mathbb{E}_{Q}[log(\frac{Q(Z|X)}{P(Z)})],$

where $\mathbb{E}_{Q}$ denotes the expectation over $Q(Z|X)$. Thus, the
learning objective is to maximize the ELBO (Eq. (LABEL:eqn:elbo)), which is
achieved by maximizing $\mathbb{E}_{Q}[logP(Y|X,Z)+logP(X|Z)]$ and minimizing
$\mathbb{E}_{Q}[log({Q(Z|X)}/{P(Z)})]$.

Firstly, to maximize $P(X|Z)$, another 3D-UNet decoder is used to take $z_{i}$
as input and the mean square error ($L_{MSE}$) is used as the loss function,
which aims to reconstruct each slice of the input volume. Secondly, to
maximize the probability $P(Y|X,Z)$, another branch can be built, which
contains a 3D-UNet encoder and a 3D-UNet decoder as well. Since only parts of
data have labels, $P(Y|X,Z)$ can also be rewritten as $P(Y_{L}|X_{L},Z)$. In
this branch, the 3D-UNet encoder receives $X_{L}$ to extract features that are
further combined with their corresponding $Z$ to be input to the 3D-UNet
decoder. Then, $Y_{L}$ is leveraged to calculate the dice loss $L_{Dice}$ [12]
and cross-entropy loss $L_{CE}$ with $X_{L}$ to maximize the probability
$P(Y_{L}|X_{L},Z)$. Thirdly, minimizing $\mathbb{E}_{Q}[log({Q(Z|X)}/{P(Z)})]$
is to use $Q(Z|X)$ to approximate the prior distribution of $Z$, and in most
cases, assuming that $Q(Z|X)$ and $P(Z)$ follow the same probability
distribution can make the computation convenient. The following theorem (with
proof in the supplementary material) is useful for constructing the
distribution of $Q(Z|X)$.

Theorem 1. _The product of any number $n\geq 1$ of multivariate Gaussian
probability density functions with precision $\Lambda_{i}$ and mean $\mu_{i}$
results in an unnormalized Gaussian curve with precision $\Lambda_{*}$ $=$
$\sum_{i=0}^{n-1}\Lambda_{i}$ and mean $\mu_{*}\,=$
$\Lambda_{*}^{-1}(\sum_{i=0}^{n-1}\Lambda_{i}\mu_{i})$._

Although an unnormalized Gaussian curve (denoted $f(x)$) is obtained by
multiplying several multivariate Gaussian PDFs, a Gaussian distribution can
still be obtained by simply normalizing $f(x)$ with $\int f(x)dx$ that does
not contain $x$. Therefore, based on Theorem 1, we can suppose that $Q(Z|X)$
follows a multivariate Gaussian distribution. In addition, we also assume that
the prior distribution $P(Z)$ follows a multivariate standard normal
distribution. To minimize $\mathbb{E}_{Q}[log({Q(Z|X)}/{P(Z)})]$, we propose a
new loss function as follows (with proof in the supplementary material).

Corollary 1. _Minimizing $\mathbb{E}_{Q}[log({Q(Z|X)}/{P(Z)})]$ is equivalent
to minimizing_

$\displaystyle L_{KL[Q(Z|X)||P(Z)]}=$ (6) $\displaystyle 1/2\cdot(-log\
det[(\sum\nolimits_{i=0}^{n-1}\Lambda_{i})^{-1}]+tr[(\sum\nolimits_{i=0}^{n-1}\Lambda_{i})^{-1}]\
+$
$\displaystyle(\sum\nolimits_{i=0}^{n-1}\Lambda_{i}\mu_{i})^{T}(\sum\nolimits_{i=0}^{n-1}\Lambda_{i})^{-2}(\sum\nolimits_{i=0}^{n-1}\Lambda_{i}\mu_{i})-D),$

_where $n$ denotes the number of slices, and $D$ denotes the number of
dimensions of mean vectors._

To summarize, maximizing ELBO (Eq. (LABEL:eqn:elbo)) is equivalent to
optimizing LRL and minimizing the loss function:

$\displaystyle L_{ELBO}=$ (7)
$\displaystyle\lambda_{1}L_{CE}+\lambda_{2}L_{Dice}+\lambda_{3}L_{MSE}+\lambda_{4}L_{KL[Q(Z|X)||P(Z)]},$

where $\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, and $\lambda_{4}$ are
coefficients. Significantly, cVAE [16] has some similar properties when
compared with our LRL architecture, and it is important to emphasize the
differences between cVAE-based models and our LRL, which mainly contain three
aspects:

  * •

Different ELBO: cVAE is mainly used for generation tasks, in which the data
distribution $P(X)$ is considered, and labels $Y$ are not involved into the
ELBO of cVAE. In contrast, our LRL concerns the joint distribution of $X$ and
$Y$, leading to a different ELBO (Eq. (LABEL:eqn:elbo)).

  * •

Independent Components Assumption: cVAE assumes that the components in the
latent representation of an input image are independent from each other. As a
result, the encoder of cVAE only gives a mean vector and a variance vector for
a latent representation333In practice, the encoder of cVAE gives log-variance
vectors instead of variance vectors.. However, our LRL does not follow this
assumption for each slice, and the encoder gives a mean vector and a
covariance matrix for the latent representation of each slice444Although the
independent components assumption is widely used, it could be too strong and
may degrade the performance..

  * •

Joint Variational Distribution: cVAE is mainly used for 2D inputs. In
contrast, our LRL is specifically designed for 3D-CNN-based architectures for
processing volumetric inputs. Therefore, since each input volume contains
several slices, the variational distribution becomes a joint distribution over
their latent representations, resulting in a different formulation of
Kullback–Leibler (KL) divergence between the joint variational distribution
and the prior distribution (Eq. (LABEL:eqn:kl1)).

Once LRL is trained well, the joint distribution of $X$ and $Y$ will be
captured, and the second step is to learn the posterior distribution
$P(W|X,Y)$. Therefore, input volumes and corresponding labels are generated
from the fixed LRL to train the regular 3D-UNet with MC dropout with two loss
functions that are widely used in medical image segmentation, namely, the dice
loss $L_{Dice}$ and the cross-entropy loss $L_{CE}$. In this paper, we use
$L_{Seg}$ to denote their summation:

$L_{Seg}=\beta_{1}L_{CE}+\beta_{2}L_{Dice},$ (8)

where $\beta_{1}$ and $\beta_{2}$ are the coefficient of the two loss terms in
$L_{Seg}$. In the real implementation, the new generated pseudo labels will be
combined with the unlabeled data and labeled data to train the regular 3D-UNet
with MC dropout. According to Eq. (4), only the the regular 3D-UNet with MC
dropout is used in the test phase. In addition, as we do a full Bayesian
inference during testing, it is convenient to obtain the voxel-wise epistemic
uncertainty for each predicted result by calculating the entropy
($-\sum_{c=0}^{C-1}p_{c}log_{2}p_{c}$) voxel-by-voxel, where $C$ is the number
of classes.

| KiTS19 | Atrial Segmentation Challenge
---|---|---
Scans used | Metrics | Scans used | Metrics
Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$ | Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$
3D-UNet w/ MC dropout | 160 | 0 | 0.940 ($\pm$0.004) | 0.888 ($\pm$0.005) | 3.66 ($\pm$0.26) | 0.78 ($\pm$0.11) | 80 | 0 | 0.915 ($\pm$0.006) | 0.832 ($\pm$0.007) | 3.89 ($\pm$0.17) | 1.25 ($\pm$0.14)
16 | 0 | 0.838 ($\pm$0.044) | 0.745 ($\pm$0.051) | 11.36 ($\pm$1.89) | 3.30 ($\pm$0.78) | 16 | 0 | 0.834 ($\pm$0.044) | 0.720 ($\pm$0.049) | 8.97 ($\pm$1.27) | 2.48 ($\pm$0.64)
4 | 0 | 0.710 ($\pm$0.051) | 0.578 ($\pm$0.065) | 19.56 ($\pm$1.95) | 6.11 ($\pm$0.92) | 8 | 0 | 0.801 ($\pm$0.057) | 0.676 ($\pm$0.051) | 11.40 ($\pm$1.53) | 3.27 ($\pm$0.71)
Volume-based LRL | 16 | 144 | 0.892 ($\pm$0.016) | 0.823 ($\pm$0.021) | 7.47 ($\pm$0.39) | 1.88 ($\pm$0.27) | 16 | 64 | 0.871 ($\pm$0.010) | 0.784 ($\pm$0.017) | 5.28 ($\pm$0.37) | 1.89 ($\pm$0.16)
IC-LRL | 16 | 144 | 0.900 ($\pm$0.011) | 0.828 ($\pm$0.012) | 6.99 ($\pm$0.26) | 1.75 ($\pm$0.22) | 16 | 64 | 0.882 ($\pm$0.014) | 0.796 ($\pm$0.026) | 4.69 ($\pm$0.22) | 1.66 ($\pm$0.12)
GBDL (LRL) | 16 | 144 |  0.911 ($\pm$0.010) |  0.840 ($\pm$0.013) |  6.38 ($\pm$0.37) |  1.51 ($\pm$0.22) | 16 | 64 |  0.894 ($\pm$0.008) |  0.822 ($\pm$0.011) |  4.03 ($\pm$0.41) |  1.48 ($\pm$0.14)
Volume-based LRL | 4 | 156 | 0.883 ($\pm$0.011) | 0.810 ($\pm$0.014) | 8.32 ($\pm$0.42) | 1.99 ($\pm$0.24) | 8 | 72 | 0.865 ($\pm$0.017) | 0.773 ($\pm$0.013) | 6.81 ($\pm$0.33) | 2.49 ($\pm$0.13)
IC-LRL | 4 | 156 | 0.889 ($\pm$0.014) | 0.814 ($\pm$0.014) | 8.01 ($\pm$0.31) | 2.01 ($\pm$0.12) | 8 | 72 | 0.871 ($\pm$0.009) | 0.779 ($\pm$0.024) | 5.96 ($\pm$0.22) | 1.97 ($\pm$0.17)
GBDL (LRL) | 4 | 156 |  0.898 ($\pm$0.008) |  0.821 ($\pm$0.011) |  6.85 ($\pm$0.44) |  1.78 ($\pm$0.23) | 8 | 72 |  0.884 ($\pm$0.007) |  0.792 ($\pm$0.012) |  5.89 ($\pm$0.31) |  1.60 ($\pm$0.15)

Table 1: Ablation studies of different LRL variants on the KiTS19 dataset and
the Atrial Segmentation Challenge dataset. For each of our results, the “mean
($\pm$std)” is reported. The first row shows the upper bound performance,
i.e., all training data with their labels are used.

| KiTS19 | Atrial Segmentation Challenge
---|---|---
Scans used | Metrics | Scans used | Metrics
Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$ | Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$
w/o $L_{MSE}$ | 16 | 144 | 0.877 ($\pm$0.014) | 0.789 ($\pm$0.019) | 8.89 ($\pm$0.33) | 1.96 ($\pm$0.19) | 16 | 64 | 0.873 ($\pm$0.012) | 0.781 ($\pm$0.011) | 4.76 ($\pm$0.46) | 1.66 ($\pm$0.17)
w/o $L_{KL}$ | 16 | 144 | 0.868 ($\pm$0.013) | 0.776 ($\pm$0.009) | 8.99 ($\pm$0.29) | 2.03 ($\pm$0.17) | 16 | 64 | 0.876 ($\pm$0.014) | 0.775 ($\pm$0.016) | 5.12 ($\pm$0.55) | 1.87 ($\pm$0.15)
Shared Encoder | 16 | 144 | 0.854 ($\pm$0.019) | 0.742 ($\pm$0.021) | 10.04 ($\pm$0.67) | 2.56 ($\pm$0.27) | 16 | 64 | 0.862 ($\pm$0.021) | 0.749 ($\pm$0.025) | 6.33 ($\pm$0.73) | 2.21 ($\pm$0.31)
GBDL | 16 | 144 |  0.911 ($\pm$0.010) |  0.840 ($\pm$0.013) |  6.38 ($\pm$0.37) |  1.51 ($\pm$0.22) | 16 | 64 |  0.894 ($\pm$0.008) |  0.822 ($\pm$0.011) |  4.03 ($\pm$0.41) |  1.48 ($\pm$0.14)
w/o $L_{MSE}$ | 4 | 156 | 0.864 ($\pm$0.016) | 0.769 ($\pm$0.014) | 9.03 ($\pm$0.39) | 2.05 ($\pm$0.18) | 8 | 72 | 0.867 ($\pm$0.011) | 0.771 ($\pm$0.017) | 4.96 ($\pm$0.34) | 1.85 ($\pm$0.20)
w/o $L_{KL}$ | 4 | 156 | 0.860 ($\pm$0.017) | 0.769 ($\pm$0.019) | 9.32 ($\pm$0.46) | 2.23 ($\pm$0.31) | 8 | 72 | 0.854 ($\pm$0.015) | 0.762 ($\pm$0.018) | 5.33 ($\pm$0.42) | 2.12 ($\pm$0.33)
Shared Encoder | 4 | 156 | 0.843 ($\pm$0.021) | 0.734 ($\pm$0.024) | 11.08 ($\pm$0.67) | 3.46 ($\pm$0.44) | 8 | 72 | 0.831 ($\pm$0.019) | 0.724 ($\pm$0.012) | 7.19 ($\pm$0.68) | 2.86 ($\pm$0.41)
GBDL | 4 | 156 |  0.898 ($\pm$0.008) |  0.821 ($\pm$0.011) |  6.85 ($\pm$0.44) |  1.78 ($\pm$0.23) | 8 | 72 |  0.884 ($\pm$0.007) |  0.792 ($\pm$0.012) |  5.89 ($\pm$0.31) |  1.60 ($\pm$0.15)

Table 2: Ablation studies of different parts in GBDL on the KiTS19 dataset and
the Atrial Segmentation Challenge dataset. For each of our results, the “mean
($\pm$std)” is reported.

##  4 Experiments

We now report on experiments of the proposed GBDL on three public medical
benchmarks. To save space, implementation details are shown in the
supplementary material.

###  4.1 Datasets

The KiTS19 dataset555https://kits19.grand-challenge.org/ is a kidney tumor
segmentation dataset, which has 210 labeled 3D computed tomography (CT) scans
for training and validation. We followed the settings in previous works [18,
3] to use 160 CT scans for training and 50 CT scans for testing. The 3D scans
centering at the kidney region were used, and we used the soft tissue CT
window range of [-100, 250] HU for the scans.

The Atrial Segmentation Challenge dataset [20] includes 100 3D gadolinium-
enhanced magnetic resonance imaging scans (GE-MRIs) with labels. We followed
previous works [18, 20, 10, 6, 9] to use 80 samples for training and the other
20 samples for testing.

The Liver Segmentation dataset has been released by the MICCAI 2018 medical
segmentation decathlon challenge.666http://medicaldecathlon.com/index.html The
dataset has 131 training and 70 testing points, and the ground-truth labels of
the testing data are not released. Therefore, we used those 131 CTs in our
experiments, in which 100 CT scans and 31 CT scans were used for training and
testing respectively. The 3D scans centering at liver regions were utilized,
and we used the soft tissue CT window range of [-100, 250] HU for the scans.

###  4.2 Evaluation Metrics

Four evaluation metrics are used as our evaluation indicators, namely, Dice
Score, Jaccard Score, 95$\%$ Hausdorff Distance (95HD), and Average
SurfaceDistance (ASD). Dice Score and Jaccard Score mainly compute the
percentage of overlap between two object regions. ASD computes the average
distance between the boundaries of two object regions, while 95HD measures the
closest point distance between two object regions.

###  4.3 Ablation Studies

We conducted ablation studies on the KiTS19 and the Atrial Segmentation
Challenge dataset. We comprehensively searched the hyperparameters
($\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, $\lambda_{4}$, $\beta_{1}$, and
$\beta_{2}$), and the number of sampled latent representations ($M$) for each
feedforward pass. Due to the space limit, the search process is shown in the
supplementary material. In all the following experiments, we fixed
$\lambda_{1}$, $\lambda_{2}$, $\lambda_{3}$, $\lambda_{4}$, $\beta_{1}$,
$\beta_{2}$, and $M$ to 1.0, 2.0, 1.0, 0.005, 1.0, 2.0, and 5 respectively.

Firstly, we evaluated the baseline model, i.e., the regular 3D-UNet with MC
dropout that is trained with a limited number of labeled data, and we also
evaluated the upper bound of the performance by utilizing all data to train
the model. As Table 1 shows, when the number of labeled data is decreased, the
performance of the regular 3D-UNet drops drastically.

| Scans used | Metrics
---|---|---
Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$
UA-MT [21] | 16 | 144 | 0.883 | 0.802 | 9.46 | 2.89
SASSNet [9] | 16 | 144 | 0.891 | 0.822 | 7.54 | 2.41
Double-UA [18] | 16 | 144 | 0.895 | 0.828 | 7.42 | 2.16
Tripled-UA [17] | 16 | 144 | 0.887 | 0.815 | 7.55 | 2.12
CoraNet [15] | 16 | 144 | 0.898 | 0.820 | 7.23 | 1.89
UA-MT* [21] | 16 | 144 | 0.878 | 0.797 | 8.93 | 2.76
Double-UA* [18] | 16 | 144 | 0.899 | 0.823 | 7.55 | 2.21
Tripled-UA* [17] | 16 | 144 | 0.886 | 0.819 | 7.43 | 2.26
CoraNet* [15] | 16 | 144 | 0.894 | 0.820 | 7.37 | 1.92
GBDL | 16 | 144 | 0.911 | 0.840 | 6.38 | 1.51
UA-MT [21] | 4 | 156 | 0.871 | 0.787 | 11.74 | 3.56
SASSNet [9] | 4 | 156 | 0.888 | 0.816 | 8.32 | 2.44
Double-UA [18] | 4 | 156 | 0.887 | 0.817 | 8.04 | 2.34
Tripled-UA [17] | 4 | 156 | 0.878 | 0.813 | 7.94 | 2.42
CoraNet [15] | 4 | 156 | 0.882 | 0.814 | 8.21 | 2.44
Double-UA* [18] | 4 | 156 | 0.890 | 0.819 | 7.93 | 2.33
UA-MT* [21] | 4 | 156 | 0.874 | 0.790 | 11.33 | 3.21
Tripled-UA* [17] | 4 | 156 | 0.882 | 0.810 | 7.81 | 2.47
CoraNet* [15] | 4 | 156 | 0.886 | 0.812 | 8.43 | 2.39
GBDL | 4 | 156 | 0.898 | 0.821 | 6.85 | 1.78

Table 3: Comparison with state-of-the-art semi-supervised segmentation methods
on the KiTS19 dataset. “*” denotes previous methods based on the regular
3D-UNet with MC dropout.

| Scans used | Metrics
---|---|---
Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$
UA-MT [21] | 16 | 64 | 0.889 | 0.802 | 7.32 | 2.26
SASSNet [9] | 16 | 64 | 0.895 | 0.812 | 8.24 | 2.20
Double-UA [18] | 16 | 64 | 0.897 | 0.814 | 7.04 | 2.03
Tripled-UA [17] | 16 | 64 | 0.893 | 0.810 | 7.42 | 2.21
CoraNet [15] | 16 | 64 | 0.887 | 0.811 | 7.55 | 2.45
Reciprocal Learning [22] | 16 | 64 | 0.901 | 0.820 | 6.70 | 2.13
DTC [10] | 16 | 64 | 0.894 | 0.810 | 7.32 | 2.10
3D Graph-S2 Net [6] | 16 | 64 | 0.898 | 0.817 | 6.68 | 2.12
LG-ER-MT [5] | 16 | 64 | 0.896 | 0.813 | 7.16 | 2.06
Double-UA* [18] | 16 | 64 | 0.894 | 0.809 | 6.16 | 2.28
UA-MT* [21] | 16 | 64 | 0.891 | 0.793 | 6.44 | 2.39
Tripled-UA* [17] | 16 | 64 | 0.889 | 0.809 | 6.88 | 2.48
CoraNet* [15] | 16 | 64 | 0.883 | 0.805 | 6.73 | 2.67
GBDL | 16 | 64 | 0.894 | 0.822 | 4.03 | 1.48
UA-MT [21] | 8 | 72 | 0.843 | 0.735 | 13.83 | 3.36
SASSNet [9] | 8 | 72 | 0.873 | 0.777 | 9.62 | 2.55
Double-UA [18] | 8 | 72 | 0.859 | 0.758 | 12.67 | 3.31
Tripled-UA [17] | 8 | 72 | 0.868 | 0.768 | 10.42 | 2.98
CoraNet [15] | 8 | 72 | 0.866 | 0.781 | 12.11 | 2.40
Reciprocal Learning [22] | 8 | 72 | 0.862 | 0.760 | 11.23 | 2.66
DTC [10] | 8 | 72 | 0.875 | 0.782 | 8.23 | 2.36
LG-ER-MT [5] | 8 | 72 | 0.855 | 0.751 | 13.29 | 3.77
3D Graph-S2 Net [6] | 8 | 72 | 0.879 | 0.789 | 8.99 | 2.32
Double-UA* [18] | 8 | 72 | 0.864 | 0.767 | 10.99 | 3.02
UA-MT* [21] | 8 | 72 | 0.847 | 0.744 | 12.32 | 3.20
Tripled-UA* [17] | 8 | 72 | 0.868 | 0.760 | 9.73 | 3.31
CoraNet* [15] | 8 | 72 | 0.861 | 0.770 | 11.32 | 2.46
GBDL | 8 | 72 | 0.884 | 0.792 | 5.89 | 1.60

Table 4: Comparison with state-of-the-art semi-supervised segmentation methods
on the Atrial Segmentation Challenge. “*” denotes previous methods based on
the regular 3D-UNet with MC dropout.

Second, as for an input volume, LRL assumes that the latent representation of
each slice follows a multivariate Gaussian distribution, and therefore, the
latent representation of the whole volume is sampled from a joint Gaussian
distribution that is obtained by multiplying the Gaussian distribution PDFs of
those slices. However, LRL can also be designed to learn the latent
representation of the whole volume, i.e., estimating the joint Gaussian
distribution directly, and we denote this variant as “volume-based LRL”.
Compared to the original LRL, this variant is much closer to a 3D-CNN-based
cVAE, since each input volume is treated as a whole, and the volume-based LRL
represents the whole volume with a latent representation, similar to the
original cVAE that represents the whole input image with a latent
representation. In this case, $L_{KL[Q(Z|X)||P(Z)]}$ degenerates to
$1/2\cdot(-log\,det[\Lambda_{v}^{-1}]+tr[\Lambda_{v}^{-1}]+\mu_{v}^{T}\mu_{v}-{D})$,
where $\mu_{v}$ and $\Lambda_{v}$ are the mean vector and the precision matrix
of the joint Gaussian distribution, respectively. We evaluated such volume-
based LRL, and the results are displayed in Table 1. For a fair comparison,
the $\mu_{v}$ and $\Lambda_{v}$ are of the same number of dimensions as the
$\mu_{i}$ and $\Lambda_{i}$ in the original LRL. The results show that the
volume-based LRL performs worse than the original LRL that is presented in
Section 3.2, demonstrating that the original LRL can better model the joint
distribution. Based on this observation, we hypothesize that keeping a
difference among slices is important, since slices and their corresponding
segmentation masks are different from each other. The original LRL adds
perturbation (random noise) to each slice, which is beneficial to keeping such
a difference. However, the volume-based LRL only adds random noise to the
latent representation of the whole volume, which goes against to keeping such
a difference.

Third, considering that the independent components assumption is widely used
in cVAE, we also evaluated LRL variant with this assumption, and we denote
this variant as “IC-LRL”. In this case, $L_{KL[Q(Z|X)||P(Z)]}$ degenerates to
$1/2\cdot\sum_{D}[-log\frac{1}{\sum_{i=0}^{n-1}\frac{1}{\sigma_{i}^{2}}}+\frac{1}{\sum_{i=0}^{n-1}\frac{1}{\sigma_{i}^{2}}}+(\frac{1}{\sum_{i=0}^{n-1}\frac{1}{\sigma_{i}^{2}}}\sum_{i=0}^{n-1}\frac{\mu_{i}}{\sigma_{i}^{2}})^{2}-1]$,
where $\mu_{i}$ and $\sigma_{i}$ are the mean vector and the variance vector
of the Gaussian distribution for each slice in the input volume, respectively.
The results shown in Table 1 indicate that the IC-LRL performs worse than our
original LRL, verifying that the independent components assumption harms the
performance of LRL.

Furthermore, we analyzed the two important parts of the proposed LRL, i.e.,
the reconstruction loss $L_{MSE}$ and the term $L_{KL[Q(Z|X)||P(Z))]}$. We
removed these two loss terms,777Removing $L_{MSE}$ is equivalent to removing
the reconstruction path, i.e., when we did an ablation study about removing
the reconstruction loss $L_{MSE}$, we also removed the upper 3D-UNet decoder
of LRL in Figure LABEL:eqn:overall. respectively, to study their impact on
GBDL. As Table 2 shows, when $L_{MSE}$ or $L_{KL[Q(Z|X)||P(Z))]}$ is removed,
a performance decline can be observed, demonstrating the importance of these
two terms and the correctness of our Bayesian formulation.

| Scans used | Metrics
---|---|---
Labeled | Unlabeled | Dice $\uparrow$ | Jaccard $\uparrow$ | 95HD $\downarrow$ | ASD $\downarrow$
3D-UNet w/ MC dropout | 100 | 0 | 0.950 | 0.899 | 6.04 | 1.53
UA-MT [21] | 5 | 95 | 0.920 | 0.867 | 13.21 | 4.54
Double-UA [18] | 5 | 95 | 0.927 | 0.878 | 12.11 | 4.19
Tripled-UA [17] | 5 | 95 | 0.921 | 0.869 | 11.77 | 3.62
CoraNet [15] | 5 | 95 | 0.923 | 0.877 | 10.84 | 4.28
GBDL | 5 | 95 | 0.935 | 0.884 | 7.89 | 2.42

Table 5: Comparison with state-of-the-art semi-supervised segmentation methods
on the Liver Segmentation dataset. Previous methods in this table are also
based on the regular 3D-UNet with MC dropout.

Method | KiTS19 | Atrial
---|---|---
16 labeleld | 4 labeleld | 16 labeleld | 8 labeleld
UA-MT [21] | 0.868 | 0.841 | 0.851 | 0.836
Double-UA [18] | 0.881 | 0.864 | 0.861 | 0.842
Triplet-UA [17] | 0.874 | 0.852 | 0.864 | 0.847
CoraNet [15] | 0.882 | 0.866 | 0.867 | 0.855
GBDL | 0.897 | 0.877 | 0.882 | 0.864

Table 6:  PAvPU of different Bayesian deep learning methods. Figure 2: The
influence of the number of feedforward passes on: (a) time costs (ms) for
processing a $128\times 128\times 32$ volume input on a single GeForce GTX
1080 Ti; (b) Dice score; and (c) PAvPU metric. The horizontal axis refers to
the number of feedforward passes, namely, $T$.

Besides, as the two 3D-UNet encoders in LRL are of the same architecture (see
the supplementary material), and both aim at extracting features from input
volumes, we evaluated a case where the two encoders share their weights.
Nevertheless, sharing the weights degrades the performance, as shown in Table
2. Similarly, although some layers of the two 3D-UNet decoders in the LRL also
have the same configuration, sharing their weights would clearly lead to a
lower performance, as the respective objectives of them are different, leading
to different gradient descent directions.

Finally, when compared with the baseline model, the proposed GBDL can improve
the performance with a huge margin with respect to the four assessment metrics
for both datasets. So, the proposed method is a good solution for semi-
supervised volumetric medical imaging segmentation.

###  4.4 Comparison with State-of-the-Art Methods

The proposed GBDL is compared with previous state-of-the-art methods, and
since previous works provide no standard deviation, we also only report the
mean values of the four evaluation metrics in Tables 3, 4, and 5. Since most
previous methods are based on VNet, for fair comparison, our regular 3D-UNet
with MC dropout is designed to have similar numbers of parameters.
Furthermore, we also re-implement previous Bayesian deep learning methods
based on the regular 3D-UNet with MC dropout, and according to the results in
Tables 3 and 4, their performance are similar to their VNet counterparts,
meaning that the comparison between our method and previous methods is
relatively fair.

By Tables 3, 4, and 5, the performance of GBDL is higher than previous state-
of-the-art results on the three public medical datasets. When the number of
labeled data decreases, the performance gap between our GBDL and previous
state-of-the-art methods becomes larger, which indicates that our method is
more robust when fewer labeled data were used for training. More importantly,
GBDL performs better than all previous Bayesian deep-learning-based methods
[21, 18, 15, 17], which verifies that GBDL is better than the teacher-student
architecture used in them and that the generative learning paradigm is more
suitable for pseudo-label generation and solving the semi-supervised
segmentation task than the discriminative counterpart.

Furthermore, we compare the patch accuracy vs. patch uncertainty (PAvPU)
metric [13] for GBDL and previous Bayesian-based methods to evaluate their
uncertainty estimation. By Table 6, GBDL can output more reliable and well-
calibrated uncertainty estimates than other methods, so that clinicians can do
a post-processing and refine the results in practice based on the uncertainty
given by GBDL.

Finally, since full Bayesian inference is time-consuming, we show the
relationship among the number of feedforward passes $T$, the time consumption,
and the performance in Figure 2. By Figure 2(a), the time costs rise with
increasing $T$, but by Figure 2(b) and (c), the Dice score and the PAvPU
metric are not greatly impacted by $T$. Thus, $T=5$ can be chosen for saving
time to get the uncertainty in practice, with only a minor performance
degradation.

##  5 Summary and Outlook

In this paper, we rethink the issues in previous Bayesian deep-learning-based
methods for semi-supervised volumetric medical image segmentation, and have
designed a new generative Bayesian deep learning (GBDL) architecture to solve
them. The proposed method outperforms previous state-of-the-art methods on
three public medical benchmarks, showing its effectiveness for handling the
data with only a very limited number of annotations.

However, in GBDL, there is still a limitation with regard to the time
consumption for full Bayesian inference and uncertainty estimation. By Figure
2(a), several numbers of feedforward passes will lead to a high cost in terms
of time. Since Bayesian deep learning is important for some safety-critical
areas, such as medical diagnosis and autonomous driving, future research can
try to find better methods to improve the inference speed for GBDL. Moreover,
since GBDL is a new proposed architecture, future research can also consider
to combine other previous deep-learning-based methods (such as different
training and learning strategies) with GBDL to further improve the
performance.

#### Acknowledgments.

This work was partially supported by the Alan Turing Institute under the EPSRC
grant EP/N510129/1, by the AXA Research Fund, by the EPSRC grant EP/R013667/1,
and by the EU TAILOR grant. We also acknowledge the use of the EPSRC-funded
Tier 2 facility JADE (EP/P020275/1) and GPU computing support by Scan
Computers International Ltd.

## References

  * [1] Wenjia Bai, Ozan Oktay, Matthew Sinclair, Hideaki Suzuki, Martin Rajchl, Giacomo Tarroni, Ben Glocker, Andrew King, Paul M Matthews, and Daniel Rueckert.  Semi-supervised learning for network-based cardiac MR image segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 253–260. Springer, 2017.
  * [2] Gerda Bortsova, Florian Dubost, Laurens Hogeweg, Ioannis Katramados, and Marleen de Bruijne.  Semi-supervised medical image segmentation via learning consistency under transformations.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 810–818. Springer, 2019.
  * [3] Kang Fang and Wu-Jun Li.  DMNet: Difference minimization network for semi-supervised segmentation in medical images.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 532–541. Springer, 2020.
  * [4] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In International Conference on Machine Learning, pages 1050–1059. PMLR, 2016.
  * [5] Wenlong Hang, Wei Feng, Shuang Liang, Lequan Yu, Qiong Wang, Kup-Sze Choi, and Jing Qin.  Local and global structure-aware entropy regularized mean teacher model for 3D left atrium segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 562–571. Springer, 2020.
  * [6] Huimin Huang, Nan Zhou, Lanfen Lin, Hongjie Hu, Yutaro Iwamoto, Xian-Hua Han, Yen-Wei Chen, and Ruofeng Tong.  3D Graph-S2Net: Shape-aware self-ensembling network for semi-supervised segmentation with bilateral graph convolution.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 416–427. Springer, 2021.
  * [7] Alex Kendall and Yarin Gal.  What uncertainties do we need in bayesian deep learning for computer vision?  In Advances in Neural Information Processing Systems, 2017.
  * [8] Philipp Krähenbühl and Vladlen Koltun.  Efficient inference in fully connected CRFs with gaussian edge potentials.  volume 24, pages 109–117, 2011.
  * [9] Shuailin Li, Chuyu Zhang, and Xuming He.  Shape-aware semi-supervised 3D semantic segmentation for medical images.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 552–561. Springer, 2020.
  * [10] Xiangde Luo, Jieneng Chen, Tao Song, and Guotai Wang.  Semi-supervised medical image segmentation through dual-task consistency.  In Proceedings of the AAAI Conference on Artificial Intelligence, pages 8801–8809, 2021.
  * [11] Xiangde Luo, Wenjun Liao, Jieneng Chen, Tao Song, Yinan Chen, Shichuan Zhang, Nianyong Chen, Guotai Wang, and Shaoting Zhang.  Efficient semi-supervised gross target volume of nasopharyngeal carcinoma segmentation via uncertainty rectified pyramid consistency.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 318–329. Springer, 2021.
  * [12] Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi.  V-Net: Fully convolutional neural networks for volumetric medical image segmentation.  In International Conference on 3D Vision (3DV), pages 565–571. IEEE, 2016.
  * [13] Jishnu Mukhoti and Yarin Gal.  Evaluating bayesian deep learning methods for semantic segmentation.  arXiv:1811.12709, 2018.
  * [14] Suman Sedai, Bhavna Antony, Ravneet Rai, Katie Jones, Hiroshi Ishikawa, Joel Schuman, Wollstein Gadi, and Rahil Garnavi.  Uncertainty guided semi-supervised segmentation of retinal layers in oct images.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 282–290. Springer, 2019.
  * [15] Yinghuan Shi, Jian Zhang, Tong Ling, Jiwen Lu, Yefeng Zheng, Qian Yu, Lei Qi, and Yang Gao.  Inconsistency-aware uncertainty estimation for semi-supervised medical image segmentation.  IEEE Transactions on Medical Imaging, 2021.
  * [16] Kihyuk Sohn, Honglak Lee, and Xinchen Yan.  Learning structured output representation using deep conditional generative models.  volume 28, pages 3483–3491, 2015.
  * [17] Kaiping Wang, Bo Zhan, Chen Zu, Xi Wu, Jiliu Zhou, Luping Zhou, and Yan Wang.  Tripled-uncertainty guided mean teacher model for semi-supervised medical image segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 450–460. Springer, 2021.
  * [18] Yixin Wang, Yao Zhang, Jiang Tian, Cheng Zhong, Zhongchao Shi, Yang Zhang, and Zhiqiang He.  Double-uncertainty weighted method for semi-supervised learning.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 542–551. Springer, 2020.
  * [19] Yutong Xie, Jianpeng Zhang, Zhibin Liao, Johan Verjans, Chunhua Shen, and Yong Xia.  Pairwise relation learning for semi-supervised gland segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 417–427. Springer, 2020.
  * [20] Zhaohan Xiong, Qing Xia, Zhiqiang Hu, Ning Huang, Cheng Bian, Yefeng Zheng, Sulaiman Vesal, Nishant Ravikumar, Andreas Maier, Xin Yang, et al.  A global benchmark of algorithms for segmenting the left atrium from late gadolinium-enhanced cardiac magnetic resonance imaging.  Medical Image Analysis, 67:101832, 2020.
  * [21] Lequan Yu, Shujun Wang, Xiaomeng Li, Chi-Wing Fu, and Pheng-Ann Heng.  Uncertainty-aware self-ensembling model for semi-supervised 3D left atrium segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 605–613. Springer, 2019.
  * [22] Xiangyun Zeng, Rian Huang, Yuming Zhong, Dong Sun, Chu Han, Di Lin, Dong Ni, and Yi Wang.  Reciprocal learning for semi-supervised segmentation.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 352–361. Springer, 2021.