跳转至

Ideal: In-DEpth ALignment Makes A Discrete Representation AutoEncoder

会议: NeurIPS 2026
arXiv: 2606.11096
代码: https://github.com/Row11n/IDEAL
领域: 图像生成
关键词: 视觉分词器, 离散自编码器, 视觉基础模型, 自回归图像生成, 跨层特征对齐

一句话总结

针对视觉基础模型(VFM)深层语义特征缺少底层细节导致重建模糊的问题,IDEAL 提出在离散向量量化前通过交叉注意力融合浅层外观与深层语义特征,并在解码端引入双分支深浅特征对齐与 DINO 判别器,在 ImageNet 上达成 0.61 rFID 重建与 1.89 gFID 自回归生成 SOTA。

研究背景与动机

构建兼具高语义表达能力与高保真像素重建能力的潜在空间,是视觉生成模型(尤其是自回归视觉生成)的基础难题。近年来基于预训练视觉基础模型(VFM,如 SigLIP2、DINOv2)的表征自编码器(RAE)展现出巨大潜力,通过复用冻结的 VFM 潜在表征指导生成,显著提升了图像合成质量与收敛效率。然而现有基于 VFM 的图像分词器普遍面临严重的重建保真度瓶颈:主流 VFM 的预训练目标是语义判别而非像素级保真重建,其深层特征高度抽象于高层语义类别,对局部的颜色、纹理、高频几何结构等低层视觉属性极不敏感。

这种语义抽象与保真重建的矛盾,在离散自回归图像生成管线中被进一步急剧放大。自回归生成需要将连续潜变量离散化为有限 Codebook 索引,向量量化(VQ)操作不可避免地丢失大量细粒度信息,若潜空间本身在量化前就缺乏底层细节,量化后缺失的低层像素信息在解码阶段将彻底不可恢复,导致重构图像出现结构扭曲与模糊。本文通过对 SigLIP2 各层特征展开深度线性探测(Linear Probing)与重建实验发现:浅层特征在像素重建上具备极高保真度(第 8 层 rFID 达 0.69,但 Top-1 准确率仅 28.66%),而深层特征语义优异却重建极差(第 24 层 Top-1 达 83.43%,但 rFID 恶化至 0.85)。浅层局部外观与深层全局语义呈现出严格的层级互补特性。

本文的切入角度是:摒弃传统分词器只承诺抽取单一 VFM 层的做法,直接在量化前建立跨深度的互补融合机制,把浅层视觉线索无缝注入深层语义表征。核心 idea:提出 IDEAL 深度对齐离散表征自编码框架,在量化前以轻量交叉注意力融合 VFM 浅层外观与深层语义特征构建统一离散 Token,在解码端通过双分支特征头同时显式监督浅层与深层特征的重构对齐,并在对抗训练中引入冻结 DINO 判别器,实现兼顾高保真像素重构、原生零样本语义兼容与最强自回归生成质量的三位一体统一表征。

方法详解

整体框架

IDEAL 的核心流水线围绕冻结 VFM 编码器、量化前跨层注意力融合、向量量化、双分支特征解码器以及轻量级像素解码器展开。输入图像首先送入冻结的 SigLIP2 骨干网络,分别提取第 8 层的浅层特征与第 24 层的深层特征;随后使用轻量级交叉注意力模块将浅层特征的局部结构融入以深层特征为 Query 的统一表示中并投影至低维空间进行离散向量量化;解量化后的离散向量经 6 层 ViT 特征解码器映射出全局隐变量,再由两个独立的线性头分别重构原始深层与浅层 VFM 特征;最后,仅将重构出的高质量深层语义特征输入轻量 CNN 像素解码器合成最终图像。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 x"] --> B["冻结 VFM 编码器<br/>SigLIP2 提取第 8 层与第 24 层"]
    B --> C["跨层注意力融合 AttnFuse<br/>深层作 Query / 浅层作 Key & Value"]
    C --> D["低维投影与向量量化 VQ<br/>l2 归一化码本生成离散 Token"]
    D --> E["双分支特征解码器 D_feat<br/>ViT 骨干并行重构深层与浅层特征"]
    E --> F["深浅特征几何对齐监督<br/>L2 距离与余弦相似度联合损失"]
    E --> G["轻量像素解码器 D_pixel<br/>基于重构深层特征生成图像"]
    G --> H["语义对抗与自编码联合优化<br/>冻结 DINO 判别器指导像素保真"]

关键设计

1. 跨层注意力融合(Cross-Attention Fusion):量化前补足局部高频结构

针对传统 VFM 分词器直接量化深层特征导致高频细节在瓶颈层永久丢失的痛点,IDEAL 引入基于注意力的跨层融合模块 AttnFuse。模型从冻结的 VFM \(\Phi(\cdot)\) 中并行抽取浅层特征 \(f^{(s)} = \Phi_{\ell_s}(x)\)(默认第 8 层)与深层特征 \(f^{(d)} = \Phi_{\ell_d}(x)\)(默认第 24 层)。由于 ViT 各层特征序列长度与空间分辨率完全一致(\(f^{(s)}, f^{(d)} \in \mathbb{R}^{B \times L \times D}\)),无需额外的空间插值或重采样。模块以深层特征作为 Query,以浅层特征作为 Key 和 Value,并保留深层特征原本的归一化层、对浅层特征赋予可学习归一化层:

\[z = \mathrm{AttnFuse}\left(f^{(d)}, f^{(s)}\right) = \mathrm{FFN}\left(\mathrm{CrossAttention}\left(Q=f^{(d)}, K=f^{(s)}, V=f^{(s)}\right)\right)\]

这种以深层为导向的查询机制使得深层语义能够自适应地从浅层特征中检索并吸纳边界、纹理等局部外观信息,既克服了线性拼接/逐元素相加对语义结构的破坏,又避免了量化前信息的不可逆丢失。

2. 稳定低维向量量化(Stabilized Vector Quantization):规避码本坍塌并支撑自回归序列建模

针对离散分词器在高维表征下极易发生码本坍塌(Codebook Collapse)与利用率不足的问题,IDEAL 采用了降维因子分解与 Codebook 向量 \(\ell_2\) 归一化策略。连续融合表征 \(z\) 首先通过降维映射到低维量化空间(\(d=64\)),然后在归一化的 Codebook \(\mathcal{C} = \{c_k\}_{k=1}^{K}\)(大小 \(K=16384\))中进行最近邻匹配,获得空间离散索引序列 \(y = \{k_i\}\):

\[k_i = \arg\min_{k \in \{1,\dots,K\}} \left\lVert z_i - c_k \right\rVert_2\]

解量化后通过升维恢复至原始特征维度 \(\tilde{z}\)。该设计配合 Commitment 损失使得 IDEAL 在 16384 的大码本容量下实现了 100% 的码本利用率,输出紧凑而高利用率的离散索引序列,非常利于后续自回归 Transformer 进行 Next-Token 预测。

3. 双分支特征重构与深浅层几何对齐(Dual-Head In-Depth Alignment):双向约束特征解码

针对单一像素重建损失难以约束解码特征保持 VFM 原始语义流形的问题,IDEAL 设计了带有两路预测头的 6 层 ViT 特征解码器 \(D_{\mathrm{feat}}\),并在序列中追加 1 个 [CLS] 标记与 4 个 Register Token 捕捉全局上下文。解码隐藏层 \(g = D_{\mathrm{feat}}(\tilde{z})\) 经两个独立线性头分别重构深层特征 \(\hat{f}^{(d)}\) 与浅层特征 \(\hat{f}^{(s)}\)。为了严格对齐原始 VFM 特征空间,设计了结合 \(\ell_2\) 空间几何距离与余弦方向一致性的复合对齐损失:

\[\mathcal{L}_{\mathrm{deep}} = \left\lVert \hat{f}^{(d)} - f^{(d)} \right\rVert_2^2 + \left(1 - \cos\left(\hat{f}^{(d)}, f^{(d)}\right)\right)\]
\[\mathcal{L}_{\mathrm{shallow}} = \left\lVert \hat{f}^{(s)} - f^{(s)} \right\rVert_2^2 + \left(1 - \cos\left(\hat{f}^{(s)}, f^{(s)}\right)\right)\]

浅层对齐损失显式督促解码特征保留低层纹理几何,而深层对齐损失则确保解出的 \(\hat{f}^{(d)}\) 与原始 SigLIP2 文本嵌入空间保持直接对齐,无需任何额外对比学习即可直接支持零样本图文理解。最后,重构出的高质量深层特征 \(\hat{f}^{(d)}\) 独占输入轻量级 CNN 像素解码器 \(D_{\mathrm{pixel}}\) 完成像素合成。

4. 语义感知对抗监督(DINO-Guided Adversarial Training):提升视觉真实度

传统 VQ-GAN 采用的 PatchGAN 判别器往往聚焦于局部高频边缘,容易引发伪影和结构性不自然。IDEAL 摒弃了传统的 PatchGAN,引入冻结的 DINOv1-s 骨干作为对抗判别器,采用 Hinge 对抗损失:

\[\mathcal{L}_{\mathrm{G}}(\hat{x}) = - \mathbb{E}\left[ D_{\mathrm{DINO}}(\hat{x}) \right]\]

冻结的 DINO 判别器兼具丰富的全局语义和感知敏感性,能提供更有语义结构指导意义的高层对抗梯度,有效约束生成图像在高阶几何与精细纹理上的真实感。

损失函数 / 训练策略

分词器端到端总训练目标包含四个部分:自编码重构损失 \(\mathcal{L}_{\mathrm{AE}}\)、标准向量量化损失 \(\mathcal{L}_{\mathrm{VQ}}\)、深层对齐损失 \(\mathcal{L}_{\mathrm{deep}}\) 以及浅层对齐损失 \(\mathcal{L}_{\mathrm{shallow}}\):

\[\mathcal{L} = \mathcal{L}_{\mathrm{AE}} + \mathcal{L}_{\mathrm{VQ}} + \mathcal{L}_{\mathrm{deep}} + \mathcal{L}_{\mathrm{shallow}}\]

其中 \(\mathcal{L}_{\mathrm{AE}} = \mathcal{L}_2(x, \hat{x}) + \mathcal{L}_{\mathrm{LPIPS}}(x, \hat{x}) + \lambda_G \mathcal{L}_G(\hat{x})\)。模型在 ImageNet-1K 上以 384×384 分辨率训练 50 个 epoch,全局 batch size 为 256,采用 AdamW 优化器(学习率 1e-4,余弦衰减退火,weight decay 为 0.05)。在下游自回归建模中,离散索引经 Ten-Crop 离线预提取以极大提高训练吞吐量,自回归模型配备 2D RoPE 旋转位置编码,通过交叉熵损失进行 Next-Token 预测。

实验关键数据

主实验

在 ImageNet-1K 基准上,IDEAL 在图像离散分词重建和条件自回归图像生成上均大幅刷新业界记录。

表 1:ImageNet 系统级图像重建性能与码本利用率对比(评价尺寸统一为 256×256)

方法类别 模型 码本容量 (Size) 码本维度 (Dim) 分词分辨率 重建 rFID ↓ 重建 rIS ↑ 码本利用率 (%)
传统分词器 TiTok 8192 64 256 1.05 191.5 100.0%
传统分词器 ImageFolder 32768 32 256 0.69 201.5 100.0%
传统分词器 VQGAN (OpenImages) 16384 4 256 1.19 – –
传统分词器 ViT-VQGAN 8192 32 256 1.28 192.3 95.0%
传统分词器 VAR 4096 32 256 0.92 196.0 100.0%
传统分词器 LlamaGen 16384 8 384 0.95 197.3 99.7%
VFM 分词器 VQRAE 16384 1536 256 1.31 – –
VFM 分词器 DINO-Tok 16384×2 832 256 1.15 – –
VFM 分词器 VFMTok 16384 12 336 0.89 215.4 100.0%
本文方法 IDEAL (Ours) 16384 64 384 0.61 230.4 100.0%

表 2:ImageNet 256×256 类条件自回归生成对比(带 CFG)

生成范式 模型 参数量 训练轮数 gFID ↓ sFID ↓ gIS ↑ 精度 Pre. ↑ 召回 Rec. ↑
Diffusion DiT-XL/2 675M 1600 2.27 4.60 278.2 0.83 0.57
Diffusion SiT-XL/2 675M 1600 2.06 4.50 270.3 0.82 0.59
Diffusion FasterDiT 675M 400 2.03 4.63 264.0 0.81 0.60
Masked MaskGiT-re 227M 555 4.02 – 355.6 – –
Next-Scale VAR 310M 350 3.30 – 274.4 0.84 0.51
AR (Base) LlamaGen-B 111M 300 6.09 7.24 182.5 0.85 0.42
AR (Base) VFMTok-B 111M 300 3.43 5.88 252.2 0.85 0.53
AR (Base) IDEAL-B 111M 300 3.38 5.18 219.8 0.84 0.51
AR (Large) LlamaGen-L 343M 300 3.07 6.09 256.1 0.83 0.52
AR (Large) VFMTok-L 343M 300 2.75 5.58 278.8 0.84 0.57
AR (Large) IDEAL-L 343M 300 2.26 5.10 219.7 0.81 0.58
AR (XXL) VFMTok-XXL 1.4B 200 2.19 5.53 278.0 0.83 0.60
AR (XXL) IDEAL-XXL 1.4B 200 1.95 4.81 260.2 0.83 0.59
AR (3B) VFMTok-3B 3.1B 200 2.07 6.23 280.4 0.81 0.62
AR (3B) IDEAL-3B 3.1B 200 1.89 5.08 270.8 0.83 0.59

消融实验

表 3:IDEAL 核心组件消融(特征融合方式、浅层对齐监督与 VFM 骨干)

消融维度 配置变体 重建 rFID ↓ 重建 rIS ↑ 说明
融合机制 交叉注意力融合 (Attention) 0.61 230.4 完整模型:保真度与语义综合最优
融合机制 线性投影拼接 (Linear) 0.63 225.9 语义保持能力下降
融合机制 无浅层融合 (None) 0.85 231.1 重建质量严重恶化(rFID +0.24)
浅层对齐 含浅层监督 \(\mathcal{L}_{\mathrm{shallow}}\) 0.61 230.4 完整模型
浅层对齐 去除浅层监督 \(\mathrm{w/o}\ \mathcal{L}_{\mathrm{shallow}}\) 0.66 229.4 缺乏显式指导,重建保真度掉点
VFM 骨干 SigLIP2-Large (默认) 0.61 230.4 兼顾保真与原生图文交互零样本分类
VFM 骨干 DINOv2 0.60 227.0 自监督特征重构略好,但语义迁移弱
VFM 骨干 DINOv3 0.54 227.9 强视觉特征下重建保真度极高

关键发现

  • 跨层融合是提升重构保真度的核心前提:去掉浅层融合后,rFID 从 0.61 骤降至 0.85,这强有力地证明了深层特征单靠自编码解码器无法凭空还原出丢失的高频几何细节,必须在量化前注入浅层外观。
  • 语义保持与生成兼备:解量化后的深层特征在零样本 ImageNet-1K 分类中获得了 80.89% 的 Top-1 准确率,极其逼近未量化的 SigLIP2 原始特征(83.23%),并完美继承了图文交互能力;在 MLLM(LLaMA-3-8B)下游评测中,SEED 得分达 68.02,显著超越原生 SigLIP2(58.24)。
  • 自回归生成效率与质量双赢:IDEAL-3B 取得 1.89 gFID,首次在自回归框架下突破 1.9 大关;IDEAL-L(343M,300 epochs)以更少的参数与训练周期比肩训练了 1600 epochs 的 675M 扩散模型(DiT-XL/2 的 2.27 gFID)。

亮点与洞察

  • 抓住 VFM 深度层级的互补天性:突破了以往 VFM 分词器在单一层特征上纠结的局限,通过系统探测证明浅层负责保真度、深层负责高语义,设计量化前注意力融合将两者统一。
  • 双分支解码与冻结 DINO 对抗机制:不仅在潜变量空间实现了深浅特征的双向几何对齐,还利用 DINO 语义判别器替代生硬的 PatchGAN,兼顾了生成像真度与语义合理性。
  • 离线预提取极大解放自回归训练开销:通过离线预提 token,避免了类似 VFMTok 在训练中动态前向分词器的巨大计算浪费,极易规模化扩展到多模态基准与十亿级参数量。

局限与展望

  • 训练数据域覆盖局限:分词器目前仅在 ImageNet 上训练,对特写人脸、密集文字以及长尾复杂场景的重建与生成仍存在一定伪影和形变。
  • 语义评测主要局限于分类与 VQA:对更细粒度密集感知任务(如语义分割、密集检测)的潜空间保持能力有待深入验证。
  • 视频生成的扩展:当前架构专注于 2D 静态图像,未来可引入时间因果注意力或 3D 卷积进一步扩展到高保真视频自回归建模。

相关工作与启发

  • vs VFMTok: VFMTok 利用可变形注意力对多尺度 VFM 特征做分词,且在 AR 训练中动态在线抽取 Token;IDEAL 创新性地从网络“深度层级”挖掘浅层高保真与深层语义的互补性,并通过离线提取极大提升了自回归吞吐量,在 rFID(0.61 vs 0.89)和 gFID(1.89 vs 2.07)上均全面占优。
  • vs LlamaGen: LlamaGen 基于传统 VQGAN 范式构建离散词表;IDEAL 基于 VFM 预训练特征赋予了潜空间极强的语义先验,在相同 Base 规格下将 gFID 从 6.09 提升至 3.38,重构 rFID 从 0.95 降至 0.61。
  • vs RAE / VQRAE: RAE 针对扩散模型且直接使用高维深层特征,VQRAE 维度高达 1536 极难优化;IDEAL 通过低维投影(\(d=64\))与跨层对齐,攻克了 VFM 离散自编码重建细节丢失与码本坍塌的两大难题。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从深浅跨层互补角度系统攻克 VFM 离散表征分词器的重构细节丢失难题。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖层级探测、重建保真、零样本图文对齐、MLLM下游及大规模自回归生成评测。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机严密,逻辑链清晰,实验数据与消融设计完整扎实。
  • 价值: ⭐⭐⭐⭐⭐ 为离散视觉表征自编码与下一代大一统多模态自回归生成确立了全新的标杆范式。