Local Canonicalization in Machine Learning
- Local canonicalization is a design pattern that maps data to canonical coordinates defined relative to local entities (sequences, nodes, or agents) to remove nuisance transformations.
- It leverages both analytic methods (e.g., Generalized Procrustes Analysis) and learned equivariant frameworks to align data across diverse applications such as NRSfM, Lorentz networks, and swarm control.
- Empirical studies show that incorporating local canonicalization improves accuracy, robustness, and computational efficiency in tasks including 3D reconstruction, high-energy physics, and graph learning.
Local canonicalization denotes a family of procedures that maps data to canonical coordinates defined relative to a sequence, node, agent, or individual test input, rather than relative to a single dataset-wide reference frame. In recent machine learning literature, this idea appears as per-sequence Generalized Procrustes alignment for deep non-rigid structure-from-motion, equivariantly predicted local frames for Lorentz- and Euclidean-equivariant networks, affine-invariant canonicalization for point clouds, agent-centered coordinate systems for swarm control, and test-time optimization over transformation parameters for robustness to affine distortions (Deng et al., 2024, Spinner et al., 26 May 2025, Di et al., 2023, Wang et al., 17 Sep 2025, Lindner et al., 23 Jun 2026, Gerhartz et al., 30 Sep 2025). Across these settings, the recurring objective is to remove nuisance transformations or rigid ambiguities before or during downstream inference, while preserving the task-relevant structure.
1. Conceptual scope and relation to global canonicalization
Local canonicalization is not a single algorithm but a design pattern. In deep NRSfM, the canonical object is an entire sequence: each test or training sequence is aligned separately to its own canonical frame via Generalized Procrustes Analysis, in contrast to per-dataset canonicalization methods that force a single global canonical frame for all sequences (Deng et al., 2024). In Lorentz-equivariant learning and O(3)-equivariant graph networks, the canonical object is usually a node-wise local frame, predicted equivariantly from local or relational geometry and used to convert tensorial features into invariant coordinates before applying a generic backbone (Spinner et al., 26 May 2025, Gerhartz et al., 30 Sep 2025). In swarm control, the canonical object is an agent-centered observation in which the focal agent is translated to the origin and rotated so that its velocity lies on the -axis (Wang et al., 17 Sep 2025). In ShapeMatcher, by contrast, the method “does not attach an orthonormal triad to each point”; the locality lies in point-wise features that are invariant to global translation, rotation, and scale, while the explicit canonical frame is shared across the whole point cloud (Di et al., 2023). In zero-shot test-time canonicalization, the canonical object is the transformed test sample that minimizes an OOD score over a prescribed transformation family (Lindner et al., 23 Jun 2026).
This literature also distinguishes local canonicalization from several adjacent ideas. “Global” canonicalization or alignment aligns an entire object once to a reference such as PCA axes, whereas local canonicalization uses sequence-wise or node-wise reference systems that adapt to the current sample (Gerhartz et al., 30 Sep 2025). Traditional equivariant GNNs and transformers bake the group action into every layer through specialized tensor-product operations; local canonicalization instead converts data into frame-invariant coordinates and then reuses ordinary MLPs, attention blocks, or message-passing layers (Spinner et al., 26 May 2025, Gerhartz et al., 30 Sep 2025). In LLoCa, data augmentation is explicitly identified as a special choice of reference frame rather than a distinct principle: fixing for a random global recovers Lorentz data augmentation, while exact equivariance requires that the frame itself be learned equivariantly from the data (Spinner et al., 26 May 2025).
| Setting | Canonical object | Mechanism |
|---|---|---|
| Deep NRSfM | Sequence-wise canonical shapes | GPA alignment |
| Lorentz-equivariant networks | Node-wise local Lorentz frame | Predicted tetrads |
| ShapeMatcher | Shared canonical point-cloud frame | Invariant features plus |
| Swarm control | Agent-centered observation | Inverse pose action |
| Test-time robustness | Canonicalized input sample | OOD-score minimization |
| O(3) GNNs | Node-wise local frame | Predicted |
2. Mathematical structure
A canonicalization can be formalized as a map such that, for a transformation group , one has
When a classifier is composed with 0, the resulting predictor 1 is invariant under 2 (Lindner et al., 23 Jun 2026). This formulation is used directly for test-time canonicalization, where one fixes a parameterized family 3 and solves
4
with the energy instantiated by an OOD score (Lindner et al., 23 Jun 2026).
In equivariant architectures, the same principle is expressed through local frames and representation maps. In LLoCa, a local Lorentz frame 5 satisfies 6, where 7 is the Minkowski metric, and under a global Lorentz transformation 8 the frame transforms as 9. A local four-vector 0 is then invariant, since 1. More generally, tensor features are canonicalized as 2, processed by an arbitrary backbone, and mapped back to global coordinates by 3, which guarantees exact equivariance (Spinner et al., 26 May 2025).
The O(3) graph formulation uses the same algebra in node-wise form. With node features 4 and predicted local frames 5, the canonicalized features are
6
Messages from 7 to 8 are built from the pair
9
which is invariant under a global rotation. Message passing can therefore be implemented by an ordinary MLP or attention block while preserving exact equivariance (Gerhartz et al., 30 Sep 2025).
Other domains instantiate the same pattern with different transformation groups. ShapeMatcher explicitly reposes a point cloud by
0
where 1, 2, and 3, while the learned point-wise features are invariant to global translation, rotation, and scaling (Di et al., 2023). LEGO defines, for each agent 4, a canonicalization map
5
which produces the agent-centered observation
6
The local policy outputs 7, and the final action is de-canonicalized by 8, yielding an 9-equivariant policy (Wang et al., 17 Sep 2025).
3. Construction of local frames and canonical representations
The literature differs most sharply in how the canonical frame or canonical representation is obtained. In deep NRSfM, the canonicalization is analytic rather than learned: after a sequence-level context layer produces refined shapes 0, a GPA layer computes
1
and outputs 2. Translation is implicitly removed by centering each 3 before alignment (Deng et al., 2024).
LLoCa learns local Lorentz transformations through a small equivariant “Frames-Net” MLP 4. For each node 5, it predicts three four-vectors
6
where 7 are four-momenta and 8 are Minkowski inner products. These vectors are turned into an orthonormal tetrad by first defining the boost 9 that sends 0 to 1 with 2, then applying 3 to 4 and Gram–Schmidt orthogonalizing their spatial parts to build a spatial rotation 5, and finally setting 6 (Spinner et al., 26 May 2025).
In the O(3) setting, local frames are predicted by a small equivariant network 7 satisfying
8
with the equivariance property 9. The paper emphasizes that this can be implemented by a DG-equivariant MLP that outputs a 0 orthogonal matrix (Gerhartz et al., 30 Sep 2025).
ShapeMatcher uses a different strategy. A Vector-Neuron MLP outputs
1
where 2 is guaranteed 3-invariant. A scale normalization head then predicts
4
so that 5 is invariant under anisotropic or isotropic rescalings. The resulting features are point-wise and permutation-invariant, and no explicit 6-NN or radius search is used (Di et al., 2023).
LEGO constructs the local frame from physically interpretable geometry. The agent-centric 7-axis is the normalized velocity
8
the reference vector is 9 with center of mass 0, and the 1-axis is
2
where 3. The rotation matrix is then 4 (Wang et al., 17 Sep 2025).
4. Integration with learning pipelines
Local canonicalization is usually not an isolated preprocessing step; it is integrated into end-to-end pipelines with task-specific losses. In the NRSfM pipeline, a shared encoder 5 and two heads 6 first predict per-frame shapes and rotations,
7
The stacked shapes are mapped into a feature space 8, processed by a gated-Toeplitz module 9 that uses both frame similarity and temporal order 0 to produce a learnable coefficient matrix 1, reconstructed as 2, mapped back to 3D as 3, and then canonicalized by GPA. Training combines the reprojection loss 4 and the nuclear norm 5 on canonicalized shapes. Because the GPA layer has no parameters and uses a one-step gradient approximation, it disappears at test time; only the backbone and context layers remain (Deng et al., 2024).
LLoCa makes local canonicalization a wrapper around arbitrary backbones. The pipeline is: predict a local frame 6, canonicalize features by 7, process the invariant features with an arbitrary backbone network such as a transformer or GNN, and back-transform the outputs by 8. For message passing, messages are built in the sender’s local frame and transferred to the receiver’s local frame through the invariant frame-shift 9, so that
0
A special case is Lorentz-equivariant scaled dot-product attention, in which the keys and values are moved between local frames by 1 while the scalar products can be computed in the global frame (Spinner et al., 26 May 2025).
The O(3) formulation follows the same architecture at graph level. Features are canonicalized once per node, neighbor features are re-rotated by 2, local relative vectors 3 are encoded, and standard dense layers perform the actual message update. The framework compares scalar messages, Cartesian tensor representations, irreducible Wigner-4 representations, and a learned MLP representation, all within the same local-canonicalization backbone (Gerhartz et al., 30 Sep 2025).
ShapeMatcher integrates canonicalization with segmentation, retrieval, and deformation. The canonicalization branch is supervised by a reconstruction branch that decodes 5 back to 6 in canonical space using Chamfer distance, and by an orthogonality penalty 7. The canonicalization loss is
8
Training also includes cross-task consistency losses for partial and full inputs:
9
which couple canonicalization to segmentation and learned part centers (Di et al., 2023).
Zero-shot test-time canonicalization is distinctive because the classifier is frozen. The canonicalizer is attached at inference time, evaluates an energy 00 for sampled or refined 01, and transforms only the input. The framework benchmarks approximately twenty training-free OOD scores and nine search algorithms, and introduces two gates: a selection gate
02
and an acceptance gate
03
These are designed to avoid unnecessary transformations on already aligned inputs (Lindner et al., 23 Jun 2026).
5. Empirical behavior across domains
The empirical record shows that local canonicalization is effective in several distinct regimes, although the gains depend on how canonicalization is combined with other inductive biases. In deep NRSfM, the sequence-to-sequence pipeline with local canonicalization and sequence modeling reports, on Human3.6M / InterHand2.6M / 3DPW, the following reconstruction errors for “Ours”: GT-H36M MPJPE 04 and Stress 05, HR-H36M MPJPE 06 and Stress 07, I26M MPJPE 08 and Stress 09, and 3DPW MPJPE 10 and Stress 11. The corresponding baselines “Seq2Seq” and “MHR” are worse on all reported columns (Deng et al., 2024). The Human3.6M ablation further reports MPJPE/Stress of 12 without the GPA layer, 13 without the context layer, 14 without the nuclear norm, and 15 for the full model, confirming that local canonicalization, subspace-plus-temporal modeling, and low-rank regularization are jointly required (Deng et al., 2024).
LLoCa reports improvements on both classification and regression in high-energy physics. On JetClass, LLoCa-Transformer reaches accuracy 16 and AUC 17, compared with 18 and 19 for a vanilla transformer, and the paper states that LLoCa-Transformer is approximately 20 faster at training and uses 21–22 fewer FLOPs than L-GATr (Spinner et al., 26 May 2025). On QFT amplitude regression, LLoCa-Transformer reports MSE 23, compared with 24 for a baseline transformer and 25 for prior SOTA L-GATr, while using 26 M FLOPs versus 27 M FLOPs for L-GATr (Spinner et al., 26 May 2025).
ShapeMatcher’s ablation on PartNet isolates the contribution of canonicalization intrinsics under arbitrary poses. The average Chamfer distance decreases from 28 with no disentanglement, to 29 with translation only, to 30 with translation and rotation, and to 31 with translation, rotation, and scale. The paper interprets this as an approximately 32 drop in Chamfer error. Full ShapeMatcher achieves average CDs as low as 33 on PartNet and approximately 34 on ComplementMe (Di et al., 2023).
Zero-shot test-time canonicalization also reports large robustness gains without retraining. On scaled, rotated, and sheared test sets, accuracy improves from 35 to 36 on MNIST, from 37 to 38 on EMNIST, from 39 to 40 on TU Berlin, and from 41 to 42 on ModelNet10. On SI-Score with ResNet-50, rotation accuracy improves from 43 to 44 using PC-kNN. With a fixed budget of 45 model calls, RS-LR achieves the highest accuracy and smallest scoring distance, and threshold sweeps recover more than 46 of original in-distribution accuracy while still boosting distorted-input accuracy by 47–48 points (Lindner et al., 23 Jun 2026).
In O(3)-equivariant graph networks, local canonicalization changes the runtime profile relative to specialized tensor-product architectures. On QM9, Equiformer runs at 49 it/s on RTX 6000, while LoCaFormer reports 50 it/s for scalar messages, 51 it/s for Cartesian tensors, 52 it/s for irreps, and 53 it/s for the learned MLP representation. On N-methylacetamide, tensorial message variants outperform scalar and MLP variants on tensorial targets: for dipole 54, LoCaFormer gives RMSE 55 for Cartesian tensors and irreps, compared with 56 for scalar messages; for polarizability 57, the corresponding RMSEs are 58, 59, and 60 (Gerhartz et al., 30 Sep 2025).
6. Misconceptions, limitations, and representation trade-offs
A recurrent misconception is that local canonicalization always means attaching an explicit orthonormal frame to every local datum. ShapeMatcher is an explicit counterexample: it uses a shared global 61 to repose the entire point cloud, while the “locality” resides in point-wise affine-invariant features rather than in per-point rigid frames (Di et al., 2023). Another misconception is that local canonicalization is equivalent to global canonicalization. The NRSfM results argue against this by replacing per-dataset canonicalization with per-sequence GPA alignment; the stated motivation is that local alignment adapts to different motions in different sequences and avoids bias or overfitting to a single global frame (Deng et al., 2024).
The relationship to data augmentation is subtler. LLoCa explicitly shows that augmentation can be written as a special stochastic choice of reference frames, but exact equivariance emerges only when the local frame is itself learned equivariantly from the data (Spinner et al., 26 May 2025). This suggests that canonicalization and augmentation are not mutually exclusive categories; one can be viewed as a degenerate or approximate instance of the other.
Several papers also expose domain-specific limitations. LEGO assumes agents are in the plane; extending to 3D requires “a choice of canonical frame beyond a single ‘heading + CoM’ vector.” When 62, the method falls back to a fixed global axis, which “arbitrarily breaks symmetry in that degenerate case,” although the paper states that this is rare in practice (Wang et al., 17 Sep 2025). In test-time canonicalization, transforming an already aligned input can hurt in-distribution accuracy, which is why the method introduces OOD-based gating and an additional acceptance check (Lindner et al., 23 Jun 2026). In the O(3) graph setting, local canonicalization does not eliminate representation choices: scalar, Cartesian tensor, irreducible, and learned MLP representations have different complexity and accuracy profiles, and the paper notes that in the very low-data regime data-augmented models can sometimes have lower absolute error despite the stronger inductive bias of exact equivariance (Gerhartz et al., 30 Sep 2025).
A plausible implication of these results is that “local canonicalization” is best understood not as a single invariant recipe but as a representation strategy. The common pattern is exact or approximate removal of nuisance transformations in a local reference system; the substantive differences lie in how the reference system is defined, whether canonicalization is analytic or learned, whether it acts on inputs or intermediate features, and whether outputs must be mapped back to the global frame. Across current formulations, those design choices determine both the theoretical guarantees and the empirical trade-off between flexibility, computational cost, and task-specific accuracy.