UniWeTok:面向统一多模态大语言模型的超大规模二进制统一Tokenizer UniWeTok: An Unified Binary Tokenizer with Codebook Size 2^{128} for Unified Multimodal Large Language Model
提出2^128码本的统一离散tokenizer,在理解、生成、编辑任务上均达到SOTA
前置知识
VQ-VAE / VQGAN
向量量化变分自编码器,通过编码器将图像压缩为离散token序列,再用解码器重建图像。VQGAN在此基础上加入感知损失和GAN损失提升视觉质量。核心思想是维护一个码本(codebook),将连续的编码向量映射到最近的码本向量,实现离散化。这种离散表示便于后续的自回归建模,但传统VQ方法受限于码本大小和利用率问题。
UniWeTok是在VQGAN范式上的重大改进,理解传统VQ方法的局限性是理解本文动机的关键。
Lookup-Free Quantization (LFQ) / Binary Quantization
传统VQ需要计算编码向量与所有码本向量的距离,计算成本随码本大小线性增长。LFQ通过将量化过程设计为无需查找操作的形式来解决这个问题。BSQ(二进制量化)进一步将每个维度量化为+1或-1,码本大小为2^d(d为维度)。这样量化过程简化为简单的sign操作,极大降低了计算成本。WeTok在此基础上提出Group-Wise LFQ,将特征分组后分别量化。
UniWeTok使用2^128的超大码本,理解二进制量化的工作原理对于理解这种大规模码本的可行性和优势至关重要。
语义蒸馏 (Semantic Distillation)
在tokenizer训练中引入预训练的语义编码器(如SigLIP)作为教师模型,通过余弦相似度损失将tokenizer编码器的特征与教师模型的语义特征对齐。目的是让原本只关注像素重建的tokenizer获得语义理解能力。这种对齐可以在编码器输出的不同阶段进行:量化前(Pre-distillation)和量化后(Post-distillation)。
UniWeTok的核心创新之一是Pre-Post Distillation策略,理解语义蒸馏的基本原理是理解该策略有效性的前提。
统一多模态大语言模型 (Unified MLLM)
同时具备视觉理解和视觉生成能力的多模态大语言模型。与传统MLLM只做理解不同,统一MLLM需要在同一个框架内处理图像理解、文本到图像生成、图像编辑等多种任务。挑战在于理解和生成对视觉表示的要求存在冲突:理解需要高层语义信息,生成需要低层细节和生成先验。
UniWeTok的目标就是作为统一MLLM的视觉tokenizer,解决理解和生成之间的矛盾。
研究动机
统一多模态大语言模型(Unified MLLM)需要一种视觉表示能够同时支持高保真重建、复杂语义提取和生成建模。然而现有视觉tokenizer难以在单一框架内满足这些相互矛盾的需求。具体来说,传统VQGAN方法受限于码本大小(通常16384),重建质量有限,rFID在2-5之间。而基于连续表示的方法(如VAE的连续潜变量)在自回归生成时容易出现错误累积和模式坍塌。虽然近期BSQ、WeTok等方法将码本扩展到2^128以上,显著提升了重建能力(rFID降至1以下),但这些大规模码本tokenizer面临两个关键问题:一是语义提取能力极弱,几乎无法提取任何语义信息(零样本分类准确率接近0%),无法用于多模态理解任务;二是超大码本增加了下游生成任务的难度,现有方法如Infinity和BitDance在文本到图像建模时仍将码本限制在2^32。
本文的目标是本文的核心目标是构建一个统一的离散视觉tokenizer,使其能够:(1) 在32倍空间下采样率下实现高保真图像重建,rFID低于1.0;(2) 具备强大的语义提取能力,零样本分类准确率超过50%;(3) 生成的离散token适合下游自回归生成任务,在ImageNet类条件生成上达到SOTA(FID低于1.5);(4) 训练成本显著低于现有方法,训练token数量控制在33B以内。最终目标是用这个tokenizer构建统一MLLM,在理解、生成、编辑三大任务上均展现竞争力。
与已有工作不同的是,本文的独特切入点在于发现并解决了三个被忽视的关键问题。首先,现有大规模码本tokenizer的语义缺失问题根源在于commitment loss将编码器输出强制锚定到-1或+1附近,与token entropy loss的优化目标产生根本冲突。其次,作者发现可以通过设计激活函数(SigLu)将特征空间约束在[-1,1]区间,使token entropy loss与commitment loss等价,从而消除冲突。第三,作者发现生成感知先验(Generative-Aware Prior)不仅提升生成能力,反而还能增强理解能力,这与直觉相反但实验验证了其有效性。这种对优化冲突的深入分析和巧妙解决是本文区别于WeTok、BSQ等前作的核心差异。
核心方法
UniWeTok的方法可以类比为给一个高速打印机装上智能芯片:原有的打印机(WeTok基础架构)速度很快(32倍压缩,64个token表示一张图),但只会'复印'像素,不懂内容含义。UniWeTok通过三个维度的改造让它既会复印又会理解:第一,安装'语义芯片'(Pre-Post Distillation),让编码器学会提取图像的语义信息;第二,安装'创意芯片'(Generative-Aware Prior),让token空间具备生成建模的先验知识;第三,重新设计'电路板'(SigLu激活 + 混合架构),解决新芯片安装后的信号冲突问题。技术路线上,UniWeTok基于WeTok的Group-Wise Lookup-Free Quantization框架,将图像通过编码器压缩为连续特征,经Group量化得到二进制token(每个token为128位二进制码),再通过解码器重建图像。关键创新是在这个框架中嵌入语义蒸馏和生成先验两个辅助目标,通过精心设计的激活函数和训练策略实现多目标和谐优化。
UniWeTok的核心创新是发现了大规模二进制码本tokenizer中语义提取失败的根本原因,并提出了优雅的解决方案。在传统设置中,commitment loss $\alpha \|U_G - \text{sg}[Q]\|^2$ 强制编码器输出 $U_G$ 靠近量化后的二进制码 $Q \in \{-1, +1\}^{d'}$,这将特征空间锁定在超立方体的顶点附近。而token entropy loss $\mathcal{L}_{\text{Token Entropy}} = \frac{1}{hw} \sum_{i,j,k} H(q_G(c_k | U_G[i,j,k]))$ 则驱使 $U_G$ 趋向正负无穷以最大化熵。这两个目标产生根本冲突:一个要把特征拉向有限区间,另一个要把特征推向无穷。在这种冲突下,任何语义蒸馏都无法稳定收敛。UniWeTok提出的SigLu激活函数 $\text{SigLu}(x) = \frac{1-e^x}{1+e^x}$ 将编码器输出天然约束在[-1,1]区间。在这一约束下,token entropy loss等价于commitment loss,因此可以移除commitment loss(设$\alpha=0$),仅保留token entropy loss。这一设计不仅解决了优化冲突,还使得语义蒸馏(无论是量化前还是量化后)都能稳定训练。
方法步骤详情
UniWeTok的完整流程包括三个阶段的训练。第一阶段(大规模预训练):在256×256固定分辨率下,使用DataComp-1B数据集进行预训练。图像经过编码器(由残差块和Transformer块组成的混合架构)提取特征,得到连续表示 $U = E(I) \in \mathbb{R}^{h \times w \times d}$。沿通道维度将 $U$ 分组为 $U_G \in \mathbb{R}^{h \times w \times g \times d'}$,通过sign操作量化得到二进制码 $Q$。训练损失包含五部分:重建损失、感知损失、GAN损失、token entropy损失,加上语义蒸馏损失(Pre-Post Distillation)和生成感知先验损失(Generative-Aware Prior)。第二阶段(多分辨率继续预训练):同时在多种分辨率(256×256到1024×1024)下训练,使tokenizer适应不同分辨率输入。第三阶段(文本-人脸退火):在人脸和文本等感知敏感数据上进行退火训练,提升对这些场景的处理能力。在推理阶段,只需encoder和decoder,不需要语义教师和生成先验模型。
技术新颖性
UniWeTok的技术新颖性体现在多个层面。首先,SigLu激活函数是一个简单但深刻的创新:它不仅解决了优化冲突,还通过将特征空间约束到[-1,1]区间,使得token entropy loss和commitment loss统一,这是首次在大规模二进制tokenizer中实现这种等价关系。其次,Pre-Post Distillation策略不同于传统的单一蒸馏:同时对量化前特征 $U_G$ 和量化后特征 $U_Q$ 进行语义对齐,确保离散token本身也携带语义信息,这对下游理解任务至关重要。第三,Generative-Aware Prior的设计巧妙地将一个轻量级生成模型(仅8.6M参数的BitDance-T)嵌入tokenizer训练,让tokenizer在训练阶段就能感知生成目标。实验表明这个先验不仅提升生成能力,还意外地增强了理解能力。最后,三阶段课程学习策略(固定分辨率→多分辨率→敏感场景退火)系统性地解决了tokenizer的分辨率泛化和场景适应问题。
实验结果
UniWeTok在多个维度上取得了突破性成果。在ImageNet类条件生成任务上,UniWeTok-H达到1.38的FID,超越了REPA的1.42,同时训练token数量仅为33B(REPA需要262B),实现了近8倍的训练效率提升。更关键的是,UniWeTok只需生成64个token就能表示一张256×256的图像,相比主流tokenizer的256个token减少了75%的推理token数量。在重建质量上,UniWeTok在ImageNet 50K验证集上达到0.79的rFID、23.26dB的PSNR,且码本利用率达到100%。在通用域评估中,UniWeTok在MS-COCO上的rFID为6.18,ImageNet上为1.18(256×256分辨率下),均优于大多数现有tokenizer。在多模态理解方面,UniWeTok-Chat基于Qwen3-8B在SEEDB上达到69.3%、POPE上85.6%、VQAv2上75.8%,在SQA上达到80.3%,表现具有竞争力。在生成质量上,UniWeTok-Gen在DPG-Bench上总分86.63,超越FLUX.1 Dev的83.84和Janus-Pro的84.19。在图像编辑任务上,UniWeTok-Edit的GEdit Overall Score为5.09(英文),首次以自回归模型身份超越了扩散模型OmniGen(5.06)。消融实验显示:Pre-Post Distillation将零样本Top-1准确率从0%提升到51.32%;SigLu激活使得后蒸馏有效(从0.10%到41.51%);Generative-Aware Prior将gFID从2.66改善到2.38;混合架构相比纯CNN或纯Transformer在重建和语义上都更优。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet类条件生成 | FID↓ | 1.38 (UniWeTok-H) | REPA: 1.42 | 2.8% (且训练token从262B降至33B) |
| ImageNet图像重建 | rFID↓ | 0.79 | FlowMo-Lo: 0.95 | 16.8% |
| DPG-Bench文本到图像生成 | Overall Score↑ | 86.63 | FLUX.1 Dev: 83.84 | 3.3% |
| GEdit-Bench图像编辑 | G_O (Overall)↑ | 5.09 | OmniGen: 5.06 | 0.6% |
| SEEDB多模态理解 | Score↑ | 69.3 | Emu3-Chat: 68.2 | 1.6% |
| SQA科学问答 | Score↑ | 80.3 | Emu3-Chat: 89.2 | -9.9% (差距较大) |
局限与改进
尽管UniWeTok取得了显著成果,但仍存在几个值得关注的局限性。首先,在科学问答(SQA)任务上,UniWeTok-Chat的80.3%得分明显低于Emu3-Chat的89.2%,说明在需要复杂推理的场景下,统一tokenizer的理解能力仍有提升空间。其次,图像编辑任务虽然总体分数超越OmniGen,但与商业模型GPT-4o(7.53)和Gemini 2.0(6.32)相比仍有较大差距,且与开源扩散模型Step1X-Edit(6.70)和BAGEL(6.52)相比也有明显差距,说明自回归编辑范式的能力上限可能受限。第三,论文主要在ImageNet和DataComp-1B上训练和评估,对于更复杂的真实世界场景(如高分辨率遥感图像、医学图像等)的泛化能力未被验证。第四,UniWeTok使用32倍下采样率,对于需要精细空间定位的任务(如目标检测、实例分割),64个token的空间分辨率可能不足。最后,虽然训练效率显著提升,但完整的三阶段训练流程仍然需要大量计算资源,论文未提供具体的GPU小时数和碳排放数据。
独立分析的弱点
从独立分析的角度,UniWeTok存在几个可改进的弱点。第一,语义蒸馏依赖外部预训练的ViT-SO400M-16-SigLIP2-384作为教师模型,这引入了对特定教师模型的依赖:如果教师模型本身存在偏差或覆盖不足,这些缺陷会传递给tokenizer。改进方向是探索自监督语义学习或多教师集成策略,减少对单一教师的依赖。第二,三阶段训练流程的设计相对启发式,缺乏理论指导。特别是第三阶段的人脸和文本退火,为什么选择这些特定场景?如何自动识别和选择'敏感场景'?改进方向可以是开发自动化的课程学习策略,根据重建质量动态调整训练数据分布。第三,Generative-Aware Prior使用BitDance-T这个特定的生成架构,其设计选择(如扩散目标、序列长度等)可能不是最优的。可以探索更先进的轻量级生成先验,或让先验模型与主tokenizer联合训练。第四,64个token对于高分辨率图像的细粒度表示可能不足,特别是在文本渲染场景下。改进方向是探索自适应token数量的策略,根据图像复杂度动态调整压缩率。
未来方向
基于UniWeTok的成果,未来研究可以从多个方向延伸。首先,最直接的方向是将UniWeTok扩展到视频领域,将时间维度纳入tokenizer设计,构建统一的视频-图像-文本tokenizer。这需要解决时序一致性和长视频压缩的挑战。其次,探索UniWeTok在更多下游任务中的应用,如视觉问答、图像描述、视觉推理等,特别是需要精细空间理解的任务。第三,研究码本大小与任务性能的关系,是否存在一个最优的码本大小平衡点,使得重建、理解、生成三方面都达到最佳?第四,将UniWeTok与更先进的LLM架构(如Mixture-of-Experts)结合,探索大规模统一MLLM的scaling law。第五,研究tokenizer的可解释性:二进制token的每一位是否对应某种可解释的视觉概念?这种理解可能带来更好的tokenizer设计。最后,探索UniWeTok在边缘设备上的部署,64个token的表示方式可能特别适合移动端应用。
复现评估
UniWeTok的复现条件相对友好。代码和模型已在GitHub开源(https://github.com/shallowdream204/BitDance),提供了完整的训练和推理代码。数据方面,使用公开的ImageNet和DataComp-1B数据集,这些数据集广泛可用。算力需求方面,论文报告了33B训练token的规模,相比REPA的262B显著降低,但具体GPU配置和训练时间未详细说明。从消融实验的设置来看,基础实验使用250K步、batch size 128的配置,在8×A100上可能需要数天时间。大规模训练的计算需求需要参考BitDance的原始配置。复现难度主要在于三阶段训练流程的超参数调优,特别是第三阶段退火策略的数据配比。此外,语义蒸馏需要预训练的SigLIP模型,这也是额外的依赖。总体而言,对于有充足计算资源的研究团队,复现UniWeTok是可行的,但完整的三阶段训练可能需要仔细调参才能达到论文报告的性能。
论文图表