About the aul_loss

[Bian-jh]

Can I add aul_loss to rate_distortion_loss as a total loss and then use one optimizer to optimize the entire network?

[YodaEmbedding]

• When minimizing loss (i.e. training the main network) , EntropyBottleneck.quantiles is unused, and therefore receives no gradients.
• Minimizing aux_loss trains only the EntropyBottleneck.quantiles (since stop_gradient=True).

The above two losses and their gradient updates are therefore independent. That is, once the gradients are computed (via loss.backward(); aux_loss.backward()), performing one optimizer step for minimizing loss and one optimizer step for minimizing aux_loss is equivalent to performing one optimizer step for minimizing total_loss = loss + aux_loss. (Assuming the LRs are equal.)

So... yes.

[Bian-jh]

• When minimizing loss (i.e. training the main network) , EntropyBottleneck.quantiles is unused, and therefore receives no gradients.

  • Minimizing aux_loss trains only the EntropyBottleneck.quantiles (since stop_gradient=True).

The above two losses and their gradient updates are therefore independent. That is, performing one optimizer step for minimizing loss and one optimizer step for minimizing aux_loss is equivalent to performing one optimizer step of total_loss = loss + aux_loss. (Assuming the LRs are equal.)

So... yes.

Thank you for your immediate reply!！