[論文レビュー] ViT-AE++: Improving Vision Transformer Autoencoder for Self-supervised Medical Image Representations

Chinmay Prabhakar, Hongwei Li+4 more|arXiv (Cornell University)|Jan 18, 2023

AI in cancer detection被引用数 10

ひとこと要約

ViT-AE++ は補助的再構成損失とコントラスト損失を用いてVision Transformer Autoencoderを強化し、2Dおよび3D医用画像の自己教師付き表現をより良く学習し、3D体積にも拡張します。

ABSTRACT

Self-supervised learning has attracted increasing attention as it learns data-driven representation from data without annotations. Vision transformer-based autoencoder (ViT-AE) by He et al. (2021) is a recent self-supervised learning technique that employs a patch-masking strategy to learn a meaningful latent space. In this paper, we focus on improving ViT-AE (nicknamed ViT-AE++) for a more effective representation of 2D and 3D medical images. We propose two new loss functions to enhance the representation during training. The first loss term aims to improve self-reconstruction by considering the structured dependencies and indirectly improving the representation. The second loss term leverages contrastive loss to optimize the representation from two randomly masked views directly. We extended ViT-AE++ to a 3D fashion for volumetric medical images as an independent contribution. We extensively evaluate ViT-AE++ on both natural images and medical images, demonstrating consistent improvement over vanilla ViT-AE and its superiority over other contrastive learning approaches. Codes are here: https://github.com/chinmay5/vit_ae_plus_plus.git.

研究の動機と目的

• 医療画像におけるラベルが乏しい場合の自己教師付き表現学習の改善を動機付ける。
• 構造情報とエッジ情報を捉える補助再構成損失を導入する。
• オートエンコーダベースの学習とコントラスト学習のパラダイムを組み合わせ、表現を直接最適化する。
• ViT-AE++ を3D体積データに拡張し、MRI/CT様モダリティを扱う。
• ViT-AEと比較して一貫した性能向上を示し、コントラスト法と競合する結果を示す。

提案手法

• パッチをランダムにマスクし、欠損パッチには MASK トークンを用いたViTベースのオートエンコーダを使用する。
• 補助損失を導入する：L_per（VGGベースの知覚損失）と L_edge（3D Sobel エッジ損失）で構造情報を捕捉する。
• 2つのランダムにマスクされたビュー間でコントラスト損失 L_CL を導入し、表現を直接最適化する。
• 全体目的関数へ損失を組み合わせる：L = L_rec + λ1 L_per + λ2 L_edge + L_CL ただし λ2 はスケジュール付きに減衰させる。
• 3Dパッチ、3D位置エンコーディング、3D Transformerデコーダを用いてフレームワークを3Dへ拡張する。
• 線形プロービングで複数のデータセットを評価し、下流タスク指標としてAUC/accuracyを報告する。

実験結果

リサーチクエスチョン

• RQ1 補助的な構造認識再構成はピクセル単位の損失を超えて意味内容を捉えるのを改善するか？
• RQ2 オートエンコード再構成と併用したコントラスト目的が医用画像の表現をより良くするか？
• RQ3 ViT-AE++ の3D適応は体積医療データで改善を維持するか？
• RQ4 マスキング比率とエッジ/知覚損失の重みは学習の安定性と表現品質にどのような影響を与えるか？

主な発見

Method	CIFAR-10	CIFAR-100	Tiny ImageNet-100	Chest X-ray
ViT-AE	95.40	78.82	72.09	95.60

• ViT-AE++ は2Dデータセットで一貫して vanilla ViT-AE を上回る（CIFAR-10: 95.40 vs 94.10; CIFAR-100: 78.82 vs 75.61; Tiny ImageNet-100: 72.09 vs 70.42; Chest X-ray: 95.60 vs 95.20）。
• 3Dデータセットでは、ViT-AE++ は BraTS でコントラスト法に対して競合〜優れた結果を達成し、EGD ではそれを上回る（BraTS: 0.767 AUC; EGD: 0.846 AUC）。
• アブレーション研究は、補助エッジ損失と知覚損失とコントラスト損失が ViT-AE を上回る改善を共同で推進することを示す。
• マスク比率 p とエッジ損失重み λ2 の減衰は性能と学習安定性にとって重要で、最適な p は約 0.75、λ2 の線形減衰が学習を安定化する。
• 3D ViT-AE++ は医用データセットで線形プロービングによる特徴抽出機として強い下流性能を示し、最新手法と競合性を保つ。