BraTS: U-Net vs. TransUNet

The problem with most comparisons

In many segmentation codebases, an architecture change arrives bundled with changes to cropping, slice sampling, intensity normalization, and the loss. When the Dice score moves, there is no way to say which change moved it. So the reported conclusion — "the transformer backbone wins" — may actually be a statement about preprocessing.

I built a single training entrypoint that instantiates three model families over the same data path, the same split rule, and the same metric:

a classic U-Net baseline with BatchNorm double-convolution blocks;
a configurable "flex" U-Net — tunable depth, base width, normalization (Batch/Group/Instance/Layer/none), bilinear vs. transposed-conv upsampling, optional residual blocks;
a TransUNet variant with a hybrid ResNet50 + ViT-B/16 encoder and a convolutional decoder, plus an optional lightweight self-attention block that refines one skip connection before concatenation.

Each checkpoint stores the full architecture configuration and the crop coordinates used during training, so evaluation scripts reconstruct the exact model without manual bookkeeping.

Data path

Every subject carries four co-registered MRI modalities (FLAIR, T1, T1-CE, T2) and a tumor mask; inputs are 2D axial slices, x ∈ R^4×H×W, with four classes after remapping the BraTS label 4 to 3.

Automatic brain-region cropping — one crop box computed from non-zero voxels across subjects and modalities, disabled automatically if the bounds are degenerate.
Slice indexing that discards empty slices and stores a foreground flag per slice, later used for oversampling.
Per-slice z-score intensity normalization, with an optional non-linear contrast mapping to test whether stronger normalization helps.
Weighted foreground oversampling (WeightedRandomSampler, α = 5) so tumor-bearing slices are seen more often, against the severe tumor/background skew.
Fixed seed and validation ratio so every variant is compared under the same split.

Objective: class-weighted cross-entropy plus a present-class foreground Dice term — Dice is computed only over foreground classes that actually appear in the ground truth of the batch, because naive Dice is inflated by absent classes. The Dice weight is warmed up over the first few epochs for stability. Training used Adam, mixed precision, and gradient accumulation under a 15-epoch budget, from scratch.

Results

Best checkpoints on the BraTS validation slices. Foreground (FG) Dice is averaged over labels 1–3, and only over slices where the class is present. Same split rule and metric for every row.
Method	Input res.	mean CE ↓	FG Dice ↑	Dice 1	Dice 2	Dice 3
Original U-Net	crop	0.0437	0.7326	0.6826	0.6800	0.8351
Flex U-Net (norm + depth only)	crop	0.3850	0.6323	0.5395	0.6475	0.7100
Flex U-Net (loss improved)	crop	0.1952	0.7033	0.6048	0.7100	0.7952
TransUNet baseline (R50-ViT-B/16)	320²	0.0455	0.7575	0.6750	0.7652	0.8322
TransUNet (skip refinement + aug)	224²	0.0411	0.7347	0.6493	0.7455	0.8094

What the numbers actually say

1. Loss and sampling design matter as much as the backbone

Making the U-Net deeper and swapping BatchNorm for GroupNorm — with the loss untouched — hurt performance badly (0.6323 vs. 0.7326 for the plain baseline, with mean CE rising to 0.3850). More capacity without a matching objective got stuck in a worse regime. Reworking only the objective and sampler on that same architecture — present-class foreground Dice with warmup, class weighting, foreground oversampling — recovered roughly +7 Dice points (to 0.7033) and cut mean CE by about 0.19.

2. The transformer's advantage is local to one subregion

TransUNet took the best overall FG Dice (0.7575, about +2.5 points over the U-Net), but not uniformly. On the compact structures the tuned U-Net is competitive or better (Dice 3: 0.8351 vs. 0.8322). Essentially all of the gain sits in label 2 — the more spatially extended, diffuse subregion — at 0.7652 vs. 0.6800. Global self-attention helps where long-range context is the binding constraint, and not much elsewhere. Reporting only the aggregate FG Dice would have hidden this.

3. Better calibration can cost Dice

My decoder-side skip refinement with stronger augmentation improved mean cross-entropy (0.0455 → 0.0411) while lowering FG Dice (0.7575 → 0.7347). The modified model is better calibrated on average but more conservative on small lesions — which is exactly what hurts on rare classes. Two reasonable metrics, opposite verdicts.

4. Class imbalance is still the binding constraint

Every method does best on the largest structure (label 3) and worst on label 1, tracking the foreground/background skew of BraTS slices rather than anything architectural.

← Back to home