# Probabilistic Spatial Transformer Networks

arXiv: https://arxiv.org/abs/2004.03637 · HTML: https://mlsanity.com/paper/2004.03637/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 7 Apr 2020.

## Built on by (most cited first, 7 of 7)

- [VTAE: Variational Transformer Autoencoder with Manifolds Learning](https://mlsanity.com/paper/2304.00948/) (2023)
- [Trainable Highly-expressive Activation Functions](https://mlsanity.com/paper/2407.07564/) (2024)
- [SpaceJAM: a Lightweight and Regularization-free Method for Fast Joint Alignment of Images](https://mlsanity.com/paper/2407.11850/) (2024)
- [Diffeomorphic Temporal Alignment Nets for Time-series Joint Alignment and Averaging](https://mlsanity.com/paper/2502.06591/) (2025)
- [Geometrically Constrained and Token-Based Probabilistic Spatial Transformers](https://mlsanity.com/paper/2509.11218/) (2025)
- [Learning to Transform for Generalizable Instance-wise Invariance](https://mlsanity.com/paper/2309.16672/) (2023)
- [Saccadic Vision for Fine-Grained Visual Classification](https://mlsanity.com/paper/2509.15688/) (2025)

## Abstract (arXiv)

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.