Equivariant Adaptation of Large Pretrained Models

Summary

The authors tackle an important problem: equivariant NNs are very useful because of capturing important symmetries of the data that lead to higher sample efficiency and more accurate and robust predictions. Yet, EqNNs are very hard to scale because of their specialized architectures. In particular, they cannot benefit from large scale pre-traininig. The authors propose to use the canonical localization network (an EqNN) from [11] to equip pre-trained model with equivariance. The authors train the localization NN jointly with the prediction network. They test both on the original and augmented test set (with various transformations). The authors identify that directly applying the approach from [11] does not work because the localization network can transform the input into atypical views for the training dataset. Then they propose a KL regularization term that encourages a delta function at the identify transformation. Their approach works effectively both for discrete and continuous groups.

Strengths

S1: The idea of equipping pre-trained models with equivariance is very timely and potentially useful for large-scale applications. S2: The regularization term is well-motivated and effective. S3: The experiments are thorough, including both discrete and continuous groups.

Weaknesses

W1: It would be useful to study your approach in more realistic setting where you do not have to augment the test dataset but rather take an existing test dataset. I'd recommend standard out-of-distribution evaluation on IN-{adversarial, sketch, renditions}, INv2 etc. Another domain for experimentation is in Transfer Learning from IN-1K pertaining. There are some standard datasets, such as Flowers-102, Food, etc. W2: Since you want to integrate equiavriance with pre-trained models, I think you should also consider SSL pre-trained checkpoints, since they are the foundation of SOTA pre-training (Pre-training on IN-1K may not give you the best features for downstream tasks). Some pre-trained checkpoints are available on github, e.g. https://github.com/rdangovs/essl/tree/main/imagenet/simclr. W3: As you acknowledged, your model is constrained by the EqNN architectures for the canonical localization network. I wonder if you could consider in your main text discussion more flexible (albeit non-exact) approaches for encouraging symmetry in predictions, such as MSE (https://arxiv.org/abs/2303.02484) and E-SSL (Equivariant Contrastive Learning).

Questions

See the above Weaknesses. My comments should be treated as questions that could improve the quality of your work. I think you have an awesome paper, but I also think that addressing W1-3 could make your work even more impactful.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately addressed the limitations of their work.

Summary

The paper proposed a fine-tuning based adaptation method for equivariance in pre-trained models. The core of the method is the canonicalization function, which is trained to normalize the input to a common orientation. The paper investigates challenges associated with training canonicalization functions and proposes a canonicalization prior regularization to alleviate those challenges. Experiments in image and point-cloud domains demonstrate that fine-tuning a backbone network with the proposed method improves equivariant properties of a backbone model.

Strengths

1. The paper is very clearly written and easy to follow. All parts of the paper/method are nicely motivated. 2. The paper tackles an important problem of efficient equivariance in neural networks. This problem is especially relevant with regard to large models. 3. The paper provides valuable insights into a practical application of canonicalization functions to achieve equivariance/invariance.

Weaknesses

In no particular order. 1. The paper argues about the equivariant adaptation, but only invariance (and only to the rotation group) is demonstrated in the experiments. 2. Although there is a conceptual novelty in the paper (in the part of analyzing the challenge of applying canonicalization functions), the proposed adaptation approach relies heavily on the prior work of [11]. With this, the methodological novelty of the proposed approach is rather limited. 3. The proposed method requires fine-tuning the whole model together with the canonicalization function. This can be troublesome, especially for large models. With this, I am also wondering if the "equivariant adaptation" is the right name, because the essence of the method seems to be more in equivariant fine-tuning. It would be nice to have experiments, where the backbone model is frozen, and only the canonicalization function is trained. 4. The paper claims out-of-distribution robustness as a key benefit of the method (L224). However, it is not clear how the canonicalization function is itself robust to the distribution shift. If the canonicalization function is trained on cifar10 and tested on cifar100, for example, will it still be able to deliver reasonable distribution over a group orbit? In that sense, it is important to note that the proposed method does not provide guaranteed equivariance as equivariant networks do. 5. I am not sure that comparing methods based on raw Accuracy is a fair evaluation protocol. If the goal of equivariant adaptation is to robustify a model (to a symmetry group of interest), then it seems more appropriate to keep track of the relative change of accuracy from original to G-averaged test sets. Otherwise, it is hard to disentangle if the performance improvement is due to better equivariance or due to improved data augmentation (and a higher accuracy on the non-transformed dataset as a result). This seems to be an important part of the comparison, which is missing from the paper. For example, consider Table 2 results for VIT on Cifar10, the LC model loses 0.2% of its Accuracy, while Regularized LC loses 1% of its Accuracy. Which one would we call more robust? Others: 1. Possibly useful related work to resolve the issue discussed in L208 - 219. Moskalev et al. LieGG: Studying Learned Lie Group Generators. NeurIPS22 https://arxiv.org/abs/2210.04345 2. Missed related work. Tai et al. Equivariant Transformer Networks. ICML19. https://arxiv.org/abs/1901.11399 UPD: Authors addressed most of my concerns in the rebuttal.

Questions

I suggest authors address weaknesses for the rebuttal.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitation section nicely highlights limitations of the proposed approach.

Summary

Paper proposes to enforce equivariance/invariance in pretrained models by using the recently proposed canonicalization functions that transform a given input to its canonical form. This allows one to enforce symmetry during the finetuning stage without the need for retraining the model. Paper uses a prior so that the canonicalization network is biased to output identity transformation for the images in the finetuning dataset (i.e., a canonical orientation is defined as the orientation in the finetuning dataset). Empirically, this is shown to be more robust than without the prior.

Strengths

Imposing invariances (via canonicalization) during finetuning rather than pretraining can be very beneficial as different finetuning datasets can have different symmetries. Authors clearly motivate the challenges of using canonicalization naively for a pretrained model that do not exist when training from scratch (e.g., Figure 3) and propose a simple regularization to solve it. Experiments show better performance across multiple finetuning datasets and pretrained models.

Weaknesses

W1. Details of the overall loss function is not presented or is assumed to be known to the reader. It is harder to evaluate the impact of the prior without these details. W2. In Section 3.1, the first drawback of Learned Canonicalization is its undesired augmentation effect during the initial phase of finetuning. - However, this should exist even for the proposed prior-regularized method because the network still begins with a random initialization. - To solve this, one probably has to encode the prior of outputting identity within the architecture (for example, by defining c(x):= I + r(x) where I is the identity and r(x) initialized with very small weights). W3. The assumption that the orientations in the finetuning and pretraining datasets are "similar" (line 153) is not well defined. It may be more reasonable to apply the prior using (a small subset of) the pretraining dataset. This will ensure that the finetuning images match the orientation that the model was pre-trained on. W4. Experiments only evaluate invariance and not equivariance in its general form. I think it is important to demonstrate the advantage of imposing end-to-end equivariance over a pretrained model (as this is different from imposing equivariance in its intermediate representations). W5. Paper can be strengthened with experiments on other transformations than rotations (but not strictly required).

Questions

Q1. Section 4.1: why is the rotation augmentation baseline applied over all angles from -180 to 180 (lines 235-236)? I believe the group is known to be C8 for the LC and prior-regularized LC methods, and should be used for the baseline as well. Q2. Section 4.1: Were rotations part of the data augmentations that were performed on the model during its pretraining on ImageNet? Q3. Section 4.2: VN/CN/PCN- in Table 3 are not defined; I am not sure what these baselines are. Which dataset were the PointNet and DGCNN models pretrained on? Do all the other methods in Table 3 use pretrained models? Q4. Line 247 says that proposed approach tries to map all images in the "original training dataset" to identity. Should it say "finetuning dataset" instead?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

Yes.

Summary

The work proposes a method for making large pre-trained models equivariant to specified group actions. The method uses the canonicalization approach, where the input is oriented to a specific orientation before being fed to the pre-trained model with the help of a trainable canonicalization network. The authors use a prior distribution over the group elements to regularize the output of the canonicalization network. The proposed technique shows robustness with respect to the group actions in different downstream tasks, such as classification and part segmentation.

Strengths

1. The paper pursues a novel approach to make pre-trained models scale equivariant. 2. The proposed technique is simple and effective, making it well-suited for practical adaptation. 3. The work is well-written and easy to follow.

Weaknesses

1. Experimental Setup: One major drawback of this work is the experimental setup. The authors used priors over the discrete group elements (i.e., 8 discrete rotations of C8 for the image domain) for the experiments. During the evaluation on the test set, the images were augmented by the action of group C8. This approach is not appropriate because we are providing the model with information about the test set augmentation during training. On the other hand, for the data augmentation setup, the baseline models were fine-tuned using random rotations between -180 to 180 degrees (i.e., all possible rotations). In this case, no information about the test set augmentation was given to the baseline. This difference in augmentation approaches makes the comparisons inappropriate. An appropriate comparison should either: 1. Use a uniform continuous prior over the group elements. 2. Perform data augmentation only with the discrete group actions while training the baseline. Option 1 is more practical because in the real world, we often do not know specific priors over the group actions. 2. Recent studies have shown that fine-tuning large pre-trained models on small datasets hampers downstream tasks [1] and makes the model vulnerable to out-of-distribution data. Experiments with frozen pre-trained backbones and trainable classification/segmentation modules would be more justifiable. 1.Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution

Questions

1. Should the Kronecker delta functions in line 172 be normalized by the number of group elements for discrete priors? 2. What is the specific structure of the **canonicalization** module used in this work?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Response to the Rebuttal

I thank the authors for the detailed response and additional experiments. It seems finetuning the baselines with C8 augmentation reduced the performance gap between the proposed method and the baseline. And should be added to the main text. I am increasing my score.

Thanks for the clarifications.

I read the whole rebuttal and would like to maintain my score. A revised paper that incorporates the discussions from the rebuttal would be suitable for publication. Small comment: It would be nice to see OOD tasks in future work because they are more realistic and the extension beyond groups is a meaningful one.