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)

First, there is no  z = z_mean in the original tutorials because if you set z to z_mean, your latent variable will always be deterministic, which is not a desired property we want from VAEs. I  think you;'re doing it to convert to a regular AE.
Second, I am not sure what your data is, but VAE's Gaussian prior and posterior do not seem to be a good fit with your data or as you mentioned posterior collapse may be the problem. You can look at this paper to train your VAEs or implement beta VAEs.
VAE's main aim is to generate new data, while AE's main purpose is to reconstruct the data. Depending on your prior/posterior, VAE could perform better or worse than the AE.
Hope this helps!
