Joint-OCTAMamba: Two-Stage OCTA Segmentation
- Joint-OCTAMamba is a two-stage OCTA segmentation network that jointly and decoupledly processes retinal vessel and FAZ segmentation.
- The framework utilizes architectural decoupling and explicit loss reweighting to address task imbalance between fine vessel extraction and precise FAZ boundary localization.
- Key modules such as HDFE, VMAF, and CFEB enhance feature representation, yielding state-of-the-art performance on OCTA-500 with fewer parameters than traditional CNNs and Transformers.
Searching arXiv for the named papers and closely related work to ground the article in current preprints. arXiv search query: "Joint-OCTAMamba OCTA segmentation Mamba" Joint-OCTAMamba is an OCTA joint segmentation network based on feature enhanced mamba for the simultaneous segmentation of retinal vessels (RV) and the foveal avascular zone (FAZ) from Optical Coherence Tomography Angiography (OCTA). Its defining characteristic is a unified but decoupled two-stage cascade: Stage 1 produces an RV probability map, and Stage 2 concatenates that map with the original OCTA image, applies ROI center-cropping, and segments the FAZ. The framework is motivated by the coexistence of fine, tortuous vessels, OCTA noise and artifacts, and the small annular FAZ boundary, and it addresses the task imbalance reported for earlier joint models through architectural decoupling and explicit loss reweighting (Liu et al., 15 Sep 2025).
1. Clinical scope and problem setting
OCTA visualizes retinal microvasculature and the FAZ, both of which are important in diagnosing and monitoring diabetic retinopathy, glaucoma, and AMD. In this setting, the target is not only pixelwise vessel extraction but also reliable delineation of the FAZ, whose compact, hollow-center geometry differs markedly from the dense, filamentary morphology of the vascular tree.
The technical difficulty arises from two different segmentation regimes coexisting in the same image. RV segmentation must preserve fine, tortuous vessels in the presence of noise and artifacts. FAZ segmentation, by contrast, depends on accurate boundary localization for a small annular structure and therefore requires both global context and boundary precision. Prior 2D CNN-based methods, including U-Net-style systems, are described as limited by insufficient long-range dependency modeling, while existing joint segmentation models are reported to exhibit performance imbalance between tasks or to rely on heavy 3D processing.
Within this problem formulation, Joint-OCTAMamba treats joint OCTA analysis as a coupled but non-symmetric problem. This indicates that “joint” does not denote a single undifferentiated decoder; rather, it denotes a coordinated multi-task system in which RV estimation explicitly informs FAZ estimation.
2. Two-stage cascade and network organization
Joint-OCTAMamba is organized as a cascaded two-stage network composed of RVMamba and FAZMamba (Liu et al., 15 Sep 2025). Stage 1 takes an OCTA image and outputs an RV probability map, denoted RVMap. Stage 2 concatenates the original OCTA image with RVMap channel-wise, center-crops a region of interest of size , and outputs the FAZ mask. Both subnetworks follow an encoder–bottleneck–decoder U-Net skeleton, but they differ in the task-specific enhancement modules inserted into the bottleneck and gating path.
The shared structural template is consistent across both stages. A Compact Multi-Branch Fusion (CMBF) input preprocessor expands features to 32 channels, splits them into four branches, and fuses them with a convolution. The encoder uses three down-sampling blocks. The bottleneck contains a task-specific Mamba block. The decoder uses three up-sampling blocks with Attentional Feature Fusion and a final conv head, FnConv, for logits.
| Component | Input and output | Distinctive modules |
|---|---|---|
| Stage-1 RVMamba | OCTA RVMap | CMBF, HDFE, VMAF, V-SS2D |
| Stage-2 FAZMamba | OCTA + RVMap, ROI FAZ mask | CMBF, HDFE, CFEB, V-SS2D |
This arrangement is central to the framework’s handling of task imbalance. FAZ segmentation is not learned as a competing output head over the same representation used for vessels; instead, it is learned after vessel estimation, with the RVMap acting as an explicit prior. The paper characterizes this as a decoupled cascade that improves FAZ delineation, leverages RV priors to reduce ambiguity, and mitigates joint-task imbalance via architectural decoupling and explicit loss reweighting.
3. Feature-enhanced Mamba components
The phrase “feature enhanced Mamba” refers to the placement of strong, task-specific feature extractors ahead of and within the Mamba state-space blocks, together with task-aware gates that modulate the selective scan. The key modules are HDFE, VMAF, and CFEB.
The Hybrid Directional Feature Extractor (HDFE) is used in both branches ahead of the Mamba core. For input , it first applies convolution, PReLU, and GhostModule to produce a compact but expressive base representation. It then splits into two streams. The direction-aware stream uses dual Dynamic Snake Convolutions with and kernels to encode anisotropy and orientation, followed by inter-branch feature interaction. The multi-scale stream uses dilated convolutions with rates and a Feature Pyramid Enhance module. The outputs are adaptively fused, re-weighted by an ECA attention layer, and residual-added. In RV segmentation this supports topology-aware vasculature mapping; in FAZ segmentation it provides orientation- and scale-aware features for annular boundary modeling.
Vessel Multi-Attention Fusion (VMAF) appears only in RVMamba. It consists of three coordinated attention mechanisms. VCAB performs channel attention using global avg/max pooling and a local-global pooling path, followed by a lightweight conv stack and Sigmoid. VSAB performs spatial attention using avg/max pooled maps, directional horizontal/vertical convolutions, and diagonal convolutions obtained via rotation. VSTAB predicts centerline, bifurcation, and width-variation maps through dedicated or dilated convolutions and Sigmoid. These three maps are fused by learnable weights, refined by an EnhanceModule, and residually added. The result is a vessel-aware, topology-sensitive gating signal for the Mamba selective scan.
Compact FAZ Enhancement Block (CFEB) replaces VMAF in FAZMamba. CFEB uses a 0 grouped convolution with dilation 1 to capture annular FAZ context, followed by SE recalibration. An optional circular mask centered at 2 with radius 3 amplifies foveal features. The output is fused with the input through a residual with scale 4. The paper explicitly notes that VMAF hurts FAZ because it biases the representation toward filamentary structures; CFEB is introduced to inject FAZ-centric priors instead.
The task-specific Mamba blocks inherit this distinction. RVMambaBlock is described as HDFE 5 V-SS2D, with the gate path routed through VMAF. FAZMambaBlock keeps the same macro-pipeline but replaces VMAF with CFEB. In both cases, the main path applies 2D depthwise convolution for local mixing followed by selective scan for long-range spatial-state modeling, and fusion uses a SiLU-gated multiplicative interaction between the main-path output and the gate, followed by residual stabilization and an appended LSA module.
4. State-space formulation, joint objective, and imbalance mitigation
Joint-OCTAMamba adopts the Mamba state-space model as a linear-time alternative to quadratic self-attention (Liu et al., 15 Sep 2025). The continuous-time formulation is
6
where 7 is the input, 8 the hidden state, and 9 the output. After discretization with step 0,
1
with 2 and 3 obtained from 4, for example by zero-order hold, including the common expressions 5 and 6.
In the 2D setting, the paper’s V-SS2D applies local 2D depthwise convolution to produce 7, then performs a global scan over spatially flattened features in 8 with 9, or through separable scanning along rows and columns. Task-aware gating is generated by VMAF for RV or CFEB for FAZ. A representative formulation is
0
where 1 is a nonlinearity such as SiLU and 2 denotes element-wise multiplication. This couples global spatial-state modeling with explicit morphological priors.
The loss design is explicitly asymmetric across tasks. Dice loss is used for both tasks,
3
For RV, Tversky loss and a differentiable Hausdorff approximation are added, together with a distance-transform-based boundary loss. The branch-wise losses are
4
5
The total objective is
6
Static task weights are set to 7, 8 for 3 mm scans and 9, 0 for 6 mm scans. This is the paper’s direct mechanism for compensating for FAZ’s smaller target size and different convergence behavior. Together with the decoupled cascade and RV-prior injection, it forms the main answer to the performance imbalance problem.
5. Data, training protocol, quantitative performance, and ablations
The reported experiments use the official 3 mm and 6 mm splits of OCTA-500, with both RV and FAZ annotations (Liu et al., 15 Sep 2025). Preprocessing differs by task: full-resolution OCTA inputs are used for RV, while FAZ uses a center-cropped ROI of 1 in Stage 2. Test-time predictions are ensembled by TTA over the original, horizontally flipped, and vertically flipped inputs. Data augmentation is implemented with Albumentations at 2, including random brightness/contrast shifts, CLAHE, 3 rotation, horizontal/vertical flips, and piecewise affine transforms applied consistently to images and masks.
Optimization uses AdamW with OneCycleLR, maximum learning rate 4, initial learning rate 5, and a 10-epoch warm-up followed by cosine decay. Training runs for 100 epochs on a single NVIDIA V100 32GB with batch size 2 and fixed seed 0. During the first 70% of epochs the model trains without evaluation; during the final 30%, validation and checkpoint saving occur every 5 epochs based on best RV-Dice, FAZ-Dice, and their average. A single optimizer updates the end-to-end cascade.
Model complexity is reported as 13.22M trainable parameters, compared with U-Net at 118.42M, FARGO at 328.14M, Joint-Seg at 154.16M, IPN-V2 at 195.64M, and Joint-OCTAFormer at 22.88M. The paper attributes the efficiency to linear-time SSM global modeling combined with task-specific enhancement blocks rather than large Transformer stacks.
The principal quantitative results on OCTA-500 are as follows.
| Task and scan size | Dice | IoU |
|---|---|---|
| RV, 3 mm | 6 | 7 |
| RV, 6 mm | 8 | 9 |
| FAZ, 3 mm | 0 | 1 |
| FAZ, 6 mm | 2 | 3 |
These numbers are described as top or near-top across tasks and scales. The paper further states that, versus prior OCTA Mamba for vessels, Joint-OCTAMamba markedly improves RV Dice/IoU while adding state-of-the-art FAZ performance, and that against CNN-, Transformer-, and joint-segmentation baselines such as U-Net, FARGO, Joint-Seg, OCT2Former, SAM-OCTA/SAM-OCTA2, and PL-Joint-Seg, it matches or surpasses the best reported numbers with far fewer parameters.
Ablation studies identify the functional roles of the principal modules. In RVMamba, HDFE, VMAF, and CMBF each improve RV and tend to improve FAZ when used alone; the combination of all three yields the best RV Dice/IoU, reported as 4. In Stage 2, CFEB, ROI cropping, and the RV prior each improve FAZ, and their combination yields the best FAZ Dice/IoU, reported as 5. The negative effect of VMAF on FAZ is treated as evidence that vessel-specific gating is not transferable to annular-region segmentation.
The codebase is available at https://github.com/lc-sfis/Joint-OCTAMamba, with the typical workflow consisting of dependency installation, OCTA-500 placement in the prescribed directory structure, configuration of 3 mm or 6 mm training, end-to-end training with a single optimizer, and inference with TTA.
6. Relation to OCTAMamba, limitations, and nomenclature
Joint-OCTAMamba is closely related to the earlier single-task OCTA vessel segmenter "OCTAMamba: A State-Space Model Approach for Precision OCTA Vasculature Segmentation" (Zou et al., 2024). That precursor addressed vasculature segmentation alone through a U-shaped encoder–decoder augmented by Mamba-based visual state-space modeling, QSEME for local feature enrichment, MSDAM for multi-scale vessel processing, and FFRM for noise suppression and target highlighting. It reported 84.50 Dice, 73.23 IoU, and 84.00 Sensitivity on OCTA_3M; 82.31 Dice, 70.03 IoU, and 82.75 Sensitivity on OCTA_6M; and 90.04 Dice, 82.03 IoU, and 88.86 Sensitivity on ROSSA, with 3.57M parameters. Its stated future directions included joint modeling of FAZ and retinal layers and OCT/OCTA fusion. Joint-OCTAMamba can therefore be understood as a later OCTA-specific generalization from vessel-only state-space segmentation to decoupled joint RV–FAZ segmentation.
The paper also states clear limitations. The two-stage cascade introduces design complexity and depends on a high-quality RV prior and accurate ROI selection. Off-center FAZs or severe artifacts may reduce the gains from the foveal mask and ROI heuristic. The loss weighting parameter 6 is dataset-dependent and requires tuning. Evaluation is centered on Dice and IoU, and future work is proposed around FAZ morphology, including area, perimeter, and circularity, as well as boundary distances such as Hausdorff for FAZ and cross-dataset generalization.
A further point of clarification concerns nomenclature. In the description of the 2026 preprint "DenVisCoM: Dense Vision Correspondence Mamba for Efficient and Real-time Optical Flow and Stereo Estimation" (Anand et al., 2 Feb 2026), the phrase “Joint-OCTAMamba” is used to denote a unified optical correspondence and triangulation via attention+Mamba for real-time joint optical flow and stereo disparity. This suggests a nomenclatural overlap rather than an identity of methods. In the OCTA literature, however, Joint-OCTAMamba refers specifically to the two-stage OCTA joint segmentation framework built from RVMamba and FAZMamba for retinal vessel and FAZ segmentation.