MLE / Research · ML Domain / ML System Design

VQ-VAE:最近码本量化、straight-through与commitment

离散表示需要同时训练encoder、decoder与码本,生成另需latent prior。

题目重点:离散表示需要同时训练encoder、decoder与码本,生成另需latent prior。

题意与假设

这道练习题保留了面经中值得训练的核心约束,并将需要澄清的条件单独列出。现场拿到变体时,先和面试官确认这些条件。

  • encoder给z_e,码本e_1..e_K,按欧氏平方距离选最近code。
  • 本文比较标准loss/straight-through形式;EMA码本更新是另一实现,不能同时误用两套更新。

解题路线

  1. k*=argmin||z_e−e_k||²,z_q=e_k*,decoder读取量化表示。硬argmin不可直接微分,用straight-through让重建梯度传给z_e。
  2. loss为重建项+||sg(z_e)−e||²+β||z_e−sg(e)||²。第二项更新码本,第三项让encoder贴近选中码,sg阻断相应梯度。
  3. 码本使用率/perplexity和dead code需要监控;较大K不自动更好。对新样本看重建和表示可用性,生成还要拟合离散code的先验,不能随机独立抽码保证合理。

手动走一遍

码本(0,0),(2,0),z_e=(1.8,.1)选第二码;量化误差(.2²+.1²)=.05。重建梯度经straight-through到encoder,码本loss又把第二码向z_e移动。所有输入选同一个码时压缩看似极强,但表征失去信息,检查decoder是否只学总体均值。

规模与取舍

M位置、K码、d维朴素量化O(MKd),码本存储O(Kd);近似检索可降成本但可能改变选择/训练。

易错点与追问

容易失分的边界

  • 把codebook loss和commitment的stop-gradient放反会更新错对象。
  • 离散latent不等于没有posterior collapse/码本崩塌风险。

继续深挖

  • EMA更新为什么需统计选中次数和向量和?
  • 按量化分辨率/码本大小怎样权衡码率、重建和下游性能?

技术参考

✓已复制微信号!