[論文レビュー] 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%の推論オーバーヘッドで、分布外回転に対して最先端のロバスト性を達成し、一般化性能を顕著に向上させる。

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)）を用いて対称性群を定義し、正準化ネットワークが群変換されたすべての入力を同一の正準方向にマッピングするように学習する。
• 複数回の入力変換に対する前向きパスを必要とする手法とは異なり、1回の前向きパスで済ませることで、計算効率が向上する。
• プラグインモジュールとして適用可能である。事前学習モデルは正準化された入力でのみ微調整され、正準化ネットワークは事前損失を用いて共同で学習される。
• 標準ベンチマーク（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 この手法は、画像や点群といった異なるモodal（モダリティ）およびSO(2)、SO(3)といった異なる変換群に一般化可能か？
• RQ5 大規模モデルにおいて、分布外変換に対するロバスト性と計算コストのトレードオフはいかほどか？

主な発見

• 提案手法は、COCOおよびShapeNet-partベンチマークで、画像および点群タスクの両方において、既存手法を上回る最先端の回転に対するロバスト性を達成した。
• 正準化関数に事前損失を追加することで、特にvanilla PointNetやDGCNNが大幅な精度低下を示すSO(3)/SO(3)設定において、性能が顕著に向上した。
• 推論時間はわずか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