Variational Autoencoder not able to reconstruct outputs, though an Autoencoder with a similar architecture works I am trying to use a variational autoencoder-like architecture that converts images of a dataset that I created myself to an equivalent compact representation. Below is my code for the model architecture (adapted from Keras tutorials here):
# Encoder
latent_dim = 256

encoder_inputs = keras.Input(shape=(40, 60, 3))
x = keras.layers.Conv2D(32, (4, 6), activation="relu", strides=(4, 6), padding="same")(encoder_inputs)
x = keras.layers.Conv2D(64, (4, 6), activation="relu", strides=(4, 6), padding="same")(x)
x = keras.layers.Flatten()(x)
x = keras.layers.Dense(64, activation="relu")(x)
z_mean = keras.layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = keras.layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])
# z = z_mean
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")

# Decoder
latent_inputs = keras.Input(shape=(latent_dim,))
x = keras.layers.Dense(15 * 10 * num_classes, activation="relu")(latent_inputs)
x = keras.layers.Reshape((10, 15, num_classes))(x)
decoder_outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

Moreover, below is the entire model with loss calculation:
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")
        self.categorical_accuracy_tracker = keras.metrics.Mean(
            name="categorical_accuracy"
        )
        self.cce = tf.keras.losses.CategoricalCrossentropy()

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
            self.categorical_accuracy_tracker,
        ]

    def train_step(self, data):
        x, y, sample_weight = data_adapter.unpack_x_y_sample_weight(data)

        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(x)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    self.cce(y_true=y, y_pred=reconstruction, sample_weight=sample_weight)
                )
            )
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss + kl_loss
            ca = tf.keras.metrics.categorical_accuracy(y, reconstruction)


        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        self.categorical_accuracy_tracker.update_state(ca)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
            "categorical_accuracy" : self.categorical_accuracy_tracker.result(),
        }

Now, when I use this model on my data, it seems like the model is not learning anything because it generates the same output for all test and even training images. However, when I use the same architecture and training hyperparameters and uncomment the line # z = z_mean, the performance becomes a lot better, and the model generates different outputs for test images. I think anything that is possible to do with an autoencoder must also be possible with a VAE. However, I was not able to achieve a good performance with the VAE by training for more epochs or using different learning rates. Why is there this noticeable difference between them (around 30% difference in categorical accuracy), and how can I improve my VAE's performance?
To show the performance differences between the two, I have attached the loss/epoch diagrams below (loss indicates the summation of reconstruction and KL-divergence losses):
VAE:

AE (i.e., without sampling or directly using the mean):

Edit: Could this be the result of the posterior collapse problem?(see here)

I finally figured out what the problem was. The reconstruction loss is using tf.keras.losses.CategoricalCrossentropy(), which by default uses SUM_OVER_BATCH_SIZE as the reduction method. So, to solve this problem, I had to simply change the reduction method to None instead.
