Context versus foreground ratio

A green foreground object inside grey background anatomy, shown inside nested crop windows of 32, 48, 64, 80 and 96 voxels. The left of the diagram is labelled lower context, higher FBR and a wider FBR range while training; the right is labelled higher context, lower FBR and a narrower FBR range. — The trade-off. Patch size grows left to right: context grows with it, the proportion of the patch that is foreground shrinks, and so does the range of foreground ratios the network ever encounters while training.

Abstract

Modern 3D medical image segmentation is done with a sliding window, because GPU memory will not hold a whole volume. That presents a trade-off between the amount of global context the network sees at once and the proportion of foreground voxels available in each training sample. Unets are known to perform worse with low context, but enlarging the context comes at the cost of heavy class imbalance between background (typically very large) and foreground (much smaller).

We study the behaviour of a transformer-based model (UNETR) and an attention-gated model (attention-Unet) alongside a vanilla Unet across this trade-off, using a controlled synthetic data set and the spleen segmentation task from the Medical Segmentation Decathlon. Beyond showing that all three prefer more global context rather than a bigger foreground-to-background ratio, we find that UNETR and attention-Unet are markedly less robust than the vanilla Unet to drifts between the foreground ratios seen in training and those seen at test time.

What we found

Context wins, for all three architectures. Larger patches score better on both tasks despite the worse class balance. The gain is largest where the architecture is weakest: attention-Unet climbs from 0.087 to 0.635 mean Dice on spleen between 32³ and 96³.
UNETR and attention-Unet are the fragile ones. On test cases whose foreground ratio falls outside the training range, their Dice drops by 0.26 to 0.84. A vanilla Unet in the same setting loses 0.014 to 0.023 — more than an order of magnitude less.
Patch size decides how much drift a model tolerates. It sets not just the mean foreground ratio but the whole range the network is exposed to, and the two architectures with the more generic inductive biases depend on having seen a ratio before.
The accuracy ranking and the robustness ranking are not the same. At 96³ UNETR matches the vanilla Unet in-range, at 0.994 Dice on the synthetic task, and loses fifteen times as much outside it.

Where the trade-off comes from

The case for a small patch is class balance: in the spleen volumes of the Medical Segmentation Decathlon, at best one voxel in a hundred belongs to the organ, and a tighter crop around it is far less skewed. The case for a large patch is that segmentation is partly a question of where, and a network looking through a 32-voxel window cannot see where it is. Patch size raises context and lowers the foreground-to-background ratio, and no setting maximises both.

The range matters as much as the ratio

There is a third consequence, and it turned out to be the one that mattered. Training draws patches half centred on foreground and half on background, so each patch size induces its own distribution of foreground ratios. We logged every ratio actually drawn. Widening the patch does not merely lower the mean — it collapses the spread.

Synthetic task, vanilla Unet, from the logged training patches. Bars are the full observed range, not a quantile.

Read down the bars and the consequence is plain. A network trained at 32³ has seen every foreground ratio there is. One trained at 96³ has seen about a tenth of that span, and at test time 249 of the 300 test cases fall outside it. Whether that matters turns out to depend entirely on the architecture.

The setup

Task	Data	Split	Foreground
Synthetic	100 volumes of 96³ voxels, uniform noise background at variance 0.8	70 train / 30 validation, plus an independent 100-volume test set	One sphere; radius 25–35 voxels in training, 5–48 at test, deliberately reaching past what training showed
Clinical	Task09_Spleen, Medical Segmentation Decathlon, chosen for its wide FBR range	26 train / 5 validation / 10 test	The organ; at best one voxel in a hundred

The architectures

Three models, implemented with MONAI. Vanilla Unet and attention-Unet share every hyperparameter they can — widths (16, 32, 64, 128, 256), strides (2, 2, 2, 2) — so the only difference between them is the attention gating on the skip connections. That pairing is what isolates the effect.

Model	What it changes
vanilla-Unet	Fully convolutional encoder–decoder with identity skips, 2 residual units, batch norm
UNETR	The encoder is replaced by a transformer: feature size 16, hidden size 768, 3072 MLP dimensions, 12 heads, perceptron position embedding, instance norm, no dropout
attention-Unet	Attention gates on the skip connections; otherwise identical to the vanilla Unet

Everything else is held fixed so nothing else can explain the differences: Adam at 1e−4 with cosine decay, batch size 2, no augmentation beyond the random crop and no post-processing. Five patch sizes (32, 48, 64, 80, 96), symmetric in 3D, three seeds per cell.

Watching it happen

Three clips, each looping. The first shows what the trade-off is; the other two show what it does to the three architectures. Every clip sweeps the patch size from 32³ up to 96³ and back.

The panels are dense. Use the fullscreen control on any clip to read the scores.

What the network actually seesOne synthetic volume. Because training centres half its patches on foreground and half on background, both are shown: solid outline for the foreground-centred patch, dashed for the background-centred one. At 32³ the foreground patch is solid object (FBR 1.000) and the background patch nearly solid noise (0.035) — neither shows the network a boundary in context. At 96³ the patch is the whole volume, the two coincide, and the ratio settles at 0.128. The right-hand panel is the chart above, with the current patch size highlighted.

Synthetic task — a sphere in noiseEvery held-out test case as a point: foreground ratio across, Dice up, for the three architectures side by side. The shaded band is the training FBR range at the current patch size; filled points fall inside it, hollow grey points are cases the network never saw the ratio of. Watch the band narrow as the patch grows, and watch the hollow points on its left. The vanilla Unet holds them near 1.0 throughout. UNETR and attention-Unet drop them off a cliff: the same networks that score 0.99 in-range score near zero just outside it.

Spleen CT — the clinical caseNote the horizontal axis: the whole test set lives below a foreground proportion of 0.02, which is the class-imbalance problem in one glance. Every panel reads “nothing outside”, and that is not a bug — all ten test volumes sit inside the range training covered, which is why the abstract reports a plain mean here rather than a split. The architecture ordering survives anyway.

Neither of the behaviour clips runs a model. The weights behind the abstract were not archived, so the clips replay the measured per-case metrics; the first clip shows regenerated phantoms and logged training ratios. This is a limit of what survived, not a design choice.

Results

Mean Dice over three seeds per cell, as the patch size grows from 32³ to 96³. Lines that climb are architectures that need the context.

vanilla-Unet UNETR attention-Unet

Synthetic task	32³	48³	64³	80³	96³
vanilla-Unet	0.982 (0.024)	0.986 (0.016)	0.994 (0.002)	0.994 (0.001)	0.994 (0.001)
drop outside band	—	—	−0.023	−0.014	−0.018
UNETR	0.677 (0.351)	0.644 (0.391)	0.958 (0.121)	0.992 (0.009)	0.994 (0.003)
drop outside band	—	—	−0.734	−0.370	−0.276
attention-Unet	0.633 (0.343)	0.664 (0.373)	0.948 (0.133)	0.993 (0.005)	0.970 (0.016)
drop outside band	—	—	−0.837	−0.300	−0.256

Mean (sd) Dice on the test cases whose foreground ratio falls inside the training range, and beneath each, the change in mean Dice on the cases outside it. Blank where the training range covers the whole test set, so nothing falls outside.

Spleen CT	32³	48³	64³	80³	96³
vanilla-Unet	0.721 (0.131)	0.908 (0.045)	0.928 (0.041)	0.922 (0.048)	0.932 (0.043)
UNETR	0.482 (0.158)	0.766 (0.179)	0.800 (0.187)	0.852 (0.116)	0.915 (0.042)
attention-Unet	0.087 (0.052)	0.102 (0.061)	0.384 (0.314)	0.582 (0.326)	0.635 (0.327)

Mean (sd) Dice over the ten held-out volumes. No split here: every test volume's foreground ratio sits inside the range training covered.

Read the drop rows of the first table and the pattern is hard to miss. Every architecture improves with context, and by 96³ the vanilla Unet and UNETR are indistinguishable in-range at 0.994. Step outside the range and they are not remotely comparable: −0.018 against −0.276.

What to do with this

Pick the patch size as large as memory allows. Across three architectures and two tasks, context beat class balance every time. The class-imbalance argument for small patches did not survive contact with the data.
Then check what FBR range that leaves you. Patch size sets the range as a side effect, and the range is what determines how much test-time drift the model survives. It costs nothing to log the foreground ratio of every training patch, and it tells you which part of the test distribution you are extrapolating into.
Report performance against foreground ratio, not just a mean. A model at 0.99 in-range and near zero just outside it averages to something respectable and is not a respectable model.
Treat the vanilla Unet as a real baseline. The more elaborate architectures matched it at best in-range and were an order of magnitude more fragile outside it. For deployment across scanners and sites, that trade is rarely worth making.

Citation

@inproceedings{kamath2022context,
  title     = {How do 3D image segmentation networks behave across
               the context versus foreground ratio trade-off?},
  author    = {Kamath, Amith and Suter, Yannick and You, Suhang and
               M{\"u}ller, Michael and Willmann, Jonas and
               Andratschke, Nicolaus and Reyes, Mauricio},
  booktitle = {Medical Imaging meets NeurIPS Workshop,
               36th Conference on Neural Information Processing Systems},
  year      = {2022}
}