Fix: Transformer encoder - apply LayerNorm after residual

[sonali-kumari1]

This PR fixes the Transformer encoder to properly apply Layer Normalization after the residual connections, in accordance with the specification from the Attention is All You Need" (Vaswani et al., 2017) paper.

Previously, Layer Normalization was incorrectly applied before the residual connection, which violated the design outlined in Section 3.1 of "Attention is All You Need," where it states: "the output of each sub-layer isLayerNorm(x + Sublayer(x))".

This update ensures that Layer Normalization is applied after the residual addition in both the attention and feed-forward sub-layers.

Fixes : #1256

[gemini-code-assist]

Summary of Changes

Hello @sonali-kumari1, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request implements a crucial fix to the Transformer encoder's architecture, specifically addressing the placement of Layer Normalization. The change ensures that Layer Normalization is consistently applied after residual connections, bringing the implementation in line with the foundational 'Attention is All You Need' paper. This improves the correctness and potentially the stability of the Transformer model within the timeseries_classification_transformer.py example.

Highlights

• Transformer Architecture Correction: I've corrected the transformer_encoder function to apply Layer Normalization after the residual connections, aligning with the original 'Attention is All You Need' paper's specification.
• Multi-Head Attention Sub-layer: For the multi-head attention sub-layer, Layer Normalization is now applied to the sum of the attention output and the original input (res = x + inputs), rather than to the attention output (x) before the residual addition.
• Feed-Forward Sub-layer: In the feed-forward sub-layer, the input to the first Conv1D layer now correctly uses the normalized output from the preceding attention block. Additionally, the residual connection is properly calculated (res = x + res), and Layer Normalization is applied to this new residual sum before returning the final output.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request corrects the application of Layer Normalization in the Transformer encoder to follow the Post-LN scheme. There is a bug in the second residual connection for the feed-forward network. The review provides a correction to ensure both sub-layers are implemented correctly according to the paper's specification, and also improves the code's readability.

[sonali-kumari1]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly adjusts the Transformer encoder to apply Layer Normalization after the residual connection, aligning the implementation with the original "Attention is All You Need" paper. The logical fix is sound. My review includes a couple of suggestions to improve code clarity and conciseness by simplifying the data flow and reducing the number of temporary variables in the transformer_encoder function. These changes are aimed at enhancing maintainability.

[sonali-kumari1]

/gemini review

[gemini-code-assist]

Code Review

This pull request fixes the Transformer encoder to correctly apply Layer Normalization after the residual connections, aligning it with the original paper's specification. A suggestion has been provided to refactor the implementation slightly for improved readability.

[sachinprasadhs]

There is no issue with the existing implementation.

The implementation in https://keras.io/examples/timeseries/timeseries_classification_transformer is also correct, it is an improvement to the original paper by having Layer Normalization inside the block so that the training will be stable.

Here is a blog explaining the difference between both the approaches https://sh-tsang.medium.com/review-pre-ln-transformer-on-layer-normalization-in-the-transformer-architecture-b6c91a89e9ab

May be you can just add a note explaining this variation.

[sachinprasadhs]

Thanks, could you please make the changes to .md and .ipynb files as well, you don't have to render all the files, you can just make manual changes to those files.