Keras shared layers with different trainable flags
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 Keras, a shared layer means one layer instance reused in multiple places. Because it is the same object with the same weights, it does not support conflicting trainable settings at the same time. If one path sees the layer as trainable, every use of that shared layer sees the same choice.
Why the trainable Flag Is Global to the Layer
The trainable property belongs to the layer instance, not to each call site. That is the key point that makes this topic confusing.
If you write a model like a Siamese network, both branches may reuse the same Dense, Conv2D, or embedding layer. The weights are shared on purpose, which means the optimizer sees one shared set of variables. Keras cannot update those weights for branch A while freezing the exact same variables for branch B in the same training step.
A Shared Layer Example
Here is a normal shared-layer pattern:
This is true weight sharing. Both branches use the same variables inside shared_dense.
If you do this:
the layer becomes frozen for both branches, because there is still only one layer instance.
What You Cannot Do
You cannot keep one branch trainable and another branch frozen while both branches point to the same Keras layer instance in the same model. That would require one set of weights to be simultaneously updated and not updated, which is internally inconsistent.
This is why code that tries to "share a layer but freeze only one use of it" is conceptually asking for two different behaviors from one object.
Correct Alternatives
If you need one path to learn and another path to stay fixed, you have a few valid options.
Option 1: Use Two Layer Instances
Create two separate layers. If they should start with the same weights, copy the weights once and then train only one of them:
Now the branches begin identically, but they are no longer truly shared. That is often the correct design when the training behavior must differ.
Option 2: Stop the Gradient on One Path
If the forward computation should be identical but only one path should contribute gradients, you can stop gradients for that path:
This is not the same as making the layer half-trainable. The shared weights are still a single set of variables, but one branch no longer sends gradient information backward.
Option 3: Use Separate Models for Separate Training Phases
Another common pattern is to freeze the shared layer for one training phase, then unfreeze it later:
Notice the recompilation. In Keras, changing trainable after compilation requires recompiling so the optimizer sees the new trainable-variable set correctly.
Common Pitfalls
The biggest mistake is assuming the trainable flag is attached to the branch rather than the shared layer instance. It is not.
Another issue is copying weights into a second layer and still calling that "shared weights." Once the layers are separate objects, the weights are merely initialized the same way.
Teams also forget to recompile after changing trainable. Without recompilation, training behavior may not match the updated setting.
Summary
- A shared Keras layer has one global
trainablesetting because it is one layer instance. - You cannot make the same shared weights trainable on one branch and frozen on another in the same model.
- Use separate layer instances if the branches truly need different training behavior.
- Use
tf.stop_gradientwhen one path should not contribute gradients. - Recompile the model after changing
trainablefor a new training phase.
Related reading
- Keras shows no Improvements to training speed with GPU partial GPU usage?
- Keras Shuffling dataset while using LSTM
- Keras, sparse matrix issue
- Keras split train test set when using ImageDataGenerator
- Keras taking very long time to make first prediction following model.load
- Keras 'Tensor' object has no attribute 'ndim
- keras tensorboard plot train and validation scalars in a same figure
- Keras Tensorflow - Exception while predicting from multiple threads
.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.