Geometrically Constrained and Token-Based Probabilistic Spatial Transformers

Johann Schmidt; Tom Siegl; Martin Becker; Sebastian Stober

cs.CVSep 14, 2025

Geometrically Constrained and Token-Based Probabilistic Spatial Transformers

Authors: Johann Schmidt, Tom Siegl, Martin Becker, Sebastian Stober

Organizations: Artificial Intelligence Lab, Otto-von-Guericke University Magdeburg, Germany · University of Rostock, Germany · Becker Lab, University of Marburg, Germany

Abstract

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.

Figures & tables

Figure 1 : Our method integrates a Spatial Transformer module directly into the Vision Transformer pipeline. Unlike traditional STNs, we regress geometric primitives (e.g., rotation, scale) individually under adaptable constraints to ensure training stability and prevent degenerate transformations. By sharing weights between the localization network and the classification backbone, the framework eliminates feature redundancy and maintains high computational efficiency.

Figure 2 : Our model first encodes the input image into patch tokens using a frozen backbone tokenizer. A shared localization encoder processes these tokens, and separate regression heads predict affine transformation parameters for rotation, scaling, shearing, translation, and reflection (each is optional). The resulting affine matrix rectifies the image, which is then re-tokenized and fed into a downstream classifier.

Figure 3 : Average top-1 accuracy on both moth benchmark test sets versus the rotoscaled test sets. Besides the vanilla SwinB classifier, an augmented version was trained and evaluated and a set of different STNs (dots) including ours (stars).

Method	EU	ECU
Vanilla	90.2	61.7
Augmented	95.1 (+4.9)	70.7 (+9.0)
STN [ 24 ]	94.6 (+4.4)	65.4 (+3.7)
Head-STN [ 4 ]	95.3 (+5.1)	69.8 (+8.1)
Polar [ 14 ]	94.7 (+4.5)	67.6 (+5.9)
LogPolar [ 14 ]	93.8 (+3.6)	65.6 (+3.9)

Table 1 : Average top-1 test accuracies on Ecuador-Moth and EU-Moth with shearing augmentation. Relative gains compared to vanilla baseline in parentheses. All methods use pre-trained Swin-Base.

Dataset	Model	Acc	AUC
	MedViT-S	$41.50$	$0.697$
RetinaMNIST	Ours (Gauss)	$52.90\pm 0.63$	$0.752\pm 0.003$
	Ours (Gamma)	$\bm{55.35}\pm 0.58$	$\bm{0.754}\pm 0.004$
	MedViT-S	$78.80$	$0.950$
OCTMNIST	Ours (Gauss)	$75.30\pm 0.00$	$\bm{0.956}\pm 0.000$
	Ours (Gamma)	$\bm{83.20}\pm 0.39$	$\bm{0.956}\pm 0.002$

Table 2 : Accuracy (Acc) and ROC AUC (AUC) on MedMNIST with statistics over 5 runs. Color encodes whether our approach consistently under- or overperforms vanilla MedViT-S or results are mixed .

Figure 4 : Normalized entropy of the predicted rotation posterior per RotMNIST class, using a frozen ImageNet-1k ViT-B/8 as the downstream classifier. Each point is a test image, colored by how the rectification altered its prediction relative to the uncorrected baseline. Stars mark the per-class mean entropy of the fixed and broken subsets, with error bars showing one standard deviation. Percentages above each class give the conditional flip rates: fixed as a fraction of base-misclassified inputs, broken as a fraction of base-correct inputs. Entropy is a monotone function of the predicted rotation standard deviation, so higher values indicate greater geometric uncertainty.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Method	EU		ECU
Method	Vanilla	RotoScale	Vanilla	RotoScale
Vanilla	96.3	90.6	76.2	62.8
Augmented	97.1 $(+0.8)$	95.1 $(+4.5)$	73.0 $(-3.2)$	71.0 $(+8.2)$
TTA [ 29 ]	96.3 $(\pm 0.0)$	92.6 $(+2.0)$	76.2 $(\pm 0.0)$	62.2 $(-0.6)$
ITS [ 45 ]	94.4 $(-1.9)$	94.7 $(+4.1)$	60.9 $(-15.3)$	54.5 $(-7.3)$
CanPrior (G) [ 39 , 6 ]	94.6 $(-1.7)$	90.4 $(-0.2)$	67.0 $(-9.2)$	60.1 $(-2.7)$

Appendix

Table 3 : Average top-1 test accuracies on EU and ECU. Relative gains compared to vanilla baseline in parentheses. Extension of our benchmark on the biodiversity monitoring task.

Explore similar work

On Preserving Geometrical Invariance for Superpixel Image Classification using Graph Transformer
Jul 5, 2026Sarabeshwar Balaji, Shubham Mohanty, Akash AnilCnn-Transformer Tradeoff Image Classification
Indian Institute of Science Education and Research Bhopal, Bhopal, Madhya Pradesh, 462066, India
Active Spatial Guidance: Eliminating Injected Positional Mechanisms in Vision Transformers
Jul 1, 2026Cong Liu, Xiaofang Li, Simon X. YangVision Transformer Vision Encoders
School of Engineering, Yancheng Institute of Technology, No. 1 Xiwang Middle Avenue, Yancheng, 224051, Jiangsu, PR China · School of Computer Information Engineering, Changzhou Institute of Technology, No. 666 Liaohe Road, Xinbei District, Changzhou, 213032, Jiangsu, PR China · Advanced Robotics and Intelligent Systems Laboratory, School of Engineering, University of Guelph, 50 Stone Road East, Guelph, N1G 2W1, Ontario, Canada
INTCORT: Training-Free Spatial Reasoning Enhancement for Vision-Language Models via Input Transformations and Confidence Routing
Sep 21, 2026Haoran Sun, Jingqi Xu, Yanhui Li +3Spatial Reasoning Transform
The University of Hong Kong · University of Southern California · China Telecom +3

Jul 5, 2026, cs.LGJ/K move · Enter open · S save

On Preserving Geometrical Invariance for Superpixel Image Classification using Graph Transformer

Sarabeshwar Balaji, Shubham Mohanty, Akash Anil

Indian Institute of Science Education and Research Bhopal, Bhopal, Madhya Pradesh, 462066, India

Cnn-Transformer Tradeoff Image Classification Multiplex Graph Transformers Convolutional Neural Networks Vision Transformer Graph Neural Networks

Open paper AskPDF arXiv

Convolutional Neural Network (CNN) and Vision Transformer (ViT) for image classification exploit a dense grid of pixels containing redundant information. Consequently, for a larger image dataset, CNNs and ViTs face deployability challenges due to high computational complexity. Representing images as graphs of superpixels offers an efficient alternative that preserves key information while eliminating pixel-level redundancy. Graph Neural Networks (GNNs) have been utilized on such graphs to perform image classification. However, GNNs are known to struggle with capturing long-range dependencies which is important in the domain of image classification. Furthermore, a majority of these superpixel-based image classification approaches do not explicitly preserve translation/rotation invariance. Nevertheless, preserving translation/rotation invariance is important for robust image classification. Thus, this paper proposes SuperGT, a Graph Transformer-based framework for image classification, which captures the long range dependencies, along with a pre-processing scheme that preserves translation/rotation invariance. We evaluate SuperGT on CIFAR-10 dataset and observe that it performs significantly better than many baselines. Furthermore, we note that the overall performance of SuperGT is comparable to the previous state-of-the-art model, namely, ShapeGNN, without relying on coordinates of the boundary points of each superpixel required by ShapeGNN.

Abstract

Explore similar work

On Preserving Geometrical Invariance for Superpixel Image Classification using Graph Transformer

Active Spatial Guidance: Eliminating Injected Positional Mechanisms in Vision Transformers

INTCORT: Training-Free Spatial Reasoning Enhancement for Vision-Language Models via Input Transformations and Confidence Routing