# TokenUnify: Scaling Up Autoregressive Pretraining for Neuron Segmentation

arXiv: https://arxiv.org/abs/2405.16847 · HTML: https://mlsanity.com/paper/2405.16847/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 66 references that are in the map.
- First posted on arXiv 27 May 2024.

## Builds on (most cited first, 30 of 36)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Flamingo: a Visual Language Model for Few-Shot Learning](https://mlsanity.com/paper/2204.14198/) (2022)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation](https://mlsanity.com/paper/2401.13560/) (2024)
- [MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis](https://mlsanity.com/paper/2211.09117/) (2022)
- [Sequential Modeling Enables Scalable Learning for Large Vision Models](https://mlsanity.com/paper/2312.00785/) (2023)
- [The pitfalls of next-token prediction](https://mlsanity.com/paper/2403.06963/) (2024)
- [Scalable Pre-training of Large Autoregressive Image Models](https://mlsanity.com/paper/2401.08541/) (2024)
- [The effectiveness of MAE pre-pretraining for billion-scale pretraining](https://mlsanity.com/paper/2303.13496/) (2023)
- [Superhuman Accuracy on the SNEMI3D Connectomics Challenge](https://mlsanity.com/paper/1706.00120/) (2017)
- [Deception Detection from Linguistic and Physiological Data Streams Using Bimodal Convolutional Neural Networks](https://mlsanity.com/paper/2311.10944/) (2023)
- [Contextual Hourglass Network for Semantic Segmentation of High Resolution Aerial Imagery](https://mlsanity.com/paper/1810.12813/) (2018)
- [BIMCV-R: A Landmark Dataset for 3D CT Text-Image Retrieval](https://mlsanity.com/paper/2403.15992/) (2024)
- [Bridging the Domain Gap: Self-Supervised 3D Scene Understanding with Foundation Models](https://mlsanity.com/paper/2305.08776/) (2023)
- [Automated Data Curation for Robust Language Model Fine-Tuning](https://mlsanity.com/paper/2403.12776/) (2024)
- [Multimodal Semi-Supervised Learning for 3D Objects](https://mlsanity.com/paper/2110.11601/) (2021)
- [Real-Time Pill Identification for the Visually Impaired Using Deep Learning](https://mlsanity.com/paper/2405.05983/) (2024)
- [The physics of brain network structure, function, and control](https://mlsanity.com/paper/1809.06441/) (2018)
- [Self-Supervised Neuron Segmentation with Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2310.04148/) (2023)
- [Learning Multiscale Consistency for Self-supervised Electron Microscopy Instance Segmentation](https://mlsanity.com/paper/2308.09917/) (2023)

## Built on by (most cited first, 7 of 7)

- [Mamba in Vision: A Comprehensive Survey of Techniques and Applications](https://mlsanity.com/paper/2410.03105/) (2024)
- [MSP-MVS: Multi-Granularity Segmentation Prior Guided Multi-View Stereo](https://mlsanity.com/paper/2407.19323/) (2024)
- [DVP-MVS: Synergize Depth-Edge and Visibility Prior for Multi-View Stereo](https://mlsanity.com/paper/2412.11578/) (2024)
- [CoC-VLA: Delving into Adversarial Domain Transfer for Explainable Autonomous Driving via Chain-of-Causality Visual-Language-Action Model](https://mlsanity.com/paper/2511.19914/) (2025)
- [FGNet: Leveraging Feature-Guided Attention to Refine SAM2 for 3D EM Neuron Segmentation](https://mlsanity.com/paper/2511.13063/) (2025)
- [STeInFormer: Spatial-Temporal Interaction Transformer Architecture for Remote Sensing Change Detection](https://mlsanity.com/paper/2412.17247/) (2024)
- [Separators in Enhancing Autoregressive Pretraining for Vision Mamba](https://mlsanity.com/paper/2603.03806/) (2026)

## Abstract (arXiv)

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.