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. Code is available in GitHub.

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 [12,14]) 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.

Fig. 2.

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

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.

Related to embryonic cartilage analysis, knee or hip cartilage segmentation are well-studied problems [1, 6–9, 13, 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.

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.

Fig. 1.

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

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 $\left(d_{\mathit{\text{model}}}\right)$ 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,\dots ,E^N)$. Learnable patch position encoding $\left(po{s}_i\right)$ for each $i^{th}$ patch, and sinusoidal relative slice spatial encoding $\left(e_{sp}\right)$ (explained in Section 2.3) are added to the patch embeddings to obtain:

$$e=\left[E^1;E^2;\dots ;E^N\right]+\left[{\mathit{\text{pos}}}_1;{\mathit{\text{pos}}}_2;\dots ;{pos}_N\right]+e_{sp}.$$ (1)

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

$$z_0=\left[E^a;\widetilde{E^1};\widetilde{E^2};\dots ;\widetilde{E^N}\right]$$ (2)

where $\widetilde{E^i}$ represents the $i^{\mathit{\text{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 $\left(z_2,z_3,z_5,z_6\right)$ are passed to the decoder via skip connections after removing the age token feature $\left(z_i\right[0\left]\right)$ 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 $\left(z_6\right)$ 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 $\left(e_{sp}\right)$ is added to each patch embedding. We also tried learnable spatial encoding but observed that sinusoid spatial encoding works slightly better.

$$\begin{gathered} {\mathit{\text{Spatial}}\mathit{\text{Encoding}}}_{(\mathit{\text{loc}},2i)}=sin\left(\mathit{\text{loc}}/{10000}^{2i/d_{\mathit{\text{model}}}}\right), \\ {\mathit{\text{Spatial}}\mathit{\text{Encoding}}}_{(\mathit{\text{loc}},2i+1)}=cos\left(\mathit{\text{loc}}/{10000}^{2i/d_{\mathit{\text{model}}}}\right). \end{gathered}$$ (3)

3. DATASETS AND IMPLEMENTATION DETAILS

Dataset:

Mouse embryos were stained with phosphotungstic acid (PTA), and whole body micro-CT volumes (~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 ${\mathit{\text{Fgfr}}3}^{Y367C/+}$ 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 (~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¹ 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]. 512×512 sized image crops are obtained with augmentations including random cropping, a non-linear Bézier curve intensity transformation, horizontal/vertical flips, and 90° 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 1 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.

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.

Model	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 ±1.3	63.1 ±1.4	39.7 ±2.2	47.9 ±4.0	89.6 ±0.7	71.7 ±4.9	32.5 ±2.1	79.0 ±1.3	83.0 ±0.8
Res2Unet	69.5 ±0.9	62.0 ±1.8	26.2 ±4.0	46.3 ±2.6	89.2 ±0.9	68.3 ±1.4	33.5 ±2.3	81.3 ±0.7	83.7 ±0.9
UNETR	68.2 ±0.2	61.0 ±1.9	36.1 ±2.3	36.1 ±5.2	90.1 ±0.5	63.7 ±1.7	31.3 ±2.6	78.1 ±0.3	82.2 ±0.5

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.

Table 2.

Joint Age-group Training and Cross Mutation Generability

Model	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 ±0.5	89.4 ±0.9	70.5 ±1.9	46.8 ±2.0	89.4 ±0.6	84.7 ± 0.8	71.0 ±1.3	89.8 ±0.4	85.0 ±1.2	67.1 ±2.3	65.6 ±2.5
Res2Unet	69.9 ±1.3	89.5 ±0.8	67.8 ±1.2	45.2 ±0.7	89.7 ±0.6	85.0 ±0.3	72.0 ±0.4	90.3 ±0.8	84.1 ±0.3	79.4 ±2.2	78.8 ±1.0
UNETR	69.9 ±0.6	90.1 ±0.4	64.6 ±1.3	47.2 ±1.9	90.5 ±0.8	84.0 ±1.2	70.2 ±0.2	90.6 ±0.5	83.6 ±0.6	78.1 ±1.2	75.1 ±3.3
ours	72.5 ±0.6	90.9 ±0.3	70.0 ±0.5	48.1 ±2.3	91.4 ±0.3	84.6 ±1.0	73.3 ±0.8	91.7 ±0.3	85.4 ±0.2	82.3 ±1.1	84.0 ±1.6

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.

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.

Fig. 3.

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

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.

ViT Encoder	Spatial Encoding.	Age Emb.	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.