[Paper Review] Equivariant Adaptation of Large Pretrained Models

Arnab Kumar Mondal, Siba Smarak Panigrahi+3 more|arXiv (Cornell University)|Oct 2, 2023

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper proposes a lightweight, plug-and-play method to make large pretrained deep learning models equivariant to transformations like rotations by using a learnable canonicalization network informed by dataset priors. By aligning the canonical forms with the training distribution, the approach achieves state-of-the-art robustness to out-of-distribution rotations with only 0.3% extra parameters and 7.3% inference overhead, significantly improving generalization without retraining the backbone.

ABSTRACT

Equivariant networks are specifically designed to ensure consistent behavior with respect to a set of input transformations, leading to higher sample efficiency and more accurate and robust predictions. However, redesigning each component of prevalent deep neural network architectures to achieve chosen equivariance is a difficult problem and can result in a computationally expensive network during both training and inference. A recently proposed alternative towards equivariance that removes the architectural constraints is to use a simple canonicalization network that transforms the input to a canonical form before feeding it to an unconstrained prediction network. We show here that this approach can effectively be used to make a large pretrained network equivariant. However, we observe that the produced canonical orientations can be misaligned with those of the training distribution, hindering performance. Using dataset-dependent priors to inform the canonicalization function, we are able to make large pretrained models equivariant while maintaining their performance. This significantly improves the robustness of these models to deterministic transformations of the data, such as rotations. We believe this equivariant adaptation of large pretrained models can help their domain-specific applications with known symmetry priors.

Motivation & Objective

• To address the challenge of making large pretrained models equivariant to transformations like rotations without architectural redesign or high computational cost.
• To overcome the distribution misalignment issue in canonicalization-based equivariance, where naive canonicalization degrades performance by shifting input distribution.
• To enable robust generalization of foundation models to out-of-distribution transformations by decoupling equivariance from the main prediction network.
• To demonstrate that dataset-dependent priors in the canonicalization function are essential for preserving the performance of pretrained models under equivariant adaptation.
• To provide a practical, efficient, and generalizable framework for adapting large models to symmetry groups in diverse domains like images and point clouds.

Proposed method

• A learnable canonicalization network transforms input data into a canonical form before passing it to a pretrained prediction network, decoupling equivariance from the main model.
• The canonicalization function is regularized using a prior loss that matches the distribution of predicted canonical forms to the training data distribution, ensuring in-distribution inputs to the backbone.
• The method uses a group representation of transformations (e.g., SO(2), SO(3)) to define the symmetry group, with the canonicalization network learning to map all group-transformed inputs to a single canonical orientation.
• The approach avoids symmetrization by using a single forward pass per input, making it computationally efficient compared to methods requiring multiple passes over transformed inputs.
• The method is applied as a plug-in module: the pretrained model is fine-tuned only on the canonicalized inputs, with the canonicalization network trained jointly using the prior loss.
• The framework is evaluated on image and point cloud tasks using standard benchmarks (COCO, ShapeNet), with ablation studies confirming the necessity of prior regularization.

Figure 1 : Predicted masks from the Segment Anything Model (SAM) [ 21 ] , showcasing both the original model and our proposed equivariant adaptation for 90 ∘ counter-clockwise rotated input images taken from the COCO 2017 dataset [ 22 ] . Our method makes SAM equivariant to the group of $90^{\circ}$

Experimental results

Research questions

• RQ1 Can a lightweight canonicalization module make large pretrained models equivariant to rotations without architectural changes or significant performance drop?
• RQ2 Does misalignment between the canonicalization distribution and the training data distribution degrade the performance of pretrained models in equivariant adaptation?
• RQ3 How effective is prior regularization in aligning the canonical forms with the data distribution to preserve model accuracy?
• RQ4 Can this method generalize across different modalities (images, point clouds) and transformation groups (SO(2), SO(3))?
• RQ5 What is the trade-off between robustness to out-of-distribution transformations and computational cost in large-scale models?

Key findings

• The proposed method achieves state-of-the-art robustness to rotations in both image and point cloud tasks, outperforming existing methods on COCO and ShapeNet-part benchmarks.
• Adding the prior loss to the canonicalization function significantly improves performance, especially in the SO(3)/SO(3) setting, where vanilla PointNet and DGCNN suffer large accuracy drops.
• The method increases inference time by only 7.3% and adds just 0.3% parameters, demonstrating high efficiency for large models.
• Empirical results show that without prior regularization, canonicalization leads to distribution shift, degrading the performance of the pretrained model despite achieving equivariance.
• The method enables strong generalization to arbitrary rotations (SO(3)) while maintaining high accuracy, even when the backbone was not trained with rotation augmentation.
• Ablation studies confirm that prior regularization is essential—models without it fail to preserve the original model’s in-domain performance.

Figure 2 : Training and inference with our proposed regularized canonicalization method. The canonicalization function outputs a distribution over image orientations used to canonicalize the input image. Additionally, during training, this predicted distribution is regularized to match the orientati