11institutetext: State Key Laboratory for Novel Software Technology, Nanjing
University 22institutetext: National Institute of Healthcare Data Science,
Nanjing University
22email<EMAIL_ADDRESS><EMAIL_ADDRESS>33institutetext: School of Computer Science and Engineering, Southeast
University
33email<EMAIL_ADDRESS>44institutetext: School of Data and Computer
Science, Shandong Women’s University
44email<EMAIL_ADDRESS>
# The Devil is in the Statistics: Mitigating and Exploiting Statistics
Difference for Generalizable Semi-supervised Medical Image Segmentation

Muyang Qiu 1122 Jian Zhang 1122 Lei Qi 33
Qian Yu 44 Yinghuan Shi Corresponding Author.1122 Yang Gao 1122

###### Abstract

Despite the recent success of domain generalization in medical image
segmentation, voxel-wise annotation for all source domains remains a huge
burden. Semi-supervised domain generalization has been proposed very recently
to combat this challenge by leveraging limited labeled data along with
abundant unlabeled data collected from multiple medical institutions,
depending on precisely harnessing unlabeled data while improving
generalization simultaneously. In this work, we observe that domain shifts
between medical institutions cause disparate feature statistics, which
significantly deteriorates pseudo-label quality due to an unexpected
normalization process. Nevertheless, this phenomenon could be exploited to
facilitate unseen domain generalization. Therefore, we propose 1) multiple
statistics-individual branches to mitigate the impact of domain shifts for
reliable pseudo-labels and 2) one statistics-aggregated branch for domain-
invariant feature learning. Furthermore, to simulate unseen domains with
statistics difference, we approach this from two aspects, _i.e_., a
perturbation with histogram matching at image level and a random batch
normalization selection strategy at feature level, producing diverse
statistics to expand the training distribution. Evaluation results on three
medical image datasets demonstrate the effectiveness of our method compared
with recent SOTA methods. The code is available at
https://github.com/qiumuyang/SIAB.

###### Keywords:

Medical image segmentation Semi-supervised domain generalization

##  1 Introduction

Medical image segmentation has witnessed substantial advancements recently,
driven by the advent of deep convolutional neural networks [26, 23, 6].
However, the efficacy of existing methods can be impeded by the distribution
shift between training and test datasets due to inherent differences in
imaging protocols, modalities, and patient populations across medical
institutions, which is frequently encountered in real-world clinical
scenarios. Therefore, domain generalization (DG) has been proposed to tackle
this challenge [33, 10], aiming to train models that are generalizable to
unseen domains only using data collected from multiple relevant source
domains.

Although existing DG methods present practical solutions to combat domain
shifts, they often require extensive annotations for each source domain. This
is particularly expensive in medical image segmentation due to the need for
precise delineation and expert knowledge [31]. Semi-supervised domain
generalization (SSDG) [43, 18], harmonizing the strengths of semi-supervised
learning (SSL) [8] and DG, has emerged as a promising new direction to
alleviate this annotation burden. It extends the capabilities of conventional
DG by leveraging the abundant but unlabeled information within each source
domain.

(a)

(b)

(c)

Figure 1: (a) Images sampled from 4 domains of the Fundus dataset. (b)
Differences in feature statistics captured by batch normalization layers on
the Fundus dataset. (c) Comparison of pseudo-label quality (measured by dice
similarity coefficient against ground truth) between separate and mixed domain
feature normalization.

Typically, applying well-established SSL or DG medical image segmentation
methods to SSDG exposes limitations arising from the unique challenges posed
by SSDG. DG methods, although focusing on tackling domain shifts and enhancing
generalization, usually rely on full label information [34, 16] or indirectly
utilize unlabeled data through self-supervised tasks [44]. SSL methods, while
specialized in exploiting unlabeled data, are not explicitly designed to
handle the complexities of multiple domains and may suffer from the
distribution shift [2]. To tackle the aforementioned issues, the main goals of
SSDG are twofold: _improving the quality of pseudo-labels under the domain
shift_ and simultaneously _promoting generalization ability utilizing all
available data_. Prior works for SSDG are primarily built upon SSL methods
[43, 18, 40], guided by the belief that a more robust decision boundary can be
learned by exploiting all available data. For instance, Yao  _et al_. [40]
incorporated a Fourier transform-based confidence-aware module into CPS [7],
filtering high-confidence pixels to obtain more reliable pseudo supervision
signals for unlabeled data in medical image segmentation. Despite their
success, the enhancement of pseudo-labels is achieved through post-processing
of the model’s outputs, refining these predictions as a secondary step, which
limits the quality of the initial results. This raises a crucial question: _Is
there a more direct and graceful approach to seamlessly mitigate the impact of
domain shifts throughout the prediction process?_

It is well known that batch normalization (BN) plays an important role
throughout deep models by accelerating convergence and stabilizing training
[12]. Its function lies in normalizing input features via batch statistics,
_i.e_., mean and variance. However, our observation reveals a significant
challenge when mini-batches are drawn from diverse domains: domain shifts
induce substantial differences in feature statistics within these batches, as
depicted in Figs. 1(a) and 1(b). For multi-site medical images, domain shifts
mainly manifest as grayscale differences. While these differences might be
imperceptible to the human eye, they are prominently captured by the BN.
Adopting one single BN for all source domains disregards these differences,
resulting in a mismatch between the normalization statistics optimal for each
domain and the statistics actually used. Consequently, it adversely affects
the normalized features, leading to a negative impact on pseudo-labels, as in
Fig. 1(c).

Motivated by this, we propose statistics-individual branches (SIBs) to
separately normalize image features for each medical institution, alleviating
the degradation of pseudo-label quality caused by the undesired normalization
process. We equip each SIB with dedicated BNs, thereby enabling independent
capture and adaptation to the statistical nuances of features from each
distinct domain. This design directly diminishes the impact of domain shifts
during prediction, minimizing the need for additional output processing.

Besides, we introduce a statistics-aggregated branch (SAB) as a complement to
SIBs. Trained on data from all medical institutions, SAB is tasked with
generating predictions consistent with those of SIBs, aimed at learning
domain-invariant features and contributing to improved generalization
performance.

While statistics discrepancies can be detrimental to pseudo-labels, we
identify an opportunity to boost overall generalization by artificially
simulating domain shifts during the training stage. Therefore, we further
introduce additional perturbations to create statistical variations. SAB is
then required to maintain consistency despite both suboptimal normalization
statistics aggregated from multiple domains and these perturbations,
facilitating robust feature learning. We integrate histogram matching at the
image level for grayscale distribution variations. Leveraging SIBs, we develop
a feature-level perturbation method utilizing affine parameters in various
BNs. These parameters, adjusting the mean and variance of features, implicitly
model the unique feature distribution of their respective domains, hence ideal
for introducing targeted perturbations.

Our contributions can be summarized as follows:

  * •

Based on the pros and cons of domain statistics difference, we propose a
multi-branch framework for SSDG medical image segmentation, utilizing both
statistics-individual and statistics-aggregated branches to collaboratively
preserve pseudo-labels and improve generalization.

  * •

We propose additional multi-level perturbation streams at both the image and
feature levels to simulate domain shifts, expanding the training distribution
to achieve better robustness for the statistics-aggregated branch.

  * •

We demonstrate the effectiveness of our method on three medical image
segmentation benchmarks by comparison with several classic and state-of-the-
art methods.

##  2 Related Work

Semi-supervised Medical Image Segmentation. Existing methods for semi-
supervised medical image segmentation can be broadly categorized into two main
streams: self-training [13, 1, 28] and consistency regularization [4, 20, 41].
Self-training methods involve iterative training by generating pseudo-labels
or using other techniques. For instance, Wang  _et al_. [36] developed a
mutual correction framework to explore network bias correction for pseudo-
labels. Consistency regularization-based methods generate diverse predictions
through input or feature perturbations, multiple tasks, or network structure,
and enforce consistency among them. Li  _et al_. [15] perturbed inputs using
various transforms to encourage consistency between the student and teacher
model. Luo  _et al_. [20] established dual tasks for segmentation and level-
set regression. Wu  _et al_. [38] integrated an extra decoder to the
segmentation network to obtain discrepancies. Bai  _et al_. [2] highlighted
the distribution gap between labeled and unlabeled data, tackling it with a
bidirectional copy-paste strategy. Most semi-supervised methods neglect
distribution differences between domains, risking performance degradation
under domain shifts. However, our method takes this into account and minimizes
the impact of domain shifts during pseudo-label generation.

Domain Generalization. Extensive efforts have been made in the field of DG.
Several methods [30, 21] incorporate data augmentation techniques to broaden
the data distribution and enhance the robustness of models. Additionally,
advances in episodic training strategy have driven the emergence of meta-
learning-based methods [14, 9, 17, 35]. Liu _et al_. [16] utilized this
strategy with continuous frequency space interpolation in a federated DG
scenario. Learning domain-invariant features by minimizing the domain
discrepancy across multiple source domains is another common approach [34].
Another category of DG methods constructs self-supervised tasks for
regularization. Zhou  _et al_. [44] introduced an image restoration task after
distortion in the frequency domain. While self-supervised tasks offer certain
advantages in SSDG by leveraging unlabeled data, DG methods still fall behind
in harnessing unlabeled data reliably. Besides, several methods employ
additional network structures such as convolution and normalization layers. Hu
_et al_. [11] tackled DG using a dynamic convolution with parameters
conditioned on predicted domain code and global features. Seo  _et al_. [27]
incorporated a mixture of batch and instance normalization for each domain and
averaged their outputs as the prediction for the unseen domain. In contrast,
our design introduces a different structure, including both individual batch
normalization layers to enhance pseudo-labels and an inter-domain shared layer
for generalization, which is aligned with the two major objectives of SSDG.

Figure 2: Overview of our method. The left figure illustrates the structure of
the statistics-individual and aggregated branches. The right figure
illustrates the consistency regularization framework extended with diverse
perturbations. $u_{w},u_{s},u_{h}$ denote the weak, strong, and histogram
matching-augmented unlabeled images respectively, and $p_{w},p_{s},p_{h}$ are
the corresponding predictions. $p_{r}$ is the prediction of $u_{w}$ through
random forward. Arrows symbolize the forward processes based on the proposed
branches during prediction. Dashed lines signify that one branch is randomly
selected each time. $\bigoplus$ denotes element-wise addition.

##  3 Method

We start by presenting the preliminaries, followed by the introduction of our
method in Sec. 3.2 and Sec. 3.3. Finally, we briefly discuss the complexity of
our method in Sec. 3.4. Fig. 2 provides an overview of the method.

###  3.1 Preliminaries

We denote the training set consisting of $K$ source domains as
$\mathcal{S}=\\{\mathcal{D}^{1},\mathcal{D}^{2},$
$\cdots,\mathcal{D}^{K}\\}$. Each source domain consists of the labeled part
$\mathcal{D}^{k}_{l}=\\{(x^{k}_{i},y^{k}_{i})\\}^{N^{k}_{l}}_{i=1}$ and the
unlabeled part $\mathcal{D}^{k}_{u}=\\{u^{k}_{i}\\}^{N^{k}_{u}}_{i=1}$, where
$x^{k}_{i},y^{k}_{i},u^{k}_{i}$ are the $i$-th labeled image, segmentation map
and unlabeled image, respectively. $N^{k}_{l}$ and $N^{k}_{u}$ are the number
of labeled and unlabeled samples in the $k$-th domain. We expect to train a
model generalizable to an unseen test domain, exploiting the unlabeled images
with limited labeled data in the presence of domain shifts.

We first set up a mean-teacher-based [32] semi-supervised baseline method for
the leverage of unlabeled data. The overall objective function for the
baseline can be formulated as

$\mathcal{L}=\mathcal{L}_{x}+\mathcal{L}_{u},$ (1)

where $\mathcal{L}_{x}$ is the supervised loss, usually a combination of
cross-entropy loss and Dice loss [23] for medical image segmentation, and
$\mathcal{L}_{u}$ is the unsupervised loss. We embrace the “weak-to-strong”
consistency regularization [29, 39] for the unsupervised loss, given that
strong augmentation changes the image appearance (in brightness, contrast,
saturation, _etc_.) by color jittering, aligning well with the manifestation
of domain shifts in multi-site medical images.

To be specific, we first randomly sample a pair of weak and strong
augmentation, denoted as $\mathcal{A}_{w}$ and $\mathcal{A}_{s}$,
respectively. Subsequently, utilizing the teacher model $\hat{f}$ which is
updated through exponential moving average (EMA) using the student’s weights,
we generate predictions for the weak-augmented unlabeled image
$\mathcal{A}_{w}(u)$. Finally, we guide the student model $f$ to predict the
strong-augmented counterpart $\mathcal{A}_{s}(u)$ with the pseudo-label
obtained from the teacher model, by enforcing the cross-entropy consistency
loss $\mathrm{H}\left(\cdot,\cdot\right)$ between them:

$\mathcal{L}_{u}=\mathrm{H}\left(\hat{f}\left(\mathcal{A}_{w}(u)\right),f\left(\mathcal{A}_{s}(u)\right)\right).$
(2)

In practice, we use one-hot pseudo-label and confidence thresholding:

$\displaystyle\mathcal{L}_{u}$
$\displaystyle=\mathbbm{1}\left(\max\left(p_{w}\right)\geqslant\tau\right)\mathrm{H}\left(\hat{p}_{w},f\left(\mathcal{A}_{s}(u)\right)\right),$
(3) $\displaystyle p_{w}$
$\displaystyle=\hat{f}\left(\mathcal{A}_{w}(u)\right),\quad\hat{p}_{w}=\mathop{\mathrm{argmax}}(p_{w}),$

where $\mathbbm{1}(\cdot)$ is the indicator function, $\tau$ denotes the
confidence threshold, $p_{w}$ is the prediction of the weakly augmented
unlabeled image, and $\hat{p}_{w}$ is the one-hot pseudo-label, respectively.
We will omit the indicator function term and abbreviate it as the form of Eq.
2 in the following sections.

###  3.2 Statistics-individual and Statistics-aggregated Branches

Statistics-individual branches for pseudo-labeling. To alleviate the impact of
domain shifts on pseudo-labels, we propose the incorporation of statistics-
individual branches. Each branch includes dedicated BN layers throughout the
forward propagation process to capture domain statistics individually while
sharing the remaining network parameters across all branches. For the batch
activation $\mathbf{x}_{d}\in\mathbb{R}^{N\times C\times H\times W}$ from
domain $d$, where $N,C,H,W$ are the batch size, number of channels, height and
width of the feature map, respectively, BNs in the $d$-th SIB normalizes it by
channel-wise mean $\mu_{d}$ and variance $\sigma^{2}_{d}$:

$\mathrm{Norm}_{\mathrm{ind}}(\mathbf{x}_{d})=\gamma_{d}\frac{\mathbf{x}_{d}-\mu_{d}}{\sqrt{\sigma_{d}^{2}+\epsilon}}+\beta_{d},$
(4)

where $\gamma_{d}$ and $\beta_{d}$ are the affine parameters, $\epsilon$ is a
small positive constant to avoid division by zero, and the mean and variance
of channel $c$ are calculated by:

$\displaystyle\mu_{d,c}$
$\displaystyle=\frac{1}{NHW}\sum\limits_{n=1}^{N}\sum\limits_{h=1}^{H}\sum\limits_{w=1}^{W}\mathbf{x}_{d}[n,c,h,w],$
(5) $\displaystyle\sigma_{d,c}^{2}$
$\displaystyle=\frac{1}{NHW}\sum\limits_{n=1}^{N}\sum\limits_{h=1}^{H}\sum\limits_{w=1}^{W}\Big{(}\mathbf{x}_{d}[n,c,h,w]-\mu_{d,c}\Big{)}^{2}.$

We exclusively forward the weak-augmented images through SIBs, as weak
augmentations (mainly geometric transformations) minimally affect the style
and statistical properties of input images. Thus, images within each domain
maintain a relatively consistent distribution without disturbing the
individual branches. Conversely, strong augmentations introduce significant
statistical variations as discussed in [42], making the augmented images
unsuitable for the SIBs.

Statistics-aggregated branch for generalization. We then introduce a SAB to
complement the SIBs. Considering the complexity of the input distribution it
handles, which includes samples from various domains and those affected by
strong augmentation, we leverage both batch and instance statistics for
normalization. This approach suppresses instance-specific style information
with instance statistics while preserving inter-instance discriminability with
batch statistics. We adopt a linear combination of them [24, 27] as the final
normalization statistics in the SAB. Denoting the instance mean and variance
as $\mu_{\mathrm{inst}}$ and $\sigma_{\mathrm{inst}}^{2}$, the normalization
process can be formulated as:

$\begin{gathered}\mathrm{Norm}_{\mathrm{agg}}(\mathbf{x}_{d})=\gamma_{\,\mathrm{agg}}\frac{\mathbf{x}_{d}-\mu_{\mathrm{agg}}}{\sqrt{\sigma_{\mathrm{agg}}^{2}+\epsilon}}+\beta_{\mathrm{agg}},\\\
\mu_{\mathrm{agg}}=\alpha\odot\mu_{d}+(1-\alpha)\odot\mu_{\mathrm{inst}},\\\
\sigma^{2}_{\mathrm{agg}}=\alpha\odot\sigma^{2}_{d}+(1-\alpha)\odot\sigma^{2}_{\mathrm{inst}},\\\
\end{gathered}$ (6)

where $\gamma_{\,\mathrm{agg}}$ and $\beta_{\mathrm{agg}}$ are the affine
parameters. $\odot$ means element-wise multiplication and $\alpha\in[0,1]^{C}$
is the learnable mixing coefficient for each channel.

Consistency learning. We denote predicting the sample using its respective SIB
as individual-forward ($\mathcal{IF}$) and using SAB as aggregated-forward
($\mathcal{AF}$). We then denote $m^{\mathcal{F}}(x)$ as the output of
applying the forward strategy $\mathcal{F}$ to model $m$ with input $x$. The
procedure of weak-to-strong consistency learning can be improved with the
proposed statistics-individual and aggregated branches as follows. We first
acquire predictions for weak-augmented unlabeled images using $\mathcal{IF}$
and $\mathcal{AF}$ via the teacher model $\hat{f}$:

$\begin{gathered}p_{w}^{\mathcal{IF}}=\hat{f}^{\mathcal{IF}}(\mathcal{A}_{w}(u)),\quad
p_{w}^{\mathcal{AF}}=\hat{f}^{\mathcal{AF}}(\mathcal{A}_{w}(u)).\end{gathered}$
(7)

We then ensemble them for better pseudo-labels:

$p_{w}=\left(p_{w}^{\mathcal{IF}}+p_{w}^{\mathcal{AF}}\right)/2,\quad\hat{p}_{w}=\mathop{\mathrm{argmax}}(p_{w}).$
(8)

Finally, we encourage a consistent prediction of strong-augmented images from
SAB of the student model $f$ in alignment with the pseudo-label:

$p_{s}=f^{\mathcal{AF}}(\mathcal{A}_{s}(u)),\quad\mathcal{L}_{s}=\mathrm{H}(\hat{p}_{w},p_{s}).$
(9)

Supervised loss. We train both statistics-individual and aggregated branches
with cross-entropy and Dice loss using weak-augmented labeled images to ensure
stable and accurate statistics information.

Figure 3: Illustration of histogram matching. The top row represents the
images from the Prostate dataset. The bottom row visualizes their grayscale
distribution.

Figure 4: T-SNE visualization of image features after applying different
perturbations. Different colors represent different domains. Best viewed in
color.

###  3.3 Multi-level Consistency Regularization

To facilitate unseen domain generalization, we aim to fully exploit the
potential of statistics differences, simulating various distributions by
introducing multi-level perturbations and consistency regularization.

Strong and style augmentation at the image level. We have previously discussed
strong augmentations, which involve random appearance perturbations to
diversify training data distribution and assist the model in adapting to
domain shifts. However, an excessive level of perturbation resulting from
randomness can compromise the semantic integrity of images, impeding effective
semantic learning. To overcome this limitation, we leverage the diverse
distribution information from existing source domains for a controllable
perturbation. Specifically, we employ histogram matching [25] to align the
grayscale distribution between images, as illustrated in Fig. 4. We then
introduce a new regularization stream named “weak-to-style” into the
consistency regularization framework. By exposing the model to images with
similar semantics but diverse styles across multiple domains, it is
anticipated that the model will learn domain-invariant features, robust
against various perturbations introduced by both color jittering and histogram
matching. Furthermore, while color jittering diversifies source domain
distribution through randomness, histogram matching bridges the distribution
gap between source domains by leveraging known information. Both augmentation
strategies synergize effectively, forming a complementary approach to
enhancing model robustness. Denoting the histogram matching operation as
$\mathcal{A}_{h}$, the loss function of weak-to-style consistency
regularization is similar to Eq. 9:

$p_{h}=f^{\mathcal{AF}}(\mathcal{A}_{h}(u)),\quad\mathcal{L}_{h}=\mathrm{H}(\hat{p}_{w},p_{h}).$
(10)

Normalization-selection-based perturbation at the feature level. Despite style
diversity being improved by strong and style augmentations, both consistency
regularization streams are confined to the image level, hindering the model
from exploring a wider range of perturbations.

To this end, we propose to further leverage SIBs for feature-level statistics
perturbation. Affine parameters within distinct BN layers inherently
encapsulate variations in features across diverse domains. Hence, opting for
the BN layer and utilizing its affine parameters to introduce feature
perturbations becomes a reasonable strategy. To be specific, for the batch
activation $\mathbf{x}_{d}$ from domain $d$ which should have been normalized
by the BNs in the $d$-th SIB, we randomly select another branch $d^{\prime}$,
using its affine parameters to substitute the ones in Eq. 4 as the source of
perturbation with a probability of $p$, which controls the intensity of the
perturbation:

$\begin{gathered}\mathrm{Norm}_{\mathrm{rand}}(\mathbf{x}_{d})=\gamma_{d^{\prime}}\frac{\mathbf{x}_{d}-\mu_{d}}{\sqrt{\sigma_{d}^{2}+\epsilon}}+\beta_{d^{\prime}},\\\
d^{\prime}\neq d,d^{\prime}\in\\{1,2,\cdots,K\\}.\end{gathered}$ (11)

We stop the gradient for BN affine parameters during backward propagation,
regardless of whether the perturbation is activated, to prevent interference
with SIBs.

The feature perturbed forward process with the random BN selection strategy is
denoted as random-forward ($\mathcal{RF}$). Its corresponding loss
$\mathcal{L}_{r}$, is defined as:

$p_{r}=f^{\mathcal{RF}}(\mathcal{A}_{w}(u)),\quad\mathcal{L}_{r}=\mathrm{H}(\hat{p}_{w},p_{r}).$
(12)

Fig. 4 illustrates the t-SNE visualization results of the image features
extracted from the Fundus dataset. It showcases the effects of various
perturbations in introducing statistics discrepancies and simulating domain
shifts.

In summary, the unsupervised loss contains three terms, each representing one
stream of consistency regularization:

$\mathcal{L}_{u}=\mathcal{L}_{s}+\lambda_{h}\mathcal{L}_{h}+\lambda_{r}\mathcal{L}_{r},$
(13)

where $\lambda_{h},\lambda_{r}$ balance the weight between different
consistency streams.

###  3.4 Complexity Discussion

The primary purpose of the SIBs during training is to provide accurate pseudo-
labels for SAB. These branches are tailored to their respective domains and
thus may not generalize well to unseen distributions. Consequently, during the
inference stage, we discard them to attain a parameter size similar to that of
a conventional model structure. Quantitative results are reported in Sec. 4.4.

Table 1: Performance comparisons with the state-of-the-art methods on the
Prostate dataset. The best performance is marked as bold.

Prostate | Dice (%) | ASD (voxel)
---|---|---
A | B | C | D | E | F | Avg. | A | B | C | D | E | F | Avg.
30% Labeled
MT [32] |  | 79.20 | 57.20 | 47.03 | 86.51 | 12.81 | 84.59 | 61.23 | 1.17 | 1.59 | 6.73 | 3.81 | 43.65 | 0.84 | 9.63
UA-MT [41] |  | 83.14 | 64.80 | 68.97 | 87.64 | 26.57 | 85.47 | 69.43 | 0.99 | 1.40 | 8.38 | 1.28 | 20.94 | 0.97 | 5.66
SS-Net [37] |  | 88.49 | 86.27 | 78.76 | 89.75 | 53.27 | 87.85 | 80.73 | 0.76 | 0.87 | 7.39 | 1.39 | 4.52 | 3.52 | 3.08
BCP [2] |  | 88.38 | 87.11 | 64.01 | 89.59 | 69.12 | 84.89 | 80.52 | 3.08 | 0.74 | 19.29 | 1.13 | 5.94 | 11.09 | 6.88
CauSSL [22] |  | 83.67 | 64.34 | 48.97 | 84.37 | 29.11 | 70.27 | 63.46 | 3.55 | 2.73 | 18.29 | 2.15 | 18.15 | 1.57 | 7.74
FedDG [16] |  | 83.12 | 71.15 | 72.62 | 84.67 | 68.45 | 88.53 | 78.09 | 7.34 | 3.74 | 11.01 | 4.04 | 5.26 | 2.18 | 5.60
RAM-DSIR [44] |  | 83.16 | 83.18 | 67.76 | 88.11 | 79.34 | 73.20 | 79.12 | 1.55 | 0.84 | 6.21 | 0.98 | 3.36 | 1.37 | 2.39
CACPS [40] |  |  90.96$\pm$2.96 |  88.18$\pm$4.60 |  79.20$\pm$14.3 |  89.16$\pm$3.46 |  81.12$\pm$11.7 |  90.19$\pm$3.02 | 86.47 |  1.47$\pm$1.58 |  0.66$\pm$0.33 |  3.65$\pm$2.30 |  1.73$\pm$2.21 |  4.21$\pm$3.06 |  1.01$\pm$1.09 | 2.12
Ours |  |  90.58$\pm$2.47 |  90.12$\pm$2.95 |  85.72$\pm$7.11 |  89.87$\pm$2.90 |  86.06$\pm$4.52 |  90.01$\pm$2.77 | 88.73 |  0.74$\pm$0.46 |  0.57$\pm$0.23 |  3.14$\pm$2.35 |  0.72$\pm$0.33 |  5.81$\pm$5.14 |  0.63$\pm$0.23 | 1.93

Figure 5: Visualization results from the Prostate dataset.

##  4 Experiment

###  4.1 Datasets

We evaluate our method on three datasets, namely, Prostate [17], Fundus [34],
and M&Ms [5]. The first two datasets are commonly used in conventional DG
medical image segmentation tasks [34, 16, 44], while the other in SSDG [18,
40].

Prostate dataset collects prostate T2-weighted MRI images from 6 medical
centers for prostate segmentation. The original 3D volumes are sliced in the
axial plane and resized to $288\times 288$ for training. At test time, slices
not containing any prostate region are discarded, following [17, 44].

Fundus dataset collects retinal fundus images from 4 medical centers for optic
cup and disc (OC and OD) segmentation. The original resolution of images is
$800\times 800$ and we resize them to $256\times 256$ for memory efficiency at
train time. At test time, the input images are downsampled and the predicted
masks are then upsampled for evaluation.

M&Ms dataset collects MRI cardiac images from 4 magnetic resonance scanner
vendors. Each subject is presented in four dimensions (3D volumes with phase
information) and only the end-systole with the end-diastole phases are
annotated. Each volume is sliced, center-cropped, and then resized to
$288\times 288$ as the network input. We follow [40] to split the labeled and
unlabeled data.

###  4.2 Experimental Setting

Following the established convention in DG, we apply the leave-one-domain-out
strategy to evaluate the performance of a method, _i.e_., choose one domain as
the unseen test domain and the rest for training. For each training source
domain, we use the same unlabeled split consistently among different methods.

Implementation details. We employ U-net [26] as the basic segmentation model
and implement it with PyTorch on two NVIDIA GeForce RTX 3090 GPUs. The
training iterations for each dataset are $2\times 10^{4},2\times
10^{4},4\times 10^{4}$, and the batch sizes are $12,10,18$, respectively. We
adopt AdamW [19] with a learning rate of $1\times 10^{-3}$ and a weight decay
of 0.01 as the optimizer. We set the parameters $\lambda_{h}$ and
$\lambda_{r}$ in Eq. 13 as 1.0 and 0.2, respectively. Weak augmentation
employed in our consistency regularization framework consists of random
rotation and flip, and strong augmentation includes color jittering and blur.

Evaluation metrics. We adopt two metrics for evaluation. Dice similarity
coefficient (Dice) measures the overlap between two binary masks. Average
surface distance (ASD) measures the average distance between the surfaces of
two segmented objects. Higher Dice or lower ASD indicates better performance.

Table 2: Performance (OC / OD) comparisons with the state-of-the-art methods
on the Fundus dataset. The best performance is marked as bold.

Fundus | Dice (%) | ASD (pixel)
---|---|---
A | B | C | D | Avg. | A | B | C | D | Avg.
30% Labeled
MT [32] |  | 59.50 / 82.25 | 60.69 / 80.02 | 76.77 / 87.77 | 63.30 / 75.55 | 73.23 | 11.31 / 0.04 | 4.09 / 2.98 | 3.36 / 2.28 | 91.29 / 73.06 | 23.55
UA-MT [41] |  | 64.57 / 87.36 | 56.29 / 77.18 | 75.97 / 87.44 | 64.10 / 66.97 | 72.49 | 9.34 / 0.06 | 7.22 / 5.35 | 2.98 / 1.48 | 56.05 / 66.51 | 18.62
SS-Net [37] |  | 71.64 / 92.37 | 57.84 / 69.03 | 80.81 / 90.75 | 63.65 / 65.12 | 73.90 | 6.12 / 0.07 |  2.60 / 1.60 | 1.65 / 0.30 | 47.83 / 66.92 | 15.89
BCP [2] |  |  77.83 / 95.14 | 64.22 / 74.58 | 81.61 / 91.04 | 80.25 / 90.86 | 81.94 | 4.43 / 0.21 | 4.46 / 3.27 | 2.57 / 0.65 | 2.11 / 1.46 | 2.39
CauSSL [22] |  | 65.05 / 90.13 | 60.89 / 76.96 | 74.30 / 86.64 | 56.94 / 69.29 | 72.53 | 9.48 / 1.04 | 6.54 / 4.33 | 2.92 / 0.99 | 47.67 / 46.52 | 14.94
FedDG [16] |  |  80.46 / 93.76 | 66.60 / 84.61 | 82.52 / 90.90 | 79.62 / 90.50 | 83.62 |  3.61 / 0.29 | 8.48 / 5.75 | 1.34 / 0.23 | 4.60 / 6.09 | 3.80
RAM-DSIR [44] |  | 74.65 / 94.38 | 67.64 / 83.92 | 80.22 / 89.69 | 63.68 / 89.64 | 80.48 | 5.46 / 0.04 | 8.84 / 6.42 |  1.19 / 0.02 | 1.74 / 1.42 | 3.14
CACPS [40] |  | 73.03 / 93.39 | 69.81 / 84.53 | 82.62 / 91.52 | 84.95 / 93.10 | 84.12 | 5.54 / 0.04 | 4.36 / 2.13 | 1.28 / 0.03 | 1.67 / 0.85 | 1.99
Ours |  | 74.05 / 93.85 |  73.53 / 88.32 |  82.85 / 91.75 |  86.68 / 93.79 | 85.60 |  5.99 / 0.04 | 4.21 / 2.70 | 1.25 / 0.07 |  1.42 / 0.79 | 2.06

Figure 6: Visualization results from the Fundus dataset. Blue contours for
optic cup and green for optic disc.

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

Compared methods. We compare our method with classic and state-of-the-art
(SOTA) SSL, DG, and SSDG methods for medical image segmentation. SSL methods
utilize both labeled and unlabeled data from multiple domains without
distinction. DG methods, typically for fully supervised scenarios, do not
utilize unlabeled data unless incorporating self-supervised tasks.

Results on the Prostate dataset. We allocate 30% labeled data and 70%
unlabeled for each domain in the Prostate dataset. The results are listed in
Tab. 1. SSL methods exhibit limited performance on challenging domains,
particularly domain E. DG methods, while effective on some domains, cannot
fully leverage unlabeled data to achieve competitive overall results. Our
method achieves the best Dice and ASD, with an improvement of 8.00% and 9.61%
in terms of average Dice compared to the best-performing SSL and DG methods,
respectively. Additionally, we outperform CACPS, the previous SOTA in SSDG, by
2.26%. Segmentation examples are shown in Fig. 5.

Results on the Fundus dataset. In Tab. 2, we report the comparison results on
the Fundus dataset. Among SSL methods, BCP stands out by creating challenging
edges through copy-paste. DG methods surpass most of the SSL methods except
BCP, likely due to the regular shapes of the optic cup and disc, making the
task less challenging and allowing effective learning with even a few labeled
data. RAM-DSIR, despite utilizing unlabeled data, underperforms compared to
FedDG, possibly due to the consistency loss without ground truth guidance. Our
method achieves the best Dice while CACPS achieves the best ASD. Fig. 6
displays the visualization results.

Table 3: Performance comparisons with the state-of-the-art methods on the M&Ms
dataset. The best performance is marked as bold (same for the following
tables).

M&Ms | Dice (%) | ASD (voxel)
---|---|---
A | B | C | D | Avg. | A | B | C | D | Avg.
2% Labeled
SS-Net [37] |  | 52.71 | 70.63 | 71.47 | 77.72 | 68.13 | 2.53 | 0.66 | 1.40 | 0.70 | 1.32
BCP [2] |  | 47.62 | 83.57 | 80.00 | 81.45 | 73.16 | 1.83 | 1.54 | 1.97 | 2.66 | 2.00
FedDG [16] |  | 60.04 | 65.40 | 67.46 | 61.61 | 63.63 | 4.07 | 3.71 | 7.41 | 6.89 | 5.52
RAM-DSIR [44] |  | 24.13 | 68.43 | 52.16 | 46.41 | 47.78 | 7.59 | 1.97 | 6.52 | 1.96 | 4.51
CACPS [40] |  |  82.35$\pm$6.24 |  82.84$\pm$7.59 |  86.31$\pm$5.47 |  86.58$\pm$4.78 | 84.52 |  1.32$\pm$0.41 |  1.13$\pm$0.87 |  1.53$\pm$1.94 |  0.73$\pm$0.66 | 1.18
Ours |  |  84.92$\pm$4.19 |  86.80$\pm$3.78 |  84.17$\pm$5.18 |  86.83$\pm$2.07 | 85.68 |  0.54$\pm$0.28 |  0.42$\pm$0.23 |  0.94$\pm$1.80 |  0.47$\pm$0.21 | 0.59
5% Labeled
MT [32] |  | 30.59 | 66.46 | 70.61 | 61.88 | 57.39 | 13.37 | 1.60 | 7.41 | 1.82 | 6.05
UA-MT [41] |  | 51.76 | 45.98 | 28.43 | 48.78 | 43.73 | 10.80 | 3.41 | 10.37 | 5.52 | 7.53
SS-Net [37] |  | 67.66 | 86.38 | 82.70 | 83.85 | 80.15 | 1.89 | 0.90 | 2.97 | 1.03 | 1.70
BCP [2] |  | 70.15 | 87.26 | 82.71 | 82.53 | 80.66 | 2.16 | 1.32 | 3.61 | 1.29 | 2.09
FedDG [16] |  | 75.44 | 83.76 | 81.07 | 76.89 | 79.29 | 1.70 | 1.92 | 2.57 | 4.93 | 2.78
RAM-DSIR [44] |  | 53.93 | 81.66 | 74.22 | 82.97 | 73.19 | 2.67 | 0.78 | 0.66 | 1.22 | 1.33
CACPS [40] |  |  83.30$\pm$5.83 |  85.04$\pm$6.49 |  87.14$\pm$4.74 |  87.38$\pm$4.49 | 85.81 |  1.01$\pm$0.34 |  0.61$\pm$0.75 |  1.26$\pm$1.86 |  0.36$\pm$0.56 | 0.81
Ours |  |  85.27$\pm$3.54 |  87.84$\pm$3.07 |  86.03$\pm$4.20 |  87.57$\pm$2.25 | 86.68 |  0.51$\pm$0.31 |  0.40$\pm$0.16 |  0.85$\pm$1.64 |  0.37$\pm$0.08 | 0.53

Figure 7: Visualization results from the M&Ms dataset.

Results on the M&Ms dataset. We conduct experiments on the M&Ms dataset using
2% and 5% split ratios following [40]. MT and UA-MT fail to generalize on the
unseen domain on 2% labeled data. Performance improves with more labeled
samples for all methods, underscoring the importance of accurate supervision
and better pseudo-label quality. Our method outperforms all others on averaged
metrics for both data splits. Quantitative and visualized results are
presented in Tab. 3 and Fig. 7 respectively.

###  4.4 Ablation Studies

Effectiveness of each component. We perform an ablation study to evaluate the
effectiveness of each component on the Fundus dataset. It can be observed from
Tab. 9 that the statistics-individual and aggregated branches (SIAB) improve
the baseline method across all four domains. Additionally, the incorporation
of histogram matching-based style augmentation (Style) contributes to enhanced
generalization. Further improvements are observed with perturbations at the
feature level through the random BN selection strategy (RandBN). Combining
these three components forms the final framework of our method.

Normalization statistics in statistics-aggregated branch. In Tab. 9, we
present experiment results about the choice of statistics used for
normalization in SAB. Performance when solely using instance statistics for
normalization is slightly inferior to using batch statistics. We believe this
could be attributed to a small amount of task-relevant information contained
within the entire batch, and normalizing with instance statistics dilutes such
information. The integration of batch and instance statistics combines the
individual strengths of both.

Table 4: Ablation study (Dice %) on the Fundus dataset.

SIAB | Style | RandBN | A | B | C | D | Avg.
---|---|---|---|---|---|---|---
|  |  | 80.60 | 79.79 | 84.90 | 87.84 | 83.28
✓ |  |  | 82.65 | 81.53 | 85.57 | 88.49 | 84.56
| ✓ |  | 81.95 | 78.53 | 85.92 | 89.12 | 83.88
✓ | ✓ |  | 83.37 | 80.80 | 87.97 | 88.75 | 85.22
✓ | ✓ | ✓ | 83.95 | 80.93 | 87.30 | 90.24 | 85.60

Table 5: Comparison (Dice %) of different statistics for normalization in the
SAB on the Fundus dataset.

Statistics | A | B | C | D | Avg.
---|---|---|---|---|---
batch | 81.73 | 80.09 | 84.50 | 87.69 | 83.50
instance | 81.37 | 79.70 | 84.42 | 87.77 | 83.31
batch & instance | 82.65 | 81.53 | 85.57 | 88.49 | 84.56

Table 6: Comparison of different weights for RandBN on the Fundus dataset.

$\lambda_{r}$ | 0 | 0.05 | 0.1 | 0.2 | 0.5 | 1.0
---|---|---|---|---|---|---
Dice (%) | 85.22 | 85.35 | 85.39 | 85.60 | 85.57 | 85.50

Table 7: Comparison (Dice %) of different image-level consistency
regularization on the Fundus dataset.

Aug1 | Aug2 | A | B | C | D | Avg.
---|---|---|---|---|---|---
strong | - | 82.42 | 80.98 | 84.46 | 89.06 | 84.23
style | - | 82.11 | 78.94 | 85.02 | 89.28 | 83.84
strong | strong | 83.27 | 80.75 | 85.37 | 89.66 | 84.76
strong | style | 83.95 | 80.93 | 87.30 | 90.24 | 85.60

Table 8: Comparison of different probability $p$ for RandBN on the Fundus
dataset.

$p$ | 0 | 0.2 | 0.4 | 0.6 | 0.8 | 1.0
---|---|---|---|---|---|---
Dice (%) | 85.18 | 85.42 | 85.44 | 85.48 | 85.60 | 85.63

Table 9: Runtime and learnable parameter comparison.

Prostate | SS-Net | BCP | FedDG | CACPS | Ours
---|---|---|---|---|---
Runtime | 2.71h | 1.16h | 4.16h | 3.31h | 4.02h
#Par (Train) | 1.83M | 1.81M | 3.18M | 3.63M | 1.83M
#Par (Test) | 1.81M | 1.81M | 3.18M | 1.81M | 1.81M

Complementarity between the strong and style augmentation. We evaluate the
effectiveness of distinct consistency regularization strategies, as summarized
in Tab. 9. Applying both the strong and style augmentation for consistency
learning surpasses each individual stream. We also assess the performance by
replacing style augmentation with another strong augmentation, resulting in
lower outcomes. This suggests that the performance enhancement is not solely
due to duplicating the augmentation stream but, more significantly, to the
complementary relationship between strong and style augmentation.

Balance of different consistency regularization. We apply the same weighting
parameter ($\lambda_{h}=1.0$) for weak-to-style consistency regularization as
we do for weak-to-strong regularization, considering that both perturbations
are performed at the image level. We experiment to examine the effect of
feature-level consistency regularization across varying values of
$\lambda_{r}$. Table 9 shows that the averaged Dice score reaches its peak at
85.60% with $\lambda_{r}$ set to 0.2. However, further increasing
$\lambda_{r}$ may result in a slight performance decrease.

Intensity of RandBN perturbation. We experiment with the sensitivity of our
method to variations in feature-level perturbation intensity. Results in Table
9 indicate that, for all $p>0$, performance improves compared to the case when
$p=0$, which represents vanilla consistency regularization without
perturbation. This suggests that the effectiveness of RandBN perturbation is
not trivially due to the introduction of an extra stream of regularization. As
perturbation intensity increases, the model’s generalization performance
improves.

Effect of the number of medical institutions. To investigate how the number of
medical institutions affects the overall performance, we experiment on the
Prostate dataset with fewer source domains and more target domains in Tab. 10.
Despite performance drop with fewer source domains, our method outperforms
CACPS more significantly than when using all 5 domains.

Table 10: Experiments with fewer medical institutions on the Prostate dataset.

Source | A D E | A E F | C D E | A B | B F | Avg.
---|---|---|---|---|---|---
Target | B | C | F | B | C | D | A | B | F | C | D | E | F | A | C | D | E
CACPS | 86.4 | 82.3 | 85.8 | 88.0 | 76.4 | 89.1 | 80.1 | 81.3 | 87.2 | 83.5 | 90.1 | 24.7 | 82.2 | 80.2 | 49.7 | 82.4 | 46.4 | 76.22
Ours | 89.6 | 83.5 | 89.6 | 90.7 | 80.9 | 89.9 | 89.5 | 86.2 | 84.5 | 84.7 | 89.2 | 85.3 | 86.8 | 86.7 | 74.7 | 85.8 | 76.0 | 85.51

Discussion on runtime and parameter efficiency. We compare the training time
and learnable parameter size in Tab. 9. Despite longer training time due to
our multi-branch structure, we believe the cost is acceptable. The quantity of
learnable parameters in our method is also relatively low.

Due to the length limitation, please refer to supplementary materials for
additional experiments on the ACDC[3] dataset and other hyperparameters.

##  5 Conclusion

In this paper, we propose a novel framework tackling semi-supervised domain
generalization in medical image segmentation. Our statistics-individual and
statistics-aggregated branches effectively mitigate the adverse effects of
statistics differences on pseudo-labels, while simultaneously leveraging them
to enhance generalization performance. Additionally, we introduce
perturbations at both image and feature levels, expanding the consistency
regularization for a broader perturbation exploration. Evaluation on multiple
benchmarks demonstrates its superiority over other SOTA methods.

## Acknowledgements

This work is supported by the NSFC Project (62222604, 62206052, 62192783),
Jiangsu Natural Science Foundation Project (BK20210224), China Postdoctoral
Science Foundation (2024M750424), the Fundamental Research Funds for the
Central Universities (020214380120), the State Key Laboratory Funds for Key
Project (ZZKT2024A14), and Shandong Natural Science Foundation (ZR2023MF037).

## References

  * [1] Bai, W., Oktay, O., Sinclair, M., Suzuki, H., Rajchl, M., Tarroni, G., Glocker, B., King, A., Matthews, P.M., Rueckert, D.: Semi-supervised learning for network-based cardiac mr image segmentation. In: Medical Image Computing and Computer Assisted Intervention. pp. 253–260. Springer (2017)
  * [2] Bai, Y., Chen, D., Li, Q., Shen, W., Wang, Y.: Bidirectional copy-paste for semi-supervised medical image segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 11514–11524 (2023)
  * [3] Bernard, O., Lalande, A., Zotti, C., Cervenansky, F., Yang, X., Heng, P.A., Cetin, I., Lekadir, K., Camara, O., Ballester, M.A.G., et al.: Deep learning techniques for automatic mri cardiac multi-structures segmentation and diagnosis: is the problem solved? IEEE Transactions on Medical Imaging 37(11), 2514–2525 (2018)
  * [4] Cai, H., Li, S., Qi, L., Yu, Q., Shi, Y., Gao, Y.: Orthogonal annotation benefits barely-supervised medical image segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 3302–3311 (2023)
  * [5] Campello, V.M., Gkontra, P., Izquierdo, C., Martin-Isla, C., Sojoudi, A., Full, P.M., Maier-Hein, K., Zhang, Y., He, Z., Ma, J., et al.: Multi-centre, multi-vendor and multi-disease cardiac segmentation: the m&ms challenge. IEEE Transactions on Medical Imaging 40(12), 3543–3554 (2021)
  * [6] Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., Zhou, Y.: Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04306 (2021)
  * [7] Chen, X., Yuan, Y., Zeng, G., Wang, J.: Semi-supervised semantic segmentation with cross pseudo supervision. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 2613–2622 (2021)
  * [8] Cheplygina, V., de Bruijne, M., Pluim, J.P.: Not-so-supervised: a survey of semi-supervised, multi-instance, and transfer learning in medical image analysis. Medical Image Analysis 54, 280–296 (2019)
  * [9] Dou, Q., Coelho de Castro, D., Kamnitsas, K., Glocker, B.: Domain generalization via model-agnostic learning of semantic features. Advances in Neural Information Processing Systems 32 (2019)
  * [10] Guo, J., Qi, L., Shi, Y.: Domaindrop: Suppressing domain-sensitive channels for domain generalization. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 19114–19124 (2023)
  * [11] Hu, S., Liao, Z., Zhang, J., Xia, Y.: Domain and content adaptive convolution based multi-source domain generalization for medical image segmentation. IEEE Transactions on Medical Imaging 42(1), 233–244 (2022)
  * [12] Ioffe, S., Szegedy, C.: Batch normalization: Accelerating deep network training by reducing internal covariate shift. In: International Conference on Machine Learning. pp. 448–456. pmlr (2015)
  * [13] Lee, D.H., et al.: Pseudo-label: The simple and efficient semi-supervised learning method for deep neural networks. In: Workshop on challenges in representation learning, ICML. vol. 3, p. 896. Atlanta (2013)
  * [14] Li, D., Yang, Y., Song, Y.Z., Hospedales, T.: Learning to generalize: Meta-learning for domain generalization. In: Proceedings of the AAAI conference on artificial intelligence. vol. 32 (2018)
  * [15] Li, X., Yu, L., Chen, H., Fu, C.W., Xing, L., Heng, P.A.: Transformation-consistent self-ensembling model for semisupervised medical image segmentation. IEEE Transactions on Neural Networks and Learning Systems 32(2) (2020)
  * [16] Liu, Q., Chen, C., Qin, J., Dou, Q., Heng, P.A.: Feddg: Federated domain generalization on medical image segmentation via episodic learning in continuous frequency space. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 1013–1023 (2021)
  * [17] Liu, Q., Dou, Q., Heng, P.A.: Shape-aware meta-learning for generalizing prostate mri segmentation to unseen domains. In: Medical Image Computing and Computer Assisted Intervention. pp. 475–485. Springer (2020)
  * [18] Liu, X., Thermos, S., O’Neil, A., Tsaftaris, S.A.: Semi-supervised meta-learning with disentanglement for domain-generalised medical image segmentation. In: Medical Image Computing and Computer Assisted Intervention. Springer (2021)
  * [19] Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. In: International Conference on Learning Representations (2019)
  * [20] Luo, X., Chen, J., Song, T., Wang, G.: Semi-supervised medical image segmentation through dual-task consistency. In: Proceedings of the AAAI conference on artificial intelligence. vol. 35, pp. 8801–8809 (2021)
  * [21] Lyu, J., Zhang, Y., Huang, Y., Lin, L., Cheng, P., Tang, X.: Aadg: automatic augmentation for domain generalization on retinal image segmentation. IEEE Transactions on Medical Imaging 41(12), 3699–3711 (2022)
  * [22] Miao, J., Chen, C., Liu, F., Wei, H., Heng, P.A.: Caussl: Causality-inspired semi-supervised learning for medical image segmentation. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 21426–21437 (2023)
  * [23] Milletari, F., Navab, N., Ahmadi, S.A.: V-net: Fully convolutional neural networks for volumetric medical image segmentation. In: 2016 fourth international conference on 3D vision (3DV). pp. 565–571. IEEE (2016)
  * [24] Nam, H., Kim, H.E.: Batch-instance normalization for adaptively style-invariant neural networks. Advances in Neural Information Processing Systems 31 (2018)
  * [25] Pizer, S.M., Amburn, E.P., Austin, J.D., Cromartie, R., Geselowitz, A., Greer, T., ter Haar Romeny, B., Zimmerman, J.B., Zuiderveld, K.: Adaptive histogram equalization and its variations. Computer vision, graphics, and image processing 39(3), 355–368 (1987)
  * [26] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer Assisted Intervention. pp. 234–241. Springer (2015)
  * [27] Seo, S., Suh, Y., Kim, D., Kim, G., Han, J., Han, B.: Learning to optimize domain specific normalization for domain generalization. In: Proceedings of the European Conference on Computer Vision. pp. 68–83. Springer (2020)
  * [28] Shi, Y., Zhang, J., Ling, T., Lu, J., Zheng, Y., Yu, Q., Qi, L., Gao, Y.: Inconsistency-aware uncertainty estimation for semi-supervised medical image segmentation. IEEE Transactions on Medical Imaging 41(3), 608–620 (2021)
  * [29] Sohn, K., Berthelot, D., Carlini, N., Zhang, Z., Zhang, H., Raffel, C.A., Cubuk, E.D., Kurakin, A., Li, C.L.: Fixmatch: Simplifying semi-supervised learning with consistency and confidence. Advances in Neural Information Processing Systems 33, 596–608 (2020)
  * [30] Su, Z., Yao, K., Yang, X., Huang, K., Wang, Q., Sun, J.: Rethinking data augmentation for single-source domain generalization in medical image segmentation. In: Proceedings of the AAAI Conference on Artificial Intelligence. vol. 37, pp. 2366–2374 (2023)
  * [31] Tajbakhsh, N., Jeyaseelan, L., Li, Q., Chiang, J.N., Wu, Z., Ding, X.: Embracing imperfect datasets: A review of deep learning solutions for medical image segmentation. Medical Image Analysis 63, 101693 (2020)
  * [32] Tarvainen, A., Valpola, H.: Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results. Advances in Neural Information Processing Systems 30 (2017)
  * [33] Wang, J., Lan, C., Liu, C., Ouyang, Y., Qin, T., Lu, W., Chen, Y., Zeng, W., Yu, P.: Generalizing to unseen domains: A survey on domain generalization. IEEE Transactions on Knowledge and Data Engineering (2022)
  * [34] Wang, S., Yu, L., Li, K., Yang, X., Fu, C.W., Heng, P.A.: Dofe: Domain-oriented feature embedding for generalizable fundus image segmentation on unseen datasets. IEEE Transactions on Medical Imaging 39(12), 4237–4248 (2020)
  * [35] Wang, X., Zhang, J., Qi, L., Shi, Y.: Generalizable decision boundaries: Dualistic meta-learning for open set domain generalization. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 11564–11573 (2023)
  * [36] Wang, Y., Xiao, B., Bi, X., Li, W., Gao, X.: Mcf: Mutual correction framework for semi-supervised medical image segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 15651–15660 (2023)
  * [37] Wu, Y., Wu, Z., Wu, Q., Ge, Z., Cai, J.: Exploring smoothness and class-separation for semi-supervised medical image segmentation. In: Medical Image Computing and Computer Assisted Intervention. pp. 34–43. Springer (2022)
  * [38] Wu, Y., Xu, M., Ge, Z., Cai, J., Zhang, L.: Semi-supervised left atrium segmentation with mutual consistency training. In: Medical Image Computing and Computer Assisted Intervention. pp. 297–306. Springer (2021)
  * [39] Yang, L., Qi, L., Feng, L., Zhang, W., Shi, Y.: Revisiting weak-to-strong consistency in semi-supervised semantic segmentation. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 7236–7246 (2023)
  * [40] Yao, H., Hu, X., Li, X.: Enhancing pseudo label quality for semi-supervised domain-generalized medical image segmentation. In: Proceedings of the AAAI Conference on Artificial Intelligence. vol. 36, pp. 3099–3107 (2022)
  * [41] Yu, L., Wang, S., Li, X., Fu, C.W., Heng, P.A.: Uncertainty-aware self-ensembling model for semi-supervised 3d left atrium segmentation. In: Medical Image Computing and Computer Assisted Intervention. pp. 605–613. Springer (2019)
  * [42] Yuan, J., Liu, Y., Shen, C., Wang, Z., Li, H.: A simple baseline for semi-supervised semantic segmentation with strong data augmentation. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 8229–8238 (2021)
  * [43] Zhou, K., Loy, C.C., Liu, Z.: Semi-supervised domain generalization with stochastic stylematch. International Journal of Computer Vision pp. 1–11 (2023)
  * [44] Zhou, Z., Qi, L., Shi, Y.: Generalizable medical image segmentation via random amplitude mixup and domain-specific image restoration. In: Proceedings of the European Conference on Computer Vision. pp. 420–436. Springer (2022)

The Devil is in the Statistics: Mitigating and Exploiting Statistics
Difference for Generalizable Semi-supervised Medical Image Segmentation
\- Supplementary Material -

##  Appendix A Additional Implementation Details

We implement the statistics-individual and statistics-aggregated branches in a
non-intrusive manner. With just one line of code, we can replace the BN layers
in most networks with statistics-individual and aggregated branches without
the need to modify the source code defining the network, boosting efficiency
while avoiding potential pitfalls linked to modifying the original code.
Specifically, our process involves traversing and replacing BN modules. To
enable the modules to accept additional parameters indicating domain
identifiers, we leveraged PyTorch’s forward hook mechanism. In Algorithm 1, we
provide the pseudocode for the consistency regularization on unlabeled data,
incorporating the integration of the statistics-individual and aggregated
branches.

For the learnable mixing coefficient $\alpha$ in Eq. 6, we utilize the sigmoid
function to ensure its range stays within 0 to 1. Additionally, we employ a
learning rate 10 times larger than other parameters with weight decay set to 0
for updating the mixing coefficient. The mixing coefficients on the Prostate,
Fundus, and M&Ms datasets are empirically initialized as 0.999, 0.5, and 0.999
respectively. The confidence threshold $\tau$ in Eq. 3 is consistently set to
$0.95$ across all datasets.

Algorithm 1 Pseudocode for unsupervised loss

⬇

# fs: student segmentation network

# ft: teacher segmentation network

# aug_w/aug_s/aug_h: weak/strong/hist-match augmentation

# k: number of source domains

# thresh: confidence threshold

# lambda_h/lambda_r: weights of consistency

# set up statistics-individual and aggregated branches

# the last domain denotes the statistics-aggregated branch

⬇

SIAB.convert(fs, n_domains=k + 1) # single line for model conversion

SIAB.convert(ft, n_domains=k + 1)

⬇

for u, domain_u in loader_u:

# domain_u contains domain labels for each sample

u_w = aug_w(u)

u_s = aug_s(u_w)

u_h = aug_h(u_w)

# obtain predictions with SIAB

p_w_ind = ft(u_w, domain_id=domain_u)

p_w_agg = ft(u_w, domain_id=k + 1)

p_w = (p_w_ind + p_w_agg) / 2

p_s = fs(u_s, domain_id=k + 1)

p_h = fs(u_h, domain_id=k + 1)

with SIAB.stop_bn_gradient(fs):

p_r = fs(u_w, domain_id=domain_u, rand=True)

# pseudo label and thresholding

conf, pl = p_w.detach().softmax(dim=1).max(dim=1)

mask = (conf > thresh).float()

# loss

criterion = nn.CrossEntropyLoss(reduction=’none’)

loss_s = (criterion(p_s, pl) * mask).mean()

loss_h = (criterion(p_h, pl) * mask).mean()

loss_r = (criterion(p_r, pl) * mask).mean()

loss_u = loss_s + lambda_h * loss_h + lambda_r + loss_r

##  Appendix B Additional Experiments

Experiments on typical semi-supervised medical image segmentation tasks. Since
our method can also be applied to semi-supervised tasks, we experiment on the
ACDC [3] dataset and compare the performance against the state-of-the-art
semi-supervised medical image segmentation methods. While specifically
designed for semi-supervised domain generalization with data from multiple
medical institutions, our method attains competitive results on typical semi-
supervised tasks.

Table 11: Comparisons with the state-of-the-art methods on the ACDC dataset. ACDC (Dice, %) | UA-MT | SS-Net | BCP | CauSSL | Ours
---|---|---|---|---|---
3 cases (5%) | 46.04 | 65.83 | 87.59 | - | 86.72
7 cases (10%) | 81.65 | 86.78 | 88.84 | 89.66 | 88.67

Experiments with reduced training iterations. We provide the performance of
our method using different training time on the 30% labeled Prostate dataset
and compare it to the previous SOTA CACPS. As shown below, our method
consistently outperforms CACPS. Notably, even with just 2 hours of training
time, our method still exhibits a significant advantage of 1.5%.

Table 12: Results with reduced training iterations on the Prostate dataset. Prostate (30%) | CACPS | Ours | Ours | Ours
---|---|---|---|---
Time (h) | 3.31 | 2.01 | 3.22 | 4.02
Dice (%) | 86.47 |  87.98 +1.51 |  88.58 +2.11 |  88.73 +2.26

Effect of more weight factors. We implicitly assigned equal weights to the
factors in the following equations:

  * •

Unsupervised loss: $\mathcal{L}=\mathcal{L}_{x}+\lambda_{u}\mathcal{L}_{u},$

  * •

Pseudo-label ensemble: $p_{w}=t\cdot p_{w}^{\mathcal{IF}}+(1-t)\cdot
p_{w}^{\mathcal{AF}},$

  * •

Supervised loss:
$\mathcal{L}_{x}=\mathcal{L}_{\mathcal{IF}}+\lambda_{\mathcal{AF}}\mathcal{L}_{\mathcal{AF}},$

where $\mathcal{L}_{\mathcal{IF}}$ and $\mathcal{L}_{\mathcal{AF}}$ denote the
supervised loss for statistics-individual and aggregated branches,
respectively. We empirically investigate the impact of these weight factors on
the 30% labeled Fundus dataset. The results are shown below with default
weight factors marked in blue and best performances in bold. Our method
demonstrates robustness against variation with different values of weighting
factors. Setting the weights equally achieves near-optimal performance, with
marginal gains observed upon further adjustment.

Table 13: Experiments on different weight factors on the Fundus dataset. $\lambda_{u}$ | 0.3 | 0.5 | 0.7 | 1.0 | 1.5 | 2.0
---|---|---|---|---|---|---
Dice (%) | 85.27 | 85.63 | 85.74 | 85.60 | 85.39 | 85.44

$t$ | 0.25 | 0.5 | 0.75
---|---|---|---
Dice (%) | 85.31 | 85.60 | 85.34

$\lambda_{\mathcal{AF}}$ | 0.5 | 1.0 | 1.5 | 2.0
---|---|---|---|---
Dice (%) | 84.97 | 85.60 | 85.69 | 85.34

Effect of confidence threshold. We evaluate the sensitivity of our method
against varying confidence thresholds on the Fundus dataset as follows. The
performance variation with different values of $\tau$ is minimal.

Table 14: Experiments on confidence threshold on the Fundus dataset. $\tau$ | 0.85 | 0.90 | 0.95 | 0.98 | 0.99
---|---|---|---|---|---
Dice (%) | 85.30 | 85.39 | 85.60 | 85.46 | 85.34

##  Appendix C Limitations

In the previous state-of-the-art method, CACPS [40], for medical image
segmentation in SSDG, domain labels were not utilized, while our approach
introduces statistics-individual branches that necessitate the use of domain
labels. We believe that obtaining domain labels is relatively straightforward.
For instance, in collaborative training involving multiple medical
institutions, using medical institution information as domain identifiers is a
viable approach. Therefore, it is reasonable to utilize domain labels in SSDG,
following the conventional practices in domain generalization.


# Reduced Contraction Costs of Corner-Transfer Methods for PEPS

Wangwei Lan  Glen Evenbly  School of Physics, Georgia Institute of Technology,
Atlanta, GA 30332, USA<EMAIL_ADDRESS>
###### Abstract

We propose a pair of approximations that allows the leading order
computational cost of contracting an infinite projected entangled-pair state
(iPEPS) to be reduced from $\mathcal{O}(\chi^{3}D^{6})$ to
$\mathcal{O}(\chi^{3}D^{3})$ when using a corner-transfer approach. The first
approximation involves (i) reducing the environment needed for truncation of
the boundary tensors (ii) relies on the sequential contraction and truncation
of bra and ket indices, rather than doing both together as with the
established algorithm. To verify the algorithm, we perform benchmark
simulations over square lattice Heisenberg model and obtain results that are
comparable to the standard iPEPS algorithm. The improvement in computational
cost enables us to perform large bond dimension calculations, extending its
potential to solve challenging problems.

##  I Introduction

In recent years, tensor network algorithms TN1 ; TN2 ; TN3  such as Matrix
Product States (MPS) MPS1 ; MPS2 ; MPS3 , Projected Entangled Pair States
(PEPS) PEPS1 ; PEPS2 ; PEPS3  and Multiscale Entanglement Renormalization
Ansatz (MERA) 2DMERA1 ; 2DMERA2 ; MERA  has been extensively used to study
quantum many-body problems. Among these methods, infinite PEPS (iPEPS) is a
wave function ansatz designed for 2D quantum systems in the thermodynamic
limits. The success of iPEPS originated from its internal spatial structures
which captures the so called area law of entanglement entropy AreaLaw1 .
However, the massive computational cost in terms of the virtual bond dimension
$D$ and boundary bond dimensions $\chi$ of environment tensors limit its
application to challenging problems. For example, the cost scales as
$\mathcal{O}(\chi^{3}D^{6})$ in the most widely used iPEPS algorithm for
square lattice system iPEPS1 . Consequently, calculations are limited to
$D\approx 7$ for non-symmetric iPEPS and $D\approx 12$ for SU(2) symmetric
iPEPS iPEPS2 . But many problems are still controversial under these bond
dimensions. Thus, reducing the cost become extremely important.

In this paper, we first point out that simplified environment tensors can be
used during the so-called move steps (described below), which reduces the time
cost from $\mathcal{O}(\chi^{3}D^{6})$ to $\mathcal{O}(\chi^{3}D^{4})$. We
then further propose a sequential-order contraction method which only slightly
modify previous algorithms, but can further reduce the cost to
$\mathcal{O}(\chi^{3}D^{3})$. At the same time, the memory cost reduces from
$\mathcal{O}(\chi^{2}D^{4})$ to $\mathcal{O}(\chi^{2}D^{3})$.

The paper is organized as follows. In section II, we briefly introduce iPEPS
algorithm. In section III we show how to use the reduced environments and in
section IV we explain in detail our proposed method. Benchmark results are
presented in section V. Finally, discussions and conclusions are presented in
section VI.

Figure 1: (a) Wave function ansatz $\ket{\psi}$ is represented by contractions
of infinite numbers of 5-index tensors. A unit cell of 2 tensors are used in
this paper. (b) Transfer tensor $a$ and $b$ are contracted by contracting
physical dimensions of local tensor $A$ and $B$. (c) Norm
$\bra{\psi}\ket{\psi}$ is defined as contractions of all physical index
between $\ket{\psi}$ with its complex conjugate $\bra{\psi}$. It is a infinite
tensor network formed by infinite numbers of $a$ and $b$. (d) An effective
environment for $2\times 2$ cell, where the exact environment is approximated
by a few small index tensors contracting with each other. The accuracy of the
environments is controlled by boundary dimension $\chi$.  Figure 2: (a) left
move update environment $\\{C_{b}^{1},T_{b}^{1},T_{a}^{1},C_{a}^{4}\\}$ from
$\\{C_{a}^{1},T_{a}^{1},T_{b}^{1},C_{b}^{1},T_{a}^{2},a,b,T_{b}^{4}\\}$. (b)
In order to avoid the curse of dimensionality of the bond dimension,
isometries $\\{v_{ab},v_{ba}\\}$ are introduced to keep the maximum bond
dimension to $\chi$. (c) $\\{v_{ab},v_{ab}^{\dagger}\\}$ and
$\\{v_{ba},v_{ba}^{\dagger}\\}$ are obtained through SVD of a sub tensor
network from FIG 1(d).

##  II Standard iPEPS method

iPEPS is a tensor network wave function ansatz for two dimensional lattice
models. The ansatz is formed by infinite numbers of local tensors contracting
at certain pattern. For simplicity reason, we focus our discussion on square
lattice in this paper, as shown in FIG 1(a). We set the unit cell to include 2
local tensors $\\{A,B\\}$, which is compatible with the ground state of our
benchmark model. Each lattice site in the ansatz is represented by a 5-index
tensor, either $A$ or $B$. Each tensor has one physical index, corresponding
to the local Hilbert space and four virtual index that connect to the nearest-
neighbor tensors in the lattice. The bond dimension of $A$ and $B$ are
$(d,D,D,D,D)$, where $d$ is the dimension of physical index and $D$ is
dimension of virtual index that controls the accuracy of the ansatz.

Tensors $A$ and $B$ are optimized either through variational method iPEPS3 ;
iPEPS4  or imaginary time evolution methodSuzuki1 ; Suzuki2 ; Suzuki3 ;
Suzuki4 . A common task for optimizing the ansatz in both ways is to compute
the environment, which is defined as contractions of all tensors except for
specific locations at $\bra{\psi}\ket{\psi}$, see FIG 1(c). The environment is
approximated by 12 environment tensors in order to be evaluated
systematically, namely
$\\{C_{a}^{1},C_{b}^{2},C_{a}^{3},C_{b}^{4},T_{a}^{1},T_{b}^{1},T_{a}^{1},T_{a}^{2},T_{a}^{3},T_{b}^{3},T_{a}^{4},T_{b}^{4}\\}$
in FIG 1(d). Note tensors $\\{C_{a}^{1},C_{b}^{2},C_{a}^{3},C_{b}^{4}\\}$
mimics the corner part of the environment, and tensors
$\\{T_{a}^{1},T_{a}^{2},T_{a}^{3},T_{a}^{4},T_{b}^{1},T_{b}^{2},T_{b}^{3},T_{b}^{4}\\}$
account for rows or columns in FIG 1 (d). Exact computation of environments
suffer from the curse of dimensionality, i.e. the bond dimension of the exact
environment tensors $\\{C_{a}^{i},T_{a}^{i},i=1,...,4\\}$ are exponentially
large. Thus, in order for the environment to be evaluated efficiently,
approximations are needed to truncate the dimension of these tensors such that
the boundary bond dimensions of the environment tensors are limited to $\chi$.

The effective environment tensors can be computed using various approaches.
Two main stream method are TEBD and corner transfer matrix (CTM) methods. In
this paper, we followed the CTM method to construct environment tensors, where
the environments are updated iteratively, as described in Ref CTMtn, ;
CTMstat, . Each iterative step includes four moves: left, right, up, down-
move. For concreteness, we only elaborate left move here. Moves in other
directions can be adjusted accordingly. To begin with, we take left most
environment tensors in FIG 1(d), i.e.
$\\{C_{a}^{1},T_{a}^{1},T_{b}^{1},C_{b}^{4}\\}$ and then insert one column of
$\\{T_{a}^{2},a,b,T_{b}^{4}\\}$ to the right side as shown in FIG 2(a). In
principle, the tensor pairs
$\\{(C_{a}^{1},T_{a}^{2}),(T_{a}^{1},a),(T_{b}^{1},b),(C_{b}^{4},T_{b}^{4})\\}$
forms the environments $\\{C_{b}^{1},T_{b}^{1},T_{a}^{1},C_{a}^{4}\\}$.
However, after the absorption, the boundary bond dimension grows from $\chi$
to $\chi D^{2}$. Thus, we need to find a proper way to reduce the dimension
from $\chi D^{2}$ to $\chi$ without affect the overall accuracy. To do that,
we introduce two projectors $P_{ab}=v_{ab}v_{ab}^{\dagger}$ and
$P_{ba}=v_{ba}v_{ba}^{\dagger}$ at shown in FIG 2(b), where
$\\{v_{ab},v_{ab}^{\dagger}\\}$ and $\\{v_{ba},v_{ba}^{\dagger}\\}$ are
4-index isometries with bond dimension $(\chi,D,D,\chi)$. The bond dimension
reduction is then fulfilled by projecting first 3-index of $v_{ab}$ and
$v_{ba}$ into a subspace expanded by the last index of $v_{ab}$ and $v_{ba}$.
We then get the final updated environment tensors
$\\{C_{b}^{1},T_{b}^{1},T_{a}^{1},C_{a}^{1}\\}$. Continue the same process by
inserting another column of tensors $\\{T_{b}^{2},b,a,T_{a}^{4}\\}$ and then
get updated $\\{C_{a}^{1},T_{a}^{1},T_{b}^{1},C_{b}^{1}\\}$. That completes a
full left move.

Isometries $\\{v_{ba},v_{ab}\\}$ can be calculated in various ways. One
straightforward way to obtain $v_{ba}$ is shown in FIG 2(d), a singular value
decomposition (SVD) is applied to a sub tensor network extracted from FIG
1(a). Then the unitary matrix $U$ is reshaped to $v_{ba}$, with truncation
applied corresponding to the largest $\chi$ singular values. Isometry $v_{ab}$
can be obtained in a similar way.

In practice, after applying moves in four directions $N_{CTM}\sim 10$-$30$
times, the environment tensors will become sufficiently converged such that
they can effectively represent for the full environment of infinite lattice
sites. The time cost for tensor contractions and svd in FIG 3(a) and FIG 3 (d)
are $\mathcal{O}(\chi^{3}D^{6})$ and the memory cost are
$\mathcal{O}(\chi^{2}D^{4})$. Since CTM method is the most time consuming
parts in iPEPS algorithms, the entire cost scales as
$\mathcal{O}(N_{CTM}\chi^{3}D^{6})$ . In most cases, the boundary dimension
$\chi\sim D^{2}$ yield pretty good convergence for physical quantities. Thus,
the cost for total calculations can also be addressed as roughly
$\mathcal{O}(N_{CTM}D^{12})$ for time and $\mathcal{O}(D^{8})$ for memory.

##  III Reduced Environment

In the last section, we obtain isometries $\\{v_{ab},v_{ba}\\}$ by taking SVD
of a sub tensor network of 8 individual tensors extracted from FIG 1(d).
However, in a fully converged iPEPS ansatz, the right half of sub tensor
networks $\\{b,T_{b}^{2},C_{b}^{2},T_{b}^{3}\\}$ can be approximated by
environment tensors $\\{C_{a}^{2},T_{a}^{3},v\\}$, as shown in FIG 3(a). The
isometry $v$ is from right move, and won’t be used in this step. We now can
implement SVD similarly as in the standard iPEPS case to obtain $v_{ab}$. The
only difference is that less environment tensors are used during the process.
Clearly, smaller numbers of environments helps to reduce the time cost to
$\mathcal{O}(\chi^{3}D^{4})$. We refer this case as reduced-env. In FIG 3(c),
we make a further approximation to the sub tensor network. Such approximation
can be fulfilled at the same leading cost. Although this step is not necessary
here, it will become important in our proposed method in the next section.

##  IV Sequential Order Contractions

As mentioned above, the bond dimensions of $v_{ab}$ and $v_{ba}$ are
$(\chi,D,D,\chi)$, therefore, any contractions with $v_{ab}$ (or $v_{ba}$)
will inevitably result into a $\mathcal{O}(\chi^{3}D^{4})$ time cost, which is
the leading cost in the algorithm. Following the common idea in tensor network
field of using multiple small index tensors to approximate large index tensors
without affect the accuracy significantly, it is possible to use small index
tensors to approximate $v_{ab}$ and $v_{ba}$ in this scenerio as well. In FIG
4(c), a decomposition of the 4-index isometry $v_{ba}$ into two 3-index
isometries $v_{ba}^{1}$ and $v_{ba}^{2}$ are introduced. By construction,
$v_{ba}^{1}$ contracts with boundary index of $T_{a}^{1}$ and virtual index of
$A$ while $v_{ba}^{2}$ contracts with boundary index and virtual index of
$A^{\dagger}$. Consequentially, if the maximum dimension of
$\\{v_{ba}^{1},v_{ba}^{2}\\}$ are bounded to $\chi$, the computational cost
for updating environment tensors is then reduced to
$\mathcal{O}(\chi^{3}D^{3})$.

In reality, isometries $\\{v_{ba}^{1},v_{ba}^{2}\\}$ are computed without
considering $v_{ba}$. Similar to the standard iPEPS case, $v_{ba}^{1}$ and
$v_{ba}^{2}$ are obtained through SVD. As shown in FIG 4(a), instead of
reshaping the sub tensor network to matrix of dimension $\chi D^{2}\times\chi$
in FIG 3(b), we reshape it to matrix of dimension $\chi D\times\chi D$. Then,
we perform a SVD on the reshaped matrix. The resultant unitary matrix
$U_{d}^{1}$ is reshaped to isometry $v_{ba}^{1}$ by truncating according to
the largest $\chi^{\prime}$ singular values. Note the bond dimension of
$v_{ba}^{1}$ is $(\chi,D,\chi^{\prime})$, where $\chi^{\prime}$ controls the
accuracy of truncation in this step. After obtaining $v_{ba}^{1}$, we contract
$v_{ba}^{1}$ with environments as in FIG 4(b), and perform an additional SVD
on $(\chi D,\chi)$ matrix to obtain unitary matrix $U_{d}^{2}$, which serves
as $v_{ba}^{2}$ after reshaping. $\\{v_{ab}^{1},v_{ab}^{2}\\}$ are constructed
with the same method. In practice, $\chi^{\prime}=\chi$ is sufficient to
obtain good results. If the system is strongly correlated, we might choose
slightly larger $\chi^{\prime}$. It is easy to verify that the leading cost is
$\mathcal{O}(\chi^{3}D^{3})$ in steps mentioned above. We refer this case as
sequential-order method.

We have used an approximated environments
$\\{{C^{\prime}}_{b}^{1},{T^{\prime}}_{b}^{2},{C^{\prime}}_{b}^{2}\\}$ to
obtain $\\{v_{ab}^{1},v_{ab}^{2}\\}$ instead of environments from FIG 3(b). It
is due to the fact that the cost of contracting environments in FIG 3(b) is
already $\mathcal{O}(\chi^{3}D^{4})$, which break the computational gain from
decomposing the isometries $v_{ba}$ into $v_{ba}^{1}$ and $v_{ba}^{2}$.
Luckily, this cost $\mathcal{O}(\chi^{3}D^{4})$ can be avoided since
$\\{{C^{\prime}}_{b}^{1},{T^{\prime}}_{b}^{2},{C^{\prime}}_{b}^{2}\\}$ can be
obtained with high accuracy at the cost of $\mathcal{O}(\chi^{3}D^{3})$
(Appendix). Unlike other environment tensors that will be iteratively updated
and reused,
$\\{{C^{\prime}}_{b}^{1},{T^{\prime}}_{b}^{2},{C^{\prime}}_{b}^{2}\\}$ are
only constructed for the calculation of $v_{ab}^{1}$ and $v_{ab}^{2}$.

Figure 3: (a) Right corner of the network
$\\{C_{b}^{2},b,T_{b}^{2},T_{b}^{3}\\}$ is approximated by environment tensors
$\\{C_{b}^{2},T_{b}^{3},v\\}$. (b) Isometry $v_{ab}$ is calculated by
performing SVD over a sub tensor network. (c) A further approximation on the
environment tensors is applied, the environments can be approximated by
$\\{{C^{\prime}}_{b}^{1},{T^{\prime}}_{b}^{2},{C^{\prime}}_{b}^{2}\\}$.
Figure 4: (a),(b) $\\{v_{ba}^{1},v_{ba}^{2}\\}$
($\\{v_{ab}^{1},v_{ab}^{2}\\}$) are calculated through SVD. (b) Environment
tensors are updated similarly as the standard iPEPS cases. The computational
cost for both (a) and (b) are $\mathcal{O}(\chi^{3}D^{3})$.

##  V Results

In this paper, we run benchmarks on square-lattice quantum Hamiltonian models
to verify the algorithm. We used imaginary-time evolution method with a
second-order Suzuki-Trotter decomposition to optimize iPEPS tensors for ground
state. To reduce total computational time, we applied the fast full update
scheme (FFU) as described in Ref iPEPS1  to avoid recomputing the effective
environments after each update step. In the simulation, we have set the ground
state wave function to be an infinite tensor network composed of a two-site
unit cell $A$ and $B$. Clearly, the system has translational symmetry. In
order for the algorithm to converge faster, we start optimization with a large
imaginary time $\delta\approx 0.1$ in the first few updating iterations. After
the wave function converges, we decrease $\delta$ until it reaches $\delta\sim
10^{-4}$. In order for the simulation to be more stable, we also implemented
the gauge-fixing step as described in Ref iPEPS1 .

Our first benchmark model is nearest-neighbor Heisenberg model on square
lattice.

$\displaystyle\mathcal{\hat{H}}=\sum_{<i,j>}\hat{h}_{i,j}=-\sum_{i,j}(\hat{S}_{i}^{x}\hat{S}_{j}^{x}+\hat{S}_{i}^{y}\hat{S}_{j}^{y}+\hat{S}_{i}^{z}\hat{S}_{j}^{z})$
(1)

where $\langle i,j\rangle$ means nearest-neighbor sites and
$\hat{S}_{i}^{x},\hat{S}_{i}^{y},\hat{S}_{i}^{z}$ are Pauli matrices. In our
simulation, we specifically set $\chi=10(\mathcal{D}+1)$ for $\mathcal{D}\leq
10$ and $\chi=\mathcal{D}^{2}$ for $\mathcal{D}\geq 11$ such that $\chi$
dependence on the final results becomes minimum. Note, in the sequential-order
case, we have introduced two extra boundary dimensions $\chi^{\prime}$ and
$\chi^{\prime\prime}$, which in principle could be set independently. However,
for the sake of concreteness, we let them be the same
$\chi=\chi^{\prime}=\chi^{\prime}$. As we shall see later, such a choice is
good enough to reach the desired accuracy. Another important thing worth
mentioning is that in order for the comparison to be reasonable, all
calculation details are exactly the same for all cases, except for the ways to
calculate the isometries.

In FIG. 5, we check the energy and convergence of staggered magnetization with
respect to $\chi$ when $\mathcal{D}=5$. In all cases, the dependence of energy
on the boundary dimension $\chi$ is systematically improved with increasing
$\chi$. Although there is a small discrepancy when $\chi<20$, our methods
produce exactly the same results in all simulations for large enough $\chi$.
Since energy will converge only in the large $\chi$ limit, it is acceptable
that there is a small disturbance when $\chi$ is small. FIG. 5 also suggests
that when $\chi\geq 30$, the energy and magnetization only improve slightly,
providing a numerical proof that $\chi\sim\mathcal{D}^{2}$ is a sufficient
condition for the convergence of physical quantities.

Figure 5: (Color online) $\chi$ dependence of energy per site and staggered
magnetization for different iPEPS methods ($\mathcal{D}=5$). Accuracy is
systematically improved as $\chi$ increased for both physical quantities.
Meanwhile, all methods reach the same accuracy as dimension $\chi$ becomes
relatively large.

In FIG. 6, we show the relative error $\epsilon$ for energy and staggered
magnetization as a function of $\mathcal{D}$. To quantify the error, we used
$E_{0}=-0.669437$ from accurate quantum Monte Carlo calculations QMC  as a
reference. It is clear that the results of energy accuracy and stagger
magnetism results are systematically improved with increasing $\mathcal{D}$,
which is consistent with the basic intuitions of the iPEPS ansatz.
Furthermore, we also observed that the cases of reduced-env and sequential-
order cases yield comparable results with standard iPEPS at various virtual
dimensions, which validate the usage of reduced environments and dual
isometries $\\{v_{ba}^{1},v_{ba}^{2}\\}$ ($\\{v_{ab}^{1},v_{ab}^{2}\\}$).
Contrary to naive expectations, our results show that there is no significant
difference between standard iPEPS and sequential-order iPEPS even though
multiple approximations are introduced in order to reduce the cost.

Figure 6: (Color online) (a) Relative energy error for different iPEPS
methods. Both sequential-order and reduced-env achieve the same accuracy as
the standard iPEPS case. (b) Staggered magnetism follows trend as relative
error.

Finally, in FIG. 7, we compare the running time of five imaginary time steps
for different proposals at each bond dimension. As we can see, due to extra
approximations mentioned above to avoid $\mathcal{O}(\chi^{3}\mathcal{D}^{4})$
computational complexity, sequential-order method is slower than reduced-env
method for $\mathcal{D}\leq 7$. But, for $\mathcal{D}\geq 8$, sequential-order
method become much faster and the speed-up factor becomes more prominent as
the $\mathcal{D}$ increased. Due to the speed-up, we can run more accurate
simulations on $\mathcal{D}=14,\chi=196$ in a reasonable time.

Figure 7: (Color online) Actual running time in seconds of 5 imaginary-time
steps for different iPEPS methods. sequential-order iPEPS is slower than
reduced-env iPEPS when $\mathcal{D}\leq 7$ due to the extra approximations
introduced in section III. When $\mathcal{D}\geq 8$, sequential-order iPEPS
becomes significantly faster than other implementations.

Our second benchmark model is the ferromagnetic quantum Ising model with
transverse magnetic field,

$\displaystyle\mathcal{\hat{H}}=-\sum_{<i,j>}\hat{S}^{z}_{i}\hat{S}^{z}_{j}-h\sum_{i}\hat{S}^{x}_{i}$
(2)

In this simulation, we choose the bond dimension $\mathcal{D}=3$ and the
boundary dimension $\chi=30$ such that all algorithms converge in physical
quantities. Like in the previous benchmark, we set all computational details
to be exactly the same. We first computed the order parameter
$m_{z}=\langle\hat{S}^{z}\rangle$ at different magnetic fields $h$ for the
quantum Ising model. As shown in FIG. 8 (a), not only do all algorithms
produce exactly the same order parameter, but they also correctly indicate the
phase transition point at $h\approx 3.04$. In FIG. 8 (b), we calculate the
correlation function
$S_{zz}(x)=\langle\hat{S}^{z}_{l}\hat{S}^{z}_{l+x}\rangle-\langle\hat{S}^{z}_{l}\rangle\langle\hat{S}^{z}_{l+x}\rangle$
at the critical point $h=3.04$. Once again, all three algorithms produce
exactly the same results.

Figure 8: (Color online) (a) Order parameter $m_{z}=\langle S^{z}\rangle$ are
calculated at different magnetic field $h$ for quantum Ising model under
transverse magnetic field. (b) Correlation function $S_{zz}(x)$ at different
positions are also calculated. In both (a) and (b), three different approaches
produce the same results provided all the calculation details are the same for
different simulations.

##  VI Discussion

In this paper, we point out that a reduced environment can be used to compute
isometries in the iPEPS algorithm. Meanwhile, we also propose to use two
isometries instead of one in the CTM steps. The computational cost is reduced
from $\mathcal{O}(\chi^{3}D^{6})$ to $\mathcal{O}(\chi^{3}D^{3})$. Thus, we
are able to run accurate simulations with much larger $D$ and $\chi$ (maximum
$D=14$).

We realize that single layer iPEPS with the same leading cost has already been
proposed in the literature iPEPSsingle . However, our algorithm still benefits
in several ways. First, the unit cell size in our proposal is the same as
standard double layer iPEPS whereas Ref iPEPSsingle,  need to expand the unit
cell to 4 times larger. Moreover, given the same $D$ and $\chi$, our method is
able to reach roughly the same accuracy as standard double layer iPEPS, while
the single layer iPEPS from Ref iPEPSsingle,  generally requires larger $\chi$
to obtain similar accuracy.

We note that our method can be adjusted easily to the existing iPEPS method.
Besides, although we obtain ground state through imaginary time evolution
methods, our proposal should be able to be used in other updating schemes as
long as CTM is applied to renormalize the environment tensors. As is well
known to tensor network community, large computational speed-up can be
achieved by including the abelian and non-abelian symmetries. It is also
interesting to implement our method in symmetric case in the future.

##  Appendix A Approximations

As is mentioned in the context, the cost of contracting environment tensors in
FIG 9(a) is already $\mathcal{O}(\chi^{3}D^{4})$. To avoid this cost, several
additional isometries $\\{v_{l}^{1},v_{l}^{2},v_{r}^{1},v_{r}^{2}\\}$ are
inserted into specific locations, as shown in FIG 9 (b). Although, thoese
isometries has the same structure as $\\{v_{ab}^{1},v_{ab}^{2}\\}$, they can
be obtained in different ways. In FIG 9(c), we compute
$\\{v_{l}^{1},v_{l}^{2}\\}$ by only considering the environment tensors
$\\{C_{a}^{1},T_{a}^{1}\\}$ from left part. In reality,
$\\{v_{l}^{1},v_{l}^{2}\\}$ works pretty good when they have the same bond
dimension as $\\{v_{ab}^{1},v_{ab}^{2}\\}$.

Figure 9: (a) A further approximation on the environment tensors is applied.
(b) two pairs of isometries $\\{v_{l}^{1},v_{l}^{2}\\}$ and
$\\{v_{r}^{1},v_{r}^{2}\\}$ are inserted into designated positions. (c)
$\\{v_{l}^{1},v_{l}^{2}\\}$ are calculated such that the approximation is
valid to high accuracy.

## References

  * (1) J. I. Cirac and F. Verstraete, Renormalization and tensor product states in spin chains and lattices, J. Phys. A: Math. Theor. 42, 504004 (2009).
  * (2) R. Orus, A practical introduction to tensor networks: Matrix product states and projected entangled pair states, Ann. Phys. 349, 117 (2014).
  * (3) J. C. Bridgeman and C. T. Chubb, Hand-waving and Interpretive Dance: An Introductory Course on Tensor Networks, J. Phys. A: Math. Theor. 50, 223001 (2017).
  * (4) M. Fannes, B. Nachtergaele, and R. F. Werner, Finitely correlated states on quantum spin chains, Commun. Math. Phys. 144, 443 (1992).
  * (5) S. Ostlund and S. Rommer, Thermodynamic limit of density matrix renormalization, Phys. Rev. Lett. 75, 3537 (1995).
  * (6) U. Schollwöck, The density-matrix renormalization group in the age of matrix product states, Ann. Phys. 326, 96 (2011).
  * (7) G. Vidal, A class of quantum many-body states that can be efficiently simulated, Phys. Rev. Lett. 101, 110501 (2008).
  * (8) L. Cincio, J. Dziarmaga, and M. M. Rams, Multiscale entanglement renormalization ansatz in two dimensions: quantum Ising model, Phys. Rev. Lett. 100, 240603 (2008).
  * (9) G. Evenbly and G. Vidal, Entanglement renormalization in two spatial dimensions, Phys. Rev. Lett. 102, 180406 (2009).
  * (10) F. Verstraete and J. I. Cirac, Renormalization algorithms for Quantum-Many Body Systems in two and higher dimensions, arXiv:cond-mat/0407066 (2004).
  * (11) F. Verstraete, J.I. Cirac, and V. Murg, Matrix Product States, Projected Entangled Pair States, and variational renormalization group methods for quantum spin systems, Adv. Phys. 57, 143 (2008).
  * (12) J. Jordan, R. Orus, G. Vidal, F. Verstraete, and J. I. Cirac, Classical simulation of infinite-size quantum lattice systems in two spatial dimensions, Phys. Rev. Lett. 101, 250602 (2008).
  * (13) H. N. Phien, J. A. Bengua, H. D. Tuan, P. Corboz, and R. Orús, Infinite projected entangled pair states algorithm improved: Fast full update and gauge fixing, Phys. Rev. B 92,035142 (2015).
  * (14) C. Hubig, Abelian and non-abelian symmetries in infinite projected entangled pair states, arxiv: 1808.10804v1 (2018)
  * (15) L. Vanderstraeten, J. Haegeman, P. Corboz, and F. Verstraete, Gradient methods for variational optimization of projected entangled-pair states, Phys. Rev. B 94,155123 (2016).
  * (16) P. Corboz, Variational optimization with infinite projected entangled-pair states, Phys. Rev. B 94,035133 (2016).
  * (17) R. Haghshenas, S. S. Gong, and D. N. Sheng, Single-layer tensor network study of the Heisenberg model with chiral interactions on a kagome lattice, Phys. Rev. B 99,174423 (2019).
  * (18) R. Orús, Exploring corner transfer matrices and corner tensors for the classical simulation of quantum lattice systems, Phys. Rev. B 85, 205117 (2012).
  * (19) T. Nishino, and K. Okunishi, Corner Transfer Matrix Renormalization Group Method J. Phys. Soc. Jpn. 65, pp. 891-894 (1996).
  * (20) M. Suzuki, Fractal decomposition of exponential operators with applications to many-body theories and Monte Carlo simulations, Phys. Lett. A, 146, 6, (1990) 319-323.
  * (21) M. Suzuki, General theory of fractal path integrals with applications to many‐body theories and statistical physics, J. Math. Phys. 32, 2, (1991) 400-407.
  * (22) H. F. Trotter, On the product of semi-groups of operators, Proc. Amer. Math. Soc. 10 (1959), 545-551
  * (23) M. Suzuki, Pair-Product Model of Heisenberg Ferromagnets, J. Phys. Soc. Jpn. 21, pp. 2274-2290 (1966)
  * (24) M. Suzuki, Relationship between d-Dimensional Quanta! Spin Systems and (d+ I) -Dimensions, Prog. Theor. Phys, 1454–1469, 56, (1976)
  * (25) Y. Shi, L. Duan and G. Vidal, Classical simulation of quantum many-body systems with a tree tensor network, Phys. Rev. A 74, 022320 (2006).
  * (26) L. Tagliacozzo, G. Evenbly and G. Vidal, Simulation of two-dimensional quantum systems using a tree tensor network that exploits the entropic area law, Phys. Rev. B 80, 235127 (2009).
  * (27) G. Vidal, J. I. Latorre, E. Rico, and A. Kitaev, Entanglement in quantum critical phenomena, Phys. Rev. Lett. 90, 227902 (2003).
  * (28) J. I. Latorre, E. Rico, and G. Vidal, Ground state entanglement in quantum spin chains, Quant. Inf. Comput. 4, 48-92 (2004).
  * (29) J. Eisert, M. Cramer, and M. B. Plenio, Area laws for the entanglement entropy, Rev. Mod. Phys. 82,277 (2010).
  * (30) A. W. Sandvik, Finite-size scaling of the ground-state parameters of the two-dimensional Heisenberg model, Phys. Rev. B 56, 11678 (1997).