Zenn
📄

ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoder

2023/01/31に公開
画像認識
論文
tech

論文

ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoder
https://arxiv.org/pdf/2301.00808.pdf

結論

  • ConvNeXt V2はConvNeXtに、自己教師あり学習の要素として、Masked Autoencoders(MAE)を結合したもの
  • ConvNeXtとMAEはアーキテクチャが大きく異なるため、feature collapseが発生してしまい簡単には結合できない。Global Response Normalization（GRN） layerを加え、内部チャネルを強調することで解決
  • ImageNet classification, COCO detection, ADE20K segmentationなどのベンチマークを更新した

背景・課題

画像認識のシステムにおいて、事前学習済みのlarge-scaleモデルを使うことは必要不可欠になっている。そして、そのシステムのパフォーマンスを左右する要素は３つあり、ニューラルネットワークアーキテクチャの選び方、ネットワークの学習方法、学習に使うデータの選び方である。

ニューラルネットワークのアーキテクチャ

ConvNetsの大規模化やTransformerアーキテクチャの出現によって、画像認識における精度や、モデルやデータのスケーリング面も向上してきている。一方、ConvNeXtでは、設計空間を見直すことで、純粋なConvNetsであっても精度やスケーリング性を改良することができることが示されている。
つまり、性能向上においては、アーキテクチャの大規模化だけではなく、設計空間を見直すことも非常に重要であることがわかる。しかし、こういった設計空間を見直す方法は、依然としてImageNetを用いた教師あり学習の性能をベンチマークとしている。

ネットワークの学習方法と学習に使うデータ

画像認識における学習方法は、ラベルを使った教師あり学習から、自己教師あり学習を事前に行う形式へとシフトしている。特に、Masked Autoencoderは画像認識領域での成功から、人気なアプローチになっている。しかし、事前学習をするときには、その後適用するアーキテクチャのデザインをほぼ決めてしまう。例えば、MAEの場合はVison Transformerのアーキテクチャを使って、自己教師あり学習を行う。
自分が使いたいアーキテクチャと、事前学習済みのフレームワークを結合できればよいが、かなりチャレンジングなことである。MAEはtransformerアーキテクチャ用にendode-decoderが最適化されたデザインになっており、ConvNetsとは大きくアーキテクチャが異なるためである。加えて、アーキテクチャと学習目的の関係が違う場合、パフォーマンスが向上しない場合もある。
https://arxiv.org/abs/2206.07700

Global Response Normalization layer(GRN layer)の導入

ConvNeXtの学習コンフィグを変えた時の特徴量空間を分析することで、masked inputの学習時にfeature collapseが潜在的に起きていることを特定した。そのため、Global Response Normalization layerを導入し、内部のチャネルの競合関係を強調させることで解決した。これにより、アーキテクチャデザインの異なるものを修正し、再利用することができるようになった。

ConvNets

1980年代にはじめて導入され、back-propagationを使った学習により数々の発展をしてきた。また、それらの検証は主にImageNet datasetによる教師あり学習によって明らかにされてきた。

ConvNeXt

ConvNetの設計空間を再検討したものであり、Transformer系アーキテクチャと遜色ないレベルの性能を達成したものである。

https://arxiv.org/abs/2201.03545

Masked Autoencoders(MAE)

Masked Autoencoderは、Vision Transformer(ViT)の画像認識のための事前学習として、入力画像のパッチをランダムにマスクし、元画像を復元できるよう学習する技術

https://arxiv.org/abs/2111.06377

オリジナルのMAEは、Encoder-Decoderの非対称性から直接は適用できなかったため、ConvNetsを意識しているMCMAEを使ったようである。

https://arxiv.org/abs/2205.03892

Fully Convolutional Masked Autoencoder

Fully Convolutional Masked AutoEncoder(FCMAE)のdesignは以下のような構成

masking strategy

32×32のpatchのうち、60%をmaskingするようにし、augumentationは最小限にする。

Encoder design

エンコーダーのデザインは慎重に考えなければならない。Transformer-baseのモデルの場合、maskingされたpatchの特徴は、上手く伝搬できる。一方でConvNets-baseの場合、maskingされた情報はただカット＆ペーストされて畳み込まれてしまう。そのため、Masked Autoencoderでは、75%というhigh ratioのmaskingによって表現力を増強していたが、ConvNetsでは、上手くいかない可能性が高い。
この解決には、3Dのpoint cloudのような"“sparse data perspective"における学習がいいヒントになった。それにより、maskingされた画像は、2D空間におけるsparse配列とみなすことができ、学習時にsparse conv layerを組み込むことにした。

Decoder design

デコーダーのデザインは、プレーンなConvNeXt blockを入れるだけのシンプルなものになった。Encoder-Decoderの非対称性が気になったらしき、複雑にすることも考えたらしいが、結果的にはシンプルな方が精度がよかったらしい。

Reconstruction target

target imageとreconstructedはMSElossを利用した。

design evaluation

Pre-TrainとFine-TuningではImageNet-1Kを800~1000epoch学習させ、validationでは、224×224のcenter-cropした画像を使って評価した。

  • sparce conv layerの性能
    sparce conv layerを入れることで、大幅に性能が向上していることがわかる。
  • decoder designの性能比較
    grayでハイライトされている条件を今回は採用。
  • 教師あり学習と事前学習の性能
    ConvNeXtで扱っていた教師あり学習と、自己教師あり学習のFCMAEで、Sup, 100epと同じ条件で比較した結果(Sup, 300epはおまけ)。これまでのtransformer-baseの論文とは対称的で、事前学習のほうが精度がよくなっているところも面白いポイント

Global Response Normalization

Feature collapse

FCMAEにおいて、ConvNeXtV1-Baseのモデルではチャネル間の特徴量の多くが消えてしまったり、飽和してしまうような"feature collapse"が発生していた。

暗いpatchは特徴量が消失しており、黄色のpatchは特徴量が飽和していることを示す。ConvNeXtV1-Baseのモデルは、特徴量の多くが消滅もしくは飽和してしまっている。

この現象を定量的に評価するため、ImageNet-1K validationをする際に、FCMAE、ConvNeXt supervided model、MAE pre-trained ViT modelそれぞれのlayerにおいて、特徴量のcosine distanceを計測した。

Feature Cosine Distanceは1に近いほど、多様的な特徴量を捉えていることを示す。
ConvNeXtV1-Baseのモデルは、実際に"feature collapse"がかなり発生していることが伺える。

こういった、特徴量の多様性を強化には、response normalizationが有効であり、本論文ではGlobal Response Normalization（GRN）を実装した。GRNの構成要素は大きく３つある。

  1. Global Feature Aggregation
    まず第一に、空間的な特徴量を作用素G(·)によってチャネル数まで集約させてしまうことである。

    実装する際は、L2normで集約することが一番精度が高かったようである。

  2. Feature Normalization
    次に、G(·)によって集約された特徴量を標準化させる。標準化によって、特徴量間が相互的に抑制させてしまうことを防ぐことができる。

    実装時は、標準化ではなく、Divisive normalizationが一番精度がよかったようである。

  3. Feature Calibration
    最後に、1. 2.を作用させた値によって、入力値をキャリブレーションする。

    キャリブレーションは、残差を足すような形が最も精度が高かったので、それを採用している。

GRNの実装コードはかなりシンプル

ConvNeXt V2 Blockは、ConvNeXt V1 BlockのLayerScaleをGRNに変更しただけのシンプルな形になっている。

ImageNet Experiments

ConvNeXtとMAEのシナジー

BackboneをConvNeXt V1, V2それぞれに対してBaseとLarge準備し、学習方法は教師あり学習もしくはFCMAEによる自己教師あり学習の2種類で実験をした結果が以下の表である。ConvNeXt V2とFCMAE(GRN)のシナジーによって、精度が大幅に向上していることがわかる。

Model Scaling

更に、Atto ModelからHuge Modelまでのモデルキャパシティに関するスケーリング実験も行っている。Huge Modelまでのスケーリング性能が良いことが伺える。

既存モデルとの比較結果

ViTやSwinTransformerと提案手法を比較した結果。Baseサイズでは、最も性能が良いが、HugeサイズになるとViTの方が性能が良い結果であった。

ImageNet-22K intermediate fine-tuning

BEiTの学習は、PreTraining->ImageNet-22K fine-tuning->ImageNet-1K fine-tuningという3step構造のテクニックを使うことで精度を向上させている。この手法に加え、画像サイズを大きくした場合を本論文では検証している。
https://arxiv.org/pdf/2106.08254.pdf

検証結果から、画像サイズ512×512、intermediate fine-tuningの条件で、state-of-the-artを達成することができた。（ちょっと無理やりな気もするけど・・・）

Transfer Learning Experiments

本手法のConNeXt V2 + FCMAEの転移学習の性能を検証している。
COCO datasetにおけるObeject DetectionとSegmentationの実行結果が以下である。SwinTransformer-Baseのモデルと比べて性能が向上していることがわかる。

ADE20K datasetにおけるSemantic segmentationの実行結果は以下である。（こちらも比較対象はSwinTransformer）Object Detection taskと同様で、本手法によって精度が向上していることがわかる。

まとめと所感

PreTrainingモデルによって、その後のモデルのデザインがある程度決まってしまうことは大きな課題であったと思う。特にViTのアーキテクチャを踏襲したようなMAEは、ConvNet系には直感的に相性が悪そうなイメージがあり、実際に単純に結合するだけでは精度はでないらしい。ただ、そういった状況でも、特徴量がどんな状況になってしまっているかを確認することが非常に大事であるということを、改めて感じた（気がする）。

Discussion