Class-Specific Segmentation Heads (CSH)

• Class-Specific Segmentation Heads (CSH) are specialized modules that assign dedicated, independently parameterized segmentation heads to each semantic class to address challenges like catastrophic forgetting and class imbalance.
• CSH decouples gradient flows using class-specific losses and efficient architectural designs, enabling scalable extension to new classes with minimal computational cost.
• Empirical results across medical imaging, weakly supervised, and zero-shot segmentation show improved class retention, segmentation accuracy, and parameter efficiency.

Class-Specific Segmentation Heads (CSH) are architectural modules that allocate a dedicated, independently parameterized segmentation head for each semantic class in a multi-class (or multi-label) segmentation network. Introduced to address catastrophic forgetting and target-class imbalance in continual, weakly-supervised, semi-supervised, and universal segmentation regimes, CSH decouples gradient flows and enables scalable, flexible extension to new object or anatomical classes with minimal computational cost. The approach has been substantiated across domains including medical imaging, semantic segmentation under weak supervision, and zero-shot adaptation, consistently demonstrating improved class retention, segmentation accuracy, and parameter efficiency (Zhang et al., 2023, Wang et al., 1 Apr 2025, Liu et al., 2024).

1. Motivations and Theoretical Underpinnings

Conventional multi-class segmentation networks utilize a single shared output layer (e.g., a $1\times1\times1$ convolution) with Softmax or multi-Sigmoid outputs, where all classes share kernel parameters. As a result, when new classes are introduced—particularly in continual learning—parameter updates for the shared head can degrade predictions for previously learned classes, a phenomenon known as "catastrophic interference". Additionally, large or dominant classes tend to overshadow minority classes due to non-uniform gradient coverage, exacerbating label imbalance (Zhang et al., 2023, Wang et al., 1 Apr 2025).

CSH resolves these issues by:

• Parameter isolation: Each class $k$ is attached to its own head $f_k(\cdot;\theta_k)$, whose parameters $\theta_k$ are not shared or overwritten by other classes.
• Independent optimization: Per-class heads are trained with class-specific losses and supervision, ensuring distinct and robust memory traces for each semantic entity.
• Flexible extensibility: New heads can be appended as new classes arrive, without modification to existing heads or retraining the entire network (Zhang et al., 2023, Liu et al., 2024, Wang et al., 1 Apr 2025).

Theoretical analysis demonstrates that, in the "divide-and-conquer" configuration of one-vs-rest specialists, the class-participation weights in loss computation are regularized closer to uniform, directly mitigating large-class dominance (Wang et al., 1 Apr 2025). Specifically, if $p_i$ is the pixel fraction for class $i$, then in the CSH specialist regime the deviation from uniform participation is shrunk by a factor dependent on class cardinality.

2. Architectural Designs

CSH modules are typically appended atop a shared encoder-decoder backbone (e.g., Swin-UNETR, U-Net, SegFormer, ViT-B), which extracts a dense feature representation $F$. For each class, a lightweight head is instantiated, with architectural implementations varying across literature:

• CNN-based CSH: Three-layer stacks of $3\times3\times3$ or $1\times1\times1$ convolutions with small channel widths (e.g., $[8 \to 8 \to 1]$ for 3D CT; total $\sim$1.5k parameters/head). Each outputs a Sigmoid-activated binary mask, allowing overlapping regions (Zhang et al., 2023, Liu et al., 2024).
• Specialist branches: Each class receives a "projector" (e.g., $3\times3$ conv + BN + ReLU) and a segmentation head producing three logits: background, class-$k$, other (Wang et al., 1 Apr 2025).
• Transformer/ViT-based CSH: Each class is assigned a dedicated [CLS] token in the input sequence; token specialization is enforced by random masking and attention-head gating, yielding class-discriminative self-attention maps suitable as pseudo-masks (Hanna et al., 9 Jul 2025).
• Language-driven parameterization: Class-specific heads are parameterized or modulated via MLPs acting on concatenated global image features and CLIP (or LLM) text embeddings of class descriptions, providing semantic priors and directly aligning kernels with natural language class semantics (Zhang et al., 2023, Liu et al., 2024, Chen et al., 27 Jun 2025).

Example CSH Layer Structure for 3D Medical Imaging

Conv Layer	Input Channels	Output Channels	Kernel Size	Nonlinearity
1	$C$	$8$	$1 \times 1 \times 1$	ReLU
2	$8$	$8$	$1 \times 1 \times 1$	ReLU
3	$8$	$1$	$1 \times 1 \times 1$	Sigmoid

(Liu et al., 2024)

3. Training Objectives and Pseudo-Labeling

Each CSH is trained on the available binary or multi-label mask for its respective class. Supervisory regimes include:

• Binary cross-entropy (BCE): Each head is optimized via class-specific BCE, using either one-hot ground-truth or soft pseudo-labels ($\hat{p}^{k}_{t-1}$ in continual learning) for absent classes in the current batch (Zhang et al., 2023, Liu et al., 2024).
• Dice loss: For robust overlap maximization, Dice loss is summed across heads when pixel-level annotation is available (Liu et al., 2024, Wang et al., 1 Apr 2025).
• Masked backpropagation: For partially labeled datasets, only task-specific heads are updated per sample; heads for absent classes receive zero gradient (Liu et al., 2024).

In continual/extension scenarios, pseudo-labels from previous model checkpoints are employed for old classes, while new classes use ground-truth annotations. This approach enables continual learning without retaining old data (Zhang et al., 2023).

• Transformer-based CSH: Multi-label BCE is applied to class logits computed from [CLS] tokens, supplemented by head pruning regularizers and random token masking to enforce class-to-token correspondence (Hanna et al., 9 Jul 2025).
• Cross-consistency losses: In collaborative architectures (e.g., CGS), consistency between a generalist segmentation head and CSH specialists is enforced via consensus pseudo-labels and cross-entropy/Dice terms (Wang et al., 1 Apr 2025).

4. Integration of Language and Vision Semantics

CSH architectures often leverage LLMs to provide semantic priors and encode class relationships:

• CLIP-driven kernels: Text embeddings $\omega_k$ of class descriptions (e.g., “a computerized tomography of a [CLASS]”) are extracted using CLIP or other vision-LLMs and fused with global image features to modulate MLP parameter generation for each head (Zhang et al., 2023, Liu et al., 2024).
• Semantic alignment: Head output features are aligned with text and visual prototypes, enforced by explicit KL divergence losses or InfoNCE losses to harmonize dense features with global and class-level CLIP representations (Chen et al., 27 Jun 2025).

This design not only boosts the discrimination and transferability of CSHs, but also enables zero-shot generalization to novel classes described by language inputs (Chen et al., 27 Jun 2025).

5. Applications in Continual, Semi-Supervised, and Weakly Supervised Segmentation

CSH have been central to a variety of segmentation architectures:

• Continual learning in medical imaging: By maintaining parameter isolation, CSH achieve superior retention and extension: for abdominal organ/tumor segmentation, CSH outperformed LwF, ILT, and PLOP in step-wise average Dice (e.g., 0.787 for 14 classes vs. ≤0.777 for competing distillation methods) and incur only minimal GFLOP overhead per new class (Zhang et al., 2023, Liu et al., 2024).
• Collaborative semi-supervised segmentation: In CGS, combining a generalist head with specialist CSHs and cross-consistency losses, model performance on imbalanced multi-target tasks is improved (e.g., ACDC dataset DSC +1.56% vs FixMatch baseline), and participation gap is theoretically bounded (Wang et al., 1 Apr 2025).
• Weakly supervised segmentation: CSH-ViT allocates one [CLS] token per class, using random masking, head pruning, and register tokens to yield sharply localized, class-specific pseudo-masks directly from attention maps, achieving mIoU gains on MS COCO/VOC (≥+1.1%) (Hanna et al., 9 Jul 2025).
• Zero-shot segmentation: Frozen CLIP-derived heads align visual features to text-guided prototypes, with empirical hIoU gains of 1–2% on ZS-COCO/PASCAL benchmarks attributable to the CSH module alone (Chen et al., 27 Jun 2025).
• Extreme small-sample regimes: In classification-via-segmentation frameworks, each class-specific segmentation output is averaged and softmaxed to yield classification logits, greatly boosting low-data classification accuracy (e.g., +15–30% on MNIST/CIFAR compared to linear heads) (Mojab et al., 2021).

6. Practical Considerations and Computational Efficiency

The parameter and computational overhead per class in CSH is minimal: instances report overheads on the order of 0.1 GFLOPs per class (relative to several hundred GFLOPs in the backbone) and total per-head parameters in the $10^2$–$10^3$ range (Zhang et al., 2023, Liu et al., 2024). For transformer-based CSH, sequence length increases with class cardinality, but head gating/pruning and register tokens help maintain throughput (Hanna et al., 9 Jul 2025). Specialists in multi-target medical segmentation are discarded at inference, so memory/compute costs vanish at test time (Wang et al., 1 Apr 2025).

Training typically employs AdamW (e.g., lr $10^{-4}$), with per-head training losses; batch sizes and optimization hyperparameters follow backbone conventions. For class extension and continual use, MLP-based parameter generators (modulated by CLIP/LLM text embeddings) enable heads to be attached on-the-fly (Zhang et al., 2023, Liu et al., 2024).

7. Benchmark Performance, Limitations, and Prospects

Empirically, CSH architectures consistently outperform one-headed and shared-head baselines in class retention (Dice, hIoU) and rare-class segmentation, as well as enabling continual or extensible segmentation without catastrophic forgetting. Representative performance deltas include +0.9% to +2.3% hIoU in zero-shot settings (Chen et al., 27 Jun 2025), +1–2% Dice in collaborative semi-supervised medical segmentation (Wang et al., 1 Apr 2025), and ranking #1 in multi-task decathlon challenges for universal medical segmentation (Liu et al., 2024).

Limitations include the linear growth in parameter count and, for transformer models, in input sequence length, which may be problematic for extremely large class vocabularies, though token sharing or dynamic instantiation has been proposed as future solutions (Hanna et al., 9 Jul 2025).

A plausible implication is that CSH will enable more scalable, flexible, and robust universal segmentation backbones, particularly where dynamic class inventories, strong language priors, or label-partial datasets are routine. The design paradigm of class-specific, language-modulated, independently supervised heads now underpins state-of-the-art continual and semi-supervised segmentation systems across medical and natural image domains.

---

References:

(Zhang et al., 2023, Hanna et al., 9 Jul 2025, Chen et al., 27 Jun 2025, Liu et al., 2024, Mojab et al., 2021, Wang et al., 1 Apr 2025)