Voxel-wise Cross-Volume Representation Learning for 3D Neuron Reconstruction

Following SimSiam [5], the symmetrised similarity loss between an anchor-voxel $j$ and a pair-voxel $k$ is defined with cos operator as:

where $\texttt{cos[}u,v\texttt{]}=\frac{u}{\parallel u\parallel_{2}}\cdot\frac{v}{\parallel v\parallel_{2}}$ denotes the cosine similarity between two embeddings. The total similarity loss comparing $N$ voxel pairs is formulated as:

Together with loss $L^{CE}$, the encoded feature embedding for each voxel is expected to contain the intrinsic semantic cues shared among voxels of the same category. As the point-wise cross-entropy loss is complimentary to the representation learning loss [23], we keep the weight of these two losses the same. The segmentation loss is then formulated as $L^{seg}=L^{CE}+L^{SIM}$.

Given that latent codes of voxels being misclassified by the base segmentation model are more important for the representation learning [23], we design three different strategies to sample $N$ anchor-voxels from each point pool $\mathcal{P}_{*}$:

• •

  Random sampling ($\mathcal{AS}_{\mathbf{random}}$): Randomly sample $N$ anchor-voxels from the whole point pool $\mathcal{P}_{*}$;

• •

  Purely hard sampling ($\mathcal{AS}_{\mathbf{PH}}$): Randomly sample $N$ anchor-voxels from a subset of point pool $\mathcal{P}_{*}$. The subset is the collection of voxels whose prediction is wrong;

• •

  Hybrid combination sampling ($\mathcal{AS}_{\mathbf{hybrid}}$): $\frac{N}{2}$ anchor-voxels are randomly sampled from the whole point pool $\mathcal{P}_{*}$ while the rest $\frac{N}{2}$ anchor-voxels are randomly sampled from the subset stated in $\mathcal{AS}_{\mathbf{PH}}$.

For each selected anchor-voxel, a pair-voxel is sampled from the same point pool to feed together into the Siamese networks. Apart from the candidate points in point pool $\mathcal{P}_{*}$, we propose a virtual point dsc${}_{\mathcal{P}_{\mathcal{N}}}$ and dsc${}_{\mathcal{P}_{\mathcal{B}}}$ as point descriptor for point pool $\mathcal{P}_{\mathcal{N}}$ and $\mathcal{P}_{\mathcal{B}}$, respectively to represent an aggregation of semantic feature for each class. We design two ways of computing such a virtual point:

• •

  Relaxed: Average pooling of the entire point pool;

• •

  Strict: Average pooling of the subset of correctly classified points from the entire point pool.

To avoid outliers and stabilise the learning, inspired by Momentum²Teacher [11], we propose to use a momentum update mechanism to keep the semantic information of past pool descriptors. Formally, the pool descriptor for each pool is defined as:

where $k$ is the current iteration. $\alpha$ is the momentum coefficient and decreases from $\alpha_{base}=1$ to 0 with cosine scheduling policy [11] defined as $\alpha=\alpha_{base}\texttimes(\cos(\frac{\pi k}{K})+1)/2$, where $K$ is the total number of iterations.

Our studies on 3D neuron reconstruction were conducted on the publicly available 42-volume Janelia dataset developed for the BigNeuron project [14], which was further divided into 35, 3, and 4 samples as training, validation, and testing set, respectively. We applied random horizontal and vertical flipping, rotation, and cropping, as data augmentation techniques, to amplify our training set with 3D patches of size $128\times 128\times 64$. Given that the number of neuron voxels is dramatically smaller than that of the background voxels, to ensure the existence of foreground neuron voxels, we re-choose the patch until the foreground voxels make up over 0.1% of the whole patch.

All the models involved in the experiments were implemented in PyTorch 1.5 and trained from scratch for 222 epochs. A model was saved when it reached a better F1-score on the validation set. We use Adam as the optimizer with the learning rate of $1\times 10^{-3}$ and the weight decay of $1\times 10^{-4}$. The batch size $M$ is set as 4. We choose $\mathcal{AS}_{\mathbf{hybrid}}$ with the number of anchor-voxels $N$ as 512 and strict pool descriptor with MoUpdate mechanism. The hidden layer dimension is 512 and 128 for the projector and predictor, respectively. $d$ and $d_{p}$ are set as 128 and 512, respectively. To make fair comparison, all the 3D U-Net based segmentation models including Li2019 [9] and MKF-Net [19] were implemented with feature dimensions of 16, 32, 64, and 128 for respective layer from top to bottom. As for 3D LinkNet [2], in addition, we replaced the head and final blocks with convolutional layer without spatial changes to keep the same 3 times downsampling of feature maps.

The quantitative segmentation result is presented in Table 1. Our proposed method achieves the best F1 score among all the other state-of-the-art segmentation methods and improves the performance of the base U-Net by 2.32%. It is also noticeable that our model includes no additional cost during inference. Our proposed VCV-RL module can be applied on the encoder output of encoder-decoder architecture easily. When applied on the 3D LinkNet, VCV-RL can enhance the segmentation performance by almost 2%.

To validate whether the proposed segmentation method can facilitate the tracing algorithm, we choose the state-of-the-art tracing algorithm APP2 [25] as the main tracer. As in [10], we perform the tracing algorithm on adjusted input volume using the probability map predicted from each segmentation model. As presented in Table 2, without extra overhead during inference, our proposed method combined with the tracer achieves the best quantitative tracing results on all the metrics among all the other deep learning based reconstruction methods except for the precision. The reason why plain APP2 reaches such high recall is that it overtraces the neuron structures, which is likely to include more real neuron points. Fig. 2 displays the enhanced segmentation results after the image adjustment operation proposed in [10] in the first row. The second row presents the tracing results after applying APP2 [25] on the segmented images produced by different segmentation methods. Our proposed method combined with APP2 achieves competitive tracing result. We note that joint training of segmentation and tracing may further improve the 3D neuron reconstruction performance [18].

As presented in Table 4, model B, C, and D.1 reach better F1 score than Model A, which demonstrates the effect of representation learning in improving learning ability of the base U-Net model and the superiority of hybrid anchor sampling strategy. The reason why model D.2 and D.3 outperform D.1 is because the existence of the proposed pool descriptor $\texttt{dsc}_{\mathcal{P}}$ can help improve the generality of the latent space. The strict way of computing $\texttt{dsc}_{\mathcal{P}}$ can further enhance the segmentation performance. We also try to remove the momentum update mechanism of $\texttt{dsc}_{\mathcal{P}}$ and the experiment result of model E demonstrates the importance of past information storage. In addition, we conducted experiments on the proposed method with different number of anchor-voxels $N$. The result is presented in Fig. 4. When $N$ is 512, the F1-score is largest.