TokenUnify: Scalable Autoregressive Visual Pre-training with Mixture Token Prediction

自回归下一个标记预测是大规模语言模型的标准预训练方法，但由于图像数据的非顺序性，导致累积误差，从而阻碍其应用于视觉任务。大多数视觉模型采用基于掩码自编码器（MAE）的预训练方法，但面临可扩展性问题。为了解决这些挑战，我们引入了一种新的预训练方法TokenUnify，它集成了随机标记预测、下一个标记预测和下一个所有标记预测。我们提供了理论证据，证明TokenUnify缓解了视觉自回归中的累积误差。与TokenUnify合作，我们组装了一个超高分辨率的大规模电子显微镜（EM）图像数据集，非常适合创建空间相关的长序列。该数据集包括超过1.2亿个带注释的体素，是迄今为止最大的神经元分割数据集，并为实验验证提供了统一的基准。在这个数据集上利用天蛇网络进行长序列建模，TokenUnify不仅减少了计算复杂度，而且在下游EM神经元分割任务上相比现有方法提高了45\%的分割性能。此外，TokenUnify展示了优于MAE和传统自回归方法的可扩展性，有效地弥合了语言和视觉模型预训练策略之间的差距。代码可在\url{https://github.com/ydchen0806/TokenUnify}上获得。

本论文旨在解决视觉任务中基于自回归模型的预训练方法应用的问题，即由于图像数据的非顺序性导致累积误差，而基于掩码自编码器的预训练方法又存在可扩展性问题。

本文提出了一种新的预训练方法TokenUnify，它结合了随机token预测、下一个token预测和下一个所有token预测，以缓解视觉自回归中的累积误差。同时，TokenUnify还展示了优于MAE和传统自回归方法的可扩展性，有效地弥合了语言和视觉模型预训练策略之间的差距。

本文提供了理论证据，证明了TokenUnify缓解了视觉自回归中的累积误差。同时，作者还构建了一个大规模的电子显微镜图像数据集，并使用Mamba网络进行了实验验证。实验结果表明，TokenUnify不仅降低了计算复杂度，还显著提高了下游任务（EM神经元分割）的分割性能。此外，作者还提供了开源代码。

在最近的相关研究中，还有一些基于自回归模型的预训练方法，例如BERT、GPT等。此外，还有一些基于掩码自编码器的预训练方法，例如Masked Convolutional Autoencoder（MCAE）和Masked Spatial-Channel-Wise Autoencoder（MSCAE）等。