SSHNN: Semi-Supervised Hybrid NAS Network for Echocardiographic Image Segmentation

To reduce memory cost, partial channel connections [18] are adopted, where $1/n$ portion of dimensional features are sent to the cell and the rest features remain unchanged. Moreover, continuous relaxation [19] is reused as differentiable search space is pursued and then the stochastic gradient descent (SGD) can be applied. Thus, the output tensor of block:

$\begin{aligned} H_{i}^{l}&=\sum_{H_{j}^{l}\in\mathcal{I}}\sum_{O_{k}^{l}\in\mathcal{O}}\frac{exp\{\alpha_{j\rightarrow i}^{k}\}}{\sum_{m=1}^{|\mathcal{O}|}exp\{\alpha_{j\rightarrow i}^{m}\}}\cdot O_{k}^{l}(Part_{j\rightarrow i}\circ H_{j}^{l})\\ &+(1-Part_{j\rightarrow i})\circ H_{j}^{l}\end{aligned}$

(1)

where $Part_{j\rightarrow i}$ is sampling mask for channel selection, and $\alpha$ are normalized scalars, called architecture parameters, measuring the weight of candidate operation. To ensure the differentiable back propagation, softmax is implemented. Finally, the output tensor $H^{l}=Concat(\{H_{i}^{l}|i<=B\})$, where $Concat(\cdot)$ is concatenation and $B$ is the number of blocks. Cell level search is denoted as $H^{l}=Cell(H^{l-1},H^{l-2};\alpha)$.
Out Layer: Layer level search aims at finding the optimal network backbone within network search space for specific dataset, with the specific procedure is combine features from different resolutions for better feature extraction. In past researches [15, 20, 21], normal algorithm uses another architecture parameter $\beta$ for feature fusion by linear combination. However, the amount of $\beta$ is a not large, which is usually in the hundreds, but are responsible for complex feature fusion, especially after applying partial channel connection. By estimation in Sec. 3.4, linear combination cannot fully fuse multi-scale features since hundreds of channels represent different features, which needs more parameters to elaborate.

Thus, we replace normalized scalars with convolution kernels to realize higher degree of flexibility, shown in Fig. 1(d, $Left$). Furthermore, long and short skip connection is utilized due to an incredible ability in image segmentation [7, 21]. The whole representation of layer design:

$\begin{aligned} {}^{s}H^{l}=&\ Conv(Concat(Down(Cell(^{\frac{s}{2}}H^{l-1},^{s}H^{l-2};\alpha)),\\ &Up(Cell(^{2s}H^{l-1},^{s}H^{l-2};\alpha)),Cell(^{s}H^{l-1},^{s}H^{l-2};\alpha),\\ &\{^{s}H^{l^{\prime}}\in{{}^{s}H}|l^{\prime}<l\}))\end{aligned}$

(2)

where $Conv(\cdot)$ denotes the convolution, transforming fused features to the same channel counts of $s$-resolution, $Down(\cdot)$ denotes downsampling, and $Up(\cdot)$ denotes upsampling.

Rather than simply concatenating multi-scale features and then processing by convolution layers to revert features to the original image size after searching optimal network structure, where these layers act as decoder and HNAS acts as encoder essentially, we apply ViT to add global context and use U-shape decoder structure, shown in Fig. 1(a).

Define input image $\mathbf{x}\in\mathbb{R}^{H\times W\times C}$, where $H\times W$ denotes spatial resolution and $C$ denotes channel counts. First, perform tokenization [22]. Reshape $\mathbf{x}$ into $\{\mathbf{x}_{p}^{i}\in\mathbb{R}^{P\times P\times C}|i=1,\dots,N\}$ by $P\times P$ convolution (stride = $P$), where $N=\frac{HW}{P^{2}}$. Second, patch embedding and transformer. Map the patches $\mathbf{x}_{p}$ into a D-dimensional embedding space and add specific position embedding to keep position information, and then apply Transformers, including Multi-head Self-Attention (MSA) and Multi-layer Perceptron (MLP) blocks:

	$\displaystyle\small\mathbf{z}_{l}^{*}$	$\displaystyle=MSA(LN(\mathbf{z}_{l-1}))+\mathbf{z}_{l-1}$		(3)
	$\displaystyle\mathbf{z}_{l}$	$\displaystyle=MLP(LN(\mathbf{z}_{l}^{*}))+\mathbf{z}_{l}^{*}$		(4)

where $LN(\cdot)$ represents the layer normalization.

Thus, Transformer output is $\mathbf{z}_{l}\in\mathbb{R}^{\frac{HW}{P^{2}}\times D}$ and then we reshape it into $\mathbb{R}^{\frac{H}{P}\times\frac{W}{P}\times D}$ for decoder. Herein, we need to combine the global features extracted from Transformers with the local features extracted from NAS. To avoid losing low-level or high-level details, use U-shape decoder is necessary.

We first use a 2D convolution to adjust the channels of encoded features from Transformer to have the same number of channels as ${}^{s}H^{l}$ for feature aggregation. Since we do not expect that transformers bring excessive parameters, $\frac{H}{P}$ and $\frac{W}{P}$ are small and then should be upsampled to $s$-resolution spatial size. After reshaping and resamping, concatenate two parts of features, followed with a convolution layer to match the number of channels with the next resolution ${}^{\frac{s}{2}}H^{l}$ channel counts, and then attain the output of current resolution.

Similarly, “Upsample-Concatenation-Convolution” (shown in Fig. 1(c)) is implemented again for the next resolution features until combining $s=4$ resolution features. For segmentation purposes, finally a upsample layer and a convolution layer are used to recover the features to the full resolution and specific number of classes for predicting the dense output.

As medical image dataset always has a small amount of labeled images and a large amount of unlabeled data. Thus, we try to combine NAS with Mean Teacher method [23], to enhance the effectiveness and generalisation of the model.

For Mean Teacher, we have two models: $Student$ $f(\theta_{s})$ and $Teacher$ $f(\theta_{t})$, where $\theta_{s}$ and $\theta_{t}$ denote the network parameters of $Student$ and $Teacher$, respectively. $\theta_{t}$ are updated by the exponential moving average (EMA) of $\theta_{s}$:

$$\theta_{t,i}=\alpha\theta_{s,i-1}+(1-\alpha)\theta_{s,i}$$

(5)

where $\alpha$ is a hyperparameter controlling the updating speed, $i$ is iteration times. Given labeled dataset $\mathcal{D}_{l}$ and unlabeled dataset $\mathcal{D}_{u}$, define the consistency regularization to update $\theta_{s}$:

$\begin{aligned} \mathcal{L}_{c}=\frac{\sum_{\mathcal{D}_{l}}{l}_{MSE}(S(p_{l}^{t}),S(p_{l}^{s}))}{m}+\frac{\sum_{\mathcal{D}_{u}}{l}_{MSE}(S(p_{u}^{t}),S(p_{u}^{s}))}{n}\end{aligned}$

(6)

where $p_{l}^{t}$ and $p_{l}^{s}$ are the outputs of $Teacher$ and $Student$ on labeled images, $p_{u}^{t}$ and $p_{u}^{s}$ are the outputs on unlabeled images, $l_{MSE}$ is the Mean-Square Error, and $S(\cdot)$ is softmax function applied in channel dimension for scale controlling.

The supervised loss of $Student$ is computed by the output of $Teacher$ and the ground truth of labeled images:

$$\mathcal{L}_{s}=\frac{1}{m}\sum_{\mathcal{D}_{l}}l_{CE}(p_{l}^{t},y_{l})$$

(7)

where $l_{CE}$ is cross entropy loss function. Thus, the total loss used to train the $\theta_{s}$ is the sum of the consistency regularization and the supervised loss: $\mathcal{L}_{total}=\lambda_{0}\mathcal{L}_{s}+\lambda_{1}\mathcal{L}_{c}$, where $\lambda_{0}=1$ and $\lambda_{1}$ follows exponential ramp-up function in [24].

To construct a differentiable computation graph, we use continuous relaxation for architecture parameters $\alpha$ controlling the connection in inner cell and differential convolution layers for architecture parameters $\gamma$ controlling the connection in outer layer, which makes gradient descent possible. In the training, we split the labeled and unlabeled training dataset into two disjoint sets respectively: $\mathcal{D}_{l,A}$ and $\mathcal{D}_{l,B}$, $\mathcal{D}_{u,A}$ and $\mathcal{D}_{u,B}$. The optimization in each epoch can be summarized as:

1. 1.

   Update weight parameters $w$ of $f(\theta_{s})$ by
   $\nabla_{w}\mathcal{L}_{\mathcal{D}_{l,A},\mathcal{D}_{u,A}}(w,\alpha,\gamma)$ using $\mathcal{D}_{l,A}$ and $\mathcal{D}_{u,A}$

2. 2.

   Update weight parameters $w$ of $f({\theta_{t}})$ by EMA

3. 3.

   Update architecture parameters $\alpha$ and $\gamma$ of $f(\theta_{s})$ by $\nabla_{\alpha,\gamma}\mathcal{L}_{\mathcal{D}_{l,B},\mathcal{D}_{u,B}}(w,\alpha,\gamma)$ using $\mathcal{D}_{l,B}$ and $\mathcal{D}_{u,B}$

4. 4.

   Update architecture parameters $\alpha$ and $\gamma$ of $f(\theta_{t})$ by EMA

where $\mathcal{L}$ is the total loss on the segmentation mini-batch.

CAMUS [1] dataset is used to evaluate the performance of SSHNN, which is an open large-scale dataset in 2D echocardiography and were collected from 500 patients. Noticed that one four-chamber and one two-chamber view sequences were collected from each patient, lasting around 20 pictures long without labels, except for the moment of end diastole (ED) and end systole (ES). Therefore, accessible labelled dataset has 2000 echocardiographic images used for supervised loss calculation, and unlabeled dataset has around 19000 echocardiographic images. Segmentation labels have four types by manual annotations are the left ventricle endocardium, the myocardium, the left atrium and background. In this paper, we adopt the Intersection over Union (IoU), Dice Coefficient (Dice) and Parameters (Params) as the evaluation metrics.

We use images of size $256\times 256$ as the network input. Each cell has $B=5$ blocks. For the partial channel connections, $n=4$. Use filter multiplier $F$ and layers $L$ to control the complexity of network and $F$ is initialized as 8. Thus, in $s=4$, there are $B\times F\times\frac{s}{4}=40$ filters. When reducing spatial size from $s\rightarrow 2s$, the number of filters doubles. We use 4 transformers to extract global context. Experiments were carried out on a Nvidia RTX3090Ti GPU. For weights $w$ and architecture $\gamma$, applied optimizer is SGD with momentum 0.9, weight decay 0.0003 and initial value 0.01. For architecture $\alpha$, Adam [27] is applied with learning rate 0.003 and weight decay 0.001. The total number of epochs is 40 and architecture is optimized after 10 epochs since unstable update of $w$ makes local optima. $Teacher$ is updated synchronously with $Student$ for stable training. Finally, use 10% labeled images in validation and the rest in training with unlabeled images.

To demonstrate the effectiveness of SSHNN, we compare it with SOTA approaches. Table 1 and Fig .2 respectively elaborate the numerical and visual results for each network on the CAMUS dataset. Mean and standard deviation values for each metric are obtained from cross-validating on the 10 folds of the dataset. From SSHNN series, obviously larger $L$ brings higher model capacity, leading to better segmentation performance at the cost of more parameters (lower speed). Compared with Transfuse, which has suboptimal Dice in Table 1, SSHNN-L gets 0.98% Dice improvement and 1.05% IoU enhancement. Similarly compared, SSHNN-M also has the better Dice with fewer parameters, verifying the superiority.

Shown in Table 2, we conducted multiple sets of tests to evaluate the impact of convolution fusion, U-shape decoder & Transformer and the fraction of labeled images. First, SSHNN-L with convolution fusion has a 1.00% higher Dice than without convolution fusion when unlabeled image versus labeled image $\frac{N_{u}}{N_{l}}=10$. Moreover, use U-shaped decoder with transformer also increases the Dice by 1.33% compared with no changes at this fraction, verifying the effectiveness of them. Obviously, combine convolution and U-shaped decoder with Transformer works better as SSHNN-L with them remains 0.921 Dice when $\frac{N_{u}}{N_{l}}=10$, aligning with the objective of semi-supervision of learning with minimal supervision.

Note that when $\frac{N_{u}}{N_{l}}=2$, the model outperforms the others as unlabeled data provides extra information and does not mask the characteristics of labeled data. The results of SSHNN-S and SSHNN-M prove the robustness of our method again.

Last but not least, Table 3 discusses the effect of $F$, showing that higher model capacity promotes segmentation ability with expensive memory cost, which is a trade-off problem.