Beyond Random Sampling: Distribution-Aware Alignment for Semi-Supervised Medical Image Segmentation

Conference: ECCV2026
Official Paper: 5497
Paper: Official PDF
Code: https://github.com/ywher/DAA4SSMIS
Area: Medical Imaging
Keywords: Semi-Supervised Medical Image Segmentation, Distribution-Aware Sample Selection, Semantic Memory Bank, Copy-Paste, Progressive Training

TL;DR

The framework selects annotation candidates using frozen visual features and density-weighted K-Center, then trains with memory-guided Copy-Paste and progressive branch activation, reaching 87.3% DSC on PROMISE with two labeled cases; the gain combines changed sample selection and changed training.

Background & Motivation

Semi-supervised medical image segmentation commonly treats a small labeled subset and a large unlabeled subset as samples from the same distribution. With enough annotations this can be reasonable, but a budget of only a few cases makes random coverage unreliable. The selected cases may occupy similar anatomical or acquisition-related clusters, leaving other structures without a useful supervision anchor. Even a strong method such as UniMatch V2 can then propagate predictions based on a narrow view of the training population.

Copy-Paste augmentation helps move reliable labels into unlabeled images, but the foreground in medical scans is often small and missing from the current mini-batch. Random mixing can therefore be dominated by background, while aggressive use of early pseudo-labels reinforces incorrect boundaries. This paper connects two decisions usually handled separately: which cases deserve annotation and how their supervision should reach the remaining data. Its source and target domains mean the labeled and unlabeled subsets, not necessarily an explicitly constructed cross-hospital adaptation benchmark.

The approach first improves the coverage of reliable supervision and then improves its propagation during training. Core idea: choose representative annotation anchors with distribution-aware sampling, then use foreground-aware historical memory and an easy-to-hard branch schedule to turn those anchors into more reliable semi-supervised learning.

Method

Overall Architecture

The input is a training pool of medical images and a fixed annotation budget; the output is a segmentation network that predicts dense masks for new images. Offline, Distribution-Aware Sample Selection chooses images or patient cases in a frozen DINOv2 feature space; these are annotated while the remainder stays unlabeled. Online, Memory-Guided Copy-Paste constructs training examples, and Progressive Branch Activation controls when labeled, mixed, and unlabeled training paths become active. These paths belong to a teacher–student training framework, rather than three independent models that must be ensembled at deployment.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Training pool and<br/>annotation budget"] --> B["Distribution-Aware<br/>Sample Selection"]
    B --> C["Annotate selected data;<br/>retain unlabeled remainder"]
    C --> D["Memory-Guided<br/>Copy-Paste"]
    D --> E["Progressive<br/>Branch Activation"]
    E --> F["Teacher–student training;<br/>output segmentation model"]

The teacher generates pseudo-labels and is updated through an exponential moving average of student parameters. The segmentation model uses a DINOv2-Small encoder and a DPT decoder; freezing the offline feature extractor does not mean freezing the encoder during segmentation training. Sample selection, memory retrieval, and the progressive schedule are preparation-time or training-time operations. The claimed absence of inference overhead means no additional overhead from these mechanisms relative to the segmentation baseline, not computation-free inference.

Key Designs

1. Distribution-Aware Sample Selection: spend scarce annotations on different structures, not redundant random cases

Selection needs a representation that can compare anatomy without first learning from the small annotation set. The method extracts patch tokens from four intermediate layers of a frozen vision foundation model, globally averages tokens within each layer, concatenates the results, and applies L2 normalization. With DINOv2-Small, four 384-dimensional layer descriptors produce a 1536-dimensional slice representation. This retains a mixture of local-texture and higher-level semantic statistics, not the complete spatial token grid. The rationale is to avoid relying only on the final layer or a single CLS token when selecting images for a boundary-sensitive task.

For volumetric cases, averaging every slice into one vector could dilute a small target with uninformative slices. The case-sequence normalization instead uses structural centering: a fixed number of central slices forms a standardized case-level feature tensor. Selection therefore remains a patient-level decision rather than independent slice selection, although this operation is not anatomical registration. It assumes that central regions contain important anatomy; peripheral lesions are an obvious situation requiring further validation. The main text does not fully specify the reference slice count or handling of shorter volumes, so no interpolation or padding rule is inferred here.

Density-K-Center then modifies a coverage-oriented selector with local typicality. Ordinary farthest-point selection can spend its budget on isolated artifacts, whereas selecting only dense regions can repeatedly choose nearly identical cases. The proposed selector estimates density from 20 nearest neighbors and constrains density weights to the range 0.3–1.0, making atypicality a soft penalty rather than an exclusion rule. It initializes with the sample farthest from the global feature mean, then greedily selects candidates with the largest density-weighted coverage distance until the annotation budget is exhausted. The intended effect is to cover previously unrepresented regions while preferring anchors that also represent neighboring data.

This is one-shot screening before annotation, not online active learning that repeatedly trains a model and requests more labels. A project with an already fixed historical annotation set cannot automatically claim the full selection benefit. The farthest-from-mean initialization also means density weighting should not be described as a hard guarantee against every outlier at every step. The cached density and greedy-selection equations are extraction-corrupted; the mechanism above follows Section 3.2 prose without inventing their exact algebra.

2. Memory-Guided Copy-Paste: make historical foreground available beyond the current mini-batch

Training divides each image into a \(7\times7\) patch grid and constructs examples through three paths. The labeled branch augments only labeled data, the unlabeled branch augments unlabeled data using teacher pseudo-labels, and the mixed branch performs bidirectional Copy-Paste between the two subsets. These paths respectively provide reliable supervision, use of unlabeled data, and intermediate examples that connect the subsets. Alignment is encouraged through augmented examples and their training losses; there is no additional adversarial domain classifier in the described method.

A small mini-batch may contain too few usable foreground patches, regardless of how its patches are rearranged. MCP therefore stores historical patches in separate foreground and background memory subsets, with pixels, masks, class probabilities, mean confidence, and source identifiers in each entry. Separate updates prevent abundant background patches from continually displacing rare lesions or organ regions. At capacity, a hybrid policy evicts half of the removed entries by FIFO for freshness and half by lowest confidence for quality. The reported capacity parameter is 1024; this shorthand does not establish an exact implementation-level memory footprint.

Augmentation uses memory retrieval with probability 0.2 and otherwise samples from the current mini-batch, retaining exposure to current data. Retrieval considers foreground/background identity and searches for a source patch whose class-probability distribution has minimum KL divergence from the target distribution. A foreground lesion from memory can also replace a background region to create examples containing the target. Consequently, semantic matching should not be interpreted as requiring background-to-background replacement in every operation. It constrains source choice using class information and probability distributions rather than relying only on random spatial mixing.

The important change is that scarce foreground observed in previous batches remains accessible as a supervision resource. However, similar class distributions do not guarantee compatible anatomical positions, shapes, or imaging styles; semantic retrieval is not a geometric-validity guarantee. The cached KL equation is incomplete, so its argument direction and full candidate constraints cannot be checked and are not reconstructed here.

3. Progressive Branch Activation: learn from reliable labels before relying on mixed and fully pseudo-labeled examples

Better patch retrieval cannot make an immature teacher immediately trustworthy. Activating every path from the beginning would expose the student to noisy supervision through unlabeled examples, mixed examples, and historical memory simultaneously. The method reserves the first 10% of training for progression: only the labeled branch runs during the first half, the mixed branch joins during the second half, and all branches run afterward. In total-training terms, this means roughly 0–5% labeled-only, 5–10% labeled plus mixed, and all paths after 10%. This is a fixed time schedule, not confidence-triggered adaptive switching.

Each active branch combines Dice and cross-entropy losses, weighting ground-truth regions by 1.0 and pseudo-labeled regions by 0.5. Dice emphasizes overlap under imbalance, cross-entropy supplies pixel-level class supervision, and pseudo-label downweighting reduces the influence of uncertain targets. Once all branches are active, their branch-level losses have equal weights, while the different region-level confidence weights remain in effect. Equal branch weights therefore do not mean that every pixel has identical supervision weight.

Loss & Training

Implementation uses AdamW with weight decay 0.01, encoder learning rate \(5\times10^{-6}\), decoder learning rate \(2\times10^{-4}\), and polynomial decay. Inputs are resized to \(518\times518\), labeled and unlabeled mini-batches are balanced, and training uses two RTX 4090 GPUs; the implementation section also lists OHEM. The main text does not completely unpack how OHEM combines with every branch loss, nor does it provide all training-duration and EMA details needed for reproduction. Volumetric datasets are still trained slice by slice, with predictions concatenated back into volumes for case-level evaluation.

Key Experimental Results

Main Results

The evaluation covers four 2D and two 3D datasets; the selected comparison below uses UniV2 because it shares the DINOv2-S encoder family. BUSI uses a 623/157 train/test image split and PMTCXR uses 2379/290; PROMISE uses 35/5/10 train/validation/test cases and ACDC uses 70/10/20. DSC is higher-is-better, and gains are percentage points, not relative percentages; scores come from Tables 1 and 2.

Dataset and annotation budget	UniV2 DSC (%) ↑	Ours DSC (%) ↑	Gain
PROMISE, 1/16, 2 cases	82.0	87.3	+5.3
ACDC, 1/20, 3 cases	88.6	90.4	+1.8
BUSI, 1/16, 39 images	74.6	79.8	+5.2
PMTCXR, 1/16, 149 images	47.6	52.9	+5.3

At the same BUSI budget, HD95 falls from 48.8 for UniV2 to 36.2, a reduction of 12.6 pixels in Table 2, providing boundary evidence alongside overlap gains. Predictions are resized to \(256\times256\) for ACDC and \(224\times224\) for the other datasets before evaluation; boundary errors must be interpreted under that protocol. Matching the encoder family does not isolate MCP: the proposed method also changes which training samples receive labels.

Ablation Study

The following selected Table 3 results use BUSI at the 1/16 annotation budget with the same 623/157 train/test split. CP denotes basic three-branch Copy-Paste, MB the semantic memory bank, Prog. progressive activation, and SS distribution-aware selection; lower boundary distances are better.

Configuration	DSC (%) ↑	ASD ↓	HD95 ↓
UniV2 baseline	74.62	36.66	48.78
+ CP	75.58	32.29	44.17
+ CP + MB	77.35	29.13	40.75
+ CP + MB + Prog.	78.06	31.30	40.05
+ SS only	78.06	30.55	40.92
+ CP + MB + Prog. + SS	79.80	27.00	36.20

Key Findings

• Adding the memory bank after CP improves DSC by 1.77 percentage points, supporting historical foreground retrieval without establishing it as the dominant component in every setting.
• Full online MCP and selection alone both reach 78.06% DSC; their combination reaches 79.80%, another 1.74 points over MCP, supporting complementarity.
• Progressive activation improves DSC but changes ASD from 29.13 to 31.30, so not every component improves every metric monotonically.
• Table 5 does not show DKC beating K-Center on every metric at every budget; the default was chosen using average performance across several low-label regimes.

Highlights & Insights

• Annotation selection becomes part of the algorithm rather than a fixed preprocessing assumption. Under an extremely small budget, a case that covers new structure can matter more than more augmentation of redundant cases.
• Foreground/background separation allocates memory resources, not merely additional storage. It preserves access to minority targets when later mini-batches would otherwise contain little usable foreground.
• Introducing mixed supervision before fully unlabeled augmentation makes reliable masks a transition into pseudo-label learning. The schedule is simple to integrate, but its timing still needs validation on the intended dataset.

Limitations & Future Work

• Author-stated: Current 3D processing remains slice-wise; direct volumetric MCP and adaptive memory management are future directions. Training has overhead, while inference adds no extra modules.
• Reader assessment: Selection assumes access to a training pool before annotation choices are fixed. Historical labeled sets, peripheral rare lesions, or modalities poorly represented by the foundation model may limit applicability.
• Evidence boundary: Main tables do not report multi-seed confidence intervals, so point estimates are not statistical proof of stability. Broader feature-space coverage also does not establish coverage of clinical patient populations.
• Reproducibility: The local source contains full main-paper methods and experiments but no separate supplement. Several equations are extraction-corrupted, and case-length handling and full pseudocode cannot be independently verified; missing rules are not invented.
• Resource status: The code URL comes from a future-release statement in the paper. Its current availability was not checked online, and a listed URL does not establish runnable public code.

Related Work & Insights

• Versus UniMatch V2: The framework retains a strong pretrained encoder and semi-supervised segmentation setting while moving part of the improvement upstream to sample selection and adding historical retrieval and scheduling. The main gain is not simply a larger-backbone comparison.
• Versus BCP, MiDSS/TP-RAM, ABD, and RCP: Copy-Paste and intermediate-domain learning already exist in these lines of work. The emphasis here is cross-batch foreground availability and historical semantic retrieval, not claiming bidirectional mixing itself as new.
• Versus K-Center and active learning: Density weighting softens the pull of extreme outliers, and one-shot offline ranking avoids repeated train-and-query rounds. It also cannot adapt its selections to segmentation errors discovered later in training.

Rating

• Novelty: 3.5/5 — A targeted combination of selection, memory augmentation, and curriculum scheduling rather than wholly new primitives.
• Experimental Thoroughness: 4/5 — Six datasets and multiple component, selection, and mixing analyses, with uncertainty reporting and implementation gaps remaining.
• Writing Quality: 3.5/5 — A clear main argument with some details deferred to supplementary material; local equation extraction damage is not treated as an author writing defect.
• Value: 4/5 — Useful for medical segmentation projects that can plan their annotation budget, without establishing clinical effectiveness.