VariViT: 可変画像サイズに対応するビジョン・トランスフォーマー

技術的サマリー（日本語）

本論文は、医療画像解析におけるビジョン・トランスフォーマー（ViT）の課題を解決し、可変サイズの画像入力を効率的に処理する新しいアーキテクチャ「VariViT」を提案しています。以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

1. 背景と課題

従来の Vision Transformer (ViT) は、画像を固定サイズのパッチに分割し、自己注意メカニズムを用いて長距離依存関係を学習する優れたアーキテクチャです。しかし、医療画像（特に腫瘍など）の解析においては以下の重大な課題が存在します。

• 固定サイズ入力による情報の劣化: 医療画像は腫瘍の形状やサイズが患者間で大きく異なります。ViT に投入するためには、画像をリサイズ、パディング、またはクロップする必要があります。
  • リサイズ: 画像を固定サイズにリサイズすると、人工的なアーティファクトが発生したり、重要な情報が失われたりします。
  • 固定クロップ: 腫瘍の大きさに関わらず同じサイズのバウンディングボックスを切り出すと、腫瘍が小さい場合は背景（正常組織）の比率が高くなり、計算リソースの無駄や、モデルが背景に注意を向けてしまう原因となります（図 1 参照）。
• 計算コストと精度のトレードオフ: 大きな画像を処理すると計算コストが膨大になり、小さな画像にリサイズすると情報が失われるというジレンマがあります。

2. 提案手法：VariViT

VariViT は、ViT の利点を維持しつつ、可変サイズの画像入力を直接処理できるように設計されたモデルです。主な技術的革新は以下の 2 点です。

A. 中心選択（Center and Select）による位置埋め込みの再調整

従来の ViT や FlexiViT は、異なるサイズへの対応のために位置埋め込み（Positional Embedding）を線形補間（Interpolation）で調整しますが、これには計算コストの増大や情報損失のリスクがあります。

VariViT は、腫瘍の中心座標が 3D クロップの中心に整合されているという特性を利用し、以下の「Center and Select」戦略を採用します。

1. 最大サイズの埋め込み初期化: 最も大きな画像サイズに対応する固定の正弦波（sinusoidal）位置埋め込みを 3D グリッドとして初期化します。
2. 中心からの選択: 入力画像のサイズが小さい場合、補間を行わず、初期化されたグリッドの「中心」から、現在の画像サイズに相当する部分のサブセットを直接抽出（Select）します。
3. 効果: これにより、補間による近似誤差を避け、低計算コストで元の位置情報を保持したまま、可変パッチ数（シーケンス長）に対応できます。

B. 効率的なバッチング戦略

異なるサイズの画像を同一バッチで処理する際の計算効率を向上させるため、2 つの戦略を提案しています。

1. カスタムバッチサンプリング (CBS): 画像サイズが同一のサンプルを同一バッチにグループ化します。これにより、バッチ内の計算を最適化し、トレーニング時間を短縮します。
2. 勾配蓄積 (Gradient Accumulation, GA): バッチ内で画像サイズが異なる場合でも、ミニバッチごとに勾配を計算し、指定された回数（更新間隔）後に重みを更新します。これにより、異なるサイズの画像を混在させた学習が可能になり、メモリ効率と計算速度を両立します。

3. 主要な貢献

1. 柔軟な位置埋め込み戦略: 画像サイズに応じて動的に調整可能で、補間を必要としない「Center and Select」手法の提案。
2. 計算効率化のバッチング: 可変サイズ画像を効率的に処理するための新しいバッチング戦略（CBS と GA）の実装。
3. 医療タスクでの実証: 2 つの 3D 脳 MRI データセットを用いた検証により、グリオーマの遺伝子型予測（IDH 変異状態）および脳腫瘍の分類（原発性 vs 転移性）において、従来の ViT や ResNet を上回る性能を示しました。

4. 実験結果

2 つの主要なタスクにおいて、VariViT はベースラインモデル（3D ResNet-18, 3D Vanilla ViT, Pix2Struct）と比較して優れた性能を発揮しました。

• IDH 変異状態予測（グリオーマデータセット）:

  • F1 スコア: VariViT-GA は 75.5% を達成（Vanilla ViT: 74.4%, ResNet-18: 71.6%）。
  • AUC: 0.942（最良）。
  • トレーニング時間: 従来のアーキテクチャと比較して最大 30% 短縮されました。
  • t-SNE 可視化により、VariViT がより明確なクラスター分離を実現していることが確認されました。

• 脳腫瘍分類（原発性 vs 転移性）:

  • F1 スコア: VariViT-GA は 76.3% を達成。
  • MCC (Matthews Correlation Coefficient): 0.706 と、他のモデルを上回る性能を示しました。
  • VariViT-CBS も Pix2Struct よりも優れた性能を示し、トレーニング時間が短縮されました。

• アブレーション研究:

  • 位置埋め込みの手法比較において、「Center and Select」が補間ベースの手法（Interp Fixed/Learned）や独立した固定埋め込みよりも優れた結果をもたらしました。
  • 多クラス分類タスク（グリオーマの 3 亜型分類）においても、同様の性能と高速な学習時間を維持しました。

5. 意義と結論

VariViT は、医療画像の「不規則な形状」と「可変なサイズ」という特性に特化して設計された Vision Transformer です。

• 情報損失の回避: 画像を無理にリサイズせず、関心領域（ROI）に焦点を当てたクロップを直接処理することで、診断に重要な微細な特徴を保持します。
• 計算効率: 補間やパディングに依存しない位置埋め込みの調整と、効率的なバッチングにより、トレーニングおよび推論の高速化を実現しました。
• 汎用性: 脳 MRI での検証を超え、他のモダリティや解剖学的領域への適用も可能であり、大規模データセットや高解像度画像における ViT の実用性を高める重要なステップとなります。

本論文は、医療画像解析における ViT の実用障壁を下げ、より精度が高く効率的な診断支援システムの構築に寄与する可能性を示しています。