Learning to Perceive "Where": Spatial Pretext Tasks for Robust Self-Supervised Learning

Abstract: Existing self-supervised learning (SSL) methods primarily learn object-invariant representations but often neglect the spatial structure and relationships among object parts. To address this limitation, we introduce Spatial Prediction (SP), a spatially aware pretext regression task that predicts the relative position and scale between a pair of disentangled local views from the same image. By modeling part-to-part relationships in a continuous geometric space, SP encourages representations to capture fine-grained spatial dependencies beyond invariant categorical semantics, thereby learning the compositional structure of visual scenes. SP is implemented as a decoupled plug-in and can be seamlessly integrated into diverse SSL frameworks. Extensive experiments show consistent improvements across image recognition, fine-grained classification, semantic segmentation, and depth estimation, as well as substantial gains in out-of-distribution robustness for object recognition. To evaluate spatial reasoning, we introduce (1) a position and scale prediction task on image patch pairs and (2) a jigsaw understanding task requiring patch reordering and recognition after reconstruction. Strong performance on these tasks indicates improved spatial structure and geometric awareness. Overall, explicitly modeling spatial information provides an effective inductive bias for SSL, leading to more structured representations and better generalization. Code and models will be released.

• The paper introduces SP as a novel pretext task by regressing relative position and scale to improve spatial reasoning in self-supervised vision.
• It integrates SP with major SSL frameworks like MoCo v3, DINO, and MAE, showing significant gains in classification, segmentation, and robustness tasks.
• Experimental results demonstrate enhanced geometric supervision, with improved fine-grained discrimination and reduced spatial regression errors across diverse benchmarks.

Spatial Prediction as an Explicit Inductive Bias for Self-Supervised Vision

Background and Motivation

Self-supervised learning (SSL) for vision has predominantly focused on either enforcing invariance to image transformations (e.g., via contrastive or distillation methods) or reconstructing local regions from masked or corrupted inputs. However, these approaches often neglect spatial structure and the geometric relationships among object parts. As a result, models tend to learn global or patch-level representations that are semantically robust, but under-constrained in terms of spatial reasoning, limiting their effectiveness for tasks requiring part-to-part compositionality or geometric understanding.

The Spatial Prediction (SP) Pretext Task

The authors introduce the Spatial Prediction (SP) pretext task to directly address the spatial reasoning deficit in current SSL paradigms. SP operates by sampling two disentangled local views from a single image and requiring the model to regress the relative position and scale of one view with respect to the other in a continuous, normalized coordinate system. This regression-based objective is intended to encourage representations that encode precise part-to-part dependencies—providing a structured geometric inductive bias that is largely absent from prior invariance-based or reconstruction-based methods.

SP is implemented as an auxiliary loss alongside the main SSL objective and is designed as an architecture-agnostic, plug-and-play module. It introduces a spatial reasoning branch that processes shared ViT features from the original encoder, with a novel projection-free cross-attention mechanism that concatenates feature tokens for spatial prediction, supervised by $\ell_2$ losses on relative position and scale. Importantly, the view pair sampling adheres to principled constraints to avoid degenerate or trivial solutions, ensuring that the spatial supervision remains informative and balanced.

Experimental Protocol and Benchmarks

To validate SP, the authors integrate it into three major SSL frameworks—MoCo v3, DINO, and MAE—across standard Vision Transformer (ViT) backbones. The evaluation suite is comprehensive, covering seven downstream tasks over eleven datasets:

• In-domain image classification: CIFAR-100, ImageNet-1K.
• Robustness to corruption/occlusion: ImageNet-C, ImageNet-R, ImageNet-Sketch, and a synthetic occlusion benchmark.
• Fine-grained transfer learning: Flowers-102, DTD, Food-101.
• Dense prediction tasks: semantic segmentation (PASCAL VOC) and depth estimation (NYU v2).
• New spatial reasoning benchmarks: regression of patch-relative position/scale and a jigsaw permutation task for patch reordering and recognition.

Metrics include Top-1 accuracy for classification, mean Intersection over Union (mIoU) for segmentation, RMSE for depth estimation, and $\ell_2$ regression error for spatial tasks.

Numerical Results and Contrasting Claims

SP-augmented SSL models demonstrate consistent performance improvements across all major tasks and metrics relative to their baselines:

• Classification: On CIFAR-100 and ImageNet-1K, SP yields robust gains in Top-1 accuracy for MAE, MoCo v3, and DINO.
• Robustness: SP drives significant reductions in mean Corruption Error (mCE) on ImageNet-C and boosts accuracy on domain/appearance-shifted setups (e.g., +2.8% MoCo v3 on ImageNet-R), evidencing improved structural and shape-based instead of texture-biased recognition.
• Fine-grained and transfer: Gains are particularly pronounced on tasks where object part configuration is discriminative (e.g., Flowers), with up to +25% Top-1 accuracy improvement for MoCo v3 + SP.
• Semantic segmentation/depth: SP consistently raises mIoU for segmentation and lowers depth estimation RMSE, indicating improved geometric and spatial grounding.
• Spatial reasoning: The SP branch dramatically lowers $\ell_2$ errors in position and scale regression, and jigsaw permutation accuracy rises markedly (e.g., MAE baseline 77.95% vs. MAE + SP 98.58%).

A central claim—supported by strong experimental evidence—is that explicit spatial pretext objectives do not impair, but systematically enhance both semantic quality and spatial reasoning of learned representations. This is in contrast to prior conjectures that geometric information might be diluted in the pursuit of strong invariance-driven SSL.

Ablation and Mechanistic Analysis

Ablation studies reveal the following insights:

• Joint supervision of position and scale produces the most effective representations, outperforming either alone.
• Parameter-free cross-attention—using raw ViT tokens without added projections—yields better spatial supervision than more heavily parameterized attention modules.
• The effect of spatial loss weights is systematically explored and found to improve spatial generalization while maintaining or improving semantic performance.

Theoretical and Practical Implications

The introduction of SP fundamentally extends the class of pretext tasks for SSL by making explicit geometric dependency modeling a first-class training signal. The results imply that compositional, part-aware visual knowledge can be induced in a scalable, architecture-compatible manner, resulting in representations that:

• Better generalize to tasks requiring spatial, compositional, and geometric reasoning.
• Show heightened robustness to texture-degrading perturbations and occlusion.
• Demonstrate improved fine-grained discrimination, which is often limited by existing invariance-focused SSL.

Practically, SP's architecture-agnosticity and plug-in design suggest compatibility with emerging backbone architectures and SSL pipelines, enabling its straightforward adoption in large-scale pretraining or continual learning scenarios.

Theoretically, these results underscore the necessity of revisiting the invariance paradigm in SSL: balancing transformation invariance with equivariance and explicit geometric supervision is critical to learn representations that are robust, spatially grounded, and transferrable.

Future Directions

The current SP paradigm is restricted to 2D spatial relationships. Extensions to 3D geometry, multi-object compositionality, temporal dynamics, and integration with cross-modal (e.g., audio-visual) reasoning are highlighted as pertinent areas. Additionally, refining spatial sampling strategies or combining SP with hierarchical object-centric modeling may yield further advances in foundation model capabilities.

Conclusion

The Spatial Prediction pretext task presents a general, effective approach for incorporating spatial reasoning inductive bias into vision SSL. Through explicit geometric supervision, SP yields representations with superior performance across classification, robustness, transfer learning, segmentation, depth estimation, and direct spatial reasoning. These findings illuminate the interplay between spatial and semantic learning in vision models, providing a foundation for future research on spatially, temporally, and physically grounded foundation models in artificial intelligence.

---

Reference: "Learning to Perceive 'Where': Spatial Pretext Tasks for Robust Self-Supervised Learning" (2605.09963)