We used a subset of the PUMA dataset containing H&E-stained melanoma regions of interest. The tissue segmentation task maps every pixel to Tumor, Stroma, or Other, while the nuclei classification task assigns extracted nuclei patches to Tumor, Lymphocyte, or Histiocyte. Both settings are affected by limited data and substantial class imbalance.
For tissue segmentation, we compared an end-to-end U-Net with the same architecture initialized through noisy image-reconstruction pre-training. For nuclei classification, we compared an end-to-end EfficientNet-B0 with frozen encoders learned through SimCLR and supervised contrastive learning. This allowed us to study when pre-training provides a useful starting point and when a model still needs to adapt its features to the downstream task.
Figure 1. Per-image tissue coverage in the training set. The Other class is sparse and heavily skewed toward zero coverage.
The segmentation split contains 163 training, 20 validation, and 22 test images at $1024\times1024$ pixels. Tumor appears in every training image and covers $68.95\%$ of an image on average. Stroma covers $25.47\%$, while Other has a mean coverage of only $5.60\%$ and is absent from many images.
We therefore trained on class-aware $512\times512$ patches rather than sampling uniformly from the complete images. We filtered patches with too little annotated tissue and constructed batches to include minority-class content. The same concern shaped the nuclei dataset: we limited the contribution from any single region of interest so densely annotated slides could not dominate training.
Our segmentation model follows the standard U-Net encoder-decoder design with four resolution stages, a $1024$-channel bottleneck, and skip connections that recover spatial detail during decoding. We use Group Normalization because the $512\times512$ patches limit the batch size to four, making batch statistics less reliable. The resulting network has $31.04$M trainable parameters, substantially fewer than the approximately $125$M-parameter baseline.
We train with a balanced combination of weighted cross-entropy and Dice loss. Cross-entropy provides stable pixel-wise supervision, while Dice directly encourages overlap for smaller classes:
On the test set, the end-to-end U-Net achieves a mean Dice score of $0.5851$, mean IoU of $0.4916$, and pixel accuracy of $0.8332$. Tumor is the strongest class with Dice $0.8524$, followed by Stroma at $0.6142$. Other remains the most difficult class at $0.2885$ because it covers smaller, more heterogeneous regions.
Figure 2. Training behavior and per-class segmentation metrics for the end-to-end U-Net.
We pre-trained the encoder through noisy image reconstruction, asking an autoencoder without skip connections to recover clean histology patches from corrupted inputs. Removing skip connections forces information through the bottleneck rather than allowing the decoder to copy local detail directly. The reconstruction objective combines pixel-level error with structural similarity so that the model preserves tissue texture and nuclei boundaries.
Original validation patches.
Reconstructed validation patches.
A frozen pre-trained encoder reaches a test Dice score of $0.5545$, below the end-to-end model. Allowing the encoder to fine-tune at a reduced learning rate of $10^{-5}$ raises the score to $0.5946$, the strongest segmentation result and a $0.1276$ improvement over the $0.4670$ baseline. The gain over the end-to-end U-Net is small, but the comparison shows that reconstruction features are most useful as an initialization rather than as a fixed representation.
For nuclei classification, we extract $100\times100$ patches centered on annotated nuclei and resize them to $224\times224$. We fine-tune an ImageNet-pretrained EfficientNet-B0 with a lightweight task-specific classifier, moderate geometric and color augmentation, label smoothing, AdamW, cosine annealing, and early stopping.
The selected model reaches $77.93\%$ test accuracy, improving on the $70.83\%$ baseline while using $4.34$M rather than approximately $5$M trainable parameters. Tumor and Lymphocyte obtain F1-scores of $0.8183$ and $0.8006$. Histiocyte is more difficult at $0.6937$, reflecting its visual overlap with the other classes and its smaller representation in the test set.
Figure 3. Representative confident correct predictions and high-confidence mistakes from the end-to-end nuclei classifier.
We also pre-trained EfficientNet-B0 encoders with SimCLR and supervised contrastive learning. SimCLR treats two augmentations of the same patch as a positive pair, while supervised contrastive learning pulls together all examples from the same nuclei class. After pre-training, we froze each encoder and trained only a $3{,}843$-parameter linear classification head.
The supervised objective produces more distinct t-SNE clusters and substantially stronger same-class similarity than SimCLR. However, the downstream results remain below the baseline: SimCLR reaches $64.79\%$ accuracy and SupCon reaches $69.64\%$. The encoders are exceptionally parameter-efficient during downstream training, but freezing them prevents the learned features from adapting to the classification objective.
SimCLR embeddings.
Supervised contrastive embeddings.
For both tasks, the strongest results come from models whose full representations can adapt to the downstream objective. Fine-tuned reconstruction pre-training gives the best tissue segmentation result, although only slightly ahead of a carefully trained end-to-end U-Net. For nuclei classification, the end-to-end EfficientNet-B0 clearly outperforms frozen SimCLR and SupCon encoders. Pre-training provides useful structure, but class-aware data construction, suitable losses, and downstream fine-tuning remain essential under limited and imbalanced histopathology data.