Efficient Neuron Segmentation in Electron Microscopy by Affinity-Guided Queries
Authors: Hang Chen
Chufeng Tang
Xiao Li
Xiaolin Hu
Venue: ICLR 2025 Poster (2025)
Keywords: Neuron Segmentation; Biomedical Image Segmentation; Electron Microscopy Image
Primary area: applications to neuroscience & cognitive science
TL;DR: This paper propose a query-based model for neuron segmentation in 3D volume EM images and achieved superior results as well as a 2-3x speedup.

Abstract: Accurate segmentation of neurons in electron microscopy (EM) images plays a crucial role in understanding the intricate wiring patterns of the brain. Existing automatic neuron segmentation methods rely on traditional clustering algorithms, where affinities are predicted first, and then watershed and post-processing algorithms are applied to yield segmentation results. Due to the nature of watershed algorithm, this paradigm has deficiency in both prediction quality and speed. Inspired by recent advances in natural image segmentation, we propose to use query-based methods to address the problem because they do not necessitate watershed algorithms. However, we find that directly applying existing query-based methods faces great challenges due to the large memory requirement of the 3D data and considerably different morphology of neurons. To tackle these challenges, we introduce affinity-guided queries and integrate them into a lightweight query-based framework. Specifically, we first predict affinities with a lightweight branch, which provides coarse neuron structure information. The affinities are then used to construct affinity-guided queries, facilitating segmentation with bottom-up cues. These queries, along with additional learnable queries, interact with the image features to directly predict the final segmentation results. Experiments on benchmark datasets demonstrated that our method achieved better results over state-of-the-art methods with a 2$\sim$3$\times$ speedup in inference. Code is available at https://github.com/chenhang98/AGQ.

Published as a conference paper at ICLR 2025
EFFICIENT NEURON SEGMENTATION IN ELECTRON MICROSCOPY BY AFFINITY-GUDDED QUERIES
Hang Chen', Chufeng Tang', Xiao Li', Xiaolin Hu
Department of Computer Science and Technology, Institute for AI, BNRist, Tsinghua University, Beijing 100084, China
Tsinghua Laboratory of Brain and Intelligence (THBI),
IDG/MeGovern Institute for Brain Research, Tsinghua University, Beijing 100084, China
3. Chinese Institute for Brain Research (CIBR), Beijing 100010, China
(chenhang20, 1ie0,,ixxxe 1ixiao20 (aro200mmaai @maalms ts tsinghua edu, en; chufeng t foxmail. com x xlhu@tsinghua edu.on
ABSTRACT
Accurate segmentation of neurons in electron microscopy (EM) images plays a crucial role in understanding the intricate wiring patterns of the brain. Existing automatic neuron segmentation methods rely on traditional clustering algorithms,, where affinitics are predicted first, and then watershed and post-processing al- gorithms are applied to yield segmentation results. Due to the nature of wa- tershed algorithm, this paradigm has deficiency in both prediction quality and speed. Inspired by recent advances in natural image segmentation, we propose to use query-based methods to address the problem because they do not neces- sitate watershed algorithms. However, we find that directly applying existing query-based methods faces great challenges due to the large memory require- ment of the 3D data and considerably different morphology of neurons. To tackle these challenges, we introduce affinity-guided queries and integrate them into a lightweight query-based framework. Specifically, we first predict affinities with a lightweight branch, which provides coarse neuron structure information. The affinities are then used construct affinity-guided queries, facilitating segmenta- tion with bottom-up cues. These queries, along with additional learnable queries, interact with the image features directly predict the final segmentation results. Experiments on benchmark datasets demonstrated that our method achieved better results over -fate-of-the-art methods with a 2~3x x speedup inference. Code available at https:::p / htt.b com/ hatmm//hhhnaag//Agg
INTRODUCTION
Neuron segmentation in electron microscopy (EM) images has significant scientific importance for connectomics research, which enables the comprehensive reconstruction of neural connections the nervous system and provides insights into the functioning of the brain (Winding et. [023 Scheffer
et al., 2020 Dorkenwald et al., 2022 Peddie et al. 2022). Reconstructing large-scale structures typically involves handling vast amounts of data, ranging from terabytes to petabytes (Scheffer al., 2020). Given the intricate ttree-dimensional and densely packed nature of neurons, the manual reconstruction process is extremely laborious and often requires years of effort from human experts to reconstruct entire brain (e.g.. Drosophila larval brain (Winding et al., 2023
In recent years, many methods (Funke et al. 2019 Lee et al. 2017 Knowles- Knowles-Barley et al. 2016 Huang et al. 2022b Chen et al.) 2024) have been proposed to address the challenges of neuron segmentation by employing deep neural networks as an alternative to labor-intensive manual pro- cedures. These methods usually follow bottom-up segmentation paradigm, as depicted in Fig- ure [ (a). Initially, 3D affinities, which indicate the probabilities of neighboring voxels belonging to the same neuron, are predicted using neural network (e.g., 3D U-Net (Çiçek al. 2016)) and then transformed into fragments via a clustering algorithm where watershed (Meyer 1992 is the
"Corresponding author.

---

Published as a conference paper at ICLR 2025
Volume Image
Affinities
Fragments
Segmentations
Volume Image
Segmentations
Querics
U-Net
Watershed
Agglomeration
Query-based Model
(a) Previous
(b) Ours
Comparison of the proposed model with previous neuron segmentation methods. (a) Previous methods utilize neural networks predict affinities, and then derive the final segmentation results via watershed and agglomeration. (b) The proposed query-based model directly predicts the segmentation results, which is more concise and efficient. Best viewed in digital with zoom-in.
Figure 1:
de-facto practice. Subsequently, these fragments are merged into final neuron segmentation results through an agglomeration algorithm (Funke et al. 2019), which merges fragments based their affinities at interfaces. However, these methods, heavily relying on traditional clustering algorithms such watershed, suffer from several drawbacks that limit their application: (1) The watershed algorithm usually produces fragments with unnatural boundaries such as straight lines, referred to as artifacts, due competition between seeds, leading inaccurate segmentation results (as shown in Section 4.4); (2) These methods may incorrectly merge different neurons and make them difficult to distinguish in the agglomeration step; (3) The watershed algorithm lacks parallelizability, mak- ing it unsuitable for GPU acceleration, resulting in inefficiency. Moreover, the large amount (e.g.. volume (Huang et al. (022bb) of generated fragments slows down the
more than 100 thousands subsequent cost-processing steps such agglomeration. Despite recent efforts to improve neuron segmentation through various techniques (Chen et al.) 2023; Funke et al. 2019 Huang et al. 2022a
Liu et al. 2023b 2022 (see Section 2). they still adhere to the traditional paradigm and, as result,
encounter the aforementioned problems. Note that there are also methods @Januzzwski et al. 2018
Meirovitch al. 2019) that work without using watershed, which discuss in detail in Section
Neuron segmentation can be considered as a form of instance segmentation which has been ex- tensively studied in computer vision. But it has notable differences from conventional instance segmentation. Firstly, the input data involves 3D volume images rather than 2D images. Secondly, the target objects in neuron segmentation are 3D neurons rather than natural objects with distinctive appearance. Given the morphologically complex and homogeneous nature of neurons, which exhibit dense distributions and similar appearances, distinguishing different neurons presents significant challenge. For example, rbgion-based methods like Mask R-CNN (He et al., 2017) may encounter difficulty in addressing this problem due to the ambiguity of bounding boxes, as pointed out in Huang et al. (2022b). Recently, query-based instance segmentation methods have shown promise in natural images, offering both efficiency and accuracy advantages (Cheng et al. 2202bb a Zhang et al. 2021). However, these methods cannot be directly applied to neuron segmentation because their specific design for 2D natural images faces problems such as unaffordable memory overhead when ported 3D data.
In this work, we seek to explore the possibility of the query-based model in neuron segmentation, which could potentially address the limitations associated with current methods. The motivation is that the query-based model can directly predict the final segmentation results, thus avoiding the inaccuracy and inefficiency caused by watershed. However, we found that transferring the query- based methods to neuron segmentation is nontrivial. First, existing query-based models are tailored for 2D natural images and cannot be directly transferred 3D neuron segmentation due to the un- affordable GPU memory overhead. Second, the distinct characteristics of neurons in EM, such as their homogeneous appearance and intricate morphology, which differ significantly from objects in natural images, could hinder the learning of queries. To tackle these challenges, we propose an efficient EM neuron segmentation model (Figure (b)) with Affinity-Gudded Query (AGQ). AGQ is constructed based the predicted affinities from lightweight affinity branch. It incorporates bottom-up cues, allowing the model to initiate predictions from coarse approximations rather than learning from scratch, thereby significantly reducing the learning difficulty. AGQ and learnable queries interact with the volume image features and with each other, capturing the diverse morpho- logical characteristics of neurons and complementing the predictions of each other.
2

---

Published as a conference paper at ICLR 2025
We conducted experiments benchmark datasets AC3/AC4 (Kasthuri et al. 2015) and ZE- BRAFINCH (Kornfeld et al. 2017). The results demonstrated that our method achieved superior results in terms of both accuracy and efficiency. In comparison to state-of-the-ar methods, our method achieved significantly lower errors with 200 ~ 300% speedup. Note this speedup is signifi- cant considering that a volume for connectomics research could reach the mm3 level (Zheng al.) 2018 Winding et al. 2023; Shapson-Cor et al. 2024) and processing mm'? volume requires hundreds of GPU hours plus hundreds of CPU hours (SheriSnn al. 2023).
The main contributions of the paper are summarized as follows:
By introducing fffinity-guiedd queries, we devise a lightweight query-based neuron seg- mentation method, eliminating the clustering algorithm used in traditional methods.
Extensive experiments validated the effectiveness of our method, improving both the aceu- racy and efficiency of neuron segmentation over previous methods.
2 RELATED WORK
Neuron Segmentation. Previous methods mainly follow bottom-up segmentation fashion (Funke ctt al. 2019 Lee et al. 2017 Huang et al.) 20226). They use a neural network (e.g., a 3D U-Net) to predict affinities, followed by the watershed algorithm to transform affinities into fragments and an agglomeration algorithm (Funke et al.] 2019) to obtain the final results. The watershed and agglom- eration algorithms require processing on the CPU, causing inefficiency. Recent advances include improving visual representation (Chen et al. 2023 Shhridann et al.) 2023), devising superior loss functions (Funke et al., 2019). implementing label efficient supervision (Huang et al. 2022a Liu itt al. 2023b). employing learnable agglomeration modules (Liu et al. (2022), and improving the U-Net architecture (Luo et al., 2024] [Sun al.,
2023). APVÍT (Sun et al.) 2023) advanced U-Net by improving the ViT architecture with tailored designs such as learnable prompt base. However, these methods still follow the previous segmentation paradigm and suffer from inaccuracy and inef- ficiency. FFN (Januszewski et al. (2018) processes one individual neuron at time. Despite being a new practice, FFN two orders of magnitude slower than affinity- methods and has similar performance (Sheridan et al., 2023). 3C (Meirovitch al. 2019) introduced combinational encod- ing to process multiple neurons in parallel, but required multiple runs and could only merge those neurons, making it more of a substitute for agglomeration. Instead, the proposed method directly predict the final segmentation results via a query-based paradigm, which is concise and efficient.
Instance Segmentation. The mainstream methods in instance segmentation include two-stage (He et al. 2017) and one-stage ones (Tian al. 2020; Wang et al., 2020a b). However, these methods were not suitable for the neuron segmentation. Two-stage methods rely on boxes to distinguish instances and cannot handle dense and widely distributed neurons. One-stage methods could not guarantee recall. Recently, query-based instance segmentation algorithms (Fang et 2021 (02T::Chen) Cheng et al., [222a]] [Zhang al.) (202)) have emerged, which directly predict objects. However, these methods cannot be directly applied to neuron segmentation because their specific design for 2D natural images faces problems such as unaffordable memory overhead when ported to 3D data. Mo- tivated by their success in natural images, we explore a lightweight query-baser model coupled with affinity-guideer queries specially tailored for neuron segmentation task. It is worth noting that al- though object detectors like Deformable DETR (two-stage version) (Zhu et al. 2021) also involve predicted queries, we address significantly different tasks and problems, with plenty of design dif- ferences (e.g., their predicted queries are implicit and do not work with learnable queries).
3D Medical and Biological Image Segmentation. The segmentation of organs or lesions from CT (Wang et al., 2019; Zhou et al.) 2019) or MRI {Ji et al. 2022 Zeng et al. 2020) images is the primary focus of 3D medical image segmentation. These tasks mainly involve semantic-level segmentation, which differs from the instance-level segmentation discussed in this paper. Building upon the 3D U-Net (Çiçek et al.) 2016). various studies have improved the segmentation quality by refining the model structure, module design (Hatamizadeh et al.
2022
2021 Peiris et al. 2022
Shaker et al. 20244) and label efficiency (Wu et al. [222 Zhou et al.
2023) The 3D biological image
segmentation tasks in electron microscopy include: mitochondria segmentation (Pan et al.
2023
Mai et al. 2023 Lin et al. 2021), synapse detection (Lin al. 2021 and soma segmentation
3

---

Published as a conference paper at ICLR 2025
al. 2023a). They typically follow the bottom up fashion, where semantic-level segmentation is first predicted and then converted into a instance-level segmentation by post-processing
3 METHOD
In this section, we first introduce the preliminaries involving the problem definition of EM neuron segmentation and the previous methods. We then present brief overview of our method, including the problem formulation and the framework. Subsequently, we detail the design and function of each module, as well as the training objectives. Finally, we describe how to apply our method on the large 3D volume EM images via block assembly.
3.1 PRELIMINARIES
3.1.1 PROBLEM DEFINITION
Given volume EM image RDxHxW " the neuron of all voxels, denoted as S E NDxHxx where D, H, W denote the depth, height, and width of the image, respectively. Each id in S indicates the segmentation result of specific neuron, and id indicates the background.
segmentation algorithm is required predict
3.1.2 PREVIOUS METHODS
As shown in Figure ! [a), previous methods typically employ deep neural networks (e.g., 3D U- Net) to predict affinities A E [0, 13x first. These affinities represent the probabilities that neighboring voxel" voxel=bbelong belong to the same neuron (Gao et al. 2019 Ke et al. 2018 Funke Funke et al. 2019)
(see Appendix A for more explanation). Subsequently, a seeded watershed algorithm (Meyer 1992)
is applied to extract the fragments G based on the predicted affinities:
G (atershed(A) E NDx#XW
Finally, the fragments G obtained from the seeded watershed algorithm merged using ag- glomeration algorithm based on the statistics of the affinities at their interfaces, yielding the final segmentation result:
Agglomeration(C,.) NDDxxW
The agglomeration usually only merges fragments and cannot dissolve incorrect mergers. Thus, the seeds of the watershed are usually densely picked, incurring significant redundancy. On the one hand, the redundancy causes unnatural boundaries such as straight lines, referred artifacts, due to the competition between seeds; while on the other hand, it slows down the overall inference speed.
3.2 QUERY- BASED NEURON SEGMENTATION FRAMEWORK
To address the above issues, we aim to devise a concise and efficient query-based neuron segmen- tation method. We adopt new modeling manner and directly predict N potential segmentation probability maps P E (0, 1) XxDxXxW based on N queries. During inference, we obtain the segmentation results S, by simply taking a argmax on the first dimension of obtain the seg- mentation results (i.e., neuron id);
S. = =rg max (Pr...) NDDHXW
(3)
Without the watershed algorithm, this modeling simplifies the inference process considerably and could achieve faster inference than previous methods.
Figure instantiates several query-based models, consisting 3D U-Net backbone and decoders predict segmentation results. We create baselines based
mainstream decoder design (Figure 2 (a) and (b)) from state-of-the-art query-based models on natural images (Zhang et al.) 2021 Cheng et al., 2021: 2022a). These baselines have been adapted for the 3D data of neuron segmentation to make the GPU memory usage affordable. However, we found that neither baseline
Fach voxel has six neighbors, of which only three need he predicted due duplication.
4

---

Published as a conference paper at ICLR 2025
LQ Tranaformer Decoder
LO
Query
Volume Image
50 U-t
Seg
Volume Image
U-Net
Seg.
Volume Image
30 Net
C-hir
Seg.
(a) Transformer Decoder
(b) K-Net Decoder
(c) Ours
Figure 2:
Comparison of different query-based models. LQ and AGQ denotes learnable queries and affinity-guicen queries, respectively.
could tackle the task-specific challenges neuron segmentation (see Section 4.5.1 The reason could be that query-based methods natural images typically involve only learnable queries that are randomly initialized and updated by gradient descent. These queries lack inductive bias and does not discriminate well between neurons which have the same appearance and complex structure. This may underlie the difficulty encountered by these methods. In view of this, we propose to embed certain prior knowledge of affinities into the generation of queries. We chose K-Net decoder as shown in Figure (c), since our experiments showed that it outperformed the Transformer decoder.
3.3 OVERVIEW OF THE PROPOSED METHOD
Figure shows our framework. which consists of 3D U-Net backbone network for feature extrac- tion, a module to generate affinity-guides queries, and 3D neuron decoder for directly predicting segmentation results. We then describe each module in detail.
Lattinity
Connected
Mw Allinity-Giieed Queriies
1W Learnable Queries
Companents
Branch
Affinities
Mooarse
Seg
3D Neuron Decoder
Qo
feeture
Bipartite Matching
P6
3D U-Net
Pred.
GT
Volume Image Features
Multiply
Concat.
Supervision
Overview the proposed method. Qo and Po stand for the initial queries and segmen- tation probability maps, respectively. Mcoarse denotes coarse binary masks. Laftiaity Lleature Leeg are losses for predicted affinities, image features and segmentation probability maps, respectively.
Figure 3:
3.4
3D NEURON DECODER
Our efficient 3D neuron decoder is inspired by the query-based model for natural images (Zhang al. 2021). The decoder comprises multiple stages with the same structure, as depicted in Figure Each stage incorporates three inputs:
The image features F E BCXDxxXW where C' denotes the feature dimension
The segmentation probability maps from the previous stage Pi-1 E (0, NxDHHxW,
The neuron queries from the previous stage Qi-1 N where i denotes the current stage (see Section 3.5 for the definition of Po and Qo).
With these inputs, we initially update the neuron queries using image features and segmentation predictions from the previous stage:
Q = "hyConv(Q:--1 P-11FT)
TThe dimension of batch size is omitted for simplicity.

---

Published as a conference paper at ICLR 2025
DyConv:
SelfAtt:
Image Features
Queries Qi-1 Kest
Queries
Queries
P
a R
Lineary
Ou E R Wood
Softmax
Seg. Prob. Maps
Sey. Probo Maps
(0.1) (XXXIN
) YOUUSAA
Figure 4: The detailed structure of the 3D Neuron Decoder. K' denotes the number of stages.
where DyCony, is a dynamic convolution layer (Fang et al. 2021), and P.-1FT E RN The enhanced queries can better capture the structural information of the morphologically complex and homogeneous neurons. Subsequently, these queries are processed through a self-attention layer SelfAtt; (Vaswani et al. 2017):
Q: = =eeelAtt;
(5)
where the queries interact with each other to better distinguish among the neurons (e.g., deciding which query predicts which neuron). Finally, these queries interact with the image features to gen- erate the segmentation probability maps for the current stage,
= Softmax (Linear,
where Linear, Lin) E DX and softmax is applied along its first dimension. The Linear; represents linear layer. This global and direct prediction approach can mitigate potential systematic artifacts that might arise from the watershed algorithm.
3.5 AFFINITY-GUIDEL QUERY
To enhance the learning efficiency of the 3D neuron decoder, propose the affinity-guided queries. These queries are designed to include coarse neuron structure information, allowing the decoder to initiate the process from it rather than from scratch. Initially, we utilize affinity branch composed of multiple convolutional layers to predict affinities based on image features F,
= Affnnhyy E (0,11)xDxxxNN
(7)
Subsequently, the predicted affinities are averaged along the first dimension to derive (0, 1) 1xxx> which is then thresholded to create 3D binary mask. We use the mean value of across all dimensions, refered as a, as the threshold. This is to ensure that there is at least one foreground connected component during training. Following this, fast connected component extraction algorithm (Wu et al. 2005) is employed to generate coarse masks:
Mccarse = ConnectedComponeet
a).
(8)
where Mcoarse E {0, N Nx comprises collection of (denoted N..) boolean masks. By multiplying and averaging each of these masks with the image features, we obtain the respective affinity-guided queries Qaffnity XC where C² denotes the feature dimension. These queries encapsulate the rudimentary neuron structure details, thereby facilitating the subsequent stages and reducing their reliance training data. This enhancement leads to improved segmentation accuracy.
Due to the possibility of incorrect merging in Mcoarse. the number of affinity-guidee queries could be less than the actual number of neurons, resulting in missing neurons not being predicted. Thus we also use N₁ learnable queries Qlearuable E R M We concatenate them with affinity-guided queries together as the initial queries of the 3D neuron decoder (i.e., N N. + Ni),
Qo = Concat(Qulmit Qlearnable) E [[N,+N))xCC
(9)
Note that No varies between passes through the network while N; is kept constant. The learnable queries are expected to interact with affinity-guided queries through self-attention layer (in Equa- tion (5) to recognize and predict the missing neurons. The initial segmentation probability maps Po are created by taking the produc of Qo (post linear layer) and image features F, which is then passed through softmax function, i.e., Po Daatmax(Lincar,(Q)
"The two matrices have the shapes N' C and C (D H W), with reshaping omitted for simplicity.

---

Published as a conference paper at ICLR 2025
3.6 TRAINING OBJECTIVES
During training, we employ three loss functions to supervise the predicted affinities, the image features, and the final segmentation results, respectively.
L Luffinity Licature Leag:
(10)
where Laffinity denotes the binary cross-entropy loss, Lieature denotes the contrastive loss, and Leeg a weighted combination dice loss and cross-entropy loss. See Appendix B for more details.
3.7 BLOCK ASSEMBLY
Due the extensive coverage of volumetric EM images across large continuous regions (e.g., 100x 1024 1024 voxels), it is impractical to directly input the entire volume into the network. Current methods utilize sliding-window technique, which generates affinity on individual cropped blocks (e.g.. 17x257 x257 voxels) and then assembles them together for watershed and agglomer- ation. The sliding stride is smaller than the block's s size (typically half) to reduce border effects. In method, the predictions of the blocks also require assembly. To achieve this, reuse the agglomeration function, but for merging the results of distinct blocks instead of fragments. Another difference is that do not necessitate overlap to address border effects, which is an advantage over previous methods. More details are discussed in Appendix
4 EXPERIMENT
4.1 DATASETS AND METRICS
Datasets. To validate the effectiveness of our method, we conducted experiments on the beneh- mark datasets AC3/AC4 (Kasthuri et al.) 2015) and ZEBRAFINCH (Kornfeld et al.) [017). AC3 and AC4 are two human-labeled sub-volumes extracted from the mouse somatosensory cortex dataset (Kasthuri et al., 2015). with dimensions of 256 x 1024 1024 and 100 x 1024 x 1024, respectively. The ZEBRAFINCH dataset contains 33 volumes (approximately 150 x 150 x 150), of which we used 30 volumes training set and 3 volumes as test More details about the datasets can be found in Appendix D
Metrics. We utilized the widely used Variation of Information (VOI) (Nunez-Iglesias et al. 2013) and Adapted Rand Error (Arand) Aarganda-Carrera can further divided into two types of errors, VOIsplit and mOlmerge. which refer to the errors of segmenting single neuron into multiple segments and merging multiple neurons into the same segment, respectively. Both VOI and Arand indicate errors, with a lower value indicating better accuracy. Further elaboration can be found in Appendix F
al. 2015) as the evaluation metrics. VOI
4.2 IMPLEMENTATION DETAILS
By default, we employed a 3D neuron decoder with two stages (i.e., K = 2 in Figure 4 comprising 100 learnable queries and a maximum of 100 affinity-giided queries. For feature extraction, we adopted 3D U-Net with ResBlock (He et al.) 2016) (i.e., ResUNet (Xiao al. 2018)). We used Adam optimizer and trained 20k iterations by default. Training was performed on 8x NVIDIA 3090 GPUs and costed about 40 hours. More implementation details provided in Appendix F
4.3 QUANTITATIVE RESULTS
We first conducted experiments on the AC3/AC4 dataset and presented the results in Table The results were reproduced utilizing the pytorch.connect omi os codebase (Lin et al. 2021 (e.g., ResUNet (Xiao al. 2018) and SwinUNETR (Hatamizadeh et al.] 2021)) referenced from pub- lished papers (e.g., PEA (Huang et al. 2022b) and Frag ViT (Luo et al. (024)). Our method was compared with these watershed-basee methods in terms of both accuracy and efficiency.
Regarding accuracy, our method generally outperformed these methods, exhibiting superior perfor- mance in metrics such as the lowest VOT and competitive Arand. This shows the strong potential of

---

Published as a conference paper at ICLR 2025
Table 1: Results the AC3/AC4 dataset. "." indicates that the value is unavailable. Inference times were tested with an NVIDIA 3090 GPU and 64 Intel Xeon Gold CPUs, which represent the time (in seconds) to process the full test set. "*' indicates results obtained from published papers (Huang et al. 2022b
Luo et al. 2024h. VOI, and VOIm are short for VO and VOImerge- respectively.
Metrics
Inference Time
VOL,
VOL..
VOI
Arand
Model Watershed
Agg.
Total
Xiao UNt 2018
al.
1.037
0.258
1.295
0.154
81.2
30.8
35.8
147.8
SeUNet (Ln et al 2021)
1.031
0.251
1.282
0.156
83.1
30.7
36.0
149.8
UNETR Hatammzadeh el al. 2022
2.750
0.281
3.031
0.220
37.0
41.4
133.0
211.4
SwinUNETK (Hatamizadeh of 2021
1,238
0.191
1.429
0.110
80.0
32.5
43.1
155.6
LSD (Sheridan et al. 2023
1.448
0.229
1.677
0.134
229.9
29.7
26.5
286.1
ML-De (De Brabandere
2017)
1.575"
0.615"
2.190"
0.196"
SuperFluman [Lee el al. 2017
1.145"
0.263"
1.408"
0.122"
52.3
28.4
19.4
100.1
MALA (Funke ct al. 2019)
1,304
0.242*
1.546"
0.120'
PEA (Huang et al. 20226
0.852
0.232"
1.084"
0.094"
60.2
37.1
25.4
122.7
Frag VTI (Luo
2024
0.868
0.191*
1.054"
0.093"
>60.2
255
122.7
APVIT
2023)
0.767
0.209"
0.976"
0.078
60.2
-=7.11
2554
122.7
AGQ (ours)
0.677
0.290
0.967
0.095
27.6
N/A
6.1
33.7
Results on the ZEBRAFINCH dataset. Results three test volumes are reported. Inference times were tested with an NVIDIA 3090 GPU and 64 Intel Xeon Gold CPUs, which represent the time (in seconds) required to process all the test volumes.
Table 2:
Volume-1
Volume-2
Volume-3
Infer. Time
VOI
Arand
VOI
Arand
VOI
Arand
Total
ResUNet (Xioo al. 2018
0.272
0.013
1.967
0.135
1.003
0.154
21.1
SeUNet (Ln el al. 2021
0.271
0.012
1.894
0.152
1.021
0.162
21.0
UNETR Hatamizadeh èl al. 2022
3.013
0.403
3.394
0.529
2.665
0.360
17.7
SwinUNETR [Haaamzzadeh el 2021
0.744
0.206
2.022
0.154
0.693
0.035
34.9
SuperHuman (Lec ct al. 2017)
0.268
0.011
1.867
0.107
0.721
0.042
16.6
PEA (Huang 20226)
0.208
0.009
1.715
0.122
0.534
0.028
17.4
AGQ (ours)
0.177
0.008
1.278
0.134
0.338
0.024
6.4
our method, as a brand new neuron segmentation method. The achievement of the lowest VOIsplit (indicating errors in predicting neuron multiple segments) aligns with our motivation to avoid the extensive fragments associated with the watershed algorithm.
Regarding efficiency, our method significantly surpassed previous methods, due its concise neuron segmentation modeling and framework. Firstly, our method did not generate segmentation results from affinities, eliminating the need for overlap prediction address border effects (as discussed in Section 3.7). As result, the model's inference time was reduced. Secondly, with its compact modeling, our method predicted segmentation results without the need of watershed, thus entirely saved this time. Finally. despite reusing the agglomeration function (as explained in Section our results were relatively complete with fewer regions to process, leading to significant time savings in the agglomeration step. Overall, our method achieved a speedup of about 340% over ResUNet and 260% over PEA. Regarding Frag ViT and APVIT, the inference time was estimated, as the code was unavailable. The estimation and post- processogn method as PEA but with a heavier network.
made due Frag ViT and APVIT employing the same watershed
On ZEBRAFINCH dataset (Table 2 our model also exhibited general better accuracy the three test volumes, with speedup of about 200%. For example, our method achieved significantly lower VOI all test volume and 230% and 170% faster than ResUNet and PEA, respectively. For deployment, our model consumed 6.5G GPU and 5.4G CPU memory during inference on these datasets, suitable for most computing devices.
8

---

Published as a conference paper at ICLR 2025
Volume Image
GroundTrud
Net
SuperHuumen
PISA
AGQ (Ours)
Ground Truth
SuperHumen
(Quis)
Figure 5: Visualization comparison of segmentation results obtained from different methods. The top two rows present a 2D slice alongside two zoomed-in regions, highlighting artifacts with white arrows. The third row illustrates neuron structures in 3D space, with colored arrows indicating locations that are prone to failed reconstruction. Best viewed in digital with zoom-in.
4.4 QUALITATIVE RESULTS
We present qualitative comparison of segmentation results on AC3/AC4 obtained using previous methods and our method in Figure 5 The first two rows display the results on 2D section, il- lustrating that our method can avoid the artifacts present in previous methods. the second row, indicated by white arrows, the watershed-basee methods generated numerous unnatural boundaries (i.e., artifacts). In contrast, our method employed a learning-based approach to predict the final result directly, effectively reducing this issue.
The third row in Figure [ showcases the 3D morphology of several predicted neurons. Highlighted by colored arrows, method excels in reconstructing the neuronal structure. We observed that tra- ditional methods struggle with locations characterized by thin branches or myelin sheaths, indicating the limitations of relying solely affinities. Predicted affinities in these areas may be uncertain or distorted. Some failure cases of our method are presented and discussed in Appendix G
To further analyze the roles of different queries in prediction, visualized the intermediate outputs of the 3D neuron decoder in the Appendix H As shown in Figure TO the integration of learnable queries gradually refines the segmentation, effectively recovering missing incorrectly merged neurons (highlighted by yellow arrows). This analysis highlights the distinct roles of the two query types: celinity-guided queries provide initial predictions, while learnable queries assist in correcting segmentation errors, particularly in resolving neuron mergers.
4.5 ABLATION STUDY
We conducted ablation experiments on the AC3/AC4 dataset to analyze the individual effects each module. We report both the metries blocks and the full test set (i.e., before and after block assembly). Note that the models were trained for 10k iterations to save training costs, thus there is discrepancy with the results in Table More ablation experiments be found in Appendix |
4.5. DESIGN OF NEURON DECODER
As shown in Table the K-Net style decoder achieved better results, probably due to its more explicit way of extracting instance features. Consequently, we chose the K-Net style decoder as a baseline and incorporated the proposed affinity-guidde queries to achieve the best results.

---

Published as a conference paper at ICLR 2025
Table 3: Design of neuron decoder.
Block
Full
VOI
Arand
VOlapii
VOlmerge
VOI
Arand
Transformer Decoder
1.245
0.184
1.695
1.138
2.833
0.395
K-Net Decoder
0.710
0.087
1.131
0.562
1.693
0.160
Ours
0,538
0.065
0.781
0.397
1.177
0.141
4.5.2 EFFECT OF AFFINITY- GUIDED AND LEARNABLE QUERIES
We studied the effectiveness of affinity-guided queries (AGQ) in Table Completely excluding AGQ significantly deteriorated results (denoted as L.Q in the table). Merely increasing the number of learnable queries (i.e., Double IQQ) did not lead to improved results, suggesting that enhancing segmentation quality solely by increasing query quantity is insufficient.
The effect of learnable queries (LQ) was validated by comparing the results with and without them, shown in Table (AGQ v.s. LQ AGQ). The incorporation of learnable queries notably enhanced the segmentation results, particularly in mitigating merging errors (i.e., merger This finding aligned with our motivation that learnable queries could retrieve incorrectly merged neurons.
Table 4: Effect of alfinity-guided queries.
Block
Full
VOI
Arand
split
VOl
merge
VOI
Arand
LQ
0.710
0.087
1.131
0.562
1.693
0.160
Double IQQ
0.695
0.082
1.078
0.650
1.729
0.209
AGQ
0.836
0.102
0.822
0.881
2.010
0.229
LQ+AGQ
0.538
0.065
0.781
0.397
1.177
0.141
4.5.3 AFFINITY GUIDANCE
We compare different ways of leveraging affinity guidance to validate whether affinity-guided query is the optimal. Two naive baselines are concatenating affinities coarse segmentation onto image features, which can fuse information from affinities into image features instead of queries. As shown in Table AGQ yielded the best results, showing its effect in leveraging bottom-up cues.
Table Results of different approaches of using bottom-up cues.
Block
Full
VOI
Arand
VO
VOI
merge
VOI
Arand
Concat affinities
0.836
0.102
0.836
0.102
2.010
0.229
Concat coarse segmentation
0.869
0.104
0.975
0.917
1.893
0.213
Affinity guided queries
0.538
0.065
0.781
0.397
1.177
0.141
CONCLUSION AND DISCUSSION
We propose a new method named AGQ for neuron segmentation in 3D volume EM images. Diverg- ing from previous methods that heavily rely on watershed, our method offers a more concise and efficient solution. Extensive experiments validated the effectiveness of our method, which achieved superior results and 2~3x speedup. We hope our method could advance the development of large-scale neuron reconstruction, thereby deepening includes exploring improved ways to handle large volumes and focused improvements for failure cases (failure cases and limitations are discussed in Appendix
understanding of the brain. Future work
ACKNOWLEDGMENNT
We thank Gang Zhang for helpful discussions and Lesi Wei for developing the drawing scripts. This work was supported in part by the National Key Research and Development Program of China (No. 20212ID02(0331) and the National Natural Science Foundation of China (No. U2341228)
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ignacio Arganda-Carreras, Srinivas C Turaga, Daniel R Berger, Dan Ciresan, Alessandro Giusti, Luca M Gambardella, Jürgen Schmidhuber, Dmitry Laptev, Sarvesh Dwivedi, Joachim M Buh- mann, et al. Crowdsourcing the creation of image segmentation algorithms for connectomics. Frontiers neuroanatomy, 9:142, 2015.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In ECCV, pp. 213-229, 2020.
Hang Chen, Chufeng Tang, and Xiaolin Hu. Dense contrastive loss for instance segmentation. In BMVC, pp. 1062, 2022.
Yinda Chen, Wei Huang, Shenglong Zhou, Qi Chen, and Zhiwei Xiong. Self-supervised neuron segmentation with multi-agent reinforcement learning. IJCAI, pp. 609 609-617, 617, 2023.
Yinda Chen, Wei Huang, Xiaoyu Liu, Shiyu Deng, Qi Chen, and Zhiwei Xiong. Learning Multiscale Consistency for Self-supervised Electron Microscopy Instance Segmentation. In ICASSP, pp- 1566-1570, 2024.
Bowen Cheng, Alexander G. Schwing, and Alexander Kirillov. Per-pixel classification is not all you need for semantic segmentation. In NeurIPS, pp. 17864-17875, 2021.
Bowen Cheng, Ishan Misra, Alexander G Schwing, Alexander Kirillov, and Rohit Girdhar. Masked- attention mask transformer for universal image segmentation. In CVPR, pp. 1290-1299) 2022a.
Tianheng Cheng, Xinggang Wang. Shaoyu Chen, Wengiang Zhang, Qian Zhang, Chang Huang, Zhaoxiang Zhang. and Wenyu Liu. Sparse instance activation for real-time instance segmentation. In CVPR, pp. 4433 4442, 2022b.
Özgün Çiçek, Ahmed Abdulkadir, Soeren S. Lienkamp, Thomas Brox, and Olaf Ronneberger. 3D U-Net: Learning dense volumetric segmentation from sparse annotation. In MICCAI, pp. 424- 432, 2016.
Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint 151 2015.
Bert De Brabandere, Davy Neven, and Luc Van Gool. Semantic instance segmentation with a discriminative loss function. arXiv preprint arXiv:170..0255// 2017.
Sven Dorkenwald, Claire E McKellar, Thomas Maerina, Nico Kemnitz, Kisuk Lee, Ran Lu, Jing- peng Wu, Sergiy Popovych, Eric Mitchell, Barak Nehoran, et al. Flywire: online community for whole-brain connectomies. Nature Methods, 19:119- 19:119-128, 2022.
Yuxin Fang, Shusheng Yang, Xinggang Wang, Yu Li, Chen Fang, Ying Shan, Bin Feng, and Wenyu Liu. Instances as queries. In ICCV, pp. 6890-6899, 2021.
Jan Funke, Fabian Tschopp, William Grisaitis, Arlo Sheridan, Chandan Singh, Stephan Saalfeld, and Srinivas C. Turaga. Large Scale Image Segmentation with Structured Loss Based Deep Learning for Connectome Reconstruction. IEEE TPAMI, 41:1669-1680, 2019.
Naiyu Gao, Yanhu Shan, Yupei Wang, Xin Zhao, Yinan Yu, Ming Yang, and Kaiqi Huang. Ssap: Single-shot instance segmentation with affinity pyramid. In ICCV, pp. 642-651, 2019.
Ali Hatamizadeh, Vishwesh Nath, Yucheng Tang, Dong Yang, Holger R Roth, and Daguang Xu. Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. In MICCAI Brainlesion Workshop, pp. 272-284, 2021.
Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Land- man, Holger R Roth, and Daguang Xu. Unetr: Transformers for 3d medical image segmentation. In WACV, pp. 574-584, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, pp. 770-778, 2016.
11

---

Published as a conference paper at ICLR 2025
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In ICCV, pp. 2961- 2969, 2017.
Wei Huang, Chang Chen, Zhiwei Xiong, Yueyi Zhang, Xuejin Chen, Xiaoyan Sun, and Feng Wu. Semi-supervised neuron segmentation via reinforced consistency learning. IEEE TMI, 41(11): 3016-3028, 2022a.
Wei Huang, Shiyu Deng, Chang Chen, Xueyang Fu, and Zhiwei Xiong. Learning to Model Pixel- Embedded Affinity for Homogeneous Instance Segmentation. AAAI, pp. 1007-1015, 2022b.
Michat Januszewski, Jörgen Kornfeld, Peter Li, Art Pope, Tim Blakely, Larry Lindsey, Jeremy Maitin-Shepard, Mike Tyka. Winfried Denk, and Viren Jain. High-precision automated recon- struction of neurons with flood-filling networks. Nature methods, 15(8):605-610, 2018.
Yuanfeng Ji, Haotian Bai, Chongjian Ge, Jie Yang, Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhanng, Wanling Ma, Xiang Wan, al. Amos: A large-scale abdominal multi-organ benchmark for versatile medical image segmentation. NeurlPS, 35:36722-36732, 2022.
Narayanan Kasthuri. Kenneth Jeffrey Hayworth, Daniel Raimund Berger, Richard Lee Schalek, José Angel Conchello, Seymour Knowles-Barles Dongil Lee, Amelio Vázquez-Reina Verena Kaynig, Thouis Raymond Jones, et al. Saturated reconstruction of volume of neocortex. Cell, 162:648-660 661, 2015.
Tsung-Wor Wei Ke, Jyh-Jing Hwang, Ziwei Liu, and Stella Yu. Adaptive affinity fields for semantic segmentation. In ECCV, pp. 587-602, 2018. 2018.
Seymour Bhowles-Barley, Verena Kaynig, Thouis Ray Jones, Alyssa Wilson, Joshua Morgan, Dongil Lee. Daniel Berger, Narayanan Kasthuri, Jeff W. Lichtman, and Hanspeter Pfister. RhoanaNet Pipeline: Dense Automatic Neural Annotation. arXiv preprint 2016.
Jörgen Kornfeld, Sam E Benezra, Rajeevan T Narayanan, Fabian Svara, Robert Egger, Marcel Ober- laender, Winfried Denk, and Michael A Long. Em connectomics reveals axonal target variation in sequennc--enneratnng network. Elife, 6:e24364, 2017.
Kisuk Lee, Jonathan Zung, Peter Li, Viren Jain, and H. Sebastian Seung. Superhuman Accuracy on the SNEMI3D Connectomics Challenge. arXiv preprint 1706.0100 2017.
Zudi Lin, Donglai Wei, Jeff Lichtman, and Hanspeter Pfister. Pytorch connectomics: A scalable and flexible segmentation framework for connectomics. arXiv preprint arXiv:21 2021.
Rosanne Liu, Joel Lehman, Piero Molino, Felipe Petroski Such, Eric Frank. Alex Sergeev, and Jason Yosinski. An intriguing failing of convolutional neural networks and the coordconv solution. NeurIPS, 31, 2018.
Xiaoyu Liu, Wei Huang, Yueyi Zhang, and Zhiwei Xiong. Biological instance segmentation with superpixel-guided graph. In IJCAI, 2022.
Xiaoyu Liu, Bo Hu, Mingxing Li, Wei Huang, Yueyi Zhang, and Zhiwei Xiong. A soma segmenta- tion benchmark in full adult fly brain. CVPR, pp. 7402-7411, 2023a.
Xiaoyu Liu, Wei Huang, Zhiwei Xiong, Shenglong Zhou, Yueyi Zhang, Xuejin Chen, Zheng-Jun Zha, and Feng Wu. Learning cross-representation affinity consistency for sparsely supervised biomedical instance segmentation. In ICCV, pp. 211072211 2023b.
Naisong Luo, Rui Sun, Yuwen Pan, Tianzhu Zhang, and Feng Wu. Electron microscopy images as set fragments for mitochondrial segmentation. In AAAI, pp. 3981-3989, 2024.
Huayu Mai, Rui Sun, Tianzhu Zhang, Zhiwei Xiong, and Feng Wu. Dualrel: Semi-supervised mitochondria segmentation from A prototype perspective. In CVPR. pp. 19617-19626. 2023.
Yaron Meirovitch, Lu Mi, Hayk Saribekyan, Alexander Matveev, David Rolnick, and Nir Shavit. Cross- closscllasifiiattinn clustering: An efficient multi-object tracking technique for 3-d instance seg- mentation in connectomics. In CVPR, pp. 8425-8435, 2019.
12

---

Published as a conference paper at ICLR 2025
Fernand Meyer. Color image segmentation. In ICIP, pp. 303-306, 1992.
Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi. V-net: Fully convolutional neural net- works for volumetric medical image segmentation. In 3DV, pp. 565-571, 2016.
Juan Nunez-Iglesias, Ryan Kennedy, Toufiq Parag, Jianbo Shi, and Dmitri B Chklovskii. Machine learning of hierarchical clustering to segment 2d and 3d images. PloS one, 8:e71715, 2013.
Yuwen Pan, Naisong Luo, Rui Sun, Meng Meng. Tianzhu Zhang, Zhiwei Xiong, and Yongdong Zhang. Adaptive template transformer for mitochondria segmentation in electron microscopy images. In ICCV, pp. 21417-21427, 2023.
Christopher Peddie, Christel Genoud, Anna Kreshuk, Kimberly Meechan, Kristina D Micheva, Kedar Narayan, Constantin Pape, Robert Parton, Nicole L Schieber, Yannick Schwab, et al. Volume electron microscopy. Nature Reviews Methods Primers, 2:51, 2022.
Himashi Peiris, Munawar Hayat, Zhaolin Chen, Gary Egan, and Mehrtash Harandi. A robust volu- metric transformer for accurate 3d tumor segmentation. In MICCAI, pp- 162-172, 2022.
Louis K Scheffer, C Shan Xu, Michal Januszewski, Zhiyuan Lu, Shin-ya Takemura, Kenneth J Hayworth, Gary B Huang, Kazunori Shinomiya, Jeremy Maitlin- Shelind- Stuart Berg, et al. A connectome and analysis of the adult drosophila central brain. Elife, 9:e57443, 2020.
Abdelrahman M Shaker, Muhammad Maaz, Hanoona Rasheed, Salman Khan, HinggHsuan Yang, and Fahad Shahbaz Khan. Unetr++: delving into efficient and accurate 3d medical image seg- mentation. IEEE TMI, 2024.
Alexander Shapson-Cod Michal Januszewski. Daniel R. Berger, Art Pope. Yuelong Wu, Tim Blakely, Richard L. Schalek, Peter Li, Shuohong Wang, Jeremy Maitin-Shepard, Neha Karlu- pia, Sven Dorkenwald, Evelina Sjostedt, Laramie Leavitt, Dongil Lee, Jakob Troidl, Forrest Coll- man, Luke Bailey, Angerica Fitzmaurice, Rohin Kar, Benjamin Field, Hank Wu, Julian Wagner- Carena, David Aley, Joanna Lau, Zudi Lin, Donglai Wei, Hanspeter Pfister, Adi Peleg, Viren Jain, and Jeff W. Lichtman. A petavoxel fragment of human cerebral cortex reconstructed at nanoscale resolution. Science, 384(6696), 2024.
Arlo Sheridan, Tri M Nguyen, Diptodip Deb, Wei-Chung Allen Lee, Stephan Saalfeld, Srinivas C Turaga, Uri Manor, and Jan Funke. Local shape descriptors for neuron segmentation. Nature methods, 20(2):295-303, 2023.
Rui Sun, Naisong Luo, Yuwen Pan, Huayu Mai, Tianzhu Zhang, Zhiwei Xiong, and Feng Wu. Appearance prompt vision transformer for connectome reconstruction. In ISCAI, pp. 1423-1431, 2023.
Christian Szegedy. Vincent Vanhoucke. Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In CVPR, pp. 2818-2826, 2016.
Zhi Tian, Chunhua Shen, and Hao Chen. Conditional convolutions for instance segmentation. In ECCV, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurIPS, pp. 6008, 2017.
Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen. Max-deeplab: End- to-end panoptic segmentation with mask transformers. In CVPR, pp. 5463-5474, 2021.
Xinlong Wang, Tao Kong, Chunhua Shen, Yuning Jiang, and Lei Li. Solo: Segmenting objects by locations. In ECCV, 2020a.
Xinlong Wang, Rufeng Zhang, Tao Kong, Lei Li, and Chunhua Shen. Solov2: Dynamic and fast instance segmentation. In NeurlPS, 2020b.
Yan Wang, Yuyin Zhou, Wei Shen, Seyoun Park, Elliot K Fishman, and Alan L Yuille. Abdominal multi-organ segmentation with organ-attention networks and statistical fusion. Medical image analysis, -528 2019.
13

---

Published as a conference paper at ICLR 2025
Michael Winding, Benjamin D Pedigo, Christopher L Barnes, Heather G Patsolic, Youngser Park, Tom Kazimiers, Akira Fushiki, Ingrid Andrade, Avinash Khandelwal, Javier Valdes- Valdes-Aleman, et al. The connectome of insect brain. Science, 379:cadd9330, 2023.
Kesheng Wu, Ekow Otoo, and Arie Shoshani. Optimizing connected component labeling algo- rithms. In Medical Imaging 2005: Image Processing, volume 5747, 2005.
Linshan Wu, Jiaxin Zhuang, and Hao Chen. Voco: A simple-yet-effective volume contrastive learn- ing framework for 3d medical image analysis. In CVPR, 2024.
Yuxin Wu and Kaiming He. Group normalization. ECCV, pp. 3-19, 2018.
Xiao Xiao, Shen Lian, Zhiming Luo, and Shaozi Li. Weighted res-unet for high-quality retina vessel segmentation. In ITME, pp- 327-331, 2018.
Chenyi Zeng, Lin Gu, Zhenzhong Liu, and Shen Zhao. Review of deep learning approaches for the segmentation of multiple sclerosis lesions on brain mri. Frontiers in Neuroinformaties, 14: 610967, 2020.
Wenwei Zhang, Jiangmiao Pang, Kai Chen, and Chen Change Loy. K-net: Towards unified image 2021.
segmentation. In NeurIPS, pp.
Zhihao Zheng, J Scott Lauritzen, Eric Perlman, Camenzind G Robinson, Matthew Nichols, Daniel Milkie, Omar Torrens, John Price, Corey B Fisher, Nadiya Sharifi, et al. A complete electron microscopy volume of the brain of adult drosophila melanogaster. Cell, 174:730-743, 2018.
Hong- Zhou, Chixiang Lu, Chaoqi Chen, Sibei Yang, and Yizhou Yu. unified visual information preservation framework for self self-supervised pre-training in medical image analysis. 2023.
Yuyin Zhou, Zhe Li, Song Bai, Chong Wang. Xinlei Chen, Mei Han, Elliot Fishman, and Alan L Yuille. aware neural network for supprvly- multi-organ segmentation. In CVPR, pp. 10672-10681, 2019.
Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai. Deformable DETR: deformable transformers for end-to-end object detection. In ICLR, 2021.
14

---

Published as a conference paper at ICLR 2025
A MORE EXPLANATIONS ABOUT AFFINITIES
The affinities indicate that whether the adjacent voxels belong to the same neuron, which have shape of D H x W where D, H, W is the depth, height and width of the volume image, respectively. For each voxel, these three values indicate the connection between the current voxel (i,j, k) and its three neighboring voxels: (i 1, (i, 1, k), and (i, k- 1), respectively. Note that due to redundancy, the voxel (i. k)) does not need to predict the connection on the other side (e.g., the connection between (2. j. k) and kk can be predicted by the voxel (i+11, j. k)).
The corresponding ground truth affinity value is 1 if and only if the neighboring voxels both belong to the same neuron, and 0 in all other cases (e.g.. belonging different neurons or belonging to background). We present a 2D example in Figure 66 demonstrating the cases of different boundary widths between two neurons (neuron A and neuron B).
width=0
width=1
width=2
1
1
voxel neuron
voxel neuron
voxel backgroud
affinity=0
1
affinity=1
Figure 6: A 2D example demonstrating the cases of different boundary widths between two neurons (neuron and neuron B).
B MORE DETAILS ON TRAINING OBJECTIVES
B. SEGMENTATION LOSS Leeg
To enable the end-to-end training, following DETR (Carion et al. 2020), we first identify the cor- respondence between predictions and the ground-truth masks via the bipartite matching and then calculate the segmentation loss based on these matched pairs. The predictions are represented as segmentation probability maps
(0, 10,1) Xx The ground-truth masks are denoted as MGT E {0, 11 x (padded with empty ) The optimal bipartite matching between the predictions and the ground-truth masks, symbolized as a permutation of N indices, is deter- mined using the Hungarian algorithm (Carion et al. 2020),
= = min Σ DICE M MCT),
(11)
where DICE denotes the dice coefficient Miilletari et al. 2016). Upon establishing the correspon- dence o, the segmentation loss is computed by combining dice loss and cross-entropy loss with weighted summation:
Leeg = ADICE CDICE MT + ACE CC)
(12)
where the loss weights are specified as ADICE = 3 and ACE = 0.3, Pa denotes the rearrangement of P according & in the first dimension. Given that the decoder outputs multiple results from various stages (i.e., Po, Pi,.... PH in Section 3.4). the segmentation loss is calculated individually for each output and then summed up the total segmentation loss.
B.2 FEATURE Loss Licature
To enhance the differentiation among neurons in image features F' E DXCDDXXW we utilize contrastive loss (Wang et al. 2021 Chen et al. 2022) supervise these features. Initially, we
Note that N is usually larger than the number of neurons.
15

---

Published as a conference paper at ICLR 2025
compute the summed features T E NxC for each neuron,
E
M i.d.b,a P F,,,h
(13)
&SDASHHWWW
The features T RNCC is then normalized along the channel dimension to obtain T IR R Subsequently, the feature loss is computed as follows,
E. M i.d.h.w exp(Ti:: F.d.wo/T) exp((Ti P.a,.//T)
Lleature = -feature
Σ
log
(14)
&SD.ASH.WWW
where Afeature 0.1 represents the loss weight, and 0.3 denotes the temperature. Notably, the calculation excludes the padded empty mask. By incorporating this loss, similar features are en- couraged for voxels of the neuron, while different features promoted for those belonging to different neurons. This facilitates the decoder in effectively discerning between individual neurons.
B.3 AFFINITY Loss Laftinity
The predicted affinities A E (0, 1x3xDxHxW were trained using binary cross-entropy loss with the ground truth AGT E (0, (0,1)sxDxHHB 333DDHxW Note that AGT is obtained from the segmentation ground-
truth masks (i.e., MGT in Appendix B.1 which follows previous work (Lin et al. 2021 Huang
al. 20226). The affinity loss is
atfinity
Xaffinity
LBCB (Ac (A, AGT),
(15)
where the loss weight Aamutty Additionally, We adopt label smoothing technique (Szegedy al. 2016) with 10-5 prevent overfitting.
C MORE DETAILS ON BLOCK ASSEMBLY
C.I FULL PIPELINE COMPARISON
Watershed
(a) Previous
Querr baned
Apgioennt
Model
(b) Ours
Figure 7: Full pipeline comparison two neuron segmentation paradigms. (a) Previous methods utilize sliding-window blocks and then assembles them together for watershed and agglomeration algorithms. (b) The proposed query-basee model directly predicts the segmentation results. The results of the distinct blocks are merged by reusing the agglomerate function. Best viewed in digital with zoom-in.
window technique for inference. which generates affinities individual cropped
In Figure 7 we present a comparison between the complete pipeline the previous method and that of our method. Previous methods involved predicting affinities individual blocks, which
16

---

Published as a conference paper at ICLR 2025
then assembled for watershed and agglomeration. In contrast, method directly predicts the seg- mentation results on the blocks. However, assembling different blocks together poses a challenge due to the lack of natural correspondence in neuron id across blocks, unlike affinities. To address this problem. we utilize the agglomeration function to merge neurons that are separated by block boundaries. Note that opt reuse the agglomeration function just for implementation simplic- ity. Given that our task is less complex than what agglomeration typically handles (as shown in Figure a[)),, there might be more efficient approaches block assembly that could be explored. We show quantitatively in Table the changes in metrics before and after agglomeration (i.e., block assembly). Agglomeration significanll reduced split errors, due neurons from different blocks being assembled, but also introduced merge errors to some extent.
Table 6: Results before and after agglomeration.
Agglomeration
VOIppls 3.606
VOL
VOI
Arand
Before
0.181
3,786
0.874
After
0.681
0.267
0.947
0.089
C.2 OVERLAP PREDICTION AND BORDER EFFECTS
key difference between our method and previous methods is the significant reliance of the lat- ter affinities. This reliance has resulted in common practice of predicting with overlap, such as employing a sliding window with a stride equal to half of the input size, to mitigate the impact of affinities' border effects (as depicted in Figure 8 In contrast, this strategy does apply to our method, as we directly predict the segmentation results without relying on affinities. Table 7 provides quantitative analysis of the effect of this strategy (i.e., overlap prediction) on previous methods. On one hand, the absence of this strategy led to varying decreases in accuracy for previ- ous methods, indicating their reliance on overlap prediction. On the other hand, even without the strategy, the previous methods were still suboptimal in efficiency due to the latency of watershed and agglomeration. This observation suggests that the efficiency improvements our method are attributed to various factors.
Pred. Atfinity w/ Overlap
Pred Segmentation
Pred. Affinity w/o Ovedap
Pred Segmentation
Volume Imega
Groand Truth
w/ Ovarlap
w/o Overiap
Figure 8: Illustration of the border effects. The predictions of PEA with and without overlap were compared. As highlighted by the red arrows, without overlap, the predicted affinities were notably misaligned at the block borders, leading to poor segmentation results.
Table 7: Comparison between methods with and without overlap. / denotes using sliding-window with overlap, which is the default setting of previous methods. Inference times were tested with an NVIDIA 3090 GPU and 64 Intel Xeon Gold CPUs, which represent the time (in seconds) required process the full test set (AC3).
Metries
Inference Time
Overlap
VOlaphi
VOlmerge
VO1
Arand
Model
Watershed
Agg.
Total
ResUNet ResUNet
1.037
0.258
1.295
0.154
81.2
30.8
35.8
147.8
1.075
0.324
1.399
0.178
17.6
38.4
38.2
94.2
PEA PEA
0.852
0.232
1.084
0.094
60.2
37.1
25.4
122.7
1.085
0.583
1.668
0.218
11.2
42.7
27.5
81.4
AGQ (ours)
0.677
0.290
0.967
0.095
27.6
N/A
33.7
17

---

Published as a conference paper at ICLR 2025
D MORE DETAILS ON DATASETS
For AC3/AC4 dataset, following previous work (Huang et al. 2022b Luo et al. 2024 Arganda-
Carreras al. 2015), utilized the top 80 sections of AC4 the training set, the subsequent 20 sections as the validation set, and the top 100 sections of AC3 as the test set for the benchmark. Each voxel AC3/AC4 represents a physical resolution of x 6 nm The ZEBRAFINCH dataset contains 33 densely labeled volumes of zebra finch brain, 29 of which have shape of 150 x 150 x 150, and the other of which are 128 256 x 256. The physical resolution of each voxel is 20 x9 nm... We selected 30 of the volumes as the training set and the remaining 3 as the test set. The three test volumes are as follows, which covers both sizes:
Volume-1: id gt 23734 3 y4 - 4465 x2209-2353 shape 150 150 150; shape 128 256 256;
Volume-2: id 111a::01100000011111000013113333
Volume-3: id gt 22868-3012 y5744-5894 x5157-5307, shape 150 x 150 150.
E
MORE DETAILS ON METRICS
E VARIATION OF INFORMATION (VOI)
The VOI calculates the conditional entropy between the predicted segmentation Spred and the ground-truth segmentation SGT- Specifically, (S (Ser) reflects the amount of over- segmentation (VOlsplit). i.e., predicting one neuron multiple segments, while H (Sct)) Sprea) reflects the amount of under-segmmenattion (VOImerge). i.e. predicting multiple neurons as a single segment. The overall metric VOI is the sum of them,
VOI = VOIsplit + VOT merge
(SPrea ScT) H(ScT)SPred
(16)
E.2 ADAPTED RAND ERROR (ARAND)
The Arand is caleulated as:
Arand = 2 p
(17)
+
where Pi. denotes the probability that voxel is labeled as i in the predicted segmentation Spred and j in the ground-truie -ruuh segmentation SGr, & is the probability that voxel is labeled as k in Spred, while tę is the probability that a voxel is labeled as k in SGT- Voxels with id 0 in the SeT are ignored in the calculation. In other words, Arand is one minus the harmonic average precision and recall.
F MORE IMPLEMENTATION DETAILS
The learning rate followed the cosine schedule with a base learning rate of 0.0001. The total batch size is 8 (i.e., one volume image block per GPU). For AC3/AC4 dataset, each block had dimensions of 17 x 257 257. For ZEBRAFINCH dataset, each block had dimensions of 76 150 x 150. Our code was built on the wwtoorchcoonnett.oo
(Lin et al. 2021) codebase.
In the 3D U-Net backbone, we adopted 32, 64, 96, 128, 160 channels for its different stages. More- over, Group Normalization (Wu & He 2018) and ELU (Clevert et al., 2015) were utilized, follow- ing the preset configurations www.oocc.coonettt.oooo 64 channels for all hidden layers. The self-attentio block consisted of 8 multi-heads, 64 chan- nels, Layer Normalization, and GELU activation function (Vaswani et al. 2017). Additionally, we incorporated coordinate convolution (Liu al. 2018) enrich the image features.
cs. In the 3D neuron decoder, we used
Regarding the background, during training, we predict it from the image features using two convo- lutional layers, supervised by the same dice loss and cross-entropy loss Equation (12. During inference, we excluded these convolutional layers to eliminate the background prediction. This practice aligns with previous methods that do predict the background since our observations
18

---

Published as a conference paper at ICLR 2025
indicated that predicting the background typically causes decline in metrics. Moreover, during inference, neurons with a volume smaller than 40 voxels in the prediction are also excluded to filter out potential noises.
G LIMITATIONS AND FAILURE CASES
The representative failure cases depicted in Figure Region A showcases some redundant predictions (highlighted by the red arrows), which could be caused by multiple queries attempting predict the same neuron simultancously. The error observed in region B might result from breakdown in block assembly, where parts of the same neuron in adjacent blocks were not merged correctly, leading to disjointed segments. The error denoted by the red arrow in region C could be attributed to the ambiguity of the data, where an organelle was erroneously identified as a single neuron. To tackle these issues, potential avenues for future work could involve: (1) increasing query variability, for example, by introducing more training data or auxiliary tasks, (2) developing dedicated block assembly algorithm for our method, and (3) collecting purposeful data for training fine-tuning.
Volume Image
Ground Truth
Prediction
Figure 9:
Several representative failure cases. The first row shows a 2D slice. The second row presents its corresponding zoomed-in regions. The errors highlighted with the red arrows. Best viewed in digital.
H ANALYSIS
To delve deeper into the roles of various queries in prediction, we visualized the intermediate outputs of the 3D neuron decoder in Figure I0 Mcoarse: Po. Pt. P2 denote the coarse masks and results different decoder stages, respectively (refer to Section Each color denote a mask corresponding a specific query. Initially, Mcuarse exhibits numerous merging errors (highlighted by yellow arrows) and only predicts fraction of neurons. This is due to inherent eтrors in the predicted affinities (indicated by red arrows), such as the ambiguity the boundaries of myelin sheath. This also implies that relying solely on affinities traditional methods, could raise potential issues.
Instead, our method gradually resolved these issues. As depicted by the yellow arrows in Figure 10, these missing (or incorrectly merged) neurons progressively recovered through the integration of learnable queries. This indicates the distinct roles of the two query types: affinity-guides queries provide initial predictions, while learnable queries aid in retrieving incorrectly merged neurons.
We conducted a quantitative analysis of the segmentation results evolution in the decoder, as shown in Figure 11 Starting from coarse segmentation results, errors gradually rectified in the 3D neuron decoder. This trend aligns with the findings in Figure I00 and underscores the importance of multi-stage decoder design. Additionally, this highlights advantage of our method
19

---

Published as a conference paper at ICLR 2025
Valume Image
Ground Truth
Pred. Affinities
Moarse
Po
Visualization of the segmentation at different stages in the 3D neuron decoder. The images depict 2D sections, which inherently exist 3D space. Therefore, the masks in Mcoarse may interconnect in 3D. Incorrectly merged neurons in Mcoarse are highlighted by yellow arrows, whereas issues with the predicted affinities are indicated by red arrows. Best viewed in color.
Figure 10:
previous methods: method enables error correction through multiple stages, whereas previous methods were confined to merging results together (from watershed generedede fragments to the final results).
1.4
0.16
VO] (blook)
1.2
1.0
Arand
0.8
0.10
0.08
0.6
Meuse
(a) VOI of the intermediate results
(b) Arand of the intermediate results
Figure 1I: Benchmark of the segmentation results at different stages in the 3D Neuron Decoder. The vertical axis shows the average metrics all blocks.
I MORE ABLATION STUDY
1.1 GENERATION AFFINIY-GUIDEE QUERIES
Table g compares two methods for generating the affinity-guidert queries, i.e., using different al- gorithms to derive Mooarse from the predicted affinities. The watershed algorithm showed limited enhancement in segmentation results when utilized with the predicted affinities, possibly due to its fragmented generation of Mcoarse, which offers limited support to the 3D neuron decoder. In con- trast, the proposed method of thresholding the affinities and extracting the connected components led to the creation of suitable coarse masks, significantly improving the accuracy.
Table 8: Comparison of affinity- query generation methods.
Block
Full
VOT
Arand
VOT
VOTnerge
VOI
Arand
Watershed
0.706
0.087
1.122 0.781
0.581
1.703
0.173
Connected components
0.538
0.065
0.397
1.177
0.141
I.2 EFFECT Loss FUNCTIONS
We compared the effects of different loss functions within the proposed framework and reported the results in Table 9 The absence of feature loss notably deteriorated the segmentation results, the image features' inability capture neuronal specificity posed challenges for the decoder's
20

---

Published as a conference paper at ICLR 2025
learning process. The exclusion of affinity loss resulted in poor model performance, given that the unsupervised affinity branch could potentially misguide the decoder. These results underscore the importance of these loss functions in achieving accurate segmentation results.
Table 9: Effect of loss functions. Feat. and Aff. denote feature loss and affinity loss, respectively.
Block
Full
Feat.
Aff.
VOI
Arand
VOLsplit
VOlmerge
VO1
Arand
0.737
0.090
1.175
0.708
1.883
0.238
0.748
0,088
3.601
0.778
4.379
0.867
0.538
0.065
0.781
0.397
1.177
0.141
1.3 NUMBER OF STAGES
We explored in Table 10 whether increasing the decoder's stages would enhance the segmentation accuracy. The analysis indicated that more decoder stages did not substantially improve accuracy; instead, it reached a saturation point. This observation suggests that additional stages may exceed the model's capacity beyond the constraints of the available training data.
Table 10: Number of stages. #stages denotes the number of stages.
Block
Full
#Stages
VOI
Arand
VOlspit
VOlmerge
VOI
Arand
2
0.538
0.065
0.781 0.858
0.397 0.367
1.177
0.141
4
0.569
0.067
1.225
0.138
J VISUALIZATION COMPARISON OF AXON
Figure 12 illustrates comparison of the reconstruction results for axon. Our method demon- strated superior reconstruction quality when compared with previous methods. As indicated by the red arrows, the previous method lost certain details, potentially attributed to the common encase- ment axons in myelin sheaths and it is more difficult to predict affinities for myelin sheaths. Our method excels in reconstructing the axon, suggesting its efficacy in addressing this particular challenge.
Grand/Tuth
RusUNNet
SuprrRarrmm
AGO (Oums
Figure 12: Visualization comparison of axon reconstructed by different methods. Red arrows indi- cate errors. Best viewed in digital with zoom-in.
21


Reviews:
[{"review_id": "1P6HRKGuDV", "reviewer": "Reviewer_uA9h", "summary": "The paper presents a novel method for neuron segmentation in electron microscopy (EM) images using affinity-guided queries within a query-based framework. The approach aims to overcome the limitations of traditional methods that rely on watershed algorithms by directly predicting segmentation results, achieving improved accuracy and efficiency. The method is evaluated on benchmark datasets, showing competitive performance.", "questions": "1. How does the method handle variations in image quality or noise levels in EM datasets?\n2. What are the specific computational requirements for training and inference, and how do these scale with larger volumes?\n3. How do affinity-guided queries compare with other forms of guided queries in terms of performance and efficiency?\n4. The paper mentions that the proposed method does not require overlap prediction to address border effects. However, it is unclear how the method handles the continuity and consistency of segmentation results across block boundaries. Can the authors elaborate on this aspect and provide visual examples?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The use of affinity-guided queries in a query-based framework is a novel contribution to the field of neuron segmentation.\n2. The method demonstrates a significant speedup (2-3x) over traditional methods, a notable advantage for large-scale data processing.\n3. The query-based framework directly predicts the final segmentation results, avoiding the inaccuracy and inefficiency associated with the watershed algorithm used in traditional methods.", "weaknesses": "1. The comparison with state-of-the-art methods is limited to a few selected approaches. A more extensive comparison with a broader range of recent methods would strengthen the paper's contributions, such as [1]. \n2. The paper does not provide a strong theoretical foundation for why affinity-guided queries outperform other query-based methods, especially in terms of handling neuron morphology.  It is better to give a more detailed explanation or analysis of the mechanisms by which affinity-guided queries better capture neuron morphology compared to other query-based approaches.\n\n[1] Zhang, Yanchao, et al. \"SegNeuron: 3D Neuron Instance Segmentation in Any EM Volume with a Generalist Model.\" International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mH0yGfVLOj", "reviewer": "Reviewer_nQ34", "summary": "This paper introduces Affinity-Guided Queries (AGQ) as a way to apply query-based instance segmentation to the field of connectomics. AGQ predicts the segments directly, bypassing postprocessing steps such as watershed used in some other approaches. The method is evaluated on two small-scale datasets, showing good results in terms of efficiency and accuracy.", "questions": "- Please include VOI_split and VOI_merge breakdown in the main text instead of only in the appendices.\n- The intro, abstract and Fig. 1 make it sound as if all currently used neuron reconstruction methods rely on watershed. This is not true (the authors are aware of at least https://doi.org/10.1038/s41592-018-0049-4, which they cite later; there also exist other approaches such as https://openaccess.thecvf.com/content_CVPR_2019/html/Meirovitch_Cross-Classification_Clustering_An_Efficient_Multi-Object_Tracking_Technique_for_3-D_Instance_CVPR_2019_paper.html) and should be more clearly discussed earlier in the paper. The proposed model is far from from the first approach that directly converts 3d EM images into a segmentation.\n- 0.08 mm^3 is a fairly specific number to quote as a \"typical size of studied volume\", and quotes a 6y old paper. It would be best to cite more works here, or just refrain from making statements about what is typical altogether. I would also suggest reformulating the comment about processing requiring \"130 days of runtime\", as this is highly specific to your specific cluster/hardware configuration.\n- How often does your algorithm generate disconnected components for the same predicted object? (Eq. 3 does not guarantee connectivity)\n- What's the motivation for predicting affinities and then immediately averaging them in Eq. 7? Why not predict a boundary map?\n- Does N_a in Eq. 8 vary between passes through the network? If so, how does that interact the number of learned queries -- is N_l or N kept constant? How is N_l chosen in your work?\n- Your method does not require block overlaps to generate consistent results. But wouldn't the additional spatial context help in making the predictions better? Have you explored this?\n- What are the error bars for your results in Tab. 1?\n- How does your network capacity (parameter count, FLOPS) compare with your baselines?\n- Your examples in Fig. 5 focus on dendrites with broken off spines. Could you please also provide a representative figure for axons? (perhaps in the supplement)", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Evaluation on two datasets acquired with different EM techniques.\n- Novel approach to segmentation of 3d EM data (query-based model).\n- Good results in terms of efficiency and promising accuracy (but see comments about metrics below).\n- Ablations showing the impact of the various modules.\n- The paper is clearly written and well-organized.", "weaknesses": "- Evaluation is performed on a very small scale (tens of um^3). Experience in the field of connectomics shows that such results are often not predictive of reconstruction quality over larger volumes. One of the datasets (ZEBRAFINCH) provides full volume (1M um^3 -scale) skeleton tracings that make it possible to compute topological metrics, which can then be compared to those reported https://doi.org/10.1038/s41592-022-01711-z and https://doi.org/10.1038/s41592-018-0049-4. Having these available, would make it much clearer how the proposed method compares to the state of the art in a practical setting.\n- VOI advantage seems to be driven by VOI_split reduction at the cost of VOI_merge. If the mergers happen at the level of supervoxels, the results are in practice less useful for proofreading, where many systems do not make it possible to manually fix this type of error.\n- The model is quite architecturally complex, with different losses, modules, and an explicit affinity modeling branch.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2rLAyBLxMg", "reviewer": "Reviewer_pZEH", "summary": "This paper presents an approach for neuron segmentation in electron microscopy images by a query-based segmentation framework. Traditional methods often rely on predicting affinities followed by watershed algorithms, which can introduce artifacts and inefficiencies. The proposed method predicts affinities using a lightweight branch to obtain coarse neuron structure information, which is then used to construct Affinity-Guided Queries (AGQs). These queries, along with learnable queries, interact with image features to directly predict the final segmentation results, eliminating the need for watershed algorithms. The method achieves better accuracy and a 2–3× speedup in inference time.", "questions": "1. There are, in fact, parallel/CUDA implementations of watershed algorithms. How would the proposed method compare to these optimized implementations in terms of efficiency?\n2. For the Transformer Decoder implementation in Table 3, how do you input the image features and queries?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. Novel and intersting idea that uses a query-based framework to eliminate the need for watershed algorithms used in traditional methods. The proposed method greatly improves efficiency.\n2. Very well written manuscript. The paper clearly articulates the shortcomings of existing methods and motivates the need for the query based approach.\n3. Comprehensive ablation studies show the contributions of different modules.", "weaknesses": "1. Despite the successful and interesting use of AGQs, I have difficulties to see why similar results couldn't be achieved by inputting image features, predicted affinities, and learnable queries directly into the neuron decoder (with modified architecture). While I understand that predicted affinities are used to create AGQ, I wonder if the LQ could \"attend\" to the predicted affinity map to construct the necessary queries themselves. It appears that the current neuron decoder is attempting to \"rescue\" the low-quality segmentation by AGQ in $M_{coarse}$. \n2. The choice of $a$ in ConnectedComponent algorithm in eq 8 is not reported and discussed.\n3. The computational efficiency is evaluated in terms of inference latency. It would be great if the manuscript could educate the readers about the model's memory consumption, which is important for deployment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "55oVV8CwdG", "reviewer": "Reviewer_wCU2", "summary": "The paper presents an innovative approach to neuron segmentation in electron microscopy (EM) images, a critical task in connectomics research that aims to reconstruct neural connections in the brain. The authors propose an affinity-guided query (AGQ) method, which introduces a lightweight query-based framework to address the limitations of traditional clustering algorithms.", "questions": "Comparison with State-of-the-Art Baselines\n\nClarification on the Use of Agglomerations\n\nUse of Well-Established Benchmarks\n\n\nOriginality of Affinity-Guided Queries", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors introduce AGQ to incorporate coarse neuron structure information into the query generation process, reducing the learning difficulty and improving segmentation accuracy.", "weaknesses": "While the work is commendable and addresses significant challenges in the field, there are areas where further development and clarification could enhance the contribution and impact of the research. Here are some specific points for consideration:\n\n1. The paper could benefit from a more comprehensive comparison with the latest state-of-the-art (SOTA) methods, such as those referenced in [1][2]. Including these comparisons will provide a clearer picture of how the AGQ method stacks up against the most recent advancements in the field.\n2. Figure 1 (b) suggests that the proposed method relies on agglomeration operations from 'waterz' post-processing. This seems to contradict the claim of directly predicting segmentation results. The authors should clarify whether the proposed method truly bypasses the need for watershed-based post-processing or if it still relies on some aspects of it.\n3. The experiments presented in Table 2, while informative, may not be based on benchmarks that are widely recognized in the community. Utilizing more established benchmarks would strengthen the paper's findings and allow for better comparisons with other methods.\n4. The motivation behind the proposed affinity-guided queries appears to have similarities with modules found in other works, such as [3]. The authors should elaborate on the novel aspects of their approach and how it differs from or improves upon existing techniques.\n5. 'typically the neuron gap' in line168 is not unprofessional\n\n[1]Sheridan A, Nguyen T M, Deb D, et al. Local shape descriptors for neuron segmentation[J]. Nature methods, 2023, 20(2): 295-303.\n[2]Liu X, Cai M, Chen Y, et al. Cross-dimension affinity distillation for 3d em neuron segmentation[C]//2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE Computer Society, 2024: 11104-11113.\n[3]Sun R, Luo N, Pan Y, et al. Appearance Prompt Vision Transformer for Connectome Reconstruction[C]//IJCAI. 2023: 1423-1431.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)