Обучение не улучшает модель в байесовской нейронной сети

Question

Я работаю над задачей классификации текста, используя тензор потока вероятности. Это задача двоичной классификации, и я закодировал свои выходные данные как категориальные переменные. Это моя сеть:

def neg_log_likelihood_with_logits(y_true, y_pred):
    y_pred_dist = tfp.distributions.Categorical(logits=y_pred)
    return -tf.reduce_mean(y_pred_dist.log_prob(tf.argmax(y_true, axis=-1)))

def add_kl_weight(layer, train_size, w_value=1.0):
    w = layer.add_weight(name=layer.name+'/kl_loss_weight', shape=(),
                         initializer=tf.initializers.constant(w_value), trainable=False)
    layer.kernel_divergence_fn = get_kernel_divergence_fn(train_size, w)
    return layer

def build_bayesian_bcnn_model(input_shape, train_size):
    model_in = Input(shape=input_shape)

    dense_1 = DenseFlipout(1000, activation='relu',
                           kernel_divergence_fn=None)
    dense_1 = add_kl_weight(dense_1, train_size)
    x = dense_1(model_in)
    x = Dropout(0.5)(x)

    dense_2 = DenseFlipout(500, activation='relu',
                           kernel_divergence_fn=None)
    dense_2 = add_kl_weight(dense_2, train_size)
    x = dense_2(x)
    x = Dropout(0.5)(x)

    dense_3 = DenseFlipout(50, activation='relu',
                           kernel_divergence_fn=None)
    dense_3 = add_kl_weight(dense_3, train_size)
    x = dense_3(x)
    x = Dropout(0.5)(x)

    dense_4 = DenseFlipout(2, activation=None,
                           kernel_divergence_fn=None)
    dense_4 = add_kl_weight(dense_4, train_size)
    model_out = dense_4(x)  # logits
    model = Model(model_in, model_out)
    return model

bcnn_model = build_bayesian_bcnn_model(X_train.shape[1:], X_train.shape[0])
bcnn_model.compile(loss=neg_log_likelihood_with_logits, optimizer=Adam(1e-3), metrics=['acc'],
                   experimental_run_tf_function=False)

Когда я тренирую модель, потери начинаются примерно с 4000 и очень медленно уменьшаются. Точность сначала немного повышается, но затем она почти не меняется c. Есть ли что-то очень очевидное, что я делаю не так? Я только начинаю с байесовских нейронных сетей. Кто-нибудь может помочь?