PyTorch生成式人工智能——VQ-VAE详解与实现_vqvae

PyTorch生成式人工智能——VQ-VAE详解与实现

• • 0. 前言
  • 1. VQ-VAE 技术原理
  • • 1.1 引入离散潜变量
    • 1.2 向量量化
    • 1.3 损失函数
    • 1.4 指数滑动平均
    • 1.5 梯度直通 (Straight-Through)
  • 2. VQ-VAE 网络架构
  • 3. 实现 VQ-VAE
  • • 3.1 模型构建
    • 3.2 模型训练
  • 相关链接

0. 前言

在传统的变分自编码器 (Variational Auto-Encoder, VAE) 中，模型学习的是一个连续的潜在表示 (latent representation)。然而，对于许多模态的数据，如语言、语音或某些图像特征，离散的潜表示往往更加自然有效。
VQ-VAE (Vector Quantised-Variational AutoEncoder) 的核心思想就是将 VAE 的连续潜变量离散化。它通过学习一个码本 (Codebook) 来实现这一点，码本是一个包含有限个嵌入向量的字典。模型不是直接输出一个连续的潜在向量，而是从码本中找出与编码器输出最接近的嵌入向量来代替它。这种离散化带来了以下优势：

• 兼容性：离散的潜在空间可以很自然地与自回归模型(如 PixelCNN、Transformer )结合，用于强大的先验建模，从而生成高质量的新样本
• 计算效率：对于下游任务，处理离散的 token 通常比处理连续的向