trAIce3D: Prompt-Driven 3D Microglia Segmentation

• trAIce3D is a prompt-driven 3D segmentation framework that couples a transformer-enhanced U-Net with a two-stage pipeline to isolate microglial somas and branches in large-scale fluorescence microscopy.
• The framework employs novel Residual Cross-Attention Modules that fuse image features with prompt-encoded soma coordinates, effectively disambiguating overlapping cellular structures.
• Trained on over 41,000 annotated microglial cells across diverse brain states, trAIce3D demonstrates superior accuracy and scalability for quantitative neurodegenerative research.

trAIce3D is a prompt-driven, transformer-enhanced 3D U-Net architecture for instance-level segmentation of microglial cells from large-scale 3D fluorescence microscopy, designed to segment both somas and fine branching processes across dense, noisy, and morphologically complex image volumes (Alamalhoda et al., 30 Jul 2025). It is organized as a two-stage pipeline in which soma coordinates detected in a first stage become prompts for cell-specific branch segmentation in a second stage, thereby addressing the central difficulty of separating overlapping tree-like cellular structures in volumetric data. The reported system was trained and evaluated on 41,230 microglial cells from 230 mouse brain images spanning development, adulthood, and neurodegeneration, with the stated aim of enabling scalable analysis of complex cellular morphologies (Alamalhoda et al., 30 Jul 2025).

1. Definition, scope, and nomenclature

In the supplied literature, the exact name trAIce3D refers to the 2025 model titled “trAIce3D: A Prompt-Driven Transformer Based U-Net for Semantic Segmentation of Microglial Cells from Large-Scale 3D Microscopy Images” (Alamalhoda et al., 30 Jul 2025). Its task domain is 3D microscopy rather than robotics or dynamic scene reconstruction. Concretely, it targets instance-level segmentation of microglia, including both soma compartments and thin, tortuous processes, from large volumetric fluorescence images.

The biological motivation is explicit. Microglial morphology changes with functional state and pathology, including normal immune surveillance, response to injury or altered homeostasis, and progression of neurodegenerative diseases such as Alzheimer’s disease. Quantitative 3D morphology is therefore treated as a potential biomarker, but segmentation is difficult because the input volumes have large fields of view, dense cell populations, noisy and variable intensity, and thin, overlapping branches that cross in 3D (Alamalhoda et al., 30 Jul 2025).

A common source of confusion is that the supplied material also includes TRADE, a UAV target-tracking system (Proença et al., 2022), and TRACE, a 3D Gaussian dynamics model for multi-view videos (Li et al., 13 Aug 2025). These are distinct systems with different application domains and architectures. This suggests that the term “trAIce3D” should be understood here in the neuroimaging sense defined by the 2025 arXiv paper, not as a synonym for those other frameworks.

2. Problem formulation and two-stage pipeline

trAIce3D is explicitly formulated as a two-stage framework (Alamalhoda et al., 30 Jul 2025). The first stage performs Soma Segmentation over large 3D volumes, and the second performs Prompt-based Branch Segmentation for each detected cell.

In Stage 1, the full 3D microscopy volume is processed using overlapping tiles of size $256 \times 256 \times 16$. A 3D U-Net with a hierarchical 3D Vision Transformer encoder and a CNN decoder outputs a binary soma probability map. This map is then used to locate soma centroids. The sliding-window formulation is essential because the original volumes have median size approximately $[2304, 2560, 56]$ voxels, which is too large for direct dense processing (Alamalhoda et al., 30 Jul 2025).

In Stage 2, each detected soma coordinate serves as a prompt. A local 3D window of size $256 \times 256 \times 16$ is extracted around the target soma, and the same encoder-decoder backbone is reused, but now with prompt-conditioned skip connections. The output is a voxelwise segmentation of that cell’s soma and branches within the local window. The paper summarizes the division of labor succinctly: Stage 1 determines where the cells are, while Stage 2 segments “this cell and its processes” (Alamalhoda et al., 30 Jul 2025).

This design directly addresses overlap ambiguity. A local crop alone does not specify which cell inside a crowded neighborhood should be segmented. By contrast, a soma coordinate prompt conditions the model on a specific target instance. A plausible implication is that the architecture separates localization and assignment: Stage 1 provides coarse object proposals, while Stage 2 performs prompt-conditioned instance disambiguation.

3. Architectural design

The common backbone is a 3D ViT–U-Net hybrid. The encoder is a hierarchical 3D Vision Transformer, and the decoder is a CNN-based U-Net-style decoder (Alamalhoda et al., 30 Jul 2025).

The input tile of size $256 \times 256 \times 16$ is partitioned into non-overlapping 3D patches of size $8 \times 8 \times 2$. A 3D convolution maps each patch to an embedding whose dimension depends on model scale: 32 for trAIce3D_S, 64 for trAIce3D_M, and 128 for trAIce3D_L. Sinusoidal positional embeddings are added in 3D. The encoder contains 6 sequential transformer blocks, each using 8-head multi-head self-attention with FlashAttention, an MLP with hidden dimension 8 times the embedding size, GeLU activation, layer normalization, and residual connections (Alamalhoda et al., 30 Jul 2025). Attention is defined as

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V .$$

A hierarchical representation is formed through patch merging after transformer blocks 1, 2, and 3. Each merge reduces spatial resolution by a factor of 2 in each dimension and doubles the embedding dimension. The overall spatial reduction from input to deepest level is $8^3 = 512$, and the encoder emits four hierarchical feature maps for the decoder (Alamalhoda et al., 30 Jul 2025).

The decoder upsamples using transposed 3D convolutions, concatenates with same-scale skip features, and applies residual convolutional blocks with $3 \times 3 \times 3$ kernels, padding 1, batch normalization, GeLU, and dropout. Additional 3D convolutions and transposed convolutions followed by a sigmoid generate the voxelwise output probability map (Alamalhoda et al., 30 Jul 2025).

The principal architectural distinction between stages lies in the skip connections. In Stage 1, skip connections are plain concatenations of encoder features. In Stage 2, they are replaced by Residual Cross-Attention Modules (RCAM) that fuse image features $I$ and prompt embeddings $P$:

$[2304, 2560, 56]$0

The reported interpretation is that prompt self-attention refines prompt features, image-to-prompt cross-attention injects contextual image information into the prompt representation, and prompt-to-image cross-attention updates image features so that structures belonging to the prompted soma become more salient (Alamalhoda et al., 30 Jul 2025). This is the mechanism by which the model disambiguates overlapping branches from neighboring cells.

4. Prompt encoding and prompt-driven segmentation

The prompt in trAIce3D is not linguistic. It is a 3D soma coordinate in pixel space, normalized within the local tile as

$[2304, 2560, 56]$1

Prompt encoding uses a Fourier feature mapping with a learned Gaussian matrix $[2304, 2560, 56]$2. The encoded prompt is built as

$[2304, 2560, 56]$3

$[2304, 2560, 56]$4

$[2304, 2560, 56]$5

with final prompt embedding matrix

$[2304, 2560, 56]$6

These embeddings are injected into RCAMs at each scale (Alamalhoda et al., 30 Jul 2025).

The supplied description emphasizes that ROI cropping and prompt conditioning are not equivalent. Cropping only defines a local field of view, whereas the prompt explicitly identifies which soma within that window should anchor the segmentation. The model is reported to perform significantly worse in branch segmentation when the same backbone is used without prompt-based RCAM skip connections, especially for thin and overlapping structures (Alamalhoda et al., 30 Jul 2025). This suggests that the decisive contribution is not merely local context restriction but conditional feature routing through cross-attention.

The authors compare this prompt formulation to point prompts in interactive segmentation systems, but note that the process here is fully automatic because the prompts come from Stage 1 rather than from a human annotator (Alamalhoda et al., 30 Jul 2025). For volumetric neuroimaging, this converts an interactive segmentation primitive into an instance-level inference mechanism.

5. Training regime, dataset, and preprocessing

trAIce3D was trained on the morphOMICs dataset, attributed to Colombo et al., Nature Neuroscience 2022, comprising 41,230 annotated microglial cells from 230 mouse brain images across development, adulthood, and neurodegeneration (Alamalhoda et al., 30 Jul 2025). Voxel size was standardized to $[2304, 2560, 56]$7 per pixel. Microglial processes were traced in 3D using the Imaris filament-tracing plugin (v9.2), followed by manual verification; these traces were rasterized into voxelwise labels for somas and branches (Alamalhoda et al., 30 Jul 2025).

Volumes were divided into $[2304, 2560, 56]$8 cubes. Cubes with less than 5% microglia volume or discontinuous labels were excluded. Per-cube preprocessing consisted of histogram equalization and normalization to range $[2304, 2560, 56]$9. Sampling differed by stage: soma segmentation used overlapping cubes from whole volumes, whereas branch segmentation used cubes centered around individual microglia (Alamalhoda et al., 30 Jul 2025).

Training proceeded in two phases. First, Stage 1 learned soma segmentation. Second, Stage 2 initialized encoder and decoder weights from the trained Stage-1 model and added the Prompt Encoder and RCAM modules for branch segmentation (Alamalhoda et al., 30 Jul 2025). The paper states that encoder and decoder weights are initialized from Stage 1 and does not explicitly state that they are frozen; the supplied summary therefore interprets them as being fine-tuned jointly.

Loss functions were stage-specific. For soma segmentation,

$256 \times 256 \times 16$0

with Focal Loss

$256 \times 256 \times 16$1

using $256 \times 256 \times 16$2 and $256 \times 256 \times 16$3 (Alamalhoda et al., 30 Jul 2025). For branch segmentation,

$256 \times 256 \times 16$4

where

$256 \times 256 \times 16$5

The use of clDice is explicitly motivated as topology preservation for tubular structures such as microglial processes (Alamalhoda et al., 30 Jul 2025).

Optimization used 200 epochs for each stage, batch size 64 with gradient accumulation 2, learning rate $256 \times 256 \times 16$6 with cosine decay, 50 warmup epochs, and final learning rate $256 \times 256 \times 16$7. Training was performed in PyTorch 2.6.0 with CUDA 12.6 on two NVIDIA A100 GPUs (80 GB) using BF16 precision. Data augmentation was implemented with TorchIO v0.20.4 and included RandomAffine, RandomFlip, RandomElasticDeform, RandomNoise, RandomBlur, and RandomGamma (Alamalhoda et al., 30 Jul 2025).

6. Quantitative performance and computational profile

The reported experiments separate soma segmentation from branch segmentation. For soma segmentation, the data split was 80/15/5 on soma patches, and baselines were CellPose (cyto3) and nnU-Net (3D Cascade Full Resolution) (Alamalhoda et al., 30 Jul 2025). The three trAIce3D scales differ in embedding dimension and parameter count.

Model	Accuracy	F1 Score
CellPose	34.9%	39.0%
nnU-Net	60.5%	64.0%
trAIce3D_S	51.0%	56.8%
trAIce3D_M	78.5%	81.6%
trAIce3D_L	83.1%	87.5%

The large model, trAIce3D_L, achieved 83.1% Accuracy, 87.5% F1 Score, and 84.3% Precision, outperforming the listed baselines on this task (Alamalhoda et al., 30 Jul 2025). The supplied analysis attributes this to the hierarchical ViT encoder, larger embedding dimension, and extensive training on a large and diverse dataset.

For branch segmentation, no strong public baselines were reported, and evaluation focused on trAIce3D variants alone. Metrics were Dice Score (DS), Average Path Length Difference (APLD), and Hausdorff Distance (HD) (Alamalhoda et al., 30 Jul 2025).

Model	DS	APLD	HD
trAIce3D_S	0.31	2.90	16.53
trAIce3D_M	0.47	1.04	9.22
trAIce3D_L	0.63	0.77	5.10

These values indicate a strong scale effect: the larger the embedding dimension and total capacity, the better the preservation of thin branches and boundary accuracy (Alamalhoda et al., 30 Jul 2025).

The 5-fold cross-validation results for trAIce3D_L further quantify stability. Soma segmentation achieved Accuracy $256 \times 256 \times 16$8, F1 $256 \times 256 \times 16$9, and Precision $256 \times 256 \times 16$0. Branch segmentation achieved Dice $256 \times 256 \times 16$1, APLD $256 \times 256 \times 16$2, and HD $256 \times 256 \times 16$3 (Alamalhoda et al., 30 Jul 2025). The supplied interpretation is that this reflects generalization across different images and brain regions.

The computational profile is also explicitly reported. Parameter counts are 6.10 M for trAIce3D_S, 27.38 M for trAIce3D_M, and 109.39 M for trAIce3D_L. On an NVIDIA A100 with tile size $256 \times 256 \times 16$4, inference requires approximately 0.5 GB / 2 GB / 6 GB GPU memory and processes approximately 90 / 50 / 30 volumes per minute for S/M/L respectively (Alamalhoda et al., 30 Jul 2025). This establishes a practical throughput-quality tradeoff.

7. Interpretation, limitations, and relation to adjacent systems

The key distinguishing feature of trAIce3D is the combination of a prompt-driven design, a transformer-enhanced 3D U-Net, two-stage training with model reuse, and topology-aware branch supervision (Alamalhoda et al., 30 Jul 2025). Its intended use is large-scale quantification of microglial morphology across brain regions, sexes, and disease states, reducing manual labor and operator bias relative to semi-manual pipelines such as Imaris-based tracing (Alamalhoda et al., 30 Jul 2025).

The supplied paper explicitly states that the architecture is generic enough to extend beyond microglia to other complex branching cell types, including neurons, astrocytes, and vascular networks (Alamalhoda et al., 30 Jul 2025). This suggests a broader design pattern: prompt-conditioned volumetric instance segmentation may be effective whenever morphology is tree-like, crowded, and topologically fragile.

Several limitations are stated. Performance may drop under markedly different imaging conditions or microscopy modalities, so fine-tuning is likely needed. The fixed $256 \times 256 \times 16$5 window may be insufficient for much larger cells such as neurons spanning longer distances. Very dense microglia populations with low SNR remain difficult even with prompts. Finally, the training pipeline depends on high-quality 3D skeletons and masks from Imaris plus manual verification, which is time-consuming to generate for new tissue types or modalities (Alamalhoda et al., 30 Jul 2025).

Future directions named in the paper include Shifted Window Attention, hybrid ViT–CNN architectures such as SwinUNet3D, learned window-shifting strategies, and improved robustness to domain shifts (Alamalhoda et al., 30 Jul 2025). These are framed as architectural extensions rather than as demonstrated results.

The adjacent systems included in the supplied material clarify what trAIce3D is not. TRADE couples 2D tracking with 3D motion and ground geometry modeling for UAV target following (Proença et al., 2022), while TRACE learns per-particle translation-rotation dynamics from multi-view RGB videos using 3D Gaussian representations (Li et al., 13 Aug 2025). Both are 3D-aware systems, but neither addresses microglial segmentation. Their presence nevertheless highlights a broader methodological trend: explicit structural priors—whether motion models, physical dynamics, or prompt-conditioned attention—are being used to overcome ambiguities that purely local image-domain inference handles poorly. In trAIce3D, that prior is the soma prompt and the RCAM-based conditioning of volumetric features (Alamalhoda et al., 30 Jul 2025).