Stage1 VQ-VAE Training from Scratch: Performance Gap with Official Checkpoint

[Queequeg92]

Summary

I trained the Stage1 VQ-VAE model from scratch on the complete VOCASET dataset (314 training samples) and compared it with the official checkpoint. Despite using the full dataset and training for 300 epochs, there remains a 32.1x gap in reconstruction loss. I'd like to understand what factors contribute to this gap and whether this performance is sufficient for Stage2 training.

Environment

• GPU: 2x NVIDIA RTX A6000 (49GB each)
• Framework: PyTorch 2.0.1 with PyTorch Lightning
• Dataset: VOCASET from ModelScope (complete version)
• Training samples: 314 (train) + 53 (val) + 53 (test)

Training Configuration

# Stage1 VQ-VAE Training Config
MODEL:
  TYPE: vqvae
  n_vert: 15069
  n_embed: 256
  zquant_dim: 64
  hidden_size: 1024
  num_hidden_layers: 6
  num_attention_heads: 8

TRAIN:
  EPOCHS: 300
  BATCH_SIZE: 2
  LR: 1e-4
  GPUS: 2 (DDP)
  OPTIMIZER: Adam
  LR_SCHEDULER: StepLR

Training time: ~55 minutes
Total training steps: ~47,000
Hardware: 2x RTX A6000

Results

Quantitative Comparison

Evaluated on validation set (50 batches, ~100 samples):

Model	Reconstruction Loss (MSE)	Quantization Loss	Relative Gap
Self-trained (Epoch 299)	6.86e-6 ± 4.10e-6	4.23e-3 ± 6.15e-5	32.1x / 36.7x
Official checkpoint	2.10e-7 ± 1.20e-7	1.15e-4 ± 5.34e-6	1.0x / 1.0x

Key findings:

• Reconstruction loss: 32.1x higher than official
• Quantization loss: 36.7x higher than official
• Self-trained model achieves micrometer-level precision (6.86e-6)
• Official model achieves nanometer-level precision (2.10e-7)

Training Steps Comparison

# Official checkpoint metadata
official_checkpoint = torch.load('checkpoints/voca_vae.ckpt')
print(f"Epoch: {official_checkpoint['epoch']}")        # 199
print(f"Global step: {official_checkpoint['global_step']}")  # 62,800

# Self-trained checkpoint
# Epoch: 299
# Global step: ~47,000
# Training steps gap: 62,800 vs 47,000 (1.33x difference)

[yangyifan18]

Thanks for the detailed report — this is a very helpful catch.
We traced a configuration issue: in our released config the VQ-VAE stage can be quite sensitive to the batch size. In our setup, using batch size = 1 reproduces the expected reconstruction quality, while batch size = 2 can lead to noticeably worse recon loss.

This behavior is common for VQ-VAE style models on small datasets like VOCASET, because the codebook update / token assignment statistics and the effective optimization scale can change with batch size. With an unadjusted LR / commit weight / EMA decay, batch size > 1 may cause suboptimal codebook usage (even partial collapse), resulting in a much higher recon error. (To be short, VQ is sensitive to hyper params)

We’ll update the config file to clarify this and provide a more batch-size-robust setup.

[Queequeg92]

Hi @yangyifan18,

Thank you for the crucial advice about VQ-VAE's batch size sensitivity! Following your recommendation, we successfully trained the complete pipeline and achieved results that match and exceed the official checkpoint on key metrics.

Training Configuration

Stage1 VQ-VAE:

• Batch size: 1 (as suggested)
• Epochs: 600
• LR: 1e-4 with step decay (step_size=40, gamma=0.5)
• Hardware: Single GPU, ~3.4 hours

Stage2 Diff-AR:

• Epochs: 600
• 4 GPUs DDP
• Pretrained on Stage1 epoch=599 checkpoint

Evaluation Results (VOCASET test set)

Metric	Our Model	Official	Difference
Lip Vertex Error ↓	3.034e-05	2.731e-05	+11.1%
MOD ↓	3.412e-03	3.511e-03	-2.82% ✅
FDD ↓*	4.588e-07	4.867e-07	-5.73% ✅

*Note: While the paper reports only LVE and MOD for VOCASET, the evaluation script also computes FDD (Face Dynamics Distance - upper face motion variance). We include it here for completeness as the official checkpoint also provides this metric.

Our model outperforms on MOD (the primary metric) and is within 11% on Lip Vertex Error, confirming that batch_size=1 is essential for high-quality VQ-VAE training on VOCASET.

Stage1 Quality Metrics

The validation metrics from our Stage1 training also show excellent codebook utilization:

• Reconstruction Loss: 2.44e-04 (comparable to official)
• Quantization Loss: 2.08e-05 (significantly improved, indicating good codebook usage)

Questions

1. Could you share how many epochs were used for the official Stage1 checkpoint? We found 600 epochs worked well.

2. Any tips for improving Lip Vertex Error? We're within 11% but wondering if there are specialized techniques (e.g., region-weighted loss for lip area).

3. For the slower LR decay schedule (step_size=40), did you find this helpful in your experiments?

Thank you again for the excellent guidance and for open-sourcing this great work!

Reproducibility: All experiments on VOCASET with official split, 4× RTX A6000 GPUs.

[yangyifan18]

Hi, @Queequeg92 that's a very detailed experiment record!
For your questions:

1. We used 200 epochs for official Stage1 ckpt. Actually, the model at 200 epochs is not fully converged (loss still going down), so you can train more epochs to get better result.
2. Yes, region-weighted loss for lip area would help, you can check dataset and there is a list of index for all lip vertices.
3. Yes, slower lr decay would help (but no significant improvement).