Self-Supervised Neuron Segmentation with Multi-Agent Reinforcement Learning

1 Introduction

Neuron segmentation is a crucial task for neuroscientists that allows for the analysis of the distribution and morphology of neurons, providing valuable insights into the connectomics research [Sheridan et al., 2022; Krasowski et al., 2017]. Electron microscopy (EM) is the mainstream method for accurately identifying neural structures, but the dense nature of neurons and the presence of artifacts and deformations in EM images make the labeling process costly and decrease the credibility of existing annotation data [Deng et al., 2022; Zhou et al., 2019]. Therefore, fully supervised neuron segmentation methods meet great challenges, especially when applied to large scale EM data.

∗Corresponding Author

Self-supervised methods have emerged as a solution to the limitations of fully supervised methods, which can be roughly divided into two categories: contrastive learningbased approach and mask image model (MIM)-based approach. The former requires a large number of positive and negative samples [Chen et al., 2020; Grill et al., 2020; You et al., 2022; Chen et al., 2023] and relies heavily on data augmentation [Caron et al., 2021], making it a highcost option for 3D biomedical images. The latter aims to learn useful structural information in images by masking and recovering certain voxels, which has been recently applied to pretraining biomedical images, showing improvements in downstream tasks [Zhou et al., 2022; Tang et al., 2022; Huang et al., 2022a]. However, the highly localized and structured nature of EM data, as well as the existence of considerable noise, make it inefficient to directly employ the existing MIM in extracting useful information for neuron segmentation. It has also been observed that the masking ratio and masking strategy of MIM are highly sensitive and the optimal ones vary greatly across different datasets. Adjusting these configurations to train large models requires significant manual efforts and resources.

In this paper, targeting the neuron segmentation task, we propose a novel decision-based MIM relying on multi-agent reinforcement learning (MARL) [Littman, 1994] for automatically selecting the appropriate masking ratio and masking strategy, which consists of a target network and a policy network. Our approach partitions the input EM volume into patches and treats each patch as a basic control unit. The overall multi-agent task is modeled as a search process for patch masking strategies, where the action space for each patch is to either keep the original voxels or mask them. The feedback of the target network, in the form of the reconstruction loss, serves as the team reward signal for guiding the policy network to adaptively learn masking strategies that are beneficial for the pretraining task [Foerster et al., 2016]. To improve the stability of training and achieve optimal joint decisionmaking for the entire volume, all agent networks share parameters and are trained in parallel. Furthermore, we introduce the HOG feature as an additional reconstruction loss to enable the target network to learn more structure information. Finally, following a UNETR decoder design [Hatamizadeh et al., 2022], we add a segmentation head to the pretrained target network in the finetuning stage. Experimental results in Figure 1 show that our decision-based MIM achieves clearer reconstruction results than the original MAE [He et al., 2022] in the pretraining phase.

Overall, our main contribution lies in the following aspects:

1. We propose an efficient self-supervised method, named decision-based MIM, for EM neuron segmentation using unlabeled EM data. To the best of our knowledge, it is the first effort that large-scale transformer pretraining is conducted on this task.
2. We propose a MARL-based approach for searching the optimal masking ratio and masking strategy by treating each patch as an agent with a shared policy, effectively exploring the search space and capturing dependencies between voxels.
3. We introduce the HOG feature as an additional reconstruction loss of our decision-based MIM, improving the convergence speed of network training and the performance of the downstream segmentation task.
4. We comprehensively demonstrate the effectiveness of our proposed method on two representative EM datasets, especially against alternative self-supervised methods on the task of neuron segmentation.

2 RelatedWork

2.1 Neuron Instance Segmentation

In the field of EM image processing, neuron instance segmentation is an important task. [Turaga et al., 2010] first proposed a convolutional neural network based on affinity generation, followed by clustering voxels in the affinity graph into instances based on post-processing methods such as watershed and LMC [Beier et al., 2017]. In recent years, there have been more advanced networks for the affinity-based approach. Funke et al.[Funke et al., 2018] introduced the MALIS loss[Briggman et al., 2009] during the training process to encourage the network to generate correct topological segmentation. [Huang et al., 2022b] introduced an embedding pyramid module to simulate affinity at different scales. [Liu et al., 2022] incorporated both embedding and affinity information and combined it with the graph neural network to further improve the distinguishability of adjacent objects in the feature space. However, due to the anisotropic resolution of 3D EM images in lateral and axial directions, the usage of transformer structures remains unexplored in the field of neuron segmentation. In this paper, we use an affinity-based setup and upgrade the decoder of UNETR [Hatamizadeh et al., 2022] to adapt to the anisotropic EM features.

2.2 Mask Image Model (MIM)

The MIM is an important branch of self-supervised learning. Masked autoencoders (MAE) [He et al., 2022] used an asymmetric encoding-decoding structure, encoding only the unmasked patches and using a lightweight decoder to recover the masked patches, which greatly reduced the resources required for computation and quickly became a mainstream structure for MIM. [Feichtenhofer et al., 2022] and [Tong et al., 2022] separately validated the effectiveness of MAE on video datasets and proved that higher image masking ratios can be used in 3D datasets. [Zhou et al., 2022] was the first to introduce MAE to the medical image field. [Bachmann et al., 2022] introduced a multi-modal multi-task adaptation to MAE, resulting in better performance than the original MAE. [Gao et al., 2022] proposed a multi-scale mixed convolution to encode images and achieved improved results on fine-grained downstream tasks such as image segmentation. [Wei et al., 2022] focused on the reconstruction target of the decoder and found that reconstructing artificial features such as HOG and SIFT can facilitate the network to classify and localize objects.

In summary, existing MIM-based methods mainly focus on the design of the encoder architecture, prediction head, and prediction target. Although many works have demonstrated the impact of masking strategies on downstream tasks, there has been little research along this line, and the latest method also requires a search-based masking strategy with a fixed masking ratio [Bandara et al., 2022]. This paper proposes a novel MARL-based approach for adaptively learning the optimal masking ratio and masking strategy, making the pretrained model more robust and achieving better performance on downstream tasks.

2.3 Multi-Agent Reinforcement Learning (MARL)

In the field of deep RL, multiple agents working together can improve the efficiency and robustness of the model due to the limited observation and action space of a single agent. Given the complexity of computer vision tasks, MARL is often used to interact with the common environment to make decisions. [Liao et al., 2020] proposed a method for iteratively refining medical image segmentation using interactive MARL, where rough image segmentation is initially provided and the network is iteratively refined through user feedback in the form of rewards until the segmentation is sufficient. [Lin et al., 2021a] proposed a method for augmenting images through the use of blocks, with each block acting as an agent and working together to produce optimal data augmentation. Still, the application of MARL in computer vision suffers from the large decision space and difficulty in obtaining rewards. In this paper, to reduce the complexity of the state and the difficulty of searching for RL policies, we first segment the image into patches through a transformer-encoder and treat each patch as an agent. The action space for each patch is limited to only two options, masking or keeping the original voxels, and the rewards are obtained through the reconstruction loss of MAE. Therefore, the masking decision can be naturally modeled as a Markov process and optimized through MARL.

3 ProposedMethod

Our decision-based MIM consists of two stages of training: pretraining and finetuning. Figure 2 illustrates the overall flow of the network training. In this section, we will first introduce some basic theories of Vision Transformer (ViT) and MARL, and then explain in detail the encoders, decoders, and loss functions of the two stages, as well as specific MARL modeling methods.

3.1 Encoder-Decoder Design

We use ViT [Dosovitskiy et al., 2021] as the backbone architecture for decision-based MIM pretraining and downstream segmentation tasks. To represent high-dimensional data in a ViT, we must transform it into a sequence of patches. Given an input 3D volume x ∈ RH×W×D×C , where C is the number of channels and (H,W,D) is the resolution, we reshape it into a sequence of flattened 3D patches xp ∈ RN×(P 3 ·C) . The patch resolution is given by (P/4, P, P ), and the number of patches is calculated as N = 4HWD P 3 . These patches are then transformed into patch embeddings via a trainable linear projection.

Consistent with the MAE setup, we divide the image patches into visible and masked groups. The encoder in the MAE and ViT architectures processes only the visible blocks. To enhance the performance of our decoder, we incorporated a histogram-of-oriented-gradients (HOG) feature [Dalal and Triggs, 2005], which has been shown to improve pretraining. This is achieved by providing the decoder with various markers, including a patch representation from the encoder and learnable position embeddings. By incorporating positional embedding in all input markers, we enable the decoder to simultaneously recover both the HOG feature and the original voxels, resulting in superior performance compared to the MAE method.

In our pretraining process, we utilize a loss function that combines both the mean squared error (MSE) loss for reconstructing the original voxels and the HOG loss for recovering the HOG features. The HOG feature can be calculated using the following equation

where HOGi,j is the histogram of oriented gradients for the cell located at position (i, j), Si,j is the set of voxels in the cell (i, j), w(x) is a weighting function that assigns a weight to each voxel x in the cell, and g(x) is the gradient orientation of the voxel x.

Our overall loss function can be expressed as

where λ1 and λ2 denote the weights assigned to the MSE and HOG losses, respectively. We set λ1 to 0.1 and λ2 to 1.

3.2 DecisionModule

In our proposed decision module, we model the masking strategy of image patches in MAE as a multi-agent cooperative decision-making problem and adopt a multi-agent reinforcement learning method to solve it. Here, we will introduce in detail the modeling methods of model states, observations, actions, the design of multi-agent team rewards, and the learning method used to update the policy network.

As shown in Figure 3, given the original input batch of image x, according to the ViT setup, we divide the image into equal-sized and non-overlapping patches. Our MARL policy aims to determine the overall joint masking policy based on the current input state and the observation of each agent. Our basic setup is as follows.

State. The enhancement policy for each patch is closely related to the information of the context, so the decisionmaking process of MARL requires perceiving the semantic information of the image rather than directly inputting the patches of the raw image. In order to ensure the consistency and convergence of training, we use the target network, i.e. ViT, as the backbone to extract deep semantic features of the image. The global state at time step t is denoted as St and is visible to all agents.

Observation. In addition to capturing global information St , each agent needs to make a masking policy decision based on its own local observation. In general, the observation in MARL tasks is often a part of the global state. Considering these factors, we take the deep feature ViT(Pi) of patch Pi as the current observation value Ot i for the i-th agent. The feature extractor for local features is the same as the one for global features, both using the ViT backbone.

Action. The action of the i-th agent aims to output whether patch Pi needs to be masked. We define the action of the i-th agent as a vector Ai. The joint action space can be represented as A = {A1,A2, ...,AN}, where N represents the total number of patches. Since the action space only has two possibilities, masking or keeping the original voxels, the dimension of Ai is 2. Given the current state St and observation Ot i , each agent i will determine an action ai(St ,Ot i ) ∈ Ai based on the current policy. The final output is the global joint action at = {at1 ∪ at2 , ..., ∪atN }. After all patches have taken their corresponding actions through the decision policy, the time step is updated to t = t + 1 and we obtain the enhanced volume through the decision module.

Rewards. Rewards are introduced in our MARL decisionmaking process in order to guide the agents to learn expected behaviors that will improve the main task's performance through more reasonable masking ratios and masking strategies, allowing the target net to better learn semantic information in the volume. Previous works [Lin et al., 2021b; Zhang et al., 2019] attempted to increase the training loss of the target network based on rewards in order to generate deeper, more difficult-to-learn features. Inspired by these works, we refine the reward design based on the MAE pretraining paradigm. By comparing the loss difference between the data obtained from the masked data x·at-1 generated by the previous time step's decision module and the data obtained from the current decision module x·at , we compute the reward for the MARL policy. This encourages higher training loss during the MARL decision-making process, as shown by the equation

In the above equation, Lpretrain represents the reconstruction loss generated by the target network, and ϕ denotes the target network. The accumulated reward of one sequence is

where T represents the desired time step length to be calculated, the discount factor γ takes a value in (0, 1], and ¯ rt is the mean rewards at time t.

Policy Learning. Considering that the action space for MARL decisions is discrete, we utilize the widely used Asynchronous Advantage Actor-Critic (A2C) algorithm [Mnih et al., 2016] to performMARL policy learning. Since the search space for actions is not large, we use simple convolution and pooling modules to adjust the Actor and Critic networks to the structure shown in Figure 3. The policy network is divided into an Actor and a Critic, which are adapted to the RL training algorithm. The Actor network learns a discrete control policy π(ati |St , Ot i ), while the Critic network aims to estimate the value of the state Vπ(St ). We model the centralized action value function Q, which takes in the state information S and the actions of all agents, and outputs a Q value for the team, given by

where at represents the joint action of all agents, defined as a = {ai, · · · , aN}, and Rt is the long-term discounted reward, given by equation 4. The advantage function on the policy is then given by

where Aπ (St , at ) is the advantage of taking action at given state St at time step t, V π (St ) is the current state estimate output by Critic. It indicates that the actual accumulated reward is independent of the state and reduces the variance of the gradient. We use θp and θv to denote the parameters of the Actor and Critic, respectively. The squared value of the dominance function Aπ is taken as the loss function to update θv as

To further achieve cooperative capability, the loss function of the updated Actor θp is defined as

where πθ(at | St ) is the Actor output, that is, the probability of taking each action ati . The Actor and Critic are jointly trained in an end-to-end manner.

3.3 Neuron Instance SegmentationMethod

The UNETR model is specifically designed for 3D segmentation tasks, such as organ segmentation. It is built upon the pretrained ViT encoder from decision-based MIM and a randomly-initialized convolutional decoder. The UNETR architecture is inspired by the U-Net model [Ronneberger et al., 2015], with the inclusion of skip connections between the encoder and decoder at multiple resolutions. The input to the UNETR decoder is a sequence of representations, denoted as z3, z6, z9, z12, which are reshaped to restore the spatial dimension of HP × WP × DP ×K, where K is the feature channel.

Starting from the deepest feature, z12, each representation is processed through a varying number of deconvolutional blocks to increase its resolution by a specific factor. For example, z12 and z9 are upsampled by a factor of ×2, z6 is upsampled by a factor of ×4, and z3 is upsampled by a factor of ×8. Then, representations at the same spatial resolution, such as z12 and z9, are concatenated and further upsampled to match the shape of a shallower feature. This process of concatenation and upsampling is repeated until the full resolution of the input is restored. Finally, the output layer combines the upsampled feature and the original full-resolution input to predict the segmentation map.

In this paper, for the segmentation task of EM neurons, we use a combination of affinity map-based and post-processing approaches. We first create an affinity graph based on the voxel affinity. This graph serves as the foundation for our post-processing techniques, which utilize both waterz [Funke et al., 2018] and LMC [Beier et al., 2017] to cluster the affinity map and produce the final neuron segmentation results. Affinity-based methods have proven to be highly effective in accurately segmenting and analyzing these complex structures.

4 Experiments

4.1 Training Strategy

Our strategy consists of two phases: pretraining and finetuning. In the pretraining phase, to improve the training efficiency of the framework, we first pretrain decision-based MIM for 100k iterations, synchronously updating the parameters of MAE and the policy network in the decision module, and then fix the parameters of policy network and only update the parameters of MAE for another 100k iterations. In the fine-tuning phase, we load the pretrained ViT weights into the model for the downstream task and train for 200k iterations.

We use the Adam optimizer in both the pretraining and fine-tuning phases, with β1 = 0.9, β2 = 0.999. The only difference lies in the pretraining process, where we set the learning rate to 0.0001 and perform batch size 16 pretraining on 8 RTX 3090s. In the fine-tuning phase, we adopt a Layerwise Learning Rate Decay (LLRD) training method, which adjusts the learning rate layer by layer during training. We set the learning rate of the last layer's parameters to 0.001 and the learning rate of the previous layer's parameters to 0.95 times the learning rate of the next layer's parameters. We conduct batch size 8 fine-tuning on 2 RTX 3090s.

4.2 Datasets and EvaluationMetrics

FAFB. The Full Adult Fly Brain (FAFB) dataset [Zheng et al., 2018] is a highly valuable resource for neuroinformatics research, offering a comprehensive and detailed view of the neural architecture of the Drosophila melanogaster (fruit fly) brain. With a size of approximately 40 terabytes, this dataset features high-resolution 3D images with a resolution of approximately 4 nanometers per pixel, as well as manually annotated segmentation data identifying various brain structures such as neurons, glial cells, blood vessels, synapses, and other neuropil regions.

In our work, we utilize the FAFB dataset as a key component in our pretraining process. To optimize the efficiency of our model, we first downsample the original dataset by a factor of 4, carefully curating a selection of 60G images that exhibit exceptional imaging quality from the entire dataset. This strategic selection ensures that our model is trained on the most accurate and reliable data possible, setting the foundation for its future performance.

CREMI. The CREMI dataset is derived from the FAFB dataset, which has three manually labeled subvolumes from drosophila brain, of which CREMI A has more regular forms, CREMI C has higher size disparities, and CREMI B has inbetween segmentation difficulty. Each sub-volume has 125 slices of 1250×1250 images, and we choose the first 60 slices for training, 15 slices for validation, and the remaining 50

slices for testing, which are utilized to validate our method's performance on segmentation tasks of varying difficulty.

AC3/AC4. AC3/AC4 [Kasthuri et al., 2015] are mouse somatosensory cortex datasets with 256 and 100 successive EM images (1024×1024), respectively. The first 80 slices of AC3 are used as the training set, the following 20 slices as the validation set, and the first 50 slices of AC4 are used as the testing set.

EvaluationMetrics. We are more interested in the model's performance in the downstream task in the self-supervised training. To assess the influence of segmentation on EM neurons, we primarily use Variation of Information (VOI) [Nunez-Iglesias et al., 2013] and Adapted Rand Error (Arand) [Arganda-Carreras et al., 2015] metrics. Smaller VOI and ARAND values represent better segmentation results.

4.3 Experimental Results

Decision Making Process. Decision-based MIM is pretrained using the FAFB dataset. During pretraining, the policy network in the decision module is also updated, and the real-time change in masking ratio is recorded as in Figure 4. The decision-making process starts with random decisions, resulting in an overall masking ratio of around 0.5. Then, using the reconstruction loss as a reward, the decision-making principles of the agents are updated. After 50k iterations, the decision module converges and produces better reconstruction results with less lost information. During pretraining, it is found that starting with a lower masking ratio is more beneficial for the EM dataset, and gradually increasing the masking ratio helps the network's learning process progress from easy to difficult datasets, which is effective for both upstream reconstruction and downstream segmentation tasks.

Results on CREMI. We compare our method to the original MAE pretraining method [He et al., 2022] and the Dino pretraining method [Caron et al., 2021] based on contrastive learning, as well as two affinity-based Unet structures for instance segmentation, superhuman [Lee et al., 2017] and MALA [Funke et al., 2018]. The outcomes of the trials are provided in Table 1 after applying both waterz and LMC post-processing procedures. Figure 5 shows the visualization results. Compared to existing self-supervised methods, our decision-based MIM approach shows a significant improvement for the downstream task and it also performs better than commonly used CNN-based methods in neuron segmentation, as can be seen in Table 1.

Results on AC3/AC4. In order to verify the robustness and generalizability of our pretraining method, we further conduct experiments on mouse cortical neural cells. The results in Table 1 indicate that our approach leads to a clear improvement in the downstream segmentation task, even when applied across different species.

4.4 Ablation Study

We use UNETR to perform a comprehensive ablation study on the AC4 dataset.

Effectiveness of the multi-task reconstruction. We conduct ablation experiments on the model's reconstruction targets. As shown in Table 2, it demonstrates that utilizing HOG and MSE losses together for reconstruction in the MIM outperforms using MSE or HOG alone.

Effectiveness of the decision module. We compare our proposed method to a straightforward solution that manually adjusts the masking ratio. The ablation results, as shown in Table 3, demonstrate that our approach not only eliminates the need for manual adjustment of masking ratios but also outperforms the best results achieved through manual adjustment.

5 Conclusion

In this paper, we propose a decision-based MIM approach for neuron segmentation. Our method eliminates the need for manual adjustment of masking ratios and masking strategies, using multi-agent cooperation to search for the optimal solution. Additionally, during the pretraining process, we incorporate a multi-task reconstruction and utilize HOG features to enhance the model's learning ability. Our method is validated on a variety of EM neuron datasets to demonstrate its generalizability.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant 62021001.