[论文解读] 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 Science被引用 3

一句话总结

本文提出一种轻量级、即插即用的方法，通过利用数据集先验信息的可学习规范化网络，使大规模预训练深度学习模型对旋转等变换实现等变性。通过将规范化形式与训练数据分布对齐，该方法在仅增加0.3%参数量和7.3%推理开销的前提下，实现了SOTA级别的分布外旋转鲁棒性，显著提升了泛化能力，且无需微调主干网络。

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.

研究动机与目标

• 解决在不进行网络架构重构或高计算成本的前提下，使大规模预训练模型对旋转等变换实现等变性的挑战。
• 克服基于规范化方法的等变性中因分布错位导致的性能下降问题，即朴素规范化会因输入分布偏移而损害性能。
• 通过将等变性与主预测网络解耦，实现基础模型对分布外变换的鲁棒泛化能力。
• 证明规范化函数中依赖数据集的先验信息对在等变适应下保持预训练模型性能至关重要。
• 提供一种实用、高效且可泛化的框架，用于将大规模模型适配至图像和点云等多样化领域中的对称群。

提出的方法

• 一个可学习的规范化网络在将输入数据传递给预训练预测网络前，将其转换为规范化形式，从而将等变性与主模型解耦。
• 通过一种匹配预测规范化形式分布与训练数据分布的先验损失对规范化函数进行正则化，确保主干网络接收到的数据分布保持在分布内。
• 该方法使用变换的群表示（如SO(2)、SO(3)）来定义对称群，规范化网络学习将所有群变换后的输入映射到单一规范化方向。
• 通过每次输入仅进行一次前向传播避免对称化操作，相比需要对多个变换输入进行多次前向传播的方法，计算效率更高。
• 该方法作为插件模块应用：仅对规范化后的输入进行主模型微调，同时联合训练规范化网络与先验损失。
• 在COCO、ShapeNet等标准基准上对图像和点云任务进行了评估，并通过消融实验验证了先验正则化的必要性。

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}$

实验结果

研究问题

• RQ1 轻量级规范化模块是否能在不改变网络架构或显著性能下降的前提下，使大规模预训练模型对旋转实现等变性？
• RQ2 若规范化分布与训练数据分布不匹配，是否会导致预训练模型在等变适应中性能下降？
• RQ3 先验正则化在将规范化形式与数据分布对齐以保持模型准确率方面有多有效？
• RQ4 该方法是否能在不同模态（图像、点云）和变换群（SO(2)、SO(3)）之间实现泛化？
• RQ5 在大规模模型中，对分布外变换的鲁棒性与计算成本之间存在怎样的权衡？

主要发现

• 所提方法在图像和点云任务中均实现了SOTA级别的旋转鲁棒性，在COCO和ShapeNet-part基准上优于现有方法。
• 在规范化函数中加入先验损失可显著提升性能，尤其在SO(3)/SO(3)设置下，原始PointNet和DGCNN的准确率出现大幅下降。
• 该方法仅使推理时间增加7.3%，参数量增加0.3%，充分体现了在大规模模型中的高效性。
• 实证结果表明，若缺乏先验正则化，规范化会导致分布偏移，从而即使实现了等变性，也会损害预训练模型的性能。
• 该方法在保持高准确率的同时，实现了对任意旋转（SO(3)）的强泛化能力，即使主干网络未经过旋转增强训练。
• 消融实验确认先验正则化至关重要——缺乏该正则化的模型无法保持原始模型在分布内任务中的性能。

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