LGQ: Learning Discretization Geometry for Scalable and Stable Image Tokenization

本文提出了可学习几何量化（LGQ），一种通过端到端学习离散化几何结构、结合温度控制软分配与正则化机制，从而在保持高生成质量的同时显著提升代码利用率并解决现有量化器在大规模词表下不稳定问题的离散图像分词方法。

要点

这篇论文介绍了一种名为 LGQ（可学习几何量化）的新方法，它旨在解决图像生成模型中一个非常棘手的问题：如何把复杂的图片“压缩”成简单的数字代码，同时还能完美地还原出来。

为了让你更容易理解，我们可以把整个过程想象成**“给图书馆里的书编目录”**。

1. 背景：为什么要“压缩”图片？

现在的 AI（比如生成图片的模型）非常强大，但它们处理原始图片就像是在处理一堆乱糟糟的沙子，太复杂、太占地方了。

• Tokenization（分词/编码）：就是要把这些“沙子”变成一个个整齐的“积木块”（代码）。AI 只需要记住这些积木块的编号，就能在脑海里重建出图片。
• 目标：积木块要越少越好（省空间），但拼出来的图要越像原图越好（保质量）。

2. 旧方法的困境：两个极端的“死胡同”

以前的方法主要有两种，但它们都有明显的缺点：

• 方法 A：硬碰硬的“最近邻居”法（传统 VQ）

  • 比喻：想象你有一堆形状各异的石头（图片特征），你要把它们扔进一个个固定的**“石头坑”**（代码本）里。规则是：扔进离它最近的那个坑。
  • 问题：随着石头越来越多，有些坑会被塞爆，而有些坑永远空着（没人用）。更糟糕的是，那些空着的坑因为没人扔石头进去，就永远学不到东西，最后整个系统变得很脆弱，稍微放大一点规模就崩溃了。这就像只有几个热门仓库在加班，其他仓库都倒闭了。

• 方法 B：死板的“格子”法（FSQ）

  • 比喻：为了避免上面的问题，这种方法直接把空间切成了整齐的方格（像棋盘一样）。不管石头是什么形状，都强行塞进最近的格子里。
  • 问题：虽然每个格子都有石头，不会倒闭，但这太死板了！现实中的石头形状千奇百怪，强行塞进方格里，要么塞不进去，要么塞得很难看（失真）。这就像非要用方形的盒子去装圆形的西瓜，要么装不下，要么把西瓜挤变形了。

3. LGQ 的解决方案：聪明的“软着陆”

LGQ 提出了一种全新的思路，它结合了上述两种方法的优点，并加入了一个**“智能温度调节器”**。

• 核心创新：从“硬选”变成“软选”

  • 比喻：以前是“非此即彼”，石头必须跳进一个坑。LGQ 说：“别急，先让石头悬浮在几个坑的上方，根据距离远近，分配不同的‘注意力’（概率）。”
  • 温度（Temperature）：这就好比**“热度”**。
    • 训练初期（高温）：石头很“软”，可以同时在几个坑之间摇摆。这让所有的坑都能收到信号，知道怎么调整自己的位置来更好地接住石头。就像所有仓库管理员都在互相学习，调整货架位置。
    • 训练后期（低温）：随着训练进行，温度慢慢降低，石头变“硬”了，最终稳稳地落入最合适的那一个坑里。
  • 结果：既保证了所有坑都能被利用（不会倒闭），又保证了最终的选择是最精准的（不会塞变形）。

• 双重保险（正则化）

  • 为了防止石头还在犹豫不决（太软）或者只挤在一个坑里（太硬），LGQ 加了两个“小规矩”：
    1. 自信度规则：强迫石头最终要做出明确的选择（不能模棱两可）。
    2. 公平分配规则：强迫石头不要全挤在一个坑里，要均匀地分布到各个坑，避免某些坑累死，某些坑闲死。

4. 实际效果：更聪明、更省钱

论文在 ImageNet（一个巨大的图片数据库）上做了测试，结果非常惊人：

• 质量更高：还原出来的图片比以前的方法更清晰、更像原图（rFID 分数更低）。
• 效率更高：以前需要动用16,000 个仓库（代码本）才能勉强拼好，LGQ 只需要动用8,000 个（甚至更少）就能达到同样的效果，而且拼得更好。
• 更稳定：不管图片库怎么变大，LGQ 都能稳稳地工作，不会像旧方法那样容易崩溃。

总结

LGQ 就像是一个“会自我进化的智能仓库管理系统”。

它不再死板地规定“石头必须进哪个坑”，而是让仓库的位置（代码本）自己去适应石头的形状。在训练过程中，它先让所有仓库都动起来学习（软分配），最后再精准定位（硬分配）。

一句话概括：LGQ 让 AI 在压缩图片时，不再是用“方盒子装圆西瓜”，而是学会了**“根据西瓜的形状，现场定制最合适的盒子”**，既省空间，又保质量，还不会让仓库倒闭。

技术摘要

这是一篇关于LGQ (Learning Discretization Geometry) 的论文技术总结，该论文提出了一种用于可扩展且稳定图像 Tokenization 的新型离散化方法。

1. 研究背景与问题 (Problem)

离散图像 Tokenization 是大规模视觉生成模型（如 VQ-VAE）中的核心瓶颈。现有的量化方法面临一个持久的权衡困境：

• 向量量化 (Vector Quantization, VQ)：能够学习灵活的几何结构，但存在严重的表示崩溃 (Representation Collapse) 问题。随着词汇表（Codebook）规模的扩大，未使用的条目保持冻结，导致梯度更新不均衡，需要复杂的启发式方法（如代码重种子、承诺损失）来缓解，且难以扩展。
• 结构化标量量化 (Structured Scalar Quantization, 如 FSQ)：通过固定的轴对齐网格实现稳定的、近乎完全的代码利用率。然而，其几何结构是固定的，假设潜在维度是各向同性和均匀分布的。面对真实世界中异质、各向异性的潜在分布时，这种固定网格会导致容量分配效率低下，无法适应数据流形。

核心问题：如何设计一种 Tokenizer，既能像 VQ 一样学习适应数据的离散化几何结构，又能像 FSQ 一样保持训练稳定性和高利用率，从而解决扩展性差和表示崩溃的问题？

2. 方法论 (Methodology)

作者提出了 LGQ (Learnable Geometric Quantization)，一种端到端学习离散化几何的离散图像 Tokenizer。

核心机制

• 软分配到硬选择 (Soft-to-Hard Assignments)：
  • LGQ 摒弃了传统的硬最近邻查找（Hard Nearest-Neighbor Lookup），转而使用温度控制的软分配 (Temperature-controlled Soft Assignments)。
  • 对于每个潜在向量 $z_{e,t}$，计算其与所有码本条目 $c_k$ 的欧氏距离，并通过 Gibbs 分布（Softmax）计算分配概率：
    $$p_{t,k} \propto \exp(-\|z_{e,t} - c_k\|^2 / \tau)$$
  • 这种软分配对应于各向同性高斯混合模型中的后验责任（Posterior Responsibilities），是变分自由能目标（期望距离 + 熵）的最小化结果。
• 直通估计器 (Straight-Through Estimator, STE)：
  • 训练阶段：使用软分配概率进行可微分训练，梯度可以流向所有码本条目，消除了 VQ 中“仅更新最近邻”导致的梯度不连续和死码问题。
  • 推理阶段：使用直通估计器将软分配转换为硬离散索引（即选择概率最大的码本条目），恢复离散的 Token 表示。
  • 理论保证：证明了当温度 $\tau \to 0$ 时，软分配收敛于硬最近邻量化。

正则化策略 (Regularization)

为了防止崩溃并鼓励平衡的利用率，LGQ 引入了两个互补的正则化项：

1. Token 级峰值正则化 (Token-level Peakedness Regularizer)：惩罚高熵分配，鼓励自信（接近 One-hot）的选择，同时保留训练初期的平滑性。
2. 全局使用正则化 (Global Usage Regularizer)：最小化经验码使用分布的 $L_2$ 范数平方。这惩罚了集中在少数代码上的分配，鼓励码本条目在整个词汇表中的平衡利用，防止代码崩溃。

训练流程

• 采用 VQGAN 风格的编码器 - 解码器架构。
• 温度参数 $\tau$ 从 1.0 线性退火至 0.1，使分配从平滑逐渐过渡到确定性的硬量化。

3. 主要贡献 (Key Contributions)

1. 提出 LGQ 框架：一种通过温度控制的 Gibbs 分配学习离散化几何的离散 Tokenizer，成功 bridging（桥接）了 VQ 的几何灵活性和 FSQ 的鲁棒性。
2. 理论分析与稳定性：
   • 基于变分自由能公式化，提供了软分配向硬量化收敛的理论证明。
   • 证明了在受限范数下的 Lipschitz 连续性，确保优化稳定性。
   • 设计了防止崩溃的正则化机制（峰值性和使用率）。
3. 实证验证：在 ImageNet 数据集上，针对不同词汇表大小进行了验证。LGQ 实现了具有竞争力的率失真（Rate-Distortion）性能，训练过程稳定，且随着词汇表增大，利用率保持平衡，未出现崩溃。

4. 实验结果 (Results)

实验在 ImageNet (128x128) 上进行，对比了 VQ, FSQ, LFQ, SimVQ 等基线模型。

• 重建质量：
  • LGQ 在 rFID (110.64) 上优于所有基线（FSQ: 125.56, VQ: 121.26, SimVQ: 117.77）。
  • 在 SSIM (0.6335) 和 LPIPS (0.4864) 上也取得了最佳或极具竞争力的表现。
• 码本利用率与效率：
  • 关键发现：FSQ 和 SimVQ 通过构造实现了近 100% 的码本利用率，但这并未带来更好的重建质量。
  • LGQ 仅激活了约 50% 的码本条目（8,199/16,384），却实现了比激活几乎所有条目的基线更好的重建质量。
  • 在更大的词汇表 ($K=65,536$) 下，LGQ 依然保持较低的 rFID (111.08) 和合理的利用率 (22.5%)，而传统 VQ 在大规模下利用率急剧下降至 8.2%。
• 率失真权衡 (Rate-Distortion Trade-off)：
  • LGQ 在更低的有效表示率（Effective Representation Rate）下实现了更低的失真。这表明 LGQ 学习到了与数据潜在分布更对齐的离散化几何，能够更有效地分配离散容量，而不是盲目地饱和整个码本。
• 几何适应性：
  • 可视化显示，LGQ 的码本中心在训练过程中发生了显著的、结构化的漂移（Drift），表明其几何结构是自适应学习的，而非固定网格。

5. 意义与影响 (Significance)

• 范式转变：将量化视为“几何学习”问题，而非简单的查找表或固定网格。LGQ 证明了通过可微分的软分配和正则化，可以消除 VQ 中的崩溃问题，同时避免 FSQ 的刚性限制。
• 系统优势：LGQ 可以作为现有离散 Tokenizer 的即插即用（Drop-in）替代方案。它消除了对承诺损失（Commitment Loss）和代码重种子（Code Reseeding）等复杂启发式方法的依赖，使得大规模词汇表的训练更加高效和稳定。
• 核心洞察：高利用率并不等同于高质量的离散化。有效的 Tokenization 应根据经验潜在分布来分配离散容量。LGQ 通过集中容量在最需要的地方，实现了更优的率失真前沿。
• 未来展望：该方法为设计可扩展的离散表示提供了原则性基础，有望扩展到视频、多模态生成以及作为 Transformer 先验的潜在空间建模。

总结：LGQ 通过引入端到端学习的离散化几何和温度控制的软分配机制，成功解决了离散 Tokenization 中扩展性与稳定性之间的矛盾，在保持高重建质量的同时，显著提高了离散容量的利用效率。