Probabilistic Spatial Transformer Networks

Pola Schwöbel1

Frederik Warburg1

Martin Jørgensen2

Kristoffer H. Madsen1, 3Søren Hauberg1

1

Section for Cognitive Systems, DTU Compute, Technical University of Denmark, Copenhagen, Denmark
Machine Learning Research Group, Department of Engineering Science, University of Oxford, Oxford, UK
3
Danish Research Centre for Magnetic Resonance, Centre for Functional and Diagnostic Imaging and Research,
Copenhagen University Hospital Hvidovre, Hvidovre, Denmark

Abstract

Negative log-likelihood

Spatial Transformer Networks (STNs) estimate image transformations that can improve downstreamtasks by ‘zooming in’ on relevant regions in an image. However, STNs are hard to train and sensitive
to mis-predictions of transformations. To circumvent these limitations, we propose a probabilistic
extension that estimates a stochastic transformation rather than a deterministic one. Marginalizing
transformations allows us to consider each imageat multiple poses, which makes the localization
task easier and the training more robust. As an
additional benefit, the stochastic transformations
act as a localized, learned data augmentation that
improves the downstream tasks. We show acrossstandard imaging benchmarks and on a challenging
real-world dataset that these two properties leadto improved classification performance, robustness
and model calibration. We further demonstrate thatthe approach generalizes to non-visual domains byimproving model performance on time-series data.

1

STN

A

A

B
B
C
C

A

P-STN

arXiv:2004.03637v2 [cs.LG] 15 Jun 2022

2

A

B
B

C
C

Figure 1: The Probabilistic Spatial Transformer Network(P-STN) marginalizes over a distribution of possible input
transformations. By ‘looking in multiple places’ we hope tostabilize the brittle nature of the regular spatial transformer:
The P-STN loss landscape is significantly more smooth andwith fewer local minima compared to the STN.

INTRODUCTION
difficult to train and often do not live up to their promise.

The Spatial Transformer Network (STN) [Jaderberg et al.,2015] predicts a transformation on input data in order tosimplify a downstream task. For example, a neural network
might benefit from e.g. ‘zooming in’ on relevant parts ofan image, remove unwarranted image rotations, or timenormalize sequence data before making predictions. In principle, this can improve robustness, interpretability and efficiency of the model. However, in practice, the situation
is not as ideal. Both at training and test time, the STN issensitive to small mis-predictions of transformations. For
example, if the STN zooms in on the wrong part of an image,then the signal is lost for the downstream task, e.g. see crop
A and C in Fig. 1. The empirical impact is that STNs are

From a probabilistic perspective, this sensitivity has an obvious solution: we should estimate the posterior over theapplied transformation and marginalize accordingly. This
amounts to ‘trying many different transformations’, andshould improve robustness. It is exactly this approach weinvestigate.
STNs consist of two parts. A localization network performs
the transformation task, i.e. it estimates the transformation
parameters θ for a given image I and applies the corresponding transformation Tθ (I). A standard neural network
performs the downstream task on the transformed image, i.e.
computing p(y|Tθ (I)). Since we are concerned with classi-

Accepted for the 38th Conference on Uncertainty in Artificial Intelligence (UAI 2022).

fication tasks, we will refer to the latter as the classifier, butnote that the approach generalizes to other tasks.
In our probabilistic STN (P-STN), we estimate a distribuRtion over transformations that we marginalize: p(y|I) =p(y|Tθ (I))dθ. We approximate this intractable integral
via Monte Carlo, i.e. we sample transformations. Thesetransformation samples produce different transformed versions of the input image, {Tθs (I)}s=1...S . The classifier
makes predictions on all samples, and we aggregate thepredictions. Figure 2 shows the model architecture.
We hypothesize that marginalizing image transformation
has benefits for both parts of the model. For the localization
network, our model gets to ‘try many different transformations’ through random sampling. This should improve
the localization. Secondly, the classifier now gets presented
with different transformed versions of the input image
through Monte Carlo samples {Tθs (I)}s=1...S . Interestingly,
this corresponds to a type of data augmentation, whichshould improve classification.
We verify these hypotheses by making the following contributions:
1. We develop the Probabilistic Spatial Transformer; a hierarchical Bayesian model over image transformations.
2. We perform variational inference to fit the transformation model as well as downstream model end-to-end,
using only label information.
3. We experimentally demonstrate that our model
achieves better localization, increased classification
accuracy (resulting from learned per-image data
augmentation) and improved calibration.

2

RELATED WORK

Spatial transformer networks apply a spatial transformation to the input data as part of an end-to-end trained model[Jaderberg et al., 2015]. The transformation parameters
are estimated from each input separately through a neuralnetwork. Most commonly, STNs implement simple affine
transformations, such that the network can learn to zoom
in on relevant parts of an image before solving the task
at hand. STNs have shown themselves to be useful forboth generative and discriminative tasks, and have seen
applications to different data modalities [Jaderberg et al.,2015, Detlefsen and Hauberg, 2019, Detlefsen et al., 2018,
Shapira Weber et al., 2019, Sønderby et al., 2015, Lin andLucey, 2016, Kanazawa et al., 2016]. We propose a probabilistic extension of this idea, replacing the usual likelihood
maximization with marginalization over transformations.
Bayesian deep learning aims to solve probabilistic computations in deep neural networks. Priors are put on weights
and marginalized at training and test time, often yielding

useful uncertainties in the posterior predictive. The required
computations are in general intractable, and approaches
differ mainly in the type of approximation to the weight posterior. Gal and Ghahramani [2016] propose to view dropout
as a Bernoulli approximation to the weight posterior (i.e.
randomly switching each weight on or off). The Laplace approximation [MacKay, 1992, Daxberger et al., 2021] places
a Gaussian posterior over a trained neural network’s weights.
Another generally successful way to obtain predictive uncertainties is to simply train an ensemble of models. Originally proposed as an alternative to Bayesian DL [Lakshminarayanan et al., 2017], the approach can be interpreted
in the Bayesian framework by interpreting the weights ofthe trained ensemble members as samples from a weight
posterior [Gustafsson et al., 2020]. Similar to our method,
Blundell et al. [2015] choose a variational approach witha simple Gaussian mean field posterior over weights. Our
approach differs from standard Bayesian DL in that weare not reasoning about distributions over neural network
weights p(w), but instead a subnetwork’s (i.e. the localizer’s) outputs p(θ). Drawing from the posterior over imagetransformations, we effectively recover data augmentation.
Data augmentation (DA) is a useful way to increase theamount of available data [LeCun et al., 1995, Krizhevsky
et al., 2012]. DA requires prior knowledge about the structure of the data: the target y is assumed to be invariant tocertain transformations of the observation I. Invariance assumptions are usually straight forward for natural images.
Thus, DA is common for image data, where the transformation family is often chosen to be rotations, scalings, andsimilar [Goodfellow et al., 2009, Baird, 1992, Simard et al.,2003, Krizhevsky et al., 2012, Loosli et al., 2007]. The general trend is that, beyond ‘intuitive’ data such as images,
gathering an invariance prior is difficult, and DA is often
hard to realize through manual tuning.
Learned data augmentation provides a more principledapproach to artificially extending datasets. Hauberg et al.
[2016] estimate an augmentation scheme from the training
data via pre-aligning images in an unsupervised manner.
The approach allows for significantly more complex transformations than the usual affine family, but the unsupervised
nature and the implied two-step training process render theapproach suboptimal. Similarly, Cubuk et al. [2019, 2020]
use reinforcement learning and grid search to learn data
augmentation schemes, but rely on validation data rather
than an end-to-end formulation.
Learning data augmentation end-to-end requires a loss function suitable for model selection, as we are effectively trying
to learn an inductive bias. Based on this realization, Van der
Wilk et al. [2018] learn DA end-to-end in Gaussian processes (GPs) via the marginal likelihood, a suitable loss
for model selection and thus invariance learning [MacKay,2003]. The marginal likelihood is hard to compute for NNs,so Schwöbel et al. [2022] extend this idea to NNs by consid-

Figure 2: The P-STN pipeline. From the observed image I, a distribution of transformations is estimated. Samples from this
distribution are applied to the observed image to produce augmented samples, which are fed to a classifier that averages
across samples. In the deterministic STN case, the localizer only computes one transformation θ(I), which can be thought ofas the maximum likelihood solution. Instead of the multiple transformation samples, we obtain a single Tθ (I) in this case.
ering a deep kernel model, i.e. a neural network with a GP in
the last layer. Benton et al. [2020] instead use the standard,
maximum likelihood loss and explicitly regularize towards
non-zero augmentations. Our model differs from existingdata augmentation approaches — learned and non-learned— in that we estimate local, i.e. per-image transformations
instead of a global augmentation scheme.

divided into subsets and an affine matrix is defined on eachcell c of such a tessellation. Each affine matrix Aθc induces
a vector field mapping each point x ∈ c to a new positionv θc : x 7→ Aθc x. These velocity fields are then integrated
to form a trajectory for each image point x
φθ (x; 1) = x +Z 1v θ (φ(x; τ ))dτ.

03

BACKGROUNDThe STN localiser module estimates a transformation θ(x)that transforms a coordinate grid and interpolates an image
accordingly. The classifier module takes the transformed image and computes p(y|Tθ (x)). Both the localizer and classifier are neural networks. The STN can be trained end-to-end
with only label information as long as the image transformations are parameterized in a differentiable manner.
Affine transformations are a simple class of transformations that can be differentiably parameterized. We limit
ourselves to the subset of affine transformations containing rotation, isotropic scaling and translation in x and y. In
two dimensions (and the corresponding three-dimensional
homogeneous coordinates), we thus learn θ = (r, s, tx , ty )which parameterizes the affine matrix
s · cos r −s · sin r tx
Aθ =  s · sin r s · cos r ty  ∈ R3×3 , s > 0. (1)001Since det(Aθ ) = s2 , the constraint s > 0 ensures invertibility and can be implemented as seen in Detlefsen et al. [2018].
In practice, the STN estimates well-behaved, non-collapsingtransformations without implementing the constraint explicitly. Tθ (I) is applied by transforming a grid of the target
image size by Aθ and interpolating the source image at theresulting coordinates (see Jaderberg et al. [2015] for details).
Diffeomorphic transformations (i.e. transformations thatare differentiable, invertible and possess a differentiable inverse) are more general than affine transformations, and are
not limited to the spatial domain. Freifeld et al. [2017] construct diffeomorphisms from continuous piecewise-affine
velocity fields as follows. The transformation domain Ω isGiven boundary and invertibility constraints [Freifeld et al.,2017], such a collection of affine matrices {Aθc }c⊂Ω defines
a diffeomorphic transformation T θ : x 7→ φθ (x, 1).
The libcpab library [Detlefsen, 2018] provides an efficient
implementation for this approach, specifically optimized
for use in a deep learning context where fast gradientevaluations are crucial. The author successfully employs
CPAB-transformations within a Spatial Transformer Network [Detlefsen et al., 2018].

4

PROBABILISTIC SPATIALTRANSFORMER NETWORKThe P-STN is a probabilistic extension of the STN, where
we replace the deterministic transformation θ(I) with a posterior over transformations p(θ|I). Figure 2 illustrates the
proposed pipeline. We assume observed data of the form
D = {yi , Ii }N
i=1 , where y is the target variable (e.g. classlabel), and I are observations of the covariates. For presentation purposes, we will consider the latter to be images, but
the approach applies to any spatio-temporal data.
4.1

THE MODEL

Recall that STNs are trained end-to-end for the downstreamtask using only label information. Thus, while we observe y,θ is a latent variable. We model it to be governed by a second
latent variable λ. λ is a precision parameter, effectively stopping the localization distribution (i.e. the amount of ‘data
augmentation’ we introduce) from collapsing. The necessity for non-collapsing augmentation is discussed in Benton
et al. [2020], Van der Wilk et al. [2018] and Schwöbel et al.
[2022].

I

λ

y

θ

We derive our lower bound using Jensen’s inequality
ZZ
log p(y, I) = log
p(y, I, θ, λ)dθdλ
(8)


ZZ
p(y, I, θ, λ)≥log
q(θ, λ)dθdλ(9)q(θ, λ)ZZp(y|I, θ)p(λ)p(I)=logp(θ|λ, I)q(λ)dθdλq(λ)

NFigure 3: A graphical representation of the model structure.
Grey nodes are observables and white are latents.

= Eq(θ,λ) log p(y|I, θ) + log p(I) − KL(q(λ)kp(λ)) .
{z}|classification loss

(10)
We wish to infer the latent variables in a Bayesian manner.
This entails computing the (log-)marginal likelihood of theobserved
ZZ
log p(I, y) = log
p(I, y, θ, λ)dθdλ.
(2)Thus, our evidence lower bound (ELBO) objective function (10), consists of two terms: a classification loss and a
KL-term controlling the distance of the approximate posterior to the prior. During inference, we can disregard log p(I)as it does not depend on parameters of interest.

We let the joint distribution factorize as (see Fig. 3)p(y, I, θ, λ) = p(y|I, θ, λ)p(I, θ, λ)= p(y|I, θ)p(θ|λ, I)p(λ)p(I).

(3)
(4)

Notice p(I) is unaffected by model parameters λ and θ, andin this sense can be specified without affecting the model.
The distribution over θ depends on observed covariates in
the following wayp(θ|λ, I) = N (θ|µ(I), 1/λ),(5)where µ(I) is a function parametrised by a neural network,
i.e. µ(I) := µΦ (I) for model parameters Φ. The prior over
λ is a Gamma distribution, i.e.
p(λi ) = Γ(α0 , β0 ).

(6)

We note here that there is one λi associated toeach observation,and they are assumed to factorize:
QNp(λ) =
i=1 p(λi ). This choice of conjugate priors
for variance estimation is similar to [Stirn and Knowles,
2020, Takahashi et al., 2018, Detlefsen et al., 2019]. Finally,
we assume that, conditional on I and
QNθ, we have marginal
independence in y, i.e. p(y|I, θ) = i=1 p(yi |Ii , θi ).
4.2

4.3

INFERENCE

The choice of variational posterior implies the following forthe classification loss
Eq(θ,λ) log p(y|I, θ)
(11)
ZZ
=
log p(y|I, θ)q(θ, λ)dθdλ
(12)
ZZ
=
log p(y|I, θ)p(θ|λ, I)q(λ)dθdλ(13)ZZ= log p(y|I, θ) N (θ|µ(I), λ)Γ(λ|α, β(I))dλdθZ= log p(y|I, θ)t2α (θ|µ(I)), β(I)(14)α )dθ.
Here t denotes a scaled and location-shifted Student’s tdistribution with mean µ(I), scaling β, and α degrees of
freedom. For clarity, the marginalized q(θ) is t-distributed.
Here p(y|I, θ) is what previously was referred to asp(y|Tθ (I)), i.e. the classifier conditioned the transformed I.
We approximate Eq. 14 using an unbiased estimate
S

VARIATIONAL APPROXIMATION

Eq(θ,λ) log p(yi |Ii , θi ) ≈1X
log p(yi |Ii , θi,s ),
S s=1

(15)

The integral equation (2) for the marginal likelihood is intractable and, thus, the posterior p(λ, θ|I, y) is too. We derive a lower bound on the log marginal likelihood to utilize
variational inference [Blei et al., 2017]. We choose the variational approximation q of the posterior p(θ, λ|I, y) as

and backpropagate through neural networks µ(I) and β(I)with the reparametrization trick. In all experiments αi = 1.

q(θ, λ) := p(θ|λ, I)q(λ).

Combining terms, the final ELBO we maximize becomes(7)

Here p(θ|λ, I) is given as before and q(λ) :=QN
i=1 Γ (αi , β(Ii )). In our approximation, β is a neural network: hence, we use amortized inference in a similar way to
the VAE model [Kingma and Welling, 2014].

i)with θi,s ∼ t2αi (·|µ(Ii )), β(Iαi )Lp,q (I, y) ≈

N
S
X
1X
i=1

S s=1

log p(yi |Ii , θi,s )− KL (q(λ)||p(λ)) + const,

(16)

(17)

which is readily optimized using any gradient-based method.
The KL-term is analytically tractable and differentiable between two gamma distributions.
In practice, following Higgins et al. [2016] we introduce
a weight parameter w to the KL-term. This requires us totune w but in turn makes the model robust to the choiceof prior. We perform a grid-search on a validation set tofind the optimal w. Alternatively, we could have done a gridsearch over β0 ; instead we chose α0 = β0 = 1 for all experiments. Similar to Kingma and Welling [2014], we often
find it sufficient to draw only S = 1 samples during training. Note that our model naturally implies marginalization,and correspondingly data augmentation, at test-time as wellas the usual training time. At test time, we draw S = 10transformation samples.

5

EXPERIMENTS & RESULTS

Our model consists of two parts, the classifier p(y|Tθ (I))
and the probabilistic localizer estimating the distribution
over transformations. In the following experiments, we
aim to disentangle our model’s benefits for localization(Sec. 5.1), classification (Sec. 5.2) and calibration (Sec. 5.3).
The probabilistic localizer estimates q(θ)=
t2 (θ|µ(I), β(I)), i.e. in practice we implement a mean anda variance network, µ(I) and β(I), respectively (see Fig. 2
for the architecture). We employ a small convolutional
network (Conv2d, Maxpool2d, ReLU, Conv2d,
Maxpool2d, ReLU) followed by two fully connected
layers for both the localizer and classifier unless statedotherwise. The P-STN localizer has two heads; one for themean and one for the variance. The number of parameters
is stated in each experimental subsection. Unless statedotherwise, we keep the number of parameters constant, i.e.
when adding a localization network we remove the extra
parameters from the classifier for fair comparison.
Our model is implemented in PyTorch and experiments are run on 12 GB Nvidia Titan X GPUs.
The code is available at https://github.com/
FrederikWarburg/pSTN-baselines.
5.1

MARGINALIZING TRANSFORMATIONSIMPROVES LOCALIZATION ACCURACY

The appeal of STN models is that they are trained end-toend, i.e. based only on labels for the downstream task, andnot the transformations. This same property, however, iswhat makes the STN hard to fit. The only signal we obtain isthrough the supervised downstream task (i.e. the classification labels) and thus gradient information is sparse. We will
now investigate whether estimating a posterior over transformations and marginalizing, i.e. ‘getting to try multiple
transformations’, simplifies the task as suggested by Fig. 1.

Figure 4: Rotated MNIST experiment. Left panel: Groundtruth transformation (rotation angles in radians) against recovered transformations (mean). Top right: Example images
from the data set and samples from the P-STN localizer. The
localizer learns to pose-normalize. Bottom right: Outputs ofthe variance network. When the transformation recovery ispoor (the error ε is above the median, in orange) the variances are slightly higher than when the localization works
well (blue).
In order to disentangle the localization from the classification task, we construct the following experiments. We firsttrain a CNN on a pose-normalized dataset (regular MNISTand Fashion MNIST). We then generate a new dataset byrandomly sampling transformations θtrue and applying themto the MNIST images. Saving these transformations provides us with ground truth. We freeze the CNN weights andtrain STN and P-STN with this fixed classifier, effectively
learning to recover and ‘undo’ the true transformations.
5.1.1

Rotated MNIST

From this data-generating process, we obtain a rotated version of the MNIST dataset (i.e. regular MNIST with groundtruth transformations given by rotation angles, θtrue (I) =rtrue (I)). See Fig. 4, top right panel for example data.
Our CNN classifier (28k weights) obtains 99.4% test accuracy on MNIST and 41.2% on rotated MNIST (frozen
weights, no re-training). The STN and P-STN (S = 10 training samples, w = 3 · 10−5 , same CNN classifier as before+72k params in the localizer) both learn to pose-normalize,
i.e. to recover these transformations to a satisfactory degree.
When training the localizers only (classifier weights remainfrozen as described above), the STN test acc. is 76.13%, and82.98% for the P-STN. We compute the expected averagetransformation error on the N = 10k rotated MNIST test
images asN

1 Xε=kθtrue (Ii ) − µ(Ii )k
N i=1

mod π.

(18)

CNN
STN
P-STN
Figure 6: The P-STN learns to localize traffic signs in the challenging MTSD
dataset. At test time, we sample 10 transformations as shown with the various
bounding boxes overlaid the images. These learned variations improve the finalclassification.
We get ε = 0.76 for the STN and ε = 0.59 for the P-STN.
The P-STN outperforms the STN, i.e. modeling uncertainty
in the transformations helps in the localization task.
Uncertainty. The bottom right panel of Fig. 4 shows ahistogram of β(I), i.e. the localizer variance (or, correspondingly, the magnitude of augmentation) per image. In orange,
we plot variances for images where pose-normalization isdifficult (the transformation error ε is larger than the median). In blue, we plot variances for images that are correctly
pose-normalized (transformation error ε smaller than the
median). The poorly localized images are, on average, assigned 17% larger variances β(I). The localizer uncertainty
and thus the amount of data augmentation applied is somewhat meaningful, corresponding to the difficulty of the task.
5.1.2

Random placement FashionMNISTWe repeat a similar experiment on the slightly more challenging FashionMNIST dataset [Xiao et al., 2017] . The
CNN baseline accuracy is 90.63% (same model as above
with 28k parameters). We then randomly sample an x and ycoordinate and place the FashionMNIST accordingly on a
black background, after downscaling it by 50%. No rotation
is applied, i.e. θtrue = [0, 0.5, txtrue , tytrue ].

Figure 7: Random Placement Fashion MNIST. Input images(left) and transformed samples Tθs (I) as learned by the PSTN. The P-STN learns to correctly pose-normalize andzoom into the relevant part of the image. The samples look
like plausible candidates for a data augmentation scheme.
We will explore this in Sec. 5.2.

Acc. ↑

NLL ↓

76.0
90.6
92.2

0.49
0.31
0.29

Table 1: Accuracy(Acc.) and negative
log-likelihood (NLL) for
CNN, STN and P-STN.

Like in the previous experiment, both localizers successfully recover θtrue , with the P-STN (S = 10 training samples,w = 3e−05, same classifier as before +193k weights in thelocalizer) doing slightly better than its deterministic counterpart: test accuracies are 84.99% and 84.41%, respectively.
Inspecting the transformation posterior and the resulting
samples Tθs (Ii ), we find that these look visually pleasing,
and, as hypothesized, might be promising candidates for adata augmentation scheme. We will explore this in Sec. 5.2.
5.1.3

Mapillary street signs

Detection and classification of objects in images have many
applications, e.g. for autonomous vehicles, detecting traffic
signs is crucial. We compare a top-performing classifier, an
STN and our P-STN on the challenging Mapillary Traffic
Sign Dataset (MTSD) [Ertler et al., 2019].
To focus this comparison, we select images that contain
only one traffic sign. We obtain this subset by selecting all
bounding boxes that do not intersect with other boundingboxes plus a margin of 150 px to each side. We further selectthe ten most common classes from this subset. This gives us
a training set of 4698 images and a test set of 500 images.
Figure 6 shows example images from the chosen subset.
Our classifier is a ResNet18 pre-trained on ImageNet, where
we replace the last fully connected layer. We use the same
ResNet for the localizers in the STN and P-STN, where we
similarly replace the last layer. As before, we wish to study
the behavior of the localizers. Therefore, we again start bytraining a classifier on the ground-truth bounding boxes. We
then initialize the classifier module of the STN and P-STN
with this pre-trained classifier and freeze the weights of theclassifier. We train the localizers of the STN and P-STN for
60 epochs with learning rate 10−4 and kl weight w = 10−7 .
Figure 6 shows that the P-STN learns to localize the traffic
signs. At test time, we sample 10 transformations illustrated
by the multiple overlaying bounding boxes.
Table 1 shows that both the STN and P-STN clearly outperform the baseline classifier when trained on the full images.
Even though the STN and P-STN have exactly the sameclassifier, the P-STN achieves better performance becauseof the ensemble of classified transformations.

CNN
affine STN
affine P-STN
optimal w

MNIST30

MNIST100

MNIST1000

MNIST3000

MNIST1000070.12 ± 2.46
69.26 ± 4.53
81.00 ± 3.92
0.001

87.29 ± 0.58
82.16 ± 2.30
92.70 ± 0.74
0.000395.80 ± 0.33
92.05 ± 0.58
96.62 ± 0.58
0.000197.48 ± 0.21
94.71 ± 0.22
97.33 ± 0.17
0.0000397.82 ± 0.34 96.96 ± 0.20
97.63 ± 0.23
0.00001

Table 2: The performance of a CNN, STN and P-STN on differently sized MNIST datasets. Bold numbers indicate that amodel is significantly better than the runner up under a two sample t-test at p = 0.05.

0.95
0.90
0.85
0.80
0.75
0.70
0.65

MN

IST

10

00

0

00
IST
30
MN

00
IST
10
MN

10
IST
MN

IST
MN

0

P-STN
STN
CNN
30

Accuracy

MNIST subsets

Figure 8: Performances of P-STN, STN and CNN on
MNIST subsets (mean ± one STD across five folds).

5.2

architecture as described at the top of Sec. 5. We use theAdam optimizer with weight decay 0.01 and the default
parameters of its PyTorch implementation. The images arecolor-normalized. We repeat the experiment 5 times, eachtime with a different k-image subset of the MNIST dataset,
and we report ± one standard deviation in tables and error
bars. From Table 2 and Fig. 8, we see that the P-STN outperforms both the STN and CNN on the small dataset sizes.
For the larger datasets, the differences vanish. This supports
our hypothesis: data augmentation is especially useful whendata is a limited resource. This intuition is also supportedby the optimal KL-weights (Table 2, bottom row) that wedetermine via grid search on validation data. For smaller
datasets, larger w and thus more regularization towards thevariance prior (away from 0) are beneficial.
The fact that the STN performs less well than the standard
CNN on this data set might be explained by the fact that theimages are already nearly perfectly pose-normalized, andwrong transformations can be detrimental.

MARGINALIZING TRANSFORMATIONSIMPROVES CLASSIFICATION ACCURACY
5.2.2

We have argued that marginalizing transformations via samples corresponds to learned, localized data augmentations(the samples Tθs (I)). We will now investigate whether theseaugmentations are indeed helpful in the downstream task,
i.e. whether they improve classification performance.

5.2.1

MNIST and subsets

We compare the performance of our P-STN against a standard convolutional neural network (CNN) and a regular STN
on MNIST. The standard MNIST images are centered andpose-normalized, so the localization task is easy. Improved
classifier performance can thus be viewed as an indicatorfor having learned a useful data augmentation scheme.
Data augmentation is particularly important when training
data is scarce, so we evaluate the models on small subsets
of MNIST: MNIST30 contains 30 images (i.e. 3 per class),
MNIST100, MNIST1000, MNIST3000 and MNIST10000.
STN and P-STN parameterize affine transformations, i.e.
the learned θ is interpreted as the full affine matrix as described in Sec. 3. All models have roughly 28k parameters,

UCR time-series dataset

For some data modalities, such as time-series, it is nottrivial to craft a useful data augmentation scheme. In thisexperiment, we show that the P-STN can learn a useful,
non-trivial data augmentation scheme that increases performance compared to a standard STN on time-series data. The
UCR dataset [Dau et al., 2018] is composed of 108 smaller
datasets, where each dataset contains univariate time-series.
The FordA dataset, for example, contains measurements ofengine noise over time and the goal is to classify whether or
not the car is faulty. We select 5 of those subsets, each large
enough to divide into training and validation sets (75/25%),
which we use to find the optimal w via grid-search; thoseare [0.0001, 1e − 05, 0.001, 0.0, 0.0001]. We draw S = 10
training samples. The test-set is pre-defined by the dataset
curators. Learning rate and optimizer are the same as inSec. 5.2.1, but we do not perform normalization. All models
have approximately one million parameters. Table 3 showsthat the P-STN achieves higher mean accuracy than boththe STN and the CNN, indicating that we can automaticallylearn a useful data augmentation scheme for time-series.
We verify this qualitatively in Fig. 10, which shows an

Figure 10: Examples of augmentations for a timeseries from the FaceAll dataset. The top plotshows the original time-series and the bottom
plot shows three augmented versions of the timeseries.

CNN

STN

P-STN

FaceAllTwoPatterns
wafer
uWaveGestureLib.*
PhalangesOutlC.**

80.83 ± 0.62
97.92 ± 0.53
99.63 ± 0.05
74.15 ± 1.27
79.88 ± 1.32

82.28 ± 0.42
99.79 ± 0.04
99.18 ± 0.17 79.77 ± 0.42 82.26 ± 0.98

84.31 ± 0.75
99.96 ± 0.04
98.86 ± 0.20
81.13 ± 0.46
81.66 ± 0.59

Mean

86.48

88.65

89.18

Table 3: Accuracies on a subset of the UCR timeseries dataset (fulldataset names are *uWaveGestureLibrary and **PhalangesOutlinesCorrect). ±1 STD is reported after 5 repetitions. Bold numbers
indicate that a model is significantly better than the runner up undera two sample t-test at p = 0.05.

MNIST100 Calibration

1.0

Accuracy

0.8

We will investigate how well the uncertainty in this distribution matches the quality of predictions. Fig. 11 shows acalibration plot for the MNIST100 subset classification task
from Sec. 5.2.1 for the CNN, STN and P-STN for two different w-parameters; w = 0.0003 yields the best performance(reported in Table 2) and w = 0.0001 yields the best calibration. The expected calibration errors [Guo et al., 2017,
Küppers et al., 2020, 2021] are CNN: 0.0743 ± 0.0094,
STN: 0.1160 ± 0.0205, P-STN, w = 0.0003 (optimal performance model): 0.0567 ± 0.0065, P-STN, w = 0.0001(optimal calibration model): 0.0271 ± 0.0088. We report
the mean over 5 folds, ± one STD. The P-STN significantly
improves calibration in the downstream classification task.

CNN
STN
P-STN w = 0.0003
P-STN w = 0.0001

0.6
0.4
0.2
0.0
0.0

0.2

0.4

0.6

Confidence

0.8

1.0

Figure 11: Calibration plots for CNN, STN and two P-STN
models. One with KL-weight yielding optimal performance(w = 0.0003) and one with KL-weight yielding optimal
calibration (w = 0.0001). Both P-STN models are better
calibrated than CNN and STN.
example of the learned data augmentation. We see that themodel does not simply apply a global transformation, butlearns to augment the time-series more in some intervals,
such as in [60; 110], and augment the time-series less inother intervals, such as in [0; 50].
5.3

MARGINALIZING TRANSFORMATIONSIMPROVES CALIBRATIONIn Sec. 5.1, we have seen that harder images on average
have larger transformation uncertainties. We now investigate whether those meaningful localization uncertainties
translate into meaningful uncertainties downstream, i.e. in
the calibration of our classifier. At test-time, we evaluate
ZS
1X
p(y|I) = p(y|I, θ)q(θ)dθ ≈p(y|Tθs (I)). (19)
S s=1

5.4

A TYPICAL FAILURE MODE IN STNS

STNs are trained end-to-end, and with only label information available. Thus, the aim is to learn the optimal transformation for solving the downstream task. Depending on thecomplexity of the downstream task and the classification
model, it might not be necessary to transform the input at
all, i.e. it might be possible to solve the downstream task
on the original input image. Indeed, this is a failure mode
we observe in practice — often, the localizer simply learns
the identity transform while the classifier learns to classify
the non-transformed image. Using more complex classifier architectures makes the STN more prone to this failuremode. This has been observed by other authors [Finnveden et al., 2021], and we investigate the problem in theexperiment in Fig. 12. We start by training differently-sizedneural networks on MNIST (black, one layer on the x-axis
is [Linear, ReLU, Dropout]). We compare the performance of this model with (P-)STN models trained on
rotated MNIST, test accuracies are plotted in the left panelof the figure. If the localization task is performed perfectly,
the (P-)STN models should be able to recover the accuracy
on the original, non-rotated dataset. In the right panel, weplot the variance of the (mean) transformations learned bythe (P-)STN models. Values close to 0 indicate that the
localizer does not transform the image, i.e. it learns theAccuracy

0.9
0.8
MNIST, NN
rotMNIST, P-STN
rotMNIST, STN
rotMNIST, NN

0.7
0.6
1 layers

2 layers

3 layers

Model size

4 layers

5 layers

Variance of (I) (radians)

1.0

1.25

rotMNIST, STN
rotMNIST, PSTN

1.0
0.75
0.5
0.25
0.0
1 layers

2 layers

3 layers

Model size

4 layers

5 layers

Figure 12: Left: Test accuracies for standard NN and (P-)STNs of different depths trained on rotated MNIST, as well as
NN baseline on original MNIST (black). The STN (green) model does not usually recover the original images and thus
behaves more like a standard NN (red) in most runs. P-STN (blue) un-transforms at least some of the rotations and is closer
in accuracy to the NN on original MNIST (black). Right: The variance of the learned transformations as a function of model
depth. The STN learns the identity for deeper downstream models (this is consistent with the test accuracies we see on theleft). P-STN learns to un-transform better, at least when the classifier is simple. For bigger classifiers it predicts the identity
transform as well, but performs relatively well nonetheless (see left panel). We report medians ±1 median absolute deviation
over 5 folds.
identity transform. Larger values indicate that the localizer
learns transformations. Median results are reported over 5
runs, error bars correspond to one mean absolute deviation.
As hypothesized, for larger classifiers the localizers do nottransform the images. Due to the increased capacity of themodel, we nonetheless achieve decent classification accuracies (left panel). The P-STN learns to localize the rotatedimages somewhat successfully (large variance in the rightpanel, and high accuracy on the left) for smaller classifiers.
The STN does not localize the images as well, most runs
behave like the standard NN on rotMNIST (red), predicting
identity transformations only. We conclude that, thanks toit ‘trying out multiple transformations’, the P-STN avoids
this failure mode to an extent. We also note that this property, while useful, is somewhat orthogonal to our interestin this work, and we have avoided the failure mode in theexperiments of Sec. 5.1 by considering models with fixed,
pre-trained classifiers.

6

CONCLUSION

We have introduced a probabilistic extension to the spatial transformer network (STN) [Jaderberg et al., 2015]. Our
work took motivation from the empirical observation that the
STN is often brittle to train, as a poorly predicted transformation may prevent the model from getting any gradient signal,
resulting in divergent optimization. Our probabilistic STN(P-STN) instead approximates the posterior distribution oftransformations using amortized variational inference, andmarginalizes accordingly. As is common, marginalization
improves the robustness of the model.
Empirically, we note the following advantages of the prob-

abilistic formulation over the deterministic. Firstly, the performance of the localization network is improved, since theMonte Carlo marginalization effectively amounts to trying
many different transformations. Secondly, the probabilistic
formulation improves the overall model performance, since
the sampled transformations act as data augmentation bothduring training and during testing. The resulting ensemble ofpredictions is more accurate and better calibrated than common classifiers as well as the original spatial transformer.
Acknowledgements
MJ was supported by a research grant from the Carlsberg
Foundation (CF20-0370). SH was supported by research
grants (15334, 42062) from VILLUM FONDEN. This
project has also received funding from the European Research Council (ERC) under the European Union’s Horizon
2020 research and innovation programme (grant agreement
No. 757360). This work was funded in part by the Novo
Nordisk Foundation through the Center for Basic Machine
Learning Research in Life Science (NNF20OC0062606).
References
Henry S Baird. Document image defect models. In SDIA,
pages 546–556. Springer, 1992.
Gregory Benton, Marc Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neuralnetworks. arXiv preprint arXiv:2010.11882, 2020.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe.
Variational inference: A review for statisticians. ArXiv,
abs/1601.00670, 2017.

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu,
and Daan Wierstra. Weight uncertainty in neural networks. arXiv preprint arXiv:1505.05424, 2015.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of theIEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 113–123, 2019.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V
Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of theIEEE/CVF Conference on Computer Vision and Pattern
Recognition Workshops, pages 702–703, 2020.
Hoang Anh Dau, Eamonn Keogh, Kaveh Kamgar, ChinChia Michael Yeh, Yan Zhu, Shaghayegh Gharghabi,
Chotirat Ann Ratanamahatana, Yanping, Bing Hu, Nurjahan Begum, Anthony Bagnall, Abdullah Mueen, Gustavo
Batista, and Hexagon-ML. The ucr time series classification archive, October 2018. https://www.cs.ucr.
edu/~eamonn/time_series_data_2018/.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer,
Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless bayesian deep learning.
Advances in Neural Information Processing Systems, 34,
2021.
Nicki Skafte Detlefsen. libcpab. https://github.
com/SkafteNicki/libcpab, 2018.
Nicki Skafte Detlefsen and Søren Hauberg. Explicit disentanglement of appearance and perspective in generativemodels. In Advances in Neural Information Processing
Systems (NeurIPS), 2019.
Nicki Skafte Detlefsen, Oren Freifeld, and Søren Hauberg.
Deep diffeomorphic transformer networks. In 2018
IEEE/CVF Conference on Computer Vision and Pattern
Recognition, pages 4403–4412, June 2018.
Nicki Skafte Detlefsen, Martin Jørgensen, and Søren
Hauberg. Reliable training and estimation of variancenetworks. In 33rd Conference on Neural Information
Processing Systems, 2019.
Christian Ertler, Jerneja Mislej, Tobias Ollmann, Lorenzo
Porzi, and Yubin Kuang. Traffic sign detection and classification around the world. CoRR, abs/1909.04422, 2019.
URL http://arxiv.org/abs/1909.04422.
Lukas Finnveden, Ylva Jansson, and Tony Lindeberg. Understanding when spatial transformer networks do notsupport invariance, and what to do about it. In 2020 25th
International Conference on Pattern Recognition (ICPR),
pages 3427–3434. IEEE, 2021.

Oren Freifeld, Søren Hauberg, Kayhan Batmanghelich, andJohn W. Fisher. Transformations based on continuous
piecewise-affine velocity fields. IEEE Transactions on
Pattern Analysis and Machine Intelligence, 2017.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian
approximation: Representing model uncertainty in deeplearning. In international conference on machine learning, pages 1050–1059. PMLR, 2016.
Ian Goodfellow, Honglak Lee, Quoc V Le, Andrew Saxe,and Andrew Y Ng. Measuring invariances in deep networks. In NIPS, pages 646–654, 2009.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. In
Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1321–1330. JMLR. org,
2017.
Fredrik K Gustafsson, Martin Danelljan, and Thomas B
Schon. Evaluating scalable bayesian deep learning methods for robust computer vision. In Proceedings of theIEEE/CVF conference on computer vision and patternrecognition workshops, pages 318–319, 2020.
Søren Hauberg, Oren Freifeld, Anders Boesen Lindbo
Larsen, John W. Fisher, and Lars Kai Hansen. Dreaming more data: Class-dependent distributions over diffeomorphisms for learned data augmentation. In Artificial
Intelligence and Statistics, pages 342–350, 2016.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess,Xavier Glorot, Matthew Botvinick, Shakir Mohamed, andAlexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. 2016.
Max Jaderberg, Karen Simonyan, Andrew Zisserman, andKoray Kavukcuoglu. Spatial transformer networks. In Advances in Neural Information Processing Systems, pages
2017–2025, 2015.
Angjoo Kanazawa, David W Jacobs, and Manmohan Chandraker. Warpnet: Weakly supervised matching for singleview reconstruction. In CVPR, 2016.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada,
April 14-16, 2014, Conference Track Proceedings, 2014.
URL http://arxiv.org/abs/1312.6114.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.
Imagenet classification with deep convolutional neuralnetworks. In F. Pereira, C. J. C. Burges, L. Bottou, andK. Q. Weinberger, editors, Advances in Neural Information Processing Systems 25, pages 1097–1105. Curran
Associates, Inc., 2012.

Fabian Küppers, Jan Kronenberger, Amirhossein Shantia,
and Anselm Haselhoff. Multivariate confidence calibration for object detection. In The IEEE/CVF Conference
on Computer Vision and Pattern Recognition (CVPR)
Workshops, June 2020.
Fabian Küppers, Jan Kronenberger, Jonas Schneider, andAnselm Haselhoff. Bayesian confidence calibration for
epistemic uncertainty modelling. In Proceedings of theIEEE Intelligent Vehicles Symposium (IV), July 2021.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles
Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural
Information Processing Systems, pages 6402–6413, 2017.
Yann LeCun, LD Jackel, Leon Bottou, A Brunot, Corinna
Cortes, JS Denker, Harris Drucker, I Guyon, UA Muller,
Eduard Sackinger, et al. Comparison of learning algorithms for handwritten digit recognition. In International
conference on artificial neural networks, volume 60,
pages 53–60. Perth, Australia, 1995.
Chen-Hsuan Lin and Simon Lucey. Inverse compositional
spatial transformer networks. CoRR, abs/1612.03897,
2016.
URL http://arxiv.org/abs/1612.
03897.
Gaëlle Loosli, Stéphane Canu, and Léon Bottou. Training invariant support vector machines using selective sampling.
Large scale kernel machines, pages 301–320, 2007.
David J C MacKay. Model comparison and Occam’s razor.
Information Theory, Inference and Learning Algorithms,
pages 343–355, 2003.
David JC MacKay. Bayesian interpolation. Neural computation, 4(3):415–447, 1992.
Pola Schwöbel, Martin Jørgensen, Sebastian W Ober, andMark Van Der Wilk. Last layer marginal likelihood forinvariance learning. In International Conference on Artificial Intelligence and Statistics, pages 3542–3555. PMLR,
2022.
Ron A Shapira Weber, Matan Eyal, Nicki Skafte, Oren
Shriki, and Oren Freifeld. Diffeomorphic temporal alignment nets. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d’ Alché-Buc, E. Fox, and R. Garnett, editors,
Advances in Neural Information Processing Systems 32,
pages 6570–6581. 2019.
Patrice Y Simard, Dave Steinkraus, and John C Platt. Bestpractices for convolutional neural networks applied tovisual document analysis. In 2013 12th International
Conference on Document Analysis and Recognition, volume 2, pages 958–958. IEEE Computer Society, 2003.

Søren Kaae Sønderby, Casper Kaae Sønderby, Lars Maaløe,
and Ole Winther. Recurrent spatial transformer networks.
arXiv preprint arXiv:1509.05329, 2015.
Andrew Stirn and David A Knowles. Variational variance:
Simple and reliable predictive variance parameterization.
arXiv e-prints, pages arXiv–2006, 2020.
Hiroshi Takahashi, Tomoharu Iwata, Yuki Yamanaka,
Masanori Yamada, and Satoshi Yagi. Student-t variational autoencoder for robust density estimation. In IJCAI,
pages 2696–2702, 2018.
Mark Van der Wilk, Matthias Bauer, ST John, and James
Hensman. Learning invariances using the marginal likelihood. In Advances in Neural Information Processing
Systems, pages 9938–9948, 2018.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashionmnist: a novel image dataset for benchmarking machinelearning algorithms. arXiv preprint arXiv:1708.07747,
2017.