IRGen: Generative Modeling for Image Retrieval¶

会议: ECCV 2024
arXiv: 2303.10126
代码: GitHub (有)
领域: 图像生成
关键词: 图像检索, 自回归模型, 生成式检索, 语义分词器, 序列到序列

一句话总结¶

将图像检索重新定义为生成式建模任务，提出 IRGen——一个序列到序列模型，通过语义图像分词器将图像转化为简短的离散 token 序列，然后自回归地生成查询图像最近邻的标识符，实现端到端可微分的检索并在三个标准基准上达到 SOTA。

研究背景与动机¶

传统图像检索系统由两个独立阶段组成：

特征表示学习：将图像编码为特征向量

近似最近邻（ANN）搜索：在特征空间中高效查找最相似的数据库图像

两个阶段是独立优化的，实际部署时需要仔细调参才能让两者协调工作。虽然有联合学习的尝试（如深度量化/哈希），但仍依赖后处理的非穷举搜索，无法实现真正的端到端优化。

核心洞察：生成式建模已在翻译、对话、图像生成等领域取得巨大成功，其核心思想是将任务统一为序列生成问题。图像检索也可以被重新定义为：给定查询图像，生成其最近邻图像的标识符序列。

这带来两个关键技术挑战： 1. 如何将图像转化为足够简短的语义 token 序列（保证搜索效率） 2. 如何让这些 token 包含足够的语义信息（保证检索准确性）

现有图像分词器（VQ-VAE、RQ-VAE）为图像生成设计，不适合检索任务——序列太长（如 256 tokens）、编码低级细节而非语义信息。

方法详解¶

整体框架¶

IRGen 包含两个核心组件：

语义图像分词器（Semantic Image Tokenizer）：将图像转化为简短的离散语义 token 序列作为标识符
编码器-解码器（Encoder-Decoder）：标准 Transformer 架构，输入查询图像，自回归生成最近邻图像的标识符

关键设计¶

1. 语义图像分词器¶

为什么现有分词器不行？

特性	VQ-VAE / RQ-VAE	IRGen 分词器
目标	像素级重建	语义级分类
输入特征	空间 patch 嵌入	全局 class token
序列长度	256（8×8×4）	仅 M 个（如 4-8）
编码内容	低级细节、纹理	高级语义信息

核心设计思路：

(a) 使用全局特征而非空间特征

不使用 ViT 的 spatial token（会产生长序列），而使用 class token $\mathbf{f}_{cls}$ 作为图像表示。这将序列长度从 64 个 token 压缩到 1 个 token，然后通过残差量化展开为 $M$ 个离散码。

(b) 残差量化（Residual Quantization, RQ）

使用 $M$ 个码本，每个含 $L$ 个元素。递归地将 $\mathbf{f}_{cls}$ 映射为 $M$ 个有序离散码 $\{l_1, l_2, \ldots, l_M\}$：

\[l_m = \arg\min_{l \in [L]} \|\mathbf{r}_{m-1} - \mathbf{c}_{ml}\|_2^2$$ $$\mathbf{r}_m = \mathbf{r}_{m-1} - \mathbf{c}_{ml_m}\]

初始残差 $\mathbf{r}_0 = \mathbf{f}_{cls}$。这种递归过程天然兼容自回归生成——每个 token 依赖前面的 token。

(c) 语义监督而非重建监督

关键创新：不用像素重建损失，而用分类损失训练分词器！对原始嵌入和各级重建嵌入都施加分类损失：

\[\mathcal{L} = \mathcal{L}_{cls}(\mathbf{f}_{cls}) + \lambda_1\sum_{m=1}^{M}\mathcal{L}_{cls}(\hat{\mathbf{f}}_{cls}^{\le m}) + \lambda_2\sum_{m=1}^{M}\|\mathbf{r}_m\|_2^2\]

其中 $\hat{\mathbf{f}}_{cls}^{\le m} = \sum_{i=1}^{m}\mathbf{c}_{il_i}$ 是使用前 $m$ 个码重建的嵌入。 - 第一项：保证原始特征的语义质量 - 第二项：保证各级前缀码都包含语义信息（层级化语义） - 第三项：最小化量化残差

训练采用交替优化更新码本和网络，量化梯度用 straight-through estimator。

2. 编码器-解码器自回归检索¶

建立好离散标识符后，训练序列到序列模型：

训练过程： - 输入：图像对 $(x_1, x_2)$，其中 $x_2$ 是 $x_1$ 的最近邻 - 编码器 $\mathbb{E}$（ViT-Base）将 $x_1$ 编码为查询嵌入 - 解码器 $\mathbb{D}$（标准 Transformer，含因果自注意力 + 交叉注意力 + MLP）自回归预测 $x_2$ 的标识符序列

训练目标——最大化标识符的条件似然：

\[p(l_1, \ldots, l_M | x_1, \theta) = \prod_{m=1}^{M} p(l_i | x_1, l_1, \ldots, l_{m-1}, \theta)\]

使用 softmax 交叉熵损失在 $M$ 个离散 token 的词汇表上训练。

注意：解码器仅在离散变量上操作，不涉及任何视觉内容——这使得模型可以利用标准 Transformer 的成熟基础设施进行扩展。

3. 束搜索（Beam Search）检索¶

Top-1 检索：贪心解码，逐步生成概率最高的 token

Top-K 检索：使用束搜索（beam size = K）

具体流程： 1. 初始化 beam 为起始 token 2. 扩展候选序列，为每个候选生成可能的下一个 token 3. 按序列概率（各 token 概率之积）排序，保留 top-K 4. 重复直到解码完最后一个 token

前缀树约束：并非所有生成的标识符都对应数据库中的有效图像。构建包含所有有效码的前缀树（prefix tree），在搜索过程中约束仅考虑有效的下一个 token，大幅提高效率。

束搜索 vs ANN 搜索的核心区别： - ANN：使用预定义的距离度量计算查询-节点距离 - 束搜索：通过可微分的神经网络计算分数，条件在查询上——这使得整个检索过程可以端到端优化

损失函数 / 训练策略¶

分词器训练： - 分类损失 $\mathcal{L}_{cls}$（softmax 交叉熵）+ 量化残差最小化 - 交替优化码本和编码器网络

检索模型训练： - 自回归交叉熵损失：预测最近邻标识符的下一个 token - 训练数据：从数据库中挖掘近邻图像对

实验关键数据¶

主实验¶

三个标准基准上的 Precision 对比（所有方法使用相同数据处理和可比模型大小）：

方法	搜索方式	In-shop P@1	In-shop P@10	CUB200 P@1	CUB200 P@2	Cars196 P@1	Cars196 P@2
FT-CLIP	线性扫描	91.4	66.8	79.2	77.6	88.4	87.7
CGD	线性扫描	83.2	47.8	76.7	75.5	87.1	86.1
IRT_R	线性扫描	92.7	59.6	79.3	77.7	75.6	73.1
FT-CLIP	SPANN	90.2	62.9	78.5	77.6	88.6	88.1
IRGen	束搜索	92.4	87.0	82.7	82.7	90.1	89.9

关键指标提升（vs 最佳基线，包括线性扫描）： - In-shop P@10: +20.2%（87.0 vs 66.8） - CUB200 P@2: +6.0%（82.7 vs 77.7） - Cars196 P@2: +2.4%（89.9 vs 87.7）

尤其注意 In-shop P@10 的巨大提升——说明 IRGen 的精度在较大 K 值下不衰减（从 P@1=92.4 到 P@10=87.0），而传统方法急剧下降（FT-CLIP 从 91.4 降到 66.8）。

消融实验¶

不同标识符类型的 Precision 对比（In-shop 数据集）：

标识符类型	P@1	P@10	P@20
随机码	76.9	60.3	56.8
层次 K-means	86.1	67.2	62.1
RQ-VAE（重建监督）	85.2	68.0	63.5
语义分词器（本文）	92.4	87.0	86.6

语义分词器大幅超越其他标识符选择，验证了语义监督的关键性。

关键发现¶

生成式检索可以超越线性扫描：这是非常反直觉的——非穷举搜索方法通常被认为无法匹敌穷举搜索，但 IRGen 通过端到端优化打破了这一局限
P@K 的"高原效应"：IRGen 的 P@10/P@20/P@30 几乎不下降（87.0/86.6/86.5），说明束搜索生成的都是高质量近邻
语义监督 vs 重建监督：分类损失训练的分词器远超 RQ-VAE 的像素重建损失，验证了语义先验对检索的重要性
扩展到百万级数据集：在 ImageNet 和 Places365 上也保持优势，证明了方法的可扩展性
可能消除重排序阶段：精度的大幅提升使得实际工作流中传统不可或缺的 reranking 阶段可能被跳过

亮点与洞察¶

统一范式的开拓性：首次将图像检索真正重新定义为序列生成任务，端到端可微分，消除了特征提取-压缩-索引这三个阶段的显式界限
与自然语言处理的深度类比：就像 NLP 中的 GPT 模型生成下一个 token，IRGen 生成目标图像的标识符 token——实现了跨模态的方法论统一
全局 class token 的妙用：将序列长度从 256 压缩到仅 $M$ 个（通常 4-8），不仅提高了推理速度（二次降低时间复杂度），还天然编码了高级语义
前缀树的优雅搜索约束：巧妙地将有效性验证转化为树结构上的路径约束，避免了无效标识符的生成
概念上极简：所有组件都是标准 Transformer，可以直接利用现有的规模化技术

局限与展望¶

自回归推理速度：尽管序列大幅缩短，自回归生成的固有顺序性仍然是效率瓶颈
训练数据构造：需要预先挖掘近邻图像对作为训练数据，依赖现有检索系统
码本大小与表达力的权衡：$L$ 和 $M$ 的选择影响标识符的唯一性和区分力
动态数据库：当数据库增删图像时，需要更新前缀树和可能重新训练
可探索将此框架推广到跨模态检索（文本→图像）和多模态统一检索

评分¶

创新性: ★★★★★ — 开创性地将图像检索重新定义为生成任务，范式级别的创新
实验充分度: ★★★★★ — 三个标准基准 + 两个百万级数据集 + 详尽消融 + 多种搜索策略对比
写作质量: ★★★★★ — 问题动机清晰，方法描述系统，实验组织严谨
实用价值: ★★★★☆ — 概念突破性强但自回归推理速度和动态数据库更新是工程挑战