Appearance Prompt Vision Transformer for Connectome Reconstruction

1 Introduction

Neural connectivity reconstruction is a fundamental task to understand the function of biological reconstruction, which can widely promote basic scientific research including electrophysiology [Ascoli, 2002], cellular physiology [Donohue and Ascoli, 2011], genetics [Livet and Weissman, 2007], etc. 3D electron microscopy (EM) is the only available imaging instrument with the sufficient resolution to visualize and reconstruct dense neural morphology without ambiguity. However, at this resolution, even moderately small neural circuits yield numerous neuron numbers (e.g., typically hundreds of neuron instances in a single megapixel image [Meirovitch et al., 2019a]) that are prohibitively laborious for human manual annotation (e.g., normally the human labor required to reconstruct a 1003 µm3 volume is at more than 100,000 hours [Berning et al., 2015]). Recently, considerable works [Funke et al., 2018; Januszewski et al., 2018; Meirovitch et al., 2019b] have turned their attention to deep neural networks in the pursuit of automatic neural connectivity reconstruction. Since all neuron instances are of the same type (i.e., biological cells), with intricate morphology and densely intertwined branches, how to fully probe discriminative information to perform accurate neuron reconstruction is thus extremely challenging.

∗Corresponding author

To tackle the neural connectivity reconstruction problem, existing methods can be roughly categorized as object tracking based and boundary detection based paradigms. In the object tracking based paradigm [Januszewski et al., 2018; Meirovitch et al., 2016], 3D recurrent convolutional neural networks (CNNs) are trained to iteratively extend one neuron object at a time, but this process is time-consuming and inappropriate for large-scale applications. The boundary detection based paradigm [Funke et al., 2018; Lee et al., 2017] tends to adopt the CNN in pursuit of the relationship between voxel pairs for perceiving neuron boundaries, which are postprocessed to yield neuron segmentation in a gradual agglomeration manner (e.g., mutex watershed [Wolf et al., 2018]). Our work follows the boundary detection based paradigm credited to achieving competitive performance (e.g., top tier in SNEMI3D [Lee et al., 2017] neuron segmentation challenge) while sustaining an efficient pipeline for numerous neuron objects.

In the boundary detection based paradigm, two representative methods are affinity learning and metric learning conditioned on the network optimization strategy [Huang et al., 2022]. On one hand, affinity learning methods [Funke et al., 2018; Lee et al., 2017; Beier et al., 2017] transform the ground truth into affinity to constrain the network to learn the relationship between voxel pairs in an intuitive yet implicit optimization manner, seeking perceptibility to spatial position and discrimination to adjacent neuron instances with similar appearance (see Figure 1b). However, these methods directly output multi-channel maps as affinities, which tend to suffer from the absence of explicit voxel semantics, leading to confusion about long-range voxel correlation (affinity). On the other hand, the work [Lee et al., 2021] takes advantage of metric learning to pull voxels belonging to the same neuron instance together and push those of unrelated instances away on top of manual threshold based optimization, performing well in preserving instance semantic information well, which is crucial for further improvements in accuracy (see Figure 1a). However, hand-crafted threshold is fragile in flexibility and robustness for different datasets. Moreover, relying solely on the optimization of voxel embeddings inevitably compromises the spatial information, and this negative impact is inevitably amplified by inbuilt localized receptive fields of 3D CNNs. Overall, the above analysis indicates that affinity learning methods and metric learning methods are naturally complementary. The former preserves the spatial information well and possesses a more flexible yet implicit optimization, but suffers from the absence of explicit voxel semantics, while the latter is just the other way around. Besides, the inherent locality of 3D CNNs limits both formulations to modeling long-range dependencies and capturing global voxel context, leading to sub-optimal results. Therefore, it is more desirable to integrate these two formulations to exploit their complementary potential by learning long-range spatial dependencies to model spatial-aware voxel semantics in an explicit and flexible optimization strategy.

Motivated by the above discussions, we propose a coherent and unified Appearance Prompt Vision Transformer (APViT) to enable appearance knowledge conditioned on affinity to instruct voxels with explicit semantics from a global perspective based on metric learning (Figure 1c), including an extension continuity-aware attention module and an appearance prompt modulator. In the extension continuity-aware attention module, we construct hierarchical attention customized for neuron extensibility and slice continuity to learn instance voxel semantic context from a global perspective and utilize continuity priors to enhance voxel spatial awareness. (a) Neuron extensibility. Considering that interleaved different neuron instances contain intricate morphology, which tends to extend from one end of the input 3D volume to another. Therefore, to endow voxel semantics with long-range dependencies, we take advantage of the extension-aware attention mechanism to aggregate global context to each voxel position to obtain robust context-aware voxel embeddings that can adapt to extended neuronal morphology. (b) Slice continuity. Intuitively, the neural deformation across several contiguous slices is always smooth and continuous. Thus, we employ the continuity-aware attention mechanism for aggregating information in a corresponding 3D spatial neighborhood for each voxel location, aiming to empower the discrimination to neighboring neuron instances. In the appearance prompt modulator, we draw inspiration from the promptbased learning [Jia et al., 2022], which provides a general paradigm for specific knowledge learning from offline training, and leverage voxel-adaptive appearance knowledge conditioned on affinity rich in spatial information to instruct instance voxel semantics, exploiting the potential of affinity learning to complement metric learning. In specific, we prepend a set of appearance prompts encapsulated in prompt base to interact with voxel features by cross-attention mechanism to obtain voxel-adaptive appearance prompts, achieving appearance knowledge extraction. Then the resultant voxeladaptive appearance prompts are leveraged to modulate the voxel feature rich in semantic pattern conditioned on the calculated affinity to enhance spatial awareness. Besides, we impose the diversity loss to expand the discrepancy among prompts, enabling prompts to carry diverse and comprehensive knowledge for voxel appearance. For training, we optimize the model with centroid-anchored contrastive learning to well structure the voxel embedding space against the coarseness of manual threshold.

An attention-based task information modeling algorithm is proposed. To solve the problem that the average pooling operation in traditional task embedding generation methods is too coarse, this algorithm introduces the attention mechanism to capture the important difference of different samples, so as to extract more accurate task information. The algorithm utilizes learnable task vectors to store task information and uses an attention mechanism to identify and assign high weights to critical samples, then aggregate sample features to model task information. To verify the effectiveness of the proposed method, extensive experiments are carried out on several standard few-shot image classification datasets. Experimental results show that the proposed attention based task information modeling algorithm achieves better performance compared with the existing methods. Furthermore, our APViT can adapt to multiple post-processing operations (e.g., waterz [Funke et al., 2018]), in other words, it could be possible to enjoy the flexibility with a single trained model via adaptive modulation of the post-processing configuration at the test time.

The contributions of our method could be summarized as follows: (1) We propose an appearance prompt vision transformer tailored for the connectome reconstruction in a coherent and unified framework. Specifically, we design the extension continuity-aware attention module to construct hierarchical attention customized for neuron extensibility and slice continuity, the appearance prompt modulator to exploit the potential of affinity learning to complement metric learning. (2) To the best of our knowledge, this is the first work to absorb the merits of both affinity learning and metric learning formulation by learning long-range spatial dependencies to model spatially-aware voxel semantics in an explicit and flexible optimization strategy. (3) Extensive experimental results on multiple challenging benchmarks demonstrate that our APViT achieves consistent improvements with huge flexibility under the same post-processing strategy.

2 RelatedWork

2.1 Connectome Reconstruction

The reconstruction of connectomes has tremendous biological significance for studying neuronal morphology and activity. Deep learning-based methods have paved the way for research, which can be roughly divided into two categories: object tracking based methods and boundary detection based methods. Among them, the object tracking based methods [Januszewski et al., 2018; Meirovitch et al., 2016] only reconstruct a single neuron at a time, which is inefficient and time-consuming. In contrast, boundary detection based methods exhibit superior performance. Among them, [Funke et al., 2018; Lee et al., 2017] take advantage of affinity learning to separate connectomes with similar appearances. [Lee et al., 2021] aggregates voxels belonging to the same connectome with the help of metric learning and utilizes handcrafted threshold to optimize the reconstruction result. However, the affinity learning paradigm tends to suffer from the absence of explicit voxel semantics, while the metrics learning one inevitably drops the spatial information, and it is greatly affected by the artificial threshold. [Huang et al., 2022] absorbs the merits of both affinity learning and metric learning methods but fails to model long-range dependencies due to the locality of 3D CNNs. Different from those methods, we propose to integrate affinity learning and metric learning via a unified Appearance Prompt Vision Transformer to alleviate the above problems and accomplish the task of connectome reconstruction in an explicit yet flexible manner.

2.2 Vision Transformer and Prompt Learning

Vision Transformer. Transformer was originally introduced in [Vaswani et al., 2017] for machine translation. Many efforts [Sun et al., 2021; Wang et al., 2022; Mai et al., 2023; Luo et al., 2023; Wang et al., 2023; Chen and Lian, 2022] have also been made to apply it to vision tasks, including object detection, image classification and image segmentation. ViT [Dosovitskiy and Beyer, 2020] applies a transformer architecture on sequences of image patches to capture global cues for image classification tasks, building a new foundation for numerous vision tasks. Besides, emerging from transformer, prompt-based learning has been proven effective in NLP tasks by importing language instruction (prompt) to the input text so that the language model can perform well for the downstream tasks. For example, VPT [Jia et al., 2022]

dynamically learns a set of trainable prompts to acquire taskspecific information. Nonetheless, the remaining problem is that it is not suitable to learn generic prompts for scenario adaptation. Hence we design a prompt-aware transformer to model adaptive prompts for different connectomes.

Prompt Learning. Prompting [Liu et al., 2021] originally refers to inserting a few instructions to the input sentences in the NLP tasks [Gao et al., 2021]. Many recent works [Li and Liang, ; Gu et al., 2022] propose to exploit the prompting techniques to deal with different downstream tasks or domains with the combination of transformers without optimizing all of the parameters. In this paper, we prepend a set of appearance prompts to modulate the voxel embedding for better instructing instance voxel semantics.

3 Method

In this section, we first present the overall architecture of the Appearance Prompt Vision Transformer (APViT) as shown in Figure 2, and then describe each component in detail.

3.1 Overview

APViT enables appearance knowledge conditioned on affinity to instruct voxels from a global perspective based on metric learning, that is, exploits the complementary potential of metric learning and affinity learning, and has four stages (indexed with i). Each stage of APViT encapsulates a patch embedding, Li extension continuity-aware attention module (ECAM, Section 3.2), and an appearance prompt modulator (APM, Section 3.3). For training, we optimize the model with centroid-anchored contrastive learning (Section 3.4) to well structure the voxel embedding space against the coarseness of manual threshold in previous work.

In specific, given the input neuron volume I ∈ RD×H×W , where D, H, and W denote depth, height and width, respectively. In the first stage, we first divide I into D × H2 × W2 patches, each of size 1 × 2 × 2. Then, we feed the flattened patches to a linear projection and obtain embedded patches of size DHW 22 ×C1. After that, the embedded patches are passed through a APViT block, and the output is reshaped to a feature map F1 of size D× H2 × W2 ×C1. In the same way, at the beginning of each stage i, using the feature map from the previous stage as input, we obtain the feature map Fi of size D× H Pi × WPi × Ci, where Pi = 2i , and i = {1, 2, 3, 4}.

3.2 Extension Continuity-aware AttentionModule

Considering that all previous methods based on metric learning and affinity learning employ 3D convolutions, however, the inherent locality of 3D CNNs limits both formulations to modeling long-range dependencies and capturing global voxel context, leading to sub-optimal results. Therefore, we develop an extension continuity-aware attention module to construct hierarchical attention customized for neuron extensibility and slice continuity to learn instance voxel semantic context from a global perspective and utilize continuity priors to enhance voxel spatial awareness.

Extension-aware Attention. Considering that interleaved different neuron instances contain intricate morphology, which tends to extend from one end of the input 3D volume to another, we specially design an Extension-aware Attention Module (AttE) to obtain semantic-rich voxel embedding with long-range dependencies. Specifically, given the feature map F ∈ RD× HP × WP ×C (omit the subscript for brevity), we first flatten the spatial dimensions to produce a feature sequence ˜F = R DHW P2 ×C . Queries, keys and values arise from the feature sequence as follow:

where WQ ∈ RC×Ck , WK ∈ RC×Ck , WV ∈ RC×Cv are linear projections. Then we can calculate the extensionattention weight matrix S ∈ R DHW P2 × DHW P2 with the scaled dot-product and the output are computed by S-weighted summation on value V:

where √C is a scaling factor for stabilizing the training and the ⊤ denotes the transpose operation. Following the standard transformer[Vaswani et al., 2017], the Eq. 2 is implemented with the multi-head mechanism and the feed-forward network (FFN) is further applied to obtain the final output.

Continuity-aware Attention. Intuitively, the neural deformation across several contiguous slices is always smooth and continuous. Thus, we employ the Continuity-aware Attention mechanism (AttC) for aggregating information in a corresponding 3D spatial neighborhood for each voxel location, aiming to empower the discrimination to neighboring neuron instances. Predefining a spatial window σ(λ) centered at voxel λ with size of z × p × p, we denotes Fσ(λ) ∈ Rzp2×C as the features of neighbourhood voxels and fλ as the feature of voxel λ. Then, denoting fλ as query, Fσ(λ) as keys and values, we can get q, K, V by:

Similar to Eq. 2, the output of AttC can be calculated by:

Treating D × HP × HW voxels for F in the same way, we can get AttC(F). The above two attention layers work in parallel and the output of them are added and fed into next extension continuity-aware attention module.

3.3 Appearance Prompt Modulator

In order to exploit the potential of affinity learning to complement metric learning, we propose an appearance prompt modulator, leveraging voxel-adaptive appearance knowledge conditioned on affinity to instruct instance voxel semantics. We introduce a prompt base Pi = {pin }Ni n=1 , where Ni refers to the number of appearance prompts for stage i ∈ {1, 2, 3, 4}. The prompts will be leveraged to extract appearance knowledge by interaction with voxel features, then modulate the voxel feature conditioned on the calculated affinity, which highly represents spatial association between voxels. In specific, as shown in Figure 3, for an arbitrary voxel a with feature f ia , we calculate its affinity with adjacent voxel b, c, d at 3 directions, respectively. For simplicity, take b as example, the affinity can be formulated as

The affinity describes the similarity between two spatially adjacent voxels. If the affinity is close to 1, it indicates that the two voxels have a high probability of belonging to the same instance; otherwise, 0 indicates different instances. We use the voxel feature f ia to retrieve the appearance prompts from prompt base Pi and obtain voxel-adaptive prompts, as

The voxel-adaptive appearance prompt ˆ pia is the linear combination of the prompts conditioned on voxel a, and is utilizing to modulate the feature of voxel b as follow,

In order to enable prompts to carry diverse and comprehensive knowledge for voxel appearance, we impose the diversity loss on the prompt base P. Formally,

where the cos(·, ·) denotes cosine similarity.

3.4 Training Objectives

After consecutive four stages process, the output of each stage will be fed into an upsampling module [Hatamizadeh et al., 2022] to restore the original size and get the final voxel embedding E ∈ RD×H×W×C . To well structure the voxel embedding space, we design two centroid-anchored contrastive loss. Firstly, we calculate the centroid ei for each instance i, by averaging the embedding of the voxels belonging to instance i based on the ground truth. With the set of centroid {ei}Ni=1 (N denotes the number of instances), we can get:

where Ei denotes the set of voxel embeddings belonging to instance i, E - ={ej}Nj=1/ei, and the temperature ε controls the concentration level of representations. Intuitively, Eq. 9 enforces each voxel embedding e to be similar with its ground truth ('positive') centroid and dissimilar with other irrelevant ('negative') centroids. Another contrastive loss is proposed for compactness by directly minimizing the distance between each embedded voxel and its ground truth centroid:

Note that both e and ei are ℓ2-normalized. As a result, our overall training objective is formulated as:

where λdiv is the trade-off weight.

4 Experiments

4.1 Experimental Setup

Datasets. Two commonly used neuron datasets, named CREMI [Funke et al., ] and AC3/AC4 [Arganda-Carreras et al., 2015], are used for the evaluation of our method. CREMI dataset is divided into three sub-datasets, each consisting of two volumes of size 125 × 1250 × 1250 for training and testing, respectively. We use the volume with public ground truth as the training and testing set, which consists 100 and 25 slices, respectively. The AC3/AC4 is used for SNEMI3D challenge, where size of AC3 is 256 × 1024 × 1024 and AC4 consists of 100 × 1024 × 1024 voxels. Following the SNEMI3D challenge, We use the top 80 slices of AC4 as training set and the rest of AC4 as validation set. And the top 100 slices of AC3 are testing set.

Implementation Details. In our APViT, the number of layers is {1, 2, 4, 2}. The volume size of the input is anisotropic (18, 160, 160), and the patch size is (1, 2, 2) at each stage. During training, our model is trained with batch size of 2, using the Adam optimizer with an initial learning rate of 0.0001 for 200,000 iterations. And we constrain the output at different resolutions for each stage with GT as an auxiliary loss where we set λdiv = 0.1.

EvaluationMetrics. Following the conventions, VOI (variation of information) is deemed as the main evaluation metric. We also report ARAND (adapted Rand error) to assess the reconstruction results. Smaller values of these two metrics indicate better segmentation performance.

4.2 Comparison with State-of-the-art Methods

Tables 1 and 2 report the neuron reconstruction performance comparison of our method and several state-of-the-art methods on AC3/AC4 and CREMI datasets, respectively. As shown in Table 1, we replace the feature extractors of several methods by our APViT. The insertion of our APViT can bring significant performance improvements to different baseline methods. For example, when we adopt MALA [Funke et al., 2018] as the baseline method and LMC as the post-processing approach, the VOI and the ARAND metrics improve from 1.3857 to 1.0235 and from 0.1143 to 0.0898, respectively. From Table 2, we can observe that our proposed APViT outperforms all previous methods by a substantial margin. Specifically, APViT surpasses the second-best method (PEA) on the VOI metric by 12.01%, 13.56%, 11.81% on Cremi-A, Cremi-B, Cremi-C, respectively.

4.3 Ablation Study and Analysis

To look deeper into our method, we perform a series of ablation studies on AC3/AC4 dataset with waterz as postprocessing to validate the effectiveness of APViT, including the extension continuity-aware attention module (ECAM), the appearance prompt modulator (APM), and the centroidanchored contrastive learning. Note that we remove all proposed modules and only maintain the bald vision transformer, and take manual threshold based optimization following [Lee et al., 2021] as our baseline.

Effectiveness of Main Components. Table 4 summarizes the results of module ablation studies under different configurations. (1) We ablate the ECAM to study the importance of hierarchical attention. As deteriorated results indicate, customized for neuron extensibility and slice continuity is crucial to learn voxel semantic context from a global perspective and utilize continuity priors to enhance voxel spatial awareness. (2) Then we investigate the impact of introducing APM, and observe a absolute performance lift (from 0.1932 to 0.1313 in ARAND). The improvements can be mainly ascribed to the strong ability of the APM to leverage voxel-adaptive appearance knowledge conditioned on affinity to instruct voxel semantics, exploiting the potential of affinity learning to complement metric learning. (3) We also explore the centroid-anchored contrastive learning. When we replace the optimization strategy with manual threshold based optimization [Lee et al., 2021], the performance of the model is clearly degraded. This proves the necessity of contrastive learning to well structure the voxel embedding space against the coarseness and sensitivity of manual threshold. Without all the proposed methods, the model has degenerated into the baseline. The performance improvement of our final model over the baselines is significant.

Advantages of Our Framework. To validate the advantage of our framework tailored for connectome reconstruction, we perform an ablation study to investigate the impact of 3D CNNs (3D ResUNet), pure vision transformer (UNETR [Hatamizadeh et al., 2022]), and our APViT with the same parameters, as tabulated in Table 3. We observe that transformer-based method outperforms CNN-based methods due to long-range dependency modeling capabilities. Furthermore, our APViT achieves a significant lead, which indicates that instead of simply using the vision transformer, APViT absorbs the merits of both affinity learning and metric learning formulation to model spatially-aware voxel semantics in an explicit and flexible optimization strategy. More importantly, it could be possible to enjoy the flexibility with a single trained model via adaptive modulation of the postprocessing configuration at the test time.

Effectiveness of Extension Contiguity-aware Attention. To analyze the ECAM in depth, we ablate extension-aware attention and continuity-aware attention separately, as described in Table 5. Adding either of the two attention mechanisms contributes to a remarkable performance gain. Furthermore, these two mechanisms work in conjunction enables further performance gain, benefiting from learning instance voxel semantic context from a global perspective and utilizing continuity priors to enhance voxel spatial awareness.

Local Window Size. In Table 6, we observe that the local window size in continuity-aware attention has a large impact on reconstruction performance. Experiments show that APViT achieves the best result at the window size of 7×1×1. We conjecture that voxels in the z-stereoscopic direction contain more connectome information, including connections between slices, thus larger size along the z-axis of the local window is more beneficial for feature extraction.

Prompt Strategy. As shown in Table 7, we ablate the components inside the prompt assignment module. In fact, the introduction of a prompt can be deemed as a guidepost in artificially solving incidents, i.e., guiding information, which can strongly correct voxel features based on ECAM and enhance different instance clustering to avoid foreground and background confusion issues better, thus reducing splitting and fusion errors. Therefore, the number of prompts selected also plays an important role in the prompt assignment module, directly affecting the modeling ability of prompts. Compared with using a fixed number of prompts at each stage, the flexible, prompt allocation strategy is obviously more advantageous for the diverse neuron reconstruction task, yielding significant improvements. The principal reason for this phenomenon is that the feature resolution processed by each encoder stage is different, making it necessary to customize the number of prompts according to different stages.

4.4 Explainable Visualization Study

We visualize the reconstruction results of various methods on the test set of AC3/AC4. It can be seen that each method performs well when reconstructing relatively simple and clear neurons, however, when faced with entangled adjacent neuron individuals, the previous methods generally perform unsatisfactorily. In Figure 4, MALA produces more oversegmentation results (e.g., blue and red connectomes), separating neurons that should belong to the same label, resulting in split errors, while APViT is able to reconstruct the correct neuron. The underlying reason is that the extension continuity-aware attention module takes into account a wide range of each neuron. Meanwhile, the continuity-aware attention mechanism takes advantage of the smoothness and continuity between slices, leading to less split error in ambiguous areas. In Figure 5, it can be observed that the embedding is the clear cluster corresponding to each neuron instance. And the affinity map obtained from our voxel embeddings does not misjudge the confusing boundaries. It indicates that the appearance prompt modulator aggregates rich spatial information via the voxel-adaptive appearance prompt.

5 Conclusion

In this paper, we propose we propose a coherent and unified Appearance Prompt Vision Transformer (APViT) to enable appearance knowledge conditioned on affinity to instruct voxels with explicit semantics based on metric learning, including an extension continuity-aware attention module and an appearance prompt modulator. Extensive experimental results on challenging benchmarks show effectiveness.

Contribution Statement

Rui Sun, Naisong Luo and Yuwen Pan contributed equally to this paper.

Acknowledgments

This work was partially supported by the National Nature Science Foundation of China (Grant 62022078, 62021001).