TokenUnify: Scaling Up Autoregressive Pretraining for Neuron Segmentation

📄 arXiv: 2405.16847v2 📥 PDF

作者: Yinda Chen, Haoyuan Shi, Xiaoyu Liu, Te Shi, Ruobing Zhang, Dong Liu, Zhiwei Xiong, Feng Wu

分类: cs.CV, cs.AI

发布日期: 2024-05-27 (更新: 2025-08-25)

备注: Accepted by ICCV 2025

💡 一句话要点

TokenUnify：通过自回归预训练扩展神经元分割能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 神经元分割 自回归预训练 电子显微镜图像 分层预测编码 Mamba架构

📋 核心要点

  1. 电子显微镜图像神经元分割面临高噪声、各向异性及长程依赖等挑战，传统视觉模型难以有效应对。
  2. TokenUnify通过分层预测编码框架，结合随机token、下一个token和所有token预测，捕获多尺度依赖关系。
  3. TokenUnify利用Mamba架构和大规模EM数据集，在神经元分割任务上显著优于现有方法，提升高达44%。

📝 摘要（中文）

从电子显微镜(EM)图像中分割神经元对于理解大脑回路至关重要，但高分辨率EM图像中复杂的神经元结构带来了巨大的挑战。EM数据具有独特的特性，包括高噪声水平、各向异性体素尺寸和超长的空间依赖性，这使得传统的视觉模型显得不足。受语言模型中自回归预训练的启发，我们提出了TokenUnify，一个分层的预测编码框架，通过三个互补的学习目标来捕获多尺度的依赖关系。TokenUnify集成了随机token预测、下一个token预测和下一个所有token预测，以创建一个具有涌现特性的综合表示空间。从信息论的角度来看，这三个任务是互补的，并提供了视觉数据结构的最佳覆盖，我们的方法将长度为K的序列的自回归误差累积从O(K)降低到O(sqrt(K))。我们还引入了一个包含12亿个带注释体素的大规模EM数据集，提供了具有空间连续性的理想长序列视觉数据。利用Mamba架构的线性时间序列建模能力，TokenUnify在下游神经元分割上实现了44%的性能提升，并且优于MAE 25%。我们的方法展示了随着模型尺寸增加的卓越缩放特性，有效地弥合了语言和视觉模型的预训练策略之间的差距。

🔬 方法详解

问题定义：论文旨在解决电子显微镜（EM）图像中神经元分割的难题。现有方法在处理EM数据时面临诸多挑战，包括高噪声、各向异性体素尺寸以及神经元结构中存在的超长空间依赖性。传统视觉模型难以有效地捕捉这些复杂的特征，导致分割精度不高。

核心思路：论文的核心思路是借鉴语言模型中的自回归预训练方法，并将其应用于视觉领域。通过设计一种分层的预测编码框架，即TokenUnify，来学习EM图像中神经元结构的多尺度依赖关系。这种方法旨在创建一个更具表达能力的表示空间，从而提高神经元分割的准确性。

技术框架：TokenUnify的技术框架主要包括三个互补的学习目标：随机token预测、下一个token预测和下一个所有token预测。这些任务共同作用，旨在全面捕捉视觉数据的结构信息。此外，该框架还利用了Mamba架构的线性时间序列建模能力，以高效地处理长序列的EM数据。整体流程包括：首先，使用大规模EM数据集对模型进行预训练；然后，将预训练的模型应用于下游的神经元分割任务。

关键创新：TokenUnify的关键创新在于其分层的预测编码框架，该框架通过结合三种不同的预测任务，实现了对视觉数据结构的最佳覆盖。与传统的自回归方法相比，TokenUnify能够更有效地减少误差累积，从而提高模型的性能。此外，该论文还引入了一个大规模的EM数据集，为长序列视觉数据的研究提供了理想的平台。

关键设计：TokenUnify的关键设计包括：1) 三种互补的预测任务：随机token预测旨在学习局部特征，下一个token预测旨在学习序列依赖关系，下一个所有token预测旨在学习全局结构；2) Mamba架构的应用，利用其线性时间复杂度来处理长序列数据；3) 大规模EM数据集的构建，为模型的训练提供了充足的数据支持；4) 从信息论的角度分析了三种预测任务的互补性，并证明了该方法能够有效地减少误差累积。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TokenUnify在下游神经元分割任务上取得了显著的性能提升，相较于现有方法，性能提升高达44%，并且优于MAE 25%。实验结果表明，随着模型尺寸的增加，TokenUnify的性能能够持续提升，展现出卓越的缩放特性。这些结果验证了该方法在处理大规模视觉数据方面的有效性和潜力。

🎯 应用场景

该研究成果可广泛应用于神经科学领域，助力大脑回路的解析和神经疾病的研究。高精度的神经元分割能够帮助科学家们更深入地理解大脑的结构和功能，为开发新的治疗方法提供理论基础。此外，该方法在其他医学图像分析任务中也具有潜在的应用价值。

📄 摘要（原文）

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.