提出可微向量量化新框架,解决代码本更新不稳与利用率低问题
Generalized Radius and Integrated Codebook Transforms for Differentiable Vector Quantization
- 用基于半径的更新替代传统估计器,实现前向硬分配与反向可微统一
- 在图像重建、生成和推荐任务中均降低误差,代码本利用率显著提升
- 适合需要稳定训练和高效代码利用的生成模型与表征学习场景
向量量化(VQ)通过将连续潜在变量转换为离散标记,支撑现代生成与表征模型。然而,硬最近邻分配不可微,通常使用启发式直通估计器优化,导致更新步长受量化间隙影响,各代码独立更新,造成梯度不稳定和大规模下代码本利用率严重不足。本文提出GRIT-VQ(广义半径与集成变换向量量化),一种统一的代理框架,在前向传播中保持硬分配的同时使整个过程完全可微。GRIT-VQ以几何感知的可控步长沿量化方向移动潜在变量,替代直通估计器;并通过数据无关的集成变换更新代码本,所有代码共享参数而非独立更新。理论分析揭示了GRIT-VQ引入的优化动态,确立了稳定梯度流、协同代码本演化及防止崩溃的条件。在图像重建、图像生成和推荐标记化基准上,GRIT-VQ一致降低重建误差、提升生成质量与推荐准确率,并显著提高代码本利用率。
原文摘要 · Abstract (English)
Vector quantization (VQ) underpins modern generative and representation models by turning continuous latents into discrete tokens. Yet hard nearest-neighbor assignments are non-differentiable and are typically optimized with heuristic straight-through estimators, which couple the update step size to the quantization gap and train each code in isolation, leading to unstable gradients and severe codebook under-utilization at scale. In this paper, we introduce GRIT-VQ (Generalized Radius and Integrated Transform-Vector Quantization), a unified surrogate framework that keeps hard assignments in the forward pass while making VQ fully differentiable. GRIT-VQ replaces the straight-through estimator with a radius-based update that moves latents along the quantization direction with a controllable, geometry-aware step, and applies a data-agnostic integrated transform to the codebook so that all codes are updated through shared parameters instead of independently. Our theoretical analysis clarifies the fundamental optimization dynamics introduced by GRIT-VQ, establishing conditions for stable gradient flow, coordinated codebook evolution, and reliable avoidance of collapse across a broad family of quantizers. Across image reconstruction, image generation, and recommendation tokenization benchmarks, GRIT-VQ consistently improves reconstruction error, generative quality, and recommendation accuracy while substantially increasing codebook utilization compared to existing VQ variants.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。