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

Research Question and Approach

Pretrained vision models often struggle with affine transformations like rotation, scaling, or shearing, even when the object class remains unchanged. While traditional solutions involve retraining models or building in architectural equivariance, this paper explores test-time canonicalization—a method that leaves the original classifier untouched. The authors propose reframing canonicalization as an out-of-distribution (OOD) detection problem. By treating an OOD score as an energy function, the system searches for the transformation that minimizes this score, effectively mapping the input to a canonical form that aligns with the training distribution.

Methodology

The authors systematically evaluate approximately twenty OOD scoring functions and nine search algorithms across diverse benchmarks, including handwritten characters, sketches, natural images, and 3D point clouds. To address the common issue where canonicalization degrades performance on already-aligned (in-distribution) inputs, they introduce a two-part gated mechanism: a selection gate that triggers canonicalization only for OOD inputs, and an acceptance gate that ensures the transformation actually reduces the OOD score.

The study finds that distance-based OOD scores, such as k-nearest neighbors (kNN) and prototype-based methods, consistently outperform logit-based scores when used as energy functions. Regarding optimization, the authors demonstrate that complex, bespoke search procedures are unnecessary; simpler global optimizers, specifically random search paired with local refinement, provide the best performance. The gated mechanism successfully resolves the accuracy trade-off, allowing the model to maintain high performance on standard inputs while significantly improving robustness against transformed inputs.

This work provides a flexible, training-free framework for improving the robustness of existing vision models. By decoupling the canonicalization process from the classifier architecture, it allows researchers to enhance deployed models without the high computational costs of retraining or the architectural constraints of equivariant networks. The use of OOD scores as a general-purpose energy function also bridges the gap between OOD detection research and spatial robustness.