# Segment Together: A Versatile Paradigm for Semi-Supervised Medical Image
Segmentation

Qingjie Zeng${}^{1^{*}}$ Yutong Xie${}^{2^{*}}$ Zilin Lu1 Mengkang Lu1 Yicheng
Wu3 Yong Xia${}^{1^{\dagger}}$
1 School of Computer Science and Engineering, Northwestern Polytechnical
University, China
2 Australian Institute for Machine Learning, University of Adelaide, Australia
3 Department of Data Science & AI, Faculty of Information Technology, Monash
University, Australia
<EMAIL_ADDRESS>
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Annotation scarcity has become a major obstacle for training powerful deep-
learning models for medical image segmentation, restricting their deployment
in clinical scenarios. To address it, semi-supervised learning by exploiting
abundant unlabeled data is highly desirable to boost the model training.
However, most existing works still focus on limited medical tasks and
underestimate the potential of learning across diverse tasks and multiple
datasets. Therefore, in this paper, we introduce a Versatile Semi-supervised
framework (VerSemi) to point out a new perspective that integrates various
tasks into a unified model with a broad label space, to exploit more unlabeled
data for semi-supervised medical image segmentation. Specifically, we
introduce a dynamic task-prompted design to segment various targets from
different datasets. Next, this unified model is used to identify the
foreground regions from all labeled data, to capture cross-dataset semantics.
Particularly, we create a synthetic task with a cutmix strategy to augment
foreground targets within the expanded label space. To effectively utilize
unlabeled data, we introduce a consistency constraint. This involves aligning
aggregated predictions from various tasks with those from the synthetic task,
further guiding the model in accurately segmenting foreground regions during
training. We evaluated our VerSemi model on four public benchmarking datasets.
Extensive experiments demonstrated that VerSemi can consistently outperform
the second-best method by a large margin (e.g., an average 2.69% Dice gain on
four datasets), setting new SOTA performance for semi-supervised medical image
segmentation. The code will be released.

##  1 Introduction

Medical image segmentation is a long-standing challenge [25, 7, 9, 41]. Due to
the scarcity of voxel-level annotations, semi-supervised learning (SSL), which
can learn from limited labeled and abundant unlabeled data, has been applied
to medical image segmentation tasks.

Figure 1: A brief illustration of the difference among previous SSL (a),
universal model (b) and our proposed SSL (c). First, the previous SSL learns
each model in isolation and neglects the importance of data integration.
Second, universal models, $e.g.,$ DoDnet [43], leverage diverse task prompts
to acquire knowledge from multiple tasks in a supervised manner, which lacks
the ability to handle unlabeled data when task info is unknown. By comparison,
our proposed SSL can not only complete various missions simultaneously but
also learn from unlabeled data without requiring associated task info.

Generally, there are two popular SSL paradigms, i.e., pseudo-labeling [42] and
consistency regularization [27, 32]. The former aims to find trustworthy
pseudo-labels for re-training, $e.g.$, setting an adaptive threshold during
the learning process to filter out unreliable predictions [44]. The latter
focuses on making consistent predictions for different augmentations of the
same input [26, 20]. Despite their prevalence, most SSL methods are restricted
to a specific task, where labeled and unlabeled data share an identical label
space. Due to the insufficient supervision of the single task, the learned
distributions of labeled and unlabeled data are prone to be inconsistent,
resulting in poor generalizability and limited performance of SSL approaches.
Besides, the task-specific scenario often means that a significant portion of
the unlabeled data, which might not fit perfectly into the predefined task
label space, remains underutilized.

Recently, universal models have drawn increasing research attention. They are
trained on multi-domain and/or multi-modality data for multi-task in two ways.
The first approach involves pre-training the model using task-agnostic
unlabeled data through self-supervised learning, followed by fine-tuning on
task-specific data for individual downstream tasks [30, 33, 45]. The second
method trains a model jointly using multiple task-specific data in a
supervised fashion [5, 38, 43]. Universal models have unveiled superior
performance over traditional task-specific models on a variety of tasks,
spanning both the computer vision [12, 35] and medical imaging communities
[15, 38, 33]. The success of universal models emphasizes the importance of
integrating data and tasks to improve representation learning. This
observation inspires us to amalgamate multiple datasets and tasks within the
framework of SSL. Such integration promises not only to harness an expanded
corpus of unannotated and annotated data, thereby bolstering the supervised
component of SSL, but also to substantially enhance the model’s generalization
capabilities, thereby extending its applicability across diverse domains.

In this paper, we introduce a Versatile Semi-supervised framework (VerSemi), a
novel approach that revolutionizes common SSL paradigms. Firstly, VerSemi
surpasses task-specific learning constraints by integrating multiple targets
into a unified framework. It seamlessly establishes an enhanced label space by
amalgamating pertinent task labels, and accomplishes multiple tasks
simultaneously with the assistance of a task-prompted dynamic head. Secondly,
considering that task specifics are required for prompt-driven models to
generate prompts (one-hot encoding, language description, $etc$) during the
learning process, an issue is raised that unlabeled data may not be mined if
associated task information is remained unknown (see Fig. 1). To tackle this
problem, VerSemi first constructs a synthetic task by leveraging cutmix on
labeled data. In this way, the data in the synthetic task span a diverse range
of foreground targets within the expanded label space. By joint training with
the synthetic task, VerSemi can recognize and segment all potential foreground
regions. Grounded on this ability, VerSemi simplifies learning from unlabeled
data by eliminating the need for task-specific details. This is achieved by
ensuring consistency between combined predictions from relevant tasks and
synthetic ones. Thirdly, we empirically observe that prompts oftentimes do not
work well,  _e.g_., models fail to recognize the object indicted by a specific
prompt (see Fig.2). To address this issue, an auxiliary constraint is designed
to regularize VerSemi to enhance its controllability when meeting task-
specific prompts. In addition, it is worth noting that current SSL methods can
not directly realize task-agnostic unlabeled data learning, as they either
demand a teacher model or extra sub-networks for supervision. Our
contributions are three-fold.

  * •

Different from current SSL methods that learn tasks individually, the proposed
VerSemi performs well with the new setting of integrating various pertinent
SSL tasks into a unified framework.

  * •

We achieve task-agnostic unlabeled data learning by devising a ”synthetic
task”. This design facilitates the learning of unified foreground
segmentation. With this segmentation ability as a constraint, unlabeled data
can be excavated without acquiring associated task specifics.

  * •

Extensive experiments on four public datasets validate the superiority of
VerSemi, which presents remarkable improvements compared to task-isolated SSL
models (_e.g_., BCP, CauSSL) and associated task-unified models (_e.g_., Uni-
BCP, Uni-CauSSL).

##  2 Related Work

Figure 2: The prompt-weakening phenomenon. In Case 1, DoDNet fails to predict
the region of pancreas, and can only recognize spleen voxels no matter under
the prompt of pancreas or spleen. In Case 2, DoDNet mistakenly highlights the
region of spleen when prompted by left atrium. By comparison, VerSemi has
addressed this issue by devising an auxiliary constraint. (see Section 3.2)
Figure 3: Illustration of VerSemi. VerSemi has a task-prompted dynamic head
which can flexibly process different tasks at the same time, along that an
auxiliary constraint $\mathcal{L}_{axu}$ is designed to augment the
reliability of associated task prompt. During labeled data learning, we
construct an synthetic task (Task#5), which aims to segment all the foreground
regions. As for unlabeled learning, the aggregated prediction prompted by
Task#1 $\sim$ Task#4 is forced to be consistent with the prediction prompted
by Task#5, when feeding mixed unlabeled data into the model. Therefore, the
proposed VerSemi does not require task information to learn from unlabeled
data and is more versatile.

###  2.1 Semi-supervised Learning

Semi-supervised learning (SSL) [6, 17, 7, 31] is emerged to mitigate the issue
of tedious data annotation, by learning from unlabeled data with scarce
labeled data. Many efforts are made to explore how to excavate information
from unlabeled data adequately. For instance, UPS [23] reduced unreliable
pseudo-labels by calibrating models with uncertainty. PEFAT [42] investigated
the probability distribution of pseudo-labeled data, and further proposed a
selection standard from the perspective of loss distribution. SoftMatch [4]
and FreeMatch [29] tried to address the quantity-quality trade-off issue with
adaptive threshold. However, these SSL frameworks focus on learning labeled
and unlabeled data within each single task, followed by an issue of whether
they have the ability to scale up to heterogeneous tasks. Beyond that, another
problem is raised up when learning tasks individually, that is the improvement
is constantly marginal due to insufficient representation acquired from
limited labels in each dataset. To address the mentioned issues, we advocate
learning a unified SSL model with an integrated dataset, under the new setting
of learning various pertinent SSL tasks concurrently.

###  2.2 Representation Learning with Integrated Data

To improve the model performance and representation ability, some works
propose to learn a unified model that can complete multiple tasks
simultaneously, rather than training task-specific model separately [10, 13].
For example, DoDNet [43] collected an abdominal dataset from seven partially
labeled datasets for model training, and presented better-averaged results
than training on every single dataset. CLIP-Driven Universal Model [15]
further advanced this idea by introducing CLIP embedding [22] to help the
model capture anatomical relations between different tumors and organs. UniSeg
[38] leveraged different modalities including CT, MRI and PET, whose
performance surpassed those models trained with a single modality. These
studies underscore the importance of robust data engines, emphasizing the need
to leverage as much accessible data as possible. While the majority of these
investigations focus on either fully supervised learning [38, 15] or self-
supervised learning [33], rare of them utilize simultaneously labeled and
unlabeled data collected from different tasks. By contrast, we propose VerSemi
and make an exploratory attempt.

Additionally, it is necessary to mention that compared to DoDNet [43], a
method leverages one-hot task prompt to learn from different tasks, our
proposed VerSemi differs in: (1) DoDNet is designed under fully-supervised
setting, which can not handle unlabeled data if the associated task
information are not given; and (2) there exists severe prompt-weakening
phenomenon during task learning procedure (see Fig. 2), this situation is
overlooked by prompt-driven model like DoDNet. And VerSemi tackles this issue
by designing an auxiliary constraint.

##  3 Method

Fig. 3 shows the pipeline of our proposed VerSemi model, integrating various
semi-supervised segmentation tasks from different datasets into a unified
framework. Here, Section 3.1 shows the dynamic kernel generation in our
VerSemi. Then, Sections 3.2 and 3.3 further delve into the details of
exploiting limited task-aware labeled data and task-agnostic unlabeled data
for training, respectively.

###  3.1 Dynamic Convolution with Task Prompt

Despite the great progress of deep-learning based works, it remains sub-
optimal for an individual model with fixed convolutional kernels to handle
different segmentation tasks at the same time [43, 14]. To improve the
performance, a common practice is to use a multi-head architecture, but it
suffers from severe computational overhead with the increase of on-coming
tasks, thus is not suitable when confronted with multiple tasks. To reduce the
computational burdens of different heads, we utilize dynamic filter generation
to construct the segmentation head, which can adaptively process different
tasks with task-specific prompts without extra costs. The filter generation is
defined as:

$\displaystyle w_{k}$
$\displaystyle=\psi(GAP(Embedding),[Prompt_{\\#k}];\theta_{\psi})$ (1)
$\displaystyle\mathcal{P}_{k}$
$\displaystyle=SoftMax({f}_{D}(Embedding)*w_{k}),$

where $\psi$ is one convolutional layer with parameter $\theta_{\psi}$, which
is employed to dynamically generate parameters $w_{k}$ for the current Task#k
with $[Prompt_{\\#k}]$. Here we set the prompt in a one-hot encoding format,
which is then concatenated with global averaged feature embedding before
feeding into $\psi_{\theta}$. $\mathcal{P}_{k}$ is the prediction for Task#k,
$f_{D}$ is the decoder and symbol $\ast$ represents convolution. With
$[Prompt_{\\#k}]$, VerSemi can accurately perceive the ongoing task and
flexibly adapt kernels to fit it.

Figure 4: Predictions made by VerSemi when facing synthetic data. Case 1 and
Case 3: cutmix between left atrium and lung tumor. Case 2: cutmix between
pancreas and spleen. Case 4: cutmix between spleen and left atrium.We can find
VerSemi can produce accurate masks if prompted by Task#5.

###  3.2 Task-aware Labeled Data Learning

In this work, four pertinent tasks are incorporated. Task#1 $\sim$ Task#4 are
the pancreas, left atrium, spleen and lung tumor segmentation tasks. Below, we
describe the construction and utilization of the additionally synthesized
Task#5.

Generation of Task#5. As the bottom-right of Fig. 3 shows, we construct a
synthetic task (Task#5) based on labeled data from pertinent tasks (Task#1
$\sim$ Task#4). Task#5 is built to help VerSemi achieve task-agnostic learning
from unlabeled data and guide VerSemi to segment all foreground regions when
facing mixed data (see Fig. 4). The data generation of Task#5 is formulated
as:

$\displaystyle\mathcal{X}_{syn(i,j)}^{l}$
$\displaystyle=\mathcal{X}_{i}^{l}\odot\mathcal{M}+\mathcal{X}_{j}^{l}\odot(1-\mathcal{M})$
(2) $\displaystyle\mathcal{Y}_{syn(i,j)}^{l}$
$\displaystyle=\mathcal{Y}_{i}^{l}\odot\mathcal{M}+\mathcal{Y}_{j}^{l}\odot(1-\mathcal{M}),$

where $\mathcal{X}_{syn(i,j)}^{l}$ and $\mathcal{Y}_{syn(i,j)}^{l}$ are
synthetic images and labels for Task#5. $\mathcal{M}$ is a mask with 30%
$\sim$ 70% random masked regions. Symbol $\odot$ is element-wise
multiplication. Note that $\mathcal{Y}_{i}^{l}$ and $\mathcal{Y}_{j}^{l}$ are
binary masks, $\mathcal{X}_{i}^{l}$ and $\mathcal{X}_{j}^{l}$ are images from
the $i$-th and $j$-th task, thus $\mathcal{X}_{syn(i,j)}^{l}$ can be regarded
as mixed data that contain various targets and background. For labeled data
learning (containing Task#5), Dice loss and cross-entropy loss are leveraged,
defined as:

$\begin{split}\mathcal{L}_{lab}=Dice(\mathcal{F}(\mathcal{X}_{k}^{l},[Prompt_{\\#k}];\Theta),\mathcal{Y}_{k}^{l})+\\\
CE(\mathcal{F}(\mathcal{X}_{k}^{l},[Prompt_{\\#k}];\Theta),\mathcal{Y}_{k}^{l}),\end{split}$
(3)

where $\mathcal{L}_{lab}$ is the supervised loss on labeled data. For
simplicity, we use $\mathcal{F}(\cdot;\Theta)$ to define the whole network
with parameter $\Theta$, which contains operations in Eq.1. Benefited by
Task#5, VerSemi could have a semantic perception of all other segmentation
tasks.

Enhancing the controllability of task prompt by $\mathcal{L}_{aux}$. As
indicated by Fig. 2, there exists a weakening phenomenon of task prompt when
using task-prompted dynamic head, in which we can find models (_e.g_., DoDNet
[43]) sometimes fail to recognize prescriptive task even under the control of
task-specific prompt. It might be caused by the shared semantic information
between different segmentation tasks. Therefore, to enhance the uniqueness of
the task prompt, we add an auxiliary constraint $\mathcal{L}_{aux}$ formulated
as:

$\begin{split}\mathcal{L}_{aux}=Dice(\mathcal{F}(\mathcal{X}_{syn(i,j)}^{l},[Prompt_{\\#k}];\Theta),\mathcal{Y}_{k}^{l})+\\\
CE(\mathcal{F}(\mathcal{X}_{syn(i,j)}^{l},[Prompt_{\\#k}];\Theta),\mathcal{Y}_{k}^{l})k=i,j.\end{split}$
(4)

This formula can be concluded as follows: models can only show interest in the
specifically prompted task, when facing mixed data. So far, the supervised
loss $\mathcal{L}_{sup}$ is written as:

$\mathcal{L}_{sup}=\mathcal{L}_{lab}+\mathcal{L}_{aux}.$ (5)

In this way, VerSemi devises a synthetic Task#5, which paves the way for task-
agnostic learning from the unlabeled data, as well as augmenting the
capability of task prompt with the help of $\mathcal{L}_{aux}$.

Table 1: Performance comparison on the pancreas dataset, in the scenario of
leveraging 10% and 20% labeled data. The best and second best results are
shown in red and blue, respectively. ((Dice, %); (Jaccard, %); (ASD, voxel);
(95HD, voxel).)

Method | Pancreas (10%/6 labeled data) | Pancreas (20%/12 labeled data)
---|---|---
Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$
VNet (3DV’16) [19] | 55.60 | 41.74 | 18.63 | 45.33 | 72.38 | 58.26 | 5.89 | 19.35
UA-MT (MICCAI’19) [39] | 66.34 | 53.21 | 4.57 | 17.21 | 76.10 | 62.62 | 2.43 | 10.84
DTC (AAAI’19) [16] | 69.21 | 54.06 | 5.95 | 17.21 | 78.27 | 64.75 | 2.25 | 8.36
ASE-Net (TMI’22) [14] | 71.54 | 56.82 | 5.73 | 16.33 | 79.03 | 66.57 | 2.30 | 8.62
CAML (MICCAI’23) [8] | 71.21 | 56.32 | 5.92 | 16.89 | 79.81 | 67.35 | 2.27 | 8.22
BCP (CVPR’23) [2] | 73.83 | 59.24 | 3.72 | 12.71 | 82.91 | 70.97 | 2.25 | 6.43
CauSSL (ICCV’23) [18] | 72.34 | 57.43 | 3.13 | 13.49 | 80.63 | 67.84 | 2.78 | 8.76
MagicNet (CVPR’23) [3] | 75.01 | 62.04 | 3.97 | 13.71 | 81.25 | 68.81 | 2.83 | 8.50
VerSemi | 78.08 | 64.82 | 2.33 | 8.05 | 83.27 | 71.68 | 1.40 | 5.33
VerSemi w/ Task Info | 78.62 | 64.91 | 2.28 | 7.99 | 83.55 | 71.93 | 1.35 | 5.02

Table 2: Performance comparison on the spleen dataset, in the scenario of
leveraging 10% and 20% labeled data. The best and second best results are
shown in red and blue, respectively. ((Dice, %); (Jaccard, %); (ASD, voxel);
(95HD, voxel).)

Method | Spleen (10%/3 labeled data) | Spleen (20%/6 labeled data)
---|---|---
Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$
VNet (3DV’16) [19] | 75.14 | 65.27 | 15.02 | 43.89 | 79.78 | 72.86 | 11.37 | 30.03
UA-MT (MICCAI’19) [39] | 79.63 | 68.62 | 15.94 | 44.71 | 83.11 | 75.98 | 8.92 | 25.41
DTC (AAAI’19) [16] | 80.27 | 69.00 | 14.53 | 41.56 | 84.59 | 75.91 | 9.75 | 31.77
ASE-Net (TMI’22) [14] | 80.65 | 69.48 | 14.37 | 41.31 | 85.02 | 75.68 | 12.53 | 37.26
CAML (MICCAI’23) [8] | 80.32 | 69.10 | 15.37 | 41.71 | 85.80 | 76.79 | 11.57 | 36.14
BCP (CVPR’23) [2] | 83.12 | 72.85 | 14.42 | 42.11 | 87.02 | 78.58 | 10.48 | 37.08
CauSSL (ICCV’23) [18] | 81.98 | 71.25 | 14.69 | 41.84 | 86.83 | 78.46 | 10.01 | 32.27
MagicNet (CVPR’23) [3] | 83.55 | 73.58 | 13.49 | 41.79 | 88.24 | 80.24 | 8.50 | 23.51
VerSemi | 89.34 | 81.73 | 3.12 | 9.33 | 94.62 | 89.89 | 2.40 | 7.50
VerSemi w/ Task Info | 90.10 | 82.75 | 3.09 | 9.28 | 94.67 | 89.93 | 2.35 | 7.33

###  3.3 Task-agnostic Unlabeled Data Learning

Considering task specifics are required for prompt-driven model to generate
prompt, here we place our VerSemi in a more demanding SSL context, in which
unlabeled task specifics are not desired. Below we describe how this is
achieved. Firstly, CutMix [40] is conducted on all unlabeled data, making the
input contain objects of different tasks. Then the prediction with Task#5
Prompt is forced to be consistent with the aggregated prediction using Task#1
Prompt $\sim$ Task#4 Prompt (see Fig. 7). The aggregated prediction can be
regarded as a combination of pseudo-masks for each task, and the prediction
prompted by Task#5 can be considered as a direct pseudo-mask for all tasks.
Therefore, the two predictions should be identical. We call this operation
self-consistency since no extra decoder or teacher model is required for
supervision. The entire process can be written as:

$\displaystyle\mathcal{X}_{syn(i,j)}^{u}=\mathcal{X}_{i}^{u}$
$\displaystyle\odot\mathcal{M}+\mathcal{X}_{j}^{u}\odot(1-\mathcal{M})$ (6)
$\displaystyle\mathcal{P}_{agg}=\max\limits_{k\in(1,4)}(\mathcal{F}($
$\displaystyle\mathcal{X}_{syn(i,j)}^{u},[Prompt_{\\#k}];\Theta)),$

where $\mathcal{X}_{syn(i,j)}^{u}$ are mixed unlabeled data,
$\mathcal{X}_{i}^{u}$ and $\mathcal{X}_{j}^{u}$ are randomly selected
unlabeled data. Element-wise maximization is performed to aggregate
predictions prompted by Task#1 $\sim$ Task#4, and $\mathcal{P}_{agg}$ is the
final aggregated prediction. The overall loss $\mathcal{L}_{total}$ and
unsupervised loss $\mathcal{L}_{unsup}$ are calculated as:

$\displaystyle\mathcal{L}_{total}=$
$\displaystyle\mathcal{L}_{sup}+\mathcal{L}_{unsup}$ (7)
$\displaystyle\mathcal{L}_{unsup}=Dice(\mathcal{P}_{agg},$
$\displaystyle\mathcal{F}(\mathcal{X}_{syn(i,j)}^{u},[Prompt_{\\#5}];\Theta)).$

To summarize, based on the design of semantic-aware Task#5, our VerSemi learns
from unlabeled data in a task-agnostic way, and also enhances the the
uniqueness of the task prompt with the auxiliary constraint
$\mathcal{L}_{aux}$.

##  4 Experiments and Results

###  4.1 Setup

Datasets. We report the model segmentation results on four public datasets,
including Task#1: NIH-Pancreas [24], Task#2: Left Atrium [34], Task#3: MSD-
Spleen [1] and Task#4: MSD-Lung Tumor [1]. Specifically, NIH-Pancreas contains
82 contrast-enhanced abdomen CT scans, which are split into 62/20 scans for
training/test. The Left Atrium has 100 gadolinium-enhanced MR images, in which
80/20 images are leveraged for training/test. MSD-Spleen contains 41 CT scans,
and 30/11 scans are split for training/test. MSD-Lung Tumor contains 63 CT
scans, which are divided into 50/13 scans for training/test. Among them, 10%
training data are split into a validation set to select the best model. All
methods follow the same data split for fair comparisons, with the same pre-
processing as [16, 2].

Implementation Details. Following previous works [2, 16, 3], we adopt V-Net
[19] as the baseline model for fair comparisons. We use the Adam optimizer
[11] with a learning rate of 0.001. The input size and batch size are set to
96$\times$96$\times$96 and 8, respectively. Experiments were implemented by
Pytorch [21] with four NVIDIA GeForce RTX 3080 Ti GPUs. Evaluation metrics of
Dice (%), Jaccard (%), Average Surface Distance (ASD, voxel) and 95% Hausdorff
Distance (95HD, voxel) are used here.

###  4.2 Comparison with Existing Methods

We compare our VerSemi with seven popular SSL methods, including uncertainty-
aware mean-teacher (UA-MT) [39], dual-task consistency (DTC) [16], adversarial
consistency and dynamic convolution (ASE-Net) [14], correlation-aware mutual
learning (CAML) [8], bidirectional copy-paste [2], causality-inspired semi-
supervised segmentation (CauSSL) [18] and cubic volume partition and recovery
(Magic-Net) [3]. (Results on Left Atrium dataset and Lung Tumor dataset are
provided in the supplementary)

Figure 5: Visualization results of spleen segmentation by incorporating other tasks sequentially. We can find the model tends to produce more accurate segmentation with the increase of integrated tasks, demonstrating the benefits of learning a unified model. Table 3: A case study of the impact of other tasks on one specific task. Here spleen segmentation task (Task #3) is selected as the baseline, as we find VerSemi presents remarkable improvements on this task when compared to other methods, and this experiment aims to figure out where the performance gains come from. Setting | 10% labels | 20% labels
---|---|---
Dice $\uparrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | 95HD $\downarrow$
Task#3 | 75.14 | 43.89 | 79.78 | 30.03
Task#3+#1 | 85.62 | 17.07 | 90.00 | 15.81
Task#3+#1+#2 | 88.03 | 11.06 | 92.06 | 10.06
Task#3+#1+#2+#4 | 89.34 | 9.33 | 94.62 | 7.50

Results on Pancreas Dataset. As shown in Table 1, we can find VerSemi
consistently surpasses others on all metrics under different SSL settings. For
example, VerSemi brings respectively 3.07% and 5.66 (voxels) improvements on
Dice and HD score than the second best method MagicNet with 10% labeled data
training. Besides, we can find the performance gains obtained by VerSemi are
larger than the others when leveraging fewer labels, indicating the
effectiveness of our VerSemi in annotation-scarce scenarios.

Table 4: Adapting BCP [2] and CauSSL [18] into unified SSL models. However,
severe performance degradation can be seen when comparing Uni-BCP to BCP and
Uni-CauSSL to CauSSL.

Method | Pancreas (10%/6 labeled data) | Left Atrium (10%/8 labeled data)
---|---|---
Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$
Uni-BCP | 68.59 | 53.73 | 7.33 | 20.62 | 85.73 | 75.06 | 10.17 | 30.33
Uni-CauSSL | 65.35 | 49.09 | 6.16 | 20.89 | 83.40 | 72.43 | 8.84 | 34.94
VerSemi w/o $\mathcal{L}_{aux}$ | 75.06 | 60.94 | 3.70 | 11.64 | 88.56 | 79.81 | 2.62 | 9.17
VerSemi (Ours) | 78.08 | 64.82 | 2.33 | 8.05 | 89.01 | 80.52 | 2.57 | 9.03
Method | Spleen (10%/3 labeled data) | Lung Tumor (10%/5 labeled data)
Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | Jaccard $\uparrow$ | ASD $\downarrow$ | 95HD $\downarrow$
Uni-BCP | 74.80 | 58.89 | 17.11 | 54.06 | 31.01 | 21.32 | 11.35 | 24.36
Uni-CauSSL | 73.06 | 57.85 | 18.28 | 55.51 | 25.38 | 20.20 | 15.06 | 28.72
VerSemi w/o $\mathcal{L}_{aux}$ | 87.80 | 79.19 | 3.36 | 10.10 | 34.74 | 22.55 | 12.62 | 24.77
VerSemi (Ours) | 89.34 | 81.73 | 3.12 | 9.33 | 36.90 | 28.12 | 10.87 | 23.41

Results on Spleen Dataset. Table 2 presents the results of spleen
segmentation. We can see that our VerSemi significantly outperforms all other
competitors by a large margin. For instance, compared to MagicNet, VerSemi has
respectively 32.46 (voxels) and 16.01 (voxels) performance gains on HD score
with 10% and 20% labeled data. Similarly, VerSemi surpasses CauSSL by 7.36% on
Dice score under 10% label percentage. A case study is conducted to see the
performance gains by introducing other tasks on spleen segmentation (Task#3).
As Table 3 and Fig. 5 show, there are consistent improvements by gradually
integrating data of other tasks. In particular, we can find the performance
improved greatly by integrating the pancreas segmentation (Task#1), where our
model has already outperformed other methods. Here we provide three potential
reasons for VerSemi’s high performance on this task: (1) due to extremely
limited labels (10% labels are equal to 3 labeled data), competitors fail to
generalize the representation learned from labeled data to unlabeled data, and
mistakenly predict the background as foreground (see Row 5-6 of Fig.10).
Therefore, a high HD score can be observed; (2) since the same modality,
$e.g.$, Task#1 and Task#3, VerSemi can learn modality-specific knowledge and
achieve better performance. (_e.g_., see Fig. 6, the feature embedding of
pancreas and spleen are very close in the latent space, both of them are
abdominal organs); and (3) by segmenting other organs, VerSemi can segment the
background regions and identify the adhesive boundaries in a negative learning
mechanism, so as to decrease the HD score.

Figure 6: t-SNE visualization of feature embedding for four tasks. The
implemented Uni-CauSSL, Uni-BCP and our proposed VerSemi are compared.

###  4.3 In-depth Analysis

Importance of the auxiliary constraint $\mathcal{L}_{aux}$.
$\mathcal{L}_{aux}$ plays the role of augmenting the uniqueness of task
prompts. As the last two rows of Table 4 indicates, by incorporating
$\mathcal{L}_{aux}$, VerSemi presents respectively 3.02%, 0.45%, 1.54% and
2.16% performance gains on Dice score on the pancreas, left atrium, spleen and
lung tumor tasks, when using 10% labeled data. This improvement demonstrates
the effectiveness and necessity of adding an accessory loss to constrain the
feasibility of task prompts.

Figure 7: Visualization results prompted by pertinent tasks and Task#5, when
meeting mixed unlabeled data. In this case, the spleen and left atrium are
mixed together. The inconsistent regions between aggregated prediction and
Task#5-prompted prediction are highlighted by red elliptic.

Adapting single SSL models into unified SSL models. In this experiment, we
revise CauSSL and BCP into the unified SSL settings. There are two changes
compared to their previous versions. (1) The input data cover four tasks and
are randomly fed into the model with the associated task id. (2) Changing the
number of output channels to match the number of tasks, which is different
from VerSemi as VerSemi has a dynamic task-prompted head with two output
channels. As Table 4 shows, the results produced by Uni-BCP and Uni-CauSSL are
far inferior to VerSemi, and compared to their original single model version,
significant performance degradation can be observed. For instance, according
to the averaged Dice score with 10% labeled data, BCP vs Uni-BCP (70.79% vs
65.03%), 5.76% drop can be found. And CauSSL vs Uni-CauSSL (69.60% vs 61.80%),
7.80% degradation is discovered. This phenomenon is mainly triggered by
chaotic representation learned from all task data, and also indicates that
naively learning from all tasks simultaneously is not effective and even
harmful to the single task. Moreover, we plot the t-SNE visualization of
feature embedding to have a clear view. As Fig. 6 exhibits, VerSemi presents a
distinguishable decision boundary while others show mixed and dispersed
embedding. This demonstrates that task prompts and the constraint to task
prompts ($\mathcal{L}_{aux}$) are essential when facing multiple SSL tasks, as
the former guides model to have a clear understanding of the ongoing task,
while the latter makes sure the learned representation of each single task are
discernible and concentrated.

Figure 8: (a) VerSemi is designed to learn from unlabeled data without knowing
associated task info. (b) The pipeline of VerSemi when unlabeled task info is
given.

Visualization of unlabeled data learning pipeline. Fig. 7 shows the
segmentation results prompted by pertinent tasks and synthetic Task#5. We can
see that VerSemi can clearly recognize the task-specific prompted regions,
which is largely benefited by the auxiliary constraint $\mathcal{L}_{aux}$,
for its ability to enhance the controllability of prompts. Meanwhile, we can
also find VerSemi smoothly highlights all task semantic regions under the
prompt of Task#5, demonstrating the effectiveness of learning a semantic-aware
synthetic task. By aligning the two predictions (see the bottom of Fig. 7,
Aggregated and PerdTask#5), VerSemi learns the unlabeled data in a task-
agnostic manner.

Incorporating unlabeled task information into VerSemi. To explore the upper
bound of VerSemi, we report the results when feeding task information of
unlabeled data into VerSemi. As Fig. 8 presents, VerSemi w/ task info can
directly generate predictions on the source image with task-specific prompts,
whereas VerSemi w/o task info should first generate predictions on the mixed
data with all pertinent task prompts and then aggregate them. From the last-
row results of Table 1 and Table 2 with gray background, we can find there is
an improvement (_i.e_., a 0.78% gain on the averaged Dice score with 10%
labels), demonstrating the ability to produce accurate predictions with mixed
data, as well as distinguishing task-prompted specific regions.

Table 5: Discussion of three types of prompt, in which language prompt, soft and one-hot vector prompts are compared. The averaged Dice and 95HD scores on four tasks are reported. types of prompt | 10% labels | 20% labels
---|---|---
Dice $\uparrow$ | 95HD $\downarrow$ | Dice $\uparrow$ | 95HD $\downarrow$
language (CLIP) | 67.21 | 22.78 | 74.30 | 19.29
soft vector | 70.02 | 17.66 | 77.45 | 13.83
one-hot vector | 73.33 | 12.46 | 80.99 | 8.74

Discussion of task prompt. Prompt is used as a signal to help the model
understand the ongoing task, typically language [36, 15] (using a sentence to
describe), soft vector [28, 38] (using randomly initialized learnable vector
to represent) and one-hot prompts [43, 37] are mostly employed. As Table 5
shows, the one-hot prompt performs best under SSL setting. Reasons for the
results are: (1) the embedding of language heavily relies on language or
vision-language models, which is not guaranteed to be aligned with the
extracted medical image embedding; (2) soft vector prompt works when there are
substantial paired image-label data, whereas only scarce labels are available
in the context of SSL, making it hard to adapt. By contrast, the one-hot
prompt is more explicit and empirically suitable for SSL.

Figure 9: Kernel density estimation of VerSemi and BCP when training with 10%
labels. This experiment is conducted to see the distribution of labeled and
unlabeled data. Top is the distribution generated by VerSemi and Bottom is BCP
[2]. It is clear to see that VerSemi aligns the distribution better than BCP.
Figure 10: Segmentation results produced by different methods. Row 1-2:
pancreas segmentation; Row 3-4: left atrium segmentation; Row 5-6: spleen
segmentation and Row 7-8: lung tumor segmentation.

Learned distribution on labeled and unlabeled data. Distribution mismatch
between labeled and unlabeled data is a commonly encountered issue in SSL,
which is mainly caused by unbalanced/partial distribution learned from labeled
data [42, 3]. Fig. 9 presents the kernel density estimation of VerSemi and BCP
when training with 10% label percentage. We can find that: (1) for Task#2
(left atrium segmentation) with large data scale, both VerSemi and BCP show
well-aligned distribution, which is mainly attributed to ample representation
learned from labeled data, thus models can successfully generalize to
unlabeled data and present comparable performance; (2) as for the other tasks,
severe inconsistency is observed for BCP, whereas VerSemi significantly aligns
the learned distribution. This demonstrates that properly learning tasks
concurrently is beneficial to unlabeled data mining, since the mismatch issue
between labeled and unlabeled data is largely alleviated.

Visualization of segmentation on four benchmarks. Fig. 10 shows the
segmentation results, it is clear to find that VerSemi can generate the most
accurate mask compared to competitors. For example, for spleen segmentation
(Row 5-6), other SSL methods extensively predict the background as the
foreground, whereas VerSemi successfully distinguishes the region of spleen.

##  5 Conclusion

In this paper, we have presented an effective model VerSemi for semi-
supervised medical image segmentation, with the new setting of integrating
various tasks into a unified framework. Specifically, our VerSemi deals with
different tasks in a dynamic way through the design of task prompts. A novel
contrastive constraint is proposed to improve the controllability of dynamic
task prompts, so as to distinguish different task information. Extensive
experiments on four public datasets clearly demonstrate the effectiveness of
our proposed VerSemi model, especially with limited training labels, setting
new SOTA performance for semi-supervised medical image segmentation.

Limitation and Future Work. Since our model was trained on several available
but limited datasets, the inter-dataset conflicts would unavoidably impact the
training. Future work will include the study of a de-biased strategy for
further investigation.

## References

  * Antonelli et al. [2022] Michela Antonelli, Annika Reinke, Spyridon Bakas, Keyvan Farahani, Annette Kopp-Schneider, Bennett A Landman, Geert Litjens, Bjoern Menze, Olaf Ronneberger, Ronald M Summers, et al.  The medical segmentation decathlon.  _Nat. Commun_ , 13(1):4128, 2022.
  * Bai et al. [2023] Yunhao Bai, Duowen Chen, Qingli Li, Wei Shen, and Yan Wang.  Bidirectional copy-paste for semi-supervised medical image segmentation.  In _CVPR_ , pages 11514–11524, 2023.
  * Chen et al. [2023a] Duowen Chen, Yunhao Bai, Wei Shen, Qingli Li, Lequan Yu, and Yan Wang.  Magicnet: Semi-supervised multi-organ segmentation via magic-cube partition and recovery.  In _CVPR_ , pages 23869–23878, 2023a.
  * Chen et al. [2022a] Hao Chen, Ran Tao, Yue Fan, Yidong Wang, Jindong Wang, Bernt Schiele, Xing Xie, Bhiksha Raj, and Marios Savvides.  Softmatch: Addressing the quantity-quality tradeoff in semi-supervised learning.  In _ICLR_ , 2022a.
  * Chen et al. [2023b] Jieneng Chen, Yingda Xia, Jiawen Yao, Ke Yan, Jianpeng Zhang, Le Lu, Fakai Wang, Bo Zhou, Mingyan Qiu, Qihang Yu, et al.  Towards a single unified model for effective detection, segmentation, and diagnosis of eight major cancers using a large collection of ct scans.  _arXiv preprint arXiv:2301.12291_ , 2023b.
  * Chen et al. [2022b] Yanbei Chen, Massimiliano Mancini, Xiatian Zhu, and Zeynep Akata.  Semi-supervised and unsupervised deep visual learning: A survey.  _IEEE Trans. Pattern Anal. Mach. Intell_ , 2022b.
  * Cheplygina et al. [2019] Veronika Cheplygina, Marleen de Bruijne, and Josien PW Pluim.  Not-so-supervised: a survey of semi-supervised, multi-instance, and transfer learning in med. image anal.  _Med. Image Anal_ , 54:280–296, 2019.
  * Gao et al. [2023] Shengbo Gao, Ziji Zhang, Jiechao Ma, Zihao Li, and Shu Zhang.  Correlation-aware mutual learning for semi-supervised medical image segmentation.  In _MICCAI_ , 2023.
  * Jiao et al. [2022] Rushi Jiao, Yichi Zhang, Le Ding, Rong Cai, and Jicong Zhang.  Learning with limited annotations: a survey on deep semi-supervised learning for medical image segmentation.  _arXiv preprint arXiv:2207.14191_ , 2022.
  * Kim et al. [2022] Donggyun Kim, Jinwoo Kim, Seongwoong Cho, Chong Luo, and Seunghoon Hong.  Universal few-shot learning of dense prediction tasks with visual token matching.  In _ICLR_ , 2022.
  * Kingma and Ba [2014] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kirillov et al. [2023] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al.  Segment anything.  _arXiv preprint arXiv:2304.02643_ , 2023.
  * Lee et al. [2023] Hyungyung Lee, Wonjae Kim, Jin-Hwa Kim, Tackeun Kim, Jihang Kim, Leonard Sunwoo, and Edward Choi.  Unified chest x-ray and radiology report generation model with multi-view chest x-rays.  _arXiv preprint arXiv:2302.12172_ , 2023.
  * Lei et al. [2022] Tao Lei, Dong Zhang, Xiaogang Du, Xuan Wang, Yong Wan, and Asoke K Nandi.  Semi-supervised medical image segmentation using adversarial consistency learning and dynamic convolution network.  _IEEE Trans. Med. Imaging_ , 2022.
  * Liu et al. [2023] Jie Liu, Yixiao Zhang, Jie-Neng Chen, Junfei Xiao, Yongyi Lu, Bennett A Landman, Yixuan Yuan, Alan Yuille, Yucheng Tang, and Zongwei Zhou.  Clip-driven universal model for organ segmentation and tumor detection.  In _ICCV_ , pages 21152–21164, 2023.
  * Luo et al. [2021] Xiangde Luo, Jieneng Chen, Tao Song, and Guotai Wang.  Semi-supervised medical image segmentation through dual-task consistency.  In _AAAI_ , pages 8801–8809, 2021.
  * Mey and Loog [2022] Alexander Mey and Marco Loog.  Improved generalization in semi-supervised learning: A survey of theoretical results.  _IEEE Trans. Pattern Anal. Mach. Intell_ , 45(4):4747–4767, 2022.
  * Miao et al. [2023] Juzheng Miao, Cheng Chen, Furui Liu, Wei Hao, and Pheng-Ann Heng.  Caussl: Causality-inspired semi-supervised learning for medical image segmentation.  In _ICCV_ , 2023.
  * Milletari et al. [2016] Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi.  V-net: Fully convolutional neural networks for volumetric medical image segmentation.  In _3DV_ , pages 565–571. IEEE, 2016.
  * Miyato et al. [2018] Takeru Miyato, Shin-ichi Maeda, Masanori Koyama, and Shin Ishii.  Virtual adversarial training: a regularization method for supervised and semi-supervised learning.  _IEEE Trans. Pattern Anal. Mach. Intell_ , 41(8):1979–1993, 2018.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _NeurIPS_ , 32, 2019.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _ICML_ , pages 8748–8763. PMLR, 2021.
  * Rizve et al. [2020] Mamshad Nayeem Rizve, Kevin Duarte, Yogesh S Rawat, and Mubarak Shah.  In defense of pseudo-labeling: An uncertainty-aware pseudo-label selection framework for semi-supervised learning.  In _ICLR_ , 2020.
  * Roth et al. [2015] Holger R Roth, Le Lu, Amal Farag, Hoo-Chang Shin, Jiamin Liu, Evrim B Turkbey, and Ronald M Summers.  Deeporgan: Multi-level deep convolutional networks for automated pancreas segmentation.  In _MICCAI_ , pages 556–564. Springer, 2015.
  * Shen et al. [2017] Dinggang Shen, Guorong Wu, and Heung-Il Suk.  Deep learning in med. image anal.  _Annu Rev Biomed Eng_ , 19:221–248, 2017.
  * Sohn et al. [2020] Kihyuk Sohn, David Berthelot, Nicholas Carlini, Zizhao Zhang, Han Zhang, Colin A Raffel, Ekin Dogus Cubuk, Alexey Kurakin, and Chun-Liang Li.  Fixmatch: Simplifying semi-supervised learning with consistency and confidence.  _NeurIPS_ , 33:596–608, 2020.
  * Van Engelen and Hoos [2020] Jesper E Van Engelen and Holger H Hoos.  A survey on semi-supervised learning.  _Mach Learn_ , 109(2):373–440, 2020.
  * Vu et al. [2021] Tu Vu, Brian Lester, Noah Constant, Rami Al-Rfou, and Daniel Cer.  Spot: Better frozen model adaptation through soft prompt transfer.  _arXiv preprint arXiv:2110.07904_ , 2021.
  * Wang et al. [2022] Yidong Wang, Hao Chen, Qiang Heng, Wenxin Hou, Yue Fan, Zhen Wu, Jindong Wang, Marios Savvides, Takahiro Shinozaki, Bhiksha Raj, et al.  Freematch: Self-adaptive thresholding for semi-supervised learning.  In _ICLR_ , 2022.
  * Wang et al. [2023] Zhao Wang, Chang Liu, Shaoting Zhang, and Qi Dou.  Foundation model for endoscopy video analysis via large-scale self-supervised pre-train.  _arXiv preprint arXiv:2306.16741_ , 2023.
  * Wu et al. [2022a] Yicheng Wu, Zongyuan Ge, Donghao Zhang, Minfeng Xu, Lei Zhang, Yong Xia, and Jianfei Cai.  Mutual consistency learning for semi-supervised medical image segmentation.  _Med. Image Anal_ , 81:102530, 2022a.
  * Wu et al. [2022b] Yicheng Wu, Zhonghua Wu, Qianyi Wu, Zongyuan Ge, and Jianfei Cai.  Exploring smoothness and class-separation for semi-supervised medical image segmentation.  In _MICCAI_ , pages 34–43. Springer, 2022b.
  * Xie et al. [2022] Yutong Xie, Jianpeng Zhang, Yong Xia, and Qi Wu.  Unimiss: Universal medical self-supervised learning via breaking dimensionality barrier.  In _ECCV_ , pages 558–575. Springer, 2022.
  * Xiong et al. [2021] 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.  _Med. Image Anal_ , 67:101832, 2021.
  * Xue et al. [2023] Le Xue, Mingfei Gao, Chen Xing, Roberto Martín-Martín, Jiajun Wu, Caiming Xiong, Ran Xu, Juan Carlos Niebles, and Silvio Savarese.  Ulip: Learning a unified representation of language, images, and point clouds for 3d understanding.  In _CVPR_ , pages 1179–1189, 2023.
  * Yao et al. [2023] Hantao Yao, Rui Zhang, and Changsheng Xu.  Visual-language prompt tuning with knowledge-guided context optimization.  In _CVPR_ , pages 6757–6767, 2023.
  * Yasunaga et al. [2022] Michihiro Yasunaga, Jure Leskovec, and Percy Liang.  Linkbert: Pretraining language models with document links.  _arXiv preprint arXiv:2203.15827_ , 2022.
  * Ye et al. [2023] Yiwen Ye, Yutong Xie, Jianpeng Zhang, Ziyang Chen, and Yong Xia.  Uniseg: A prompt-driven universal segmentation model as well as a strong representation learner.  _arXiv preprint arXiv:2304.03493_ , 2023.
  * Yu et al. [2019] 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 _MICCAI_ , pages 605–613. Springer, 2019.
  * Yun et al. [2019] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In _ICCV_ , pages 6023–6032, 2019.
  * Zeng et al. [2023a] Qingjie Zeng, Yutong Xie, Zilin Lu, Mengkang Lu, and Yong Xia.  Discrepancy matters: Learning from inconsistent decoder features for consistent semi-supervised medical image segmentation.  _arXiv preprint arXiv:2309.14819_ , 2023a.
  * Zeng et al. [2023b] Qingjie Zeng, Yutong Xie, Zilin Lu, and Yong Xia.  Pefat: Boosting semi-supervised medical image classification via pseudo-loss estimation and feature adversarial training.  In _CVPR_ , pages 15671–15680, 2023b.
  * Zhang et al. [2021] Jianpeng Zhang, Yutong Xie, Yong Xia, and Chunhua Shen.  Dodnet: Learning to segment multi-organ and tumors from multiple partially labeled datasets.  In _CVPR_ , pages 1195–1204, 2021.
  * Zhang et al. [2022] Wenqiao Zhang, Lei Zhu, James Hallinan, Shengyu Zhang, Andrew Makmur, Qingpeng Cai, and Beng Chin Ooi.  Boostmis: Boosting medical image semi-supervised learning with adaptive pseudo labeling and informative active annotation.  In _CVPR_ , pages 20666–20676, 2022.
  * Zhou et al. [2022] Hong-Yu Zhou, Chenyu Lian, Liansheng Wang, and Yizhou Yu.  Advancing radiograph representation learning with masked record modeling.  In _ICLR_ , 2022.


# A preliminary approach to knowledge integrity risk assessment in Wikipedia
projects

Pablo Aragón  Wikimedia FoundationBarcelonaSpain<EMAIL_ADDRESS>and
Diego Sáez-Trumper  Wikimedia FoundationBarcelonaSpain<EMAIL_ADDRESS>
(2021)

###### Abstract.

Wikipedia is one of the main repositories of free knowledge available today,
with a central role in the Web ecosystem. For this reason, it can also be a
battleground for actors trying to impose specific points of view or even
spreading disinformation online. There is a growing need to monitor its
“health” but this is not an easy task. Wikipedia exists in over 300 language
editions and each project is maintained by a different community, with their
own strengths, weaknesses and limitations. In this paper, we introduce a
taxonomy of knowledge integrity risks across Wikipedia projects and a first
set of indicators to assess internal risks related to community and content
issues, as well as external threats such as the geopolitical and media
landscape. On top of this taxonomy, we offer a preliminary analysis
illustrating how the lack of editors’ geographical diversity might represent a
knowledge integrity risk. These are the first steps of a research project to
build a Wikipedia Knowledge Integrity Risk Observatory.

Wikipedia, knowledge integrity, risk assessment

††copyright: rightsretained††journalyear: 2021††conference: MIS2’21:
Misinformation and Misbehavior Mining on the Web Workshop held in conjunction
with KDD 2021; August 15, 2021; Online††booktitle: MIS2’21: Misinformation and
Misbehavior Mining on the Web Workshop held in conjunction with KDD 2021,
August 15, 2021, Online

##  1\. Introduction

The Web has become the largest repository of knowledge ever known in just
three decades. However, we are witnessing in recent years the proliferation of
sophisticated strategies that are heavily affecting the reliability and
trustworthiness of online information. Web platforms are increasingly
encountering misinformation problems caused by deception techniques such as
astroturfing (Zhang, Jerry and Carpenter, Darrell and Ko, Myung, 2013),
harmful bots (Ferrara, Emilio and Varol, Onur and Davis, Clayton and Menczer,
Filippo and Flammini, Alessandro, 2016), computational propaganda (Woolley,
Samuel C and Howard, Philip N, 2018), sockpuppetry (Kumar, Srijan and Cheng,
Justin and Leskovec, Jure and Subrahmanian, VS, 2017), data voids
(Golebiewski, Michael and Boyd, Danah, 2018), etc.

Wikipedia, the world’s largest online encyclopedia in which millions of
volunteer contributors create and maintain free knowledge, is not free from
the aforementioned problems. Disinformation is one of its most relevant
challenges (Saez-Trumper, Diego, 2019) and some editors devote a substantial
amount of their time in patrolling tasks in order to detect vandalism and make
sure that new contributions fulfill community policies and guidelines
111https://en.wikipedia.org/wiki/Wikipedia:Policies_and_guidelines (Morgan,
Jonathan, 2019). Furthermore, _Knowledge Integrity_ is one of the strategic
programs of Wikimedia Research with the goal of identifying and addressing
threats to content on Wikipedia, increasing the capabilities of patrollers,
and providing mechanisms for assessing the reliability of sources (Zia, Leila
and Johnson, Isaac and Mansurov, Bahodir and Morgan, Jonathan and Redi, Miriam
and Saez-Trumper, Diego and Taraborelli, Dario, 2019).

Many lessons have been learnt from fighting misinformation in Wikipedia
(Kelly, Heather, 2021) and analyses of recent cases like the 2020 United
States presidential election have suggested that the platform was better
prepared than major social media outlets (Morrison, Sara, 2020). However,
there are Wikipedia editions in more than 300 languages, with very different
contexts. To provide Wikipedia communities with an actionable monitoring
system, this paper introduces a preliminary approach consisting of a taxonomy
of knowledge integrity risks in Wikipedia projects, based on a review of the
state of the art literature, and an initial set of indicators to be the core
of a Wikipedia Knowledge Integrity Risk Observatory.

##  2\. Taxonomy of knowledge integrity risks

Risks to knowledge integrity in Wikipedia can arise in many and diverse forms.
Inspired by a recent work that has proposed a taxonomy of knowledge gaps for
Wikimedia projects (Redi, Miriam and Gerlach, Martin and Johnson, Isaac and
Morgan, Jonathan and Zia, Leila, 2021), we have reviewed works by the
Wikimedia Foundation, academic researchers and journalists that provided
empirical evidence of knowledge integrity risks. Then, we have classified them
by developing a hierarchical categorical structure that is presented in Figure
1.

Figure 1. Taxonomy of knowledge integrity risks in Wikipedia.

We initially differentiate between internal and external risks according to
their origin. The former correspond to issues specific to the Wikimedia
ecosystem while the latter involve activity from other environments, both
online and offline. For internal risks, we have identified the following
categories focused on either community or content:

  * •

Community capacity: Pool of resources of the community. Risks have been found
when, given the size of the corresponding Wikipedia project (volume of
articles, edits and editors), there is a shortage of patrolling resources,
i.e., editors with elevated user rights (e.g., admins, checkusers,
oversighters) (Saez-Trumper, Diego, 2019; Morgan, Jonathan, 2019; Song,
Victoria, 2020; Sato, Yumiko, 2021) and tools (e.g., specialized bots,
scripts, MediaWiki extensions, desktop/web apps) (Saez-Trumper, Diego, 2019;
Morgan, Jonathan, 2019).

  * •

Community governance: Situations and procedures involving decision-making
within the community. The reviewed literature has identified risks like the
unavailability of local rapid-response noticeboards on smaller wikis (Morgan,
Jonathan, 2019) or the abuse of blocking practices by admins (Rogers, Richard
and Sendijarevic, Emina and others, 2012; Sato, Yumiko, 2021; Song, Victoria,
2020).

  * •

Community demographics: Characteristics of community members. Some analyses
highlight that the lack of geographical diversity might favor nationalistic
biases (Rogers, Richard and Sendijarevic, Emina and others, 2012; Sato,
Yumiko, 2021). Other relevant dimensions are editors’ age and activity since
misbehavior is often observed in editing patterns of newly created accounts
(Kumar, Srijan and Spezzano, Francesca and Subrahmanian, VS, 2015; Kumar,
Srijan and West, Robert and Leskovec, Jure, 2016; Joshi, Nikesh and Spezzano,
Francesca and Green, Mayson and Hill, Elijah, 2020) or accounts that have been
inactive for a long period to avoid certain patrolling systems or that are no
longer monitored and became hacked (Morgan, Jonathan, 2019).

  * •

Content verifiability: Usage and reliability of sources in articles. This
category is directly inspired by one the three principal core content policies
of Wikipedia (WP:V) 222https://en.wikipedia.org/wiki/Wikipedia:Verifiability
which states that readers and editors must be able to check that information
comes from a reliable source. It is referred to in several studies of content
integrity (Redi, Miriam and Fetahu, Besnik and Morgan, Jonathan and
Taraborelli, Dario, 2019; Lewoniewski, Wlodzimierz and Wecel, Krzysztof and
Abramowicz, Witold, 2019; Saez-Trumper, Diego, 2019).

  * •

Content quality: Criteria used for the assessment of article quality. Since
each Wikipedia language community decides its own standards and grading scale,
some works have explored language-agnostic signals of content quality such as
the volume of edits and editors or the appearance of specific templates
(Lewoniewski, Wlodzimierz and Wecel, Krzysztof and Abramowicz, Witold, 2017,
2019). In fact, there might exist distinctive cultural quality mechanisms as
these metrics do not always correlate with featured status of articles
(Rogers, Richard and Sendijarevic, Emina and others, 2012).

  * •

Content controversiality: Disputes between community members due to
disagreements about the content of articles. Edit wars are the best known
phenomenon that occurs when content becomes controversial (Yasseri, Taha and
Spoerri, Anselm and Graham, Mark and Kertész, János, 2014; Rogers, Richard and
Sendijarevic, Emina and others, 2012), sometimes requiring articles to be
protected (Spezzano, Francesca and Suyehira, Kelsey and Gundala, Laxmi Amulya,
2019).

For external risks, we have identified the following categories:

  * •

Media References and visits to the Wikipedia project from other external media
on the Internet. Unusual amount of traffic to specific articles coming from
social media sites or search engines may be a sign of coordinated vandalism
(Morgan, Jonathan, 2019).

  * •

Geopolitics: Political context of the community and content of the Wikipedia
project. Some well resourced interested parties (e.g., corporations, nations)
might be interested in externally-coordinated long-term disinformation
campaigns in specific projects (Morgan, Jonathan, 2019; Shubber, Kadhim,
2021).

##  3\. Initial set of indicators and preliminary analysis

To capture risks in each category of the proposed taxonomy, we have compiled
an initial set of indicators presented in Table 1. The criteria for proposing
these indicators are that they should be simple to be easily interpreted by
non-technical stakeholders, comparable across wikis, language-agnostic and
periodically updatable. For this reason, they are counts of items (e.g.,
articles, editors, edits, etc.) or distributions of items over informative
features.

Table 1. Initial set of indicators for a Wikipedia Knowledge Integrity Risk Observatory. Risk category |  Candidate indicators
---|---
Community capacity |  Number of articles, editors, active editors, editors with elevated user rights (admins, bureaucrats, checkusers, oversighters, rollbackers); ratio of active editors with elevated user rights; number of specialized patrolling tools; number of AbuseFilter rules.
Community governance |  Number of requests in steward’s noticeboard; number of global stewards knowledgeable with that language; number of requests for comment (local and meta); ratio of articles for deletion; ratio of blocked accounts (spam, long-term abuse, etc.).
Community demographics |  Distribution of views and edits by country; distribution of active editors by age, local activity and cross-wiki activity.
Content verifiability |  Distribution of articles by number of citations, number of scientific citations and number of citation and verifiability article maintenance templates, distribution of sources by reliability.
Content quality |  Ratio of stub articles; editing depth; distribution of articles by community quality grading, ORES scoring (Halfaker, Aaron and Geiger, R Stuart, 2020), number of editors, number of quality flaw templates, distribution of edits by source type (i.e., editor, newly-registered editor, admin, bot, IP).
Content controversiality |  Ratio of locked articles; distribution of articles by controversiality (Yasseri, Taha and Spoerri, Anselm and Graham, Mark and Kertész, János, 2014), distribution articles by number of comments in discussion page and n-chains in discussion pages (Laniado, David and Tasso, Riccardo and Volkovich, Yana and Kaltenbrunner, Andreas, 2011).
Media |  Distribution of mentions/references and visits by online media outlets, social media platforms and search engines.
Geopolitics |  Democratic quality scores derived from views and edits by country and well-established country democratic indexes (e.g., (Coppedge, Michael and Gerring, John and Lindberg, Staffan I and Skaaning, Svend-Erik and Teorell, Jan and Altman, David and Bernhard, Michael and Fish, M Steven and Glynn, Adam and Hicken, Allen and others, 2017; Schenkkan, Nate and Repucci, Sarah, 2019)).

To illustrate the value of the indicators for knowledge integrity risk
assessment in Wikipedia, we provide an example on community demographics, in
particular, geographical diversity. Figure 2 shows the entropy value of the
distributions of number of edits and views by country of the language editions
with over 500K articles. The data has been collected from November 2018 to
April 2021. On the one hand, we observe large entropy values for both edits
and views in the Arabic, English and Spanish editions, i.e., global
communities. On the other hand, other large language editions like the
Italian, Indonesian, Polish, Korean or Vietnamese Wikipedia lack that
geographical diversity. We should highlight the extraordinarily low entropy of
views of the Japanese Wikipedia, which supports one of the main causes
attributed to misinformation incidents in this edition (Sato, Yumiko, 2021).
We also notice the misalignment between high edit entropy and low view entropy
values in Cebuano and Waray-Waray editions, which might be the result of the
large fraction of content produced by bots distributed around the world (Song,
Victoria, 2020). It is also remarkable the misalignment of the Egyptian Arabic
Wikipedia with much larger entropy values for views than edits.

Figure 2. Entropy values ($S$) of the distributions of the number of edits and
views by country of the Wikipedia language editions, identified by the ISO
639-1 code, with over 500K articles. The graph includes a linear regression
model fit.

##  4\. Discussion and future work

This article has presented a preliminary approach to knowledge integrity risk
assessment in Wikipedia projects. We have covered the first steps of an
ongoing process with the ultimate goal of building a Wikipedia Knowledge
Integrity Risk Observatory. The taxonomy relies on risks detected in previous
work on knowledge integrity in Wikipedia, nevertheless, it could be enriched
through the review of additional literature on risk assessment in other web
platforms.

The analysis shown in this paper has focused exclusively on community
demographics. We will extend this work by implementing the rest of the
indicators of this and other categories from the taxonomy to assess their
informative value. As mentioned above, the current indicators are essentially
item counts and distributions of items over features. Future work will also
focus on defining advanced metrics while preserving the criteria of ease of
interpretation, comparability across wikis and language-agnosticism. Also,
indicators should be periodically updatable to allow longitudinal
observations. Another future challenge will be to define indicators with finer
levels of granularity, that is to say, metrics computed not only on an entire
Wikipedia project but on categories, pages, etc.

Last but not least, following the principles of openness, transparency and
accountability of Wikimedia
333https://meta.wikimedia.org/wiki/Wikimedia_Foundation_Guiding_Principles, we
expect to release the Wikipedia Knowledge Integrity Risk Observatory as a
dashboard available to the global movement of volunteers. Therefore, we will
also focus on designing an open technological infrastructure to provide
Wikimedia communities with valuable information on knowledge integrity.

## References

  * (1)
  * Coppedge, Michael and Gerring, John and Lindberg, Staffan I and Skaaning, Svend-Erik and Teorell, Jan and Altman, David and Bernhard, Michael and Fish, M Steven and Glynn, Adam and Hicken, Allen and others (2017) Coppedge, Michael and Gerring, John and Lindberg, Staffan I and Skaaning, Svend-Erik and Teorell, Jan and Altman, David and Bernhard, Michael and Fish, M Steven and Glynn, Adam and Hicken, Allen and others. 2017.  V-dem dataset v7.  (2017).
  * Ferrara, Emilio and Varol, Onur and Davis, Clayton and Menczer, Filippo and Flammini, Alessandro (2016) Ferrara, Emilio and Varol, Onur and Davis, Clayton and Menczer, Filippo and Flammini, Alessandro. 2016\.  The rise of social bots.  _Communications of the ACM_ 59, 7 (2016), 96–104.
  * Golebiewski, Michael and Boyd, Danah (2018) Golebiewski, Michael and Boyd, Danah. 2018\.  Data voids: Where missing data can easily be exploited.  (2018).
  * Halfaker, Aaron and Geiger, R Stuart (2020) Halfaker, Aaron and Geiger, R Stuart. 2020\.  Ores: Lowering barriers with participatory machine learning in wikipedia.  _Proceedings of the ACM on Human-Computer Interaction_ 4, CSCW2 (2020), 1–37.
  * Joshi, Nikesh and Spezzano, Francesca and Green, Mayson and Hill, Elijah (2020) Joshi, Nikesh and Spezzano, Francesca and Green, Mayson and Hill, Elijah. 2020.  Detecting Undisclosed Paid Editing in Wikipedia. In _Proceedings of The Web Conference 2020_. 2899–2905.
  * Kelly, Heather (2021) Kelly, Heather. 2021\.  On its 20th birthday, Wikipedia might be the safest place online.  {https://www.washingtonpost.com/technology/2021/01/15/wikipedia-20-year-anniversary/}
  * Kumar, Srijan and Cheng, Justin and Leskovec, Jure and Subrahmanian, VS (2017) Kumar, Srijan and Cheng, Justin and Leskovec, Jure and Subrahmanian, VS. 2017.  An army of me: Sockpuppets in online discussion communities. In _Proceedings of the 26th International Conference on World Wide Web_. 857–866.
  * Kumar, Srijan and Spezzano, Francesca and Subrahmanian, VS (2015) Kumar, Srijan and Spezzano, Francesca and Subrahmanian, VS. 2015.  Vews: A wikipedia vandal early warning system. In _Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining_. 607–616.
  * Kumar, Srijan and West, Robert and Leskovec, Jure (2016) Kumar, Srijan and West, Robert and Leskovec, Jure. 2016.  Disinformation on the web: Impact, characteristics, and detection of wikipedia hoaxes. In _Proceedings of the 25th international conference on World Wide Web_. 591–602.
  * Laniado, David and Tasso, Riccardo and Volkovich, Yana and Kaltenbrunner, Andreas (2011) Laniado, David and Tasso, Riccardo and Volkovich, Yana and Kaltenbrunner, Andreas. 2011\.  When the Wikipedians talk: Network and tree structure of Wikipedia discussion pages. In _Proceedings of the International AAAI Conference on Web and Social Media_ , Vol. 5.
  * Lewoniewski, Wlodzimierz and Wecel, Krzysztof and Abramowicz, Witold (2017) Lewoniewski, Wlodzimierz and Wecel, Krzysztof and Abramowicz, Witold. 2017.  Relative quality and popularity evaluation of multilingual Wikipedia articles. In _Informatics_ , Vol. 4. Multidisciplinary Digital Publishing Institute, 43.
  * Lewoniewski, Wlodzimierz and Wecel, Krzysztof and Abramowicz, Witold (2019) Lewoniewski, Wlodzimierz and Wecel, Krzysztof and Abramowicz, Witold. 2019.  Multilingual ranking of Wikipedia articles with quality and popularity assessment in different topics.  _Computers_ 8, 3 (2019), 60.
  * Morgan, Jonathan (2019) Morgan, Jonathan. 2019\.  Research: Patrolling on Wikipedia.  _Report-Meta, Sep_ (2019).
  * Morrison, Sara (2020) Morrison, Sara. 2020\.  How Wikipedia is preparing for Election Day.  {https://www.vox.com/recode/2020/11/2/21541880/wikipedia-presidential-election-misinformation-social-media}
  * Redi, Miriam and Fetahu, Besnik and Morgan, Jonathan and Taraborelli, Dario (2019) Redi, Miriam and Fetahu, Besnik and Morgan, Jonathan and Taraborelli, Dario. 2019.  Citation needed: A taxonomy and algorithmic assessment of Wikipedia’s verifiability. In _The World Wide Web Conference_. 1567–1578.
  * Redi, Miriam and Gerlach, Martin and Johnson, Isaac and Morgan, Jonathan and Zia, Leila (2021) Redi, Miriam and Gerlach, Martin and Johnson, Isaac and Morgan, Jonathan and Zia, Leila. 2021\.  A Taxonomy of Knowledge Gaps for Wikimedia Projects (Second Draft).  _arXiv preprint arXiv:2008.12314_ (2021).
  * Rogers, Richard and Sendijarevic, Emina and others (2012) Rogers, Richard and Sendijarevic, Emina and others. 2012.  Neutral or National Point of View? A Comparison of Srebrenica articles across Wikipedia’s language versions. In _unpublished conference paper, Wikipedia Academy, Berlin, Germany_ , Vol. 29.
  * Saez-Trumper, Diego (2019) Saez-Trumper, Diego. 2019\.  Online disinformation and the role of wikipedia.  _arXiv preprint arXiv:1910.12596_ (2019).
  * Sato, Yumiko (2021) Sato, Yumiko. 2021\.  Non-English Editions of Wikipedia Have a Misinformation Problem.  {https://slate.com/technology/2021/03/japanese-wikipedia-misinformation-non-english-editions.html}
  * Schenkkan, Nate and Repucci, Sarah (2019) Schenkkan, Nate and Repucci, Sarah. 2019\.  The freedom house survey for 2018: democracy in retreat.  _Journal of Democracy_ 30, 2 (2019), 100–114.
  * Shubber, Kadhim (2021) Shubber, Kadhim. 2021\.  Russia caught editing Wikipedia entry about MH17.  {https://www.wired.co.uk/article/russia-edits-mh17-wikipedia-article}
  * Song, Victoria (2020) Song, Victoria. 2020\.  A Teen Threw Scots Wiki Into Chaos and It Highlights a Massive Problem With Wikipedia.  {https://www.gizmodo.com.au/2020/08/a-teen-threw-scots-wiki-into-chaos-and-it-highlights-a-massive-problem-with-wikipedia/}
  * Spezzano, Francesca and Suyehira, Kelsey and Gundala, Laxmi Amulya (2019) Spezzano, Francesca and Suyehira, Kelsey and Gundala, Laxmi Amulya. 2019.  Detecting pages to protect in Wikipedia across multiple languages.  _Social Network Analysis and Mining_ 9, 1 (2019), 1–16.
  * Woolley, Samuel C and Howard, Philip N (2018) Woolley, Samuel C and Howard, Philip N. 2018\.  _Computational propaganda: political parties, politicians, and political manipulation on social media_.  Oxford University Press.
  * Yasseri, Taha and Spoerri, Anselm and Graham, Mark and Kertész, János (2014) Yasseri, Taha and Spoerri, Anselm and Graham, Mark and Kertész, János. 2014\.  The most controversial topics in Wikipedia.  _Global Wikipedia: International and cross-cultural issues in online collaboration_ 25 (2014), 25–48.
  * Zhang, Jerry and Carpenter, Darrell and Ko, Myung (2013) Zhang, Jerry and Carpenter, Darrell and Ko, Myung. 2013.  Online astroturfing: A theoretical perspective.  (2013).
  * Zia, Leila and Johnson, Isaac and Mansurov, Bahodir and Morgan, Jonathan and Redi, Miriam and Saez-Trumper, Diego and Taraborelli, Dario (2019) Zia, Leila and Johnson, Isaac and Mansurov, Bahodir and Morgan, Jonathan and Redi, Miriam and Saez-Trumper, Diego and Taraborelli, Dario. 2019.  Knowledge Integrity.  {https://doi.org/10.6084/m9.figshare.7704626}


# Self-Supervised Music Source Separation Using Vector-Quantized Source
Category Estimates

Marco Pasini∗
Queen Mary University
London, UK
&Stefan Lattner∗
Sony Computer Science Laboratories
Paris, France
&György Fazekas
Queen Mary University
London, UK

###### Abstract

Music source separation is focused on extracting distinct sonic elements from
composite tracks. Historically, many methods have been grounded in supervised
learning, necessitating labeled data, which is occasionally constrained in its
diversity. More recent methods have delved into N-shot techniques that utilize
one or more audio samples to aid in the separation. However, a challenge with
some of these methods is the necessity for an audio query during inference,
making them less suited for genres with varied timbres and effects. This paper
offers a proof-of-concept for a self-supervised music source separation system
that eliminates the need for audio queries at inference time. In the training
phase, while it adopts a query-based approach, we introduce a modification by
substituting the continuous embedding of query audios with Vector Quantized
(VQ) representations. Trained end-to-end with up to N classes as determined by
the VQ’s codebook size, the model seeks to effectively categorise instrument
classes. During inference, the input is partitioned into N sources, with some
potentially left unutilized based on the mix’s instrument makeup. This
methodology suggests an alternative avenue for considering source separation
across diverse music genres. We provide examples and additional results online
111https://anonymous2732.github.io/vqss/.

**footnotetext: Equal contribution

##  1 Introduction

Source separation is a crucial domain in signal processing, aiming to extract
individual components from mixed signals. This challenge has been widely
addressed in the audio domain across various fields, including universal
source separation, (multi-)speaker separation, and music source separation
(MSS). MSS stands as a significant sub-domain, with the goal of isolating
individual instruments, vocals, or other sonic elements from a composite
track. Traditionally, MSS has been driven by supervised methods that
necessitate large amounts of labeled stems. However, dedicated, openly
available datasets are rare, and limited in the number of instrument classes.
As a result, the default benchmark in MSS only involves four classes: vocals,
drums, bass, and other.

In order to lift this limitation, methods for unsupervised, semi-supervised,
or self-supervised source separation have been investigated in recent years. A
particular trend in semi- and self-supervised methods involves zero- one- or
few-shot learning approaches. These models typically rely on one or multiple
audio samples from a specific instrument class to guide the separation
process. In practice, it is easier to find datasets of unlabelled or weakly
labeled musical sounds, contained in public collections like AudioSet [1],
Freesound [2], or proprietary datasets. Consequently, using audio queries
instead of instrument labels makes it possible to train self-supervised source
separation models on such datasets (e.g., by using different segments of a
stem as query and separation target).

However, the necessity of having the respective audio query available at
inference time is a limiting factor in N-shot source separation. Particularly
in contemporary popular or electronic music, the constituents can have
peculiar timbres, like different synthesizers, guitar amplifiers, unusual
samples, or sound effects. Therefore, ideally, a source separation system
should be able to separate different instrument classes at inference time
without requiring additional information.

This work introduces a proof-of-concept self-supervised music source separator
that does not require audio queries at inference time. The method is a self-
supervised query-based source separation system at train time, but we replace
the continuous embedding space of the query audios with VQ embeddings. The
model is trained end-to-end using a maximum of N classes, which is the
codebook size of the VQ layer, to learn a helpful clustering of the respective
instrument classes. During inference, the model divides the input mix into N
sources. Depending on the presence of the different instrument classes in the
mix, some outputs may remain empty.

##  2 Method

We propose a self-supervised MSS system that does not require audio queries
during inference. The model is trained in a query-based fashion using random
crops of stems as queries and targets. However, instead of using a continuous
embedding space for the queries, we introduce a vector quantization layer to
categorize the instrument classes into a discrete codebook. The overall
architecture consists of three main components: A Style Encoder $E$ that
embeds the audio query into a latent space, a Vector Quantizer $Q$ that
quantizes the continuous embedding into a discrete codebook, and a Generator
$G$ that generates the separated source conditioned on the quantized code.

Let $x$ be an unlabelled audio stem. We first randomly crop $x$ two times to
obtain a target excerpt $x_{targ}$ and a reference excerpt $x_{ref}$. The
target excerpt is linearly mixed with $k\sim\mathcal{U}(0,K)$ additional
random stems to obtain the mixture $x_{mix}$. The reference excerpt $x_{ref}$
is fed into the Style Encoder $E$ to produce a $d$-dimensional embedding
$z_{ref}=E(x_{ref})$. This continuous embedding is then quantized by a Vector
Quantization layer $Q$ to obtain a quantized embedding
$\hat{z}_{ref}=Q(z_{ref})$ corresponding to the closest entry in a codebook of
size $N$, representing the maximum number of different sources allowed by the
system. The Generator $G$ takes as input the mixture $x_{mix}$ and the
quantized embedding $\hat{z}_{ref}$, and is tasked with separating the target
source $x_{targ}$ from the mixture:

$\hat{x}_{targ}=G(x_{mix},\hat{z}_{ref}).$

We parameterize the Generator with a U-Net [3] model which takes magnitude
spectrograms as input and outputs magnitude and phase spectrograms which are
then combined into a waveform sample via the iSTFT operator. The model is
inspired by the Autoencoder proposed in [4] and is thus trained in a similar
fashion. Differently from [4], the model is trained fully end-to-end using the
following objectives: The Reconstruction loss uses the same reconstruction
objective $\mathcal{L}_{rec}$ proposed in [4], which consists of a L1 loss
between log-spectrograms and a multi-scale spectral loss [5, 6] between the
output $\hat{x}_{targ}$ and target $x_{targ}$. Details about the
reconstruction objective can be found in the original paper. The Adversarial
loss utilizes the hinge adversarial loss $\mathcal{L}_{adv}$ [7] to implicitly
model the phase spectrogram as:

$\displaystyle\mathcal{L}_{adv}=\phantom{+}$
$\displaystyle\mathbb{E}_{x_{targ}}[\min(0,-1+D(\log(|\text{STFT}(x_{targ})|)))]$
$\displaystyle+$
$\displaystyle\mathbb{E}_{\hat{x}_{targ}}[\min(0,-1-D(\log(|\text{STFT}(\hat{x}_{targ})|)))],$

where $D$ is a Discriminator model. Lastly, the Commitment loss is the VQ-
specific objective $\mathcal{L}_{vq}$ that ensures the encoder outputs are
close to the quantized codes from the codebook:
$\mathcal{L}_{vq}=||z_{ref}-\text{sg}(\hat{z}_{ref})||_{2}^{2}$, with sg
denoting the stop-gradient operator. The codebook loss is realized through
exponential moving average (EMA) updates, as suggested in [8, 9].

The overall training loss is a weighted sum of these objectives:

$\mathcal{L}=\mathcal{L}_{adv}+\lambda_{rec}\mathcal{L}_{rec}+\lambda_{vq}\mathcal{L}_{vq}.$

At inference time, the model takes as input the mixture $x_{mix}$ and each of
the final quantized embeddings and iteratively outputs $N$ estimated sources,
without requiring any audio query. Depending on the actual instruments present
in the mixture, some of the outputs may be silent.

This approach eliminates the need for audio queries during inference. The VQ-
based categorization acts as an instrument classifier that is learned in a
self-supervised manner through the source separation task. Furthermore, by
training fully end-to-end, the model learns a discrete embedding space
specifically tailored for the task of source separation.

##  3 Implementation Details

The Generator, Style Encoder and Critic are fully convolutional models using
standard residual blocks [10] with group normalization layers [11] with 8
groups. The Generator is based on a U-Net architecture where the conditioning
information is introduced via FiLM layers [12] at every block. The Style
Encoder takes mel-spectrograms with 128 frequency bins as input and produces a
single $512$-dimensional vector via a Global Average Pooling (GAP) layer. The
VQ layer is initialized with $N=16$ embeddings via $10$ iterations of k-means
clustering of the first input training batch. As proposed by [13], we use
cosine similarity instead of Euclidean distance and we factorize the codes in
a lower dimensional space with dimensionality of $8$ to improve codebook
usage. The codebook loss is implemented via EMA with a factor of $0.99$.
Regarding the STFT parameters used during training, we adopt the ones proposed
by [4]. We choose $\lambda_{rec}=2.5$ and $\lambda_{vq}=100$. During training
we use random crops of $1.5$ seconds with a sampling rate of $44.1$ kHz to
train the system. While we ensure that all $x_{ref}$ crops are not completely
silent, we allow $x_{targ}$ crops to be silent, since it allows the model to
produce a silent sample when the quantized embedding of a source that is not
present in the input mix is used. We train the model for 1 million iterations
with a batch size of $32$, using Adam [14] with $\beta_{1}=0.5$ and
$\beta_{2}=0.9$ as the optimizer.

Figure 1: Mel-spectrograms of an example input mix and separated outputs. For
each separation from 0 to 15 the Generator is conditioned on the corresponding
quantized embedding. By cross-referencing this visualization with the
histogram in Figure 2, it is possible to recognize specific instruments, such
as Vocals for cluster number 14.

Figure 2: Distribution of the different classes of the test set for each
quantized embedding found by the system. We notice a clear clustering of
certain sources (OH, Bass, Vocals), while for other classes (Toms, Snare,
Ride) the distribution is shared across different clusters.

##  4 Experiments and Results

| OH | Misc | Crash 1 | Vocals | Kick | Toms | Crash 2 | Bass | Snare | Gtr | Ride
---|---|---|---|---|---|---|---|---|---|---|---
L1 | 14.51 | 16.55 | 19.30 | 15.06 | 24.84 | 23.60 | 19.39 | 16.68 | 15.70 | 14.98 | 17.75
L1 rand | 23.64 | 22.47 | 22.92 | 22.25 | 23.85 | 31.14 | 23.97 | 22.61 | 23.18 | 23.40 | 21.65

Table 1: L1 distance between log-spectrograms of separated and true sources
for each class in the test set. For comparison, we also calculate the L1
distance when the model is conditioned on a random quantized embedding instead
of the target one. The difficulty of the model to correctly separate ’Kick’
from other drums results in high L1 for that class.

We train the system on an internal collection of 20,000 unlabelled stems. To
evaluate the separation performance, we use an internal collection of 1,000
labelled stems with the following classes: OH, Misc, Crash 1, Crash 2, Vocals,
Kick, Toms, Bass, Snare, Gtr, Ride. We first visualize in Figure 2 the
distribution of the different classes for each final quantized embedding, or
cluster, found by the system. It is worth noting how the data distribution of
training samples differs from the limited labelled collection used for
testing. For example, training stems may not contain clearly separated
percussive instruments, despite them being separated in different classes for
the testing samples. This can explain the sub-par clustering performance on
some specific classes used for testing. Considering the specific capability of
the system to perform separation of arbitrary learned sources, evaluation
using standard metrics in the field of source separation [15] can be
challenging, since different clusters can represent a single instrument class,
and some uncommon instrument classes can be excluded from the clustering
altogether. In Table 1 we evaluate the separation performance of the model in
terms of L1 distance between generated and target log-spectrograms for each
class in the test set. More specifically, different crops serve as the target
sample, which is mixed with 4 other sources to create the input mixture, and
as the input to the Style Encoder to produce the quantized embedding used as
conditioning. For comparison, we also calculate L1 distances when a random
quantized embedding is used as conditioning. We also visualize in Figure 2
mel-spectrograms of an example test input mix and of the corresponding
separations performed by the system, one for each of the codebook embeddings.
By cross-referencing Figure 2, we can notice specific instruments being
separated. However, we can also notice hallucination behavior that is
characteristic of the proposed system: considering that we explicitly train a
generative model to perform the task, the Generator may produce realistic
separated outputs that are not actually present in the input mix.

##  5 Related Work

Supervised neural network-based MSS has been tackled with multi-layer
perceptrons [16, 17, 18], CNNs and RNNs (and combinations thereof) [19, 20,
21, 22, 23], UNets [24, 25] and with Transformers [26]. In such supervised
approaches, a specific model is typically trained for each source type. Self-
supervised paradigms involve permutation-invariant training for universal
source separation [27], and an explicit instrument-clustering approach for MSS
[28]. Unsupervised approaches to source separation have been investigated
using mixture-invariant training [29] and unsupervised blind source separation
with Variational Autoencoders [30, 31]. Audio query-based (i.e., N-shot)
approaches exist in supervised fashions for one-shot [32] and few-shot [33]
learning. Such query-based approaches are suitable for self-supervised
training, where two segments of a stem are typically used as a separation
target and a query, leading to zero-shot [34] and one-shot [35] architectures.
Such methods typically use a single separation network with conditioning that
is derived from embeddings of an audio query. Such a single separation model
is also used in a meta-learning approach [36] and for universal source
separation with weakly labelled data [37], the latter uses conditioning from a
separate audio tagging model. The model proposed in this work is a one-hot
conditioned UNet like in [25] that is trained in a self-supervised one-shot
fashion like in [35]. In addition, we perform an unsupervised clustering to
inform the separation module like in [28]. However, in contrast to [28], we
perform this clustering in an end-to-end fashion using a VQ embedding space.

##  6 Conclusion

This work introduced a proof-of-concept for a self-supervised MSS system that
does not require audio queries during inference. The method adopts a fully
end-to-end query-based approach during training, using random crops of stems
as queries and targets. However, instead of a continuous embedding space, we
introduced a vector quantization layer to categorize unlabelled sources into a
discrete codebook. We demonstrate both the potential and current limitations
of the proposed approach. On one hand, the model shows a clear ability to
cluster certain common instrument classes fairly consistently. This indicates
the feasibility of learning meaningful representations in a completely self-
supervised manner using the source separation task as a proxy and without the
explicit use of contrastive learning techniques. On the other hand, we also
highlight some flaws that need to be addressed in future work. Firstly, the
model can exhibit a hallucination behavior, where it tends to generate
realistic outputs even for sources that are not present in the input mixture.
While this phenomenon is not uncommon in adversarial networks, it is
undesirable for source separation where the objective is to filter out
existing sources. Hallucinations could be alleviated by training on a dataset
with a more balanced class distribution or by increasing the weight of the
reconstruction loss. Secondly, the perceived audio quality of the separated
outputs is currently lacking, with noticeable artifacts. This suggests that
the model has not learned a robust generalized representation of the diverse
instrument classes.

Overall, this work should be considered a first step towards fully self-
supervised source separation models that do not rely on audio queries at test
time. The clustering-based separation indicates a promising direction, but
more research is required to improve the fidelity of the outputs and reduce
hallucination behavior. With further development, end-to-end self-supervised
source separation could provide flexibility to handle the wide range of sounds
that occur in real-world music.

## References

  * [1] Jort F. Gemmeke et al.  “Audio Set: An ontology and human-labeled dataset for audio events”  In _ICASSP_ IEEE, 2017, pp. 776–780
  * [2] Frederic Font et al.  “Freesound technical demo”  In _ACM Multimedia_ ACM, 2013, pp. 411–412
  * [3] Olaf Ronneberger et al.  “U-Net: Convolutional Networks for Biomedical Image Segmentation”  In _MICCAI_ , 2015
  * [4] Marco Pasini and Jan Schlüter  “Musika! Fast Infinite Waveform Music Generation”  In _ISMIR_ , 2022, pp. 543–550
  * [5] Jesse H. Engel et al.  “DDSP: Differentiable Digital Signal Processing”  In _ICLR_ OpenReview.net, 2020
  * [6] Antoine Caillon and Philippe Esling  “RAVE: A variational autoencoder for fast and high-quality neural audio synthesis”  In _arXiv preprint arXiv:2111.05011_ , 2021
  * [7] Jae Hyun Lim and Jong Chul Ye  “Geometric gan”  In _arXiv preprint arXiv:1705.02894_ , 2017
  * [8] Aäron Oord et al.  “Neural Discrete Representation Learning”  In _NeurIPS_ , 2017
  * [9] Ali Razavi et al.  “Generating Diverse High-Fidelity Images with VQ-VAE-2”  In _NeurIPS_ , 2019
  * [10] Kaiming He et al.  “Deep Residual Learning for Image Recognition”  In _CVPR_ , 2016
  * [11] Yuxin Wu and Kaiming He  “Group Normalization”  In _ECCV (13)_ 11217, Lecture Notes in Computer Science  Springer, 2018, pp. 3–19
  * [12] Ethan Perez et al.  “FiLM: Visual Reasoning with a General Conditioning Layer”  In _AAAI_ AAAI Press, 2018, pp. 3942–3951
  * [13] Jiahui Yu et al.  “Vector-quantized Image Modeling with Improved VQGAN”  In _ICLR_ OpenReview.net, 2022
  * [14] Diederik P. Kingma and Jimmy Ba  “Adam: A Method for Stochastic Optimization”  In _ICLR_ , 2015
  * [15] Jonathan Le Roux et al.  “SDR - Half-baked or Well Done?”  In _ICASSP_ IEEE, 2019, pp. 626–630
  * [16] Aditya Arie Nugraha et al.  “Multichannel music separation with deep neural networks”  In _EUSIPCO_ IEEE, 2016, pp. 1748–1752
  * [17] Stefan Uhlich et al.  “Deep neural network based instrument extraction from music”  In _ICASSP_ IEEE, 2015, pp. 2135–2139
  * [18] Naoya Takahashi and Yuki Mitsufuji  “Multi-Scale multi-band densenets for audio source separation”  In _WASPAA_ IEEE, 2017, pp. 21–25
  * [19] Naoya Takahashi and Yuki Mitsufuji  “D3Net: Densely connected multidilated DenseNet for music source separation”  In _CoRR_ abs/2010.01733, 2020
  * [20] Stefan Uhlich et al.  “Improving music source separation based on deep neural networks through data augmentation and network blending”  In _ICASSP_ IEEE, 2017, pp. 261–265
  * [21] Naoya Takahashi et al.  “MMDenseLSTM: An Efficient Combination of Convolutional and Recurrent Neural Networks for Audio Source Separation”  In _IWAENC_ IEEE, 2018, pp. 106–110
  * [22] Fabian-Robert Stöter et al.  “Open-Unmix - A Reference Implementation for Music Source Separation”  In _J. Open Source Softw._ 4.41, 2019, pp. 1667
  * [23] Yi Luo and Jianwei Yu  “Music Source Separation With Band-Split RNN”  In _IEEE ACM Trans. Audio Speech Lang. Process._ 31, 2023, pp. 1893–1901
  * [24] Daniel Stoller et al.  “Wave-U-Net: A Multi-Scale Neural Network for End-to-End Audio Source Separation”  In _ISMIR_ , 2018, pp. 334–340
  * [25] Woo-Sung Choi et al.  “Lasaft: Latent Source Attentive Frequency Transformation For Conditioned Source Separation”  In _ICASSP_ IEEE, 2021, pp. 171–175
  * [26] Wei Tsung Lu et al.  “Music Source Separation with Band-Split RoPE Transformer”  In _CoRR_ abs/2309.02612, 2023
  * [27] Dong Yu et al.  “Permutation invariant training of deep models for speaker-independent multi-talker speech separation”  In _ICASSP_ IEEE, 2017, pp. 241–245
  * [28] Ke Chen et al.  “Pac-HuBERT: Self-Supervised Music Source Separation Via Primitive Auditory Clustering And Hidden-Unit Bert”  In _ICASSP Workshops_ IEEE, 2023, pp. 1–5
  * [29] Scott Wisdom et al.  “Unsupervised Sound Separation Using Mixture Invariant Training”  In _NeurIPS_ , 2020
  * [30] Julian Neri et al.  “Unsupervised Blind Source Separation with Variational Auto-Encoders”  In _EUSIPCO_ IEEE, 2021, pp. 311–315
  * [31] Li Li et al.  “FastMVAE2: On Improving and Accelerating the Fast Variational Autoencoder-Based Source Separation Algorithm for Determined Mixtures”  In _IEEE ACM Trans. Audio Speech Lang. Process._ 31, 2023, pp. 96–110
  * [32] Jie Hwan Lee et al.  “Audio Query-based Music Source Separation”  In _ISMIR_ , 2019, pp. 878–885
  * [33] Yu Wang et al.  “Few-Shot Musical Source Separation”  In _ICASSP_ IEEE, 2022, pp. 121–125
  * [34] Liwei Lin et al.  “A unified model for zero-shot music source separation, transcription and synthesis”  In _ISMIR_ , 2021, pp. 381–388
  * [35] Beat Gfeller et al.  “One-Shot Conditional Audio Filtering of Arbitrary Sounds”  In _ICASSP_ IEEE, 2021, pp. 501–505
  * [36] David Samuel et al.  “Meta-Learning Extractors for Music Source Separation”  In _ICASSP_ IEEE, 2020, pp. 816–820
  * [37] Qiuqiang Kong et al.  “Universal Source Separation with Weakly Labelled Data”  In _CoRR_ abs/2305.07447, 2023