Tensorflow \`RNN\` cells weight sharing
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
Introduction
In an RNN, weight sharing across time steps is not an optional trick. It is the defining behavior of the model. In TensorFlow, the usual question is not "how do I share weights across time?" but rather "how do I avoid accidentally creating separate weights when I meant to reuse the same cell or layer?"
RNNs Share Weights Across Time by Default
If you run one RNN cell over a sequence, the same kernel and recurrent kernel are reused at every time step. That is what makes the network recurrent instead of just a stack of unrelated per-step layers.
With Keras-style TensorFlow, this happens automatically when you use an RNN layer on sequence input:
The layer processes 10 time steps, but it does not create 10 different parameter sets. One parameter set is reused across the whole sequence.
Sharing Across Two Sequence Inputs
Where people usually get confused is sharing weights across different branches, not across time. If you want two inputs to use the same RNN weights, reuse the same layer object.
Because shared_lstm is one layer instance, both branches use the same learned weights.
If you instead wrote tf.keras.layers.LSTM(32) twice, you would create two independent parameter sets.
Old-Style TensorFlow RNN Cells
In TensorFlow 1.x APIs such as tf.nn.dynamic_rnn, the same idea applies. A single cell instance implies shared weights over time.
Here the GRUCell parameters are reused at each time step automatically. You do not manually create a new cell per step.
Manual Unrolling Requires Care
If you unroll an RNN yourself in a Python loop, then variable reuse becomes your responsibility. Reusing the same cell or variable scope is what keeps the weights shared.
If you instantiate a new cell inside the loop, you are no longer building a normal recurrent model. You are building separate step-specific transformations, which changes both parameter count and behavior.
That is why higher-level APIs such as dynamic_rnn, tf.keras.layers.RNN, LSTM, or GRU are usually safer.
Why Weight Sharing Matters
Shared weights do three important things:
- keep the parameter count manageable
- let the model learn patterns that apply at any time step
- allow the model to generalize across sequence length
Without sharing, a sequence model would need different parameters for step 1, step 2, step 3, and so on, which defeats the usual purpose of an RNN.
Common Pitfalls
- Assuming you must implement time-step sharing manually in standard TensorFlow RNN APIs.
- Creating two separate RNN layer instances when you meant to share one across branches.
- Instantiating a new cell inside a manual loop and accidentally breaking weight reuse.
- Confusing sharing across time steps with sharing across multiple inputs or towers.
- Mixing old
tf.compat.v1cell APIs and modern Keras APIs without being clear about which reuse model you are following.
Summary
- In TensorFlow RNNs, weights are shared across time steps by default.
- To share weights across multiple inputs, reuse the same layer or cell instance.
- '
dynamic_rnnand Keras RNN layers handle time-step reuse for you.' - Manual unrolling only works correctly if you also reuse the same variables.
- Most weight-sharing bugs come from creating extra layer instances, not from TensorFlow failing to share automatically.
Related reading
- Tensorflow `RNN` many to many Time series for binary labels
- tensorflow running error with cublas
- Tensorflow Sampled Softmax \`Loss\` Correct Usage
- Tensorflow save the model with smallest validation error
- Tensorflow save the model with smallest validation error
- TensorFlow Saver has 5 models limit
- TensorFlow Saver has 5 models limit
- TensorFlow saving into/loading a graph from a file
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.