Test-Time Canonicalization for Robust Inference

• Test-Time Canonicalization is the process of mapping inputs to a canonical representative of their symmetry class, reducing nuisance variation at inference.
• It employs methods such as learned transformations, energy minimization, and OOD scoring to select the most typical representative among symmetries.
• By decoupling symmetry removal from the main predictor, it offers computational efficiency, simplified modeling, and robust performance across domains.

Searching arXiv for papers on test-time canonicalization and closely related canonicalization methods. Found relevant papers, including "Test-Time Canonicalization by Foundation Models for Robust Perception" (Singhal et al., 14 Jul 2025), "Zero-Shot Test-Time Canonicalization using Out-of-Distribution Scoring" (Lindner et al., 23 Jun 2026), "Improved Canonicalization for Model Agnostic Equivariance" (2405.14089), "Equivariance with Learned Canonicalization Functions" (Kaba et al., 2022), and domain-specific canonicalization papers in NRSfM, language modeling, motion retargeting, and generative modeling. Test-time canonicalization is the inference-time practice of mapping an observed input, an intermediate representation, or a generated sample to a canonical representative of its symmetry class before applying a downstream predictor or after sampling from a canonical slice. Across recent work, the object being canonicalized varies: images may be rotated or relit to a visually typical view, token prefixes may be restricted to canonical tokenizations, point clouds and skeletons may be mapped to canonical frames, predicted 3D sequences may be aligned to remove rigid ambiguity, and generative models may sample in canonical coordinates and only afterward restore symmetry by randomization over the group (Singhal et al., 14 Jul 2025, Lindner et al., 23 Jun 2026, Chatzi et al., 6 Jun 2025, Zhou et al., 16 Feb 2026). The common objective is to factor out nuisance variation without redesigning the main predictor around strict equivariance or retraining it on exhaustive augmentations.

1. Formal definition and mathematical setting

A standard formulation treats canonicalization as orbit selection under a group action. If a group $G$ acts on inputs $x\in\mathcal X$, a canonicalizer seeks a representative from the orbit $Gx$ that is consistent across transformed versions of the same underlying object. In architecture-agnostic equivariant adaptation, the canonicalized predictor is written as

$f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$

where $p$ is the downstream predictor, $c(x)$ is the estimated canonicalizing transformation, and $c'(x)$ is the corresponding output-space action (2405.14089). In learned canonicalization for equivariant learning, the same idea appears as

$\phi(x)=h'(x)\,f\!\left(h(x)^{-1}x\right),$

with $h$ a canonicalization function and $h'$ the output-side inverse transport (Kaba et al., 2022).

A complementary formulation defines canonicalization by energy minimization over transformations: $x\in\mathcal X$0 which underlies test-time search methods that score transformed candidates and select the minimum-energy representative (Singhal et al., 14 Jul 2025). A closely related zero-shot formulation writes

$x\in\mathcal X$1

and interprets $x\in\mathcal X$2 as an out-of-distribution score, so that canonicalization becomes OOD minimization over inverse transformations (Lindner et al., 23 Jun 2026).

Not all canonicalization is single-valued. A canonicalization perspective on invariant and equivariant learning defines a set-valued map

$x\in\mathcal X$3

satisfying

$x\in\mathcal X$4

with invariant prediction obtained by canonical averaging,

$x\in\mathcal X$5

This accommodates cases where symmetry or stabilizers prevent a unique canonical representative (Ma et al., 2024).

The same structural idea extends beyond geometric groups. In autoregressive language modeling, canonicalization is defined relative to a tokenizer: a next token $x\in\mathcal X$6 is admissible exactly when the extended prefix $x\in\mathcal X$7 remains canonical, and canonical decoding masks all continuations that violate this condition (Chatzi et al., 6 Jun 2025).

2. Inference-time algorithmic patterns

Recent work realizes test-time canonicalization through several distinct mechanisms. Some methods run a small learned canonicalizer once and then transform the input into a canonical frame. Others search over transformed candidates, using either a learned score, a foundation-model energy, or an OOD score. Language-model work canonicalizes incrementally at every decoding step. Generative modeling often inverts the chronology: it trains on canonicalized data, generates directly in canonical space, and restores symmetry only after sampling (Kaba et al., 2022, 2405.14089, Singhal et al., 14 Jul 2025, Lindner et al., 23 Jun 2026, Chatzi et al., 6 Jun 2025, Zhou et al., 16 Feb 2026).

Setting	Canonicalized object	Test-time mechanism
Architecture-agnostic equivariance	Input sample	Predict $x\in\mathcal X$8, apply $x\in\mathcal X$9, run backbone once
Search-based perception	Transformed image candidates	Enumerate or optimize over $Gx$0, choose minimum-energy view
OOD-based canonicalization	Inverse-transformed input	Minimize OOD score over $Gx$1
Canonical autoregressive generation	Token prefix	Mask next tokens that make the prefix non-canonical
Canonical diffusion / flow	Generated sample on slice	Sample in canonical space, optionally project during sampling, randomize over group at end

Direct learned canonicalizers are the simplest operationally. In the main learned-direct variant of canonicalization for equivariant learning, inference is one extra forward pass through a small canonicalization module plus one transformation of the input, rather than repeated evaluation over group elements (Kaba et al., 2022). By contrast, EquiOptAdapt explicitly evaluates all transformed candidates for a discrete group using a non-equivariant scorer $Gx$2, but still runs the large downstream model only once on the chosen canonicalized input (2405.14089).

Search-based perception generalizes this pattern. FoCal adopts a “Vary and Rank” pipeline: generate candidate transformations, score each candidate with a combined CLIP and diffusion energy,

$Gx$3

then choose

$Gx$4

For continuous spaces it uses Bayesian Optimization rather than exhaustive enumeration (Singhal et al., 14 Jul 2025). OOD-based canonicalization retains the same outer optimization form but broadens the energy design space by letting any OOD score serve as $Gx$5, and adds selection and acceptance gates so already aligned inputs can bypass canonicalization (Lindner et al., 23 Jun 2026).

3. Relation to invariance, equivariance, and symmetrization

Canonicalization is often presented as an alternative to equivariant architecture design rather than a replacement for symmetry as such. If the canonicalization function is itself appropriately equivariant, then the transformed predictor becomes equivariant for any downstream model $Gx$6, which is precisely why canonicalization is attractive for adapting pretrained models at inference time (2405.14089). A corresponding universality result states that with a continuous $Gx$7-equivariant canonicalizer and a universal backbone, the overall model is universal for continuous $Gx$8-equivariant functions (Kaba et al., 2022).

A more structural result is that frame averaging and canonicalization are equivalent at the level of invariant averaging. For any frame $Gx$9, there exists an orbit canonicalization $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$0 such that

$f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$1

and the induced canonicalization removes stabilizer-induced duplication by reducing the effective representation count from $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$2 to $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$3 (Ma et al., 2024). This reframes many “frame” methods as canonicalization procedures implemented in group space rather than input space.

Canonicalization also differs from augmentation and test-time augmentation. Augmentation asks the model to absorb nuisance variability during training; TTA averages predictions over multiple transformed copies at inference. Canonicalization instead attempts to identify one representative before prediction. The literature repeatedly treats this as the main practical advantage: the downstream model can remain non-equivariant and often needs to be evaluated only once at test time (2405.14089, Singhal et al., 14 Jul 2025, Schmidt et al., 9 Oct 2025).

The distinction between training-time and test-time use is nevertheless important. Some systems use canonicalization machinery only as training scaffolding. DRACO uses C3DPO-derived sparse canonicalization during training to teach a dense NOCS predictor, but at inference the network directly predicts dense canonical geometry from RGB without explicitly rerunning C3DPO (Sajnani et al., 2020). In deep NRSfM, the GPA layer is explicitly a training regularization module and “can be left out of the computation during testing,” even though the underlying alignment procedure is naturally interpretable as a per-sequence inference-time canonicalizer (Deng et al., 2024).

4. Representative domain instantiations

In robust visual perception, test-time canonicalization is used to move images toward visually typical states. FoCal canonicalizes images by optimizing over 2D rotations, 3D viewpoints, illumination shifts, and day-night variations using foundation-model priors. On Objaverse-LVIS it improves accuracy on the hardest viewpoints from $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$4 to $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$5, and on rotated COCO segmentation it matches PRLC’s mAP while improving pose accuracy (Singhal et al., 14 Jul 2025). A related zero-shot line treats transformed inputs as OOD relative to the training pose distribution and finds that distance-based OOD scores, especially kNN- and prototype-style methods, outperform prior logit-based energies, with random search plus local refinement performing best overall (Lindner et al., 23 Jun 2026).

For adapting pretrained models under discrete transformations, EquiOptAdapt canonicalizes by scoring all inverse-transformed candidates $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$6 with a non-equivariant network $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$7, selecting

$f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$8

and then running the downstream model once on $f(x) = c'(x)\, p\!\left(c(x)^{-1}x\right),$9. On rotated-image benchmarks it closes the gap between standard accuracy and transformed-test accuracy, and the paper reports that the canonicalization process is up to $p$0 faster than EquiAdapt (2405.14089).

In 3D perception and geometry, test-time canonicalization often means explicit estimation of pose, translation, and scale or direct prediction of canonical coordinates. DRACO predicts dense object-centric depth and a dense NOCS map from one or more RGB images at inference, yielding a canonical category-level coordinate map without requiring keypoints or camera poses at test time (Sajnani et al., 2020). ShapeMatcher takes a partial target point cloud $p$1, predicts

$p$2

and can explicitly form a canonicalized shape

$p$3

which then supports segmentation, retrieval, and deformation (Di et al., 2023).

Motion and sequence modeling provide a different interpretation. MoCaNet canonicalizes 2D skeleton sequences along two axes: structure canonicalization removes body-shape variation and view canonicalization removes camera pose variation, enabling retargeting from monocular 2D video to 3D motion without 3D labels (Zhu et al., 2021). Deep NRSfM canonicalizes predicted 3D shape sequences by generalized Procrustes alignment,

$p$4

so that low-rank sequence regularization acts in a coherent shared frame rather than being corrupted by framewise rotational ambiguity (Deng et al., 2024).

In language modeling, the canonicalization target is the token sequence rather than a geometric pose. Canonical autoregressive generation observes that a tokenizer presents only one canonical tokenization during training, yet a model can still generate non-canonical token sequences at inference. Its central theorem states that for BPE, if a prefix $p$5 is non-canonical, then $p$6 is also non-canonical for any token $p$7. Canonical sampling therefore masks every token that would make the current prefix non-canonical and renormalizes over the rest, guaranteeing that the full generated sequence is canonical (Chatzi et al., 6 Jun 2025).

In symmetry-aware generative modeling, canonicalization often moves from input preprocessing to generative gauge fixing. Learned canonicalization for diffusion defines

$p$8

trains a non-equivariant denoiser on $p$9, and samples directly in canonical space (Sareen et al., 14 Jan 2025). Canonical diffusion and CanonFlow make this explicit at test time: sample from a canonical slice prior, integrate reverse dynamics in canonical space, optionally reproject intermediate states with projected canonical sampling, and finally recover the invariant distribution by sampling a random group element $c(x)$0 from Haar measure and outputting $c(x)$1 (Zhou et al., 16 Feb 2026).

5. Empirical advantages and recurring benefits

The principal advantage claimed across the literature is symmetry handling without redesigning the main predictor. FoCal argues that transform-specific retraining is not necessary for robust perception and uses foundation-model priors to canonicalize inputs at test time without re-training or architectural changes (Singhal et al., 14 Jul 2025). OOD-based canonicalization is similarly designed for a fixed pretrained classifier, with robustness added post hoc by optimizing over transformations and optionally gating the procedure when the input already appears in-distribution (Lindner et al., 23 Jun 2026).

A second benefit is computational asymmetry: the canonicalizer can be much smaller or cheaper than the predictor. Learned canonicalization functions are proposed specifically as a way to obtain exact equivariance with ordinary backbones while paying the symmetry cost once up front rather than throughout the network (Kaba et al., 2022). In point-cloud and $c(x)$2-body experiments, the canonicalizer is deliberately shallow and inference overhead is reported as negligible relative to the prediction network (Kaba et al., 2022). EquiOptAdapt makes the same design choice by evaluating multiple transformed candidates only through the small canonicalizer rather than through the large pretrained model (2405.14089).

A third benefit is statistical or geometric simplification. In NRSfM, canonicalization is necessary because applying a nuclear norm directly to arbitrarily rotated sequences would inflate rank and spoil the subspace prior; alignment removes rigid inter-frame motion first (Deng et al., 2024). In canonical diffusion, the theoretical advantage is framed as elimination of symmetry-induced mixture complexity in diffusion scores and removal of a nonnegative symmetry-ambiguity term in flow-matching conditional variance, which in turn improves few-step generation quality (Zhou et al., 16 Feb 2026).

Some works also provide formal distributional guarantees. Canonical autoregressive generation proves that the token-sequence distribution induced by canonical sampling is strictly closer in KL divergence to the true training distribution than standard sampling under the theorem’s assumptions,

$c(x)$3

because non-canonical token mass is removed and redistributed onto canonical continuations (Chatzi et al., 6 Jun 2025).

6. Ambiguity, failure modes, and scope limitations

A central limitation is non-uniqueness. Inputs with nontrivial stabilizers may not admit a unique canonical form. The canonicalization literature therefore distinguishes canonicalizable from uncanonicalizable inputs and often falls back to a set of representatives or averaging when uniqueness is impossible (Ma et al., 2024). Learned canonicalization functions face a related continuity problem: small changes in nearly symmetric inputs can produce large changes in the selected canonical pose, and the literature explicitly notes that smoothness is not guaranteed (Kaba et al., 2022).

Search-based methods add their own caveats. FoCal’s exact invariance intuition assumes invertible transformations, whereas 3D viewpoint changes are not invertible in image space, so viewpoint canonicalization is best understood as approximate canonicalization by search over synthesized views (Singhal et al., 14 Jul 2025). The same method is computationally heavy, with runtime dominated in some settings by transformation generation and energy evaluation rather than the downstream predictor (Singhal et al., 14 Jul 2025). OOD-based canonicalization shows that always-on canonicalization can hurt clean accuracy because canonicalizing an already aligned input may perturb discriminative features; this motivates the paper’s selection and acceptance gates (Lindner et al., 23 Jun 2026).

Training assumptions are equally important. Canonicalization-prior training typically presumes that the dataset already shares a global canonical mode. Robust Canonicalization through Bootstrapped Data Re-Alignment argues that real-world datasets violate this assumption, so the learned canonicalizer becomes brittle; its remedy is an iterative training-time re-alignment procedure that contracts the Fréchet variance of the pose distribution under mild assumptions for compact groups (Schmidt et al., 9 Oct 2025). This suggests that test-time canonicalization quality can be limited as much by training-set alignment as by inference algorithm design.

Several papers also delimit the scope of their guarantees. Canonical autoregressive generation proves improvement in token space, not string-space quality, and explicitly does not claim a string-space KL inequality (Chatzi et al., 6 Jun 2025). The deep NRSfM paper does not actually deploy GPA at test time in its final evaluation pipeline, even though the alignment operator is defined per sequence and could be applied to a test sequence independently (Deng et al., 2024). In symmetry-aware generation, canonicalization may intentionally discard nuisance degrees of freedom and return only canonical-pose samples unless a random group action is applied afterward (Sareen et al., 14 Jan 2025, Zhou et al., 16 Feb 2026).

7. Conceptual synthesis and current directions

Across domains, test-time canonicalization now spans three recurring templates. The first is canonicalize–predict: estimate a pose, ordering, token constraint, or orbit representative and then run an ordinary model in that canonical frame (2405.14089, Chatzi et al., 6 Jun 2025). The second is search–rank–predict: generate transformed candidates, score them with a learned or pretrained energy, and select the most typical or least OOD representative before inference (Singhal et al., 14 Jul 2025, Lindner et al., 23 Jun 2026). The third is canonicalize–generate–randomize: learn a model on a canonical slice and restore symmetry only after sampling (Sareen et al., 14 Jan 2025, Zhou et al., 16 Feb 2026).

This suggests a broader conceptual shift. Canonicalization is no longer confined to rigid geometric preprocessing. It has become a general inference-time strategy for quotienting nuisance structure, whether the nuisance is camera pose, body morphology, frame ambiguity, tokenization multiplicity, or permutation–rotation redundancy in molecule generation. The strongest formal developments treat canonicalization as a complete view of frame design, a universal route to invariant and equivariant learning, or an exact factorization device for invariant generative distributions (Ma et al., 2024, Kaba et al., 2022, Zhou et al., 16 Feb 2026).

The open questions identified in the literature are correspondingly broad. Search-based perception highlights automatic transformation-family selection and more efficient optimization as unresolved practical issues (Singhal et al., 14 Jul 2025). Architecture-agnostic image canonicalization currently focuses on discrete groups and explicitly leaves continuous rotations as future work (2405.14089). Language-model work identifies practical string-level evaluation as future work (Chatzi et al., 6 Jun 2025). Generative modeling raises unresolved questions about discontinuity, stabilizers, and the choice of canonical slice (Sareen et al., 14 Jan 2025, Zhou et al., 16 Feb 2026). What is already clear, however, is that test-time canonicalization has become a unifying operational principle: rather than forcing the predictor to internalize all symmetry, it externalizes nuisance removal into an inference-time canonicalization step and lets the main model operate in a more stable coordinate system.