Introduction

Sometimes you want only part of an embedding table to keep learning while the rest stays fixed. A common example is starting from pretrained embeddings, freezing most rows, and fine-tuning only a domain-specific subset. In TensorFlow, the cleanest way to do this is usually not to mutate only part of the variable by magic, but to control which rows receive gradient updates.

The Key Idea: Mask the Gradient

An embedding matrix is usually just a trainable variable. During optimization, all rows referenced by the current batch receive gradients unless you intervene.

A practical way to update only selected rows is:

• compute the gradient for the whole embedding matrix
• zero out the rows you want frozen
• apply the masked gradient

This keeps the model structure simple while controlling which rows are allowed to change.

A Small TensorFlow Example

The following example trains only rows 1 and 3 while keeping the others fixed.

1import tensorflow as tf
2
3embedding = tf.Variable([
4    [0.1, 0.1],
5    [0.2, 0.2],
6    [0.3, 0.3],
7    [0.4, 0.4],
8], dtype=tf.float32)
9
10trainable_rows = tf.constant([0.0, 1.0, 0.0, 1.0], dtype=tf.float32)
11mask = tf.reshape(trainable_rows, (-1, 1))
12optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
13indices = tf.constant([1, 3])
14target = tf.constant([[1.0, 1.0], [2.0, 2.0]])
15
16with tf.GradientTape() as tape:
17    gathered = tf.gather(embedding, indices)
18    loss = tf.reduce_mean(tf.square(gathered - target))
19
20grad = tape.gradient(loss, embedding)
21masked_grad = grad * mask
22optimizer.apply_gradients([(masked_grad, embedding)])
23
24print(embedding.numpy())

Rows 1 and 3 can move. Rows 0 and 2 remain unchanged because their gradients were multiplied by zero.

Why This Works Better Than Manual Row Assignment

A tempting idea is to manually assign updated values only to certain rows after training. That is usually harder to reason about and less consistent with the optimizer's normal behavior.

Gradient masking keeps the optimization loop explicit:

• TensorFlow still computes the real loss
• the optimizer still applies updates
• only the forbidden rows are neutralized

That makes the training rule easier to maintain.

Another Option: Split Frozen and Trainable Embeddings

If the trainable and frozen vocabularies are clearly separate, a structural split may be cleaner than masking.

For example, keep one constant embedding table for frozen rows and one trainable table for the subset that should adapt.

1import tensorflow as tf
2
3frozen_embeddings = tf.constant([
4    [0.1, 0.1],
5    [0.3, 0.3],
6], dtype=tf.float32)
7
8trainable_embeddings = tf.Variable([
9    [0.2, 0.2],
10    [0.4, 0.4],
11], dtype=tf.float32)

This approach can be cleaner when the index mapping between frozen and trainable tokens is already well-defined, but it complicates lookup logic because you now manage two tables instead of one.

tf.stop_gradient Is Related but Different

tf.stop_gradient can prevent gradients from flowing through a tensor expression, but it is usually more convenient for graph composition than for row-level partial updates inside one embedding variable.

If your goal is "freeze some rows, train others," masked gradients or split tables are usually more direct.

When Partial Updates Are Useful

This pattern is helpful when:

• most embeddings come from pretrained vectors and should stay stable
• only new or domain-specific tokens should keep adapting
• you want to reduce destructive drift in a large embedding table

It is less useful when the whole vocabulary should learn together or when the trainable subset changes constantly in hard-to-maintain ways.

Common Pitfalls

• Assuming TensorFlow provides a built-in flag that says "train only these rows" for a normal embedding variable.
• Manually overwriting rows after optimization without thinking through how the optimizer state behaves.
• Using a mask with the wrong shape and accidentally masking columns instead of rows.
• Splitting embeddings into multiple tables without a clear index-mapping strategy.
• Forgetting that only rows touched by the batch can receive gradients in the first place.

Summary

• To update only part of an embedding matrix, control the gradients rather than hoping the optimizer will infer your intent.
• Gradient masking is a straightforward way to freeze selected rows.
• Splitting frozen and trainable embeddings can also work when the vocabulary partition is clean.
• 'tf.stop_gradient is related, but row-level masking is often clearer for this problem.'
• The best solution is the one that keeps the training rule explicit and the index mapping understandable.