Keras
Embedding Layers
Deep Learning
Neural Networks
Machine Learning

Keras embedding layers how do they work?

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

When building neural networks for tasks like text classification or recommendation systems, you often need to convert discrete categories (words, user IDs, product IDs) into continuous numerical vectors. Keras provides the Embedding layer for exactly this purpose. It acts as a trainable lookup table that maps integer indices to dense vectors of a fixed size. During training, the network learns vector representations that capture meaningful relationships between the input categories. This article explains how embedding layers work internally, walks through practical code examples, and covers common mistakes.

What an Embedding Layer Does

At its core, an embedding layer is a weight matrix of shape (input_dim, output_dim). When you pass an integer index i to the layer, it returns the i-th row of this matrix. No multiplication happens. It is a pure lookup operation, which makes it very fast.

For example, if you have a vocabulary of 10,000 words and choose an embedding dimension of 128, the embedding layer creates a matrix with 10,000 rows and 128 columns. Word index 42 maps to row 42, which is a vector of 128 numbers.

The values in this matrix start as small random numbers and are updated through backpropagation during training, just like any other layer's weights.

Creating an Embedding Layer in Keras

Here is a basic example that creates a model with an embedding layer:

python
1from tensorflow.keras.models import Sequential
2from tensorflow.keras.layers import Embedding, Flatten, Dense
3
4model = Sequential([
5    Embedding(input_dim=10000, output_dim=128, input_length=50),
6    Flatten(),
7    Dense(64, activation='relu'),
8    Dense(1, activation='sigmoid')
9])
10
11model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
12model.summary()

The three key parameters are:

input_dim is the size of the vocabulary, meaning the total number of unique tokens. Integer indices passed to this layer must be in the range 0 to input_dim - 1.

output_dim is the dimensionality of each embedding vector. Common values range from 32 for small datasets to 300 or more for large language tasks.

input_length is the length of each input sequence. If your sentences are padded to 50 tokens, set this to 50.

Understanding the Input and Output Shapes

The embedding layer expects integer input. If you pass a batch of sequences with shape (batch_size, sequence_length), the output will have shape (batch_size, sequence_length, output_dim).

python
1import numpy as np
2
3# Simulate 3 sentences, each padded to length 5
4sample_input = np.array([
5    [1, 14, 25, 0, 0],
6    [7, 3, 88, 12, 0],
7    [42, 9, 0, 0, 0]
8])
9
10embedding_layer = Embedding(input_dim=100, output_dim=8, input_length=5)
11output = embedding_layer(sample_input)
12print(output.shape)  # (3, 5, 8)

Each integer in the input gets replaced by its corresponding 8-dimensional vector, so a sequence of 5 integers becomes a sequence of 5 vectors.

Using Pretrained Embeddings

Instead of learning embeddings from scratch, you can initialize the layer with pretrained vectors like GloVe or Word2Vec. This is especially useful when your training data is small.

python
1import numpy as np
2
3# Assume you loaded a pretrained embedding matrix
4# with shape (vocab_size, embedding_dim)
5pretrained_matrix = np.random.rand(10000, 300)  # placeholder
6
7embedding_layer = Embedding(
8    input_dim=10000,
9    output_dim=300,
10    weights=[pretrained_matrix],
11    input_length=50,
12    trainable=False  # freeze the weights
13)

Setting trainable=False prevents the pretrained vectors from being updated during training. This preserves the knowledge captured in the pretrained model. If you have enough training data, you can set trainable=True to fine-tune the embeddings for your specific task.

Embedding Layers Beyond NLP

Embedding layers are not limited to text. Any categorical feature with a large number of possible values can benefit from embeddings.

Recommendation systems use embeddings for user IDs and item IDs. The model learns vectors for each user and each item, and the dot product of a user vector and an item vector predicts how much that user will like that item.

python
1from tensorflow.keras.layers import Input, Embedding, Flatten, Dot
2from tensorflow.keras.models import Model
3
4user_input = Input(shape=(1,))
5item_input = Input(shape=(1,))
6
7user_embedding = Embedding(input_dim=5000, output_dim=32)(user_input)
8item_embedding = Embedding(input_dim=10000, output_dim=32)(item_input)
9
10user_flat = Flatten()(user_embedding)
11item_flat = Flatten()(item_embedding)
12
13dot_product = Dot(axes=1)([user_flat, item_flat])
14model = Model(inputs=[user_input, item_input], outputs=dot_product)
15model.compile(optimizer='adam', loss='mse')

Categorical features in tabular data like zip codes, product categories, or device types can also be represented as embeddings. This often works better than one-hot encoding when the number of categories is large, because one-hot vectors are sparse and high-dimensional while embeddings are dense and compact.

How Embeddings Learn Meaningful Representations

During training, the loss function provides a gradient signal that flows back through the network and into the embedding matrix. Words (or categories) that appear in similar contexts will have their vectors pushed closer together in the embedding space. For example, "king" and "queen" will end up with similar vectors because they appear in similar sentence structures.

This property enables operations like vector arithmetic. The classic example is that vector("king") - vector("man") + vector("woman") produces a vector close to vector("queen").

Common Pitfalls

Passing out-of-range indices. If your vocabulary size is 10,000 (input_dim=10000), valid indices are 0 through 9,999. Passing 10,000 or higher will cause an index-out-of-bounds error that can be hard to debug in a large pipeline.

Choosing an embedding dimension that is too large. For a small vocabulary (a few hundred categories), an embedding dimension of 300 is excessive and will lead to overfitting. A rule of thumb is to start with min(50, (vocab_size // 2)) and tune from there.

Forgetting to pad sequences. The input_length parameter requires all input sequences to have the same length. Use tf.keras.preprocessing.sequence.pad_sequences to pad shorter sequences with zeros before passing them to the model.

Not masking padded values. Padding tokens (index 0) get their own embedding vector, which means the model treats them as real input. Add mask_zero=True to the Embedding layer so that downstream layers like LSTMs can ignore padded positions.

Summary

Keras embedding layers convert integer indices into dense vectors by performing a lookup in a trainable weight matrix. They are the standard way to handle categorical inputs in neural networks, whether for words in NLP, users and items in recommendation systems, or any other high-cardinality categorical feature. You can initialize them randomly and train from scratch, or load pretrained vectors and optionally freeze them. The key parameters to get right are input_dim (vocabulary size), output_dim (vector dimensionality), and input_length (sequence length), and you should always ensure your input indices are within the valid range.


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.