Confidence-Guided Elastic Augmentation

Standard elastic augmentation reduced this transformer segmentation model's Dice score by 75%. Making the augmentation confidence-aware recovered 81% of that loss.

---

The finding

Elastic deformation has been standard practice in medical image segmentation since the original U-Net paper. For CNNs it is close to free accuracy.

Applied to Swin-UNet — a segmentation model built entirely from shifted-window attention, no convolutions — on the 18-volume Synapse benchmark, it did the opposite:

Configuration	Mean Dice ↑	HD95 (mm) ↓
No augmentation	0.755	30.9
Standard elastic augmentation	0.185	99.5
Confidence-guided augmentation	0.673	32.7

A plausible mechanism: uniform spatial warping disrupts the absolute positional relationships that window attention relies on. A CNN's local receptive fields absorb this; a pure transformer learning long-range dependencies from 18 volumes does not. This is a hypothesis consistent with the ablation, not something the experiments isolate directly.

The response is not to drop augmentation but to stop applying it uniformly. Scaling deformation per-pixel by an EMA teacher's confidence map — deforming uncertain regions hard, leaving well-learned regions intact — recovers most of the lost performance at no parameter cost.

---

Method

Three modules attach to the training loop. The Swin-UNet architecture is unchanged.

EMATeacher — an exponential moving average copy of the student (momentum 0.99) used purely as a confidence oracle. One forward pass returns both logits and a per-pixel confidence map (max softmax probability). Computing them separately triples the cost of a training step; the single-pass version costs roughly 50% overhead.

ConfidenceAugmentor — builds a smooth random displacement field, then scales its magnitude per-pixel by a confidence-derived weight:

aug_mode	weight	intent
none	—	no deformation, true baseline
standard	uniform	conventional elastic augmentation
hard	1 − C(x,y)	deform where the model is uncertain
easy	C(x,y)	deform where the model is confident

hard and easy encode opposing hypotheses — hard-example mining versus overconfidence regularisation. The ablation tests both rather than assuming one.

ConfidenceWeightedLoss — adds a term penalising pixels where the teacher is confident and wrong. In this setting it hurt rather than helped; see the ablation below.

---

Ablation

Six configurations. Identical seed, hardware, data split, and architecture — only the training scheme differs. Synapse multi-organ CT, 18 train / 12 test volumes, 100 epochs (60 for E), A100-SXM4 80GB.

Run	aug_mode	λ	Mean Dice	HD95	Δ vs B
A	none	0.0	0.755	30.9	+0.570
B	standard	0.0	0.185	99.5	—
C	hard	0.5	0.584	37.1	+0.398
D	easy	0.5	0.224	83.5	+0.039
E	standard	0.5	0.164	105.9	−0.022
F	hard	0.0	0.673	32.7	+0.487

Direction matters enormously. Hard-mode (0.673) versus easy-mode (0.224) is a threefold gap. Deforming regions the model has not yet learned is productive; deforming regions it already handles is barely better than deforming at random.

The loss term backfired. Run F beat Run C by 0.089 Dice, and Run E shows the term does not help in isolation either. Reported as-is rather than dropped from the study — see Known issues.

For reference, the published Swin-UNet baseline is 0.794 Dice / 21.55 HD95. Our unaugmented run reaching 0.755 suggests the gap is training budget and GPU non-determinism rather than a broken pipeline; the authors themselves document result variation across GPU types.

---

Repository

confidence/
  ema_teacher.py       EMA teacher / confidence oracle
  conf_augmentor.py    Spatially-adaptive elastic augmentation
  conf_loss.py         DiceCE + confidence correction term
train.py               Training entrypoint
trainer.py             Training loop
docs/
  SETUP.md             Install, data, and all six run commands
  architecture.png

Quickstart, after following docs/SETUP.md:

python train.py \
    --aug_mode hard --conf_warmup 20 --conf_lambda 0.0 \
    --batch_size 24 --max_epochs 100

---

Known issues

• Teacher/label spatial misalignment. trainer.py computes teacher logits on the un-augmented image, then passes them to a loss that compares against the elastically deformed label. The two are on different spatial grids, so the "teacher is confidently wrong" mask partly fires on pixels that moved rather than pixels the teacher got wrong. This affects every run with conf_lambda > 0 (C, E) and is a likely contributor to the loss term underperforming. Diagnostic instrumentation and a fix are in progress on a branch; the results above are from the code as published here.
• Single dataset, single architecture. Whether the augmentation collapse generalises beyond the 18-volume regime is untested.
• λ = 0.5, τ = 0.7, and the 20-epoch warmup were set by reasoning rather than search.
• An earlier revision of ConfidenceAugmentor gated only the confidence weighting behind the warmup counter rather than deformation itself, so none runs silently applied uniform elastic and produced weights byte-identical to the standard baseline. Caught by hashing checkpoints across runs — worth checking if you adapt this code.

---

Credits

Originated as a CS584 (Machine Learning) project at Illinois Institute of Technology, Spring 2026, by Sre Ganesh Subramaniam Ramanathan and Sabarish Dhalayan.

Built on Swin-UNet:

@inproceedings{cao2022swinunet,
  title     = {Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation},
  author    = {Cao, Hu and Wang, Yueyue and Chen, Joy and Jiang, Dongsheng
               and Zhang, Xiaopeng and Tian, Qi and Wang, Manning},
  booktitle = {ECCV Workshops},
  year      = {2022}
}

MIT licensed. The base Swin-UNet code carries its own license — see LICENSE.