[Feature Request] Reservoir Restart & Batch Normalisation Before Flattening

[danieltudosiu]

trafficstars

There are two interesting features (low implementation overhead) from the paper Robust Training of Vector Quantized Bottleneck Models:

• 3.A Importance of proper scaling - Batch normalisation
  • This is in line with the Orthogonal Regularisation, but I am not sure of its interaction with it. Furthermore it will need DDP handling by using SyncBatchNorm more specifically the SyncBatchNorm.convert_sync_batchnorm() method
• 3.B Batch data-dependent codebook updates
  • This is an extended logic improvement of the K-mean initialization and of the Code Restarts.
  • Useful code is here