# Towards Improved Zero-shot Voice Conversion with Conditional DSVAE

###### Abstract

Disentangling content and speaking style information is essential for zero-
shot non-parallel voice conversion (VC). Our previous study investigated a
novel framework with disentangled sequential variational autoencoder (DSVAE)
as the backbone for information decomposition. We have demonstrated that
simultaneous disentangling content embedding and speaker embedding from one
utterance is feasible for zero-shot VC. In this study, we continue the
direction by raising one concern about the prior distribution of content
branch in the DSVAE baseline. We find the random initialized prior
distribution will force the content embedding to reduce the phonetic-structure
information during the learning process, which is not a desired property.
Here, we seek to achieve a better content embedding with more phonetic
information preserved. We propose conditional DSVAE, a new model that enables
content bias as a condition to the prior modeling and reshapes the content
embedding sampled from the posterior distribution. In our experiment on the
VCTK dataset, we demonstrate that content embeddings derived from the
conditional DSVAE overcome the randomness and achieve a much better phoneme
classification accuracy, a stabilized vocalization and a better zero-shot VC
performance compared with the competitive DSVAE baseline.

Index Terms: Voice Conversion, DSVAE, Representation Learning, Generative
Model, Zero-shot style transfer

##  1 Introduction

Voice Conversion (VC) is a technique that converts the non-linguistic
information of a given utterance to a target style (e.g., speaker identity,
emotion, accent or rhythm etc.), while preserving the linguistic content
information. VC has become a very active research topic in speech processing
with potential applications in privacy protection speaker de-identification,
audio editing or sing voice conversion/generation [1, 2, 3, 4].

Current VC systems embrace the technological advancements from statistical
modeling to deep learning and have made a major shift on how the pipeline
develops [1]. For example, the conventional VC approaches with parallel
training data utilize a conversion module to map source acoustic features to
target acoustic features, the source-target pair has to be aligned before the
mapping [5]. With the advent of sequence-to-sequence models even without the
alignment prerequisite, better VC performance is reported [6]. For VC with
non-parallel data, direct feature mapping method is difficult. Instead,
studies start to explicitly learn the speaking style and content
representations and train a neural network as a decoder to reconstruct the
acoustic feature, with the assumption that the decoder can also generalize
well when the content and speaker style is swapped during the conversion.
Among the approaches, phonetic posteriorgrams (PPGs) and pre-trained speaker
embeddings are widely used as the content and speaking style representations
[7, 8, 9, 10]. However, developing such system usually requires a big amount
of external data with rich transcriptions and speaker labels. The relatively
small-footprint AUTOVC and AdaIN-VC employ encoder-decoder frameworks for
zero-shot VC [11, 12]. The encoder decomposes the speaking style and the
content information into the latent embedding, and the decoder generates a
voice sample by combining both disentangled information. Nevertheless, these
models require supervisions such as positive pair of utterances (i.e., two
utterances come from the same speaker), and the systems still have to rely on
pre-trained speaker models. Progress has also been made with generative
adversarial networks (GAN) based VC systems [13, 14, 15, 16]. This categorical
of method usually assumes that the speaker of source-target VC pair is pre-
known, which limits the application of such models in the real world. At the
same time, bunch of regularization terms have to be applied in the training
process, which imposes generalization doubts to such systems for zero-shot
non-parallel VC scenarios.

Our previous study proposed a novel disentangled sequential variational
autoencoder (DSVAE) [17] as a backbone framework for zero-shot non-parallel
VC. We designed two branches in the encoder of DSVAE to hold the time-varying
and the time-invariant components, where balanced content and speaking style
information flow is achieved with the VAE training [18]. We demonstrated that
the vanilla VAE [18, 19] loss can be extended to force strong disentanglement
between speaker and content components, which is essential for the success of
challenging zero-shot non-parallel VC.

In this study, we continue the direction by further improving the disentangled
representation learning in the DSVAE framework. One major concern is raised
after we analyzed the content embedding learned from our DSVAE baseline [17].
We find that the random initialed prior distribution in the content branch of
the baseline DSVAE is not optimal to preserve the phonetic/content structure
information. The randomness of content embedding $z_{c}$ has a negative impact
to phoneme classification and VC. To cope with this issue, we propose
$conditional$ DSVAE (C-DSVAE), an improved framework that corrects the
randomness in the content prior distribution with content bias. Alternative
content biases extended from unsupervised learning, supervised learning and
self-supervised learning are explored in this portion of study. The VC
experiments on VCTK dateset demonstrate a clear stabilized vocalization and a
significantly improved performance with the new content embeddings. Phoneme
classification with $z_{c}$ also justifies the effectiveness of the proposed
model in an objective way.

Figure 1: The system diagram of conditional DSVAE.

##  2 The DSVAE Baseline

###  2.1 Related Work

DSVAE [19] was proposed as a sequential generative model that disentangles the
time-invariant information from the time-variant information in the latent
space. The original DSVAE model [19] and its early variants [20, 21] achieved
limited success in speech disentanglement. Recently, we extended the DSVAE by
balancing the information flow between speaker and content representations and
it achieved the state-of-the-art performance for zero-shot non-parallel VC
[17]. To be unified, we refer to DSVAE as the baseline we developed in [17],
although it is already very different from the previous systems [19, 20, 21].

###  2.2 Baseline Overview

The DSVAE baseline adopted here is shown in Fig. 1. Denote
$X,\hat{X},z_{s},z_{s},\theta$ as input melspectrogram, reconstructed
melspectrogram, speaker embedding, content embedding and model parameters,
respectively. The shared encoder $E_{share}$ takes $X$ as input and outputs a
latent representation, with the speaker encoder $E_{S}$ and the content
encoder $E_{C}$ modeling the posterior distribution $q_{\theta}(z_{s}|X)$ and
$q_{\theta}(z_{c}|X)$ subsequently. $z_{s}$ and $z_{c}$ are then sampled from
$q_{\theta}(z_{s}|X)$ and $q_{\theta}(z_{c}|X)$. In the next stage, the
decoder takes the concatenation of $z_{s}$ and $z_{c}$, and passes them into
decoder $D$ to reconstruct the melspectrogram $\hat{X}$, i.e.
$\hat{X}=D(z_{s},z_{c})$. The vocoder then converts $\hat{X}$ into waveform.
Both the prior distribution $p_{\theta}(z)$ and the posterior distribution
$q_{\theta}(z|X)$ are designed to follow the independence criterion, which is
similar to [19, 20, 21, 17]. Specifically, they can be factorized as Eq. (1)
and Eq. (2). Note that we use $q_{\theta}(z_{ct}|X)$ to model the content
posterior since the content encoder consists of BiLSTM modules, which is
slightly different from the streaming posterior $q_{\theta}(x_{ct}|X{<t})$
described in [19, 20], where they adopt unidirectional LSTM or RNN.

$p_{\theta}(z_{s},z_{c})=p(z_{s})p_{\theta}(z_{c})=p(z_{s})\prod_{t=1}^{T}p_{\theta}(z_{ct}|z_{c}{<t})$

(1)

$q_{\theta}(z_{s},z_{s}|X)=q_{\theta}(z_{s}|X)q_{\theta}(z_{c}|X)=q_{\theta}(z_{s}|X)\prod_{t=1}^{T}q_{\theta}(z_{ct}|X)$

(2)

###  2.3 Training and Inference

During training, the model takes fixed length of $X$ as input and optimizes
with three objectives: $\mathcal{L}_{REC}$, $\mathcal{L}_{KLD_{s}}$ and
$\mathcal{L}_{KLD_{c}}$, as Eq. (3) (4) (5). $\mathcal{L}_{REC}$ is the
reconstruction loss, which is implemented as the negative log likelihood.
$\mathcal{L}_{KLD_{s}}$ and $\mathcal{L}_{KLD_{c}}$ denote the KL divergence
for speaker and content respectively.

$\mathcal{L}_{REC}=\mathbb{E}_{p(X)}\mathbb{E}_{q_{\theta}(X|z_{s},z_{c})}[-log(q_{\theta}(X|z_{s},z_{c}))]$
(3)
$\mathcal{L}_{KLD_{s}}=\mathbb{E}_{p(X)}[KLD(q_{\theta}(z_{s}|X)||p(z_{s}))]$
(4)
$\mathcal{L}_{KLD_{c}}=\mathbb{E}_{p(X)}[KLD(q_{\theta}(z_{c}|X)||p_{\theta}(z_{c}))]$
(5)

Given $X_{1}$ as the source utterance and $X_{2}$ as the target utterance for
VC inference, the transferred sample is simply $D(z_{s2},z_{c1})$, where
$z_{s2}$ and $z_{c1}$ are sampled from $q_{\theta}(z_{s}|X_{2})$ and
$q_{\theta}(z_{c}|X_{1})$. We use a vocoder to convert the mel spec to the
waveform.

###  2.4 Implementation Details

Table. 1 provides detailed descriptions of each module of the DSVAE baseline.
For shared encoder and decoder, the instance normalization [22] is applied on
both time and frequency axis. For speaker encoder $E_{S}$, content encoder
$E_{C}$ and the content prior model $p_{z_{c}}$, two dense layers are used to
model the mean and standard deviation of the $q(z_{s}|X)$, $q(z_{ct}|X)$,
$p(z_{ct}|z_{c}{<t})$ respectively. For the prior models, $p(z_{s})$ is the
standard normal distribution and $p_{\theta}(z_{c})$ is modeled by an
autoregressive LSTM: at each time step $t$, the model generates
$p(z_{ct}|z_{c}{<t})$, from which $z_{ct}$ is sampled and taken as the input
for next time step. Note that $p_{\theta}(z_{c})$ is independent of the input
data $X$. The decoder consists of a prenet and postnet, which is introduced in
[11]. We use HiFi-GAN V1 [23] instead of WaveNet [24] as vocoder since HiFi-
GAN results in better speech quality with much faster inference speed. The
vocoder is pretrained with VCTK [25] and is not involved in the training.

Shared Encoder $E_{share}$ (Conv1D(256, 5, 2, 1)$\rightarrow$
InstanceNorm2D$\rightarrow$ ReLU)$\times$3 Speaker Encoder $E_{S}$ Content
Encoder $E_{C}$ BiLSTM(512, 2)$\rightarrow$ Pooling BiLSTM(512,
2)$\rightarrow$RNN(512, 1) Dense(64) Dense(64) Decoder-PreNet $D_{Pre}$
(InstanceNorm2D$\rightarrow$ Conv1D(512, 5, 2, 1)$\rightarrow$ ReLU)$\times$3
LSTM(512, 1) $\rightarrow$ LSTM(1024, 2) $\rightarrow$ Dense(80) Decoder-
PostNet $D_{Post}$ (Conv1D(512, 5, 2, 1)$\rightarrow$ tanh$\rightarrow$
InstanceNorm2D)$\times$4 Vocoder $D$: HiFiGAN-V1 Prior $p(z_{c})$ Prior
$p(z_{s})$ LSTM(256, 1)$\rightarrow$Dense(64) N(0,I)

Table 1: Detailed DSVAE architecture. For Conv1D, the configuration is (output
channels, kernel size, padding, stride). For LSTM/BiLTSM/RNN, the
configuration is (hidden dim, layers). For Dense layer, the configuration is
(output dim).

##  3 Conditional DSVAE

###  3.1 Conditional Prior Distribution

Ideal disentanglement requires $z_{s}$ to carry speaking style information and
$z_{c}$ to carry content information without losing the phonetic structure.
One problem for the vanilla DSVAEs [19, 20, 21, 17] is that the prior
distribution is randomly initialized, thus it does not impose any constraint
to regularize the posterior distribution. We argue that such randomness on the
content prior distribution $p_{\theta}(z_{c})$ impedes the content embedding
$z_{c}$ from learning the phonetic structure information. Since the phonetic
structure is explicitly modeled by $q_{\theta}(z_{c}|X)$, according to Eq. 5,
one of the objective is to minimize the KL divergence between
$q_{\theta}(z_{c}|X)$ and $p_{\theta}(z_{c})$. Thus, we expect that content
embedding will be significantly influenced by the prior $p_{\theta}(z_{c})$
during VAE training. In that sense, the learned phonetic structure
$q_{\theta}(z_{c}|X)$ for all utterances will also follow the prior
distribution, which does not reflect the real phonetic structure of the
utterance. Such phenomenon can be observed in Fig. 2(a) and Fig. 2(c) which
gives the t-SNE [26] visualization of $z_{c}$ comparing the learned content
embeddings from the pretrained DSVAE [17] and the raw melspectrogram of the
same utterances. It is observable that DSVAE representations are not
phonetically discrimative in comparison to melspectrogram and they actually
follow the random distribution. The aforementioned problem is detrimental to
disentanglement and will generate discontinuous speech with non-stable
vocalizations.

Our solution is that, instead of modeling $p_{\theta}(z_{c})$, we will model
the conditional content prior distribution $p_{\theta}(z_{c}|Y(X))$ such that
the prior distribution is meaningful in carrying the content information. We
call $Y(X)$ as the content bias. The expectation is that, by incorporating the
content bias into the prior distribution $p_{\theta}(z_{c})$, the posterior
distribution $q_{\theta}(z_{c}|X)$ will retain the phonetic structure of $X$.

###  3.2 Proposed C-DSVAE

Based on the aforementioned discussion, we introduce four conditional DSVAE
candidates: C-DSVAE(Align), C-DSVAE(BEST-RQ), C-DSVAE(Mel) and C-DSVAE(WavLM)
based on different content bias source.

##### C-DSVAE(Align)

In order to let $z_{c}$ or $q_{\theta}(z_{c}|X)$ to keep the phonetic
structure of the speech data $X$, the content bias $Y(X)$ is expected to carry
the fine-grained phonetic information. One natural choice is to let $Y(X)$ be
the forced alignment of $X$. To do so, we employ the Kaldi toolkit [27] to
train a monophone model with 42 phonemes to obtain the forced alignment. The
training portion of the VCTK dataset is used in the HMM training (see Sec. 4.1
for dataset split). We denote this bias as $Y_{Align}$. As an example showing
in Fig. 1, the content bias $Y_{Align}$ for the current utterance is the
forced alignment labels [1 1 5 5 34 34 2 5 5 5]. In the next step, the one-hot
vectors are derived based on these labels for each frame, and are concatenated
with the original inputs of $p_{\theta}(z_{c})$ at each time step so that the
new content prior becomes $p_{\theta}(z_{c}|Y(X))$. Such conditioned content
prior is still factorized in a streaming manner, which is described as Eq. 6.

$p_{\theta}(z_{c}|Y(X))=\prod_{t=1}^{T}P_{\theta}(z_{ct}|z_{<t},Y(X_{t}))$ (6)

Note that $Y_{Align}$ is derived in the supervised manner, which has to reply
on the audio-transcription pairs. However, transcription is not always
available in practical usage. We present a few unsupervised labeling methods
as the content bias candidates. We note all these methods as Pseudo Labeling
(PL), which is also mentioned in [28]. The essence of PL is to derive closed-
set discrete acoustic units given continuous speech input.

##### C-DSVAE(BEST-RQ)

Given the continuous representations as input, VQ-VAE [29] will derive the
corresponding quantized vector as well as discrete indices by looking up in a
closed-set codebook. We adopt BEST-RQ [30] to extract pseudo labels.
Specifically, the melspectrogram is linearly projected into frame-wise
vectors, and then nearest-neighbour search is performed within a codebook to
derive pseudo labels. Both the projection matrix and codebook are randomly
initialized and then fixed during training. We denote this bias as $Y_{BEST-
RQ}$.

##### C-DSVAE(Mel)

BEST-RQ [30] is more like an online clustering algorithm that generates the
pseudo labels without seeing the entire dataset. In contrast, kmeans is an
offline method that embraces more global information. We directly perform
kmeans on the offline melspectrogram features on the whole training data.
After that, the index of cluster center is used as the pseudo label. This is
consistent with the first step of HuBERT [28]. We denote this bias as
$Y_{Mel}$.

##### C-DSVAE(WavLM)

The problem in C-DSVAE(Mel) is that melspectrogram is noisy and not
linguistically discriminative. To handle this problem, we attempt to apply
kmeans on the pre-trained features. Specifically, we use the pre-trained WavLM
features for kmeans clustering [31]. The advantage of WavLM is that the
aforementioned bias from melspectrograms will be alleviated via iterative
clustering and the masked prediction training process. The other point is that
WavLM acts as a teacher model so that the phonetic structure knowledge can be
transferred from a larger corpus, which potentially improves the robustness
and generalization capacity. We denote this bias as $Y_{WavLM}$. We use the
WavLM Base model, which is pretrained with 960 hours of Librispeech data [32].

Kmeans++ [33] is employed for implementing clustering. The number of cluster
is set as 50 for all experiments. We still keep speaker prior $p(z_{s})$ to be
a Gaussian prior $p(z_{s})$, which is actually a common assumption in speaker
recognition.

###  3.3 Training Objective

The content conditioned KL divergence loss is shown in Eq. 7. The overall loss
is shown in Eq, 8, where $\alpha$ and $\beta$ are the factors that balances
the disentanglement [17].

$\mathcal{L}_{KLD_{c-cond}}=\mathbb{E}_{p(X)}[KLD(q_{\theta}(z_{c}|X)||p_{\theta}(z_{c}|Y(X)))]$
(7)
$\mathcal{L}_{C-DSVAE}=\mathcal{L}_{REC}+\alpha\mathcal{L}_{KLD_{s}}+\beta\mathcal{L}_{KLD_{c}-cond}$
(8)

Following [17], we use the same training configuration for all experiments:
the ADAM optimizer is used with the initial learning rate of 5e-4 [34].
Learning rate is decayed every 5 epochs with a factor of 0.95. Weight decay is
1e-4, the batch-size is 256. Both speaker embedding and frame-wise content
embedding are 64-D. $\alpha=0.01$ and $\beta=10$ are kept the same as [17].

##  4 Experiments

###  4.1 Dataset

We use VCTK corpus for experimental study [25]. 90% of the speakers are used
for training and the remaining 10% are used for evaluation [17].
Melspectrogram is used as acoustic feature with the window size/hop size of
64ms/16ms, and the feature dimension is 80. We randomly select segments of 100
frames (1.6s) from the whole utterances for training.

###  4.2 Experimental Results

####  4.2.1 Content embedding and phoneme Classification

Fig.2 demonstrates the t-SNE [26] visualizations of the content embeddings
$z_{c}$ from 6 different content embeddings. The purpose of this portion of
study is to show how much the underlying distribution of $z_{c}$ matches the
(almost) ground truth phonetic structure. As shown in Fig. 2, content
embeddings from the DSVAE baseline follow a random uniform distribution. Such
distribution is detrimental to preserve the phonetic structure of raw speech.
At the same time, melspectrogram captures phone-dependent information due to
continuous speech signal nature. C-DSVAE(BEST-RQ) employs a random labelling
process for content biasing, thus it justifies that constraint is needed for
better performance. C-DSVAE(Mel), C-DSVAE(Align) and C-DSVAE(WavLM) deliver
much desired content distributions which successfully result in phonetically
discriminative embeddings. The phonetic structure of raw speech is retained
and better disentanglement is expected.

We also perform phoneme classification to evaluate content embeddings in an
objective way. The phoneme classifer is mentioned in Sec. 2.4. The consistent
conclusion could be drawn that DSVAE and C-DSVAE(BEST-RQ) give lower accuracy.
The reason for which C-DSVAE(Mel), C-DSVAE(Align) and C-DSVAE(WavLM)
outperform melspectrogram is that the latter contains the coarse-grained
phonetic structure which can be improved via offline clustering.
C-DSVAE(Align) is better than C-DSVAE(Mel) since alignment is obtained with a
supervised alignment model. C-DSVAE(WavLM) gives the best result because the
masked language modeling and iterative clustering tend to capture better
phonetic structure where the knowledge can also be transferred from the larger
corpus.

Setting | Phn ACC %
---|---
DSVAE | 30.2
C-DSVAE(BEST-RQ) | 35.6
Melspectrogram | 44.1
C-DSVAE(Mel) | 48.2
C-DSVAE(Align) | 51.1
C-DSVAE(WavLM) | 52.8

Table 2: Phoneme Classification with content embeddings.

(a) DSVAE

(b) C-DSVAE(BEST-RQ)

(c) Melspec

(d) C-DSVAE(Mel)

(e) C-DSVAE(Align)

(f) C-DSVAE(WavLM)

Figure 2: Visualizations of learned content embeddings.

####  4.2.2 Voice conversion111Samples of voice conversion can be found at
https://jlian2.github.io/Improved-Voice-Conversion-with-Conditional-DSVAE.

We also conduct a mean opinion score (MOS) test to evaluate our system. The
evaluation corpus setup is the same as [17]. The listener needs to give a
score for each sample in a test case according to the criterion: 1 = Bad; 2 =
Poor; 3 = Fair; 4 = Good; 5 = Excellent. The final score for each model is
calculated by averaging the collected results. Table 3 shows the MOS results
of different models.

As illustrated in the table, HiFi-GAN [23] based DSVAE outperforms the WaveNet
[24] with the same acoustic features. Except for C-DSVAE(Mel), our proposed
C-DSVAEs outperform the DSVAE baseline by a large margin in terms of
naturalness and similarity under both seen to seen and unseen to unseen
scenarios, and the MOS results are consistent with phoneme experiments as
introduced in Sec 4.2.1. The only exception is C-DSVAE(Mel), which achieves
worse naturalness than C-DSVAE(BEST-RQ) and worse similarity than DSVAE
baseline, the potential reason is that speaker embeddings learned in
C-DSVAE(Mel) are not as discriminative as those in either DSVAE baseline or
other C-DSVAEs.

| seen to seen | unseen to unseen
---|---|---
model | naturalness | similarity | naturalness | similarity
AUTOVC [17] |  2.65$\pm$0.12 |  2.86$\pm$0.09 |  2.47$\pm$0.10 |  2.76$\pm$0.08
AdaIN-VC [17] |  2.98$\pm$0.09 |  3.06$\pm$0.07 |  2.72$\pm$0.11 |  2.96$\pm$0.09
DSVAE [17] |  3.40$\pm$0.07 |  3.56$\pm$0.06 |  3.22$\pm$0.09 |  3.54$\pm$0.07
DSVAE(HiFi-GAN) |  3.76$\pm$0.07 |  3.83$\pm$0.06 |  3.65$\pm$0.07 |  3.89$\pm$0.05
C-DSVAE(BEST-RQ) |  3.88$\pm$0.06 |  3.93$\pm$0.07 |  3.82$\pm$0.08 |  3.98$\pm$0.07
C-DSVAE(Mel) |  3.86$\pm$0.10 |  3.65$\pm$0.07 |  3.78$\pm$0.05 |  3.58$\pm$0.08
C-DSVAE(Align) |  4.03$\pm$0.04 |  4.12$\pm$0.07 |  3.93$\pm$0.06 |  4.06$\pm$0.07
C-DSVAE(WavLM) |  4.08$\pm$0.06 |  4.17$\pm$0.06 |  3.98$\pm$0.07 |  4.12$\pm$0.05

Table 3:  The MOS (95% CI) test on different models.

####  4.2.3 Speaker verification

We consider speaker verification as an objective measure to evaluate the VC
performance. The speaker verification accuracy measures whether the
transferred voice belongs to the target speaker. For this purpose, we generate
200 source-target pairs and produce 400 target trials from the test set. A
state-of-the-art ECAPA-TDNN [35] based speaker verification system is used to
verify the speaker identity from the transferred voices. Please refer to [36]
for more details. We employ the cosine distance scoring method to perform
verification and use 0.42 as the threshold, which is determined by a third-
party test set [36]. As shown in Table 4, voice transferred from system
C-DSVAE(WavLM) achieved the best speaker verification accuracy. The trend is
similar to the phoneme classification and VC MOS test, which indicates that
stable content embeddings with more phonetic structure information boost the
VC performance in both subjective and objective evaluations.

Setting | ACC %
---|---
DSVAE | 85.0
C-DSVAE(BEST-RQ) | 86.3
C-DSVAE(Mel) | 83.8
C-DSVAE(Align) | 91.5
C-DSVAE(WavLM) | 92.3

Table 4: Test accuracy for transferred voice verification across different
models.

##  5 Conclusion

This paper proposes C-DSVAE, a novel voice conversion system that introduces
the content bias to the prior modeling to enforce the content embeddings to
retain the phonetic structure of the raw speech. The VC experiments on VCTK
dateset demonstrate a clear stabilized vocalization and a significantly
improved performance with the new content embeddings. With these contributions
and progress, our C-DSVAE achieves state-of-the-art voice conversion
performance.

## References

  * [1] B. Sisman, J. Yamagishi, S. King, and H. Li, “An overview of voice conversion and its challenges: From statistical modeling to deep learning,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , 2020.
  * [2] F. Bahmaninezhad, C. Zhang, and J. H. Hansen, “Convolutional neural network based speaker de-identification.” in _ISCA Odyssey_ , 2018.
  * [3] L. Zhang, C. Yu, H. Lu, C. Weng, C. Zhang, Y. Wu, X. Xie, Z. Li, and D. Yu, “Durian-sc: Duration informed attention network based singing voice conversion system,” _arXiv preprint arXiv:2008.03009_ , 2020.
  * [4] Y. Gao, J. Lian, B. Raj, and R. Singh, “Detection and evaluation of human and machine generated speech in spoofing attacks on automatic speaker verification systems,” 2020. [Online]. Available: https://arxiv.org/abs/2011.03689
  * [5] D. J. Berndt and J. Clifford, “Using dynamic time warping to find patterns in time series.” in _KDD workshop_ , vol. 10, no. 16. Seattle, WA, USA:, 1994, pp. 359–370.
  * [6] J.-X. Zhang, Z.-H. Ling, L.-J. Liu, Y. Jiang, and L.-R. Dai, “Sequence-to-sequence acoustic modeling for voice conversion,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 27, no. 3, pp. 631–644, 2019.
  * [7] L. Sun, K. Li, H. Wang, S. Kang, and H. Meng, “Phonetic posteriorgrams for many-to-one voice conversion without parallel data training,” in _2016 IEEE International Conference on Multimedia and Expo (ICME)_. IEEE, 2016, pp. 1–6.
  * [8] S. Liu, J. Zhong, L. Sun, X. Wu, X. Liu, and H. Meng, “Voice conversion across arbitrary speakers based on a single target-speaker utterance,” _Proc. Interspeech 2018_ , pp. 496–500, 2018.
  * [9] H. Guo, H. Lu, N. Hu, C. Zhang, S. Yang, L. Xie, D. Su, and D. Yu, “Phonetic posteriorgrams based many-to-many singing voice conversion via adversarial training,” _arXiv preprint arXiv:2012.01837_ , 2020.
  * [10] M. Zhang, Y. Zhou, L. Zhao, and H. Li, “Transfer learning from speech synthesis to voice conversion with non-parallel training data,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 29, pp. 1290–1302, 2021.
  * [11] K. Qian, Y. Zhang, S. Chang, X. Yang, and M. Hasegawa-Johnson, “Autovc: Zero-shot voice style transfer with only autoencoder loss,” in _International Conference on Machine Learning_. PMLR, 2019, pp. 5210–5219.
  * [12] J.-c. Chou, C.-c. Yeh, and H.-y. Lee, “One-shot voice conversion by separating speaker and content representations with instance normalization,” _arXiv preprint arXiv:1904.05742_ , 2019.
  * [13] H. Kameoka, T. Kaneko, K. Tanaka, and N. Hojo, “Stargan-vc: Non-parallel many-to-many voice conversion using star generative adversarial networks,” in _2018 IEEE Spoken Language Technology Workshop (SLT)_. IEEE, 2018, pp. 266–273.
  * [14] T. Kaneko and H. Kameoka, “Cyclegan-vc: Non-parallel voice conversion using cycle-consistent adversarial networks,” in _2018 26th European Signal Processing Conference (EUSIPCO)_. IEEE, 2018, pp. 2100–2104.
  * [15] Y. A. Li, A. Zare, and N. Mesgarani, “Starganv2-vc: A diverse, unsupervised, non-parallel framework for natural-sounding voice conversion,” in _Interspeech_ , 2021.
  * [16] T. Li, Y. Liu, C. Hu, and H. Zhao, “Cvc: Contrastive learning for non-parallel voice conversion,” _arXiv preprint arXiv:2011.00782_ , 2020.
  * [17] J. Lian, C. Zhang, and D. Yu, “Robust disentangled variational speech representation learning for zero-shot voice conversion,” in _IEEE ICASSP_. IEEE, 2022.
  * [18] D. P. Kingma and M. Welling, “Auto-encoding variational bayes,” 2014.
  * [19] Y. Li and S. Mandt, “Disentangled sequential autoencoder,” _arXiv preprint arXiv:1803.02991_ , 2018.
  * [20] Y. Zhu, M. R. Min, A. Kadav, and H. P. Graf, “S3vae: Self-supervised sequential vae for representation disentanglement and data generation,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2020, pp. 6538–6547.
  * [21] J. Bai, W. Wang, and C. P. Gomes, “Contrastively disentangled sequential variational autoencoder,” _Advances in Neural Information Processing Systems_ , vol. 34, 2021.
  * [22] D. Ulyanov, A. Vedaldi, and V. Lempitsky, “Instance normalization: The missing ingredient for fast stylization,” _arXiv preprint arXiv:1607.08022_ , 2016\.
  * [23] J. Kong, J. Kim, and J. Bae, “Hifi-gan: Generative adversarial networks for efficient and high fidelity speech synthesis,” _Advances in Neural Information Processing Systems_ , vol. 33, pp. 17 022–17 033, 2020.
  * [24] A. v. d. Oord, S. Dieleman, H. Zen, K. Simonyan, O. Vinyals, A. Graves, N. Kalchbrenner, A. Senior, and K. Kavukcuoglu, “Wavenet: A generative model for raw audio,” _arXiv preprint arXiv:1609.03499_ , 2016.
  * [25] C. Veaux, J. Yamagishi, and K. MacDonald, “Cstr vctk corpus: English multi-speaker corpus for cstr voice cloning toolkit,” 2017.
  * [26] L. Van der Maaten and G. Hinton, “Visualizing data using t-sne.” _Journal of machine learning research_ , vol. 9, no. 11, 2008.
  * [27] D. Povey, A. Ghoshal, G. Boulianne, L. Burget, O. Glembek, N. Goel, M. Hannemann, P. Motlicek, Y. Qian, P. Schwarz _et al._ , “The kaldi speech recognition toolkit,” in _IEEE 2011 workshop on automatic speech recognition and understanding_ , no. CONF. IEEE Signal Processing Society, 2011.
  * [28] W.-N. Hsu, B. Bolte, Y.-H. H. Tsai, K. Lakhotia, R. Salakhutdinov, and A. Mohamed, “Hubert: Self-supervised speech representation learning by masked prediction of hidden units,” _IEEE/ACM Transactions on Audio, Speech, and Language Processing_ , vol. 29, pp. 3451–3460, 2021.
  * [29] A. Van Den Oord, O. Vinyals _et al._ , “Neural discrete representation learning,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [30] C.-C. Chiu, J. Qin, Y. Zhang, J. Yu, and Y. Wu, “Self-supervised learning with random-projection quantizer for speech recognition,” _arXiv preprint arXiv:2202.01855_ , 2022.
  * [31] S. Chen, C. Wang, Z. Chen, Y. Wu, S. Liu, Z. Chen, J. Li, N. Kanda, T. Yoshioka, X. Xiao _et al._ , “Wavlm: Large-scale self-supervised pre-training for full stack speech processing,” _arXiv preprint arXiv:2110.13900_ , 2021.
  * [32] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur, “Librispeech: an asr corpus based on public domain audio books,” in _2015 IEEE international conference on acoustics, speech and signal processing (ICASSP)_. IEEE, 2015, pp. 5206–5210.
  * [33] D. Arthur and S. Vassilvitskii, “k-means++: The advantages of careful seeding,” Stanford, Tech. Rep., 2006.
  * [34] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” 2014. [Online]. Available: https://arxiv.org/abs/1412.6980
  * [35] B. Desplanques, J. Thienpondt, and K. Demuynck, “Ecapa-tdnn: Emphasized channel attention, propagation and aggregation in tdnn based speaker verification,” _arXiv preprint arXiv:2005.07143_ , 2020.
  * [36] C. Zhang, J. Shi, C. Weng, M. Yu, and D. Yu, “Towards end-to-end speaker diarization with generalized neural speaker clustering,” in _IEEE ICASSP_. IEEE, 2022.


# ConUNETR: A conditional transformer Network for 3D Micro-CT Embryonic
Cartilage Segmentation

###### Abstract

Studying the morphological development of cartilaginous and osseous structures
is critical to the early detection of life-threatening skeletal dysmorphology.
Embryonic cartilage undergoes rapid structural changes within hours,
introducing biological variations and morphological shifts that limit the
generalization of deep learning-based segmentation models that infer across
multiple embryonic age groups. Obtaining individual models for each age group
is expensive and less effective, while direct transfer (predicting an age
unseen during training) suffers a potential performance drop due to
morphological shifts. We propose a novel Transformer-based segmentation model
with improved biological priors that better distills morphologically diverse
information through conditional mechanisms. This enables a single model to
accurately predict cartilage across multiple age groups. Experiments on the
mice cartilage dataset show the superiority of our new model compared to other
competitive segmentation models. Additional studies on a separate mice
cartilage dataset with a distinct mutation show that our model generalizes
well and effectively captures age-based cartilage morphology patterns.

Index Terms—  Cartilage Segmentation, Conditional Model, Transformers, Cranial
Dysmorphology

##  1 Introduction

Congenital anomalies (birth defects) that include cartilaginous and osseous
(bone) structures are a major cause of infant mortality and childhood
morbidity, affecting 2-3% of human neonates [11]. Accurate segmentation of
cartilaginous structures (e.g., the chondrocranium and Meckel’s cartilage [14,
12]) is central to understanding these diseases and revealing novel targets
for treatment. Embryonic mice are often sacrificed at different prenatal
periods due to the invasive nature of these studies, to study cartilage
development through contrast-enhanced micro-CT.

During embryonic growth, large and rapid changes in cartilage size and
morphology occur, making the automatic segmentation of cartilage challenging.
More specifically, many morphological variations such as discontinuous or
absent cartilage structures, poorly identifiable boundaries, and different
cartilage thickening or thinning are observed during growth. Moreover,
studying multi-aged mice with different mutations (the usual practice in this
line of work) adds more challenges because of the various morphological
differences among even same-age cartilages across mutations. Due to the
difficulty of encapsulating cartilage diversity across age groups and
mutations, separate models are often trained for each age group and each
mutation. However, this approach requires intensive annotation labor, incurs
higher training costs/efforts, and risks overfitting to a single age
group/mutation (see Figure 2) without leveraging inter-age/mutation structure
information.

Alternatively, training a single model on the combined set of all age groups
offers the benefits of an expanded training set and better generalizability.
This idea has been explored with proven results in biomedical image analysis
with “universal models” [5, 17] that attempt to train a single model to
perform well across modalities, biological structures, or image acquisition
settings. However, this approach remains unexplored in cross-sectional studies
like our setting where the goal is to model structural changes over time
rather than over different modalities and structures of interest. Further,
enabling our model to generalize to unseen age groups and other mutation
cohorts is an additional challenge that would be valuable for craniofacial
researchers.

Fig. 1: An overview of the architecture of our proposed Conditional Universal
Model (ConUNETR).

Related to embryonic cartilage analysis, knee or hip cartilage segmentation
are well-studied problems [7, 9, 1, 6, 13, 8, 18]. While these models have
successfully addressed multi-domain adaptation and proposed novel ways to add
cartilage shape priors to facilitate segmentation even in very small data
regimes, the structure of interests are well-developed in their cases and do
not contain longitudinal samples with drastic morphology variations like in
our developing embryonic cartilages cases. Embryonic mouse cartilage
segmentation has seen fewer studies, and known methods either use iterative
active learning with experts in the loop [12] or account for morphological
variations only in the skull regions of the same age mice [19]. To this end,
our work is the first to address morphology variations across multi-age and
multi-mutation embryonic cartilage data with a single “universal” model.

In this paper, we propose a new cartilage segmentation model, ConUNETR, to
enable accurate prediction across age groups and improve generalizability to
unseen ages and mutations. We leverage attention in Transformers [16] to
effectively capture global image context while informing feature extraction
components of relevant biological priors. We propose to add these priors as
conditional inputs to the Transformer encoder (i.e., ViT [2]) in the form of
tokens so that the network can extract age-relevant features during the
forward pass and model cartilage structure variations across time for improved
generalizability. Besides, we propose spatial embedding components to tackle
the lack of spatial context from non-consecutive 2D labeled slices of micro-CT
volumes.

Our main contributions are: (1) a novel conditional modeling approach using
biological priors in Transformer encoders to address morphology variations in
the multi-age embryonic cartilage dataset; (2) a lightweight 2D Transformer
segmentation network with spatial awareness for non-consecutive 2D labeled
slices of micro-CT volumes; (3) superior generalizability across age groups
and across mutations compared to popular segmentation models.

Table 1:  Individual Age-group Training and Direct Transfer: In Tables 1 and
2, the numbers represent the per volume dice scores averaged within an age
group (average across 3 runs). Grey columns represent the cases with the model
trained on age-specific training data, and white columns represent the ones
with the model transferred directly with no training data.

215mmModel | Trained on MUTA-13.5 | Trained on MUTA-14.5 | Trained on MUTA-15.5
---|---|---|---
| MUTA-13.5 | MUTA-14.5 | MUTA-15.5 |  MUTA-13.5 |  MUTA-14.5 | MUTA-15.5 | MUTA-13.5 |  MUTA-14.5 | MUTA-15.5
U-Net |  71.1 $\pm$ 1.3 | 63.1  $\pm$ 1.4 | 39.7  $\pm$ 2.2 |  47.9  $\pm$ 4.0 |  89.6  $\pm$ 0.7 | 71.7  $\pm$ 4.9 | 32.5  $\pm$ 2.1 |  79.0  $\pm$ 1.3 | 83.0  $\pm$ 0.8
Res2Unet | 69.5  $\pm$ 0.9 | 62.0  $\pm$ 1.8 | 26.2  $\pm$ 4.0 |  46.3  $\pm$ 2.6 |  89.2  $\pm$ 0.9 | 68.3  $\pm$ 1.4 | 33.5  $\pm$ 2.3 |  81.3  $\pm$ 0.7 |  83.7 $\pm$ 0.9
UNETR | 68.2  $\pm$ 0.2 | 61.0  $\pm$ 1.9 | 36.1  $\pm$ 2.3 |  36.1  $\pm$ 5.2 |  90.1 $\pm$ 0.5 | 63.7  $\pm$ 1.7 | 31.3  $\pm$ 2.6 |  78.1  $\pm$ 0.3 | 82.2  $\pm$ 0.5

Table 2:  Joint Age-group Training and Cross Mutation Generability

215mmModel | Trained on MUTA-13.5 & 14.5 | Trained on MUTA-14.5 & 15.5 | Trained on All MUTA | Trained on All MUTA
---|---|---|---|---
| MUTA-13.5 | MUTA-14.5 | MUTA-15.5 |  MUTA-13.5 |  MUTA-14.5 | MUTA-15.5 | MUTA-13.5 |  MUTA-14.5 | MUTA-15.5 | ACH-14.5 | ACH-16.5
U-Net | 70.7  $\pm$ 0.5 | 89.4  $\pm$ 0.9 |  70.5 $\pm$ 1.9 |  46.8  $\pm$ 2.0 |  89.4  $\pm$ 0.6 | 84.7  $\pm$ 0.8 | 71.0  $\pm$ 1.3 |  89.8  $\pm$ 0.4 | 85.0  $\pm$ 1.2 | 67.1  $\pm$ 2.3 | 65.6  $\pm$ 2.5
Res2Unet | 69.9  $\pm$ 1.3 | 89.5  $\pm$ 0.8 | 67.8  $\pm$ 1.2 |  45.2  $\pm$ 0.7 |  89.7  $\pm$ 0.6 |  85.0 $\pm$ 0.3 | 72.0  $\pm$ 0.4 |  90.3  $\pm$ 0.8 | 84.1  $\pm$ 0.3 | 79.4  $\pm$ 2.2 | 78.8  $\pm$ 1.0
UNETR | 69.9  $\pm$ 0.6 | 90.1  $\pm$ 0.4 | 64.6  $\pm$ 1.3 |  47.2  $\pm$ 1.9 |  90.5  $\pm$ 0.8 | 84.0  $\pm$ 1.2 | 70.2  $\pm$ 0.2 |  90.6  $\pm$ 0.5 | 83.6  $\pm$ 0.6 | 78.1  $\pm$ 1.2 | 75.1  $\pm$ 3.3
ours |  72.5 $\pm$ 0.6 |  90.9 $\pm$ 0.3 | 70.0  $\pm$ 0.5 |  48.1 $\pm$ 2.3 |  91.4 $\pm$ 0.3 | 84.6  $\pm$ 1.0 |  73.3 $\pm$ 0.8 |  91.7 $\pm$ 0.3 |  85.4 $\pm$ 0.2 |  82.3 $\pm$ 1.1 |  84.0 $\pm$ 1.6

##  2 Methodology

We aim to build a universal cartilage segmentation model for multi-age
embryonic data using biological priors and conditioning the model to learn
age-specific morphology features to benefit from joint training on the multi-
age data. Our segmentation model ConUNETR is inspired by the widely used UNETR
[4], and through careful design changes, we achieve a lightweight network
(12.3M parameters vs. UNETR’s 92.6M) to adapt for small training data regimes.
Our model consists of Transformer encoders, CNN decoders, learnable age
tokens, and spatial encoding modules as shown in Figure 1. These components
cost a few additional parameters and are seamlessly integrated with the
Transformer encoder.

###  2.1 Architecture

Our ConUNETR model consists of stacked ViT-like [2] Transformer encoders with
residual connections and CNN decoders with U-Net-like [15] skip connections.

Encoder: The encoder includes 6 stacked Transformers (6 stages), with 4
attention heads each, and maintains hidden dimensions ($d_{model}$) of 256.
Given a 2D image of size $H\times W\times C$ from a volume, we divide it into
non-overlapping patches of size $P\times P$ to obtain $N=HW/P^{2}$ patches. We
flatten these patches to attain a 1D vector that is linearly projected to a
$K$-dimensional embedding space and obtain patch embeddings
($E^{1},\ldots,E^{N}$). Learnable patch position encoding ($pos_{i}$) for each
$i^{th}$ patch, and sinusoidal relative slice spatial encoding ($e_{sp}$)
(explained in Section 2.3) are added to the patch embeddings to obtain:

$e=[E^{1};E^{2};\ldots;E^{N}]+[pos_{1};pos_{2};\ldots;pos_{N}]+e_{sp}.$

Age token $E^{a}$ (explained in Section 2.2) in the embedding space is
prepended to the embedding vector $e$ to obtain:

$z_{0}=[E^{a};\widetilde{E^{1}};\widetilde{E^{2}};\ldots;\widetilde{E^{N}}],$

where $\widetilde{E^{i}}$ represents the $i^{th}$ patch embedding with
position and spatial encoding added to it. The first stage of the stacked
Transformer encoder is fed with $z_{0}$ and it outputs $z_{1}$ after multi-
headed self-attention is applied among features of embedded patch tokens and
age tokens. Outputs of stages 2, 3, 5, and 6 ($z_{2},z_{3},z_{5},z_{6}$) are
passed to the decoder via skip connections after removing the age token
feature ($z_{i}[0]$) and only image features are fed to the decoder. These
image features have the morphology priors instilled in them due to self-
attention with age token features.

Decoder: The bottleneck layer at the end of the encoder upsamples the reshaped
encoder output ($z_{6}$) by a factor of 2 using a deconvolution layer. Multi-
scale features $z_{2},z_{3},z_{5}$ are reshaped and converted from the
embedding to the input space using convolutions and merged with the respective
decoder blocks through skip connections. The rest of the decoder consists of 4
stages of upsampling blocks that receive feature maps from the encoder which
are concatenated with the upsampled features. Softmax is applied to the output
of the final stage of the decoder to obtain the segmentation maps.

###  2.2 Age Tokens

We introduce an additional learnable token (similar in implementation to a
class token in [2]) to represent the age of samples in the dataset. $k$ age
tokens are initialized to represent $k$ ages and the one corresponding to the
age of the input image is prepended to the list of embedded patch features. We
also try injecting the biological priors of the input image throughout the
network using learnable age embeddings. Different from the age tokens which
make use of biological priors in the encoder only, age embeddings allow us to
represent concept information throughout the network by encoding this into the
image patches directly (through pixel-wise summation). We have observed that
using biological priors with tokens works better than embedding.

###  2.3 Spatial Encoding

To overcome the loss of spatial contexts when operating on some 2D slices of a
volume (due to non-consecutive slice annotations), we embed spatial encoding
into linearly projected image patches. Note that, while patch position
encoding represents the relative location of a patch within a 2D slice,
Spatial Encoding represents the relative location of a 2D slice within a
volume. We map these 2D slice locations from 1 to 100, and for each location
($loc$), build sinusoid wavelengths, forming a geometric progression from
$2\pi$ to $10000*2\pi$. The corresponding spatial encoding ($e_{sp}$) is added
to each patch embedding. We also tried learnable spatial encoding but observed
that sinusoid spatial encoding works slightly better.

$\displaystyle\textit{Spatial Encoding}_{(loc,2i)}$
$\displaystyle=sin(loc/10000^{2i/d_{model}}),$ $\displaystyle\textit{Spatial
Encoding}_{(loc,2i+1)}$ $\displaystyle=cos(loc/10000^{2i/d_{model}}).$

##  3 Datasets and Implementation Details

Dataset: Mouse embryos were stained with phosphotungstic acid (PTA), and whole
body micro-CT volumes ($\sim$1600 2D slices per volume) for mice of embryonic
(E) ages E13.5, E14.5, and E15.5 days with craniofacial mutation A (MUTA) and
ages E14.5 and E16.5 for Fgfr3Y367C/+ mouse model with achondroplasia (ACH)
mutation [12] were acquired. The training set for our experiments consists of
9 volumes (3 volumes per MUTA age) and the training ground truth consists of
cartilage masks obtained through an iterative active learning framework [12]
that produces cartilage predictions as good as experts’ hand segmentations.
Unless mentioned otherwise, we use only 5% of image slices per volume (a
common annotation budget in practice for our embryonic cartilage segmentation)
during training. The test set includes the experts’ hand segmentations on 2D
slices and comprises 10 partially labeled ($\sim$5%) volumes (2 volumes for
each age in both mutations). Thus, our total training data consists of 720
MUTA and the test set consists of 480 MUTA and 320 ACH 2D slices. We have made
the ACH dataset
111www.datacommons.psu.edu/commonswizard/MetadataDisplay.aspx?Dataset=6367
publicly available in [12].

Implementation Details: All models are trained from scratch using the AdamW
optimizer (initial learning rate of 0.0001) with weight decay of 0.001 and
Cosine Annealing learning rate scheduler. Cross-entropy loss is applied to all
the models, trained on a single NVIDIA A10 GPU for 700 epochs (batch size of
45) using AMP [10]. 512x512 sized image crops are obtained with augmentations
including random cropping, a non-linear Bézier curve intensity transformation,
horizontal/vertical flips, and $90^{\circ}$ rotations.

Baselines:  We compare the performance of our segmentation model (ViT encoder
with spatial encoding and age tokens) with three popular models: (1) A small
U-Net [15] (7.9M parameters) with CNN encoder; (2) Res2UNet (25.4M) with a
Res2Net [3] encoder; (3) Small UNETR (12.3M) [4] with Transformer [2] encoder.
A U-Net-like decoder with skip connections is added to all these encoders.

##  4 Experiments, Studies, and Results

###  4.1 Study on Same Age Group Training

Table 2 shows that the model trained on one age group does not generalize well
to other age groups without additional training data, due to age-specific
structural differences of developing cartilages. One might associate this to
the small size of the training data but our experiments (see Figure 2) suggest
that with more training data, a model trained on an individual age overfits on
that age and does not generalize well across age groups. This emphasizes the
need for joint training.

Fig. 2: Performance with additional same-age train. data (U-Net).

###  4.2 Study on Joint Training

The first two blocks in Table 2 show results with joint training on volumes
with closer age groups, and the third block shows joint training on all age
groups within the same mutation (MUTA). Except in some cases, all the models
exhibit performance gains with joint training. Cases with lower or subpar
performance improvement from individual to joint training are explained by the
difficulty in learning multi-age morphology variations. However, our model
generally outperforms the baselines in two age-group joint training and beats
all the baselines in three age-group training, indicating its stronger ability
to utilize cross-age-group information.

###  4.3 Study on Cross Mutation Generalizability

We introduce additional mutation variations using ACH volumes. Models trained
on all the three MUTA ages are applied without any ACH training data to ACH
volumes. As seen in the $4^{th}$ block of Table 2, U-Net yields poor
generalizability while Res2Unet and UNETR yield subpar performance. Our
conditional model, ConUNETR, exhibits significant improvement over these
baselines (2.9% on E14.5 age and 5.2% on E16.5 age volumes). Cartilages aged
E16.5 are well developed and easily distinguishable from background
structures, which explains the high performance on these volumes.
Nevertheless, these improvements show that our model generalizes well in
cross-mutation morphology variations.

Fig. 3:  Performance with additional joint cross-age train. data (all MUTA
ages). Average Dice score computed across all age groups.

###  4.4 Study on Performance with Increased Training Data

In Figure 3, we observe that the CNN model has a lower performance ceiling
while Transformer models make steady gains with additional joint age training
data. The performance differences between the Transformers and our conditional
Transformer model narrow as the training data size increases, suggesting that
large training data potentially negate the need for biological priors.
However, in practice, since the cartilage segmentation task is usually in a
small data regime, our proposed model finds a meaningful spot in this domain.

###  4.5 Ablation Study

Adding spatial embeddings to the ViT encoder yields 3.1% on ACH and 0.7%
performance gain on MUTA (Table 3). The two conditional components, age
embeddings and age tokens, contribute to an additional 2.7% and 3.5%
respectively on multi-age ACH, and 0.5% and 1.3% respectively on multi-age
MUTA volumes. These components add negligible differences in the total
parameter count and training time.

Table 3:  Ablation Study: Average Dice scores across MUTA and ACH age groups
obtained by the model trained on all MUTA ages.

2*  ViT Encoder |  2*  Spatial Encoding. |  2*  Age Emb. |  2*  Age Token |
---|---|---|---|---
|  |  |  |  MUTA |  ACH
✓ |  |  |  |  81.5 |  76.6
✓ | ✓ |  |  |  82.2 |  79.7
✓ | ✓ | ✓ |  |  82.7 |  82.4
✓ | ✓ |  | ✓ |  83.5 |  83.2

##  5 Conclusions

We studied the potential to use a universal model in segmenting multi-
age/mutation embryonic mouse cartilages through conditional mechanisms using
morphology and spatial priors. Our proposed lightweight segmentation model
uses components that are seamlessly integrated with the Transformer encoder
with very few additional parameters. Through multiple studies, we showed that
our method demonstrates superior generalizability on new ages/mutations (even
in lesser data regimes) compared to the popular segmentation models. We plan
to extend this work further in facilitating accurate cartilage segmentations
in weakly supervised settings and possibly on other datasets with similar
challenges.

##  6 COMPLIANCE WITH ETHICAL STANDARDS

Mice were produced, processed, and sacrificed in compliance with animal
welfare guidelines approved by the Icahn School of Medicine at Mount Sinai and
the Pennsylvania State University to create the dataset

##  7 Acknowledgments

This work was supported in part by NIH/NICDR grants R01 DE027677 to JTR and
DZC, R01 DE031439 to SMMP and EWJ, and R01 DE029832 to EWJ and SMMP.

## References

  * [1] F Ambellan, A Tack, M Ehlke, and S Zachow.  Automated segmentation of knee bone and cartilage combining statistical shape knowledge and convolutional neural networks: Data from the Osteoarthritis Initiative.  MIA, 52:109–118, 2019.
  * [2] A Dosovitskiy, L Beyer, A Kolesnikov, X Weissenborn, Dand Zhai, T Unterthiner, M Dehghani, M Minderer, G Heigold, S Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  ICLR, 2021.
  * [3] S-H Gao, M-M Cheng, K Zhao, X-Y Zhang, M-H Yang, and P Torr.  Res2Net: A new multi-scale backbone architecture.  IEEE TPAMI, 43(2):652–662, 2019.
  * [4] A Hatamizadeh, Y Tang, V Nath, D Yang, A Myronenko, B Landman, H R Roth, and D Xu.  UNETR: Transformers for 3D medical image segmentation.  In WACV, pages 574–584, 2022.
  * [5] C Huang, H Han, Q Yao, S Zhu, and S K Zhou.  3D U 2-Net: A 3D universal U-Net for multi-domain medical image segmentation.  In MICCAI, pages 291–299, 2019.
  * [6] S Li, Xand Lv, M Li, J Zhang, Y Jiang, Y Qin, H Luo, and S Yin.  SDMT: Spatial dependence multi-task transformer network for 3D knee MRI segmentation and landmark localization.  IEEE TMI, 2023.
  * [7] D Liang, J Liu, Gongning Wang, Kand Luo, W Wang, and S Li.  Position-prior clustering-based self-attention module for knee cartilage segmentation.  In MICCAI, pages 193–202, 2022.
  * [8] Y Lin, H Yao, Z Li, G Zheng, and X Li.  Calibrating label distribution for class-imbalanced barely-supervised knee segmentation.  In MICCAI, pages 109–118, 2022.
  * [9] Q Liu, Z Xu, Y Jiao, and M Niethammer.  iSegFormer: Interactive segmentation via Transformers with application to 3D knee MR images.  In MICCAI, pages 464–474, 2022.
  * [10] P Micikevicius, S Narang, J Alben, G Diamos, E Elsen, D Garcia, B Ginsburg, M Houston, O Kuchaiev, G Venkatesh, et al.  Mixed precision training.  ICLR, 2018.
  * [11] P A Mossey, E E Catilla, et al.  Global registry and database on craniofacial anomalies: Report of a WHO Registry Meeting on Craniofacial Anomalies.  WHO, 2003.
  * [12] SM Motch Perrine, N Sapkota, K Kawasaki, Y Zhang, DZ Chen, M Kawasaki, EL Durham, Y Heuzé, L Legeai-Mallet, and J T Richtsmeier.  Embryonic cranial cartilage defects in the Fgfr3 Y367C/+ mouse model of achondroplasia.  The Anatomical Record, 2023.
  * [13] Y Peng, H Zheng, P Liang, L Zhang, F Zaman, X Wu, M Sonka, and D Z Chen.  KCB-Net: A 3D knee cartilage and bone segmentation network via sparse annotation.  MIA, 82:102574, 2022.
  * [14] M K Pitirri, E L Durham, NA Romano, J I Santos, A P Coupe, H Zheng, D Z Chen, K Kawasaki, E W Jabs, J T Richtsmeier, et al.  Meckel’s cartilage in mandibular development and dysmorphogenesis.  Frontiers in Genetics, 13:871927, 2022.
  * [15] O Ronneberger, P Fischer, and T Brox.  U-Net: Convolutional networks for biomedical image segmentation.  In MICCAI, pages 234–241, 2015.
  * [16] A Vaswani, N Shazeer, Niki Parmar, L Uszkoreit, Jand Jones, A N Gomez, Ł Kaiser, and I Polosukhin.  Attention is all you need.  NIPS, 30, 2017.
  * [17] K Yan, J Cai, Y Zheng, A P Harrison, D Jin, Y Tang, Y Tang, L Huang, J Xiao, and L Lu.  Learning from multiple datasets with heterogeneous and partial labels for universal lesion detection in CT.  IEEE TMI, pages 2759–2770, 2020.
  * [18] G Zeng, F Schmaranzer, TD Lerch, A Boschung, G Zheng, J Burger, K Gerber, M Tannast, K Siebenrock, Y-J Kim, et al.  Entropy guided unsupervised domain adaptation for cross-center hip cartilage segmentation from MRI.  In MICCAI, pages 447–456, 2020.
  * [19] H Zheng, S M Motch Perrine, M K Pitirri, K Kawasaki, C Wang, J T Richtsmeier, and D Z Chen.  Cartilage segmentation in high-resolution 3D micro-CT images via uncertainty-guided self-training with very sparse annotation.  In MICCAI, pages 802–812, 2020.