Finding the 0.01%

Tumour segmentation in PET/CT scans with a loss built for extreme class imbalance · Image and Speech Processing Lab, NITK

Ayush Deshmukh · supervised by Assoc. Prof. Jeny Rajan, NIT Karnataka

Code Image & Speech Processing Lab

TL;DR In these PET/CT scans, tumour pixels are under 0.01% of the data, so a model can score almost perfectly by predicting "no tumour" everywhere. I wrote a binary weighted focal Tversky loss that makes missed tumour pixels expensive, combined it with BCE and Dice, and took a TransUNet from a Dice score of 0.41 to 0.656. It was the best of the 23 configurations I tested.

The problem

Tumour segmentation should be about the rare pixels. Here, across more than 30,000 PET and CT images, tumour made up less than one pixel in ten thousand. With standard losses the easy answer wins: predict background everywhere, get a great-looking pixel accuracy and a useless model. Accuracy doesn't matter here; what matters is whether you find the tumour, and whether what you mark as tumour really is one.

The loss

I started from the Tversky index, a Dice-like overlap score that weighs false positives and false negatives separately, and changed it in two ways:

Imbalance-aware weights, computed per batch. The false-positive and false-negative terms are scaled by the ratio of negative to positive pixels in that batch, so the rarer the tumour, the more a missed tumour pixel costs. Nothing to hand-tune per dataset.
A focal exponent on (1 − Tversky), so training keeps pushing on the hard, badly-segmented examples instead of the ones it already gets right.

On its own the Tversky term can be unstable, so the final objective is a weighted mix: 0.6 × BCE + 0.4 × weighted focal Tversky + 0.6 × Dice. BCE keeps per-pixel gradients well behaved, Dice rewards overlap, and the Tversky term is what makes the rare class count.

Setup

Model: TransUNet, built with Segmentation Models PyTorch
Data: 30,000+ PET and CT images, with heavy augmentation: flips, 90° rotations, transposes and elastic deformation
Training: SGD with an exponential learning-rate decay, and early stopping on validation loss
Evaluation: Dice and IoU, computed with MedPy

Results

Dice score up from 0.41 to 0.656
The best of 23 model and loss configurations I compared

Context

This was a research project from October 2023 to February 2024 with the Image and Speech Processing Lab at NIT Karnataka, supervised by Assoc. Prof. Jeny Rajan. It's also where I first ran into the problem that came back in the Mpox work: when the class you care about is rare, the loss function decides what the model learns.