An Embedding layer is essentially a trainable lookup table, storing one learned vector of length output_dim for every possible integer index from 0 up to input_dim minus 1 — passing it an integer, like a word's index in a vocabulary, returns that word's corresponding vector, and passing a sequence of indices returns a sequence of vectors, one per index. These vectors start out randomly initialized and are refined during training just like any other weights, gradually positioning semantically similar inputs, like related words, closer together in the embedding space as the model learns. It's the standard first layer for any model processing categorical or text data represented as integer indices.
1Understanding tf.keras.layers.Embedding()
An Embedding layer is essentially a trainable lookup table, storing one learned vector of length output_dim for every possible integer index from 0 up to input_dim minus 1 — passing it an integer, like a word's index in a vocabulary, returns that word's corresponding vector, and passing a sequence of indices returns a sequence of vectors, one per index. These vectors start out randomly initialized and are refined during training just like any other weights, gradually positioning semantically similar inputs, like related words, closer together in the embedding space as the model learns. It's the standard first layer for any model processing categorical or text data represented as integer indices.
input_dim must be at least as large as the total vocabulary size, the number of distinct possible index values — passing an index equal to or greater than input_dim raises an error, so it's worth double-checking your vocabulary size includes any reserved indices, like one for unknown or padding tokens.
import tensorflow as tf
from tensorflow.keras import layers
layer = layers.Embedding(input_dim=1000, output_dim=8)
output = layer(tf.constant([[1, 5, 9]]))
print(output.shape)2Practical Example
Here is a real-world application of tf.keras.layers.Embedding() showing how it is used in production TensorFlow code.
import tensorflow as tf
from tensorflow.keras import layers
layer = layers.Embedding(input_dim=1000, output_dim=8)
print(layer.count_params())3Best Practices
Follow these guidelines when working with tf.keras.layers.Embedding():
1. Set input_dim to your full vocabulary size, including any reserved indices for padding or unknown tokens, not just the count of 'real' distinct words
2. Use an Embedding layer, rather than one-hot encoding followed by a Dense layer, whenever a categorical feature has a large number of possible values, since it's far more memory-efficient and lets the model learn meaningful relationships between categories
3. Consider initializing an Embedding layer's weights from pretrained vectors, like GloVe or word2vec, when working with a relatively small text dataset, rather than training embeddings entirely from scratch
Tip: input_dim must be at least as large as the total vocabulary size, the number of distinct possible index values — passing an index equal to or greater than input_dim raises an error, so it's worth double-checking your vocabulary size includes any reserved indices, like one for unknown or padding tokens.
import tensorflow as tf
from tensorflow.keras import layers
layer = layers.Embedding(input_dim=1000, output_dim=8)
output = layer(tf.constant([[1, 5, 9]]))
print(output.shape)