Keras
embeddings
pretrained-embeddings
machine-learning
deep-learning

Keras initialize large embeddings layer with pretrained embeddings

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

Pretrained word vectors can give a Keras model a strong starting point, especially when your dataset is small or your domain uses ordinary language. The core job is to align your tokenizer's integer ids with rows in an embedding matrix, then load that matrix into tf.keras.layers.Embedding.

How the Embedding Layer Expects Its Weights

An embedding layer stores a two-dimensional weight matrix. Each row corresponds to one token id, and each column is one dimension of the vector space. If your vocabulary size is vocab_size and each vector has length embedding_dim, the weight matrix must have shape (vocab_size, embedding_dim).

Keras does not automatically map a GloVe or word2vec file to your tokenizer. You have to build that matrix yourself.

python
1import numpy as np
2import tensorflow as tf
3
4vocab_size = 6
5embedding_dim = 4
6
7embedding_matrix = np.array([
8    [0.0, 0.0, 0.0, 0.0],   # padding token
9    [0.1, 0.2, 0.3, 0.4],   # "cat"
10    [0.4, 0.3, 0.2, 0.1],   # "dog"
11    [0.7, 0.1, 0.2, 0.9],   # "runs"
12    [0.3, 0.8, 0.1, 0.2],   # "fast"
13    [0.0, 0.0, 0.0, 0.0],   # unknown token
14], dtype="float32")
15
16embedding_layer = tf.keras.layers.Embedding(
17    input_dim=vocab_size,
18    output_dim=embedding_dim,
19    weights=[embedding_matrix],
20    trainable=False,
21)

The weights argument receives a list because Keras layers can have multiple weight arrays. For Embedding, that list contains exactly one matrix.

Building the Embedding Matrix From a Tokenizer

In a realistic project, your tokenizer creates a dictionary from words to integer ids. You then look up each word in the pretrained embedding file and place the vector in the matching row.

python
1import numpy as np
2
3word_index = {
4    "cat": 1,
5    "dog": 2,
6    "runs": 3,
7    "fast": 4,
8}
9
10pretrained = {
11    "cat": np.array([0.1, 0.2, 0.3, 0.4], dtype="float32"),
12    "dog": np.array([0.4, 0.3, 0.2, 0.1], dtype="float32"),
13    "runs": np.array([0.7, 0.1, 0.2, 0.9], dtype="float32"),
14}
15
16embedding_dim = 4
17vocab_size = len(word_index) + 1
18embedding_matrix = np.zeros((vocab_size, embedding_dim), dtype="float32")
19
20for word, index in word_index.items():
21    vector = pretrained.get(word)
22    if vector is not None:
23        embedding_matrix[index] = vector

Notice that "fast" stays as zeros because it is not present in pretrained. That is acceptable, but you should be aware of how many tokens are missing. Large coverage gaps reduce the benefit of using pretrained vectors in the first place.

Plug the Matrix Into a Model

Once the matrix is ready, the layer fits into a normal text model. This example uses integer-tokenized sequences and a pooling layer for a simple classifier.

python
1import tensorflow as tf
2import numpy as np
3
4x_train = np.array([
5    [1, 3, 4, 0],
6    [2, 3, 0, 0],
7    [1, 4, 0, 0],
8], dtype="int32")
9
10y_train = np.array([1, 0, 1], dtype="float32")
11
12model = tf.keras.Sequential([
13    tf.keras.layers.Embedding(
14        input_dim=vocab_size,
15        output_dim=embedding_dim,
16        weights=[embedding_matrix],
17        trainable=False,
18    ),
19    tf.keras.layers.GlobalAveragePooling1D(),
20    tf.keras.layers.Dense(1, activation="sigmoid"),
21])
22
23model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
24model.fit(x_train, y_train, epochs=3, verbose=0)

If you want the model to fine-tune the vectors during training, set trainable=True. Freezing the embeddings is often a good first experiment, because it tells you whether the pretrained space already helps without adding more trainable parameters.

Working With Large Embedding Tables

The tricky part is often not syntax, but scale. A vocabulary of 500,000 words with 300-dimensional vectors consumes a lot of memory. Roughly speaking, that is 500,000 multiplied by 300 multiplied by 4 bytes for float32, which is about 600 MB just for the matrix.

A few practical ways to reduce pressure:

  • Limit the tokenizer vocabulary to the most frequent words.
  • Use only the embeddings that match your tokenizer instead of loading every vector into memory permanently.
  • Prefer float32 unless you have a measured reason to use something larger.
  • Revisit whether your task really needs a massive static vocabulary.

For many applications, keeping the top 20,000 to 100,000 tokens captures most of the useful signal.

Common Pitfalls

The most common error is an off-by-one mismatch between tokenizer ids and matrix rows. If token id 1 is the first real word, row 0 usually needs to be reserved for padding or a special token.

Another frequent issue is using the wrong embedding dimension. If the pretrained file contains 300 values per word, then output_dim must also be 300. Keras will reject a matrix whose shape does not match the declared layer size.

Case handling also matters. If your tokenizer lowercases text but your embedding lookup uses original casing, many words will appear missing even though the vectors exist. Keep tokenization and embedding lookup rules aligned.

Finally, do not assume trainable=False is always best. Frozen embeddings can stabilize training, but domain-specific tasks sometimes benefit from fine-tuning. Treat that choice as an experiment, not a rule.

Summary

  • Keras embeddings expect a weight matrix shaped like (vocab_size, embedding_dim).
  • You must align tokenizer integer ids with the correct pretrained vectors.
  • Missing words can remain zero-initialized, but low coverage limits the value of pretrained embeddings.
  • 'trainable=False freezes the vectors; trainable=True allows fine-tuning.'
  • Large vocabularies can consume significant memory, so vocabulary trimming is often necessary.

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.