ClaSP PE: Stratified Predictive Entropy for 3D Segmentation

• The paper introduces ClaSP PE, a query strategy that integrates class-stratified querying with scheduled log-scale power noising to enhance 3D segmentation active learning.
• It addresses critical issues such as severe class imbalance and early redundancy by incorporating class-aware uncertainty scores and a decaying noise schedule.
• Empirical results on nnActive benchmarks demonstrate that ClaSP PE outperforms random and uncertainty-only baselines, delivering statistically significant segmentation quality gains.

Class-stratified Scheduled Power Predictive Entropy (ClaSP PE) is a query strategy for pool-based active learning in 3D biomedical image segmentation. It was introduced to address a specific empirical impasse: existing active learning methods in this setting were unable to consistently outperform improved random sampling baselines adapted to 3D data. ClaSP PE combines two lightweight mechanisms—class-stratified querying and scheduled log-scale power noising—to target two failure modes of uncertainty-only acquisition in volumetric segmentation: class imbalance and redundancy in early selections. Within the nnActive benchmark, it is described as the only method that generally outperforms improved random baselines in terms of both segmentation quality with statistically significant gains, whilst remaining annotation efficient, and it is reported to generalize to previously unseen datasets without dataset-specific tuning (Lüth et al., 20 Jan 2026).

1. Problem setting and motivation

ClaSP PE is formulated within nnActive for pool-based active learning on 3D segmentation with nnU-Net as backbone, specifically the 3D full-resolution configuration, together with partial loss training so that only annotated voxels contribute to the loss. Query granularity is fixed-size 3D patches, with two benchmarked patch sizes, “Main” and “Patch×1/2,” the latter halving each dimension. The active learning trajectory uses 5 cycles per experiment, with nnU-Net retrained for 200 epochs per cycle in the benchmark; a post-hoc analysis on AMOS showed that 500 epochs can further boost ClaSP PE.

The method is motivated by two empirical shortfalls of standard uncertainty-based active learning in 3D biomedical segmentation. First, voxel-level class imbalance is severe: background dominates, while foreground classes differ widely in volume. When uncertainty is aggregated over whole 3D items, large and easy background regions and large organs tend to dominate the score, leaving small or rare structures underqueried and suppressing gains in mean-per-class metrics such as Dice. Second, early active learning rounds often suffer from low diversity: top-$k$ selection from peaky uncertainty distributions over highly overlapping 3D patch candidates induces redundancy, so near-duplicate regions consume annotation budget.

ClaSP PE addresses these two limitations with orthogonal mechanisms. Class-stratified querying enforces coverage of underrepresented structures using the model’s predictive probabilities rather than additional ground-truth labels. Scheduled log-scale power noising introduces controlled stochasticity into the ranking to diversify early selections and then decays toward near-deterministic exploitation in later cycles. This suggests that the method is designed not as a replacement for predictive entropy, but as a restructuring of how predictive entropy is operationalized under 3D patch-based constraints.

2. Predictive entropy scoring and class-aware stratification

The base uncertainty signal in ClaSP PE is voxel-level predictive entropy computed from a deterministic nnU-Net. No MC dropout or ensemble is employed, and no epistemic/aleatoric decomposition is used. For voxel $v$, with softmax probabilities $p(v) = (p_1(v), \ldots, p_C(v))$, the uncertainty is

$H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$

Patch-level acquisition relies on scores $u(P)$ derived from an uncertainty map $u(x)$. The exact aggregation operator is not specified in the paper; in nnActive, predictive-entropy baselines typically use mean aggregation over voxels,

$u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$

The paper notes that other aggregation operators, such as max, are possible but not described. Because the aggregation is left unspecified, mean aggregation should be understood as the concrete operationalization consistent with typical nnActive predictive-entropy baselines rather than as a universally mandated component.

Class-stratified querying is built on a class-specific patch score. For each patch $P$ and class $c$,

$u_c(P) = \bar{p}_c(P)\cdot u(P),$

where

$v$0

is the mean predicted probability of class $v$1 over voxels in the patch. This weighting biases uncertainty by estimated class presence, encouraging selection of patches relevant to class $v$2. In the segmentation tasks discussed in the paper, classes correspond to the segmentation labels, such as organs in AMOS, left and right hippocampus in MSD Hippocampus, kidney/tumor/cyst in KiTS, and cardiac structures in ACDC.

The significance of this construction is methodological rather than architectural. Stratification does not require an explicit prevalence estimator $v$3, inverse-frequency weighting, or additional annotation metadata. Instead, it uses model predictions $v$4 to stratify candidate patches per class. A plausible implication is that ClaSP PE preserves the computational simplicity of entropy-based active learning while reallocating selection pressure toward structures that are otherwise neglected by class-agnostic aggregation.

3. Budget allocation and query selection

Each active learning round selects a batch of $v$5 patches from the unlabeled pool. A fixed fraction $v$6 of this budget is class-stratified, with the paper fixing $v$7. The quota scheme is equal per class:

$v$8

for each class $v$9, with any remainder distributed arbitrarily. No inverse-frequency weighting is applied.

Selection proceeds in two branches. For each class $p(v) = (p_1(v), \ldots, p_C(v))$0, all candidate patches are ranked by $p(v) = (p_1(v), \ldots, p_C(v))$1 in descending order, and the top $p(v) = (p_1(v), \ldots, p_C(v))$2 patches are selected. For the remaining $p(v) = (p_1(v), \ldots, p_C(v))$3 budget, the pool is ranked by the class-agnostic uncertainty $p(v) = (p_1(v), \ldots, p_C(v))$4, and the top $p(v) = (p_1(v), \ldots, p_C(v))$5 patches are chosen. The two sets are then merged. Ties are broken arbitrarily, for example randomly. The implementation ensures unique patch instances, while overlapping patches are permitted by design because informativeness is localized; exact deduplication constraints are not specified. No volume-level constraints, such as per-volume caps, are imposed by the method itself.

This budget split is central to the method’s behavior. The equal per-class quota scheme does not estimate class prevalence and does not explicitly adapt to frequency differences. Instead, it guarantees some class-specific acquisition pressure regardless of prevalence. The paper’s rules-of-thumb recommend increasing $p(v) = (p_1(v), \ldots, p_C(v))$6 for severe class imbalance or small and rare targets, and note that budgets should scale with the number of classes. This suggests that the quota design is intended as a robust default rather than a closed-form optimum.

4. Scheduled log-scale power noising

The second defining component of ClaSP PE is scheduled log-scale power noising, whose objective is to inject controlled stochasticity into rank ordering early in active learning and recover near-deterministic top-$p(v) = (p_1(v), \ldots, p_C(v))$7 selection later. For each patch score $p(v) = (p_1(v), \ldots, p_C(v))$8, either $p(v) = (p_1(v), \ldots, p_C(v))$9 or $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$0, the noisy score is defined on the log-scale as

$H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$1

with

$H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$2

Equivalently, the method applies multiplicative log-Gumbel noise $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$3 with scale $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$4; larger $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$5 yields smaller noise.

The noise magnitude follows a decaying schedule,

$H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$6

with $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$7 at cycle $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$8, implying scale $H(p(v)) = -\sum_{c=1}^C p_c(v)\log p_c(v).$9, and $u(P)$0 at cycle $u(P)$1, implying scale $u(P)$2. The schedule is therefore exponential in log-space, enforcing strong noise early and weak noise late.

The paper’s interpretation is that early cycles contain many patches with similar base entropy, so adding Gumbel noise in log-space spreads ranks and encourages coverage across anatomy and image regions, thereby reducing redundancy from overlapping candidates. Later, as the model becomes more confident and uncertainty concentrates on hard regions, the small-noise regime makes selection closely match top-$u(P)$3 ranking by $u(P)$4 or $u(P)$5, emphasizing exploitation. The mechanism is explicitly linked to stochastic batch acquisition, or “power noising,” but adapted to the 3D patch setting and scheduled over active learning cycles.

5. Complete active learning round and computational profile

A ClaSP PE round takes as inputs the current model $u(P)$6, the pool of unlabeled volumes $u(P)$7, the query budget $u(P)$8, the stratification rate $u(P)$9, and scheduler parameters $u(x)$0, $u(x)$1, and $u(x)$2. The procedure is as follows: train or update $u(x)$3 on the labeled data with partial loss for 200 epochs; compute voxel-wise probabilities $u(x)$4 and entropy $u(x)$5; compute $u(x)$6, $u(x)$7, and $u(x)$8 for each candidate patch; compute $u(x)$9; perturb class-specific and class-agnostic scores using Gumbel noise; select top $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$0 patches per class; select top $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$1 remaining patches by noisy class-agnostic score; merge, enforce uniqueness, and dispatch the final $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$2 patches for annotation; then add the annotated patches to the labeled set and proceed to the next cycle (Lüth et al., 20 Jan 2026).

The method is designed for the scale constraints of 3D segmentation pools. The paper notes that KiTS median volumes yield approximately $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$3 candidate patches with the benchmark’s strides. ClaSP PE therefore avoids representation-based selection such as Core-Set, with the explicit aim of keeping runtime and memory tractable in 3D. Complexity is dominated by batched forward passes and simple reductions. Computing $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$4 and $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$5 over the pool is $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$6 per forward pass, while aggregation to $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$7 and $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$8 is linear in voxels per patch and per class. Noising and ranking are $u(P) = \frac{1}{|P|}\sum_{v \in P} H(p(v)).$9 per list, where $P$0 is the number of candidate patches.

The operational notion of annotation cost is Foreground Efficiency (FG-Eff), which relates performance to the number of queried foreground voxels using an exponential fit. Improved random baselines, specifically Foreground-aware Random 33% and 66% FG, are used to emulate practical foreground-focused sampling. Annotation units remain fixed-size 3D patches, and the paper positions this granularity as compatible with modern interactive tooling such as scribbles, sparse slice annotations, and promptable models, while partial loss allows efficient use of these sparse annotations.

6. Hyperparameters, empirical evaluation, and ablations

The nnActive defaults reported for ClaSP PE are concise. The stratification fraction is $P$1, so 66% of the batch per cycle is class-stratified. The schedule endpoints are $P$2 and $P$3, with no further tuning performed across datasets. The uncertainty aggregation used in nnActive predictive-entropy baselines is mean entropy over voxels within a patch. The class-specific score $P$4 uses no smoothing constants. Training uses 200 epochs per cycle in the benchmark, although longer training improved results for challenging multi-class tasks such as AMOS Low-Label.

Parameter	Default	Note
$P$5	0.66	66% of each batch is class-stratified
$P$6	1	Strong noise at $P$7
$P$8	100	Near-deterministic at $P$9
Cycles	5	Benchmark active learning trajectory
Training per cycle	200 epochs	500 epochs improved AMOS post-hoc
Aggregation	Mean entropy over voxels	Default in nnActive predictive-entropy baselines

The empirical benchmark comprises 24 settings defined by 4 datasets, 3 label regimes, and 2 patch sizes, with 4 seeds, nnU-Net 3D full-resolution, partial loss, 5 active learning cycles, and 200 epochs per cycle. The datasets are AMOS2022 Challenge task 2, MSD Hippocampus, KiTS2021, and ACDC. Metrics include Mean Dice, Area Under Budget Curve (AUBC), Final Dice, FG-Eff, Pairwise Penalty Matrix (PPM), and a Nemenyi post-hoc test at $c$0 for average rank groups. Comparative methods are Random, Foreground-aware Random 33% and 66% FG, Predictive Entropy, BALD, PowerBALD, SoftrankBALD, and PowerPE.

Aggregate performance over all 24 settings is reported as follows: ClaSP PE achieves AUBC $c$1 and Final Dice $c$2; Random 66% FG achieves AUBC $c$3 and Final Dice $c$4; PowerPE achieves AUBC $c$5 and Final Dice $c$6; Predictive Entropy achieves AUBC $c$7 and Final Dice $c$8; Random achieves AUBC $c$9 and Final Dice $u_c(P) = \bar{p}_c(P)\cdot u(P),$0 (Lüth et al., 20 Jan 2026). ClaSP PE is reported to lie in the distinct top-performing group for AUBC and Final Dice and to be competitive in FG-Eff. PPM indicates significant wins against most baselines, with the remaining losses against Random 66% FG concentrated on AMOS Low-Label, where the number of classes is very large relative to the budget.

Dataset-wise trends are also differentiated. KiTS shows substantial improvements for ClaSP PE. AMOS Low-Label is challenging because 15 classes interact with small budgets; longer training, specifically 500 epochs, reduces losses, and the method shows consistent wins versus Random 66% FG in the High-Label regime. For Hippocampus and ACDC, small absolute differences are commonly within error bars, which the paper uses to motivate benchmark aggregation rather than single-dataset conclusions.

The ablations disentangle the two main components. Stratification alone, denoted Cla PE, improves over standard Predictive Entropy. Using $u_c(P) = \bar{p}_c(P)\cdot u(P),$1 yields higher AUBC and Final Dice than $u_c(P) = \bar{p}_c(P)\cdot u(P),$2, with only minor FG-Eff reduction. Constant power noising, denoted ClaP PE with $u_c(P) = \bar{p}_c(P)\cdot u(P),$3, increases FG-Eff relative to stratification-only but reduces absolute AUBC and Final Dice. Decayed noising, i.e., ClaSP PE, yields the best overall trade-off: early diversity improves FG-Eff, while the late shift to exploitation boosts Dice and AUBC.

7. Generalization, deployment, limitations, and related positioning

The paper explicitly evaluates roll-out to four previously unseen datasets under a plug-and-play protocol with no dataset-specific tuning. These datasets are LiTS, WORD, Tooth Fairy 2, and MAMA MIA, using a 75% train/pool and 25% test split. Patch size is chosen by guidelines based on median connected component size, and budgets are normalized, for example 50–100 patches per class. The compared methods are ClaSP PE, PE, Random, and Random 66% FG. ClaSP PE is reported to yield robust top performance across AUBC and Final Dice, to improve FG-Eff over PE in all roll-out datasets, and never to be significantly outperformed by Random or Random 66% FG in the aggregated PPM. On MAMA MIA, which is described as large and redundant, the table reports AUBC $u_c(P) = \bar{p}_c(P)\cdot u(P),$4, Final Dice $u_c(P) = \bar{p}_c(P)\cdot u(P),$5, and FG-Eff $u_c(P) = \bar{p}_c(P)\cdot u(P),$6, underscoring the diversity benefit attributed to scheduled noising (Lüth et al., 20 Jan 2026).

Deployment is framed as minimal within nnActive. The implementation is available at the nnActive repository. The prescribed pipeline is to configure dataset preprocessing and nnU-Net auto-configuration, specify the 3D patch query size according to deployment guidelines, set the query strategy to ClaSP PE with $u_c(P) = \bar{p}_c(P)\cdot u(P),$7, $u_c(P) = \bar{p}_c(P)\cdot u(P),$8, $u_c(P) = \bar{p}_c(P)\cdot u(P),$9, and $v$00 equal to the number of active learning cycles, ensure entropy computation is enabled, train the initial model with a warm-start and partial loss trainer, and run active learning cycles while storing checkpoints per cycle. The scheduler is per-cycle rather than per-batch, and noise injection is applied at query time on candidate score lists.

Several failure modes and mitigations are given. Because entropy relies on calibration, miscalibrated probabilities can degrade selection quality; temperature scaling, longer training, and pretraining are proposed as mitigations. For extreme rare classes, equal quotas may still underselect if $v$01 is near zero; the paper recommends increasing $v$02, reducing patch size to better target small structures, raising budgets proportional to class count, and adopting longer training. Performance also depends on warm-start quality; if the initial model is too weak, stratification is less effective, so a larger starting budget or pretraining may be beneficial. Redundancy from overlapping patches is reduced but not eliminated by scheduled noising, and per-volume or per-region caps may be imposed by the surrounding framework if needed.

In relation to prior active learning strategies, ClaSP PE is positioned against uncertainty-only methods such as entropy, BALD, PowerBALD, SoftrankBALD, and PowerPE. The paper argues that these methods suffer from class imbalance and redundancy in 3D patch pools, whereas ClaSP PE compensates class imbalance through stratification and reduces early redundancy through scheduled stochasticity. It is also positioned against representativeness or core-set methods, which are described as memory- and runtime-heavy in 3D patch settings and as underperforming uncertainty-based active learning for U-Net backbones.

Two implementation details are explicitly marked as unspecified. The paper does not specify the exact patch-level aggregation operator for entropy and probabilities, though mean aggregation is used in the reported operationalization and is consistent with typical nnActive predictive-entropy baselines. Tie-breaking and uniqueness constraints are also not specified in detail; random tie-breaking and prevention of duplicate patch instances within one query batch are described as reasonable implementation choices. This suggests that ClaSP PE should be understood primarily as a query policy defined by class-stratified predictive-entropy scoring and scheduled log-scale power noising, rather than by a rigidly fixed low-level implementation.