Misconception of the transformer_encoder in notebook https://github.com/keras-team/keras-io/blob/master/examples/timeseries/timeseries_classification_transformer.py

[YacineKaci]

The normalization layer should be applied after the residual connection, not before. At least this is how it is done in the Transformer community.
So the code:
.
.
x = layers.LayerNormalization(epsilon=1e-6)(x)
res = x + inputs
.
.
x = layers.LayerNormalization(epsilon=1e-6)(x)
return x + res

Should be replaced by:
.
.
x = x + inputs
res = layers.LayerNormalization(epsilon=1e-6)(x)
.
.
x = x + res
return layers.LayerNormalization(epsilon=1e-6)(x)

Kind regards

[sonali-kumari1]

Hi @YacineKaci -

Thanks for reporting this issue. While the original Transformer paper applies LayerNorm after the residual connection (Post-LN), the Keras Transformer example applies LayerNorm before the residual (Pre-LN) for a very specific reason: to improve training stability.

As explained in this blog, the originally designed Post-LN Transformer can lead to large gradients near the output layer, often requiring a warm-up phase to avoid instability. On the other hand, the Pre-LN Transformer leads to better-behaved gradients and simplifies the training process.

Given these advantages, we'll maintain the Pre-LN approach in the Keras Transformer example. Please feel free to close this issue if it resolves your concerns.

[github-actions]

Are you satisfied with the resolution of your issue?
Yes
No