[Paper Review] Probabilistic Spatial Transformer Networks

Pola Schwöbel, Frederik Warburg+3 more|arXiv (Cornell University)|Apr 7, 2020

Energy Efficient Wireless Sensor Networks37 references4 citations

TL;DR

This paper proposes a probabilistic extension to Spatial Transformer Networks (P-STN) that models image transformations as a distribution rather than a single deterministic transformation. By marginalizing over sampled transformations via Monte Carlo inference, P-STN improves localization robustness, acts as learned data augmentation, and achieves higher classification accuracy and better model calibration than deterministic STNs across vision and time-series benchmarks.

ABSTRACT

Spatial Transformer Networks (STNs) estimate image transformations that can improve downstream tasks 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 image at 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 across standard imaging benchmarks and on a challenging real-world dataset that these two properties lead to improved classification performance, robustness and model calibration. We further demonstrate that the approach generalizes to non-visual domains by improving model performance on time-series data.

Motivation & Objective

• To address the brittleness of deterministic Spatial Transformer Networks (STNs), which fail catastrophically when transformation predictions are inaccurate.
• To improve training stability and localization performance by modeling uncertainty in image transformations through a probabilistic framework.
• To leverage stochastic transformations as implicit, per-image data augmentation to boost downstream classification performance.
• To demonstrate generalization of the approach beyond vision to time-series data, showing broader applicability.

Proposed method

• Proposes a hierarchical Bayesian model over image transformations, replacing deterministic transformation estimation with a variational posterior distribution.
• Uses amortized variational inference to jointly train the localization network and downstream classifier using only label information.
• Approximates the intractable integral of the predictive distribution via Monte Carlo sampling of transformations: p(y|I) ≈ (1/S) Σ p(y|T_θ^s(I)) over S sampled transformations.
• Applies stochastic transformations to the input image, feeding multiple transformed views to the classifier, which aggregates predictions across samples.
• Employs a differentiable spatial transformer module with a localization network that outputs parameters of a Gaussian distribution over transformations.
• Trains end-to-end using backpropagation, with the classifier loss computed over the ensemble of transformed samples.

Experimental results

Research questions

• RQ1 Can modeling transformation uncertainty improve the robustness of spatial transformer networks during training and inference?
• RQ2 Does marginalizing over stochastic transformations lead to better localization performance compared to deterministic STNs?
• RQ3 To what extent does the stochastic transformation process act as effective, learned data augmentation for downstream classification tasks?
• RQ4 Can the probabilistic STN framework generalize to non-visual domains such as time-series data?
• RQ5 How does the P-STN model compare in calibration and accuracy to standard neural networks and deterministic STNs on benchmark datasets?

Key findings

• P-STN achieves higher classification accuracy than both standard neural networks and deterministic STNs on rotated MNIST, especially when the classifier is small, due to effective localization and data augmentation.
• The P-STN model avoids the common failure mode of STNs—learning the identity transformation—by successfully localizing rotated images even with larger classifiers.
• The variance of the predicted transformations is significantly higher in P-STN than in STN, indicating more active and diverse localization behavior.
• P-STN improves model calibration, producing more reliable predictive uncertainty estimates compared to deterministic models.
• The method generalizes to time-series data, where it improves model performance, demonstrating applicability beyond vision tasks.
• The loss landscape of P-STN is smoother with fewer local minima than STN, indicating improved training stability.