[論文レビュー] TokenUnify: Scaling Up Autoregressive Pretraining for Neuron Segmentation

Yinda Chen, Haoyuan Shi+6 more|arXiv (Cornell University)|May 27, 2024

Advanced Image and Video Retrieval Techniques被引用数 4

ひとこと要約

TokenUnifyは、視覚自己回帰的モデルにおける誤差蓄積を低減するため、ランダム、次トークン、およびすべての次トークン予測を組み合わせた新しい混合トークン予測事前学習手法を導入する。Mambaアーキテクチャと1億2000万ボクセルにのぼる超高解像度3次元電子顕微鏡データセットを活用することで、従来手法に比べてニューロンセグメンテーションで45%の性能向上を達成するとともに、スケーラブルで効率的な長系列学習を実現する。

ABSTRACT

Neuron segmentation from electron microscopy (EM) volumes is crucial for understanding brain circuits, yet the complex neuronal structures in high-resolution EM images present significant challenges. EM data exhibits unique characteristics including high noise levels, anisotropic voxel dimensions, and ultra-long spatial dependencies that make traditional vision models inadequate. Inspired by autoregressive pretraining in language models, we propose TokenUnify, a hierarchical predictive coding framework that captures multi-scale dependencies through three complementary learning objectives. TokenUnify integrates random token prediction, next-token prediction, and next-all token prediction to create a comprehensive representational space with emergent properties. From an information-theoretic perspective, these three tasks are complementary and provide optimal coverage of visual data structure, with our approach reducing autoregressive error accumulation from O(K) to O(sqrt(K)) for sequences of length K. We also introduce a large-scale EM dataset with 1.2 billion annotated voxels, offering ideal long-sequence visual data with spatial continuity. Leveraging the Mamba architecture's linear-time sequence modeling capabilities, TokenUnify achieves a 44% performance improvement on downstream neuron segmentation and outperforms MAE by 25%. Our approach demonstrates superior scaling properties as model size increases, effectively bridging the gap between pretraining strategies for language and vision models.

研究の動機と目的

• 視覚的自己回帰的事前学習における誤差蓄積を解消することで、視覚モデルにおけるスケーリング則の制限を克服すること。
• 特に高解像度3次元データにおいて、長系列画像モデリングにおける計算の非効率性を克服すること。
• 画像内に空間的に相関のある長系列を構築し、効果的な自己回帰的学習を可能にすること。
• 言語と視覚の事前学習のギャップを埋め、視覚分野におけるスケーラブルで自己回帰的な戦略を可能にすること。
• 大規模かつ高解像度のEMデータセットを用いて、ニューロンセグメンテーションの統一ベンチマークを確立すること。

提案手法

• TokenUnifyは、グローバルな文脈モデリングを向上させるために、ランダムトークン予測、次トークン予測、およびすべての次トークン予測の3つの事前学習タスクを統合する。
• 自己回帰的計算を二次関数的から線形的複雑度に削減するMambaアーキテクチャを採用し、長系列処理の効率性を実現する。
• 超高解像度電子顕微鏡画像から非オーバーラップの3次元パッチを抽出し、順序付き表現にトークン化する。
• 6つの機能的マウス脳領域にまたがる1億2000万以上のアノテート済みボクセルを含む大規模3次元EMデータセットを構築し、長系列学習を支援する。
• 3つの予測タスクを統合した統一目的関数を用いて、エンドツーエンドで事前学習フレームワークを訓練する。
• 学習可能なマスクトークンとクロスアテンションデコーダーを用いて、Mambaベースのバックボーン（EMmamba）を微調整し、下流のニューロンセグメンテーションタスクに適用する。

Figure 1 : TokenUnify prediction paradigms divide the 3D EM image into non-overlapping patches, which are tokenized into a sequence. Three tasks are performed for rich 3D image representations: (a) random token prediction, (b) next token prediction, and (c) next-all token prediction.

実験結果

リサーチクエスチョン

• RQ1 混合トークン予測戦略は、視覚的自己回帰的事前学習における誤差蓄積を低減できるか？
• RQ2 複数の予測目的を組み合わせることで、視覚モデルにおける表現学習がどのように向上するか？
• RQ3 Mambaアーキテクチャは、線形的複雑度で自己回帰的視覚モデルを効果的にスケーリングできるか？
• RQ4 大規模で高解像度かつ空間的に連続的な3次元EMデータセットは、ニューロンセグメンテーションのためのより優れた長系列モデリングを可能にするか？
• RQ5 TokenUnifyは、マスク自己符号化法や標準的な自己回帰的手法に比べて、下流のセグメンテーションタスクでどの程度優れているか？

主な発見

• TokenUnifyは、既存手法と比較して、下流のEMニューロンセグメンテーションタスクで45%の相対的性能向上を達成した。
• パrameter数が少ないにもかかわらず、MAEよりも21%高い事前学習性能を達成した。
• アブレーションスタディの結果、ランダム、次、すべての次トークン予測の3つのタスクを統合した場合が最良の結果を示し、VOIとARANDスコアはそれぞれ0.9951および0.0509であった。
• MambaベースのEMmambaモデルは、TokenUnify事前学習戦略を用いた場合、Wafer4データセットでSOTA結果を達成し、VOIは0.2124 ± 0.0172、ARANDは0.0967 ± 0.0097であった。
• Mambaアーキテクチャは、計算複雑度を二次関数的（Transformer）から線形に削減することで、スケーラビリティと効率性に優れ、長系列での学習を可能にした。
• 1億2000万ボクセルにのぼる大規模EMデータセットは、今日までで最大の統一ベンチマークをニューロンセグメンテーションに提供している。

Figure 2 : illustrates the main pretraining workflow of TokenUnify. The image $\mathbf{X}$ is fed into the Tokenizer, transforming it into a long sequence of tokens $\mathbf{x}_{i}|_{i=1}^{K}$ . The predictions for the random token, next token, and next-all token are performed sequentially. The Perc