2025-11-10T02:44:53.419690

Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models

Zheng, Li
Vector Quantized Variational Autoencoders (VQ-VAEs) leverage self-supervised learning through reconstruction tasks to represent continuous vectors using the closest vectors in a codebook. However, issues such as codebook collapse persist in the VQ model. To address these issues, existing approaches employ implicit static codebooks or jointly optimize the entire codebook, but these methods constrain the codebook's learning capability, leading to reduced reconstruction quality. In this paper, we propose Group-VQ, which performs group-wise optimization on the codebook. Each group is optimized independently, with joint optimization performed within groups. This approach improves the trade-off between codebook utilization and reconstruction performance. Additionally, we introduce a training-free codebook resampling method, allowing post-training adjustment of the codebook size. In image reconstruction experiments under various settings, Group-VQ demonstrates improved performance on reconstruction metrics. And the post-training codebook sampling method achieves the desired flexibility in adjusting the codebook size.
academic

Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models

基本信息

  • 论文ID: 2510.13331
  • 标题: Group-Wise Optimization for Self-Extensible Codebooks in Vector Quantized Models
  • 作者: Hong-Kai Zheng, Piji Li (南京航空航天大学)
  • 分类: cs.CV
  • 发表时间/会议: ICLR 2026
  • 论文链接: https://arxiv.org/abs/2510.13331

摘要

Vector Quantized Variational Autoencoders (VQ-VAEs) 通过重构任务进行自监督学习，使用码本中最接近的向量来表示连续向量。然而，VQ模型中仍然存在码本崩溃等问题。为了解决这些问题，现有方法采用隐式静态码本或联合优化整个码本，但这些方法限制了码本的学习能力，导致重构质量下降。本文提出了Group-VQ，对码本进行分组优化。每个组独立优化，组内进行联合优化。这种方法改善了码本利用率和重构性能之间的权衡。此外，我们还引入了无需训练的码本重采样方法，允许训练后调整码本大小。在各种设置下的图像重构实验中，Group-VQ在重构指标上表现出改进的性能。

研究背景与动机

问题描述

Vector Quantization (VQ) 是一种将连续特征映射到离散token的技术，在VQ-VAE中被广泛应用。然而，传统VQ训练面临码本利用率低的问题，即只有部分码向量被使用和更新，导致"码本崩溃"，限制了模型的编码能力。

现有方法的局限性

  1. Vanilla VQ: 每个码向量独立更新，容易导致码本崩溃
  2. Joint VQ方法 (如SimVQ、VQGAN-LC): 通过共享参数联合优化整个码本，虽然能达到100%利用率，但限制了码本的学习能力

研究动机

作者通过实验发现，Joint VQ虽然能快速达到100%码本利用率，但在相同利用率下，其重构质量反而不如Vanilla VQ。这表明码本利用率和重构性能之间存在权衡，需要一种更好的平衡策略。

核心贡献

  1. 提出Group-VQ方法: 基于分组的码本优化方法，平衡VQ模型中的利用率和重构性能
  2. 泛化Joint VQ方法: 基于共享参数的角度重新理解Joint VQ，并引入训练后码本采样方法
  3. 无需训练的码本调整: 实现训练后灵活调整码本大小，无需重新训练模型
  4. 全面实验验证: 在图像重构任务上验证了Group-VQ和码本重采样的有效性

方法详解

任务定义

给定图像 I∈RH×W×3I \in \mathbb{R}^{H \times W \times 3}，VQ-VAE首先使用编码器得到特征图 Z∈Rh×w×dZ \in \mathbb{R}^{h \times w \times d}，然后通过量化器将每个特征向量 z∈Rdz \in \mathbb{R}^d 替换为码本 C={qi∣qi∈Rd,i=0,1,...,n−1}C = \{q_i | q_i \in \mathbb{R}^d, i = 0,1,...,n-1\} 中最近的码向量：

q=arg⁡min⁡qi∈C∥z−qi∥,i=0,1,...,n−1q = \arg\min_{q_i \in C} \|z - q_i\|, i = 0,1,...,n-1

模型架构

Group-VQ设计

Group-VQ将码本 CC 划分为 kk 个互不相交的组（子码本）：

C=⋃j=0k−1Gj,Gj∩Gj′=∅ if j≠j′C = \bigcup_{j=0}^{k-1} G_j, \quad G_j \cap G_{j'} = \emptyset \text{ if } j \neq j'

每个组 GjG_j 独立更新，组内采用联合优化。对于码向量 qjt∈Gjq_{jt} \in G_j，其梯度更新为：

∇qjtLcmt=∇qjtLj\nabla_{q_{jt}} L_{cmt} = \nabla_{q_{jt}} L_j

这确保了每个组只受其内部码向量产生的梯度影响。

码本参数化

每个组 GjG_j 通过共享参数进行参数化：

Gj=G^jWj+bjG_j = \hat{G}_j W_j + b_j

其中：

  • G^j∈Rnj×rj\hat{G}_j \in \mathbb{R}^{n_j \times r_j}: 码本核心（固定分布采样）
  • Wj∈Rrj×dW_j \in \mathbb{R}^{r_j \times d}: 投影器（可学习）
  • bj∈Rdb_j \in \mathbb{R}^d: 偏置向量

技术创新点

1. 统一视角下的分析

  • Vanilla VQ: k=nk = n，每个码向量为一组
  • Joint VQ: k=1k = 1，整个码本为一组
  • Group-VQ: 1≤k≤n1 \leq k \leq n，平衡两种极端情况

2. 码本重采样机制

利用生成式码本的特性，训练后可以重新采样码本核心：

q~=v^Wj,v^∼N(0,I)q̃ = v̂ W_j, \quad v̂ \sim \mathcal{N}(0, I)

支持两种模式：

  • 重采样: 完全替换码本
  • 自扩展: 在原码本基础上添加新的码向量

实验设置

数据集

  • ImageNet-1k: 主要数据集
  • MS-COCO: 补充验证
  • 输入分辨率: 128×128，下采样因子f=8

评价指标

  • rFID (reconstruction FID): 重构图像与原图像的分布距离
  • LPIPS(VGG16): 感知相似性
  • PSNR: 峰值信噪比
  • SSIM: 结构相似性指数

对比方法

  • VQGAN, ViT-VQGAN, VQGAN-FC
  • FSQ, LFQ (固定码本方法)
  • VQGAN-LC, SimVQ (Joint VQ方法)

实现细节

  • 学习率: 1×10⁻⁴
  • 优化器: Adam (β₁=0.5, β₂=0.9)
  • 批大小: 32/GPU
  • 硬件: NVIDIA A5000 GPU

实验结果

主要结果

ImageNet-1k上的性能对比 (码本大小65,536):

方法组数利用率rFID↓LPIPS↓PSNR↑SSIM↑
VQGAN65,5361.4%3.740.1722.200.706
SimVQ1100.0%1.990.1224.340.788
Group-VQ6499.9%1.860.1124.370.787

Group-VQ在所有指标上都达到了最佳性能，显著优于baseline方法。

消融实验

不同组数的影响:

组数13264128512
利用率100%100%100%95.6%78.8%
rFID↓6.456.056.096.116.28

实验表明32-64个组数是最优选择，平衡了码本利用率和重构性能。

码本重采样实验

码本大小调整结果:

方法码本大小rFID↓PSNR↑
Group-VQ65,5361.8724.32
+ 下采样32,7682.1624.02
+ 上采样131,0721.7924.49
+ 自扩展131,0721.7624.51

结果验证了码本重采样方法的有效性，能够灵活调整码本大小并获得预期的性能变化。

可视化分析

通过随机投影将码向量可视化到2D空间，发现：

  1. 不同组学习了不同的特征分布
  2. 组内码向量相对相似，组间差异较大
  3. 各组的统计特性（均值、方差、使用频率）存在明显差异

相关工作

VQ改进方法分类

  1. Straight-Through Estimator改进: 优化梯度传播
  2. 多索引量化: RQ-VAE, Product Quantization等
  3. 码本改进: 本文重点关注的方向

Joint VQ方法

  • VQGAN-LC: 使用预训练特征初始化+投影层
  • SimVQ: 随机初始化+矩阵重参数化
  • LFQ/FSQ: 固定码本避免崩溃

本文将这些方法统一为"通过共享参数实现的Joint VQ"，并在此基础上提出分组优化策略。

结论与讨论

主要结论

  1. 码本利用率和重构质量存在权衡: 100%利用率不一定带来最佳重构效果
  2. 分组优化是有效的平衡策略: Group-VQ通过调节组数实现灵活控制
  3. 码本重采样提供了实用价值: 训练后可灵活调整码本大小

局限性

  1. 未在生成任务上验证: 仅在重构任务上测试，缺乏生成模型的验证
  2. 组数选择需要调优: 最优组数依赖于具体任务和数据集
  3. 计算复杂度: 多组优化可能增加训练时间

未来方向

  1. 在生成模型（如自回归模型）上验证Group-VQ效果
  2. 探索自适应组数选择策略
  3. 研究Group-VQ与其他VQ改进方法的结合

深度评价

优点

  1. 理论贡献清晰: 从组优化角度统一理解现有VQ方法，提供了新的分析视角
  2. 方法简单有效: Group-VQ设计直观，易于实现和理解
  3. 实验充分: 多数据集、多架构的全面验证，消融实验详尽
  4. 实用价值高: 码本重采样方法解决了实际应用中的灵活性需求

不足

  1. 理论分析不够深入: 缺乏为什么分组优化更有效的理论解释
  2. 适用范围有限: 主要关注图像重构，在其他模态和任务上的效果未知
  3. 计算开销分析缺失: 未详细分析多组优化的计算成本

影响力

  1. 学术价值: 为VQ研究提供了新的优化思路，可能启发后续工作
  2. 实用价值: 码本重采样方法在实际部署中很有价值
  3. 可复现性: 作者承诺提供代码，有利于方法推广

适用场景

  1. 图像/视频编码: 需要高质量重构的压缩任务
  2. 多模态学习: 作为通用的向量量化组件
  3. 生成模型: 作为tokenizer为生成模型提供离散表示

参考文献

本文主要基于以下重要工作：

  1. Van Den Oord et al. (2017) - VQ-VAE原始论文
  2. Zhu et al. (2024b) - SimVQ方法
  3. Yu et al. (2023) - LFQ方法
  4. Mentzer et al. (2023) - FSQ方法

总结: 这是一篇在VQ领域具有重要贡献的论文。Group-VQ方法简单而有效，为码本优化提供了新思路。码本重采样方法具有很强的实用价值。虽然在理论分析和适用范围上还有提升空间，但整体而言是一项高质量的研究工作。