[Paper Review] 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 Techniques4 citations

TL;DR

TokenUnify introduces a novel mixed token prediction pretraining method that combines random, next, and next-all token prediction to reduce error accumulation in visual autoregressive modeling. By leveraging the Mamba architecture and a new ultra-high-resolution 3D electron microscopy dataset with 120 million annotated voxels, it achieves a 45% performance gain in neuron segmentation over prior methods while enabling scalable, efficient long-sequence learning.

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.

Motivation & Objective

• To address cumulative error in visual autoregressive pretraining, which limits scaling laws in vision models.
• To overcome computational inefficiencies in long-sequence image modeling, especially for high-resolution 3D data.
• To construct spatially correlated long sequences in images to enable effective autoregressive learning.
• To bridge the gap between language and vision pretraining by enabling scalable, autoregressive strategies in vision.
• To establish a unified benchmark for neuron segmentation using a large-scale, high-resolution EM dataset.

Proposed method

• TokenUnify integrates three pretraining tasks: random token prediction, next-token prediction, and next-all token prediction to enhance global context modeling.
• It employs the Mamba architecture, which reduces autoregressive computation from quadratic to linear complexity, enabling efficient long-sequence processing.
• The method uses non-overlapping 3D patches from ultra-high-resolution electron microscopy images, tokenized into sequential representations.
• A large-scale 3D EM dataset with over 120 million annotated voxels across six functional mouse brain regions was constructed to support long-sequence learning.
• The pretraining framework is trained end-to-end with a unified objective combining all three prediction tasks.
• The Mamba-based backbone (EMmamba) is fine-tuned for downstream neuron segmentation with a learnable mask token and cross-attention decoder.

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.

Experimental results

Research questions

• RQ1 Can a mixed token prediction strategy reduce cumulative error in visual autoregressive pretraining?
• RQ2 How does combining multiple prediction objectives improve representation learning in vision models?
• RQ3 Can the Mamba architecture effectively scale autoregressive vision models with linear complexity?
• RQ4 Does a large-scale, high-resolution, spatially continuous 3D EM dataset enable better long-sequence modeling for neuron segmentation?
• RQ5 To what extent does TokenUnify outperform masked autoencoding and standard autoregressive methods in downstream segmentation tasks?

Key findings

• TokenUnify achieved a 45% relative improvement in segmentation performance on downstream EM neuron segmentation tasks compared to existing methods.
• The method outperformed MAE by 21% in pretraining performance despite using fewer parameters.
• The ablation study confirmed that combining all three prediction tasks (random, next, next-all) yields the best results, with VOI and ARAND scores of 0.9951 and 0.0509, respectively.
• The Mamba-based EMmamba model achieved state-of-the-art results on the Wafer4 dataset, with a VOI of 0.2124 ± 0.0172 and ARAND of 0.0967 ± 0.0097 under the TokenUnify pretraining strategy.
• The Mamba architecture demonstrated superior scalability and efficiency, reducing computational complexity from quadratic (Transformer) to linear, enabling training on long sequences.
• The large-scale EM dataset with 120 million annotated voxels provides the largest unified benchmark for neuron segmentation to date.

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