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

Jul 5, 2026cs.LG

On Preserving Geometrical Invariance for Superpixel Image Classification using Graph Transformer

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.

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

Ask PDF

Jul 1, 2026cs.CV

Active Spatial Guidance: Eliminating Injected Positional Mechanisms in Vision Transformers

Vision Transformers (ViTs) commonly rely on injected positional mechanisms to address self-attention's permutation invariance. Motivated by the spatial regularities of natural images, we ask whether spatial organization can be induced from data rather than explicitly injected. Under controlled, matched from-scratch training, we propose Active Spatial Guidance (Guidance), a training-only objective that disables positional injection and applies an auxiliary 2D coordinate-regression loss to the final-layer patch tokens. The guidance head is used only during training and removed for inference; the deployed model consists of a positional-injection-free ViT encoder and the task-specific prediction module. Using DINOv3 ViT backbones, Guidance consistently improves performance on ImageNet-100 classification, ADE20K semantic segmentation, and Hypersim monocular depth estimation, outperforming strong injected baselines such as learned absolute positional embeddings and rotary positional embeddings under identical training protocols. On ImageNet-100, broader comparisons against representative injected positional designs further support Guidance's effectiveness. Guidance also improves robustness under resolution transfer, and multi-resolution training further strengthens accuracy across input sizes. Overall, our results suggest that spatial inductive bias in ViTs need not be architecturally injected, but can be shaped through training-time supervision. The code used for training and evaluation is publicly available in https://github.com/cloudlc/asg.

Cong Liu, Xiaofang Li, Simon X. Yang

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

Vision Transformer Vision Encoders Rotary Position Embedding Geometric Guidance Inductive Bias Guidance

Ask PDF

Sep 21, 2026cs.CV

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

Vision-Language Models (VLMs) have demonstrated remarkable capabilities in multimodal tasks, yet they still exhibit poor ability in spatial reasoning. Existing training-dependent and training-free enhancement methods suffer from high computational costs with catastrophic forgetting and internal mechanism interference that compromises general capabilities, respectively. In this work, we first verify two key hypotheses: appropriate geometric image transformation and query-reversal transformation can recover incorrect spatial predictions, and correct predictions exhibit higher relation-token confidence than incorrect ones. Based on these findings, we propose INTCORT, a training-free spatial reasoning enhancement framework that constructs multiple inference views through input transformations and aggregates their predictions via relation-token confidence routing, without modifying the VLM's internal mechanisms. Experimental results on several commonly-used benchmarks demonstrate that INTCORT substantially improves spatial reasoning accuracy across diverse VLMs, achieving an average improvement of 10.01% over all models and benchmarks. Compared with prior works, INTCORT achieves superior performance with improvements of up to 25.01%.

Haoran Sun, Jingqi Xu, Yanhui Li +3

The University of Hong Kong · University of Southern California · China Telecom +3

Spatial Reasoning Transform Multi-View

Ask PDF