tensorflow
rnn
weight sharing
machine learning
neural networks

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.

Practice ML system design

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:

python
1import tensorflow as tf
2
3inputs = tf.keras.Input(shape=(10, 8))
4outputs = tf.keras.layers.SimpleRNN(16)(inputs)
5model = tf.keras.Model(inputs, outputs)
6
7model.summary()

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.

python
1import tensorflow as tf
2
3shared_lstm = tf.keras.layers.LSTM(32)
4
5input_a = tf.keras.Input(shape=(20, 8))
6input_b = tf.keras.Input(shape=(20, 8))
7
8encoded_a = shared_lstm(input_a)
9encoded_b = shared_lstm(input_b)
10
11model = tf.keras.Model([input_a, input_b], [encoded_a, encoded_b])
12model.summary()

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.

python
1import tensorflow as tf
2
3tf.compat.v1.disable_eager_execution()
4
5inputs = tf.compat.v1.placeholder(tf.float32, [None, 10, 8])
6cell = tf.compat.v1.nn.rnn_cell.GRUCell(16)
7outputs, state = tf.compat.v1.nn.dynamic_rnn(cell, inputs, dtype=tf.float32)

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.v1 cell 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_rnn and 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
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.