Zero-Shot Test-Time Canonicalization using Out-of-Distribution Scoring

Abstract: Pretrained vision models often misclassify inputs that are rotated, scaled, or sheared, even though these affine transformations leave the object class unchanged. Robustness is usually restored either by building equivariance into the architecture or by retraining with augmentation, both of which require changing or retraining the model. Test-time canonicalization instead leaves the classifier untouched. It undoes the transformation of each input, mapping it to a canonical form near the training distribution before classification. Existing canonicalizers, however, rely on a narrow set of logit-based energy scores and bespoke search procedures, leaving the design space of scoring functions and optimizers unexplored. We reframe canonicalization as out-of-distribution (OOD) detection, which lets any OOD score serve as the energy minimized over transformations. Across benchmarks ranging from handwritten characters and sketches to natural images and 3D point clouds, we systematically evaluate around twenty OOD scores and nine search algorithms, finding that distance-based scores paired with random search and local refinement perform best overall. Because canonicalizing an already-aligned input can hurt accuracy, we add a gated mechanism that transforms an input only when its OOD score indicates this is needed, preserving most in-distribution accuracy while retaining the robustness gains on transformed inputs. Code is available at github.com/johschm/its.

• The paper introduces a novel approach that reframes canonicalization as an out-of-distribution detection task using OOD scores as energy functions.
• It employs global optimization strategies, notably random search with local refinement, to minimize OOD scores and significantly improve robustness on transformed data.
• A gating mechanism is implemented to selectively apply canonicalization, preserving in-distribution accuracy while enhancing performance on out-of-distribution inputs.

Zero-Shot Test-Time Canonicalization by Out-of-Distribution Scoring

Problem Setting and Motivation

Robustness to class-preserving input transformations is a critical limitation in pretrained vision models, particularly when objects are rotated, scaled, or sheared. Existing solutions, such as architectural inductive biases or data augmentation, require retraining or architecture modification, which is infeasible for deployed models. Test-time canonicalization, in contrast, offers an approach to restore robustness by transforming inputs to a canonical form close to the training distribution before inference, without modifying the downstream classifier.

Prior canonicalization methods are restricted in their energy function choices—typically relying on logit-based scores—and in their search strategies, both of which constrain their scalability and performance across diverse data modalities. This paper reframes canonicalization as an out-of-distribution (OOD) detection task, leveraging OOD scores as natural energy functions to be minimized over transformation groups. By empirically evaluating OOD scores paired with global optimization strategies, it demonstrates significant improvements in robustness, and further addresses the degradation of in-distribution (ID) accuracy via selective gating mechanisms.

Figure 1: The canonicalization pipeline: ID samples are passed directly to the downstream model; OOD samples undergo canonicalization via affine transformations to minimize OOD score, only if this improves distributional alignment.

Canonicalization via OOD Scoring

Canonicalization is formalized using group theory, with affine transformations forming matrix Lie groups acting on input spaces. Mapping an input $x$ to its canonical form requires finding the group element $g^*$ that minimizes an energy $E(g^{-1} \cdot x)$, where $E$ is ideally invariant under class-preserving transformations.

The central hypothesis is that OOD scores are suitable energy functions for canonicalization: points near the data manifold yield low OOD scores, while transformed (OOD) examples are penalized. While a unique minimum per orbit is not guaranteed, empirical results show practical viability, especially when combined with robust search methods.

Search strategies evaluated encompass random search, simulated annealing, multi-start gradient descent, and coordinate descent. Random search with local refinement emerges as the most effective trade-off between accuracy and computational cost, enabling canonicalization in diverse modalities.

Gated Canonicalization: Preserving ID Accuracy

A practical challenge is that unconditional canonicalization degrades ID accuracy—transforming already-aligned inputs distorts features unnecessarily. The proposed gating mechanism selectively applies canonicalization only when the OOD score exceeds a threshold, and accepts the canonicalized sample only when its OOD score is reduced. This two-stage gating preserves high ID accuracy while retaining robustness gains for transformed (OOD) inputs, permitting a tunable trade-off.

Extensive Empirical Evaluation

Benchmarks span MNIST, EMNIST, SI-Score (rotation), TU Berlin (sketches), and ModelNet10 (3D point clouds), utilizing downstream models from ResNet and ViT to PointNet++. Both accuracy and scoring distance (relative OOD score reduction) are used as metrics.

Retraining-heavy methods (augmentation, equivariant architectures) achieve highest robustness, but incur significant deployment costs. Among training-free methods, OOD-based canonicalization (using optimized OOD scores and search) yields superior accuracy on transformed data, notably outperforming logit-based canonicalization and standard test-time augmentation.

Figure 2: SI-Score (rotation): left shows canonicalized images (correct blue, incorrect red); right shows test accuracy across methods, with feature-based OOD scores (PC-kNN) yielding maximal gains.

Search algorithm comparison reveals that random search with local refinement is consistently optimal for canonicalization at moderate computational budgets.

Figure 3: EMNIST accuracy as a function of search budget; RS-LR (random search with local refinement) dominates at practical sample counts.

OOOD score ablation indicates that distance-based scores—kNN and prototype-based—effectively recover canonical forms and outperform logit-based and rectification-method scores across all datasets.

Figure 4: Unsupervised OOD score comparison across datasets: best test score highlighted, with kNN-type scores universally leading.

Qualitative results demonstrate correct alignment and classification after canonicalization, particularly with the best detector/search configuration.

Figure 5: Visualization of MNIST samples before and after transformation/canonicalization, showing restoration of alignment and classifiability.

Figure 6: Trade-off plots between ID and OOD accuracy with varying OOD gating threshold; substantial gains on transformed data achieved without compromising untransformed accuracy.

Budget scaling across datasets confirms RS-LR's advantage at higher budgets, while coordinate descent may be preferable in highly constrained domains such as line sketches.

Figure 7: Scaling search budget: accuracy after canonicalization for multiple algorithms across MNIST, ModelNet10, and TU Berlin.

Theoretical and Practical Implications

This framework decouples robustness from retraining and architectural design, permitting black-box deployment to any classifier. Canonicalization becomes a conditional mapping triggered by distributional misalignment, guided by a landscape of well-studied OOD scores and global optimization algorithms.

The method's flexibility extends to arbitrary transformation groups and diverse modalities. It exploits the full body of OOD detection literature, allowing domain-specific tailoring by score and search selection. The primary computational cost arises from inference-time search, which may limit real-time applications but is offset by its plug-and-play nature.

From a theoretical perspective, the framework marries group-theoretic invariance with statistical data manifold proximity, operationalizing robust perception without data or architecture constraints.

Conclusion

Zero-shot test-time canonicalization via OOD scoring systematically enhances robustness to affine transformations for vision models, requiring no retraining and no architectural modification. Distance-based OOD scores with random search and local refinement consistently outperform prior approaches, while gating mechanisms resolve the ID/OOD accuracy trade-off. The approach is applicable to any pretrained vision classifier and can be adopted across tasks, with practical runtime trade-offs as the chief limitation.

Future work could pursue sample-efficient search (evolutionary or Bayesian optimization), explore canonicalization for sub-regions/objects within data, and investigate hybrid scoring using foundation models for more complex transformations. Combining OOD scores with energy-based canonicalization and part-based sampling may further improve robustness in realistic scenarios.

(2606.24178)