[tfjs-layers] Add a Transformer layer (Multihead Attention)

[ierezell]

As the title says, this aim to add a transformer layer (based on multihead attention) in the built-in tfjs-layers.

Note : This is only one layer, the full architecture with positional embeddings, Encoder Decoder would be built on this.

• Refactor code to use tensor built-in operation to be compatible (like tf.add(a,b) => a.add(b) as other layers because add is not imported)
• Some operations are not avalaible in the tfjs-layers/src/core.ts like matMul, sqrt,softmax,cast, scalar. I don't know the best way to integrate and use them.
• Need the tf.layers.LayerNormalization and tf.layers.Dropout which I guess could be imported from the core (as the matMul).

Please tell me how to use the core functionalities to be fully compatible with the other layers. I tried it myself as a custom tfjs layer and it seems to be working (so I wanted to make a PR to have it built-in). However it seems a bit slow.

Have a great day.

PS : Sorry for the linting modifications, just pay attention to the new layer at the bottom of the file.
I also commented the shape for each variable, but it will be removed, it just help for reviewing and debuging.

---

This change is 

[google-cla]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[ierezell]

@googlebot I signed it!

[caisq]

Quick comment: TransformerLayer doesn't seem to be in the standard TensorFlow (Python) API. We only add layer types that are already in TensorFlow (Python) to the TF.js layers API.

[ierezell]

Hi @caisq, thanks for taking time to read my PR !

Would it be better if I respect the layer named MultiHeadAttention (which does the same thing) ?

It would only be minor changes.

However, I'm still not sure how to access the layers I need to make it build in tfjs.

Thanks for the feedback

[caisq]

Yes, implementing MultiHeadAttention will be good. Please note that we generally strive for high fidelity in replicating the behavior of Python. The implementation of MultiHeadAttention in Python uses the einsum op, which has been recently added to tfjs-core. But note that because the gradient of the einsum op hasn't been implemented in tfjs-core yet, the first implementation of MultiHeadAttention in tfjs most likely won't be able to support training.

[ierezell]

Hi @caisq, thanks a lot for the hindsights.

Could we start with Dense only and then create another PR when einsum will be stable for training ? Or is the python fidelity more important ?

In any case I will start to change the code to match the MultiHeadAttention behaviour.

Note that I'm doing this on my free time so I will comment here when modifications are made.

Have a great day.

[caisq]

Not sure what you mean by "start with Dense". Dense is already implemented in TF.js.

I still think there is value in adding MultiHeadAttention to TF.js at this time, because there is no specialized code for training in the code for that layer. So when einsum gradient is supported in tfjs-core in the future, training for that layer should just work. Before the gradient is supported, the layer can still support inference.

[ierezell]

Sorry if I wasn't clear, I meant using classic fully connected layers to create the MultiheadAttention so we can train it and then refactor when einsum gradient will be available.

Then I understand that you still prefer to do everything with einsum from now on even if we cannot train it.

So I will change the code to use einsum, do you have any ETA on when the gradient will be available ?

[caisq]

Sorry for the delay in response. But my take is to use the einsum implementation. This will involve less total amount of work. I don't have an ETA for the gradient of einsum yet. But einsum grad is not terribly hard to write.

[ierezell]

I replaced the K.dot(a,b, bias) with add(eisum('eq', a,b)) as required.
However I don't know if there is an equation to do the dot and the sum at once.

To be honest I find harder to read einsum than plain methods like .add or .dot why is this better ?

I also modified things to be more close to the python code (still not the same and I would like your thoughts about it).

The only pain point is the layerNormalisation, I don't know how to get it and apply it in the tfjs-layers/core.ts as it's another layer...

Then just back and forth to ensure everything is correct and match your guidelines/API requirements.

Note : I don't know which formater you're using and I'm sorry that the linting is so different.

Thanks again for your time on this PR.

Have a great day.

[ierezell]

Hello there,
Sorry to bump but is there any news on merging this PR to have MultiHead attention in JS ?

I'm willing to continue to work on it but I don't know what to change to make it "google standards compatible".

Thanks,
have a great day

[jaspermolgvits]

Hey @ierezell! Thank you for pushing this issue, I'm sure many people are waiting for this to be added to the API.
Would you happen to have a working fork of TFJS that people interested in transformers could use, while Google figures itself out? Maybe if people start using your fork, it might expedite their attention (layers), too.

[ierezell]

Hi @jaspermolgvits, I did this work as I needed Transformers in JS a moment ago...so my fork was totally usable.

I just rebased with the latest master and I will re-check this code and provide an example here (or in my forked readme).
I guess many things have changed in almost one year.

Thanks to look at this PR after so long. Please tell me if anything isn't in accord with your guidelines.

Have a great week.

[ierezell]

I'm sorry I had to recreate another PR as this one was too old and the history/changes were too big. It was easier for me just to fork again and create a new PR.

The new PR is here: #6212
it still needs a bit of love and insights from someone knowing the best practices of tfjs and I will be happy to put a few more hours to close it. (for example the best way to add missing ops and linting)

Have a great day.