Implementation Details. The AR and VQ-VAE module are jointly trained. The AR module contains 12 parameters. In VQ-VAE, the codebook size \(K\)  is set to 256 and the codebook dimension \(D\)  is set to 32.
The encoder and decoder both have 4 residual blocks, all with 32 channels. The encoder downsampling ratio is 2, and pixel shuffle is used when upsampling in the decoder. Thus, an image \(x \)  with shape \(32 \times 32 \times 3\)  would be quantized to \(16 \times 16 \times 1\)  vectors.
\(\beta \)  in  eq: vq loss is set to 0.25 as in [1]}, and \(\alpha \)  in  eq:totalloss is set to 125.
