论文
随笔
Tags
Slides
⤴
Search
## Vector Quantization
<h3 id="background">Background</h3> <ul> <li><strong>表征学习</strong>一直是深度学习的重点</li> </ul> <div class="slide-img"> <img src="https://miro.medium.com/v2/resize:fit:4416/format:webp/1*bvMhd_xpVxfJYoKXYp5hug.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div>
<h3 id="background-1">Background</h3> <ul> <li> <p>Encoder $\phi: X \rightarrow \bm{z} \in \textcolor{red}{\mathbb{R}^{d}}$ (连续空间)</p> </li> <li> <p>向量量化: $X \rightarrow \bm{c} \in \mathcal{C} = \{\bm{c}_k\}_{k=1}^K$ (离散空间)</p> </li> </ul> <p><span style="color: blue">✓</span> 离散化表示更符合人类语言和符号特性, 或许更利于生成任务</p> <p><span style="color: blue">✓</span> 更强的可解释性和控制性</p> <p><span style="color: blue">✓</span> 更好的可检索性</p>
<h3 id="vae">VAE</h3> <ol> <li> <p>Encoder $\phi$: 它将输入 $X \in \mathbb{R}^{H \times W \times 3}$ 映射到一个分布:</p> $$ \bm{z} \sim q(\bm{z}|X; \phi). $$<p>e.g., 高斯分布: $\phi(\bm{x}) \rightarrow (\bm{\mu}, \Sigma) \rightarrow \mathcal{N}(\bm{\mu}, \Sigma)$.</p> </li> <li> <p>Decoder $\Phi$: 它将隐变量 $\bm{z}$ 映射回 (通常来说) $X$ 的空间:</p> $$ p(X|\bm{z}; \Phi); $$</li> <li> <p>还有一个先验分布 $p(\bm{z})$ 用于辅助训练.</p> </li> </ol> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">van den Oord A., et al., Neural Discrete Representation Learning. NeurIPS, 2017.</p> </div>
<h3 id="vae-1">VAE</h3> <ul> <li> <p>训练目标:</p> $$ \begin{align*} -\log p(X) \le \underbrace{\mathbf{KL}(q_{\phi}\| p(\bm{z})) + \mathbb{E}_{\bm{z} \sim q_{\phi}} -\log p(X|\bm{z}; \Phi)}_{\text{negative ELBO}}. \end{align*} $$</li> <li> <p>最小化 KL 散度促进 $q_{\phi}$ 的散度</p> </li> <li> <p>$-\log p(X|\bm{z}; \Phi)$ 在高斯分布的假设下退化为重构损失:</p> $$ \mathcal{L}_{rec} = \| \Phi(\bm{z}) - X\|_2^2 $$</li> </ul>
<h3 id="vq-vae">VQ-VAE</h3> <ul> <li>VQ-VAE $\bm{z}$ 通过可训练的 Codebook 来解决实现离散化:</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250310215306.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div> <ul> <li> <p>$X \in \mathbb{R}^{H \times W \times 3} \overset{\phi}{\rightarrow} Z \in \mathbb{R}^{H' \times W' \times d} \overset{\varphi}{\rightarrow} \hat{Z} \in \mathcal{C}^{H' \times W' \times d} \rightarrow$</p> </li> <li> <p>$\varphi(\bm{z}) = \text{argmin}_{\bm{c} \in \mathcal{C}} \|\bm{c} - \bm{z}\|$.</p> </li> </ul>
<h3 id="vq-vae-1">VQ-VAE</h3> <ul> <li> <p>$Z \rightarrow \hat{Z}$ 是离散的, 无法传递梯度.</p> </li> <li> <p><strong>STE</strong> (straight-through estimator):</p> $$ \tilde{Z} \leftarrow Z + \text{sg}\big((\hat{Z} - Z)\big), \\ \mathrm{d}\tilde{Z} = \mathrm{d} Z + 0. $$</li> <li> <p>训练目标:</p> $$ \mathcal{L} = \mathcal{L}_{rec} + \underbrace{ \| \text{sg} (Z) - \hat{Z}\|_2^2 + \beta \cdot \| Z - \text{sg} (\hat{Z})\|_2^2. }_{\mathcal{L}_{commit}} $$</li> </ul>
<h3 id="vq-gan">VQ-GAN</h3> <ul> <li>图片 Token 化 + Next-token prediction $p(s_i | s_{< i}, \textcolor{red}{condition})$</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250311144000.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Esser P., et al. Taming Transformers for High-Resolution Image Synthesis. CVPR, 2021.</p> </div>
<h3 id="codebook-collapse">Codebook Collapse</h3> <ul> <li> <p>离散化操作终究是带来了训练困难:</p> <ol> <li>Codebook 中部分向量过于接近而造成的冗余</li> <li>Codebook 中部分向量由于训练始终匹配不到 $Z$ 导致的冗余</li> </ol> </li> <li> <p>一些方案:</p> <ol> <li>对于 codebook 采用 K-means ++ 初始化 [1];</li> <li>对于训练不充分的向量重新初始化 [2];</li> <li>用 Gumbel-softmax 替代 STE [3]</li> </ol> </li> </ul> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">[1] Lancucki A., et al. Robust Training of Vector Quantized Bottleneck Models. 2020.</p> <p style="margin: 2px 0;">[2] Dhariwai P., et al. Jukebox: A Generative Model for Music. 2020.</p> <p style="margin: 2px 0;">[3] Takida Y., et al. SQ-VAE: Variational Bayes on Discrete Representation with Self-annealed Stochastic Quantization. ICML, 2022.</p> </div>
<h3 id="residual-quantization-rq-vae">Residual Quantization (RQ-VAE)</h3> <ul> <li> <p>Collapse 问题通常是当 codebook size 增大的时候发生</p> </li> <li> <p>减小 size $\rightarrow$ 更差的表达能力 <strong>vs.</strong> 增大 size $\rightarrow$ Collpase</p> </li> <li> <p>RQ-VAE:</p> $$ Z \overset{\varphi}{\rightarrow} \textcolor{red}{\hat{Z}_1} \overset{Z - \hat{Z}_1}{\rightarrow} R_1 \overset{\varphi}{\rightarrow} \textcolor{red}{\hat{Z}_2} \overset{R_1 - \hat{Z}_2}{\rightarrow} R_2 \rightarrow \cdots $$</li> <li> <p>$\hat{Z} = \sum_{n}^N \hat{Z}_n$, 离散化表示 $(k_1, k_2, \ldots, k_N)$</p> </li> </ul> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Lee D., et al. Autoregressive Image Generation using Residual Quantization. CVPR, 2022.</p> </div>
<h3 id="tiger">TIGER</h3> <ul> <li> <p>传统推荐 (matching):</p> $$ \bm{e}_u^T \bm{e}_v, \quad v \in \mathcal{V}. $$</li> <li> <p>检索式推荐:</p> </li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250316175859.png" alt="Image" style="max-width: 80%; height: auto;margin: 0 auto;"> </div> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Rajput S., et al. Recommender Systems with Generative Retrieval. NeurIPS, 2023.</p> </div>
<h3 id="tiger-1">TIGER</h3> <ul> <li>检索式推荐:</li> </ul> <div class="slide-img"> <img src="https://raw.githubusercontent.com/MTandHJ/blog_source/master/images/20250316180725.png" alt="Image" style="max-width: 100%; height: auto;margin: 0 auto;"> </div>
<h3 id="fixed-codebook">Fixed Codebook</h3> <ul> <li> <p>固定 Codebook 为 (size: $|\mathcal{C}| = (2 \lfloor L / 2 \rfloor + 1)^d$):</p> $$ \mathcal{C} = \{-\lfloor L / 2 \rfloor, -\lfloor L / 2 \rfloor + 1, \ldots, 0, \ldots \lfloor L / 2 \rfloor - 1, \lfloor L / 2 \rfloor\}^{d}. $$</li> <li> <p>比如 $L = 3, d=3$:</p> $$ \mathcal{C} = \{ (-1, -1, -1), (-1, -1, 0), \ldots, (1, 1, 1) \}. $$</li> <li> <p>量化:</p> $$ \bm{\hat{z}} = \varphi \big(\tanh(\bm{z}) \big) = \textcolor{red}{\text{round}} \big( \textcolor{blue}{\tanh} (\bm{z}) \big). $$</li> </ul> <div class="slide-ref"> <div style="width: 100px; height: 1px; background: black; margin-bottom: 5px;"></div> <p style="margin: 2px 0;">Mentzer F., et al. Finite Scalar Quantization: VQ-VAE Made Simple. 2023.</p> </div>
<h3 id="总结">总结</h3> <ul> <li> <p>向量量化提供了一种 token 化的方式</p> </li> <li> <p>Codebook 的设定和学习仍存在问题</p> </li> </ul>