Has anyone experienced a sudden and drastic increase of the loss while training a Bayesian neural network?

I am trying to train the following Bayesian CNN model, which contains several Bayesian convolution and dense layers. The inputs are images with shape (64, 80, 3) and the outputs are normal distributions, one for each of the 289 output neurons (i.e. I need to make 289 predictions for each image). You can see the details of the architecture below.

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
camera (InputLayer)             [(None, 64, 80, 3)]  0                                            
__________________________________________________________________________________________________
conv2d_flipout (Conv2DFlipout)  (None, 64, 80, 8)    441         camera[0][0]                     
__________________________________________________________________________________________________
activation (Activation)         (None, 64, 80, 8)    0           conv2d_flipout[0][0]             
__________________________________________________________________________________________________
max_pooling2d (MaxPooling2D)    (None, 32, 40, 8)    0           activation[0][0]                 
__________________________________________________________________________________________________
conv2d_flipout_1 (Conv2DFlipout (None, 32, 40, 16)   2321        max_pooling2d[0][0]              
__________________________________________________________________________________________________
activation_1 (Activation)       (None, 32, 40, 16)   0           conv2d_flipout_1[0][0]           
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D)  (None, 16, 20, 16)   0           activation_1[0][0]               
__________________________________________________________________________________________________
conv2d_flipout_2 (Conv2DFlipout (None, 16, 20, 32)   9249        max_pooling2d_1[0][0]            
__________________________________________________________________________________________________
activation_2 (Activation)       (None, 16, 20, 32)   0           conv2d_flipout_2[0][0]           
__________________________________________________________________________________________________
max_pooling2d_2 (MaxPooling2D)  (None, 8, 10, 32)    0           activation_2[0][0]               
__________________________________________________________________________________________________
flatten (Flatten)               (None, 2560)         0           max_pooling2d_2[0][0]            
__________________________________________________________________________________________________
dense_flipout (DenseFlipout)    (None, 512)          2621953     flatten[0][0]                    
__________________________________________________________________________________________________
target_locations (DenseFlipout) (None, 289)          296226      dense_flipout[0][0]              
__________________________________________________________________________________________________
target_rhos (DenseFlipout)      (None, 289)          296226      dense_flipout[0][0]              
__________________________________________________________________________________________________
target (DistributionLambda)     ((None, 289), (None, 0           target_locations[0][0]           
                                                                 target_rhos[0][0]                
==================================================================================================
Total params: 3,226,416
Trainable params: 3,226,410
Non-trainable params: 6
__________________________________________________________________________________________________

The activation functions of the convolution layers and the first dense layer (512) are ReLUs, while the dense layers representing both the locations and rhos have None as activation function (i.e. no or linear activation function).

I tried to train this model. The training ELBO loss starts around 61k and the validation ELBO loss around 58k. Both the training and validation ELBO losses decrease for roughly 10 epochs. However, around the 12th, both the training and validation ELBO losses suddenly and drastically increase. The training loss goes from ~11k to ~27k, while the validation loss, at the 13th epoch, goes from ~7k to ~110k. What could have caused this? Maybe a numerical error?

The dataset is quite small: it contains slightly more than 10k instances. However, the non-Bayesian version (i.e. the same layers, but non-Bayesian ones, apart from the output layers, because, in the Bayesian case, I am trying to learn both the mean and variance of the normal distributions) of this Bayesian model performs a lot better than this Bayesian model, in terms of AUC, throughout all epochs: in the Bayesian case, the model cannot predict at all, i.e. an AUC of 0.5 or even less (which means that the model predicts more incorrectly than correctly), for almost all of the 289 outputs, while, in the non-Bayesian case, the AUC is definitely higher (i.e. in certain outputs, the AUC is like 0.8 or even 1.0).

The issue does not mention a repository file, test, or entry point to modify. Start by reproducing the Bayesian CNN training run from the architecture and ELBO values provided, then compare the epoch-12 behavior with the non-Bayesian model; done means identifying and documenting the cause of the loss spike and poor AUC.