How exactly should we scale the KL divergence of a layer when doing stochastic gradient descent?

The documentation of the property losses of the class Convolution2DFlipout states

Upon being built, this layer adds losses (accessible via the losses property) representing the divergences of kernel and/or bias surrogate posteriors and their respective priors. When doing minibatch stochastic optimization, make sure to scale this loss such that it is applied just once per epoch (e.g. if kl is the sum of losses for each element of the batch, you should pass kl / num_examples_per_epoch to your optimizer).

The documentation for Convolution2DReparameterization, DenseReparameterization, etc., states a similar thing. I've read the paper "Weight Uncertainty in Neural Networks", so I am quite familiar with the theoretical topics behind Bayesian neural networks, including the ELBO loss (which is composed of a KL divergence part and the likelihood part).

What exactly does the losses property keep track of, in mathematical terms? For example, in the first forward pass (of the first epoch) with one training example (a batch of size one), what does losses will contain? In the second forward pass with another single training example, what will losses contain? Will it accumulate the KL part of the loss for each training example fed into the network (or maybe it resets the previous losses field before calculating the new one)? If in the first forward pass we use a batch of size K > 1, rather than just 1 training example, what will losses contain? And in the second iteration with again a batch of size K > 1, what will losses contain? And in the M iteration?

The KL part of the ELBO loss does not depend on the input data, but it depends on the specifically sampled weights during the forward pass.

In your description (the documentation), the relationship between an epoch (which I assume to be a forward pass for all training examples), a batch and kl (which you assume to be the sum of the losses) is unclear. It is not even clear what kl really is. You say it is the sum of losses, but, again, it is unclear how losses is computed (for each layer and for the model and for 1, 2, or more training examples) or what it really is. Nonetheless, I know that the losses field of a model is a list of size M, where M is the number of layers of this model that contain the field losses (i.e. Bayesian layers or layers that have a regularisation term).

There's a related issue https://github.com/tensorflow/probability/issues/282, where @SiegeLordEx suggests dividing the divergence between the prior and the posterior by the total number of training examples, which is also suggested by @junpenglao in another related issue https://github.com/tensorflow/probability/issues/396. See also https://github.com/tensorflow/probability/issues/127.

Start with the documented losses property for Convolution2DFlipout, Convolution2DReparameterization, and DenseReparameterization, then review related issues 282, 396, and 127. Clarify what losses contains across forward passes and batch sizes, how the KL divergence should be scaled, and update the documentation so the behavior and epoch relationship are unambiguous.