Loss and Backprop Details

[Malik7115]

Hi,

During training the vqvae backprops on multiple losses. While inputting feature maps to the model, we are given a loss, shoud I manually backpropagate and update weights through (the good ol' loss.backward() and optimizer.step()) this or is it handled implicitly?

[lucidrains]

@Malik7115 you have to sum that loss to your final loss (it is usually called an auxiliary loss, as it is auxiliary to the main objective), and then you call backwards on that

however, there are many papers that show you don't even need the auxiliary commitment loss for convergence. you can just set commitment_weight to 0 and it should learn just fine with the exponential moving average

[hassaanmuzammil]

so basically we put our image thru an encoder, pass it to the vq model and further pass the quantized to the decoder. At the end we calculate the reconstruction loss wrt input and output of decoder, add this loss to the loss given by the vq and call loss.backward()

in summary:

enc_out = encoder(images)
quantized, indices, loss = net(enc_out)
x_hat = decoder(quantized)
reconstruction_loss = mse(x_hat, images)
loss = loss + reconstruction_loss
loss.backward()

Do i understand this correctly?

[lucidrains]

@hassaanmuzammil yup, that is correct

you can check out a working example here