# Geometrically Constrained and Token-Based Probabilistic Spatial Transformers

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 61 references that are in the map.
- First posted on arXiv 14 Sep 2025.

## Builds on (most cited first, 30 of 35)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Spatial Transformer Networks](https://mlsanity.com/paper/1506.02025/) (2015)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classification](https://mlsanity.com/paper/2110.14795/) (2021)
- [Making Convolutional Networks Shift-Invariant Again](https://mlsanity.com/paper/1904.11486/) (2019)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [Cyclical Annealing Schedule: A Simple Approach to Mitigating KL Vanishing](https://mlsanity.com/paper/1903.10145/) (2019)
- [Frame Averaging for Invariant and Equivariant Network Design](https://mlsanity.com/paper/2110.03336/) (2021)
- [Equivariance with Learned Canonicalization Functions](https://mlsanity.com/paper/2211.06489/) (2022)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [MedViT: A Robust Vision Transformer for Generalized Medical Image Classification](https://mlsanity.com/paper/2302.09462/) (2023)
- [Polar Transformer Networks](https://mlsanity.com/paper/1709.01889/) (2017)
- [Meta-Learning Symmetries by Reparameterization](https://mlsanity.com/paper/2007.02933/) (2020)
- [TI-POOLING: transformation-invariant pooling for feature learning in Convolutional Neural Networks](https://mlsanity.com/paper/1604.06318/) (2016)
- [The Effects of Regularization and Data Augmentation are Class Dependent](https://mlsanity.com/paper/2204.03632/) (2022)
- [Equivariant Adaptation of Large Pretrained Models](https://mlsanity.com/paper/2310.01647/) (2023)
- [Equivariant Transformer Networks](https://mlsanity.com/paper/1901.11399/) (2019)
- [Inverse Compositional Spatial Transformer Networks](https://mlsanity.com/paper/1612.03897/) (2016)
- [Robustness through Data Augmentation Loss Consistency](https://mlsanity.com/paper/2110.11205/) (2021)
- [Probabilistic Spatial Transformer Networks](https://mlsanity.com/paper/2004.03637/) (2020)
- [Tilt your Head: Activating the Hidden Spatial-Invariance of Classifiers](https://mlsanity.com/paper/2405.03730/) (2024)

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

- [Saccadic Vision for Fine-Grained Visual Classification](https://mlsanity.com/paper/2509.15688/) (2025)
- [Robust Canonicalization through Bootstrapped Data Re-Alignment](https://mlsanity.com/paper/2510.08178/) (2025)
- [Zero-Shot Test-Time Canonicalization using Out-of-Distribution Scoring](https://mlsanity.com/paper/2606.24178/) (2026)

## Abstract (arXiv)

Spatial transformations such as rotation and scale obscure the morphological cues needed for accurate image classification. Careful consideration is required for reliable use in high stakes settings. A model should stay robust under such transformations, expose why a correction was applied, and signal when its input is ambiguous. While geometrically equivariant architectures provide a mathematically grounded solution, they often limit model flexibility through strict symmetry constraints and incur significant computational overhead. Spatial Transformer Networks (STNs) offer a data-driven, flexible alternative for learning pseudo-equivariances to affine transformations. However, STNs have historically been restricted to convolutional architectures and suffer from training instability. To address this, we introduce a novel STN framework. It leverages the global modeling capabilities of transformers to regress the affine transformation acting on the input. For this, we decompose affine transformations into interpretable primitives, regressed under adaptable geometric constraints, thereby preventing the training instability typically caused by degenerate transformations. By sharing weights between the localization network and the classification backbone, the framework requires minimal computational overhead. Extensive experiments on challenging insect biodiversity and medical imaging benchmarks demonstrate that our approach achieves superior predictive performance under diverse spatial transformations while maintaining high efficiency. Code is available at https://github.com/johSchm/TokenSTN.