Keras
Variational Autoencoder
Text Data
Machine Learning
NLP

How to use Keras Variational Autoencoder example with text data

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Keras VAE examples are often written for MNIST images, where the input is a dense numeric tensor and reconstruction is relatively straightforward. Text data is different because it is discrete, sparse, and usually represented as token sequences or bag-of-words vectors.

That means you cannot just take an image VAE example and swap in strings. You need a text representation, a text-appropriate decoder target, and a reconstruction loss that matches that representation.

Choose a Text Representation First

There are two common starting points:

  • bag-of-words or TF-IDF vectors
  • token sequences with embeddings

For a first working text VAE, bag-of-words is usually easier because it turns each document into a fixed-size numeric vector. That matches the dense-input assumptions of many Keras VAE examples much better than raw token sequences do.

A Simple Bag-of-Words VAE

Here is a minimal example using binary bag-of-words vectors.

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5texts = [
6    "deep learning with text",
7    "variational autoencoder example",
8    "text generation with keras",
9    "neural networks for language"
10]
11
12vectorizer = layers.TextVectorization(
13    max_tokens=1000,
14    output_mode="multi_hot"
15)
16vectorizer.adapt(texts)
17
18x_train = vectorizer(tf.constant(texts))
19input_dim = x_train.shape[1]
20latent_dim = 8

This gives each text a fixed-length vector showing which vocabulary items are present.

Build the Encoder

The encoder maps the input vector to a latent mean and log-variance.

python
1encoder_inputs = keras.Input(shape=(input_dim,))
2x = layers.Dense(128, activation="relu")(encoder_inputs)
3z_mean = layers.Dense(latent_dim, name="z_mean")(x)
4z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)

The VAE then samples from that latent distribution:

python
1def sample_latent(args):
2    z_mean, z_log_var = args
3    epsilon = tf.random.normal(shape=tf.shape(z_mean))
4    return z_mean + tf.exp(0.5 * z_log_var) * epsilon
5
6z = layers.Lambda(sample_latent, name="z")([z_mean, z_log_var])

Build the Decoder

For a bag-of-words representation, the decoder can output a vector of probabilities for each vocabulary slot.

python
1decoder_inputs = keras.Input(shape=(latent_dim,))
2x = layers.Dense(128, activation="relu")(decoder_inputs)
3decoder_outputs = layers.Dense(input_dim, activation="sigmoid")(x)
4
5decoder = keras.Model(decoder_inputs, decoder_outputs, name="decoder")
6outputs = decoder(z)

Because the targets are multi-hot vectors, a sigmoid output with binary cross-entropy is a reasonable first reconstruction choice.

Wrap It in a Custom VAE Model

python
1class VAE(keras.Model):
2    def __init__(self, encoder_inputs, outputs, z_mean, z_log_var, **kwargs):
3        super().__init__(**kwargs)
4        self.model = keras.Model(encoder_inputs, outputs, name="vae")
5        self.z_mean = z_mean
6        self.z_log_var = z_log_var
7
8    def call(self, inputs):
9        outputs = self.model(inputs)
10        kl_loss = -0.5 * tf.reduce_mean(
11            1 + self.z_log_var - tf.square(self.z_mean) - tf.exp(self.z_log_var)
12        )
13        self.add_loss(kl_loss)
14        return outputs
15
16
17vae = VAE(encoder_inputs, outputs, z_mean, z_log_var)
18vae.compile(optimizer="adam", loss="binary_crossentropy")
19vae.fit(x_train, x_train, epochs=10, batch_size=2)

This is not a state-of-the-art language model. It is a workable adaptation of the image-style VAE structure to text represented as fixed-size vectors.

Why Text Is Harder Than Images

Image VAEs reconstruct continuous pixel values. Text VAEs must reconstruct discrete symbols or distributions over vocabulary items.

That creates two important differences:

  • output space is often vocabulary-sized and sparse
  • generation quality depends heavily on the text representation

If you use bag-of-words, you lose word order. If you use token sequences, the model becomes more complex because the decoder now has to generate sequence structure rather than a flat vector.

Moving From Bag-of-Words to Sequence Models

If you want a more realistic text VAE, you usually move toward:

  • an embedding layer
  • an encoder such as LSTM, GRU, or Transformer
  • a decoder that predicts token distributions over time

That is a much harder model than the standard Keras MNIST VAE example. The core VAE ideas still apply, but the reconstruction loss now becomes sequence-based categorical cross-entropy instead of simple pixel or multi-hot reconstruction.

So if your goal is “make the Keras VAE example work with text,” the honest answer is:

  • start with bag-of-words if you want a simple adaptation
  • move to sequence VAEs only if you truly need ordered text generation

Interpreting Generated Outputs

For a bag-of-words VAE, sampling from the latent space gives you a probability vector over vocabulary items, not a perfect sentence directly. You usually need a heuristic to turn that vector back into words, such as selecting the top-scoring tokens.

That can be useful for topic-like generation or document similarity experiments, but it is not equivalent to fluent sentence generation.

Common Pitfalls

One common mistake is trying to feed raw strings directly into a dense image-style VAE. Text has to be numerically encoded first.

Another mistake is using a reconstruction loss that does not match the output representation. Multi-hot vectors and token sequences need different decoder designs and losses.

It is also easy to expect high-quality sentence generation from a simple bag-of-words VAE. That representation throws away order, so the generated output is inherently limited.

Finally, many tutorials forget that the standard Keras VAE example is optimized for images, not language. Adapting the example requires changing the data representation and often the decoder objective as well.

Summary

  • You cannot use the standard image VAE example with raw text directly.
  • Start by converting text into a numeric representation, usually bag-of-words for the simplest adaptation.
  • A bag-of-words text VAE can use dense layers and binary cross-entropy reconstruction.
  • Sequence-based text VAEs are possible but substantially more complex than the image example.
  • Choose the text representation first, because it determines the decoder design and loss function.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.