Shrut Dalwadi

The question

A CNN is translation equivariant, but equivariance is not invariance. In practice we lean on data augmentation and max-pooling, and both only cover distortions close to what training saw. A Spatial Transformer Network takes a different approach: a small localisation network predicts an affine transform and warps the input into a canonical pose before the classifier sees it, learned entirely from the classification loss with no pose labels.

So there are really two questions here: how far does that robustness actually stretch as the input gets harder, and what does the localisation network end up learning to do?

Setup

MNIST, the same CNN backbone and hyperparameters (20 epochs, Adam, lr 0.001, batch 64) with and without an STN in front. Four test sets with escalating geometric distortion: clean, moderate (±20° rotation), high (±40°), extreme (±60°), with matching translation and scale ranges. Training used 40% random moderate augmentation.

Results

Test set	CNN	STN + CNN	Delta
Clean	99.41%	99.44%	+0.03 pp
Moderate	98.32%	99.01%	+0.69 pp
High	85.42%	94.43%	+9.01 pp
Extreme	47.92%	69.20%	+21.28 pp

Test accuracy across four distortion levels; the gap grows from nothing on clean data to about 21 points at the extreme level. — loader1 to loader4 on the x-axis are the clean, moderate, high and extreme test sets.

Heavily distorted digits before and after the STN’s learned warp. — Distorted inputs, and the same digits after the STN’s learned warp.

What I learned

On clean data the STN adds nothing, and that is the point: the advantage grows with distortion, from +0.03 percentage points on clean MNIST to +21.28 under extreme distortion. The warped outputs show what the STN learned to do: center shifted digits, rescale them, and partially correct rotations it never saw in training.