[论文解读] 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 架构和一个包含 12000 万个标注体素的超高分辨率 3D 电子显微镜数据集，其在神经元分割任务上相比先前方法实现了 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.

研究动机与目标

• 为解决视觉自回归预训练中的误差累积问题，该问题限制了视觉模型的扩展规律。
• 克服长序列图像建模中的计算低效问题，特别是针对高分辨率 3D 数据。
• 在图像中构建空间相关的长序列，以支持有效的自回归学习。
• 通过在视觉领域实现可扩展的自回归策略，弥合语言与视觉预训练之间的差距。
• 基于大规模、高分辨率的电子显微镜数据集，建立统一的神经元分割基准。

提出的方法

• TokenUnify 集成了三种预训练任务：随机标记预测、下一个标记预测和全部下一个标记预测，以增强全局上下文建模能力。
• 采用 Mamba 架构，将自回归计算复杂度从二次方降低至线性，从而实现高效的长序列处理。
• 使用来自超高分辨率电子显微镜图像的非重叠 3D 小块，并将其转换为序列化表示。
• 构建了一个大规模 3D 电子显微镜数据集，涵盖六个功能小鼠脑区，包含超过 12000 万个标注体素，以支持长序列学习。
• 预训练框架采用统一目标端到端训练，结合所有三种预测任务。
• 基于 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 大规模、高分辨率、空间连续的 3D 电子显微镜数据集是否能提升神经元分割的长序列建模能力？
• RQ5 在下游分割任务中，TokenUnify 相较于掩码自编码和标准自回归方法，性能提升程度如何？

主要发现

• 与现有方法相比，TokenUnify 在下游电子显微镜神经元分割任务中实现了 45% 的相对性能提升。
• 尽管参数量更少，该方法在预训练性能上仍比 MAE 提高了 21%。
• 消融实验证实，结合全部三种预测任务（随机、下一个、全部下一个）可获得最佳结果，VOI 和 ARAND 分数分别为 0.9951 和 0.0509。
• 基于 Mamba 的 EMmamba 模型在 Wafer4 数据集上达到最先进性能，VOI 为 0.2124 ± 0.0172，ARAND 为 0.0967 ± 0.0097，采用 TokenUnify 预训练策略。
• Mamba 架构展现出卓越的可扩展性和效率，将计算复杂度从二次方（Transformer）降低至线性，支持长序列训练。
• 包含 12000 万个标注体素的大规模电子显微镜数据集，为神经元分割提供了迄今为止最大的统一基准。

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