Question: Class Imbalance Handling

[vmiller987]

Hello,

I actually have an ML question and not a technical issue question for once.

I have a dataset with fairly large class imbalance. I haven't dealt with this before and I am trying to understand the proper way to handle this. Aluminum and Tungsten don't exist in this specific dataset, but I have other datasets with this material distribution as well. From my understanding, nnUNet basically ignores the missing labels (excluding the final dice calculation which always ends up as NaN due to missing labels).

# Labels
    "labels": {
        "background": 0,
        "steel": 1,
        "titanium": 2,
        "aluminum": 3,
        "tungsten": 4,
        "brass": 5
    }

# Dataset Global Class Distribution
Label 0: 2,783,707,340 voxels (99.951214%)
Label 1: 1,275,394 voxels (0.045794%)
Label 2: 39,547 voxels (0.001420%)
Label 3: 0 voxels (0.000000%)
Label 4: 0 voxels (0.000000%)
Label 5: 43,783 voxels (0.001572%)

Using nnUNet as is learns and handles steel very well. For some reason, the models generally learn titanium (IoU: ~0.9), but struggle with brass (IoU: ~0.7). I am trying to get better scores for titanium and brass, ideally over IoU of 0.94.

Researching class imbalance, it seems like there's a lot of loss functions options. I have two main questions:

1. Is there a prefered/best loss function to handle class imbalance like this?
2. Am I implementing a Weighted Cross-Entropy + Dice correctly? (Feel free to roast me) It seems like this might be a good way to handle class imbalance. My other idea was to simply try other loss functions or combinations.

class nnUNetTrainer_WCE(nnUNetTrainer):
    def _build_loss(self):
        if self.label_manager.has_regions:
            raise ValueError("Bad luck, WCE does not support regions yet")
        else:
            # Weighted Cross-Entropy + Dice
            # class frequencies from your dataset [bg, steel, titanium, alluminum, tungsten, brass]
            freqs = [0.99951214, 0.00045794, 0.00001420, 0.0, 0.0, 0.00001572]
            # median of nonzero target-class frequencies (steel, titanium, brass)
            med = float(np.median([freqs[i] for i in (1, 2, 5)]))
            # inverse frequency weights
            inv_weights = [med / f if f > 0 else 0.0 for f in freqs]
            ce_weight = torch.tensor(inv_weights, dtype=torch.float32, device=self.device)

            loss = DC_and_CE_loss(
                {
                    "batch_dice": self.configuration_manager.batch_dice,
                    "smooth": 1e-5,
                    "do_bg": False,
                    "ddp": self.is_ddp,
                },
                {
                    "weight": ce_weight,
                },
                weight_ce=1,
                weight_dice=1,
                ignore_label=self.label_manager.ignore_label,
                dice_class=MemoryEfficientSoftDiceLoss,
            )

        if self._do_i_compile():
            loss.dc = torch.compile(loss.dc)

        # we give each output a weight which decreases exponentially (division by 2) as the resolution decreases
        # this gives higher resolution outputs more weight in the loss

        if self.enable_deep_supervision:
            deep_supervision_scales = self._get_deep_supervision_scales()
            weights = np.array([1 / (2**i) for i in range(len(deep_supervision_scales))])
            if self.is_ddp and not self._do_i_compile():
                # very strange and stupid interaction. DDP crashes and complains about unused parameters due to
                # weights[-1] = 0. Interestingly this crash doesn't happen with torch.compile enabled. Strange stuff.
                # Anywho, the simple fix is to set a very low weight to this.
                weights[-1] = 1e-6
            else:
                weights[-1] = 0

            # we don't use the lowest 2 outputs. Normalize weights so that they sum to 1
            weights = weights / weights.sum()
            # now wrap the loss
            loss = DeepSupervisionWrapper(loss, weights)

        return loss

Thanks,
Vincent