ATOKEN：跨模态视觉分词器的4D统一表示与训练策略

1. 视觉分词器的技术演进与核心挑战

视觉分词技术作为计算机视觉领域的基础任务，其发展历程经历了从单模态到多模态、从单一任务到多任务的演进过程。传统视觉分词器通常专注于单一模态（如图像、视频或3D）或单一任务（如重建或理解），这种局限性导致视觉表示呈现碎片化状态，无法实现跨模态的知识迁移和协同优化。

在图像领域，早期的VAE和VQ-VAE模型通过编码器-解码器架构实现了图像到潜在空间的映射，但这些方法往往难以兼顾重建质量与语义理解。视频处理则面临更大的挑战，传统方法通常将视频视为图像序列进行处理，忽视了时间维度的连续性特征。3D视觉领域更是长期缺乏有效的通用表示方法，不同3D表示形式（如点云、网格、体素等）之间的转换成本高昂。

2. ATOKEN的架构设计与创新突破

2.1 4D稀疏表示的统一框架

ATOKEN的核心创新在于提出了统一的4D稀疏表示空间（t,x,y,z），其中不同模态自然占据各自的子空间：图像使用(t=0,z=0)的2D平面，视频扩展至时间维度(t,x,y,z=0)，3D资产则占据空间维度(t=0,x,y,z)。这种表示方式通过空间-时间块嵌入实现，每个输入被分割为非重叠的t×p×p块，然后通过共享线性层投影到初始嵌入。

技术实现上，4D RoPE（旋转位置编码）的应用尤为关键。与传统2D位置编码不同，4D RoPE在注意力层中为每个维度计算旋转矩阵，使得模型能够感知4D空间中的相对位置关系。具体公式表示为：

code复制Rθ,d = [[cosθ, -sinθ], [sinθ, cosθ]] (d ∈ {t,x,y,z})

其中θ根据位置索引和维度特征长度计算得到。这种设计既保留了SigLIP2的语义先验，又实现了跨模态的统一处理。

2.2 纯Transformer架构的实现细节

ATOKEN采用纯Transformer架构，包含27个Transformer块，隐藏维度d=1152，16个注意力头。编码器基于预训练的SigLIP2模型进行扩展，主要修改包括：

1. 将2D块嵌入泛化为时空块（t×p×p）
2. 添加4D RoPE位置编码
3. 引入稀疏注意力机制处理非规则网格数据

解码器则从零开始训练，包含特定任务头：

• 图像/视频：直接解码到像素空间
• 3D资产：解码为高斯溅射参数（位置偏移、颜色、尺度、不透明度、旋转）

3. 对抗性自由训练策略解析

3.1 损失函数的创新设计

传统视觉分词器通常依赖GAN训练来提升重建质量，但ATOKEN发现这种方案在Transformer架构中存在严重不稳定性。通过分解rFID误差，团队发现86.6%的误差源于协方差分量（纹理/风格），仅13.4%来自均值特征。这一洞察促使他们设计了包含四项关键损失的训练目标：

code复制L = 0.2*Lrec + 1.0*Lsem + 1e-8*LKL
Lrec = 1.0*L1 + 10.0*LPIPS + 1e3*LGram + 1.0*LCLIP

其中Gram矩阵损失的计算方式为：

code复制LGram = Σ∥FlFl^T - F̂lF̂l^T∥_F^2

（Fl表示第l层特征图，F̂l为重建特征）

3.2 训练稳定性保障机制

为确保大规模训练的稳定性，ATOKEN采用了多项创新技术：

1. 梯度累积：平衡不同任务的批次大小（图像理解8192样本，图像重建1024-4096，视频512-1024，3D 256-512）
2. EMA平滑：衰减率γ=0.9999
3. 分层学习率：编码器lr=0.1×基础lr
4. 分辨率自适应：感知损失在224×224计算，L1损失保持原生分辨率

4. 渐进式课程学习框架

4.1 四阶段训练流程

ATOKEN的训练分为四个渐进阶段，每个阶段都保留前一阶段的能力：

1. 图像基础（200k步）：

   • 分辨率：64-512像素
   • 潜在维度：32
   • 关键创新：4D块嵌入中的时间维度零填充

2. 视频动态（200k步）：

   • 扩展至48维潜在空间
   • 支持1024像素图像和512像素视频
   • 引入KV缓存机制处理长视频（见图6）

3. 3D几何（50k步）：

   • 分辨率提升至2048（图像）和1024（视频）
   • 3D表示为64×64×64体素网格
   • 多视角渲染聚合策略

4. 离散量化（100k步）：

   • 采用FSQ量化（8组6维，每维4级）
   • 总词汇量4096
   • 保持跨模态兼容性

4.2 跨模态协同效应

令人惊讶的是，多模态训练不仅没有造成性能损失，反而提升了单模态能力。在ImageNet上，图像重建质量随训练阶段持续提升：

• Stage 1：rFID 0.258
• Stage 2：rFID 0.246（↓4.7%）
• Stage 3：rFID 0.209（↓15.1%）

这表明不同模态间存在有益的协同效应，3D的几何理解能力和视频的时间建模能力都能增强图像表示。

5. 性能评估与下游应用

5.1 跨模态基准测试

ATOKEN在三大模态上均达到SOTA或接近SOTA水平：

5.2 实际应用场景

ATOKEN的统一表示使其在多种下游任务中表现出色：

1. 多模态LLM：将图像、视频、3D统一编码为兼容文本的token序列
2. 跨模态生成：
   • 文本→视频生成（通过离散token）
   • 图像→3D合成（连续潜在空间插值）
3. 高效检索：支持跨模态的联合语义搜索
4. 内容编辑：在共享潜在空间实现跨模态风格迁移

6. 技术局限与未来方向

尽管ATOKEN取得了显著突破，但仍存在一些待改进之处：

1. 计算资源需求：完整训练需要256块H100 GPU运行22天（约138k GPU小时）
2. 3D表示限制：当前仅支持中等分辨率（64×64×64体素）
3. 动态3D场景：尚未处理随时间变化的3D内容

未来可能的发展方向包括：

• 更高效的稀疏注意力机制
• 支持神经辐射场等新兴3D表示
• 探索与音频、触觉等更多模态的统一

在实际部署中，我们发现KV缓存机制对长视频处理尤为关键。例如处理5分钟视频（9000帧）时，通过16帧块处理配合KV缓存，可将内存占用降低8倍，同时保持时序一致性。这种优化使得ATOKEN在有限资源下也能处理实际应用中的长时内容。