跳转至

Improved Masked Image Generation with Knowledge-Augmented Token Representations

会议: AAAI 2026
arXiv: 2511.12032
代码: https://github.com/GuotaoLiang/KA-MIG
领域: 图像生成
关键词: 掩码图像生成, 知识图, 离散token, 先验知识增强, 图卷积网络

一句话总结

提出KA-MIG框架，通过从训练数据中挖掘三种token级语义先验知识图（共现图、语义相似图、位置-token不兼容图），使用图感知编码器学习增强的token表示，并通过轻量级加减融合机制注入现有MIG模型，持续提升多种骨干网络的生成质量。

研究背景与动机

掩码图像生成（MIG），以MaskGIT为代表，通过并行解码实现了采样速度和质量的良好平衡。其流程是：将图像编码为VQ-VAE的离散token序列 → 训练transformer预测被掩码的token → 迭代采样生成完整token序列。

然而MIG仍然落后于扩散模型，现有改进工作主要集中在改进解码/采样策略（如Token-Critic、DPC、Self-Guidance、Halton采样），几乎没有关注模型内部表示能力的提升。

作者识别的根本问题：现有MIG方法完全依赖transformer自身学习token之间的语义依赖，但这很困难，因为：

单个token缺乏明确的语义含义：VQ-VAE的codebook entry只是潜在空间中的向量，人类无法直接理解其含义

token序列通常很长（如256 tokens/image），长序列中的复杂关系难以被有效捕获

核心创新动机：既然token本身缺乏语义，能否从大规模训练数据中挖掘token之间的隐性结构规律，作为先验知识注入模型？

方法详解

整体框架

KA-MIG由三步组成： 1. 图构建：从训练数据中构建三种先验知识图 2. 图感知编码器：使用GCN学习增强的token和位置表示 3. 轻量融合机制：通过加减操作将先验知识注入MIG transformer

关键设计

1. 三种先验知识图的构建

（a）共现图 \(\mathcal{G}_{co}\)（正先验）

捕获局部范围内频繁共现的token对，反映潜在的空间-语义相关性。

  • 构建方式：统计训练数据中所有图像的token序列，记录每对token在一阶邻域（水平、垂直、对角方向）共同出现的频率
  • 构建加权无向图，剪枝低频边以减少噪声
  • 直觉：如果token A和token B经常在相邻位置出现，它们很可能编码了语义相关的视觉模式

（b）语义相似图 \(\mathcal{G}_s\)（正先验）

识别在图像合成上下文中语义相似（类似"同义词"）的token。

  • 核心假设：如果两个token在大量图像中的位置分布相似，它们可能表达相似的语义
  • 对每个token构建长度为N的位置分布向量（每个entry是该token出现在特定位置的频率）
  • 使用Jensen-Shannon散度衡量分布相似度
  • 每个token保留top-2最相似的token，形成有向图

验证实验极具说服力：将token(1013)替换为最相似token(463)，重建图像与原图在视觉上无差异（PSNR=35.78）；替换为最不相似的token(149)则严重降质（PSNR=18.97）。

（c）位置-token不兼容图 \(\mathcal{G}_p^c\)（负先验）

识别在特定类别下，哪些token不应出现在特定空间位置。

  • 对每个类别c，扫描所有训练图像，记录从未出现在某位置的token
  • 例如："飞机"类中，地面/草地纹理token几乎从不出现在图像上半部分
  • 帮助模型避免不合理的空间-语义组合

2. 图感知编码器（Graph-aware Encoder）

正先验处理：两个独立的3层GCN提取全局token表示 $\(C_{co} = f_{\theta_{co}}(\mathcal{G}_{co}, C), \quad C_s = f_{\theta_s}(\mathcal{G}_s, C)\)$ 其中 \(C\) 是VQ-VAE codebook embedding。

负先验处理：对每个类别c的每个位置i，聚合不兼容token的embedding均值 $\(p_i^c = \frac{1}{|\mathcal{I}_{i,j}|}\sum_{t \in \mathcal{I}_{i,j}} C_t W\)$ 得到位置embedding \(P^c \in \mathbb{R}^{N \times d}\)，编码空间约束。

3. 轻量融合机制

加法融合（正先验）：在每层transformer前增强未掩码token表示 $\(Z_{\overline{M}}^l = Z_{\overline{M}}^l + f_{pos}^l(C_{co}[Z_{\overline{M}}]) + f_{pos}^l(C_s[Z_{\overline{M}}])\)$

减法融合（负先验）：在每层抑制掩码位置上的不兼容token特征 $\(Z_M^l = Z_M^l - \alpha f_{neg}^l(P^c)\)$

\(f_{pos}\) 和 \(f_{neg}\) 均使用zero convolution实现，确保训练初期不干扰已有知识。

损失函数 / 训练策略

  • 使用标准的MIG训练目标（掩码token的负对数似然）
  • 冻结骨干网络，仅微调分类层和新增参数
  • 图特征可预计算存储，推理时只有轻量加减操作
  • 在MaskGIT、AutoNAT、TiTok三种骨干上验证

实验关键数据

主实验

ImageNet-256 类别条件生成

模型 类型 参数量 FID↓ IS↑ Prec↑ Rec↑
MaskGIT MIG 227M 6.18 182.1 0.80 0.52
MaskGIT-KA MIG 245M 5.69 170.2 0.81 0.50
AutoNAT MIG 194M 2.68 278.8 - -
AutoNAT-KA MIG 211M 2.45 274.1 0.82 0.56
TiTok-b64 MIG 177M 2.48 214.7 - -
TiTok-b64-KA MIG 194M 2.40 217.0 0.78 0.60
TiTok-s128 MIG 177M 1.97 281.8 - -
TiTok-s128-KA MIG 194M 1.90 271.9 0.78 0.61
VAR-d20 AR 600M 2.57 302.6 0.83 0.56
LDM-4 Diff. 400M 3.60 247.7 - -

MS-COCO 文本到图像生成

方法 FID↓ CLIP-Score↑
MaskGen 22.27 25.58
MaskGen + KA (Ours) 21.01 26.10

消融实验

配置 FID↓ IS↑ 说明
AutoNAT (baseline) 2.68 278.8
+ \(\mathcal{G}_s\) only 2.49 279.6 语义相似图贡献最大
+ \(\mathcal{G}_p\) only 2.51 285.6 位置不兼容图提升IS最多
+ \(\mathcal{G}_{co}\) only 2.51 282.1 共现图也有效
+ \(\mathcal{G}_s\) + \(\mathcal{G}_p\) 2.46 279.9 两两组合进一步提升
+ \(\mathcal{G}_{co}\) + \(\mathcal{G}_p\) 2.46 280.7
+ \(\mathcal{G}_{co}\) + \(\mathcal{G}_s\) 2.48 277.4
+ 三者全部 (KA-MIG) 2.45 274.1 FID最优

效率分析:

图类型 在线计算参数 预计算参数 在线TFLOPs
\(\mathcal{G}_{co}\) +16M +0.79M ~0
\(\mathcal{G}_s\) +16M +0.79M ~0
\(\mathcal{G}_p\) +15M +196M +0.06

最优策略：预计算 \(\mathcal{G}_{co}\) 和 \(\mathcal{G}_s\)（轻量级），在线计算 \(\mathcal{G}_p\)（避免存储每个类的图）。

关键发现

  1. 三种图各自有效且互补：单独使用都有提升，组合后进一步改善
  2. \(\mathcal{G}_s\) 对FID贡献最大：通过学习可互换的token模式增强了鲁棒性和多样性
  3. 长序列获益更多：MaskGIT/AutoNAT（256 tokens）比TiTok（64/128 tokens）获得更大提升，因为更长序列中的token依赖更复杂
  4. 仅增加约20M参数：轻量设计，实际推理开销极小
  5. AutoNAT-KA（2.45 FID）超越了更大的LlamaGen-XL（2.62）和VAR-d20（2.57）

亮点与洞察

  • 问题定义精准：指出MIG改进工作几乎全在采样策略，首次系统地关注内部表示能力
  • 数据驱动的先验知识挖掘非常实用：不需要外部标注或手工规则，纯粹从训练数据的统计规律中提取
  • "位置分布相似=语义相似"的假设简单但有效，验证实验（替换token重建）极具说服力
  • 加法融合正先验、减法融合负先验的设计直觉清晰，实现简单
  • 与骨干网络完全解耦：图特征可预计算，骨干冻结只微调少量参数，实用性极强

局限与展望

  • 三种图的构建都是静态的（来自训练数据统计），没有随训练动态更新
  • \(\mathcal{G}_p\) 的类条件特性导致存储开销较大（196M参数/1000类），可探索更紧凑的表示
  • 在TiTok等短token序列模型上改进相对有限，对新一代紧凑VQ方法的适用性待看
  • IS指标上三图组合未必最优（274 vs baseline 278），可能存在信息冗余
  • 未探索更复杂的图网络结构（如GAT、GraphSAGE），当前3层GCN可能不够
  • 缺少对更高分辨率（512×512）的系统评估

相关工作与启发

  • 与MaskGIT-SAG（自引导采样）和Halton采样互补：它们改进采样，KA-MIG改进表示，可以联合使用
  • 图神经网络在推荐系统中的共现建模思路被巧妙迁移到视觉token域
  • "负先验"（位置不兼容图）的思路类似对比学习中的难负样本挖掘
  • 可启发自回归图像生成（如LlamaGen、VAR）也引入类似的token先验知识
  • zero convolution的使用借鉴了ControlNet的思路，确保新增模块初始不干扰预训练模型

评分

  • 新颖性: ⭐⭐⭐⭐ — 先验知识图的构建方法有创新，但"从数据中挖掘统计规律"的思路相对直接
  • 实验充分度: ⭐⭐⭐⭐⭐ — 三种骨干网络、详细消融、效率分析、可视化验证，非常充分
  • 写作质量: ⭐⭐⭐⭐ — 结构清晰，图示丰富
  • 价值: ⭐⭐⭐⭐ — 提供了MIG改进的新方向，轻量即插即用设计实用性强

相关论文