Blog Details

We were building a multi-class staging classifier for a microscopy imaging pipeline. Its purpose was to sort tissue samples into five ordinal severity stages based on cellular morphology. The dataset had a common shape in clinical imaging projects: a dominant early-stage class (roughly 68% of samples), a healthy middle tier, and two advanced-stage classes that together barely cracked 8% of the total images. Standard practice suggested this was solvable with the right backbone and enough training time.

It wasn't. We cycled through ResNet-18, ResNet-34, and ResNet-50 variants, adjusted learning rates, swapped optimizers, and layered in the usual augmentation tricks. Despite these efforts, the model consistently predicted the majority stage for nearly everything. Validation accuracy looked respectable at a glance, around 71%, but macro-averaged F1 told a different story, hovering near 0.40. This post details what we tried, what quietly failed, and why switching the loss function did more for this classifier than any architecture change we attempted.

A Five-Stage Grading Task Sitting on a Skewed Dataset

The clinical task was ordinal, but we treated it as multi-class for this project phase: stage 1 through stage 5. Each stage was defined by increasingly abnormal cell morphology visible under magnification. Pathology labeling is expensive and slow; this explains why the dataset skewed so heavily toward early-stage samples, which are easier and faster to source and confirm. Advanced-stage cases were also rarer in the underlying population, not just underrepresented in collection. This created a biological imbalance layered on top of a sampling one.

Cross-entropy loss, in its plain form, treats every misclassified prediction with equal weight, regardless of how common or rare the true class is. When 68% of the training signal points to 'stage 1,' gradient updates reinforcing stage 1 predictions dominate the loss surface early and often. The network finds a shortcut: predict the majority class with high confidence. This makes it right most of the time on paper, even though it's clinically useless for the cases that matter most.

We noticed this first in the confusion matrix, rather than in the aggregate accuracy number. This was a lesson in itself. Aggregate accuracy sat at a believable 71%, but the confusion matrix showed stage 4 and stage 5 samples routed almost entirely into stage 1 and stage 2 buckets. A staging tool that cannot distinguish advanced disease from early disease is not a marginally imperfect tool; it is the wrong tool entirely.

Throwing Architecture at a Data Problem

Our first instinct, a common one, was that the network wasn't expressive enough to separate the subtle morphological differences between adjacent stages. We moved from ResNet-18 to ResNet-34 and then to ResNet-50, assuming more capacity would help the model find finer-grained features distinguishing stage 3 from stage 4. We also tried initializing from ImageNet-pretrained weights versus training from scratch, theorizing that low-level texture features transfer reasonably well to microscopy despite the domain gap.

Alongside backbone changes, we adjusted standard hyperparameters: learning rate schedules, batch size, weight decay, and dropout placement before the final classification head. We tried both Adam and SGD with momentum, swept learning rates across two orders of magnitude, and extended training with cosine annealing to ensure the model wasn't simply under-trained. None of these efforts changed the fundamental shape of the confusion matrix.

This point is worth stating plainly because it's a common trap: architecture search feels productive. New numbers appear with every run, creating the illusion of progress even when the underlying failure mode hasn't moved. We spent roughly three weeks on this approach before stepping back and asking a different question. It was not 'is the model big enough,' but 'is the loss function even asking the model to care about the minority classes?'

• ResNet-18, ResNet-34, and ResNet-50 backbones, both pretrained and from-scratch
• Learning rate sweeps across 1e-4 to 1e-2 with cosine annealing and step decay
• Dropout rates from 0.2 to 0.5 before the final fully connected layer
• Standard augmentation: rotation, flip, color jitter, random crop

“Cross-entropy doesn't fail loudly; it fails quietly by finding the laziest correct answer available. A bigger network just finds that same lazy answer faster.”

What Didn't Work and Why It Looked Like It Might

Before touching the loss function itself, we tried more conventional imbalance remedies. These are cheaper to test than rewriting training objectives. We tried naive oversampling of minority stages, duplicating stage 4 and stage 5 images until the training set looked artificially balanced. This helped marginally, nudging macro-F1 up by a few points. However, it also caused visible overfitting on the duplicated samples, since microscopy images from the same slide region share enough texture that oversampling amounted to memorization rather than learning generalizable features.

We also tried simple inverse-frequency class weighting within standard cross-entropy, a textbook first move. It moved the needle slightly, shifting some predictions away from the majority class. However, it overcorrected in a different way: the model began over-predicting stage 5 on genuinely ambiguous early-stage images, trading one kind of systematic error for another. The weighting was static across all samples of a class, regardless of how confidently or poorly the model was already handling individual images. This turned out to matter more than we expected.

The pattern across all these attempts was consistent: each fix addressed the symptom of imbalance in the data distribution. Yet, none addressed how the loss function scored easy versus hard examples during training. That distinction is subtle enough that it's easy to miss when focused on class counts rather than on the gradient behavior those counts produce.

Focal Loss and Stratified 5-Fold Validation, Together

Focal loss modifies cross-entropy by adding a modulating term. This term down-weights well-classified examples and keeps the gradient signal focused on hard, misclassified ones, regardless of their class. In practice, once the model becomes confidently correct on easy majority-class images, those examples stop contributing much to the loss. Harder minority-class and boundary-case images then continue driving updates. We used a gamma of 2 combined with per-class alpha weighting set inversely proportional to class frequency. This allowed us to address both the easy-example problem and the raw frequency imbalance at once, rather than picking one lever.

This alone would not have been trustworthy without changing how we validated the model. A single train/validation split with this level of imbalance can accidentally place most rare-class examples in one fold, making results look better or worse than they truly are. We moved to stratified 5-fold cross-validation, ensuring every fold preserved the same class proportions as the full dataset. We reported metrics as the mean and standard deviation across all five folds, rather than from a single run. This mattered enormously for advanced-stage classes with only a few dozen total examples, where a single unlucky split could make the model appear to have solved the problem when it had simply gotten lucky.

The combination is what actually moved results, not either piece alone. We reran the earlier class-weighted cross-entropy experiment under the same stratified 5-fold protocol. This confirmed it genuinely underperformed focal loss, not just under one unlucky split. That gave us confidence the gain came from the loss function's behavior, not from a validation artifact.

• Gamma = 2, tuned via grid search across 0.5 to 3
• Alpha weights set inversely proportional to class frequency, then fine-tuned per fold
• Stratified 5-fold cross-validation preserving per-stage class ratios in every split
• Metrics reported as mean ± standard deviation across folds, not single-run numbers

Attention Mechanisms and Early Stopping Started Fighting the Loss Function

Once focal loss was in place, we layered in a lightweight channel attention block (a squeeze-and-excitation style module) between the backbone's later stages. We expected it to help the model focus on specific morphological regions that distinguish adjacent stages. On the majority class, it helped, sharpening feature responses around nucleus density and staining patterns. However, on the minority classes, it initially made results less stable across folds. This happened because attention weights learned early in training, while the model was still dominated by easy majority-class gradients, baked in a bias toward majority-class-relevant regions before focal loss could redirect learning toward harder examples.

Early stopping compounded this in an unanticipated way. Our early stopping criterion monitored validation loss. Under focal loss, this metric can plateau early even while macro-F1 on minority classes is still improving. This occurs because the loss, by design, de-emphasizes the now-easy majority class, and harder examples improve more slowly and noisily. We were stopping training runs 8 to 12 epochs before minority-class recall had actually converged, effectively undoing part of what focal loss was trying to accomplish.

The fix involved switching the early stopping metric from validation loss to macro-averaged F1 computed per fold. We also gave the attention module a warm-up period, training it with frozen weights for the first several epochs before unfreezing it. This sequencing, an attention warm-up followed by joint fine-tuning under focal loss, produced the most stable per-fold results we saw across the entire project. It's a reminder that individual components can be reasonable choices yet still interact badly if their combined behavior on imbalanced data is not checked.

Where the Numbers Landed and What It Tells Us Going Forward

The final configuration moved macro-F1 from 0.41 to 0.68 on average across folds. This configuration included a ResNet-34 backbone, channel attention with a warm-up schedule, focal loss with gamma 2 and inverse-frequency alpha, and macro-F1-based early stopping under stratified 5-fold validation. More importantly for the clinical use case, recall on the two advanced stages went from single digits to the 50-60% range. This is the difference between a tool that is clinically dead weight and one that is a genuinely useful second reader. Overall accuracy actually dropped slightly, from 71% to about 69%. This is exactly what we would expect when a model stops taking the majority-class shortcut, and it was a number we had to explain carefully to stakeholders who were anchored to accuracy as the headline metric.

What this project really taught us is that architecture tuning and loss function design solve different problems. It is easy to spend weeks on the wrong one if you are not looking at per-class metrics from the start. A bigger ResNet can find more complex features, but it cannot decide on its own that a rare class deserves more attention than a common one. That decision must be encoded somewhere: in the loss, the sampling strategy, or the validation protocol, and ideally in more than one of those places consistently.

We now treat per-class recall and macro-F1 under stratified cross-validation as non-negotiable reporting requirements on any imbalanced classification project, computer vision or otherwise, before any architecture comparison is considered meaningful.